beaunix's picture
set img_classififer demo to hub
40b3df0 verified
Raw
History Blame Contribute Delete
1.61 kB
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()