hudaakram's picture
Upload folder using huggingface_hub
35aee79 verified
raw
history blame contribute delete
535 Bytes
import gradio as gr
from transformers import pipeline
clf = pipeline("image-classification",
model="google/vit-base-patch16-224",
device_map="auto")
def predict(img):
out = clf(img)[:5]
return {o['label']: float(o['score']) for o in out}
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=5),
title="Image Classification (ViT)"
)
if __name__ == "__main__":
demo.launch(share=True) # share=True gives you a temporary Colab link