# app.py

from transformers import pipeline
import gradio as gr

# Load the image classification pipeline with the ViT model
classifier = pipeline("image-classification", model="google/vit-base-patch16-224")

# Define the prediction function
def classify_image(img):
    results = classifier(img)
    # Format the results as a dictionary: {label: score}
    return {res['label']: round(res['score'], 4) for res in results}

# Create the Gradio interface
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=5),
    title="Image Classifier",
    description="Upload an image and see the top 5 predicted labels using ViT (google/vit-base-patch16-224)."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()