Update app.py

app.py

@@ -207,11 +207,13 @@ transform = transforms.Compose([
 ])
 
 def predict(img: Image.Image):
-    img = transform(img).unsqueeze(0).to(device)
+    img_t = transform(img).unsqueeze(0).to(device)
     with torch.no_grad():
-        out = model(img)
-        pred = out.argmax(1).item()
-    return classes[pred]
+        out = model(img_t)
+        probs = torch.softmax(out, dim=1)[0]
+        top5 = probs.topk(5)
+    result = {classes[i]: float(probs[i]) for i in top5.indices}
+    return result
 
 # ------------------------
 # Gradio interface