import gradio as gr
from transformers import pipeline

# 1. Choose a Hugging Face image classification model
#    You can change this later (for digits, beans, general images etc.)
MODEL_ID = "nateraw/vit-base-beans"  # demo model trained on 3 types of beans

# 2. Create a pipeline (this downloads and loads the model)
classifier = pipeline("image-classification", model=MODEL_ID)

# 3. Define a prediction function Gradio will use
def predict(image):
    """
    image: PIL.Image (Gradio handles conversion)
    returns: dict {label: score} for the top classes
    """
    results = classifier(image)

    # Gradio's Label component likes a dict: {label: confidence}
    output = {r["label"]: float(r["score"]) for r in results}
    return output

# 4. Build the Gradio Interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload or draw an image"),
    outputs=gr.Label(num_top_classes=3, label="Predictions"),
    title="Vision Transformer Image Classifier",
    description=(
        "Upload a small image and this app will classify it using a pretrained "
        "Vision Transformer model from Hugging Face."
    ),
    examples=[],  # you can add example images later
)

# 5. Run the app (Spaces will call launch() automatically)
if __name__ == "__main__":
    demo.launch()