import gradio as gr
from transformers import CLIPProcessor, CLIPModel

model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")


def inference(input_img, captions):
    captions_list = captions.split(",")
    inputs = processor(text=captions_list, images=input_img, return_tensors="pt", padding=True)
    outputs = model(**inputs)
    # this is the image-text similarity score
    logits_per_image = outputs.logits_per_image  
    probs = logits_per_image.softmax(dim=1).tolist()[0]
    confidences = {captions_list[i][:30]: probs[i] for i in range(len(probs))}
    return confidences

title = "CLIP Inference: Application using a pretrained CLIP model"
description = "An application using Gradio interface that accepts an image and some captions, and displays a probability score with which each caption describes the image "

examples = [["example_images/12863.jpg","photo of water, a photo of pizza, photo of a smiling lady, photo of luggage, photo of bank"],
            ["example_images/12659.jpg","person riding bicycle, person driving car, photo of traffic lights, photo of vehicle, photo of light"],
            ["example_images/12291.jpg","photo of a cat, a photo of a cat sleeping on keyboard, photo of desktop monitor, photo of typing keyboard, photo of computer mouse, photo of water glass"],
            ["example_images/12272.jpg","person playing base ball, person holding bat, photo of a net, photo of audience behind net, photo of currency"],
            ["example_images/9309.jpg","a photo of cows, a photo of grass, group of cows grazing grass, photo of electric pole, photo of trees"],
            ["example_images/3805.jpg","a photo of water, Zebras drinking water, photo of a bird swimming, photo of grass"],
            ["example_images/2788.jpg","a photo of a man dropping pigeon feed, a photo of pigeons, photo of a man feeding pigeons, photo of water, people walking, a photo of old building in the background"]
           ]

demo = gr.Interface(
    inference,
    inputs = [gr.Image(shape=(416, 416), label="Input Image"), 
              gr.Textbox(placeholder="Enter different captions for image, separated by comma")],
    outputs = [gr.Label()],
    title = title,
    description = description,
    examples = examples,
)
demo.launch(debug=True)