Update app.py

app.py

@@ -27,6 +27,7 @@ imagenet_id_to_classname = {}
 for k, v in imagenet_classnames.items():
     imagenet_id_to_classname[k] = v[1] 
 
+
 def inference(img):
     image = img
     image_transform = T.Compose(
@@ -39,28 +40,24 @@ def inference(img):
     )
     image = image_transform(image)
     
-    # The model expects inputs of shape: B x C x T x H x W
+    # The model expects inputs of shape: B x C x H x W
     image = image.unsqueeze(0)
     
     prediction = model(image)
-    prediction = F.softmax(prediction, dim=1)
+    prediction = F.softmax(prediction, dim=1).flatten()
 
-    return {imagenet_id_to_classname[str(i)]: float(prediction[0][i]) for i in range(1000)}
+    return {imagenet_id_to_classname[str(i)]: float(prediction[i]) for i in range(1000)}
     
+
 inputs = gr.inputs.Image(type='pil')
-# outputs = gr.outputs.Textbox(label="Output")
 label = gr.outputs.Label(num_top_classes=5)
 
 title = "UniFormer-S"
-
 description = "Gradio demo for UniFormer: To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
-
 article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2201.09450' target='_blank'>UniFormer: Unifying Convolution and Self-attention for Visual Recognition</a> | <a href='https://github.com/Sense-X/UniFormer' target='_blank'>Github Repo</a></p>"
 
-
 gr.Interface(
     inference, inputs, outputs=label, 
-    title=title, description=description, 
-    article=article, 
+    title=title, description=description, article=article, 
     examples=[['library.jpeg'], ['cat.png'], ['dog.png'], ['panda.png']]
     ).launch(enable_queue=True, cache_examples=True)