Update app.py

app.py

@@ -15,14 +15,13 @@ def classify_image(img):
     with torch.no_grad():
         outputs = model(**inputs)
     logits = outputs.logits
-    predicted_label = logits.argmax(-1).item()
-    label = model.config.id2label[predicted_label]
-    confidence = torch.nn.functional.softmax(logits, dim=-1)[0, predicted_label].item()
-    return f"Prediction: {label}\nConfidence: {confidence:.4f}"
+    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
+    results = {model.config.id2label[i]: float(probs[i]) for i in range(len(probs))}
+    return results
 
 # Create the Gradio interface
 image_input = gr.Image()
-label_output = gr.Label()
+label_output = gr.Label(num_top_classes=3)
 interface = gr.Interface(fn=classify_image, inputs=image_input, outputs=label_output)
 
 # Launch the interface