Update app.py

app.py

@@ -34,9 +34,14 @@ def predict(image):
         inputs = feature_extractor(images=image, return_tensors="pt")
         outputs = model(**inputs)
         logits = outputs.logits
-        predicted_class_idx = logits.argmax(-1).item()
+        probs = torch.nn.functional.softmax(logits, dim=-1)
+        top_probs, top_idxs = probs.topk(3, dim=-1)
+        top_probs = top_probs.detach().numpy()[0]
+        top_idxs = top_idxs.detach().numpy()[0]
+        top_classes = [model.config.id2label[idx] for idx in top_idxs]
+        result = {top_classes[i]: float(top_probs[i]) for i in range(3)}
         logging.info("Prediction successful.")
-        return model.config.id2label[predicted_class_idx]
+        return result
     except Exception as e:
         logging.error("Error during prediction: %s", e)
         return str(e)
@@ -45,7 +50,7 @@ def predict(image):
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Sketchpad(),
-    outputs=gr.Label(),
+    outputs=gr.JSON(),
     title="Drawing Classifier",
     description="Draw something and the model will try to identify it!"
 )