from transformers import pipeline
import gradio as gr
from pathlib import Path

examples = Path('./examples').glob('*')
examples = list(map(str,examples))


pipe = pipeline("image-classification", model="shreydan/vit-base-oxford-iiit-pets")

def predict(inp_path):
  confidences = pipe(inp_path)
  confidences = {s['label']:s['score'] for s in confidences}
  return confidences


gr.Interface(fn=predict,
             inputs=gr.Image(type="filepath"),
             outputs=gr.Label(num_top_classes=3),
             examples=examples).queue().launch()