import numpy as np
import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image



processor = AutoImageProcessor.from_pretrained("microsoft/swin-tiny-patch4-window7-224")
model = AutoModelForImageClassification.from_pretrained("microsoft/swin-tiny-patch4-window7-224")

def classifier(image):
    image = Image.open(image.raw)
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    # model predicts one of the 1000 ImageNet classes
    predicted_class_idx = logits.argmax(-1).item()
    return model.config.id2label[predicted_class_idx]


food = gr.Interface(
    fn=classifier,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title = "what's your eating?",
    description = "A simple model for food classification"
)

food.launch()