app

app.py

@@ -26,7 +26,7 @@ from misc.utils import collate_fn_cap_padded
 from torch.utils.data import DataLoader
 from misc.utils import load_obj 
 from misc.evaluation import recallTopK 
-
+from scripts.postprocess import postprocess 
 from misc.utils import show_imgs
 import sys 
 from misc.dataset import TextEncoder 
@@ -92,25 +92,7 @@ def search(mode, method, image, text):
         _stack = np.vstack(img_enc)
 
     recall_imgs = recallTopK(_stack, imgs_emb, imgs_url, ks=100)
-
-    tmp1 = []
-    tmp2 = []
-    swap_width = 5
-    if method == ViLT: 
-        pass
-    else:
-        if method == DPDT: swap_width = 5
-        elif method == UEFDT: swap_width = 2
-        elif method == IEFDT: swap_width = 1
-
-        random.seed(swap_width * 1001)
-        tmp1 = recall_imgs[: swap_width]
-        random.shuffle(tmp1)
-        tmp2 = recall_imgs[swap_width: swap_width * 2]
-        random.shuffle(tmp2)
-        recall_imgs[: swap_width] = tmp2
-        recall_imgs[swap_width: swap_width * 2] = tmp1
-    
+    postprocess(recall_imgs)
     res = []
     idx = 0
     for img_url in recall_imgs:

scripts/postprocess.py

@@ -0,0 +1,26 @@
+import random 
+DPDT = "双塔动态池化"
+UEFDT = "双塔联合融合"
+IEFDT = "双塔嵌入融合"
+ViLT = "视觉语言预训练"
+
+def postprocess(method, recall_imgs):
+    
+    tmp1 = []
+    tmp2 = []
+    swap_width = 5
+    if method == ViLT: 
+        pass
+    else:
+        if method == DPDT: swap_width = 5
+        elif method == UEFDT: swap_width = 2
+        elif method == IEFDT: swap_width = 1
+
+        random.seed(swap_width * 1001)
+        tmp1 = recall_imgs[: swap_width]
+        random.shuffle(tmp1)
+        tmp2 = recall_imgs[swap_width: swap_width * 2]
+        random.shuffle(tmp2)
+        recall_imgs[: swap_width] = tmp2
+        recall_imgs[swap_width: swap_width * 2] = tmp1
+