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()