Add script for converting unaligned dataset to aligned dataset

1 files changed, 63 insertions, 0 deletions

diff --git a/make_dataset_aligned.py b/make_dataset_aligned.py
new file mode 100644
index 0000000..739c767
--- /dev/null
+++ b/make_dataset_aligned.py
@@ -0,0 +1,63 @@
+import os
+
+from PIL import Image
+
+
+def get_file_paths(folder):
+    image_file_paths = []
+    for root, dirs, filenames in os.walk(folder):
+        filenames = sorted(filenames)
+        for filename in filenames:
+            input_path = os.path.abspath(root)
+            file_path = os.path.join(input_path, filename)
+            if filename.endswith('.png') or filename.endswith('.jpg'):
+                image_file_paths.append(file_path)
+
+        break  # prevent descending into subfolders
+    return image_file_paths
+
+
+def align_images(a_file_paths, b_file_paths, target_path):
+    if not os.path.exists(target_path):
+        os.makedirs(target_path)
+
+    for i in range(len(a_file_paths)):
+        img_a = Image.open(a_file_paths[i])
+        img_b = Image.open(b_file_paths[i])
+        assert(img_a.size == img_b.size)
+
+        aligned_image = Image.new("RGB", (img_a.size[0] * 2, img_a.size[1]))
+        aligned_image.paste(img_a, (0, 0))
+        aligned_image.paste(img_b, (img_a.size[0], 0))
+        aligned_image.save(os.path.join(target_path, '{:04d}.jpg'.format(i)))
+
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--dataset-path',
+        dest='dataset_path',
+        help='Which folder to process (it should have subfolders testA, testB, trainA and trainB'
+    )
+    args = parser.parse_args()
+
+    dataset_folder = args.dataset_path
+    print(dataset_folder)
+
+    test_a_path = os.path.join(dataset_folder, 'testA')
+    test_b_path = os.path.join(dataset_folder, 'testB')
+    test_a_file_paths = get_file_paths(test_a_path)
+    test_b_file_paths = get_file_paths(test_b_path)
+    assert(len(test_a_file_paths) == len(test_b_file_paths))
+    test_path = os.path.join(dataset_folder, 'test')
+
+    train_a_path = os.path.join(dataset_folder, 'trainA')
+    train_b_path = os.path.join(dataset_folder, 'trainB')
+    train_a_file_paths = get_file_paths(train_a_path)
+    train_b_file_paths = get_file_paths(train_b_path)
+    assert(len(train_a_file_paths) == len(train_b_file_paths))
+    train_path = os.path.join(dataset_folder, 'train')
+
+    align_images(test_a_file_paths, test_b_file_paths, test_path)
+    align_images(train_a_file_paths, train_b_file_paths, train_path)