i2i mode

app.py

@@ -58,25 +58,33 @@ def download_url_img(url):
     return False, []
 
 
-def search(mode, text):
-
-    # translator = Translator(from_lang="chinese",to_lang="english")
-    # text = translator.translate(text)
-    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 i, (caps, length) in enumerate(dataset_loader, 0):
-        input_caps = caps.to(device)
-        with torch.no_grad():
-            _, output_emb = join_emb(None, input_caps, length)
-        caps_enc.append(output_emb.cpu().data.numpy())
-
-    caps_stack = np.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(caps_stack, imgs_emb, imgs_url, ks=100)
+def search(mode, image, 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)
+        caps_enc = list()
+        for i, (caps, length) in enumerate(dataset_loader, 0):
+            input_caps = caps
+            with torch.no_grad():
+                _, output_emb = join_emb(None, input_caps, length)
+            caps_enc.append(output_emb)
+        _stack = np.vstack(caps_enc)
+        
+    elif mode == I2I:
+        dataset = torch.Tensor(image).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
+            with torch.no_grad():
+                _, output_emb = join_emb(input_imgs, None, length)
+            img_enc.append(output_emb)
+        _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
@@ -109,11 +117,13 @@ if __name__ == "__main__":
     encoder = TextEncoder()
     imgs_emb_file_path = "./coco_img_emb"
     imgs_emb, imgs_path = load_obj(imgs_emb_file_path)
+    imgs_url = [os.path.join("http://images.cocodataset.org/train2017", img_path.strip().split('_')[-1]) for img_path in imgs_path]
     print("prepare done!")
     iface = gr.Interface(
         fn=search,
         inputs=[
-            gr.inputs.Radio([T2I]),
+            gr.inputs.Radio([I2I, T2I]),
+            gr.inputs.Image(label="Image to search", optional=True),
             gr.inputs.Textbox(
                 lines=1, label="Text query", placeholder="Introduce the search text...",
             ),