diff options
Diffstat (limited to 'data/aligned_dataset.py')
| -rwxr-xr-x | data/aligned_dataset.py | 7 |
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: |