fix: load size for rectangular images, resize ref image for inference

jolibrain · Nov 3, 2023 · 965e1bf · 965e1bf
1 parent 7a70209
commit 965e1bf
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 5 deletions.
diff --git a/data/online_creation.py b/data/online_creation.py
@@ -42,7 +42,7 @@ def crop_image(
             if len(load_size) == 1:
                 load_size.extend(load_size)
             old_size = img.size
-            img = F.resize(img, load_size)
+            img = F.resize(img, (load_size[1], load_size[0]))
             new_size = img.size
             ratio_x = img.size[0] / old_size[0]
             ratio_y = img.size[1] / old_size[1]

diff --git a/scripts/gen_single_image_diffusion.py b/scripts/gen_single_image_diffusion.py
@@ -447,6 +447,9 @@ def generate(
         """if crop_width > 0 and crop_height > 0:
         mask = resize(mask).clone().detach()"""
     if ref is not None:
+        ref = cv2.resize(
+            ref, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC
+        )
         ref_tensor = tran(ref).clone().detach()
 
     if not cpu:

diff --git a/tests/test_run_diffusion_online.py b/tests/test_run_diffusion_online.py
@@ -22,8 +22,8 @@
     "data_online_creation_mask_delta_A_ratio": [[0.2, 0.2]],
     "data_online_creation_crop_size_B": 420,
     "data_online_creation_crop_delta_B": 50,
-    "data_online_creation_load_size_A": [2500, 1000],
-    "data_online_creation_load_size_B": [2500, 1000],
+    "data_online_creation_load_size_A": [1000, 2500],
+    "data_online_creation_load_size_B": [1000, 2500],
     "train_n_epochs": 1,
     "train_n_epochs_decay": 0,
     "data_max_dataset_size": 10,

diff --git a/tests/test_run_semantic_mask_online.py b/tests/test_run_semantic_mask_online.py
@@ -24,8 +24,8 @@
     "data_online_creation_mask_delta_A_ratio": [[0.2, 0.2]],
     "data_online_creation_crop_size_B": 420,
     "data_online_creation_crop_delta_B": 50,
-    "data_online_creation_load_size_A": [2500, 1000],
-    "data_online_creation_load_size_B": [2500, 1000],
+    "data_online_creation_load_size_A": [1000, 2500],
+    "data_online_creation_load_size_B": [1000, 2500],
     "data_online_context_pixels": 0,
     "train_n_epochs": 1,
     "train_n_epochs_decay": 0,