Add application file

app.py

@@ -1,32 +1,35 @@
 import requests
-
 import gradio as gr
 import torch
 from timm import create_model
 from timm.data import resolve_data_config
 from timm.data.transforms_factory import create_transform
+from PIL import Image
 
 IMAGENET_1k_URL = "https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt"
-LABELS = requests.get(IMAGENET_1k_URL).text.strip().split('\n')
-
-model = create_model('resnet50', pretrained=True)
+LABELS = requests.get(IMAGENET_1k_URL, timeout=10).text.strip().split("\n")
 
-transform = create_transform(
-    **resolve_data_config({}, model=model)
-)
-model.eval()
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def predict_fn(img):
-    img = img.convert('RGB')
-    img = transform(img).unsqueeze(0)
+model = create_model("resnet50", pretrained=True).to(device).eval()
+config = resolve_data_config({}, model=model)
+transform = create_transform(**config)
 
+def predict_fn(img: Image.Image):
+    img = img.convert("RGB")
+    x = transform(img).unsqueeze(0).to(device)
     with torch.no_grad():
-        out = model(img)
-    
-    probabilites = torch.nn.functional.softmax(out[0], dim=0)
-
-    values, indices = torch.topk(probabilites, k=5)
-
-    return {LABELS[i]: v.item() for i, v in zip(indices, values)}
+        out = model(x)[0]
+        probs = torch.nn.functional.softmax(out, dim=0)
+        values, indices = torch.topk(probs, k=5)
+    return {LABELS[i.item()]: v.item() for i, v in zip(indices, values)}
+
+demo = gr.Interface(
+    fn=predict_fn,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=5),
+    title="ResNet-50 (timm) · ImageNet Top-5"
+)
 
-gr.Interface(predict_fn, gr.inputs.Image(type='pil'), outputs='label').launch()
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -1,2 +1,2 @@
 timm
-gradio==3.50.2
+gradio