| # THIS IS THE CODE I MADE FOR THE PREDICTING | |
| import torch | |
| import torchvision | |
| from torchvision.transforms import transforms | |
| from PIL import Image | |
| model = torchvision.models.resnet50(pretrained=False) | |
| model.fc = torch.nn.Linear(in_features=2048, out_features=1) | |
| model.load_state_dict(torch.load('/content/zero_shot_classification_model.pth')) | |
| model.eval() | |
| test_transform = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) | |
| ]) | |
| test_image_path = '/content/test_image.png' | |
| test_image = Image.open(test_image_path).convert('RGB') | |
| test_image = test_transform(test_image) | |
| test_image = test_image.unsqueeze(0) | |
| with torch.no_grad(): | |
| prediction = model(test_image) | |
| probability = torch.sigmoid(prediction).item() | |
| print(f"The probability of the image being a Roblox character is: {probability}") |