diff options
Diffstat (limited to 'data/image_folder.py')
| -rwxr-xr-x | data/image_folder.py | 14 |
1 files changed, 14 insertions, 0 deletions
diff --git a/data/image_folder.py b/data/image_folder.py index df0141f..ba675a3 100755 --- a/data/image_folder.py +++ b/data/image_folder.py @@ -30,6 +30,20 @@ def make_dataset(dir): return images +def make_aligned_dataset(dir, dir_b): + images = [] + assert os.path.isdir(dir), '%s is not a valid directory' % dir + assert os.path.isdir(dir_b), '%s is not a valid directory' % dir_b + + files_A = [f for f in os.listdir(dir) if os.path.isfile(os.path.join(mypath, f) and is_image_file(f))] + files_B = [f for f in os.listdir(dir_b) if os.path.isfile(os.path.join(mypath, f) and is_image_file(f))] + images = sorted(list(set(files_A).intersection(files_B))) + # path = os.path.join(root, fname) + # images.append(path) + images_A = [os.path.join(dir, f) for f in images] + images_B = [os.path.join(dir_b, f) for f in images] + return images_A, images_B + def default_loader(path): return Image.open(path).convert('RGB') |