diff --git a/test_code/cartoonize.py b/test_code/cartoonize.py index f32e962..533b0b1 100755 --- a/test_code/cartoonize.py +++ b/test_code/cartoonize.py @@ -6,7 +6,7 @@ import guided_filter from tqdm import tqdm - +tf.compat.v1.disable_eager_execution() def resize_crop(image): h, w, c = np.shape(image) @@ -23,19 +23,19 @@ def resize_crop(image): def cartoonize(load_folder, save_folder, model_path): - input_photo = tf.placeholder(tf.float32, [1, None, None, 3]) + input_photo = tf.compat.v1.placeholder(tf.float32, [1, None, None, 3]) network_out = network.unet_generator(input_photo) final_out = guided_filter.guided_filter(input_photo, network_out, r=1, eps=5e-3) - all_vars = tf.trainable_variables() + all_vars = tf.compat.v1.trainable_variables() gene_vars = [var for var in all_vars if 'generator' in var.name] - saver = tf.train.Saver(var_list=gene_vars) + saver = tf.compat.v1.train.Saver(var_list=gene_vars) - config = tf.ConfigProto() + config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True - sess = tf.Session(config=config) + sess = tf.compat.v1.Session(config=config) - sess.run(tf.global_variables_initializer()) + sess.run(tf.compat.v1.global_variables_initializer()) saver.restore(sess, tf.train.latest_checkpoint(model_path)) name_list = os.listdir(load_folder) for name in tqdm(name_list): @@ -65,4 +65,3 @@ def cartoonize(load_folder, save_folder, model_path): cartoonize(load_folder, save_folder, model_path) - \ No newline at end of file diff --git a/test_code/guided_filter.py b/test_code/guided_filter.py index fd019d1..0ef511c 100755 --- a/test_code/guided_filter.py +++ b/test_code/guided_filter.py @@ -10,14 +10,14 @@ def tf_box_filter(x, r): weight = 1/(k_size**2) box_kernel = weight*np.ones((k_size, k_size, ch, 1)) box_kernel = np.array(box_kernel).astype(np.float32) - output = tf.nn.depthwise_conv2d(x, box_kernel, [1, 1, 1, 1], 'SAME') + output = tf.nn.depthwise_conv2d(input=x, filter=box_kernel, strides=[1, 1, 1, 1], padding='SAME') return output def guided_filter(x, y, r, eps=1e-2): - x_shape = tf.shape(x) + x_shape = tf.shape(input=x) #y_shape = tf.shape(y) N = tf_box_filter(tf.ones((1, x_shape[1], x_shape[2], 1), dtype=x.dtype), r) @@ -43,9 +43,9 @@ def fast_guided_filter(lr_x, lr_y, hr_x, r=1, eps=1e-8): #assert lr_x.shape.ndims == 4 and lr_y.shape.ndims == 4 and hr_x.shape.ndims == 4 - lr_x_shape = tf.shape(lr_x) + lr_x_shape = tf.shape(input=lr_x) #lr_y_shape = tf.shape(lr_y) - hr_x_shape = tf.shape(hr_x) + hr_x_shape = tf.shape(input=hr_x) N = tf_box_filter(tf.ones((1, lr_x_shape[1], lr_x_shape[2], 1), dtype=lr_x.dtype), r) @@ -57,8 +57,8 @@ def fast_guided_filter(lr_x, lr_y, hr_x, r=1, eps=1e-8): A = cov_xy / (var_x + eps) b = mean_y - A * mean_x - mean_A = tf.image.resize_images(A, hr_x_shape[1: 3]) - mean_b = tf.image.resize_images(b, hr_x_shape[1: 3]) + mean_A = tf.image.resize(A, hr_x_shape[1: 3]) + mean_b = tf.image.resize(b, hr_x_shape[1: 3]) output = mean_A * hr_x + mean_b @@ -69,17 +69,17 @@ def fast_guided_filter(lr_x, lr_y, hr_x, r=1, eps=1e-8): import cv2 from tqdm import tqdm - input_photo = tf.placeholder(tf.float32, [1, None, None, 3]) + input_photo = tf.compat.v1.placeholder(tf.float32, [1, None, None, 3]) #input_superpixel = tf.placeholder(tf.float32, [16, 256, 256, 3]) output = guided_filter(input_photo, input_photo, 5, eps=1) image = cv2.imread('output_figure1/cartoon2.jpg') image = image/127.5 - 1 image = np.expand_dims(image, axis=0) - config = tf.ConfigProto() + config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True - sess = tf.Session(config=config) - sess.run(tf.global_variables_initializer()) + sess = tf.compat.v1.Session(config=config) + sess.run(tf.compat.v1.global_variables_initializer()) out = sess.run(output, feed_dict={input_photo: image}) out = (np.squeeze(out)+1)*127.5 diff --git a/test_code/network.py b/test_code/network.py index 6f16cee..9aad70d 100755 --- a/test_code/network.py +++ b/test_code/network.py @@ -1,12 +1,12 @@ import tensorflow as tf import numpy as np -import tensorflow.contrib.slim as slim +import tf_slim as slim def resblock(inputs, out_channel=32, name='resblock'): - with tf.variable_scope(name): + with tf.compat.v1.variable_scope(name): x = slim.convolution2d(inputs, out_channel, [3, 3], activation_fn=None, scope='conv1') @@ -20,7 +20,7 @@ def resblock(inputs, out_channel=32, name='resblock'): def unet_generator(inputs, channel=32, num_blocks=4, name='generator', reuse=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): x0 = slim.convolution2d(inputs, channel, [7, 7], activation_fn=None) x0 = tf.nn.leaky_relu(x0) @@ -41,15 +41,15 @@ def unet_generator(inputs, channel=32, num_blocks=4, name='generator', reuse=Fal x2 = slim.convolution2d(x2, channel*2, [3, 3], activation_fn=None) x2 = tf.nn.leaky_relu(x2) - h1, w1 = tf.shape(x2)[1], tf.shape(x2)[2] - x3 = tf.image.resize_bilinear(x2, (h1*2, w1*2)) + h1, w1 = tf.shape(input=x2)[1], tf.shape(input=x2)[2] + x3 = tf.image.resize(x2, (h1*2, w1*2), method=tf.image.ResizeMethod.BILINEAR) x3 = slim.convolution2d(x3+x1, channel*2, [3, 3], activation_fn=None) x3 = tf.nn.leaky_relu(x3) x3 = slim.convolution2d(x3, channel, [3, 3], activation_fn=None) x3 = tf.nn.leaky_relu(x3) - h2, w2 = tf.shape(x3)[1], tf.shape(x3)[2] - x4 = tf.image.resize_bilinear(x3, (h2*2, w2*2)) + h2, w2 = tf.shape(input=x3)[1], tf.shape(input=x3)[2] + x4 = tf.image.resize(x3, (h2*2, w2*2), method=tf.image.ResizeMethod.BILINEAR) x4 = slim.convolution2d(x4+x0, channel, [3, 3], activation_fn=None) x4 = tf.nn.leaky_relu(x4) x4 = slim.convolution2d(x4, 3, [7, 7], activation_fn=None) diff --git a/test_code/saved_models/checkpoint b/test_code/saved_models/checkpoint old mode 100644 new mode 100755 diff --git a/test_code/saved_models/model-33999.data-00000-of-00001 b/test_code/saved_models/model-33999.data-00000-of-00001 old mode 100644 new mode 100755 diff --git a/test_code/saved_models/model-33999.index b/test_code/saved_models/model-33999.index old mode 100644 new mode 100755 diff --git a/test_code/test_images/actress2.jpg b/test_code/test_images/actress2.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/china6.jpg b/test_code/test_images/china6.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/food16.jpg b/test_code/test_images/food16.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/food6.jpg b/test_code/test_images/food6.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/liuyifei4.jpg b/test_code/test_images/liuyifei4.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/london1.jpg b/test_code/test_images/london1.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/mountain4.jpg b/test_code/test_images/mountain4.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/mountain5.jpg b/test_code/test_images/mountain5.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/national_park1.jpg b/test_code/test_images/national_park1.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/party5.jpg b/test_code/test_images/party5.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/party7.jpg b/test_code/test_images/party7.jpg old mode 100644 new mode 100755 diff --git a/train_code/guided_filter.py b/train_code/guided_filter.py old mode 100644 new mode 100755 index 340f450..5f86d86 --- a/train_code/guided_filter.py +++ b/train_code/guided_filter.py @@ -13,14 +13,14 @@ def tf_box_filter(x, r): weight = 1/((2*r+1)**2) box_kernel = weight*np.ones((2*r+1, 2*r+1, ch, 1)) box_kernel = np.array(box_kernel).astype(np.float32) - output = tf.nn.depthwise_conv2d(x, box_kernel, [1, 1, 1, 1], 'SAME') + output = tf.nn.depthwise_conv2d(input=x, filter=box_kernel, strides=[1, 1, 1, 1], padding='SAME') return output def guided_filter(x, y, r, eps=1e-2): - x_shape = tf.shape(x) + x_shape = tf.shape(input=x) #y_shape = tf.shape(y) N = tf_box_filter(tf.ones((1, x_shape[1], x_shape[2], 1), dtype=x.dtype), r) diff --git a/train_code/layers.py b/train_code/layers.py old mode 100644 new mode 100755 index e361ad5..cf1a854 --- a/train_code/layers.py +++ b/train_code/layers.py @@ -1,93 +1,92 @@ -''' -CVPR 2020 submission, Paper ID 6791 -Source code for 'Learning to Cartoonize Using White-Box Cartoon Representations' -''' - - -import tensorflow as tf -import numpy as np -import tensorflow.contrib.slim as slim - - - -def adaptive_instance_norm(content, style, epsilon=1e-5): - - c_mean, c_var = tf.nn.moments(content, axes=[1, 2], keep_dims=True) - s_mean, s_var = tf.nn.moments(style, axes=[1, 2], keep_dims=True) - c_std, s_std = tf.sqrt(c_var + epsilon), tf.sqrt(s_var + epsilon) - - return s_std * (content - c_mean) / c_std + s_mean - - - -def spectral_norm(w, iteration=1): - w_shape = w.shape.as_list() - w = tf.reshape(w, [-1, w_shape[-1]]) - - u = tf.get_variable("u", [1, w_shape[-1]], - initializer=tf.random_normal_initializer(), trainable=False) - - u_hat = u - v_hat = None - for i in range(iteration): - """ - power iteration - Usually iteration = 1 will be enough - """ - v_ = tf.matmul(u_hat, tf.transpose(w)) - v_hat = tf.nn.l2_normalize(v_) - - u_ = tf.matmul(v_hat, w) - u_hat = tf.nn.l2_normalize(u_) - - u_hat = tf.stop_gradient(u_hat) - v_hat = tf.stop_gradient(v_hat) - - sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(u_hat)) - - with tf.control_dependencies([u.assign(u_hat)]): - w_norm = w / sigma - w_norm = tf.reshape(w_norm, w_shape) - - return w_norm - - -def conv_spectral_norm(x, channel, k_size, stride=1, name='conv_snorm'): - with tf.variable_scope(name): - w = tf.get_variable("kernel", shape=[k_size[0], k_size[1], x.get_shape()[-1], channel]) - b = tf.get_variable("bias", [channel], initializer=tf.constant_initializer(0.0)) - - x = tf.nn.conv2d(input=x, filter=spectral_norm(w), strides=[1, stride, stride, 1], padding='SAME') + b - - return x - - - -def self_attention(inputs, name='attention', reuse=False): - with tf.variable_scope(name, reuse=reuse): - h, w = tf.shape(inputs)[1], tf.shape(inputs)[2] - bs, _, _, ch = inputs.get_shape().as_list() - f = slim.convolution2d(inputs, ch//8, [1, 1], activation_fn=None) - g = slim.convolution2d(inputs, ch//8, [1, 1], activation_fn=None) - s = slim.convolution2d(inputs, 1, [1, 1], activation_fn=None) - f_flatten = tf.reshape(f, shape=[f.shape[0], -1, f.shape[-1]]) - g_flatten = tf.reshape(g, shape=[g.shape[0], -1, g.shape[-1]]) - beta = tf.matmul(f_flatten, g_flatten, transpose_b=True) - beta = tf.nn.softmax(beta) - - s_flatten = tf.reshape(s, shape=[s.shape[0], -1, s.shape[-1]]) - att_map = tf.matmul(beta, s_flatten) - att_map = tf.reshape(att_map, shape=[bs, h, w, 1]) - gamma = tf.get_variable("gamma", [1], initializer=tf.constant_initializer(0.0)) - output = att_map * gamma + inputs - - return att_map, output - - - -if __name__ == '__main__': - pass - - - - \ No newline at end of file +''' +CVPR 2020 submission, Paper ID 6791 +Source code for 'Learning to Cartoonize Using White-Box Cartoon Representations' +''' + + +import tensorflow as tf +import numpy as np +import tf_slim as slim + + + +def adaptive_instance_norm(content, style, epsilon=1e-5): + + c_mean, c_var = tf.nn.moments(x=content, axes=[1, 2], keepdims=True) + s_mean, s_var = tf.nn.moments(x=style, axes=[1, 2], keepdims=True) + c_std, s_std = tf.sqrt(c_var + epsilon), tf.sqrt(s_var + epsilon) + + return s_std * (content - c_mean) / c_std + s_mean + + + +def spectral_norm(w, iteration=1): + w_shape = w.shape.as_list() + w = tf.reshape(w, [-1, w_shape[-1]]) + + u = tf.compat.v1.get_variable("u", [1, w_shape[-1]], + initializer=tf.compat.v1.random_normal_initializer(), trainable=False) + + u_hat = u + v_hat = None + for i in range(iteration): + """ + power iteration + Usually iteration = 1 will be enough + """ + v_ = tf.matmul(u_hat, tf.transpose(a=w)) + v_hat = tf.nn.l2_normalize(v_) + + u_ = tf.matmul(v_hat, w) + u_hat = tf.nn.l2_normalize(u_) + + u_hat = tf.stop_gradient(u_hat) + v_hat = tf.stop_gradient(v_hat) + + sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(a=u_hat)) + + with tf.control_dependencies([u.assign(u_hat)]): + w_norm = w / sigma + w_norm = tf.reshape(w_norm, w_shape) + + return w_norm + + +def conv_spectral_norm(x, channel, k_size, stride=1, name='conv_snorm'): + with tf.compat.v1.variable_scope(name): + w = tf.compat.v1.get_variable("kernel", shape=[k_size[0], k_size[1], x.get_shape()[-1], channel]) + b = tf.compat.v1.get_variable("bias", [channel], initializer=tf.compat.v1.constant_initializer(0.0)) + + x = tf.nn.conv2d(input=x, filters=spectral_norm(w), strides=[1, stride, stride, 1], padding='SAME') + b + + return x + + + +def self_attention(inputs, name='attention', reuse=False): + with tf.compat.v1.variable_scope(name, reuse=reuse): + h, w = tf.shape(input=inputs)[1], tf.shape(input=inputs)[2] + bs, _, _, ch = inputs.get_shape().as_list() + f = slim.convolution2d(inputs, ch//8, [1, 1], activation_fn=None) + g = slim.convolution2d(inputs, ch//8, [1, 1], activation_fn=None) + s = slim.convolution2d(inputs, 1, [1, 1], activation_fn=None) + f_flatten = tf.reshape(f, shape=[f.shape[0], -1, f.shape[-1]]) + g_flatten = tf.reshape(g, shape=[g.shape[0], -1, g.shape[-1]]) + beta = tf.matmul(f_flatten, g_flatten, transpose_b=True) + beta = tf.nn.softmax(beta) + + s_flatten = tf.reshape(s, shape=[s.shape[0], -1, s.shape[-1]]) + att_map = tf.matmul(beta, s_flatten) + att_map = tf.reshape(att_map, shape=[bs, h, w, 1]) + gamma = tf.compat.v1.get_variable("gamma", [1], initializer=tf.compat.v1.constant_initializer(0.0)) + output = att_map * gamma + inputs + + return att_map, output + + + +if __name__ == '__main__': + pass + + + diff --git a/train_code/loss.py b/train_code/loss.py old mode 100644 new mode 100755 index 889a243..6e95c43 --- a/train_code/loss.py +++ b/train_code/loss.py @@ -1,191 +1,190 @@ -''' -CVPR 2020 submission, Paper ID 6791 -Source code for 'Learning to Cartoonize Using White-Box Cartoon Representations' -''' - - -import numpy as np -import scipy.stats as st -import tensorflow as tf - - - -VGG_MEAN = [103.939, 116.779, 123.68] - - -class Vgg19: - - def __init__(self, vgg19_npy_path=None): - - self.data_dict = np.load(vgg19_npy_path, encoding='latin1', allow_pickle=True).item() - print('Finished loading vgg19.npy') - - - def build_conv4_4(self, rgb, include_fc=False): - - rgb_scaled = (rgb+1) * 127.5 - - blue, green, red = tf.split(axis=3, num_or_size_splits=3, value=rgb_scaled) - bgr = tf.concat(axis=3, values=[blue - VGG_MEAN[0], - green - VGG_MEAN[1], red - VGG_MEAN[2]]) - - self.conv1_1 = self.conv_layer(bgr, "conv1_1") - self.relu1_1 = tf.nn.relu(self.conv1_1) - self.conv1_2 = self.conv_layer(self.relu1_1, "conv1_2") - self.relu1_2 = tf.nn.relu(self.conv1_2) - self.pool1 = self.max_pool(self.relu1_2, 'pool1') - - self.conv2_1 = self.conv_layer(self.pool1, "conv2_1") - self.relu2_1 = tf.nn.relu(self.conv2_1) - self.conv2_2 = self.conv_layer(self.relu2_1, "conv2_2") - self.relu2_2 = tf.nn.relu(self.conv2_2) - self.pool2 = self.max_pool(self.relu2_2, 'pool2') - - self.conv3_1 = self.conv_layer(self.pool2, "conv3_1") - self.relu3_1 = tf.nn.relu(self.conv3_1) - self.conv3_2 = self.conv_layer(self.relu3_1, "conv3_2") - self.relu3_2 = tf.nn.relu(self.conv3_2) - self.conv3_3 = self.conv_layer(self.relu3_2, "conv3_3") - self.relu3_3 = tf.nn.relu(self.conv3_3) - self.conv3_4 = self.conv_layer(self.relu3_3, "conv3_4") - self.relu3_4 = tf.nn.relu(self.conv3_4) - self.pool3 = self.max_pool(self.relu3_4, 'pool3') - - self.conv4_1 = self.conv_layer(self.pool3, "conv4_1") - self.relu4_1 = tf.nn.relu(self.conv4_1) - self.conv4_2 = self.conv_layer(self.relu4_1, "conv4_2") - self.relu4_2 = tf.nn.relu(self.conv4_2) - self.conv4_3 = self.conv_layer(self.relu4_2, "conv4_3") - self.relu4_3 = tf.nn.relu(self.conv4_3) - self.conv4_4 = self.conv_layer(self.relu4_3, "conv4_4") - self.relu4_4 = tf.nn.relu(self.conv4_4) - self.pool4 = self.max_pool(self.relu4_4, 'pool4') - - return self.conv4_4 - - def max_pool(self, bottom, name): - return tf.nn.max_pool(bottom, ksize=[1, 2, 2, 1], - strides=[1, 2, 2, 1], padding='SAME', name=name) - - def conv_layer(self, bottom, name): - with tf.variable_scope(name): - filt = self.get_conv_filter(name) - - conv = tf.nn.conv2d(bottom, filt, [1, 1, 1, 1], padding='SAME') - - conv_biases = self.get_bias(name) - bias = tf.nn.bias_add(conv, conv_biases) - - #relu = tf.nn.relu(bias) - return bias - - - - def fc_layer(self, bottom, name): - with tf.variable_scope(name): - shape = bottom.get_shape().as_list() - dim = 1 - for d in shape[1:]: - dim *= d - x = tf.reshape(bottom, [-1, dim]) - - weights = self.get_fc_weight(name) - biases = self.get_bias(name) - - # Fully connected layer. Note that the '+' operation automatically - # broadcasts the biases. - fc = tf.nn.bias_add(tf.matmul(x, weights), biases) - - return fc - - def get_conv_filter(self, name): - return tf.constant(self.data_dict[name][0], name="filter") - - def get_bias(self, name): - return tf.constant(self.data_dict[name][1], name="biases") - - def get_fc_weight(self, name): - return tf.constant(self.data_dict[name][0], name="weights") - - - -def vggloss_4_4(image_a, image_b): - vgg_model = Vgg19('vgg19_no_fc.npy') - vgg_a = vgg_model.build_conv4_4(image_a) - vgg_b = vgg_model.build_conv4_4(image_b) - VGG_loss = tf.losses.absolute_difference(vgg_a, vgg_b) - #VGG_loss = tf.nn.l2_loss(vgg_a - vgg_b) - h, w, c= vgg_a.get_shape().as_list()[1:] - VGG_loss = tf.reduce_mean(VGG_loss)/(h*w*c) - return VGG_loss - - - -def wgan_loss(discriminator, real, fake, patch=True, - channel=32, name='discriminator', lambda_=2): - real_logits = discriminator(real, patch=patch, channel=channel, name=name, reuse=False) - fake_logits = discriminator(fake, patch=patch, channel=channel, name=name, reuse=True) - - d_loss_real = - tf.reduce_mean(real_logits) - d_loss_fake = tf.reduce_mean(fake_logits) - - d_loss = d_loss_real + d_loss_fake - g_loss = - d_loss_fake - - """ Gradient Penalty """ - # This is borrowed from https://github.com/kodalinaveen3/DRAGAN/blob/master/DRAGAN.ipynb - alpha = tf.random_uniform([tf.shape(real)[0], 1, 1, 1], minval=0.,maxval=1.) - differences = fake - real # This is different from MAGAN - interpolates = real + (alpha * differences) - inter_logit = discriminator(interpolates, channel=channel, name=name, reuse=True) - gradients = tf.gradients(inter_logit, [interpolates])[0] - slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1])) - gradient_penalty = tf.reduce_mean((slopes - 1.) ** 2) - d_loss += lambda_ * gradient_penalty - - return d_loss, g_loss - - -def gan_loss(discriminator, real, fake, scale=1,channel=32, patch=False, name='discriminator'): - - real_logit = discriminator(real, scale, channel, name=name, patch=patch, reuse=False) - fake_logit = discriminator(fake, scale, channel, name=name, patch=patch, reuse=True) - - real_logit = tf.nn.sigmoid(real_logit) - fake_logit = tf.nn.sigmoid(fake_logit) - - g_loss_blur = -tf.reduce_mean(tf.log(fake_logit)) - d_loss_blur = -tf.reduce_mean(tf.log(real_logit) + tf.log(1. - fake_logit)) - - return d_loss_blur, g_loss_blur - - - -def lsgan_loss(discriminator, real, fake, scale=1, - channel=32, patch=False, name='discriminator'): - - real_logit = discriminator(real, scale, channel, name=name, patch=patch, reuse=False) - fake_logit = discriminator(fake, scale, channel, name=name, patch=patch, reuse=True) - - g_loss = tf.reduce_mean((fake_logit - 1)**2) - d_loss = 0.5*(tf.reduce_mean((real_logit - 1)**2) + tf.reduce_mean(fake_logit**2)) - - return d_loss, g_loss - - - -def total_variation_loss(image, k_size=1): - h, w = image.get_shape().as_list()[1:3] - tv_h = tf.reduce_mean((image[:, k_size:, :, :] - image[:, :h - k_size, :, :])**2) - tv_w = tf.reduce_mean((image[:, :, k_size:, :] - image[:, :, :w - k_size, :])**2) - tv_loss = (tv_h + tv_w)/(3*h*w) - return tv_loss - - - - -if __name__ == '__main__': - pass - - - \ No newline at end of file +''' +CVPR 2020 submission, Paper ID 6791 +Source code for 'Learning to Cartoonize Using White-Box Cartoon Representations' +''' + + +import numpy as np +import scipy.stats as st +import tensorflow as tf + + + +VGG_MEAN = [103.939, 116.779, 123.68] + + +class Vgg19: + + def __init__(self, vgg19_npy_path=None): + + self.data_dict = np.load(vgg19_npy_path, encoding='latin1', allow_pickle=True).item() + print('Finished loading vgg19.npy') + + + def build_conv4_4(self, rgb, include_fc=False): + + rgb_scaled = (rgb+1) * 127.5 + + blue, green, red = tf.split(axis=3, num_or_size_splits=3, value=rgb_scaled) + bgr = tf.concat(axis=3, values=[blue - VGG_MEAN[0], + green - VGG_MEAN[1], red - VGG_MEAN[2]]) + + self.conv1_1 = self.conv_layer(bgr, "conv1_1") + self.relu1_1 = tf.nn.relu(self.conv1_1) + self.conv1_2 = self.conv_layer(self.relu1_1, "conv1_2") + self.relu1_2 = tf.nn.relu(self.conv1_2) + self.pool1 = self.max_pool(self.relu1_2, 'pool1') + + self.conv2_1 = self.conv_layer(self.pool1, "conv2_1") + self.relu2_1 = tf.nn.relu(self.conv2_1) + self.conv2_2 = self.conv_layer(self.relu2_1, "conv2_2") + self.relu2_2 = tf.nn.relu(self.conv2_2) + self.pool2 = self.max_pool(self.relu2_2, 'pool2') + + self.conv3_1 = self.conv_layer(self.pool2, "conv3_1") + self.relu3_1 = tf.nn.relu(self.conv3_1) + self.conv3_2 = self.conv_layer(self.relu3_1, "conv3_2") + self.relu3_2 = tf.nn.relu(self.conv3_2) + self.conv3_3 = self.conv_layer(self.relu3_2, "conv3_3") + self.relu3_3 = tf.nn.relu(self.conv3_3) + self.conv3_4 = self.conv_layer(self.relu3_3, "conv3_4") + self.relu3_4 = tf.nn.relu(self.conv3_4) + self.pool3 = self.max_pool(self.relu3_4, 'pool3') + + self.conv4_1 = self.conv_layer(self.pool3, "conv4_1") + self.relu4_1 = tf.nn.relu(self.conv4_1) + self.conv4_2 = self.conv_layer(self.relu4_1, "conv4_2") + self.relu4_2 = tf.nn.relu(self.conv4_2) + self.conv4_3 = self.conv_layer(self.relu4_2, "conv4_3") + self.relu4_3 = tf.nn.relu(self.conv4_3) + self.conv4_4 = self.conv_layer(self.relu4_3, "conv4_4") + self.relu4_4 = tf.nn.relu(self.conv4_4) + self.pool4 = self.max_pool(self.relu4_4, 'pool4') + + return self.conv4_4 + + def max_pool(self, bottom, name): + return tf.nn.max_pool2d(input=bottom, ksize=[1, 2, 2, 1], + strides=[1, 2, 2, 1], padding='SAME', name=name) + + def conv_layer(self, bottom, name): + with tf.compat.v1.variable_scope(name): + filt = self.get_conv_filter(name) + + conv = tf.nn.conv2d(input=bottom, filters=filt, strides=[1, 1, 1, 1], padding='SAME') + + conv_biases = self.get_bias(name) + bias = tf.nn.bias_add(conv, conv_biases) + + #relu = tf.nn.relu(bias) + return bias + + + + def fc_layer(self, bottom, name): + with tf.compat.v1.variable_scope(name): + shape = bottom.get_shape().as_list() + dim = 1 + for d in shape[1:]: + dim *= d + x = tf.reshape(bottom, [-1, dim]) + + weights = self.get_fc_weight(name) + biases = self.get_bias(name) + + # Fully connected layer. Note that the '+' operation automatically + # broadcasts the biases. + fc = tf.nn.bias_add(tf.matmul(x, weights), biases) + + return fc + + def get_conv_filter(self, name): + return tf.constant(self.data_dict[name][0], name="filter") + + def get_bias(self, name): + return tf.constant(self.data_dict[name][1], name="biases") + + def get_fc_weight(self, name): + return tf.constant(self.data_dict[name][0], name="weights") + + + +def vggloss_4_4(image_a, image_b): + vgg_model = Vgg19('vgg19_no_fc.npy') + vgg_a = vgg_model.build_conv4_4(image_a) + vgg_b = vgg_model.build_conv4_4(image_b) + VGG_loss = tf.compat.v1.losses.absolute_difference(vgg_a, vgg_b) + #VGG_loss = tf.nn.l2_loss(vgg_a - vgg_b) + h, w, c= vgg_a.get_shape().as_list()[1:] + VGG_loss = tf.reduce_mean(input_tensor=VGG_loss)/(h*w*c) + return VGG_loss + + + +def wgan_loss(discriminator, real, fake, patch=True, + channel=32, name='discriminator', lambda_=2): + real_logits = discriminator(real, patch=patch, channel=channel, name=name, reuse=False) + fake_logits = discriminator(fake, patch=patch, channel=channel, name=name, reuse=True) + + d_loss_real = - tf.reduce_mean(input_tensor=real_logits) + d_loss_fake = tf.reduce_mean(input_tensor=fake_logits) + + d_loss = d_loss_real + d_loss_fake + g_loss = - d_loss_fake + + """ Gradient Penalty """ + # This is borrowed from https://github.com/kodalinaveen3/DRAGAN/blob/master/DRAGAN.ipynb + alpha = tf.random.uniform([tf.shape(input=real)[0], 1, 1, 1], minval=0.,maxval=1.) + differences = fake - real # This is different from MAGAN + interpolates = real + (alpha * differences) + inter_logit = discriminator(interpolates, channel=channel, name=name, reuse=True) + gradients = tf.gradients(ys=inter_logit, xs=[interpolates])[0] + slopes = tf.sqrt(tf.reduce_sum(input_tensor=tf.square(gradients), axis=[1])) + gradient_penalty = tf.reduce_mean(input_tensor=(slopes - 1.) ** 2) + d_loss += lambda_ * gradient_penalty + + return d_loss, g_loss + + +def gan_loss(discriminator, real, fake, scale=1,channel=32, patch=False, name='discriminator'): + + real_logit = discriminator(real, scale, channel, name=name, patch=patch, reuse=False) + fake_logit = discriminator(fake, scale, channel, name=name, patch=patch, reuse=True) + + real_logit = tf.nn.sigmoid(real_logit) + fake_logit = tf.nn.sigmoid(fake_logit) + + g_loss_blur = -tf.reduce_mean(input_tensor=tf.math.log(fake_logit)) + d_loss_blur = -tf.reduce_mean(input_tensor=tf.math.log(real_logit) + tf.math.log(1. - fake_logit)) + + return d_loss_blur, g_loss_blur + + + +def lsgan_loss(discriminator, real, fake, scale=1, + channel=32, patch=False, name='discriminator'): + + real_logit = discriminator(real, scale, channel, name=name, patch=patch, reuse=False) + fake_logit = discriminator(fake, scale, channel, name=name, patch=patch, reuse=True) + + g_loss = tf.reduce_mean(input_tensor=(fake_logit - 1)**2) + d_loss = 0.5*(tf.reduce_mean(input_tensor=(real_logit - 1)**2) + tf.reduce_mean(input_tensor=fake_logit**2)) + + return d_loss, g_loss + + + +def total_variation_loss(image, k_size=1): + h, w = image.get_shape().as_list()[1:3] + tv_h = tf.reduce_mean(input_tensor=(image[:, k_size:, :, :] - image[:, :h - k_size, :, :])**2) + tv_w = tf.reduce_mean(input_tensor=(image[:, :, k_size:, :] - image[:, :, :w - k_size, :])**2) + tv_loss = (tv_h + tv_w)/(3*h*w) + return tv_loss + + + + +if __name__ == '__main__': + pass + + diff --git a/train_code/network.py b/train_code/network.py old mode 100644 new mode 100755 index 8527aa8..dc10944 --- a/train_code/network.py +++ b/train_code/network.py @@ -7,7 +7,7 @@ import layers import tensorflow as tf import numpy as np -import tensorflow.contrib.slim as slim +import tf_slim as slim from tqdm import tqdm @@ -15,7 +15,7 @@ def resblock(inputs, out_channel=32, name='resblock'): - with tf.variable_scope(name): + with tf.compat.v1.variable_scope(name): x = slim.convolution2d(inputs, out_channel, [3, 3], activation_fn=None, scope='conv1') @@ -28,7 +28,7 @@ def resblock(inputs, out_channel=32, name='resblock'): def generator(inputs, channel=32, num_blocks=4, name='generator', reuse=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): x = slim.convolution2d(inputs, channel, [7, 7], activation_fn=None) x = tf.nn.leaky_relu(x) @@ -60,7 +60,7 @@ def generator(inputs, channel=32, num_blocks=4, name='generator', reuse=False): def unet_generator(inputs, channel=32, num_blocks=4, name='generator', reuse=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): x0 = slim.convolution2d(inputs, channel, [7, 7], activation_fn=None) x0 = tf.nn.leaky_relu(x0) @@ -81,15 +81,15 @@ def unet_generator(inputs, channel=32, num_blocks=4, name='generator', reuse=Fal x2 = slim.convolution2d(x2, channel*2, [3, 3], activation_fn=None) x2 = tf.nn.leaky_relu(x2) - h1, w1 = tf.shape(x2)[1], tf.shape(x2)[2] - x3 = tf.image.resize_bilinear(x2, (h1*2, w1*2)) + h1, w1 = tf.shape(input=x2)[1], tf.shape(input=x2)[2] + x3 = tf.image.resize(x2, (h1*2, w1*2), method=tf.image.ResizeMethod.BILINEAR) x3 = slim.convolution2d(x3+x1, channel*2, [3, 3], activation_fn=None) x3 = tf.nn.leaky_relu(x3) x3 = slim.convolution2d(x3, channel, [3, 3], activation_fn=None) x3 = tf.nn.leaky_relu(x3) - h2, w2 = tf.shape(x3)[1], tf.shape(x3)[2] - x4 = tf.image.resize_bilinear(x3, (h2*2, w2*2)) + h2, w2 = tf.shape(input=x3)[1], tf.shape(input=x3)[2] + x4 = tf.image.resize(x3, (h2*2, w2*2), method=tf.image.ResizeMethod.BILINEAR) x4 = slim.convolution2d(x4+x0, channel, [3, 3], activation_fn=None) x4 = tf.nn.leaky_relu(x4) x4 = slim.convolution2d(x4, 3, [7, 7], activation_fn=None) @@ -101,7 +101,7 @@ def unet_generator(inputs, channel=32, num_blocks=4, name='generator', reuse=Fal def disc_bn(x, scale=1, channel=32, is_training=True, name='discriminator', patch=True, reuse=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): for idx in range(3): x = slim.convolution2d(x, channel*2**idx, [3, 3], stride=2, activation_fn=None) @@ -115,7 +115,7 @@ def disc_bn(x, scale=1, channel=32, is_training=True, if patch == True: x = slim.convolution2d(x, 1, [1, 1], activation_fn=None) else: - x = tf.reduce_mean(x, axis=[1, 2]) + x = tf.reduce_mean(input_tensor=x, axis=[1, 2]) x = slim.fully_connected(x, 1, activation_fn=None) return x @@ -124,7 +124,7 @@ def disc_bn(x, scale=1, channel=32, is_training=True, def disc_sn(x, scale=1, channel=32, patch=True, name='discriminator', reuse=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): for idx in range(3): x = layers.conv_spectral_norm(x, channel*2**idx, [3, 3], @@ -140,14 +140,14 @@ def disc_sn(x, scale=1, channel=32, patch=True, name='discriminator', reuse=Fals x = layers.conv_spectral_norm(x, 1, [1, 1], name='conv_out'.format(idx)) else: - x = tf.reduce_mean(x, axis=[1, 2]) + x = tf.reduce_mean(input_tensor=x, axis=[1, 2]) x = slim.fully_connected(x, 1, activation_fn=None) return x def disc_ln(x, channel=32, is_training=True, name='discriminator', patch=True, reuse=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): for idx in range(3): x = slim.convolution2d(x, channel*2**idx, [3, 3], stride=2, activation_fn=None) @@ -161,7 +161,7 @@ def disc_ln(x, channel=32, is_training=True, name='discriminator', patch=True, r if patch == True: x = slim.convolution2d(x, 1, [1, 1], activation_fn=None) else: - x = tf.reduce_mean(x, axis=[1, 2]) + x = tf.reduce_mean(input_tensor=x, axis=[1, 2]) x = slim.fully_connected(x, 1, activation_fn=None) return x @@ -172,4 +172,3 @@ def disc_ln(x, channel=32, is_training=True, name='discriminator', patch=True, r if __name__ == '__main__': pass - \ No newline at end of file diff --git a/train_code/pretrain.py b/train_code/pretrain.py old mode 100644 new mode 100755 index d18684b..29c9f58 --- a/train_code/pretrain.py +++ b/train_code/pretrain.py @@ -7,8 +7,7 @@ import tensorflow as tf -import tensorflow.contrib.slim as slim - +import tf_slim as slim import utils import os import numpy as np @@ -19,7 +18,6 @@ os.environ["CUDA_VISIBLE_DEVICES"]="0" - def arg_parser(): parser = argparse.ArgumentParser() parser.add_argument("--patch_size", default = 256, type = int) @@ -38,20 +36,20 @@ def arg_parser(): def train(args): - input_photo = tf.placeholder(tf.float32, [args.batch_size, + input_photo = tf.compat.v1.placeholder(tf.float32, [args.batch_size, args.patch_size, args.patch_size, 3]) output = network.unet_generator(input_photo) - recon_loss = tf.reduce_mean(tf.losses.absolute_difference(input_photo, output)) + recon_loss = tf.reduce_mean(input_tensor=tf.compat.v1.losses.absolute_difference(input_photo, output)) - all_vars = tf.trainable_variables() + all_vars = tf.compat.v1.trainable_variables() gene_vars = [var for var in all_vars if 'gene' in var.name] - update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) + update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): - optim = tf.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ + optim = tf.compat.v1.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ .minimize(recon_loss, var_list=gene_vars) @@ -60,17 +58,17 @@ def train(args): config.gpu_options.allow_growth = True sess = tf.Session(config=config) ''' - gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=args.gpu_fraction) - sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) - saver = tf.train.Saver(var_list=gene_vars, max_to_keep=20) + gpu_options = tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=args.gpu_fraction) + sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options)) + saver = tf.compat.v1.train.Saver(var_list=gene_vars, max_to_keep=20) with tf.device('/device:GPU:0'): - sess.run(tf.global_variables_initializer()) - face_photo_dir = 'dataset/photo_face' - face_photo_list = utils.load_image_list(face_photo_dir) - scenery_photo_dir = 'dataset/photo_scenery' - scenery_photo_list = utils.load_image_list(scenery_photo_dir) + sess.run(tf.compat.v1.global_variables_initializer()) + face_photo_dir = 'dataset/face_photo' + utils.load_image_list(face_photo_dir) + scenery_photo_dir = 'dataset/scenery_photo' + utils.load_image_list(scenery_photo_dir) for total_iter in tqdm(range(args.total_iter)): @@ -84,8 +82,10 @@ def train(args): if np.mod(total_iter+1, 50) == 0: + print('pretrain, iter: {}, recon_loss: {}'.format(total_iter, r_loss)) if np.mod(total_iter+1, 500 ) == 0: + saver.save(sess, args.save_dir+'save_models/model', write_meta_graph=False, global_step=total_iter) @@ -104,6 +104,7 @@ def train(args): str(total_iter)+'_scenery_result.jpg', 4) utils.write_batch_image(photo_scenery, args.save_dir+'/images', str(total_iter)+'_scenery_photo.jpg', 4) + @@ -113,4 +114,3 @@ def train(args): args = arg_parser() train(args) - \ No newline at end of file diff --git a/train_code/selective_search/__pycache__/__init__.cpython-37.pyc b/train_code/selective_search/__pycache__/__init__.cpython-37.pyc new file mode 100644 index 0000000..fc2cd4a Binary files /dev/null and b/train_code/selective_search/__pycache__/__init__.cpython-37.pyc differ diff --git a/train_code/selective_search/__pycache__/core.cpython-37.pyc b/train_code/selective_search/__pycache__/core.cpython-37.pyc new file mode 100644 index 0000000..ca9137f Binary files /dev/null and b/train_code/selective_search/__pycache__/core.cpython-37.pyc differ diff --git a/train_code/selective_search/__pycache__/measure.cpython-37.pyc b/train_code/selective_search/__pycache__/measure.cpython-37.pyc new file mode 100644 index 0000000..05e5410 Binary files /dev/null and b/train_code/selective_search/__pycache__/measure.cpython-37.pyc differ diff --git a/train_code/selective_search/__pycache__/structure.cpython-37.pyc b/train_code/selective_search/__pycache__/structure.cpython-37.pyc new file mode 100644 index 0000000..554e298 Binary files /dev/null and b/train_code/selective_search/__pycache__/structure.cpython-37.pyc differ diff --git a/train_code/selective_search/__pycache__/util.cpython-37.pyc b/train_code/selective_search/__pycache__/util.cpython-37.pyc new file mode 100644 index 0000000..f295ac1 Binary files /dev/null and b/train_code/selective_search/__pycache__/util.cpython-37.pyc differ diff --git a/train_code/selective_search/core.py b/train_code/selective_search/core.py index 502f9e9..129e8db 100755 --- a/train_code/selective_search/core.py +++ b/train_code/selective_search/core.py @@ -1,7 +1,7 @@ from joblib import Parallel, delayed from skimage.segmentation import felzenszwalb -from util import oversegmentation, switch_color_space, load_strategy -from structure import HierarchicalGrouping +from .util import oversegmentation, switch_color_space, load_strategy +from .structure import HierarchicalGrouping diff --git a/train_code/selective_search/structure.py b/train_code/selective_search/structure.py index 636bf0d..fe64462 100755 --- a/train_code/selective_search/structure.py +++ b/train_code/selective_search/structure.py @@ -2,7 +2,7 @@ from skimage.segmentation import find_boundaries from skimage.segmentation import felzenszwalb from scipy.ndimage import find_objects -import measure +from .measure import * class HierarchicalGrouping(object): @@ -14,7 +14,7 @@ def __init__(self, img, img_seg, sim_strategy): def build_regions(self): self.regions = {} - lbp_img = measure.generate_lbp_image(self.img) + lbp_img = generate_lbp_image(self.img) for label in self.labels: size = (self.img_seg == 1).sum() region_slice = find_objects(self.img_seg==label)[0] @@ -22,8 +22,8 @@ def build_regions(self): [region_slice[i].stop for i in (1,0)]) mask = self.img_seg == label - color_hist = measure.calculate_color_hist(mask, self.img) - texture_hist = measure.calculate_texture_hist(mask, lbp_img) + color_hist = calculate_color_hist(mask, self.img) + texture_hist = calculate_texture_hist(mask, lbp_img) self.regions[label] = { 'size': size, @@ -39,7 +39,7 @@ def build_region_pairs(self): neighbors = self._find_neighbors(i) for j in neighbors: if i < j: - self.s[(i,j)] = measure.calculate_sim(self.regions[i], + self.s[(i,j)] = calculate_sim(self.regions[i], self.regions[j], self.img.size, self.sim_strategy) @@ -118,7 +118,7 @@ def calculate_similarity_for_new_region(self): for j in neighbors: # i is larger than j, so use (j,i) instead - self.s[(j,i)] = measure.calculate_sim(self.regions[i], + self.s[(j,i)] = calculate_sim(self.regions[i], self.regions[j], self.img.size, self.sim_strategy) diff --git a/train_code/train.py b/train_code/train.py old mode 100644 new mode 100755 index fb3c5bb..1c3571f --- a/train_code/train.py +++ b/train_code/train.py @@ -6,21 +6,20 @@ import tensorflow as tf -import tensorflow.contrib.slim as slim - +import tf_slim as slim import utils import os import numpy as np import argparse import network import loss - +import random from tqdm import tqdm from guided_filter import guided_filter - +random.seed(0) os.environ["CUDA_VISIBLE_DEVICES"]="0" - +tf.compat.v1.disable_eager_execution() def arg_parser(): parser = argparse.ArgumentParser() parser.add_argument("--patch_size", default = 256, type = int) @@ -40,11 +39,11 @@ def arg_parser(): def train(args): - input_photo = tf.placeholder(tf.float32, [args.batch_size, + input_photo = tf.compat.v1.placeholder(tf.float32, [args.batch_size, args.patch_size, args.patch_size, 3]) - input_superpixel = tf.placeholder(tf.float32, [args.batch_size, + input_superpixel = tf.compat.v1.placeholder(tf.float32, [args.batch_size, args.patch_size, args.patch_size, 3]) - input_cartoon = tf.placeholder(tf.float32, [args.batch_size, + input_cartoon = tf.compat.v1.placeholder(tf.float32, [args.batch_size, args.patch_size, args.patch_size, 3]) output = network.unet_generator(input_photo) @@ -62,14 +61,14 @@ def train(args): scale=1, patch=True, name='disc_blur') - vgg_model = loss.Vgg19('vgg19_no_fc.npy') + vgg_model = loss.Vgg19('train_code/vgg19_no_fc.npy') vgg_photo = vgg_model.build_conv4_4(input_photo) vgg_output = vgg_model.build_conv4_4(output) vgg_superpixel = vgg_model.build_conv4_4(input_superpixel) h, w, c = vgg_photo.get_shape().as_list()[1:] - photo_loss = tf.reduce_mean(tf.losses.absolute_difference(vgg_photo, vgg_output))/(h*w*c) - superpixel_loss = tf.reduce_mean(tf.losses.absolute_difference\ + photo_loss = tf.reduce_mean(input_tensor=tf.compat.v1.losses.absolute_difference(vgg_photo, vgg_output))/(h*w*c) + superpixel_loss = tf.reduce_mean(input_tensor=tf.compat.v1.losses.absolute_difference\ (vgg_superpixel, vgg_output))/(h*w*c) recon_loss = photo_loss + superpixel_loss tv_loss = loss.total_variation_loss(output) @@ -77,29 +76,29 @@ def train(args): g_loss_total = 1e4*tv_loss + 1e-1*g_loss_blur + g_loss_gray + 2e2*recon_loss d_loss_total = d_loss_blur + d_loss_gray - all_vars = tf.trainable_variables() + all_vars = tf.compat.v1.trainable_variables() gene_vars = [var for var in all_vars if 'gene' in var.name] disc_vars = [var for var in all_vars if 'disc' in var.name] - tf.summary.scalar('tv_loss', tv_loss) - tf.summary.scalar('photo_loss', photo_loss) - tf.summary.scalar('superpixel_loss', superpixel_loss) - tf.summary.scalar('recon_loss', recon_loss) - tf.summary.scalar('d_loss_gray', d_loss_gray) - tf.summary.scalar('g_loss_gray', g_loss_gray) - tf.summary.scalar('d_loss_blur', d_loss_blur) - tf.summary.scalar('g_loss_blur', g_loss_blur) - tf.summary.scalar('d_loss_total', d_loss_total) - tf.summary.scalar('g_loss_total', g_loss_total) + tf.compat.v1.summary.scalar('tv_loss', tv_loss) + tf.compat.v1.summary.scalar('photo_loss', photo_loss) + tf.compat.v1.summary.scalar('superpixel_loss', superpixel_loss) + tf.compat.v1.summary.scalar('recon_loss', recon_loss) + tf.compat.v1.summary.scalar('d_loss_gray', d_loss_gray) + tf.compat.v1.summary.scalar('g_loss_gray', g_loss_gray) + tf.compat.v1.summary.scalar('d_loss_blur', d_loss_blur) + tf.compat.v1.summary.scalar('g_loss_blur', g_loss_blur) + tf.compat.v1.summary.scalar('d_loss_total', d_loss_total) + tf.compat.v1.summary.scalar('g_loss_total', g_loss_total) - update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) + update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): - g_optim = tf.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ + g_optim = tf.compat.v1.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ .minimize(g_loss_total, var_list=gene_vars) - d_optim = tf.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ + d_optim = tf.compat.v1.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ .minimize(d_loss_total, var_list=disc_vars) ''' @@ -107,28 +106,32 @@ def train(args): config.gpu_options.allow_growth = True sess = tf.Session(config=config) ''' - gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=args.gpu_fraction) - sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) + gpu_options = tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=args.gpu_fraction) + sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options)) - train_writer = tf.summary.FileWriter(args.save_dir+'/train_log') - summary_op = tf.summary.merge_all() - saver = tf.train.Saver(var_list=gene_vars, max_to_keep=20) + train_writer = tf.compat.v1.summary.FileWriter(args.save_dir+'/train_log') + summary_op = tf.compat.v1.summary.merge_all() + saver = tf.compat.v1.train.Saver(var_list=gene_vars, max_to_keep=20) with tf.device('/device:GPU:0'): - sess.run(tf.global_variables_initializer()) - saver.restore(sess, tf.train.latest_checkpoint('pretrain/saved_models')) + sess.run(tf.compat.v1.global_variables_initializer()) + saver.restore(sess, tf.train.latest_checkpoint('pretrainsave_models')) - face_photo_dir = 'dataset/photo_face' + face_photo_dir = 'dataset/face_photo' face_photo_list = utils.load_image_list(face_photo_dir) - scenery_photo_dir = 'dataset/photo_scenery' + scenery_photo_dir = 'dataset/scenery_photo' scenery_photo_list = utils.load_image_list(scenery_photo_dir) - face_cartoon_dir = 'dataset/cartoon_face' - face_cartoon_list = utils.load_image_list(face_cartoon_dir) - scenery_cartoon_dir = 'dataset/cartoon_scenery' - scenery_cartoon_list = utils.load_image_list(scenery_cartoon_dir) + face_cartoon_dir_kyoto_face = 'dataset/face_cartoon/kyoto_face/' + face_cartoon_list = utils.load_image_list(face_cartoon_dir_kyoto_face) + face_cartoon_dir_pa_face = 'dataset/face_cartoon/pa_face/' + face_cartoon_list.extend(utils.load_image_list(face_cartoon_dir_pa_face)) + scenery_cartoon_dir = 'dataset/scenery_cartoon/' + scenery_cartoon_list = utils.load_image_list(scenery_cartoon_dir+"hayao/") + scenery_cartoon_list.extend(utils.load_image_list(scenery_cartoon_dir+"hosoda/")) + scenery_cartoon_list = utils.load_image_list(scenery_cartoon_dir+"shinkai/") for total_iter in tqdm(range(args.total_iter)): @@ -172,10 +175,12 @@ def train(args): print('Iter: {}, d_loss: {}, g_loss: {}, recon_loss: {}'.\ format(total_iter, d_loss, g_loss, r_loss)) - if np.mod(total_iter+1, 500 ) == 0: + if np.mod(total_iter+1, 250 ) == 0: saver.save(sess, args.save_dir+'/saved_models/model', write_meta_graph=False, global_step=total_iter) - + + if np.mod(total_iter+1, 500 ) == 0: + photo_face = utils.next_batch(face_photo_list, args.batch_size) cartoon_face = utils.next_batch(face_cartoon_list, args.batch_size) photo_scenery = utils.next_batch(scenery_photo_list, args.batch_size) @@ -193,15 +198,14 @@ def train(args): str(total_iter)+'_face_result.jpg', 4) utils.write_batch_image(photo_face, args.save_dir+'/images', str(total_iter)+'_face_photo.jpg', 4) - utils.write_batch_image(result_scenery, args.save_dir+'/images', str(total_iter)+'_scenery_result.jpg', 4) utils.write_batch_image(photo_scenery, args.save_dir+'/images', str(total_iter)+'_scenery_photo.jpg', 4) + if __name__ == '__main__': args = arg_parser() train(args) - \ No newline at end of file diff --git a/train_code/utils.py b/train_code/utils.py old mode 100644 new mode 100755 index 472559c..96f9cf3 --- a/train_code/utils.py +++ b/train_code/utils.py @@ -167,4 +167,3 @@ def write_batch_image(image, save_dir, name, n): if __name__ == '__main__': pass - \ No newline at end of file