Update app.py

app.py

@@ -14,14 +14,21 @@ def predict(image):
     with torch.no_grad():
         outputs = model(**inputs).logits
     
-    predicted_class_idx = outputs.argmax(-1).item()
+    # predicted_class_idx = outputs.argmax(-1).item()
+    # predicted_class = class_names[predicted_class_idx]
+    
+    # return predicted_class
+
+    probabilities = F.softmax(outputs, dim=1)
+    confidence_score = probabilities[0][predicted_class_idx].item()
+    predicted_class_idx = probabilities.argmax(-1).item()
     predicted_class = class_names[predicted_class_idx]
     
-    return predicted_class
+    return predicted_class, confidence_score
 
 iface = gr.Interface(fn=predict, 
                      inputs=gr.Image(type="pil"), 
-                     outputs=gr.Label(num_top_classes=1),
+                     outputs=[gr.Label(num_top_classes=1),gr.Label()]
                      title="Hateful Content Detection",
                      description="Upload an image to classify hateful gestures or symbols")