|
import gradio as gr |
|
from transformers import pipeline |
|
|
|
model_names = [ |
|
"apple/mobilevit-small", |
|
"facebook/deit-base-patch16-224", |
|
"facebook/convnext-base-224", |
|
"google/vit-base-patch16-224", |
|
"google/mobilenet_v2_1.4_224", |
|
"microsoft/resnet-50", |
|
"microsoft/swin-base-patch4-window7-224", |
|
"microsoft/beit-base-patch16-224", |
|
"nvidia/mit-b0", |
|
"shi-labs/nat-base-in1k-224", |
|
"shi-labs/dinat-base-in1k-224", |
|
] |
|
|
|
|
|
def process(image_file, top_k): |
|
labels = [] |
|
for m in model_names: |
|
p = pipeline("image-classification", model=m) |
|
pred = p(image_file) |
|
labels.append({x["label"]: x["score"] for x in pred[:top_k]}) |
|
return labels |
|
|
|
|
|
|
|
image = gr.Image(type="filepath", label="Upload an image") |
|
top_k = gr.Slider(minimum=1, maximum=5, step=1, value=5, label="Top k classes") |
|
|
|
|
|
labels = [gr.Label(label=m) for m in model_names] |
|
|
|
description = "This Space lets you quickly compare the most popular image classifiers available on the hub, including the recent NAT and DINAT models. All of them have been fine-tuned on the ImageNet-1k dataset. Anecdotally, the three sample images have been generated with a Stable Diffusion model :)" |
|
|
|
iface = gr.Interface( |
|
theme="huggingface", |
|
description=description, |
|
layout="horizontal", |
|
fn=process, |
|
inputs=[image, top_k], |
|
outputs=labels, |
|
examples=[ |
|
["bike.jpg", 5], |
|
["car.jpg", 5], |
|
["food.jpg", 5], |
|
], |
|
allow_flagging="never", |
|
) |
|
|
|
iface.launch() |
|
|
