98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
"""
|
|
combine_A_and_B.py
|
|
Preprocessing for pix2pix
|
|
|
|
Copyright (c) 2016, Phillip Isola and Jun-Yan Zhu
|
|
All rights reserved.
|
|
|
|
BSD License
|
|
|
|
Redistribution and use in source and binary forms, with or without
|
|
modification, are permitted provided that the following conditions are met:
|
|
|
|
* Redistributions of source code must retain the above copyright notice, this
|
|
list of conditions and the following disclaimer.
|
|
|
|
* Redistributions in binary form must reproduce the above copyright notice,
|
|
this list of conditions and the following disclaimer in the documentation
|
|
and/or other materials provided with the distribution.
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
from multiprocessing import Pool
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
def image_write(path_A, path_B, path_AB):
|
|
im_A = cv2.imread(
|
|
path_A, 1
|
|
) # python2: cv2.CV_LOAD_IMAGE_COLOR; python3: cv2.IMREAD_COLOR
|
|
im_B = cv2.imread(
|
|
path_B, 1
|
|
) # python2: cv2.CV_LOAD_IMAGE_COLOR; python3: cv2.IMREAD_COLOR
|
|
|
|
# im_A = cv2.transpose(im_A)
|
|
# im_B = cv2.transpose(im_B)
|
|
|
|
im_AB = np.concatenate([im_A, im_B], 1)
|
|
cv2.imwrite(path_AB, im_AB)
|
|
|
|
|
|
parser = argparse.ArgumentParser("create image pairs")
|
|
parser.add_argument("--fold_A", dest="fold_A", help="input directory for image A", type=str, default="../dataset/50kshoes_edges")
|
|
parser.add_argument("--fold_B", dest="fold_B", help="input directory for image B", type=str, default="../dataset/50kshoes_jpg")
|
|
parser.add_argument("--fold_AB", dest="fold_AB", help="output directory", type=str, default="../dataset/test_AB")
|
|
parser.add_argument("--num_imgs", dest="num_imgs", help="number of images", type=int, default=1000000)
|
|
parser.add_argument("--use_AB", dest="use_AB", help="if true: (0001_A, 0001_B) to (0001_AB)", action="store_true")
|
|
parser.add_argument("--no_multiprocessing", dest="no_multiprocessing", help="If used, chooses single CPU execution instead of parallel execution", action="store_true", default=False,
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
for arg in vars(args):
|
|
print("[%s] = " % arg, getattr(args, arg))
|
|
|
|
splits = os.listdir(args.fold_A)
|
|
|
|
if not args.no_multiprocessing:
|
|
pool = Pool()
|
|
|
|
for sp in splits:
|
|
img_fold_A = os.path.join(args.fold_A, sp)
|
|
img_fold_B = os.path.join(args.fold_B, sp)
|
|
img_list = os.listdir(img_fold_A)
|
|
if args.use_AB:
|
|
img_list = [img_path for img_path in img_list if "_A." in img_path]
|
|
|
|
num_imgs = min(args.num_imgs, len(img_list))
|
|
print("split = %s, use %d/%d images" % (sp, num_imgs, len(img_list)))
|
|
img_fold_AB = os.path.join(args.fold_AB, sp)
|
|
if not os.path.isdir(img_fold_AB):
|
|
os.makedirs(img_fold_AB)
|
|
print("split = %s, number of images = %d" % (sp, num_imgs))
|
|
for n in range(num_imgs):
|
|
name_A = img_list[n]
|
|
path_A = os.path.join(img_fold_A, name_A)
|
|
if args.use_AB:
|
|
name_B = name_A.replace("_A.", "_B.")
|
|
else:
|
|
name_B = name_A
|
|
path_B = os.path.join(img_fold_B, name_B)
|
|
if os.path.isfile(path_A) and os.path.isfile(path_B):
|
|
name_AB = name_A
|
|
if args.use_AB:
|
|
name_AB = name_AB.replace("_A.", ".") # remove _A
|
|
path_AB = os.path.join(img_fold_AB, name_AB)
|
|
if not args.no_multiprocessing:
|
|
pool.apply_async(image_write, args=(path_A, path_B, path_AB))
|
|
else:
|
|
im_A = cv2.imread(path_A, 1) # python2: cv2.CV_LOAD_IMAGE_COLOR; python3: cv2.IMREAD_COLOR
|
|
im_B = cv2.imread(path_B, 1) # python2: cv2.CV_LOAD_IMAGE_COLOR; python3: cv2.IMREAD_COLOR
|
|
im_AB = np.concatenate([im_A, im_B], 1)
|
|
cv2.imwrite(path_AB, im_AB)
|
|
if not args.no_multiprocessing:
|
|
pool.close()
|
|
pool.join()
|