Update app.py

app.py

@@ -18,9 +18,8 @@ def predict(model_name, image):
     input_tensor = transform(image).unsqueeze(0)
     with torch.no_grad():
         output = model(input_tensor)
-    output_np = output[0].numpy()
-    class_ind = np.argmax(output_np)
-    return class_names[class_ind]
+    output_np = torch.softmax(output)[0].numpy()
+    return {clsname: prob for clsname, prob in zip(class_names, output_np)}
  
 interface = gr.Interface(
     fn=predict,
@@ -28,7 +27,7 @@ interface = gr.Interface(
         gr.Dropdown(label="Select Model", value="hb_former_b36", choices=["hpx_former_s18", "hpx_former_s18_384", "hb_former_s18", "c_hpx_former_s18", "hpx_a_former_s18", "hb_a_former_s18", "hpx_former_b36", "hb_former_b36"]),
         gr.Image(type="pil", label="Upload Image")
     ],
-    outputs=gr.Textbox(label="Predicted Class"),
+    outputs=gr.Label(label="Prediction", num_top_classes=10),
     title="Image Classification",
     description="Choose a model and upload an image to predict the class."
 )