Spaces:
Runtime error
Runtime error
| 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() |