from transformers import pipeline
from turtle import title
import gradio as gr
from transformers import pipeline
from PIL import Image
import numpy as np

checkpoint = "openai/clip-vit-large-patch14-336"
classifier = pipeline("zero-shot-image-classification",
                      model=checkpoint)


def shot(image, labels_text):
    PIL_image = Image.fromarray(np.uint8(image)).convert('RGB')
    labels = labels_text.split(",")
    res = classifier(images=PIL_image, 
           candidate_labels=labels,
           hypothesis_template= "This is a photo of a {}")
    return {dic["label"]: dic["score"] for dic in res}
    
interface = gr.Interface(shot, 
                     inputs=["image", "text"], 
                     outputs="label",
                    description="Add a picture and a list of labels separated by commas",
                    examples=[["tundra.jpg", "Ford F-150, RAM 1500, Tundra, GMC Sierra, Silverado"],
                              ['interior.jpeg', "Interior, Exterior"],
                              ['carplay.webp', "Carplay, Android play"]
                    ],
                    title="Zero-shot Image Classification")

interface.launch()