hudaakram commited on
Commit
97039ed
·
verified ·
1 Parent(s): f11f73b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +44 -0
app.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import time
2
+ from PIL import Image
3
+ import gradio as gr
4
+ from transformers import pipeline
5
+
6
+ MODEL_MAP = {
7
+ "ViT (Base/16, 224)": "google/vit-base-patch16-224",
8
+ "ResNet-50": "microsoft/resnet-50",
9
+ "EfficientNet-B0": "google/efficientnet-b0"
10
+ }
11
+
12
+ # Lazy-load to keep startup fast
13
+ _pipes = {}
14
+ def get_pipe(model_id: str):
15
+ if model_id not in _pipes:
16
+ _pipes[model_id] = pipeline("image-classification", model=model_id, top_k=5)
17
+ return _pipes[model_id]
18
+
19
+ def predict(img: Image.Image, model_name: str):
20
+ if img is None:
21
+ return "Upload an image.", None
22
+ model_id = MODEL_MAP[model_name]
23
+ pipe = get_pipe(model_id)
24
+ t0 = time.time()
25
+ preds = pipe(img)
26
+ latency_ms = int((time.time() - t0) * 1000)
27
+ # Clean top-k dict for Gradio Label
28
+ scores = {p["label"]: round(float(p["score"]), 3) for p in preds}
29
+ return scores, f"{model_name} • ~{latency_ms} ms"
30
+
31
+ with gr.Blocks(title="Image Classifier – Multi-Model") as demo:
32
+ gr.Markdown("# 🐶🐱 Image Classifier (Multi-Model)\nUpload an image, choose a backbone, see top-5 predictions.")
33
+ with gr.Row():
34
+ with gr.Column():
35
+ img = gr.Image(type="pil", label="Image")
36
+ model = gr.Dropdown(list(MODEL_MAP.keys()), value="ViT (Base/16, 224)", label="Backbone")
37
+ btn = gr.Button("Predict")
38
+ with gr.Column():
39
+ out = gr.Label(label="Top-5")
40
+ info = gr.Markdown()
41
+ btn.click(fn=predict, inputs=[img, model], outputs=[out, info])
42
+
43
+ if __name__ == "__main__":
44
+ demo.launch()