import gradio as gr
from transformers import pipeline
from datasets import load_dataset

# This "links" your project to a specific dataset on Hugging Face
dataset = load_dataset("huggingface/cats-image", split="test") 

# Now you can use 'dataset[0]' in your code!

# Load the pre-trained Image Classification pipeline
classifier = pipeline("image-classification", model="google/vit-base-patch16-224")

def predict(image):
    # Get predictions from the model
    results = classifier(image)
    # Reformat for Gradio's Label component: {"Label": Score}
    return {res["label"]: res["score"] for res in results}

# Define the Gradio Interface
demo = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="pil"), 
    outputs=gr.Label(num_top_classes=3),
    title="AI Image Classifier",
    description="Upload any image to see what the Vision Transformer thinks it is!"
)

if __name__ == "__main__":
    demo.launch()