Update app.py

app.py

@@ -16,6 +16,14 @@ import gradio as gr
 from PIL import Image
 from torchvision import transforms
 
+import json
+
+
+with open("in_cls_idx.json", "r") as f:
+    imagenet_id_to_name = {int(cls_id): name for cls_id, (label, name) in json.load(f).items()}
+
+
+
 
 def load_image(image_path):
     return Image.open(image_path).convert("RGB")
@@ -48,6 +56,7 @@ def visualize_and_predict(model, resolution, image_path):
         _, preds = model(image).topk(5)
     # convert preds to a Python list and remove the batch dimension
     preds = preds.tolist()[0]
+
     return preds
     
 os.system("wget https://github.com/pytorch/hub/raw/master/images/dog.jpg -O dog.jpg")
@@ -56,8 +65,7 @@ os.system("wget https://github.com/pytorch/hub/raw/master/images/dog.jpg -O dog.
 
 def inference(img):
     preds = visualize_and_predict(model, resolution, img)
-    
-    return preds
+    return [imagenet_id_to_name[cls_id] for cls_id in preds]
 
 inputs = gr.inputs.Image(type='filepath')
 outputs = gr.outputs.Textbox(label="Output")