Spaces:
Runtime error
Runtime error
| import os | |
| from PIL import Image | |
| import torch | |
| import torchvision | |
| from torch import nn | |
| import torch.nn.functional as F | |
| from torchvision.transforms import Compose, Resize, ToTensor, Normalize | |
| class FashionResnet(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.model = torchvision.models.resnet152() | |
| self.model.fc = torch.nn.Identity() | |
| def forward(self, x): | |
| return self.model(x) | |
| class FashionClassifictions(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.gender_linear1 = nn.Linear(2048, 1024) | |
| self.gender_linear2 = nn.Linear(1024, 256) | |
| self.gender_out = nn.Linear(256, 5) | |
| self.mastercat_linear1 = nn.Linear(2048, 1024) | |
| self.mastercat_linear2 = nn.Linear(1024, 256) | |
| self.mastercat_out = nn.Linear(256, 4) | |
| self.subcat_linear1 = nn.Linear(2048, 1024) | |
| self.subcat_linear2 = nn.Linear(1024, 256) | |
| self.subcat_out = nn.Linear(256, 32) | |
| self.color_linear1 = nn.Linear(2048, 1024) | |
| self.color_linear2 = nn.Linear(1024, 256) | |
| self.color_out = nn.Linear(256, 44) | |
| self.activation = nn.ReLU() | |
| self.dropout = nn.Dropout(0.3) | |
| def forward(self, out): | |
| gender_out = self.activation(self.dropout((self.gender_linear1(out)))) | |
| gender_out = self.activation(self.dropout(self.gender_linear2(gender_out))) | |
| gender_out = self.gender_out(gender_out) | |
| master_out = self.activation(self.dropout((self.mastercat_linear1(out)))) | |
| master_out = self.activation(self.dropout(self.mastercat_linear2(master_out))) | |
| master_out = self.mastercat_out(master_out) | |
| subcat_out = self.activation(self.dropout((self.subcat_linear1(out)))) | |
| subcat_out = self.activation(self.dropout(self.subcat_linear2(subcat_out))) | |
| subcat_out = self.subcat_out(subcat_out) | |
| color_out = self.activation(self.dropout((self.color_linear1(out)))) | |
| color_out = self.activation(self.dropout(self.color_linear2(color_out))) | |
| color_out = self.color_out(color_out) | |
| return gender_out, master_out, subcat_out, color_out | |
| class FashionPrediction(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.feature_model = FashionResnet() | |
| self.classification_model = FashionClassifictions() | |
| def forward(self, x, only_embedding=False): | |
| out_embed = self.feature_model(x) | |
| if only_embedding: | |
| return out_embed | |
| gender_out, master_out, subcat_out, color_out = self.classification_model(out_embed) | |
| return gender_out, master_out, subcat_out, color_out, out_embed | |
| if __name__ == '__main__': | |
| trained_model_path = os.path.join('./data/final-models/resnet_152_classification.pt') | |
| model = FashionPrediction() | |
| # print(model) | |
| model.load_state_dict(torch.load(trained_model_path, map_location=torch.device('cpu'))) | |
| model.eval() | |
| sample_image = Image.open('./data/small_images/0.jpg') | |
| transforms = Compose([Resize((232, 232)), ToTensor(), | |
| Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) | |
| transformed_image = transforms(sample_image) | |
| transformed_image = torch.unsqueeze(transformed_image, 0) | |
| print(transformed_image.shape) | |
| with torch.inference_mode(): | |
| logits = model(transformed_image, False) | |
| gender_prob = F.softmax(logits[0], dim=1) | |
| master_prob = F.softmax(logits[1], dim=1) | |
| subcat_prob = F.softmax(logits[2], dim=1) | |
| color_prob = F.softmax(logits[3], dim=1) | |
| top2_gender = torch.topk(gender_prob, 2, dim=1) | |
| top2_master = torch.topk(master_prob, 2, dim=1) | |
| top2_subcat = torch.topk(subcat_prob, 2, dim=1) | |
| top2_color = torch.topk(color_prob, 2, dim=1) | |
| all_predictions = (list(top2_gender.values.numpy().reshape(-1)), list(top2_gender.indices.numpy().reshape(-1))), \ | |
| (list(top2_master.values.numpy().reshape(-1)), list(top2_master.indices.numpy().reshape(-1))), \ | |
| (list(top2_color.values.numpy().reshape(-1)), list(top2_color.indices.numpy().reshape(-1))), \ | |
| (list(top2_subcat.values.numpy().reshape(-1)), list(top2_subcat.indices.numpy().reshape(-1))) | |
| gender_dict = {0: 'Boys', 1: 'Girls', 2: 'Men', 3: 'Unisex', 4: 'Women'} | |
| master_dict = {0: 'Accessories', 1: 'Apparel', 2: 'Footwear', 3: 'Personal Care'} | |
| subcat_dict = {0: 'Accessories', 1: 'Apparel Set', 2: 'Bags', 3: 'Belts', 4: 'Bottomwear', 5: 'Cufflinks', | |
| 6: 'Dress', 7: 'Eyes', 8: 'Eyewear', 9: 'Flip Flops', 10: 'Fragrance', 11: 'Headwear', | |
| 12: 'Innerwear', 13: 'Jewellery', 14: 'Lips', 15: 'Loungewear and Nightwear', 16: 'Makeup', | |
| 17: 'Mufflers', 18: 'Nails', 19: 'Sandal', 20: 'Saree', 21: 'Scarves', 22: 'Shoe Accessories', | |
| 23: 'Shoes', 24: 'Skin', 25: 'Skin Care', 26: 'Socks', 27: 'Stoles', 28: 'Ties', 29: 'Topwear', | |
| 30: 'Wallets', 31: 'Watches'} | |
| color_dict = {0: 'Beige', 1: 'Black', 2: 'Blue', 3: 'Bronze', 4: 'Brown', 5: 'Burgundy', 6: 'Charcoal', | |
| 7: 'Coffee Brown', 8: 'Copper', 9: 'Cream', 10: 'Gold', 11: 'Green', 12: 'Grey', 13: 'Grey Melange', | |
| 14: 'Khaki', 15: 'Lavender', 16: 'Magenta', 17: 'Maroon', 18: 'Mauve', 19: 'Metallic', 20: 'Multi', | |
| 21: 'Mushroom Brown', 22: 'Mustard', 23: 'Navy Blue', 24: 'Nude', 25: 'Off White', 26: 'Olive', | |
| 27: 'Orange', 28: 'Peach', 29: 'Pink', 30: 'Purple', 31: 'Red', 32: 'Rose', 33: 'Rust', | |
| 34: 'Sea Green', 35: 'Silver', 36: 'Skin', 37: 'Steel', 38: 'Tan', 39: 'Taupe', 40: 'Teal', | |
| 41: 'Turquoise Blue', 42: 'White', 43: 'Yellow'} | |
| print("All predictions:", all_predictions) | |
| pred_dict = { | |
| 'Predicted Master Category': (master_dict[all_predictions[1][1][0]], master_dict[all_predictions[1][1][1]]), | |
| 'Master Category Probability': all_predictions[1][0], | |
| 'Predicted Sub Category': (subcat_dict[all_predictions[3][1][0]], subcat_dict[all_predictions[3][1][1]]), | |
| 'Sub Category Probability': all_predictions[3][0], | |
| 'Predicted person type': (gender_dict[all_predictions[0][1][0]], gender_dict[all_predictions[0][1][1]]), | |
| 'Person Type Probability': all_predictions[0][0], | |
| 'Predicted Color': (color_dict[all_predictions[2][1][0]], color_dict[all_predictions[2][1][1]]), | |
| 'Color Probability': all_predictions[2][0] | |
| } | |
| print(pred_dict) | |