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


clf = pipeline(
    task="zero-shot-image-classification",
    model="openai/clip-vit-large-patch14",   # 100 % open weights
    device=0 if torch.cuda.is_available() else -1
)

def classify(image: Image.Image, labels: str) -> dict:
    """
    image : PIL image uploaded by user
    labels: comma-separated string, e.g. 'cat,dog,horse'
    """
    if image is None:
        return {}
    candidate_labels = [lbl.strip() for lbl in labels.split(",") if lbl.strip()]
    if not candidate_labels:
        return {}
    preds = clf(image, candidate_labels=candidate_labels)
    # Gradio expects {label: score, ...}
    return {p["label"]: p["score"] for p in preds}



# 3.  Gradio interface
# ------------------------------------------------------------------
demo = gr.Interface(
    fn=classify,
    inputs=[
        gr.Image(type="pil", label="Upload an image"),
        gr.Textbox(
            placeholder="cat, dog, horse",
            label="Candidate labels (comma-separated)",
        ),
    ],
    outputs=gr.Label(num_top_classes=5, label="Scores"),
    title="Zero-Shot Image Classifier 🤗 CLIP",
    description="Type any labels you want—the open-source CLIP model will score them.",
    examples=[["https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg", "cat, dog, fox"]],
)


# 4.  Launch (Spaces sets server_name & port automatically)
# ------------------------------------------------------------------
if __name__ == "__main__":
    demo.launch()