AXL14/ZeroShotClassificationDemo

1import gradio as gr2import numpy as np3from PIL import Image4from transformers import pipeline5 6# Create the pipeline object7pipe = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32")8 9# Define the function that will be used by the interface10def zero_shot_classification(image, labels_text):11    # Convert image to a PIL image object12    pil_image = Image.fromarray(np.uint8(image)).convert("RGB")13    14    # Split the labels text into a list of labels15    labels = labels_text.split(",")16    17    # Use the pipeline to classify the image with the given labels18    res = pipe(19        images=pil_image, 20        candidate_labels=labels,21        hypothesis_template= "This is a photo of a {}"22    )23    24    # Return a dictionary mapping labels to scores25    return {dic["label"]: dic["score"] for dic in res}26    27# Create the interface28iface = gr.Interface(29    zero_shot_classification, 30    ["image", "text"], 31    "label", 32    examples=[33        ["dog.jpg", "dog,cat,horse,zebra"],34    ],35    description="Please add a picture and a list of labels separated by commas to see the zero-shot classification capabilities",36    title="Zero-shot Image Classification"37)38 39# Launch the interface40iface.launch()