File size: 679 Bytes
8121a75
570913a
8121a75
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
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()