import gradio as gr
from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import requests
hotdog_url = 'https://potatorolls.com/wp-content/uploads/2020/10/Basic-Hot-Dogs-960x640.jpg'
hotdog_image = Image.open(requests.get(hotdog_url, stream=True).raw)
muffin_url = "https://www.recipetineats.com/wp-content/uploads/2023/05/Up-and-go-breakfast-muffins_9.jpg"
muffin_image = Image.open(requests.get(muffin_url, stream=True).raw)
juice_url = "https://recipes.net/wp-content/uploads/2024/01/how-to-drink-fresh-juice-1705739043.jpg"
juice_image = Image.open(requests.get(juice_url, stream=True).raw)

def snacks_classifier(input_image):

    # Init model, transforms
    processor = ViTImageProcessor.from_pretrained('yangswei/snacks_classification')
    model = ViTForImageClassification.from_pretrained('yangswei/snacks_classification')

    # inputs & outputs
    inputs = processor(images=input_image, return_tensors="pt")
    outputs = model(**inputs).logits.softmax(1)
    labels = model.config.id2label
    confidences = {labels[i]: outputs[0][i].item() for i in range(len(labels))}

    return confidences

with gr.Blocks(theme=gr.themes.Base()) as demo:
    gr.Interface(fn=snacks_classifier, inputs="image", outputs=gr.Label(num_top_classes=20, label="Prediction"),
                 examples=[hotdog_image, muffin_image, juice_image])

demo.launch()