import gradio as gr
from transformers import pipeline

# Initialize the pipeline
pipe = pipeline(
    "image-classification",
    model="ariG23498/vit_base_patch16_224.augreg2_in21k_ft_in1k.ft_food101"
)

# Function for classification
def classify(image):
    return pipe(image)[0]["label"]

# Gradio Interface with a detailed description
demo = gr.Interface(
    fn=classify,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Textbox(label="Predicted Label"),
    examples=[["./sushi.png", "sushi"]],
    title="Food Classification with ViT 🥗🍣",
    description=(
        "### Explore Food Classification with Vision Transformers (ViT) 🔍\n\n"
        "This application demonstrates the power of Vision Transformers (ViT) for food classification tasks, "
        "leveraging the pre-trained model `vit_base_patch16_224.augreg2_in21k_ft_in1k.ft_food101` fine-tuned on the Food-101 dataset. "
        "With just a few lines of code, you can integrate state-of-the-art image classification models using the Hugging Face `pipeline` API.\n\n"
        "#### How to Use:\n"
        "1. Upload an image of food (e.g., sushi, pizza, or burgers).\n"
        "2. The model will classify the image and provide the predicted label.\n"
        "3. Try the provided example for a quick start or test your own food images!\n\n"
        "#### About the Model:\n"
        "- **Model Name**: `vit_base_patch16_224.augreg2_in21k_ft_in1k.ft_food101`\n"
        "- **Dataset**: [Food-101](https://www.kaggle.com/dansbecker/food-101)\n"
        "- **Architecture**: Vision Transformers (ViT), which process images by splitting them into patches and leveraging self-attention for feature extraction.\n\n"
        "#### Learn More:\n"
        "Discover more about Vision Transformers in the [Hugging Face blog](https://huggingface.co/blog). "
        "Explore the Food-101 dataset [here](https://www.kaggle.com/dansbecker/food-101)."
    )
)

demo.launch()