|
from transformers import pipeline |
|
import gradio as gr |
|
|
|
|
|
clip_models = [ |
|
"zer0int/CLIP-GmP-ViT-L-14", |
|
"openai/clip-vit-large-patch14", |
|
"laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", |
|
] |
|
|
|
clip_checkpoint = clip_models[0] |
|
clip_detector = pipeline(model=clip_checkpoint, task="zero-shot-image-classification") |
|
|
|
def postprocess(output): |
|
return {out["label"]: float(out["score"]) for out in output} |
|
|
|
|
|
def infer(image, candidate_labels): |
|
candidate_labels = [label.lstrip(" ") for label in candidate_labels.split(",")] |
|
clip_out = clip_detector(image, candidate_labels=candidate_labels) |
|
return postprocess(clip_out) |
|
|
|
|
|
def load_clip_model(modelname): |
|
global clip_detector |
|
try: |
|
clip_detector = pipeline(model=modelname, task="zero-shot-image-classification") |
|
except Exception as e: |
|
raise gr.Error(f"Model load error: {modelname} {e}") |
|
return modelname |
|
|
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown("# Test CLIP") |
|
with gr.Row(): |
|
with gr.Column(): |
|
image_input = gr.Image(type="pil") |
|
text_input = gr.Textbox(label="Input a list of labels") |
|
model_input = gr.Dropdown(label="CLIP model", choices=clip_models, value=clip_models[0], allow_custom_value=True, interactive=True) |
|
run_button = gr.Button("Run", visible=True) |
|
|
|
with gr.Column(): |
|
clip_output = gr.Label(label = "CLIP Output", num_top_classes=3) |
|
|
|
examples = [["./baklava.jpg", "baklava, souffle, tiramisu"], ["./cheetah.jpg", "cat, dog"], ["./cat.png", "cat, dog"]] |
|
gr.Examples( |
|
examples = examples, |
|
inputs=[image_input, text_input], |
|
outputs=[clip_output], |
|
fn=infer, |
|
cache_examples=True |
|
) |
|
run_button.click(fn=infer, |
|
inputs=[image_input, text_input], |
|
outputs=[clip_output]) |
|
model_input.change(load_clip_model, [model_input], [model_input]) |
|
|
|
demo.launch() |