Spaces:
Runtime error
Runtime error
| import os | |
| from PIL import Image | |
| import faiss | |
| import streamlit as st | |
| import torch | |
| import torch.nn.functional as F | |
| from torchvision.transforms import Compose, Resize, ToTensor, Normalize | |
| from model import FashionPrediction | |
| transforms = Compose([Resize((232, 232)), ToTensor(), | |
| Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) | |
| path = './data/final-models/resnet_152_classification.pt' | |
| index_path = os.path.join('./data/index_files/resnet152_unweighted_flat.index') | |
| 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'} | |
| def load_trained_model(model_path): | |
| model = FashionPrediction() | |
| model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))) | |
| model.eval() | |
| return model | |
| def load_index(): | |
| index = faiss.read_index(index_path) | |
| return index | |
| def get_nearest_k(input_image, top_k=5): | |
| model = load_trained_model(path) | |
| transformed_image = transforms(input_image) | |
| transformed_image = torch.unsqueeze(transformed_image, 0) | |
| with torch.inference_mode(): | |
| query_embeddings = model(transformed_image, True) | |
| the_index = load_index() | |
| dist, indexes = the_index.search(query_embeddings, top_k) | |
| return dist, indexes | |
| def get_predictions(input_image): | |
| model = load_trained_model(path) | |
| transformed_image = transforms(input_image) | |
| transformed_image = torch.unsqueeze(transformed_image, 0) | |
| 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))) | |
| pred_dict = { | |
| 'Predicted Master Category': (master_dict[all_predictions[1][1][0]], master_dict[all_predictions[1][1][1]]), | |
| 'Master Category Probability': [round(prob, 3) for prob in all_predictions[1][0]], | |
| 'Predicted Sub Category': (subcat_dict[all_predictions[3][1][0]], subcat_dict[all_predictions[3][1][1]]), | |
| 'Sub Category Probability': [round(prob, 3) for prob in all_predictions[3][0]], | |
| 'Predicted person type': (gender_dict[all_predictions[0][1][0]], gender_dict[all_predictions[0][1][1]]), | |
| 'Person Type Probability': [round(prob, 3) for prob in all_predictions[0][0]], | |
| 'Predicted Color': (color_dict[all_predictions[2][1][0]], color_dict[all_predictions[2][1][1]]), | |
| 'Color Probability': [round(prob, 3) for prob in all_predictions[2][0]] | |
| } | |
| return pred_dict | |
| if __name__ == '__main__': | |
| sample_image = Image.open('./data/small_images_0_9999/0.jpg') | |
| output = get_predictions(sample_image) | |
| print(output) | |