| from collections import namedtuple
|
| from utils import *
|
| from ops import *
|
| import time
|
| from glob import glob
|
|
|
|
|
| def gaussian_noise_layer(input_layer, std):
|
| noise = tf.random.normal(
|
| shape=tf.shape(input_layer), mean=0.0, stddev=std, dtype=tf.float32
|
| )
|
| return input_layer + noise
|
|
|
|
|
| def generator_resnet(image, options, transfer=False, reuse=False, name="generator"):
|
| with tf.compat.v1.variable_scope(name):
|
| if reuse:
|
| tf.compat.v1.get_variable_scope().reuse_variables()
|
| else:
|
| assert tf.compat.v1.get_variable_scope().reuse is False
|
|
|
| def residule_block_dilated(x, dim, ks=3, s=1, name="res", down=False):
|
| if down:
|
| dim = dim * 2
|
| y = instance_norm(
|
| dilated_conv2d(x, dim, ks, s, padding="SAME", name=name + "_c1"),
|
| name + "_bn1",
|
| )
|
| y = tf.nn.relu(y)
|
| y = instance_norm(
|
| dilated_conv2d(y, dim, ks, s, padding="SAME", name=name + "_c2"),
|
| name + "_bn2",
|
| )
|
| out = y + x
|
| if down:
|
| out = tf.nn.relu(
|
| instance_norm(
|
| conv2d(out, dim // 2, 3, 1, name=name + "_down_c"),
|
| name + "_in_down",
|
| )
|
| )
|
| return out
|
|
|
| def residual_block(x_init, dim, ks=3, s=1, name="resblock", down=False):
|
| with tf.compat.v1.variable_scope(name):
|
| if down:
|
| dim = dim * 2
|
|
|
| with tf.compat.v1.variable_scope("res1"):
|
| x = instance_norm(
|
| conv2d(x_init, dim, ks, s, padding="SAME", name=name + "_c1"),
|
| name + "_in1",
|
| )
|
| x = tf.nn.relu(x)
|
|
|
| with tf.compat.v1.variable_scope("res2"):
|
|
|
| x = instance_norm(
|
| conv2d(x, dim, ks, s, padding="SAME", name=name + "_c2"),
|
| name + "_in2",
|
| )
|
|
|
| out = x + x_init
|
|
|
| if down:
|
| out = tf.nn.relu(
|
| instance_norm(
|
| conv2d(out, dim // 2, 3, 1, name=name + "_down_c"),
|
| name + "_in_down",
|
| )
|
| )
|
| return out
|
|
|
|
|
| c0 = tf.pad(image, [[0, 0], [3, 3], [3, 3], [0, 0]], "REFLECT")
|
| c1 = tf.nn.relu(
|
| instance_norm(
|
| conv2d(c0, options.gf_dim, 7, 1, padding="VALID", name="g_e1_c"),
|
| "g_e1_bn",
|
| )
|
| )
|
| c2 = tf.nn.relu(
|
| instance_norm(
|
| conv2d(c1, options.gf_dim * 2, 3, 2, name="g_e2_c"), "g_e2_bn"
|
| )
|
| )
|
| c3 = tf.nn.relu(
|
| instance_norm(
|
| conv2d(c2, options.gf_dim * 4, 3, 2, name="g_e3_c"), "g_e3_bn"
|
| )
|
| )
|
| r1 = residule_block_dilated(c3, options.gf_dim * 4, name="g_r1")
|
| r2 = residule_block_dilated(r1, options.gf_dim * 4, name="g_r2")
|
| r3 = residule_block_dilated(r2, options.gf_dim * 4, name="g_r3")
|
| r4 = residule_block_dilated(r3, options.gf_dim * 4, name="g_r4")
|
|
|
|
|
| if transfer:
|
| t1 = residual_block(r4, options.gf_dim * 4, name="g_t1")
|
| t2 = residual_block(t1, options.gf_dim * 4, name="g_t2")
|
| t3 = residual_block(t2, options.gf_dim * 4, name="g_t3")
|
| t4 = residual_block(t3, options.gf_dim * 4, name="g_t4")
|
|
|
|
|
| feature = t4
|
| else:
|
| feature = r4
|
| t4 = None
|
| down = False
|
|
|
|
|
| r6 = residule_block_dilated(feature, options.gf_dim * 4, name="g_r6")
|
| r7 = residule_block_dilated(r6, options.gf_dim * 4, name="g_r7")
|
| r8 = residule_block_dilated(r7, options.gf_dim * 4, name="g_r8")
|
| r9 = residule_block_dilated(r8, options.gf_dim * 4, name="g_r9")
|
| d1 = deconv2d(r9, options.gf_dim * 2, 3, 2, name="g_d1_dc")
|
| d1 = tf.nn.relu(instance_norm(d1, "g_d1_bn"))
|
| d2 = deconv2d(d1, options.gf_dim, 3, 2, name="g_d2_dc")
|
| d2 = tf.nn.relu(instance_norm(d2, "g_d2_bn"))
|
| d2 = tf.pad(d2, [[0, 0], [3, 3], [3, 3], [0, 0]], "REFLECT")
|
| pred = tf.nn.tanh(
|
| conv2d(d2, options.output_c_dim, 7, 1, padding="VALID", name="g_pred_c")
|
| )
|
|
|
|
|
| r5 = gaussian_noise_layer(r4, 0.02)
|
| r6_rec = residule_block_dilated(r5, options.gf_dim * 4, name="g_r6_rec")
|
| r6_rec = gaussian_noise_layer(r6_rec, 0.02)
|
| r7_rec = residule_block_dilated(r6_rec, options.gf_dim * 4, name="g_r7_rec")
|
| r8_rec = residule_block_dilated(r7_rec, options.gf_dim * 4, name="g_r8_rec")
|
| r9_rec = residule_block_dilated(r8_rec, options.gf_dim * 4, name="g_r9_rec")
|
| d1_rec = deconv2d(r9_rec, options.gf_dim * 2, 3, 2, name="g_d1_dc_rec")
|
| d1_rec = tf.nn.relu(instance_norm(d1_rec, "g_d1_bn_rec"))
|
| d2_rec = deconv2d(d1_rec, options.gf_dim, 3, 2, name="g_d2_dc_rec")
|
| d2_rec = tf.nn.relu(instance_norm(d2_rec, "g_d2_bn_rec"))
|
| d2_rec = tf.pad(d2_rec, [[0, 0], [3, 3], [3, 3], [0, 0]], "REFLECT")
|
| pred_rec = tf.nn.tanh(
|
| conv2d(
|
| d2_rec, options.output_c_dim, 7, 1, padding="VALID", name="g_pred_c_rec"
|
| )
|
| )
|
|
|
|
|
|
|
| if transfer:
|
|
|
| d_conf = deconv2d(d1, options.gf_dim, 3, 2, name="g_d_dc_conf")
|
| d_conf = tf.nn.relu(instance_norm(d_conf, "g_d_bn_conf"))
|
| d_conf = tf.pad(d_conf, [[0, 0], [3, 3], [3, 3], [0, 0]], "REFLECT")
|
| pred_conf = tf.nn.softplus(
|
| conv2d(d_conf, 1, 7, 1, padding="VALID", name="g_pred_c_conf")
|
| )
|
|
|
| else:
|
| pred_conf = None
|
|
|
| return pred, pred_rec, r4, t4, pred_conf
|
|
|
|
|
| def discriminator(image, options, n_scale=2, reuse=False, name="discriminator"):
|
| images = []
|
| for i in range(n_scale):
|
| images.append(
|
| tf.compat.v1.image.resize_bicubic(
|
| image, [get_shape(image)[1] // (2**i), get_shape(image)[2] // (2**i)]
|
| )
|
| )
|
| with tf.compat.v1.variable_scope(name):
|
| if reuse:
|
| tf.compat.v1.get_variable_scope().reuse_variables()
|
| else:
|
| assert tf.compat.v1.get_variable_scope().reuse is False
|
| images = dis_down(images, 4, 2, n_scale, options.df_dim, "d_h0_conv_scale_")
|
| images = dis_down(images, 4, 2, n_scale, options.df_dim * 2, "d_h1_conv_scale_")
|
| images = dis_down(images, 4, 2, n_scale, options.df_dim * 4, "d_h2_conv_scale_")
|
| images = dis_down(images, 4, 2, n_scale, options.df_dim * 8, "d_h3_conv_scale_")
|
| images = final_conv(images, n_scale, "d_pred_scale_")
|
| return images
|
|
|
