Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 from huggingface_hub import Repository
+import numpy as np
 
 # Retrieve the token from the environment variables
 token = os.environ.get("token")
@@ -65,15 +66,20 @@ def predict(input_img):
     with torch.inference_mode():
         img = cmp(input_img).unsqueeze(0)
         res = modelm(img.to(device))
-        prediction = le.inverse_transform(torch.argmax(res.cpu()).unsqueeze(0).numpy())[0]
-        return prediction
+        probabilities = torch.softmax(res, dim=1).cpu().numpy().flatten()
+        top_10_indices = np.argsort(probabilities)[-10:][::-1]
+        top_10_probabilities = probabilities[top_10_indices]
+        top_10_predictions = le.inverse_transform(top_10_indices)
+        
+        results = {top_10_predictions[i]: float(top_10_probabilities[i]) for i in range(10)}
+        return results
 
 gradio_app = gr.Interface(
     fn=predict,
     inputs=gr.Image(label="Upload an Image", type="pil"),
-    outputs=gr.Label(label="Location"),
+    outputs=gr.JSON(label="Top 10 Predictions with Probabilities"),
     title="Predict the Location of this Image"
 )
 
 if __name__ == "__main__":
-    gradio_app.launch()
+    gradio_app.launch()