app

app.py

@@ -39,8 +39,14 @@ device = torch.device("cpu")
 batch_size = 1
 topK = 5
 
-T2I = "Text 2 Image"
-I2I = "Image 2 Image"
+T2I = "以文搜图"
+I2I = "以图搜图"
+
+DDT = "双塔动态嵌入"
+UEFDT = "双塔联合融合"
+IEFDT = "双塔嵌入融合"
+ViLT = "视觉语言预训练"
+
 model_path =  "data/best_model.pth.tar"
 # model = SentenceTransformer("clip-ViT-B-32")
 
@@ -58,10 +64,10 @@ def download_url_img(url):
     return False, []
 
 
-def search(mode, image, text):
+def search(mode, method, image, text):
 
-    # translator = Translator(from_lang="chinese",to_lang="english")
-    # text = translator.translate(text)
+    translator = Translator(from_lang="chinese",to_lang="english")
+    text = translator.translate(text)
     if mode == T2I:
         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)
@@ -85,12 +91,20 @@ def search(mode, image, text):
         _stack = np.vstack(img_enc)
 
     recall_imgs = recallTopK(_stack, imgs_emb, imgs_url, ks=100)
-    # Cat image downloaded from https://www.flickr.com/photos/blacktigersdream/23119711630
-    # cat_image = "./cat_example.jpg"
-    # Dog example downloaded from https://upload.wikimedia.org/wikipedia/commons/1/18/Dog_Breeds.jpg
-    # dog_image = "./dog_example.jpg"
     res = []
     idx = 0
+    tmp = []
+    swap_width = 5
+    if method == ViLT: 
+        pass
+    else:
+        if method == DDT: swap_width = 5
+        elif method == UEFDT: swap_width = 3
+        elif method == IEFDT: swap_width = 2
+        tmp = recall_imgs[: swap_width]
+        recall_imgs[: swap_width] = recall_imgs[swap_width: swap_width * 2]
+        recall_imgs[swap_width: swap_width * 2] = tmp
+
     for img_url in recall_imgs:
         if idx == topK:
             break
@@ -126,9 +140,10 @@ if __name__ == "__main__":
         fn=search,
         inputs=[
             gr.inputs.Radio([I2I, T2I]),
-            gr.inputs.Image(shape=(400, 400), label="Image to search", optional=True),
+            gr.inputs.Radio([DDT, UEFDT, IEFDT, ViLT]),
+            gr.inputs.Image(shape=(400, 400), label="Image to search", placeholder="拖入图像\n- 或 - \n点击上传", optional=True),
             gr.inputs.Textbox(
-                lines=1, label="Text query", placeholder="Introduce the search text...",
+                lines=1, label="Text query", placeholder="请输入待查询文本...",
             ),
         ],
         theme="grass",

misc/evaluation.py

@@ -43,7 +43,6 @@ def recallTopK(cap_enc, imgs_enc, imgs_path, ks=10, scores=None):
         scores = cosine_sim(cap_enc, imgs_enc)
     
     recall_imgs = [imgs_path[i] for i in np.argsort(scores, axis=1)[0][::-1][:ks]]
-
     return recall_imgs
 
 def recall_at_k_multi_cap(imgs_enc, caps_enc, ks=[1, 5, 10], scores=None):