make aligned dataset

1 files changed, 4 insertions, 3 deletions

diff --git a/data/aligned_dataset.py b/data/aligned_dataset.py
index 41468d2..bb6947e 100755
--- a/data/aligned_dataset.py
+++ b/data/aligned_dataset.py
@@ -2,7 +2,7 @@
 ### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
 import os.path
 from data.base_dataset import BaseDataset, get_params, get_transform, normalize
-from data.image_folder import make_dataset
+from data.image_folder import make_dataset, make_aligned_dataset
 from PIL import Image
 
 class AlignedDataset(BaseDataset):
@@ -13,13 +13,14 @@ class AlignedDataset(BaseDataset):
         ### input A (label maps)
         dir_A = '_A' if self.opt.label_nc == 0 else '_label'
         self.dir_A = os.path.join(opt.dataroot, opt.phase + dir_A)
-        self.A_paths = sorted(make_dataset(self.dir_A))
 
         ### input B (real images)
         if opt.isTrain:
             dir_B = '_B' if self.opt.label_nc == 0 else '_img'
             self.dir_B = os.path.join(opt.dataroot, opt.phase + dir_B)  
-            self.B_paths = sorted(make_dataset(self.dir_B))
+            self.A_paths, self.B_paths = make_aligned_dataset(self.dir_A, self.dir_B)
+        else:
+            self.A_paths = sorted(make_dataset(self.dir_A))
 
         ### instance maps
         if not opt.no_instance: