From 704ff80bbaf6ca5297b136cb715c18534aa42d82 Mon Sep 17 00:00:00 2001 From: deeppomf Date: Sat, 10 Feb 2018 22:19:48 -0500 Subject: [PATCH] update --- decensor.py | 95 +++++++++++++++++++ layer.py | 95 +++++++++++++++++++ load.py | 14 +++ model.py | 162 ++++++++++++++++++++++++++++++++ models/.gitkeep | 0 test.py | 95 +++++++++++++++++++ train.py | 140 +++++++++++++++++++++++++++ training_data/.gitignore | 3 + training_data/to_npy.py | 34 +++++++ training_output_images/.gitkeep | 0 10 files changed, 638 insertions(+) create mode 100644 decensor.py create mode 100644 layer.py create mode 100644 load.py create mode 100644 model.py create mode 100644 models/.gitkeep create mode 100644 test.py create mode 100644 train.py create mode 100644 training_data/.gitignore create mode 100644 training_data/to_npy.py create mode 100644 training_output_images/.gitkeep diff --git a/decensor.py b/decensor.py new file mode 100644 index 0000000..fdc0985 --- /dev/null +++ b/decensor.py @@ -0,0 +1,95 @@ +import numpy as np +import tensorflow as tf +from PIL import Image +import tqdm +import os +import matplotlib.pyplot as plt +import sys +sys.path.append('..') +from model import Model + +IMAGE_SIZE = 128 +LOCAL_SIZE = 64 +HOLE_MIN = 24 +HOLE_MAX = 48 +BATCH_SIZE = 16 + +image_path = './lfw.npy' + +def test(): + x = 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, mask, local_x, global_completion, local_completion, is_training, batch_size=BATCH_SIZE) + sess = tf.Session() + init_op = tf.global_variables_initializer() + sess.run(init_op) + + saver = tf.train.Saver() + saver.restore(sess, './models/latest') + + x_test = np.load(test_npy) + np.random.shuffle(x_test) + x_test = np.array([a / 127.5 - 1 for a in x_test]) + + step_num = int(len(x_test) / BATCH_SIZE) + + cnt = 0 + for i in tqdm.tqdm(range(step_num)): + x_batch = x_test[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] + _, mask_batch = get_points() + completion = sess.run(model.completion, feed_dict={x: x_batch, mask: mask_batch, is_training: False}) + for i in range(BATCH_SIZE): + cnt += 1 + raw = x_batch[i] + raw = np.array((raw + 1) * 127.5, dtype=np.uint8) + masked = raw * (1 - mask_batch[i]) + np.ones_like(raw) * mask_batch[i] * 255 + img = completion[i] + img = np.array((img + 1) * 127.5, dtype=np.uint8) + dst = './output/{}.jpg'.format("{0:06d}".format(cnt)) + output_image([['Input', masked], ['Output', img], ['Ground Truth', raw]], dst) + + +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 output_image(images, dst): + fig = plt.figure() + for i, image in enumerate(images): + text, img = image + fig.add_subplot(1, 3, i + 1) + plt.imshow(img) + plt.tick_params(labelbottom='off') + plt.tick_params(labelleft='off') + plt.gca().get_xaxis().set_ticks_position('none') + plt.gca().get_yaxis().set_ticks_position('none') + plt.xlabel(text) + plt.savefig(dst) + plt.close() + + +if __name__ == '__main__': + test() + diff --git a/layer.py b/layer.py new file mode 100644 index 0000000..59e3ad8 --- /dev/null +++ b/layer.py @@ -0,0 +1,95 @@ +import tensorflow as tf + +def conv_layer(x, filter_shape, stride): + filters = tf.get_variable( + name='weight', + shape=filter_shape, + dtype=tf.float32, + initializer=tf.contrib.layers.xavier_initializer(), + trainable=True) + return tf.nn.conv2d(x, filters, [1, stride, stride, 1], padding='SAME') + + +def dilated_conv_layer(x, filter_shape, dilation): + filters = tf.get_variable( + name='weight', + shape=filter_shape, + dtype=tf.float32, + initializer=tf.contrib.layers.xavier_initializer(), + trainable=True) + return tf.nn.atrous_conv2d(x, filters, dilation, padding='SAME') + + +def deconv_layer(x, filter_shape, output_shape, stride): + filters = tf.get_variable( + name='weight', + shape=filter_shape, + dtype=tf.float32, + initializer=tf.contrib.layers.xavier_initializer(), + trainable=True) + return tf.nn.conv2d_transpose(x, filters, output_shape, [1, stride, stride, 1]) + + +def batch_normalize(x, is_training, decay=0.99, epsilon=0.001): + def bn_train(): + batch_mean, batch_var = tf.nn.moments(x, axes=[0, 1, 2]) + train_mean = tf.assign(pop_mean, pop_mean * decay + batch_mean * (1 - decay)) + train_var = tf.assign(pop_var, pop_var * decay + batch_var * (1 - decay)) + with tf.control_dependencies([train_mean, train_var]): + return tf.nn.batch_normalization(x, batch_mean, batch_var, beta, scale, epsilon) + + def bn_inference(): + return tf.nn.batch_normalization(x, pop_mean, pop_var, beta, scale, epsilon) + + dim = x.get_shape().as_list()[-1] + beta = tf.get_variable( + name='beta', + shape=[dim], + dtype=tf.float32, + initializer=tf.truncated_normal_initializer(stddev=0.0), + trainable=True) + scale = tf.get_variable( + name='scale', + shape=[dim], + dtype=tf.float32, + initializer=tf.truncated_normal_initializer(stddev=0.1), + trainable=True) + pop_mean = tf.get_variable( + name='pop_mean', + shape=[dim], + dtype=tf.float32, + initializer=tf.constant_initializer(0.0), + trainable=False) + pop_var = tf.get_variable( + name='pop_var', + shape=[dim], + dtype=tf.float32, + initializer=tf.constant_initializer(1.0), + trainable=False) + + return tf.cond(is_training, bn_train, bn_inference) + + +def flatten_layer(x): + input_shape = x.get_shape().as_list() + dim = input_shape[1] * input_shape[2] * input_shape[3] + transposed = tf.transpose(x, (0, 3, 1, 2)) + return tf.reshape(transposed, [-1, dim]) + + +def full_connection_layer(x, out_dim): + in_dim = x.get_shape().as_list()[-1] + W = tf.get_variable( + name='weight', + shape=[in_dim, out_dim], + dtype=tf.float32, + initializer=tf.truncated_normal_initializer(stddev=0.1), + trainable=True) + b = tf.get_variable( + name='bias', + shape=[out_dim], + dtype=tf.float32, + initializer=tf.constant_initializer(0.0), + trainable=True) + return tf.add(tf.matmul(x, W), b) + diff --git a/load.py b/load.py new file mode 100644 index 0000000..b84a94e --- /dev/null +++ b/load.py @@ -0,0 +1,14 @@ +import os +import numpy as np + +def load(dir_='./training_data/npy'): + x_train = np.load(os.path.join(dir_, 'x_train.npy')) + x_test = np.load(os.path.join(dir_, 'x_test.npy')) + return x_train, x_test + + +if __name__ == '__main__': + x_train, x_test = load() + print(x_train.shape) + print(x_test.shape) + diff --git a/model.py b/model.py new file mode 100644 index 0000000..f95407d --- /dev/null +++ b/model.py @@ -0,0 +1,162 @@ +from layer import * + +class Model: + def __init__(self, x, mask, local_x, global_completion, local_completion, is_training, batch_size): + self.batch_size = batch_size + self.imitation = self.generator(x * (1 - mask), 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/models/.gitkeep b/models/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/test.py b/test.py new file mode 100644 index 0000000..d9dc2ae --- /dev/null +++ b/test.py @@ -0,0 +1,95 @@ +import numpy as np +import tensorflow as tf +from PIL import Image +import tqdm +import os +import matplotlib.pyplot as plt +import sys +sys.path.append('..') +from model import Model + +IMAGE_SIZE = 128 +LOCAL_SIZE = 64 +HOLE_MIN = 24 +HOLE_MAX = 48 +BATCH_SIZE = 16 + +test_npy = './lfw.npy' + +def test(): + x = 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, mask, local_x, global_completion, local_completion, is_training, batch_size=BATCH_SIZE) + sess = tf.Session() + init_op = tf.global_variables_initializer() + sess.run(init_op) + + saver = tf.train.Saver() + saver.restore(sess, './models/latest') + + x_test = np.load(test_npy) + np.random.shuffle(x_test) + x_test = np.array([a / 127.5 - 1 for a in x_test]) + + step_num = int(len(x_test) / BATCH_SIZE) + + cnt = 0 + for i in tqdm.tqdm(range(step_num)): + x_batch = x_test[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] + _, mask_batch = get_points() + completion = sess.run(model.completion, feed_dict={x: x_batch, mask: mask_batch, is_training: False}) + for i in range(BATCH_SIZE): + cnt += 1 + raw = x_batch[i] + raw = np.array((raw + 1) * 127.5, dtype=np.uint8) + masked = raw * (1 - mask_batch[i]) + np.ones_like(raw) * mask_batch[i] * 255 + img = completion[i] + img = np.array((img + 1) * 127.5, dtype=np.uint8) + dst = './testing_output_images/{}.jpg'.format("{0:06d}".format(cnt)) + output_image([['Input', masked], ['Output', img], ['Ground Truth', raw]], dst) + + +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 output_image(images, dst): + fig = plt.figure() + for i, image in enumerate(images): + text, img = image + fig.add_subplot(1, 3, i + 1) + plt.imshow(img) + plt.tick_params(labelbottom='off') + plt.tick_params(labelleft='off') + plt.gca().get_xaxis().set_ticks_position('none') + plt.gca().get_yaxis().set_ticks_position('none') + plt.xlabel(text) + plt.savefig(dst) + plt.close() + + +if __name__ == '__main__': + test() + diff --git a/train.py b/train.py new file mode 100644 index 0000000..4d406f5 --- /dev/null +++ b/train.py @@ -0,0 +1,140 @@ +import numpy as np +import tensorflow as tf +from PIL import Image +import tqdm +from model import Model +import load + +IMAGE_SIZE = 128 +LOCAL_SIZE = 64 +HOLE_MIN = 24 +HOLE_MAX = 48 +LEARNING_RATE = 1e-3 +BATCH_SIZE = 16 +PRETRAIN_EPOCH = 100 + +def train(): + x = 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, 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() + + _, g_loss = sess.run([g_train_op, model.g_loss], feed_dict={x: x_batch, mask: mask_batch, is_training: True}) + g_loss_value += g_loss + + print('Completion loss: {}'.format(g_loss_value)) + + np.random.shuffle(x_test) + x_batch = x_test[:BATCH_SIZE] + 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)))) + + 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() + + _, g_loss, completion = sess.run([g_train_op, model.g_loss, model.completion], feed_dict={x: x_batch, mask: mask_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] + 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)))) + + 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) + + +if __name__ == '__main__': + train() + diff --git a/training_data/.gitignore b/training_data/.gitignore new file mode 100644 index 0000000..b7ae298 --- /dev/null +++ b/training_data/.gitignore @@ -0,0 +1,3 @@ +images/* +npy/* +!.gitkeep diff --git a/training_data/to_npy.py b/training_data/to_npy.py new file mode 100644 index 0000000..ca4f892 --- /dev/null +++ b/training_data/to_npy.py @@ -0,0 +1,34 @@ +import glob +import os +#import cv2 +from PIL import Image +import numpy as np + +ratio = 0.95 +image_size = 128 + +x = [] +paths = glob.glob('images/*') +for path in paths: + #img = cv2.imread(path) + #img = Image.open(path) + #img = cv2.resize(img, (image_size, image_size)) + #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + #x.append(img) + temp = Image.open(path) + keep = temp.copy() + keep = np.array(keep) + x.append(keep) + temp.close() + +x = np.array(x, dtype=np.uint8) +#np.random.shuffle(x) + +p = int(ratio * len(x)) +x_train = x[:p] +x_test = x[p:] + +if not os.path.exists('./npy'): + os.mkdir('./npy') +np.save('./npy/x_train.npy', x_train) +np.save('./npy/x_test.npy', x_test) \ No newline at end of file diff --git a/training_output_images/.gitkeep b/training_output_images/.gitkeep new file mode 100644 index 0000000..e69de29