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)
    logits_per_image = outputs.logits_per_image  # this is the image-text similarity score
    probs = logits_per_image.softmax(dim=1) 
    probabilities_percentages = ', '.join(['{:.2f}%'.format(prob.item() * 100) for prob in probs[0]])
    return probabilities_percentages

title = "TSAI S18 Assignment: Use a pretrained CLIP model and give a demo on its workig"
description = "A simple Gradio interface that accepts an image and some captions, and gives a score as to how much the caption describes the image "
 
examples = [["cats.jpg","a photo of a cat, a photo of a dog"],
            ["personBicycle.jpg","person riding bicycle, person driving car, photo of a dog"]
           ]

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.Textbox(label="Probability score of captions")],
    title = title,
    description = description,
    examples = examples,
)

demo.launch()