add search code

Author: CVxTz

code/eda.py

@@ -9,6 +9,8 @@
 
 data = pd.read_json(out_name)
 
+data = data.sample(n=10000)
+
 img_repr = data['image_repr'].tolist()
 img_repr_random = data['image_repr'].tolist()
 shuffle(img_repr_random)

code/model_triplet.py

@@ -98,12 +98,14 @@ def cap_sequences(list_sequences, max_len, append):
     return capped
 
 
-def read_img(path):
+def read_img(path, preprocess=True):
     img = cv2.imread(path)
     if img is None or img.size<10:
         img = np.zeros((222, 171))
     img = cv2.resize(img, (171, 222))
-    return preprocess_input(img)
+    if preprocess:
+        img = preprocess_input(img)
+    return img
 
 
 def gen(list_images, list_captions, batch_size=16, aug=False):

code/predict_model_triplet.py

@@ -31,37 +31,40 @@ def chunker(seq, size):
     t_model.load_weights(file_path, by_name=True)
     i_model.load_weights(file_path, by_name=True)
 
-    target_image_encoding = []
-
-    for img_paths in tqdm(chunker(list_images_test, 128), total=len(list_images_test)//128):
-        images = np.array([read_img(file_path) for file_path in img_paths])
-        e = i_model.predict(images)
-        target_image_encoding += e.tolist()
-
-    target_text_encoding = t_model.predict(np.array(captions_test), verbose=1, batch_size=128)
-
-    target_text_encoding = target_text_encoding.tolist()
-
-    df = pd.DataFrame({"images": list_images_test, "text": _captions_test, "image_repr": target_image_encoding,
-                       "text_repr": target_text_encoding})
-
-    df.to_json(out_name, orient='records')
-
-    data = json.load(open(out_name, 'r'))
-    json.dump(data, open(out_name, 'w'), indent=4)
-
-    # New queries
-
-    out_name = "../output/queries_representations.json"
-
-    _captions_test = ['blue shirt', 'red dress', 'halloween outfit', 'baggy jeans', 'pokemon']
+    # target_image_encoding = []
+    #
+    # for img_paths in tqdm(chunker(list_images_test, 128), total=len(list_images_test)//128):
+    #     images = np.array([read_img(file_path) for file_path in img_paths])
+    #     e = i_model.predict(images)
+    #     target_image_encoding += e.tolist()
+    #
+    # target_text_encoding = t_model.predict(np.array(captions_test), verbose=1, batch_size=128)
+    #
+    # target_text_encoding = target_text_encoding.tolist()
+    #
+    # df = pd.DataFrame({"images": list_images_test, "text": _captions_test, "image_repr": target_image_encoding,
+    #                    "text_repr": target_text_encoding})
+    #
+    # df.to_json(out_name, orient='records')
+    #
+    # data = json.load(open(out_name, 'r'))
+    # json.dump(data, open(out_name, 'w'), indent=4)
+    #
+    # # New queries
+    #
+    # out_name = "../output/queries_representations.json"
+
+    _captions_test = ['blue tshirt', 'blue shirt', 'red dress', 'halloween outfit', 'baggy jeans', 'ring',
+                      'Black trousers', 'heart Pendant']
 
     captions_test = [tokenize(x) for x in _captions_test]
     captions_test = map_sentences(captions_test, mapping)
     captions_test = cap_sequences(captions_test, 70, 0)
 
     target_text_encoding = t_model.predict(np.array(captions_test), verbose=1, batch_size=128)
 
+    target_text_encoding = target_text_encoding.tolist()
+
     df = pd.DataFrame({"text": _captions_test,
                        "text_repr": target_text_encoding})
 

code/search.py

@@ -0,0 +1,28 @@
+import json
+from random import shuffle
+import pandas as pd
+import numpy as np
+from matplotlib import pyplot
+from sklearn.neighbors import NearestNeighbors
+
+repr_json = "../output/test_representations.json"
+
+data = pd.read_json(repr_json)
+
+data = data.sample(n=1000)
+
+img_repr = data['image_repr'].tolist()
+text_repr = data['text_repr'].tolist()
+
+nn = NearestNeighbors(n_jobs=-1, n_neighbors=1000)
+
+nn.fit(text_repr)
+
+preds = nn.kneighbors(img_repr, return_distance=False).tolist()
+ranks = []
+
+for i, x in enumerate(preds):
+    rank = x.index(i)+1
+    ranks.append(rank)
+
+print("Average rank :", np.mean(ranks))

code/search_by_image.py

@@ -0,0 +1,57 @@
+import json
+from random import shuffle
+import pandas as pd
+import numpy as np
+from matplotlib import pyplot
+from sklearn.neighbors import NearestNeighbors
+import matplotlib.pyplot as plt
+from model_triplet import read_img
+import cv2
+from uuid import uuid4
+
+
+repr_json = "../output/test_representations.json"
+
+data = pd.read_json(repr_json)
+
+data = data.sample(n=50000)
+
+img_repr = data['image_repr'].tolist()
+img_paths = data['images'].tolist()
+text_repr = data['text_repr'].tolist()
+
+nn = NearestNeighbors(n_jobs=-1, n_neighbors=9)
+
+nn.fit(img_repr)
+
+preds = nn.kneighbors(img_repr[:100], return_distance=False).tolist()
+
+most_similar_images = []
+query_image = []
+
+
+for i, x in enumerate(preds):
+    preds_paths = [img_paths[i] for i in x]
+    query_image.append(preds_paths[0])
+    most_similar_images.append(preds_paths[1:])
+
+for q, similar in zip(query_image, most_similar_images):
+    fig, axes = plt.subplots(3, 3)
+    all_images = [q]+similar
+
+    for idx, img_path in enumerate(all_images):
+        i = idx % 3  # Get subplot row
+        j = idx // 3  # Get subplot column
+        image = read_img(img_path, preprocess=False)
+        image = image[:, :, ::-1]
+        axes[i, j].imshow(image/255)
+        axes[i, j].axis('off')
+        axes[i, j].axis('off')
+        if idx == 0:
+            axes[i, j].set_title('Query Image')
+        else:
+            axes[i, j].set_title('Result Image %s'%i)
+
+    plt.subplots_adjust(wspace=0.2, hspace=0.2)
+    plt.savefig('../output/images/%s.png'%uuid4().hex)
+

code/search_by_keywords.py

@@ -0,0 +1,59 @@
+import json
+from random import shuffle
+import pandas as pd
+import numpy as np
+from matplotlib import pyplot
+from sklearn.neighbors import NearestNeighbors
+import matplotlib.pyplot as plt
+from model_triplet import read_img
+import cv2
+from uuid import uuid4
+
+
+repr_json = "../output/test_representations.json"
+
+data = pd.read_json(repr_json)
+
+queries_repr_json = "../output/queries_representations.json"
+
+queries_data = pd.read_json(queries_repr_json)
+
+data = data.sample(n=50000)
+
+img_repr = data['image_repr'].tolist()
+img_paths = data['images'].tolist()
+text_repr = queries_data['text_repr'].tolist()
+
+nn = NearestNeighbors(n_jobs=-1, n_neighbors=9)
+
+nn.fit(img_repr)
+
+preds = nn.kneighbors(text_repr, return_distance=False).tolist()
+
+most_similar_images = []
+query_image = []
+
+
+for i, x in enumerate(preds):
+    preds_paths = [img_paths[i] for i in x]
+    most_similar_images.append(preds_paths)
+
+for q, all_images in zip(queries_data['text'], most_similar_images):
+    fig, axes = plt.subplots(3, 3)
+
+    for idx, img_path in enumerate(all_images):
+        i = idx % 3  # Get subplot row
+        j = idx // 3  # Get subplot column
+        image = read_img(img_path, preprocess=False)
+        image = image[:, :, ::-1]
+        axes[i, j].imshow(image/255)
+        axes[i, j].axis('off')
+        axes[i, j].axis('off')
+        axes[i, j].set_title('Result Image %s'%i)
+
+
+    plt.subplots_adjust(wspace=0.2, hspace=0.2)
+    fig.title('Query : %s'%q)
+
+    plt.savefig('../output/queries/%s.png'%uuid4().hex)
+