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

# Use ResNet-50 model (1000 common ImageNet categories like dog, cat, car, etc.)
classifier = pipeline("image-classification", model="microsoft/resnet-50")

def classify_image(img, top_k=3):
    """
    Takes an uploaded image, runs classification, 
    and returns the top_k labels with confidence scores.
    """
    if img is None:
        return {"Error": 1.0}
    
    results = classifier(img, top_k=top_k)
    return {r["label"]: float(r["score"]) for r in results}

# Gradio interface
demo = gr.Interface(
    fn=classify_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(1, 5, value=3, step=1, label="Top K Predictions")
    ],
    outputs=gr.Label(num_top_classes=5, label="Predictions"),
    title="Image Classification App",
    description="Upload an image and the model will predict the top objects in it."
)

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