transform

app.py

@@ -33,7 +33,7 @@ from misc.dataset import TextEncoder
 import requests 
 from io import BytesIO
 from translate import Translator
-
+from torchvision import transforms
 
 device = torch.device("cpu")
 batch_size = 1
@@ -74,13 +74,13 @@ def search(mode, image, text):
         _stack = np.vstack(caps_enc)
         
     elif mode == I2I:
-        dataset = torch.Tensor(image).unsqueeze(dim=0)
+        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
             with torch.no_grad():
-                _, output_emb = join_emb(input_imgs, None, length)
+                output_emb, _ = join_emb(input_imgs, None, None)
             img_enc.append(output_emb)
         _stack = np.vstack(img_enc)
 
@@ -118,12 +118,15 @@ if __name__ == "__main__":
     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]
+
+    normalize = transforms.Normalize(mean=[0.485 * 255, 0.456 * 255, 0.406 * 255], std=[0.229 * 255, 0.224 * 255, 0.225 * 255])
+
     print("prepare done!")
     iface = gr.Interface(
         fn=search,
         inputs=[
             gr.inputs.Radio([I2I, T2I]),
-            gr.inputs.Image(shape=(512, 512), label="Image to search", optional=True),
+            gr.inputs.Image(shape=(400, 400), label="Image to search", optional=True),
             gr.inputs.Textbox(
                 lines=1, label="Text query", placeholder="Introduce the search text...",
             ),