diff options
Diffstat (limited to 'make_dataset_aligned.py')
| -rw-r--r-- | make_dataset_aligned.py | 63 |
1 files changed, 0 insertions, 63 deletions
diff --git a/make_dataset_aligned.py b/make_dataset_aligned.py deleted file mode 100644 index 739c767..0000000 --- a/make_dataset_aligned.py +++ /dev/null @@ -1,63 +0,0 @@ -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) |
