From 05d81eaad4de764c7e6f1faca1a810675402a27e Mon Sep 17 00:00:00 2001 From: StartleStars <36647583+StartleStars@users.noreply.github.com> Date: Mon, 19 Feb 2018 21:54:52 +0000 Subject: [PATCH] Add files via upload --- model_mosaic.py | 163 ++++++++++++++++++++++++++++++++++++++++++++ train_mosaic.py | 176 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 339 insertions(+) create mode 100644 model_mosaic.py create mode 100644 train_mosaic.py diff --git a/model_mosaic.py b/model_mosaic.py new file mode 100644 index 0000000..c316179 --- /dev/null +++ b/model_mosaic.py @@ -0,0 +1,163 @@ +from layer import * + +class Model: + def __init__(self, x, mosaic, mask, local_x, global_completion, local_completion, is_training, batch_size): + self.batch_size = batch_size + self.merged = x * (1 - mask) + mosaic * (mask) + self.imitation = self.generator(self.merged, is_training) + self.completion = self.imitation * mask + x * (1 - mask) + self.real = self.discriminator(x, local_x, reuse=False) + self.fake = self.discriminator(global_completion, local_completion, reuse=True) + self.g_loss = self.calc_g_loss(x, self.completion) + self.d_loss = self.calc_d_loss(self.real, self.fake) + self.g_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='generator') + self.d_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='discriminator') + + + def generator(self, x, is_training): + with tf.variable_scope('generator'): + with tf.variable_scope('conv1'): + x = conv_layer(x, [5, 5, 3, 64], 1) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv2'): + x = conv_layer(x, [3, 3, 64, 128], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv3'): + x = conv_layer(x, [3, 3, 128, 128], 1) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv4'): + x = conv_layer(x, [3, 3, 128, 256], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv5'): + x = conv_layer(x, [3, 3, 256, 256], 1) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv6'): + x = conv_layer(x, [3, 3, 256, 256], 1) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('dilated1'): + x = dilated_conv_layer(x, [3, 3, 256, 256], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('dilated2'): + x = dilated_conv_layer(x, [3, 3, 256, 256], 4) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('dilated3'): + x = dilated_conv_layer(x, [3, 3, 256, 256], 8) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('dilated4'): + x = dilated_conv_layer(x, [3, 3, 256, 256], 16) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv7'): + x = conv_layer(x, [3, 3, 256, 256], 1) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv8'): + x = conv_layer(x, [3, 3, 256, 256], 1) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('deconv1'): + x = deconv_layer(x, [4, 4, 128, 256], [self.batch_size, 64, 64, 128], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv9'): + x = conv_layer(x, [3, 3, 128, 128], 1) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('deconv2'): + x = deconv_layer(x, [4, 4, 64, 128], [self.batch_size, 128, 128, 64], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv10'): + x = conv_layer(x, [3, 3, 64, 32], 1) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv11'): + x = conv_layer(x, [3, 3, 32, 3], 1) + x = tf.nn.tanh(x) + + return x + + + def discriminator(self, global_x, local_x, reuse): + def global_discriminator(x): + is_training = tf.constant(True) + with tf.variable_scope('global'): + with tf.variable_scope('conv1'): + x = conv_layer(x, [5, 5, 3, 64], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv2'): + x = conv_layer(x, [5, 5, 64, 128], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv3'): + x = conv_layer(x, [5, 5, 128, 256], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv4'): + x = conv_layer(x, [5, 5, 256, 512], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv5'): + x = conv_layer(x, [5, 5, 512, 512], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('fc'): + x = flatten_layer(x) + x = full_connection_layer(x, 1024) + return x + + def local_discriminator(x): + is_training = tf.constant(True) + with tf.variable_scope('local'): + with tf.variable_scope('conv1'): + x = conv_layer(x, [5, 5, 3, 64], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv2'): + x = conv_layer(x, [5, 5, 64, 128], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv3'): + x = conv_layer(x, [5, 5, 128, 256], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('conv4'): + x = conv_layer(x, [5, 5, 256, 512], 2) + x = batch_normalize(x, is_training) + x = tf.nn.relu(x) + with tf.variable_scope('fc'): + x = flatten_layer(x) + x = full_connection_layer(x, 1024) + return x + + with tf.variable_scope('discriminator', reuse=reuse): + global_output = global_discriminator(global_x) + local_output = local_discriminator(local_x) + with tf.variable_scope('concatenation'): + output = tf.concat((global_output, local_output), 1) + output = full_connection_layer(output, 1) + + return output + + + def calc_g_loss(self, x, completion): + loss = tf.nn.l2_loss(x - completion) + return tf.reduce_mean(loss) + + + def calc_d_loss(self, real, fake): + alpha = 4e-4 + d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=real, labels=tf.ones_like(real))) + d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=fake, labels=tf.zeros_like(fake))) + return tf.add(d_loss_real, d_loss_fake) * alpha + diff --git a/train_mosaic.py b/train_mosaic.py new file mode 100644 index 0000000..7f4b630 --- /dev/null +++ b/train_mosaic.py @@ -0,0 +1,176 @@ +import numpy as np +import tensorflow as tf +from PIL import Image, ImageFilter +import tqdm +from model_mosaic import Model +import load + +IMAGE_SIZE = 128 +LOCAL_SIZE = 64 +HOLE_MIN = 24 +HOLE_MAX = 48 +MOSAIC_MIN = 8 #Minimum number of mosaic squares across image +MOSAIC_MAX = 32 #Maximum number of mosaic squares across image +MOSAIC_GAUSSIAN_P = 0.5 #represent images that have been compressed post-mosaic +MOSAIC_GAUSSIAN_MIN = 0.2 +MOSAIC_GAUSSIAN_MAX = 1.2 +LEARNING_RATE = 1e-3 +BATCH_SIZE = 16 +PRETRAIN_EPOCH = 100 + +def train(): + x = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3]) + mosaic = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3]) + mask = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 1]) + local_x = tf.placeholder(tf.float32, [BATCH_SIZE, LOCAL_SIZE, LOCAL_SIZE, 3]) + global_completion = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3]) + local_completion = tf.placeholder(tf.float32, [BATCH_SIZE, LOCAL_SIZE, LOCAL_SIZE, 3]) + is_training = tf.placeholder(tf.bool, []) + + model = Model(x, mosaic, mask, local_x, global_completion, local_completion, is_training, batch_size=BATCH_SIZE) + sess = tf.Session() + global_step = tf.Variable(0, name='global_step', trainable=False) + epoch = tf.Variable(0, name='epoch', trainable=False) + + opt = tf.train.AdamOptimizer(learning_rate=LEARNING_RATE) + g_train_op = opt.minimize(model.g_loss, global_step=global_step, var_list=model.g_variables) + d_train_op = opt.minimize(model.d_loss, global_step=global_step, var_list=model.d_variables) + + init_op = tf.global_variables_initializer() + sess.run(init_op) + + if tf.train.get_checkpoint_state('./models'): + saver = tf.train.Saver() + saver.restore(sess, './models/latest') + + x_train, x_test = load.load() + x_train = np.array([a / 127.5 - 1 for a in x_train]) + x_test = np.array([a / 127.5 - 1 for a in x_test]) + + step_num = int(len(x_train) / BATCH_SIZE) + + while True: + sess.run(tf.assign(epoch, tf.add(epoch, 1))) + print('epoch: {}'.format(sess.run(epoch))) + + np.random.shuffle(x_train) + + # Completion + if sess.run(epoch) <= PRETRAIN_EPOCH: + g_loss_value = 0 + for i in tqdm.tqdm(range(step_num)): + x_batch = x_train[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] + points_batch, mask_batch = get_points() + mosaic_batch = get_mosaic(x_batch) + + _, g_loss = sess.run([g_train_op, model.g_loss], feed_dict={x: x_batch, mask: mask_batch, mosaic: mosaic_batch, is_training: True}) + g_loss_value += g_loss + + print('Completion loss: {}'.format(g_loss_value)) + + f = open("loss.csv","a+") + f.write(str(sess.run(epoch)) + "," + str(g_loss_value) + "," + "0" + "\n") + f.close() + + np.random.shuffle(x_test) + x_batch = x_test[:BATCH_SIZE] + mosaic_batch = get_mosaic(x_batch) + merged, completion = sess.run([model.merged, model.completion], feed_dict={x: x_batch, mask: mask_batch, mosaic: mosaic_batch, is_training: False}) + sample = np.array((merged[0] + 1) * 127.5, dtype=np.uint8) + result = Image.fromarray(sample) + result.save('./training_output_images/{}_0.png'.format("{0:06d}".format(sess.run(epoch)))) + sample = np.array((completion[0] + 1) * 127.5, dtype=np.uint8) + result = Image.fromarray(sample) + result.save('./training_output_images/{}_1.png'.format("{0:06d}".format(sess.run(epoch)))) + + saver = tf.train.Saver() + saver.save(sess, './models/latest', write_meta_graph=False) + if sess.run(epoch) == PRETRAIN_EPOCH: + saver.save(sess, './models/pretrained', write_meta_graph=False) + + + # Discrimitation + else: + g_loss_value = 0 + d_loss_value = 0 + for i in tqdm.tqdm(range(step_num)): + x_batch = x_train[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] + points_batch, mask_batch = get_points() + mosaic_batch = get_mosaic(x_batch) + + _, g_loss, completion = sess.run([g_train_op, model.g_loss, model.completion], feed_dict={x: x_batch, mask: mask_batch, mosaic: mosaic_batch, is_training: True}) + g_loss_value += g_loss + + local_x_batch = [] + local_completion_batch = [] + for i in range(BATCH_SIZE): + x1, y1, x2, y2 = points_batch[i] + local_x_batch.append(x_batch[i][y1:y2, x1:x2, :]) + local_completion_batch.append(completion[i][y1:y2, x1:x2, :]) + local_x_batch = np.array(local_x_batch) + local_completion_batch = np.array(local_completion_batch) + + _, d_loss = sess.run( + [d_train_op, model.d_loss], + feed_dict={x: x_batch, mask: mask_batch, local_x: local_x_batch, global_completion: completion, local_completion: local_completion_batch, is_training: True}) + d_loss_value += d_loss + + print('Completion loss: {}'.format(g_loss_value)) + print('Discriminator loss: {}'.format(d_loss_value)) + + np.random.shuffle(x_test) + x_batch = x_test[:BATCH_SIZE] + mosaic_batch = get_mosaic(x_batch) + merged, completion = sess.run([model.merged, model.completion], feed_dict={x: x_batch, mask: mask_batch, mosaic: mosaic_batch, is_training: False}) + sample = np.array((merged[0] + 1) * 127.5, dtype=np.uint8) + result = Image.fromarray(sample) + result.save('./training_output_images/{}_0.png'.format("{0:06d}".format(sess.run(epoch)))) + sample = np.array((completion[0] + 1) * 127.5, dtype=np.uint8) + result = Image.fromarray(sample) + result.save('./training_output_images/{}_1.png'.format("{0:06d}".format(sess.run(epoch)))) + + saver = tf.train.Saver() + saver.save(sess, './models/latest', write_meta_graph=False) + + +def get_points(): + points = [] + mask = [] + for i in range(BATCH_SIZE): + x1, y1 = np.random.randint(0, IMAGE_SIZE - LOCAL_SIZE + 1, 2) + x2, y2 = np.array([x1, y1]) + LOCAL_SIZE + points.append([x1, y1, x2, y2]) + + w, h = np.random.randint(HOLE_MIN, HOLE_MAX + 1, 2) + p1 = x1 + np.random.randint(0, LOCAL_SIZE - w) + q1 = y1 + np.random.randint(0, LOCAL_SIZE - h) + p2 = p1 + w + q2 = q1 + h + + m = np.zeros((IMAGE_SIZE, IMAGE_SIZE, 1), dtype=np.uint8) + m[q1:q2 + 1, p1:p2 + 1] = 1 + mask.append(m) + + return np.array(points), np.array(mask) + + +def get_mosaic(x_batch): + mosaic = [] + for i in range(BATCH_SIZE): + im = np.array((x_batch[i] + 1) * 127.5, dtype=np.uint8) + im = Image.fromarray(im) + size = np.random.randint(MOSAIC_MIN, MOSAIC_MAX) + im = im.resize((size,size),Image.LANCZOS) + im = im.resize((IMAGE_SIZE,IMAGE_SIZE),Image.NEAREST) + if np.random.rand() < MOSAIC_GAUSSIAN_P: + im = im.filter(ImageFilter.GaussianBlur(np.random.uniform(MOSAIC_GAUSSIAN_MIN, MOSAIC_GAUSSIAN_MAX))) + + mosaic.append(np.array(im)) + + mosaic = np.array([a / 127.5 - 1 for a in mosaic]) + return mosaic + + +if __name__ == '__main__': + train() +