update

app.py

@@ -48,7 +48,20 @@ 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)}
+    # 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"
+
+    return res
+
+
+def set_example_image(example: list) -> dict:
+    return gr.Image.update(value=example[0])
 
 
 demo = gr.Blocks()
@@ -79,8 +92,9 @@ with demo:
     )
 
     submit_button.click(fn=inference, inputs=input_image, outputs=label)
+    example_images.click(fn=set_example_image, inputs=example_images, outputs=example_images.components)
 
-demo.launch(enable_queue=True)
+demo.launch(enable_queue=True, cache_examples=True)