import gradio as gr
from models.blip_vqa import blip_vqa
import torch
from torchvision import transforms
from torchvision.transforms.functional import InterpolationMode

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
image_size = 480

transform = transforms.Compose([
    transforms.Resize((image_size,image_size),interpolation=InterpolationMode.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
        ])

model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_vqa_capfilt_large.pth'
model = blip_vqa(pretrained=model_url, image_size=image_size, vit='base')
model.eval()
model = model.to(device)


def pool_alarm(raw_image):
    question = 'there is someone in the pool?'
    image = transform(raw_image).unsqueeze(0).to(device)
    with torch.no_grad():
        answer = model(image, question, train=False, inference='generate')

    return 'answer: ' + answer[0]


input = gr.inputs.Image(type='pil')
output = gr.outputs.Textbox()
examples = ['alarm.jpeg', 'alarm1.jpeg', 'walk.jpeg']
title = "Pool Alarm"
description = "Using visual question answering to check if there is someone in the swimming pool"
intf = gr.Interface(fn=pool_alarm, inputs=input, outputs=output, examples=examples,
                    title=title, description=description).launch()