debug

app.py

@@ -1,24 +1,21 @@
 import os
 import json
-from typing import List
 
 import torch
 import torch.nn.functional as F
 import torchvision.transforms as T
-from PIL import Image
-from torchvision.transforms._transforms_video import NormalizeVideo
-from uniformer import uniformer_small, uniformer_base, uniformer_small_plus, uniformer_base_ls
+from uniformer import uniformer_small
 
 import gradio as gr
 
 # Device on which to run the model
 # Set to cuda to load on GPU
 device = "cpu"
-os.system("https://huggingface.co/Andy1621/uniformer/blob/main/uniformer_small_in1k.pth")
+os.system("wget https://huggingface.co/Andy1621/uniformer/blob/main/uniformer_small_in1k.pth")
 # Pick a pretrained model 
 model = uniformer_small()
 state_dict = torch.load('uniformer_small_in1k.pth', map_location='cpu')
-model.load_state_dict(state_dict)
+model.load_state_dict(state_dict['model'])
 
 # Set to eval mode and move to desired device
 model = model.to(device)
@@ -56,6 +53,7 @@ def inference(img):
     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 = []
     return "Top 5 predicted labels: %s" % ", ".join(pred_class_names)
     
 inputs = gr.inputs.Image(type='pil')