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

# Load zero-shot image classification
checkpoint = "openai/clip-vit-large-patch14"
detector = pipeline(model=checkpoint, task="zero-shot-image-classification")

# Inference function
# def classify_image(image):
#     labels = ["cat", "dog"]
#     image = image.convert("RGB")
#     results = detector(image, candidate_labels=labels)
#     return {res["label"]: round(res["score"], 3) for res in results}

def classify_image(image):
    try:
        labels = ["cat", "dog"]
        image = image.convert("RGB")
        results = detector(image, candidate_labels=labels)
        return {res["label"]: round(res["score"], 3) for res in results}
    except Exception as e:
        return {"error": str(e)}


# Gradio interface
iface = gr.Interface(fn=classify_image, inputs=gr.Image(type="pil"), outputs=gr.Label())
iface.launch()