Datatrooper/posters_classification

1import gradio as gr2import models3import torch4import torchvision.transforms as transforms5import cv26import numpy as np7 8 9# initialize the computation device10device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')11#intialize the model12model = models.model(pretrained=False, requires_grad=False).to(device)13# load the model checkpoint14checkpoint = torch.load('model.pth', map_location=device)15# load model weights state_dict16model.load_state_dict(checkpoint['model_state_dict'])17model.eval()18 19transform = transforms.Compose([20            transforms.ToPILImage(),21            transforms.ToTensor(),22            ])23 24genres = ['Action', 'Adventure', 'Animation', 'Biography', 'Comedy', 'Crime',25 'Documentary', 'Drama', 'Family', 'Fantasy', 'History', 'Horror', 'Music',26 'Musical', 'Mystery', 'N/A', 'News', 'Reality-TV', 'Romance', 'Sci-Fi', 'Short',27 'Sport', 'Thriller', 'War', 'Western']28 29 30def segment(image):31    image = np.asarray(image)32    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)33    image = transform(image)34    image = torch.tensor(image, dtype=torch.float32)35    image = image.to(device)36    image = torch.unsqueeze(image, dim=0)37    # get the predictions by passing the image through the model38    outputs = model(image)39    outputs = torch.sigmoid(outputs)40    outputs = outputs.detach().cpu()41 42    out_dict = {k: v for k, v in zip(genres, outputs.tolist()[0])}43    return out_dict44 45iface = gr.Interface(fn=segment, 46                     inputs="image", 47                     outputs="label",48                     title="Poster classification",49                     description="classify the genre of your poster by uploading an image",50                     examples=[["imagenes/tt0084058.jpg"], ["imagenes/tt0084867.jpg"], ["imagenes/tt0085121.jpg"]]).launch()