make aligned dataset

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')