SivaResearch's picture
Added Files
062f571
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
# Inputs
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")
# Output
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()