import gradio as gr
import models
import torch
import torchvision.transforms as transforms
import cv2
import numpy as np


# initialize the computation device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#intialize the model
model = models.model(pretrained=False, requires_grad=False).to(device)
# load the model checkpoint
checkpoint = torch.load('model.pth', map_location=device)
# load model weights state_dict
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.ToTensor(),
            ])

genres = ['Action', 'Adventure', 'Animation', 'Biography', 'Comedy', 'Crime',
 'Documentary', 'Drama', 'Family', 'Fantasy', 'History', 'Horror', 'Music',
 'Musical', 'Mystery', 'N/A', 'News', 'Reality-TV', 'Romance', 'Sci-Fi', 'Short',
 'Sport', 'Thriller', 'War', 'Western']


def segment(image):
    image = np.asarray(image)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = transform(image)
    image = torch.tensor(image, dtype=torch.float32)
    image = image.to(device)
    image = torch.unsqueeze(image, dim=0)
    # get the predictions by passing the image through the model
    outputs = model(image)
    outputs = torch.sigmoid(outputs)
    outputs = outputs.detach().cpu()

    out_dict = {k: v for k, v in zip(genres, outputs.tolist()[0])}
    return out_dict

iface = gr.Interface(fn=segment, 
                     inputs="image", 
                     outputs="label",
                     title="Poster classification",
                     description="classify the genre of your poster by uploading an image",
                     examples=[["imagenes/tt0084058.jpg"], ["imagenes/tt0084867.jpg"], ["imagenes/tt0085121.jpg"]]).launch()