import gradio as gr
from transformers import AutoModelForImageClassification, AutoProcessor
import torch
from PIL import Image

# Load your model and processor outside the function to avoid reloading them on each function call
model_name = "Khadidja22/my_awesome_food_model"
model = AutoModelForImageClassification.from_pretrained(model_name)
processor = AutoProcessor.from_pretrained(model_name)

def classify_image(uploaded_image):
    # Process the uploaded image
    inputs = processor(images=uploaded_image, return_tensors="pt")
    
    # Predict
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
    
    # Get the highest probability label
    predicted_label_idx = logits.argmax(-1).item()
    predicted_label = model.config.id2label[predicted_label_idx]
    
    return predicted_label

iface = gr.Interface(fn=classify_image,
                     inputs=gr.Image(),
                     outputs="text",
                     title="Food Classification",
                     description="Upload an image of food, and the model will classify it.")

iface.launch()