def build_generator(input_shape=(256, 256, 3), num_blocks=9):
"""Generator network architecture"""
x0 = layers.Input(input_shape)
x = ReflectionPadding2D(padding=(3, 3))(x0)
x = layers.Conv2D(filters=64, kernel_size=7, strides=1, kernel_initializer=RandomNormal(mean=0, stddev=0.02))(x)
x = InstanceNormalization()(x)
x = layers.ReLU()(x)
# downsample
x = layers.Conv2D(filters=128,
kernel_size=3,
strides=2,
padding='same',
kernel_initializer=RandomNormal(mean=0, stddev=0.02))(x)
x = InstanceNormalization()(x)
x = layers.ReLU()(x)
x = layers.Conv2D(filters=256,
kernel_size=3,
strides=2,
padding='same',
kernel_initializer=RandomNormal(mean=0, stddev=0.02))(x)
x = InstanceNormalization()(x)
x = layers.ReLU()(x)
# residual
for _ in range(num_blocks):
x = _resblock(x)
# upsample
x = layers.Conv2DTranspose(filters=128,
kernel_size=3,
strides=2,
padding='same',
kernel_initializer=RandomNormal(mean=0, stddev=0.02))(x)
x = InstanceNormalization()(x)
x = layers.ReLU()(x)
x = layers.Conv2DTranspose(filters=64,
kernel_size=3,
strides=2,
padding='same',
kernel_initializer=RandomNormal(mean=0, stddev=0.02))(x)
x = InstanceNormalization()(x)
x = layers.ReLU()(x)
# final
x = ReflectionPadding2D(padding=(3, 3))(x)
x = layers.Conv2D(filters=3, kernel_size=7, activation='tanh', kernel_initializer=RandomNormal(mean=0,
stddev=0.02))(x)
return Model(inputs=x0, outputs=x)
def _resblock(x0, num_filter=256, kernel_size=3):
"""Residual block architecture"""
x = ReflectionPadding2D()(x0)
x = layers.Conv2D(filters=num_filter, kernel_size=kernel_size, kernel_initializer=RandomNormal(mean=0,
stddev=0.02))(x)
x = InstanceNormalization()(x)
x = layers.ReLU()(x)
x = ReflectionPadding2D()(x)
x = layers.Conv2D(filters=num_filter, kernel_size=kernel_size, kernel_initializer=RandomNormal(mean=0,
stddev=0.02))(x)
x = InstanceNormalization()(x)
x = layers.Add()([x, x0])
return x
def build_discriminator(input_shape=(256, 256, 3)):
"""Discriminator network architecture"""
x0 = layers.Input(input_shape)
x = layers.Conv2D(filters=64,
kernel_size=4,
strides=2,
padding='same',
kernel_initializer=RandomNormal(mean=0, stddev=0.02))(x0)
x = layers.LeakyReLU(0.2)(x)
x = layers.Conv2D(filters=128,
kernel_size=4,
strides=2,
padding='same',
kernel_initializer=RandomNormal(mean=0, stddev=0.02))(x)
x = InstanceNormalization()(x)
x = layers.LeakyReLU(0.2)(x)
x = layers.Conv2D(filters=256,
kernel_size=4,
strides=2,
padding='same',
kernel_initializer=RandomNormal(mean=0, stddev=0.02))(x)
x = InstanceNormalization()(x)
x = layers.LeakyReLU(0.2)(x)
x = ReflectionPadding2D()(x)
x = layers.Conv2D(filters=512,
kernel_size=4,
strides=1,
kernel_initializer=RandomNormal(mean=0, stddev=0.02))(x)
x = InstanceNormalization()(x)
x = layers.LeakyReLU(0.2)(x)
x = ReflectionPadding2D()(x)
x = layers.Conv2D(filters=1,
kernel_size=4,
strides=1,
kernel_initializer=RandomNormal(mean=0, stddev=0.02))(x)
return Model(inputs=x0, outputs=x)
def StyleTransferNet(input_shape=(256, 256, 3), num_resblock=5):
"""Creates the Style Transfer Network.
"""
x0 = tf.keras.layers.Input(shape=input_shape)
x = ReflectionPadding2D(padding=(40, 40))(x0)
x = _conv_block(x, num_filter=32)
x = _downsample(x, num_filter=64)
x = _downsample(x, num_filter=128)
for _ in range(num_resblock):
x = _residual_block(x, num_filter=128)
x = _upsample(x, num_filter=64)
x = _upsample(x, num_filter=32)
x = _conv_block(x, num_filter=3, apply_relu=False)
x = tf.keras.layers.Activation("tanh")(x)
return tf.keras.Model(inputs=x0, outputs=x)