Upload 3 files

tools/extract_features.py

@@ -0,0 +1,63 @@
+import torch
+from torch import nn
+from torchvision import models, transforms
+
+
+# https://towardsdatascience.com/image-feature-extraction-using-pytorch-e3b327c3607a
+class FeatureExtractor(nn.Module):
+    def __init__(self, model):
+        super(FeatureExtractor, self).__init__()
+        # Extract VGG-16 Feature Layers
+        self.features = list(model.features)
+        self.features = nn.Sequential(*self.features)
+        # Extract VGG-16 Average Pooling Layer
+        self.pooling = model.avgpool
+        # Convert the image into one-dimensional vector
+        self.flatten = nn.Flatten()
+        # Extract the first part of fully-connected layer from VGG16
+        self.fc = model.classifier[0]
+
+    def forward(self, x):
+        # It will take the input 'x' until it returns the feature vector called 'out'
+        out = self.features(x)
+        out = self.pooling(out)
+        out = self.flatten(out)
+        out = self.fc(out)
+        return out
+
+# Initialize the model
+"""
+https://download.pytorch.org/models/vgg16-397923af.pth
+"""
+MODEL = models.vgg16(pretrained=True)
+NEW_MODEL = FeatureExtractor(MODEL)
+
+# Transform the image, so it becomes readable with the model
+"""
+Without center crop isn't working! 
+-> Así que se lo puse pero al final no me afecta en nada a la imagen mientras el resize sea el mismo.
+https://stackoverflow.com/questions/69334048/understanding-transforms-resize-and-centercrop-with-same-size
+https://discuss.pytorch.org/t/transforms-resize-vs-centercrop/86588/2
+"""
+TRANSFORMS = transforms.Compose([
+    transforms.ToPILImage(),
+    transforms.Resize(448),
+    transforms.CenterCrop(448),
+    transforms.ToTensor()
+])
+
+
+def extract_features_image(img):
+    # Transform the image
+    img = TRANSFORMS(img)
+    # Reshape the image. PyTorch model reads 4-dimensional tensor
+    # [batch_size, channels, width, height]
+    img = img.reshape(1, 3, 448, 448)
+
+    # We only extract features, so we don't need gradient
+    with torch.no_grad():
+        # Extract the feature from the image
+        feature = NEW_MODEL(img)
+
+    return feature.tolist()[0]
+

tools/load_database.py

@@ -0,0 +1,17 @@
+import pandas as pd
+
+# Default values
+filename = "./database/dataframes/all_database.csv"
+DATABASE = pd.read_csv(filename)
+
+
+def select_database(marca, prenda):
+    filtered_df =  DATABASE[DATABASE['tag'] == prenda]
+
+    if marca != "Ninguno":
+        filtered_df = filtered_df[filtered_df['marca'] == marca]
+    
+    # Restar index y eliminar el ex index column
+    filtered_df.reset_index(inplace=True, drop=True)
+   
+    return filtered_df

tools/search.py

@@ -0,0 +1,40 @@
+import numpy as np
+from scipy.spatial import cKDTree
+import ast
+
+
+def search_similar_products(vgg_search, df, number_of_items=5):
+    """Find similar products"""
+
+    try:
+        # Similar features
+        indexes = df.index
+        indexes_list = indexes.tolist()
+        vgg_vector_list = df["vgg_vector"].tolist()
+        vgg_vector_list_interested = [ast.literal_eval(vgg_vector_list[i]) for i in indexes_list]
+        
+        # Searching
+        matrix = np.array(vgg_vector_list_interested)
+
+        # [0]: distances
+        # [1]: indexes
+        results_indexes = cKDTree(matrix).query(vgg_search, k=number_of_items)[1]
+
+        # Make output
+        images_urls_list = []
+        for index in results_indexes:
+            image_url = df.iloc[index]['image_urls']
+            images_urls_list.append(image_url)
+
+        result = {
+            "success": True,
+            "result": images_urls_list,
+            }
+    
+    except:
+        result = {
+            "success": False,
+            "result":  [],
+            }
+
+    return result