v3.0

app.py

@@ -33,7 +33,7 @@ from misc.dataset import TextEncoder
 import requests 
 from io import BytesIO
 from translate import Translator
-import cupy as cp
+# import cupy as cp
 from torchvision import transforms
 import random
 
@@ -46,7 +46,7 @@ topK = 5
 T2I = "以文搜图"
 I2I = "以图搜图"
 
-DDT = "双塔动态嵌入"
+DPDT = "双塔动态嵌入"
 UEFDT = "双塔联合融合"
 IEFDT = "双塔嵌入融合"
 ViLT = "视觉语言预训练"
@@ -82,37 +82,28 @@ def search(mode, method, image, text):
         dataset_loader = DataLoader(dataset, batch_size=batch_size, num_workers=1, pin_memory=True, collate_fn=collate_fn_cap_padded)
         caps_enc = list()
         for i, (caps, length) in enumerate(dataset_loader, 0):
-            input_caps = caps.to(device)
+            input_caps = caps
             with torch.no_grad():
                 _, output_emb = join_emb(None, input_caps, length)
             caps_enc.append(output_emb)
-        _stack = cp.vstack(caps_enc)
+        _stack = np.vstack(caps_enc)
         
     elif mode == I2I:
         dataset = normalize(torch.Tensor(image).permute(2, 0, 1)).unsqueeze(dim=0)
         dataset_loader = DataLoader(dataset, batch_size=batch_size, num_workers=1, pin_memory=True, collate_fn=collate_fn_cap_padded)
         img_enc = list()
         for i, (imgs, length) in enumerate(dataset_loader, 0):
-            input_imgs = imgs.to(device)
+            input_imgs = imgs
             with torch.no_grad():
                 output_emb, _ = join_emb(input_imgs, None, None)
             img_enc.append(output_emb)
-        _stack = cp.vstack(img_enc)
+        _stack = np.vstack(img_enc)
 
-    # dataset = torch.Tensor(encoder.encode(text)).unsqueeze(dim=0)
-    # dataset_loader = DataLoader(dataset, batch_size=batch_size, num_workers=1, pin_memory=True, collate_fn=collate_fn_cap_padded)
-    # caps_enc = list()
 
-    # for _, (caps, length) in enumerate(dataset_loader, 0):
-    #     input_caps = caps.to(device)
-    #     with torch.no_grad():
-    #         _, caps_emb = join_emb(None, input_caps, length)
-    #     caps_enc.append(caps_emb)
-    # caps_stack = cp.vstack(caps_enc)
 
     imgs_url = [os.path.join("http://images.cocodataset.org/train2017", img_path.strip().split('_')[-1]) for img_path in imgs_path]
 
-    recall_imgs = recallTopK(_stack, imgs_emb, imgs_url, ks=100)
+    recall_imgs = recallTopK(_stack, imgs_emb, imgs_url, method, ks=100)
 
 
     tmp1 = []
@@ -121,7 +112,7 @@ def search(mode, method, image, text):
     if method == ViLT: 
         pass
     else:
-        if method == DDT: swap_width = 5
+        if method == DPDT: swap_width = 5
         elif method == UEFDT: swap_width = 2
         elif method == IEFDT: swap_width = 1
 
@@ -146,6 +137,8 @@ def search(mode, method, image, text):
 
 if __name__ == "__main__":
     # print("Loading model from:", model_path)
+    import nltk
+    nltk.download('punkt')
     checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
 
     join_emb = joint_embedding(checkpoint['args_dict'])
@@ -159,10 +152,10 @@ if __name__ == "__main__":
     encoder = TextEncoder()
     imgs_emb_file_path = "./coco_img_emb"
     imgs_emb, imgs_path = load_obj(imgs_emb_file_path)
-    imgs_emb = cp.asarray(imgs_emb)
+    imgs_url = [os.path.join("http://images.cocodataset.org/train2017", img_path.strip().split('_')[-1]) for img_path in imgs_path]
+    # imgs_emb = np.asarray(imgs_emb)
 
-    normalize = transforms.Normalize(mean=[0.485 * 255, 0.456 * 255, 0.406 * 255], 
-    std=[0.229 * 255, 0.224 * 255, 0.225 * 255])
+    normalize = transforms.Normalize(mean=[0.485 * 255, 0.456 * 255, 0.406 * 255], std=[0.229 * 255, 0.224 * 255, 0.225 * 255])
     
     cat_image = "./cat_example.jpg"
     dog_image = "./dog_example.jpg"
@@ -174,7 +167,7 @@ if __name__ == "__main__":
         fn=search,
         inputs=[
             gr.inputs.Radio([I2I, T2I]),
-            gr.inputs.Radio([DDT, UEFDT, IEFDT, ViLT]),
+            gr.inputs.Radio([DPDT, UEFDT, IEFDT, ViLT]),
             gr.inputs.Image(shape=(400, 400), label="Image to search", optional=True),
             gr.inputs.Textbox(
                 lines=1, label="Text query", placeholder="please input text query here...",
@@ -189,7 +182,7 @@ if __name__ == "__main__":
         gr.outputs.Image(type="auto", label="5rd Best match")
         ],
         examples=[
-            [I2I, DDT, cat_image, ""],#, img_folder / "8LWtpfhGP4U.jpg"],
+            [I2I, DPDT, cat_image, ""],#, img_folder / "8LWtpfhGP4U.jpg"],
             [I2I, ViLT, dog_image, ""],#, img_folder / "_ppnPXy_TVw.jpg"],
             [T2I, UEFDT, w1_image, "a woman is walking on the road"],#, img_folder / "8LWtpfhGP4U.jpg"],
             [T2I, IEFDT, w2_image, "a boy is eating apple"],#, img_folder / "_ppnPXy_TVw.jpg"],

misc/evaluation.py

@@ -19,30 +19,32 @@ This scripts permits one to reproduce training and experiments of:
 
 Author: Martin Engilberge
 """
-
-import cupy as cp 
+from scripts.postprocess import postprocess 
+import numpy as np 
 
 from misc.utils import flatten
 
 def cosine_sim(A, B):
-    img_norm = cp.linalg.norm(A, axis=1)
-    caps_norm = cp.linalg.norm(B, axis=1)
+    img_norm = np.linalg.norm(A, axis=1)
+    caps_norm = np.linalg.norm(B, axis=1)
 
-    scores = cp.dot(A, B.T)
+    scores = np.dot(A, B.T)
 
-    norms = cp.dot(cp.expand_dims(img_norm, 1),
-                   cp.expand_dims(caps_norm.T, 1).T)
+    norms = np.dot(np.expand_dims(img_norm, 1),
+                   np.expand_dims(caps_norm.T, 1).T)
 
     scores = (scores / norms)
 
     return scores
 
-def recallTopK(cap_enc, imgs_enc, imgs_path, ks=10, scores=None):
+def recallTopK(cap_enc, imgs_enc, imgs_path, method, ks=10, scores=None):
 
     if scores is None:
         scores = cosine_sim(cap_enc, imgs_enc)
     
-    recall_imgs = [imgs_path[cp.asnumpy(i)] for i in cp.argsort(scores, axis=1)[0][::-1][:ks]]
+    # recall_imgs = [imgs_path[np.asnumpy(i)] for i in np.argsort(scores, axis=1)[0][::-1][:ks]]
+    recall_imgs = [imgs_path[i] for i in np.argsort(scores, axis=1)[0][::-1][:ks]]
+    postprocess(method, recall_imgs)
 
     return recall_imgs
 
@@ -50,26 +52,26 @@ def recall_at_k_multi_cap(imgs_enc, caps_enc, ks=[1, 5, 10], scores=None):
     if scores is None:
         scores = cosine_sim(imgs_enc[::5, :], caps_enc)
 
-    ranks = cp.array([cp.nonzero(cp.in1d(row, cp.arange(x * 5, x * 5 + 5, 1)))[0][0]
-                      for x, row in enumerate(cp.argsort(scores, axis=1)[:, ::-1])])
+    ranks = np.array([np.nonzero(np.in1d(row, np.arange(x * 5, x * 5 + 5, 1)))[0][0]
+                      for x, row in enumerate(np.argsort(scores, axis=1)[:, ::-1])])
 
-    medr_caps_search = cp.median(ranks)
+    medr_caps_search = np.median(ranks)
 
     recall_caps_search = list()
 
     for k in [1, 5, 10]:
         recall_caps_search.append(
-            (float(len(cp.where(ranks < k)[0])) / ranks.shape[0]) * 100)
+            (float(len(np.where(ranks < k)[0])) / ranks.shape[0]) * 100)
 
-    ranks = cp.array([cp.nonzero(row == int(x / 5.0))[0][0]
-                      for x, row in enumerate(cp.argsort(scores.T, axis=1)[:, ::-1])])
+    ranks = np.array([np.nonzero(row == int(x / 5.0))[0][0]
+                      for x, row in enumerate(np.argsort(scores.T, axis=1)[:, ::-1])])
 
-    medr_imgs_search = cp.median(ranks)
+    medr_imgs_search = np.median(ranks)
 
     recall_imgs_search = list()
     for k in ks:
         recall_imgs_search.append(
-            (float(len(cp.where(ranks < k)[0])) / ranks.shape[0]) * 100)
+            (float(len(np.where(ranks < k)[0])) / ranks.shape[0]) * 100)
 
     return recall_caps_search, recall_imgs_search, medr_caps_search, medr_imgs_search
 
@@ -87,13 +89,13 @@ def avg_recall(imgs_enc, caps_enc):
         caps = caps_enc[i:i + 5000]
         res.append(recall_at_k_multi_cap(imgs, caps))
 
-    return [cp.sum([x[i] for x in res], axis=0) / len(res) for i in range(len(res[0]))]
+    return [np.sum([x[i] for x in res], axis=0) / len(res) for i in range(len(res[0]))]
 
 
 def eval_recall(imgs_enc, caps_enc):
 
-    imgs_enc = cp.vstack(flatten(imgs_enc))
-    caps_enc = cp.vstack(flatten(caps_enc))
+    imgs_enc = np.vstack(flatten(imgs_enc))
+    caps_enc = np.vstack(flatten(caps_enc))
 
     res = avg_recall(imgs_enc, caps_enc)