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

model_id = "prithivMLmods/Food-101-93M"
processor = AutoImageProcessor.from_pretrained(model_id)
model = AutoModelForImageClassification.from_pretrained(model_id)

def predict_food(image: Image.Image):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
    topk = torch.topk(probs, k=5)
    labels = [model.config.id2label[i.item()] for i in topk.indices]
    scores = [round((p * 100).item(), 2) for p in topk.values]
    return "\n".join(f"{lbl}: {sc}%" for lbl, sc in zip(labels, scores))

gr.Interface(
    fn=predict_food,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="🍽️ Food‑101 Food Classifier",
    description="Upload food image, outputs the top 5 dish categories."
).launch()