import gradio as gr
from transformers import pipeline
from PIL import Image

# Load model
classifier = pipeline(
    "image-classification",
    model="Nav772/vit-food-classifier"
)

def classify_food(image):
    if image is None:
        return "Please upload an image."
    
    try:
        results = classifier(image)
        
        # Format output
        output_lines = []
        for r in results[:5]:  # Top 5 predictions
            label = r["label"]
            score = r["score"]
            bar = "█" * int(score * 20)
            output_lines.append(f"{label}: {score:.1%} {bar}")
        
        return "\n".join(output_lines)
    except Exception as e:
        return f"Error processing image: {str(e)}"

demo = gr.Interface(
    fn=classify_food,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(label="Predictions", lines=6),
    title="🍕 Food Image Classifier",
    description="Upload an image of food and the model will predict what it is. Trained on 10 categories: pizza, sushi, hamburger, ice cream, steak, baklava, cheesecake, pancakes, tacos, and ramen.",
    theme="soft",
    flagging_mode="never"
)

demo.launch()