diff --git a/main.py b/main.py old mode 100644 new mode 100755 index 9bcf3b3..c02a17a --- a/main.py +++ b/main.py @@ -10,8 +10,10 @@ parser.add_argument('--epoch_step', dest='epoch_step', type=int, default=100, help='# of epoch to decay lr') parser.add_argument('--batch_size', dest='batch_size', type=int, default=1, help='# images in batch') parser.add_argument('--train_size', dest='train_size', type=int, default=1e8, help='# images used to train') -parser.add_argument('--load_size', dest='load_size', type=int, default=286, help='scale images to this size') -parser.add_argument('--fine_size', dest='fine_size', type=int, default=256, help='then crop to this size') +parser.add_argument('--load_height', dest='load_height', type=int, default=258, help='scale images to this size') +parser.add_argument('--load_width', dest='load_width', type=int, default=514, help='scale images to this size') +parser.add_argument('--fine_height', dest='fine_height', type=int, default=256, help='then crop to this size') +parser.add_argument('--fine_width', dest='fine_width', type=int, default=512, help='then crop to this size') parser.add_argument('--ngf', dest='ngf', type=int, default=64, help='# of gen filters in first conv layer') parser.add_argument('--ndf', dest='ndf', type=int, default=64, help='# of discri filters in first conv layer') parser.add_argument('--input_nc', dest='input_nc', type=int, default=3, help='# of input image channels') @@ -42,8 +44,8 @@ def main(_): if not os.path.exists(args.test_dir): os.makedirs(args.test_dir) - tfconfig = tf.ConfigProto(allow_soft_placement=True) - tfconfig.gpu_options.allow_growth = True + gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction = 0.95, allow_growth = False) + tfconfig = tf.ConfigProto(allow_soft_placement=True, gpu_options = gpu_options) with tf.Session(config=tfconfig) as sess: model = cyclegan(sess, args) model.train(args) if args.phase == 'train' \ diff --git a/main_horse.py b/main_horse.py new file mode 100755 index 0000000..0128780 --- /dev/null +++ b/main_horse.py @@ -0,0 +1,55 @@ +import argparse +import os +import tensorflow as tf +tf.set_random_seed(19) +from model import cyclegan + +parser = argparse.ArgumentParser(description='') +parser.add_argument('--dataset_dir', dest='dataset_dir', default='horse2zebra', help='path of the dataset') +parser.add_argument('--epoch', dest='epoch', type=int, default=200, help='# of epoch') +parser.add_argument('--epoch_step', dest='epoch_step', type=int, default=100, help='# of epoch to decay lr') +parser.add_argument('--batch_size', dest='batch_size', type=int, default=1, help='# images in batch') +parser.add_argument('--train_size', dest='train_size', type=int, default=1e8, help='# images used to train') +parser.add_argument('--load_height', dest='load_height', type=int, default=286, help='scale images to this size') +parser.add_argument('--load_width', dest='load_width', type=int, default=286, help='scale images to this size') +parser.add_argument('--fine_height', dest='fine_height', type=int, default=256, help='then crop to this size') +parser.add_argument('--fine_width', dest='fine_width', type=int, default=256, help='then crop to this size') +parser.add_argument('--ngf', dest='ngf', type=int, default=64, help='# of gen filters in first conv layer') +parser.add_argument('--ndf', dest='ndf', type=int, default=64, help='# of discri filters in first conv layer') +parser.add_argument('--input_nc', dest='input_nc', type=int, default=3, help='# of input image channels') +parser.add_argument('--output_nc', dest='output_nc', type=int, default=3, help='# of output image channels') +parser.add_argument('--lr', dest='lr', type=float, default=0.0002, help='initial learning rate for adam') +parser.add_argument('--beta1', dest='beta1', type=float, default=0.5, help='momentum term of adam') +parser.add_argument('--which_direction', dest='which_direction', default='AtoB', help='AtoB or BtoA') +parser.add_argument('--phase', dest='phase', default='train', help='train, test') +parser.add_argument('--save_freq', dest='save_freq', type=int, default=1000, help='save a model every save_freq iterations') +parser.add_argument('--print_freq', dest='print_freq', type=int, default=100, help='print the debug information every print_freq iterations') # default 100 +parser.add_argument('--continue_train', dest='continue_train', type=bool, default=False, help='if continue training, load the latest model: 1: true, 0: false') +parser.add_argument('--checkpoint_dir', dest='checkpoint_dir', default='./checkpoint', help='models are saved here') +parser.add_argument('--sample_dir', dest='sample_dir', default='./sample', help='sample are saved here') +parser.add_argument('--test_dir', dest='test_dir', default='./test', help='test sample are saved here') +parser.add_argument('--L1_lambda', dest='L1_lambda', type=float, default=10.0, help='weight on L1 term in objective') +parser.add_argument('--use_resnet', dest='use_resnet', type=bool, default=True, help='generation network using reidule block') +parser.add_argument('--use_lsgan', dest='use_lsgan', type=bool, default=True, help='gan loss defined in lsgan') +parser.add_argument('--max_size', dest='max_size', type=int, default=50, help='max size of image pool, 0 means do not use image pool') + +args = parser.parse_args() + + +def main(_): + if not os.path.exists(args.checkpoint_dir): + os.makedirs(args.checkpoint_dir) + if not os.path.exists(args.sample_dir): + os.makedirs(args.sample_dir) + if not os.path.exists(args.test_dir): + os.makedirs(args.test_dir) + + gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction = 0.95, allow_growth = False) + tfconfig = tf.ConfigProto(allow_soft_placement=True, gpu_options = gpu_options) + with tf.Session(config=tfconfig) as sess: + model = cyclegan(sess, args) + model.train(args) if args.phase == 'train' \ + else model.test(args) + +if __name__ == '__main__': + tf.app.run() diff --git a/model.py b/model.py old mode 100644 new mode 100755 index b04d0e8..39e207a --- a/model.py +++ b/model.py @@ -14,7 +14,8 @@ class cyclegan(object): def __init__(self, sess, args): self.sess = sess self.batch_size = args.batch_size - self.image_size = args.fine_size + self.image_height = args.fine_height + self.image_width = args.fine_width self.input_c_dim = args.input_nc self.output_c_dim = args.output_nc self.L1_lambda = args.L1_lambda @@ -30,9 +31,9 @@ def __init__(self, sess, args): else: self.criterionGAN = sce_criterion - OPTIONS = namedtuple('OPTIONS', 'batch_size image_size \ + OPTIONS = namedtuple('OPTIONS', 'batch_size image_height image_width \ gf_dim df_dim output_c_dim is_training') - self.options = OPTIONS._make((args.batch_size, args.fine_size, + self.options = OPTIONS._make((args.batch_size, args.fine_height, args.fine_width, args.ngf, args.ndf, args.output_nc, args.phase == 'train')) @@ -42,7 +43,7 @@ def __init__(self, sess, args): def _build_model(self): self.real_data = tf.placeholder(tf.float32, - [None, self.image_size, self.image_size, + [None, self.image_height, self.image_width, self.input_c_dim + self.output_c_dim], name='real_A_and_B_images') @@ -68,10 +69,10 @@ def _build_model(self): + self.L1_lambda * abs_criterion(self.real_B, self.fake_B_) self.fake_A_sample = tf.placeholder(tf.float32, - [None, self.image_size, self.image_size, + [None, self.image_height, self.image_width, self.input_c_dim], name='fake_A_sample') self.fake_B_sample = tf.placeholder(tf.float32, - [None, self.image_size, self.image_size, + [None, self.image_height, self.image_width, self.output_c_dim], name='fake_B_sample') self.DB_real = self.discriminator(self.real_B, self.options, reuse=True, name="discriminatorB") self.DA_real = self.discriminator(self.real_A, self.options, reuse=True, name="discriminatorA") @@ -104,10 +105,10 @@ def _build_model(self): ) self.test_A = tf.placeholder(tf.float32, - [None, self.image_size, self.image_size, + [None, self.image_height, self.image_width, self.input_c_dim], name='test_A') self.test_B = tf.placeholder(tf.float32, - [None, self.image_size, self.image_size, + [None, self.image_height, self.image_width, self.output_c_dim], name='test_B') self.testB = self.generator(self.test_A, self.options, True, name="generatorA2B") self.testA = self.generator(self.test_B, self.options, True, name="generatorB2A") @@ -115,7 +116,7 @@ def _build_model(self): t_vars = tf.trainable_variables() self.d_vars = [var for var in t_vars if 'discriminator' in var.name] self.g_vars = [var for var in t_vars if 'generator' in var.name] - for var in t_vars: print(var.name) + def train(self, args): """Train cyclegan""" @@ -127,7 +128,11 @@ def train(self, args): init_op = tf.global_variables_initializer() self.sess.run(init_op) - self.writer = tf.summary.FileWriter("./logs", self.sess.graph) + + log_dir = "./logs/" + self.dataset_dir + if not os.path.exists(log_dir): + os.makedirs(log_dir) + self.writer = tf.summary.FileWriter(log_dir, self.sess.graph) counter = 1 start_time = time.time() @@ -139,6 +144,7 @@ def train(self, args): print(" [!] Load failed...") for epoch in range(args.epoch): + print('The %d epoch starts...'%(epoch+1)) dataA = glob('./datasets/{}/*.*'.format(self.dataset_dir + '/trainA')) dataB = glob('./datasets/{}/*.*'.format(self.dataset_dir + '/trainB')) np.random.shuffle(dataA) @@ -149,7 +155,7 @@ def train(self, args): for idx in range(0, batch_idxs): batch_files = list(zip(dataA[idx * self.batch_size:(idx + 1) * self.batch_size], dataB[idx * self.batch_size:(idx + 1) * self.batch_size])) - batch_images = [load_train_data(batch_file, args.load_size, args.fine_size) for batch_file in batch_files] + batch_images = [load_train_data(batch_file, args.load_height, args.load_width, args.fine_height, args.fine_width) for batch_file in batch_files] batch_images = np.array(batch_images).astype(np.float32) # Update G network and record fake outputs @@ -173,14 +179,14 @@ def train(self, args): epoch, idx, batch_idxs, time.time() - start_time))) if np.mod(counter, args.print_freq) == 1: - self.sample_model(args.sample_dir, epoch, idx) + self.sample_model(args, args.sample_dir, epoch, idx) if np.mod(counter, args.save_freq) == 2: self.save(args.checkpoint_dir, counter) def save(self, checkpoint_dir, step): model_name = "cyclegan.model" - model_dir = "%s_%s" % (self.dataset_dir, self.image_size) + model_dir = "%s_%s_%s" % (self.dataset_dir, self.image_height, self.image_width) checkpoint_dir = os.path.join(checkpoint_dir, model_dir) if not os.path.exists(checkpoint_dir): @@ -193,7 +199,7 @@ def save(self, checkpoint_dir, step): def load(self, checkpoint_dir): print(" [*] Reading checkpoint...") - model_dir = "%s_%s" % (self.dataset_dir, self.image_size) + model_dir = "%s_%s_%s" % (self.dataset_dir, self.image_height, self.image_width) checkpoint_dir = os.path.join(checkpoint_dir, model_dir) ckpt = tf.train.get_checkpoint_state(checkpoint_dir) @@ -204,23 +210,24 @@ def load(self, checkpoint_dir): else: return False - def sample_model(self, sample_dir, epoch, idx): - dataA = glob('./datasets/{}/*.*'.format(self.dataset_dir + '/testA')) - dataB = glob('./datasets/{}/*.*'.format(self.dataset_dir + '/testB')) + def sample_model(self, args, sample_dir, epoch, idx): + sample_dir_new = sample_dir + '/' + self.dataset_dir + if not os.path.exists(sample_dir_new): + os.makedirs(sample_dir_new) + + dataA = glob('./datasets/{}/*.*'.format(self.dataset_dir + '/trainA')) + dataB = glob('./datasets/{}/*.*'.format(self.dataset_dir + '/trainB')) np.random.shuffle(dataA) np.random.shuffle(dataB) - batch_files = list(zip(dataA[:self.batch_size], dataB[:self.batch_size])) - sample_images = [load_train_data(batch_file, is_testing=True) for batch_file in batch_files] - sample_images = np.array(sample_images).astype(np.float32) - - fake_A, fake_B = self.sess.run( - [self.fake_A, self.fake_B], - feed_dict={self.real_data: sample_images} - ) - save_images(fake_A, [self.batch_size, 1], - './{}/A_{:02d}_{:04d}.jpg'.format(sample_dir, epoch, idx)) - save_images(fake_B, [self.batch_size, 1], - './{}/B_{:02d}_{:04d}.jpg'.format(sample_dir, epoch, idx)) + for ind in range(5): # display 5 result images in each epoch + batch_files = list(zip(dataA[ind * self.batch_size:(ind + 1) * self.batch_size], dataB[ind * self.batch_size:(ind + 1) * self.batch_size])) + sample_images = [load_train_data(batch_file, args.load_height, args.load_width, args.fine_height, args.fine_width, is_testing=True) for batch_file in batch_files] + sample_images = np.array(sample_images).astype(np.float32) + fake_A, fake_B = self.sess.run([self.fake_A, self.fake_B], feed_dict={self.real_data: sample_images}) + real_A = sample_images[:, :, :, :self.input_c_dim] + real_B = sample_images[:, :, :, self.input_c_dim:self.input_c_dim + self.output_c_dim] + concat_img = np.concatenate([real_A, real_B, fake_B, fake_A], 2) + save_images(concat_img, [self.batch_size, 1], '{}/output_{:03d}_{:02d}.jpg'.format(sample_dir_new, epoch, ind)) def test(self, args): """Test cyclegan""" @@ -249,7 +256,7 @@ def test(self, args): for sample_file in sample_files: print('Processing image: ' + sample_file) - sample_image = [load_test_data(sample_file, args.fine_size)] + sample_image = [load_test_data(sample_file, args.fine_height, args.fine_width)] sample_image = np.array(sample_image).astype(np.float32) image_path = os.path.join(args.test_dir, '{0}_{1}'.format(args.which_direction, os.path.basename(sample_file))) diff --git a/train_horse.sh b/train_horse.sh new file mode 100755 index 0000000..0a46fbd --- /dev/null +++ b/train_horse.sh @@ -0,0 +1 @@ +CUDA_VISIBLE_DEVICES=0 python3 main_horse.py --dataset_dir=horse2zebra diff --git a/utils.py b/utils.py old mode 100644 new mode 100755 index 8a54287..1a1ccdf --- a/utils.py +++ b/utils.py @@ -42,35 +42,36 @@ def __call__(self, image): else: return image -def load_test_data(image_path, fine_size=256): +def load_test_data(image_path, fine_height=256, fine_width=256): img = imread(image_path) - img = scipy.misc.imresize(img, [fine_size, fine_size]) + img = scipy.misc.imresize(img, [fine_height, fine_width]) img = img/127.5 - 1 return img -def load_train_data(image_path, load_size=286, fine_size=256, is_testing=False): +def load_train_data(image_path, load_height = 286, load_width = 286, fine_height = 256, fine_width = 256, is_testing=False): img_A = imread(image_path[0]) img_B = imread(image_path[1]) if not is_testing: - img_A = scipy.misc.imresize(img_A, [load_size, load_size]) - img_B = scipy.misc.imresize(img_B, [load_size, load_size]) - h1 = int(np.ceil(np.random.uniform(1e-2, load_size-fine_size))) - w1 = int(np.ceil(np.random.uniform(1e-2, load_size-fine_size))) - img_A = img_A[h1:h1+fine_size, w1:w1+fine_size] - img_B = img_B[h1:h1+fine_size, w1:w1+fine_size] + img_A = scipy.misc.imresize(img_A, [load_height, load_width]) + img_B = scipy.misc.imresize(img_B, [load_height, load_width]) + h1 = int(np.ceil(np.random.uniform(1e-2, load_height-fine_height))) + w1 = int(np.ceil(np.random.uniform(1e-2, load_width-fine_width))) + img_A = img_A[h1:h1+fine_height, w1:w1+fine_width] + img_B = img_B[h1:h1+fine_height, w1:w1+fine_width] if np.random.random() > 0.5: img_A = np.fliplr(img_A) img_B = np.fliplr(img_B) else: - img_A = scipy.misc.imresize(img_A, [fine_size, fine_size]) - img_B = scipy.misc.imresize(img_B, [fine_size, fine_size]) + img_A = scipy.misc.imresize(img_A, [fine_height, fine_width]) + img_B = scipy.misc.imresize(img_B, [fine_height, fine_width]) img_A = img_A/127.5 - 1. img_B = img_B/127.5 - 1. img_AB = np.concatenate((img_A, img_B), axis=2) - # img_AB shape: (fine_size, fine_size, input_c_dim + output_c_dim) + + # img_AB shape: (fine_height, fine_width, input_c_dim + output_c_dim) return img_AB # -----------------------------