paschalc24
Cleaned Up App and Added Requirements
570913a
raw
history blame contribute delete
679 Bytes
import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
image_model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224")
extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
image_pipe = pipeline("image-classification", model=image_model, feature_extractor=extractor)
# classify the image and returns the results
def classify_image(inp):
results = image_pipe(inp)
return {result["label"]: result["score"] for result in results}
gr.Interface(fn=classify_image,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=5)).launch()