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