import gradio as gr
from transformers import pipeline
import os

MODEL_OPENAI_CLIP_VIT_LARGE_PATCH_14 = "openai/clip-vit-large-patch14"
MODEL_OPENAI_CLIP_VIT_BASE_PATCH_16 = "openai/clip-vit-base-patch16"
MODEL_OPENAI_CLIP_VIT_BASE_PATCH_32 = "openai/clip-vit-base-patch32"

input_examles = [
    [
        os.path.join(os.path.dirname(__file__), "images/example-01-two-cats.jpg"), 
        'the photo contains one cat,'   \
        'the photo contains a animal,'  \
        'the photo contains two cats,'  \
        'the photo contains something else',
        MODEL_OPENAI_CLIP_VIT_LARGE_PATCH_14
    ],
    [
        os.path.join(os.path.dirname(__file__), "images/example-02-two-cats.jpg"), 
        'the photo contains one cat,'   \
        'the photo contains a animal,'  \
        'the photo contains two cats,'  \
        'the photo contains something else',
        MODEL_OPENAI_CLIP_VIT_BASE_PATCH_16
    ],
    [
        os.path.join(os.path.dirname(__file__), "images/example-03-woman-with-pink-shirt-and-a-printer.jpg"), 
        'the photo contains a man and a printer,'                       \
        'the photo contains a woman with blue shirt and a printer,'     \
        'the photo contains a woman and a printer,'                     \
        'the photo contains a printer,'                                 \
        'the photo contains a woman with pink shirt and a printer,'     \
        'the photo contains something else',
        MODEL_OPENAI_CLIP_VIT_LARGE_PATCH_14
    ],
    [
        os.path.join(os.path.dirname(__file__), "images/example-04-printer-in-the-left.jpg"), 
        'the photo contains a printer in the left,'     \
        'the photo contains a printer in the right,'    \
        'the photo contains a printer in the middle'    \
        'the photo contains something else',
        MODEL_OPENAI_CLIP_VIT_LARGE_PATCH_14
    ],
    [
        os.path.join(os.path.dirname(__file__), "images/example-05-car-in-the-left-and-motorbike-in-the-right.jpg"), 
        'the photo contains a car in the left and a motorbike in the right,'     \
        'the photo contains a car in the right and a motorbike in the left',
        MODEL_OPENAI_CLIP_VIT_LARGE_PATCH_14
    ]
]

def launch(image_input, labels_input, dropdown_input):
    
    return predict(image_input, labels_input, dropdown_input)

def predict(image_input, labels_input, dropdown_input):
    
    labels_candidate = labels_input.split(",")
    
    selected_model = f"""{dropdown_input}"""
    pipe = pipeline(task="zero-shot-image-classification",  model=selected_model)
    
    predictions = pipe(image_input, candidate_labels=labels_candidate)
    
    return {p["label"]: p["score"] for p in predictions}

# SYSTEM INPUTS
image_input = gr.Image(label="Upload iamge candidate", type="filepath")
text_input = gr.Textbox(type="text", label="Possible classes to detect.")
dropdown_input = gr.Dropdown(
    [
        MODEL_OPENAI_CLIP_VIT_LARGE_PATCH_14,
        MODEL_OPENAI_CLIP_VIT_BASE_PATCH_16,
        MODEL_OPENAI_CLIP_VIT_BASE_PATCH_32
    ], 
    label="OpenAI ClipVit Model")

# SYSTEM OUTPUTS
output_label = gr.outputs.Label()

demo = gr.Interface(
    launch,
    inputs=[image_input, text_input, dropdown_input],
    outputs=output_label,
    title="Demo aobut zero-shot-image-classification",
    examples = input_examles
)

if __name__ == "__main__":
    demo.launch()