diff --git a/config.py b/config.py index 9b8964e..48aa45c 100644 --- a/config.py +++ b/config.py @@ -18,17 +18,18 @@ parser.add_argument('--local_input_size', dest='local_input_size', default=64, h parser.add_argument('--input_channel_size', dest='input_channel_size', default=3, help='input image channel') parser.add_argument('--min_mask_size', dest='min_mask_size', default=24, help='minimum mask size') parser.add_argument('--max_mask_size', dest='max_mask_size', default=48, help='maximum mask size') -parser.add_argument('--rotate_chance', dest='rotate_chance', default=0.5, help='chance the mask will be randomly rotated') +parser.add_argument('--rotate_chance', dest='rotate_chance', default=0.7, help='chance the mask will be randomly rotated') parser.add_argument('--train_mosaic', dest ='train_mosaic', default=False, help='train neural network to decensor mosaics') # parser.add_argument('--input_dim', dest='input_dim', default=100, help='input z size') # #Training Settings parser.add_argument('--continue_training', dest='continue_training', default=False, type=str2bool, help='flag to continue training') +parser.add_argument('--training_samples_path', dest='training_samples_path', default='./training_samples/', help='samples images generated during training path') +parser.add_argument('--batch_size', dest='batch_size', default=16, help='batch size') # parser.add_argument('--data', dest='data', default='../ambientGAN_TF/data', help='cats image train path') -parser.add_argument('--batch_size', dest='batch_size', default=16, help='batch size') # parser.add_argument('--train_step', dest='train_step', default=400, help='total number of train_step') # parser.add_argument('--Tc', dest='Tc', default=100, help='Tc to train Completion Network') # parser.add_argument('--Td', dest='Td', default=1, help='Td to train Discriminator Network') @@ -49,7 +50,9 @@ parser.add_argument('--learning_rate', dest='learning_rate', default=0.001, help # parser.add_argument('--checkpoints_path', dest='checkpoints_path', default='./checkpoints/', help='saved model checkpoint path') # parser.add_argument('--graph_path', dest='graph_path', default='./graphs/', help='tensorboard graph') # parser.add_argument('--images_path', dest='images_path', default='./images/', help='result images path') -parser.add_argument('--training_samples_path', dest='training_samples_path', default='./training_samples/', help='samples images generated during training path') parser.add_argument('--testing_output_path', dest='testing_output_path', default='./testing_output/', help='output images generated from running test.py path') +parser.add_argument('--decensor_input_path', dest='decensor_input_path', default='./decensor_input/', help='input images to be decensored by decensor.py path') +parser.add_argument('--decensor_output_path', dest='decensor_output_path', default='./decensor_output/', help='output images generated from running decensor.py path') + args = parser.parse_args() \ No newline at end of file diff --git a/decensor.py b/decensor.py index c7054d0..48f5c65 100644 --- a/decensor.py +++ b/decensor.py @@ -13,16 +13,15 @@ from config import * BATCH_SIZE = 1 -image_folder = 'decensor_input_images/' mask_color = [0, 255, 0] poisson_blending_enabled = False def decensor(args): - x = tf.placeholder(tf.float32, [args.batch_size, args.image_size, args.image_size, args.input_channel_size]) - mask = tf.placeholder(tf.float32, [args.batch_size, args.image_size, args.image_size, 1]) - local_x = tf.placeholder(tf.float32, [args.batch_size, args.local_image_size, args.local_image_size, args.input_channel_size]) - global_completion = tf.placeholder(tf.float32, [args.batch_size, args.image_size, args.image_size, args.input_channel_size]) - local_completion = tf.placeholder(tf.float32, [args.batch_size, args.local_image_size, args.local_image_size, args.input_channel_size]) + x = tf.placeholder(tf.float32, [BATCH_SIZE, args.input_size, args.input_size, args.input_channel_size]) + mask = tf.placeholder(tf.float32, [BATCH_SIZE, args.input_size, args.input_size, 1]) + local_x = tf.placeholder(tf.float32, [BATCH_SIZE, args.local_input_size, args.local_input_size, args.input_channel_size]) + global_completion = tf.placeholder(tf.float32, [BATCH_SIZE, args.input_size, args.input_size, args.input_channel_size]) + local_completion = tf.placeholder(tf.float32, [BATCH_SIZE, args.local_input_size, args.local_input_size, args.input_channel_size]) is_training = tf.placeholder(tf.bool, []) model = Model(x, mask, local_x, global_completion, local_completion, is_training, batch_size=BATCH_SIZE) @@ -35,7 +34,7 @@ def decensor(args): x_decensor = [] mask_decensor = [] - for subdir, dirs, files in sorted(os.walk(image_folder)): + for subdir, dirs, files in sorted(os.walk(args.decensor_input_path)): for file in sorted(files): file_path = os.path.join(subdir, file) if os.path.isfile(file_path) and os.path.splitext(file_path)[1] == ".png": @@ -54,7 +53,6 @@ def decensor(args): mask_batch = get_mask(x_batch) completion = sess.run(model.completion, feed_dict={x: x_batch, mask: mask_batch, is_training: False}) for i in range(BATCH_SIZE): - cnt += 1 img = completion[i] img = np.array((img + 1) * 127.5, dtype=np.uint8) original = x_batch[i] @@ -62,8 +60,9 @@ def decensor(args): if (poisson_blending_enabled): img = blend(original, img, mask_batch[0,:,:,0]) output = Image.fromarray(img.astype('uint8'), 'RGB') - dst = './decensor_output_images/{}.png'.format("{0:06d}".format(cnt)) + dst = args.decensor_output_path + '{}.png'.format("{0:06d}".format(cnt)) output.save(dst) + cnt += 1 def get_mask(x_batch): points = [] @@ -71,14 +70,16 @@ def get_mask(x_batch): for i in range(BATCH_SIZE): raw = x_batch[i] raw = np.array((raw + 1) * 127.5, dtype=np.uint8) - m = np.zeros((args.image_size, args.image_size, 1), dtype=np.uint8) - for x in range(args.image_size): - for y in range(args.image_size): + m = np.zeros((args.input_size, args.input_size, 1), dtype=np.uint8) + for x in range(args.input_size): + for y in range(args.input_size): if np.array_equal(raw[x][y], mask_color): m[x, y] = 1 mask.append(m) return np.array(mask) if __name__ == '__main__': + if not os.path.exists(args.decensor_output_path): + os.makedirs(args.decensor_output_path) decensor(args) diff --git a/train.py b/train.py index 3c94270..e00f3d4 100644 --- a/train.py +++ b/train.py @@ -68,7 +68,7 @@ def train(args): completion = sess.run(model.completion, feed_dict={x: x_batch, mask: mask_batch, is_training: False}) sample = np.array((completion[0] + 1) * 127.5, dtype=np.uint8) result = Image.fromarray(sample) - result.save('./training_output_images/{}.jpg'.format("{0:06d}".format(sess.run(epoch)))) + result.save(args.training_samples_path + '/{}.jpg'.format("{0:06d}".format(sess.run(epoch)))) saver = tf.train.Saver() saver.save(sess, './models/latest', write_meta_graph=False)