Spaces:
Runtime error
Runtime error
File size: 1,366 Bytes
cb8043e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import models
import torch
import torchvision.transforms as transforms
import cv2
# 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('../outputs/model.pth')
# 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']
image = cv2.imread(f"../input/movie-classifier/Multi_Label_dataset/Images/tt0084058.jpg")
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])}
print(out_dict)
|