Update app.py

app.py

@@ -5,7 +5,6 @@ import torch
 import torch.nn.functional as F
 import torchvision.transforms as T
 from uniformer import uniformer_small
-from imagenet_class_index import imagenet_classnames
 
 import gradio as gr
 
@@ -22,6 +21,11 @@ model.load_state_dict(state_dict['model'])
 model = model.to(device)
 model = model.eval()
 
+os.system("wget https://huggingface.co/Andy1621/uniformer/blob/main/imagenet_class_index.json")
+
+with open("imagenet_class_index.json", "r") as f:
+    imagenet_classnames = json.load(f)
+
 # Create an id to label name mapping
 imagenet_id_to_classname = {}
 for k, v in imagenet_classnames.items():
@@ -44,7 +48,7 @@ def inference(img):
     # The model expects inputs of shape: B x C x T x H x W
     image = image[None, :, None, ...]
     
-    prediction = model(image, input_type="image")
+    prediction = model(image)
     prediction = F.softmax(prediction, dim=1)
     pred_classes = prediction.topk(k=5).indices