Upload 2 files

app.py

@@ -0,0 +1,26 @@
+import gradio as gr
+from PIL import Image
+from typing import List, Dict
+
+import torch
+
+class GradioApp:
+    def __init__(self) -> None:
+        self.model = torch.load('pretrained_vit.pth', map_location='cpu')
+    def predict(self, img_file: str, classes: List[str]) -> Dict[str, float]:
+        classes = ['0', '1', '2']
+        img = self.model.val_transform(Image.open(img_file)).unsqueeze(0)
+        with torch.inference_mode():
+            preds = torch.softmax(self.model(img), dim=1)[0].cpu().numpy()
+        return {classes[i] : preds[i] for i in range(len(classes))}
+    def launch(self):
+        demo = gr.Interface(
+            fn=self.predict,
+            inputs=gr.Image(type='filepath'),
+            outputs=gr.Label(num_top_classes=3),
+        )
+        demo.launch()
+
+if __name__ == '__main__':
+    app = GradioApp()
+    app.launch()

pretrained_vit.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:35925a38eead26593bb6b028900224ef80fe380334cb3c33cb4932af84f5f221
+size 346305720