Update app.py

app.py

@@ -46,14 +46,15 @@ def inference(img):
     
     prediction = model(image)
     prediction = F.softmax(prediction, dim=1)
-    pred_classes = prediction.topk(k=5).indices
-    
-    pred_class_names = [imagenet_id_to_classname[str(i.item())] for i in pred_classes[0]]
-    pred_class_probs = [prediction[0][i.item()].item() * 100 for i in pred_classes[0]]
-    res = "Top 5 predicted labels:\n"
-    for name, prob in zip(pred_class_names, pred_class_probs):
-        res += f"[{prob:2.2f}%]\t{name}\n"
-    return res
+
+    # pred_classes = prediction.topk(k=5).indices
+    # pred_class_names = [imagenet_id_to_classname[str(i.item())] for i in pred_classes[0]]
+    # pred_class_probs = [prediction[0][i.item()].item() * 100 for i in pred_classes[0]]
+    # res = "Top 5 predicted labels:\n"
+    # for name, prob in zip(pred_class_names, pred_class_probs):
+    #     res += f"[{prob:2.2f}%]\t{name}\n"
+
+    return {imagenet_id_to_classname[str(i)]: float(prediction[i]) for i in range(1000)}
     
 inputs = gr.inputs.Image(type='pil')
 outputs = gr.outputs.Textbox(label="Output")