update

app.py

@@ -48,17 +48,16 @@ def inference(img):
     prediction = model(image)
     prediction = F.softmax(prediction, dim=1).flatten()
 
-    # return {imagenet_id_to_classname[str(i)]: float(prediction[i]) for i in range(1000)}
+    # 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"
 
-    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
+    # return res
 
+    return {imagenet_id_to_classname[str(i)]: float(prediction[i]) for i in range(1000)}
 
 def set_example_image(example: list) -> dict:
     return gr.Image.update(value=example[0])
@@ -77,11 +76,11 @@ with demo:
         with gr.Row():
                 with gr.Column():
                     with gr.Row():
-                        input_image = gr.Image(label='Input Image', type='file')
+                        input_image = gr.Image(label='Input Image', type='numpy')
                     with gr.Row():
                         submit_button = gr.Button('Submit')
                 with gr.Column():
-                    label = gr.Label()
+                    label = gr.Label(num_top_classes=5)
         with gr.Row():
             example_images = gr.Dataset(components=[input_image], samples=[['library.jpeg'], ['cat.png'], ['dog.png'], ['panda.png']])