import gradio as gr
import clip
import torch

import utils

clip_model = "RN50x4"
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load(clip_model, device=device, jit=False)
model.eval()


def grad_cam_fn(text, img, saliency_layer):
    resize = model.visual.input_resolution
    img = img.resize((resize, resize))

    text_input = clip.tokenize([text]).to(device)
    text_feature = model.encode_text(text_input).float()
    image_input = preprocess(img).unsqueeze(0).to(device)

    attn_map = utils.gradCAM(
        model.visual,
        image_input,
        text_feature,
        getattr(model.visual, saliency_layer)
    )
    attn_map = attn_map.squeeze().detach().cpu().numpy()
    attn_map = utils.getAttMap(img, attn_map)

    return attn_map


interface = gr.Interface(
    fn=grad_cam_fn,
    inputs=[
        gr.inputs.Textbox(
            label="Target Text",
            lines=1),
        gr.inputs.Image(
            label='Input Image',
            image_mode="RGB",
            type='pil',
            shape=(512, 512)),
        gr.inputs.Dropdown(
            ["layer4", "layer3", "layer2", "layer1"],
            default="layer4",
            label="Saliency Layer")
    ],
    outputs=gr.outputs.Image(
        type="pil", 
        label="Attention Map"),
    examples=[
        ['a cat lying on the floor', 'assets/cat_dog.jpg', 'layer4'],
        ['a dog sitting', 'assets/cat_dog.jpg', 'layer4']
    ],
    description="OpenAI CLIP Grad CAM")
interface.launch()