networks/generator.py

from tensorflow.keras.layers import *
from tensorflow.keras.models import Model
from tensorflow_addons.layers import InstanceNormalization
from networks.layers import AdaIN, AdaptiveAttention

import numpy as np


def residual_down_block(inputs, filters, resample=True):
 x = inputs

 r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
 if resample:
 r = AveragePooling2D()(r)

 x = InstanceNormalization()(x)
 x = LeakyReLU(0.2)(x)
 x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)

 if resample:
 x = AveragePooling2D()(x)

 x = Add()([x, r])

 return x


def residual_up_block(inputs, filters, resample=True, name=None):
 x, z_id = inputs

 r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
 if resample:
 r = UpSampling2D(interpolation='bilinear')(r)

 x = InstanceNormalization()(x)
 x = AdaIN()([x, z_id])
 x = LeakyReLU(0.2)(x)
 x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)

 if resample:
 x = UpSampling2D(interpolation='bilinear')(x)

 x = Add()([x, r])

 return x


def adaptive_attention(inputs, filters, name=None):
 x_t, x_s = inputs

 m = Concatenate(axis=-1)([x_t, x_s])
 m = Conv2D(filters=filters // 4, kernel_size=3, strides=1, padding='same')(m)
 m = LeakyReLU(0.2)(m)
 m = InstanceNormalization()(m)
 m = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same', activation='sigmoid', name=name)(m)

 x = AdaptiveAttention()([m, x_t, x_s])

 return x


def adaptive_fusion_up_block(inputs, filters, resample=True, name=None):
 x_t, x_s, z_id = inputs

 x = adaptive_attention([x_t, x_s], x_t.shape[-1], name=name)

 r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
 if resample:
 r = UpSampling2D(interpolation='bilinear')(r)

 x = InstanceNormalization()(x)
 x = AdaIN()([x, z_id])
 x = LeakyReLU(0.2)(x)
 x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)

 if resample:
 x = UpSampling2D(interpolation='bilinear')(x)

 x = Add()([x, r])

 return x


def dual_adaptive_fusion_up_block(inputs, filters, resample=True, name=None):
 x_t, x_s, z_id = inputs

 x = adaptive_attention([x_t, x_s], x_t.shape[-1], name=name + '_0')
 x = adaptive_attention([x_t, x], x_t.shape[-1], name=name + '_1')

 r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
 if resample:
 r = UpSampling2D(interpolation='bilinear')(r)

 x = InstanceNormalization()(x)
 x = AdaIN()([x, z_id])
 x = LeakyReLU(0.2)(x)
 x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)

 if resample:
 x = UpSampling2D(interpolation='bilinear')(x)

 x = Add()([x, r])

 return x


def adaptive_fusion_up_block_no_add(inputs, filters, resample=True, name=None):
 x_t, x_s, z_id = inputs

 x = adaptive_attention([x_t, x_s], x_t.shape[-1], name=name)

 x = InstanceNormalization()(x)
 x = AdaIN()([x, z_id])
 x = LeakyReLU(0.2)(x)
 x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)

 return x


def adaptive_fusion_up_block_concat_baseline(inputs, filters, resample=True, name=None):
 x_t, x_s, z_id = inputs

 x = Concatenate(axis=-1)([x_t, x_s])

 r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
 if resample:
 r = UpSampling2D(interpolation='bilinear')(r)

 x = InstanceNormalization()(x)
 x = AdaIN()([x, z_id])
 x = LeakyReLU(0.2)(x)
 x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)

 if resample:
 x = UpSampling2D(interpolation='bilinear')(x)

 x = Add(name=name if name == 'final' else None)([x, r])

 return x


def adaptive_fusion_up_block_add_baseline(inputs, filters, resample=True, name=None):
 x_t, x_s, z_id = inputs

 x = Add()([x_t, x_s])

 r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
 if resample:
 r = UpSampling2D(interpolation='bilinear')(r)

 x = InstanceNormalization()(x)
 x = AdaIN()([x, z_id])
 x = LeakyReLU(0.2)(x)
 x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)

 if resample:
 x = UpSampling2D(interpolation='bilinear')(x)

 x = Add()([x, r])

 return x


def get_generator_original(mapping_depth=4, mapping_size=256):
 x_target = Input(shape=(256, 256, 3))
 z_source = Input(shape=(512,))

 z_id = z_source
 for m in range(np.max([mapping_depth - 1, 0])):
 z_id = Dense(mapping_size)(z_id)
 z_id = LeakyReLU(0.2)(z_id)
 if mapping_depth >= 1:
 z_id = Dense(mapping_size)(z_id)

 x_0 = Conv2D(filters=64, kernel_size=3, strides=1, padding='same')(x_target) # 256

 x_1 = residual_down_block(x_0, 128) # 128

 x_2 = residual_down_block(x_1, 256) # 64

 x_3 = residual_down_block(x_2, 512)

 x_4 = residual_down_block(x_3, 512)

 x_5 = residual_down_block(x_4, 512)

 x_6 = residual_down_block(x_5, 512, resample=False)

 u_5 = residual_up_block([x_6, z_id], 512, resample=False)

 u_4 = residual_up_block([u_5, z_id], 512)

 u_3 = residual_up_block([u_4, z_id], 512)

 u_2 = residual_up_block([u_3, z_id], 256) # 64

 u_1 = adaptive_fusion_up_block([x_2, u_2, z_id], 128, name='aff_attention_64x64') # 128

 u_0 = adaptive_fusion_up_block([x_1, u_1, z_id], 64, name='aff_attention_128x128') # 256

 out = adaptive_fusion_up_block([x_0, u_0, z_id], 3, resample=False, name='aff_attention_256x256')

 gen_model = Model([x_target, z_source], out)
 gen_model.summary()

 return gen_model


def make_layer(l_type, inputs, filters, resample, name=None):
 if l_type == 'affa':
 return adaptive_fusion_up_block(inputs, filters, resample=resample, name=name)
 if l_type == 'd_affa':
 return dual_adaptive_fusion_up_block(inputs, filters, resample=resample, name=name)
 elif l_type == 'concat':
 return adaptive_fusion_up_block_concat_baseline(inputs, filters, resample=resample, name=name)
 elif l_type == 'no_skip':
 return residual_up_block(inputs[1:], filters, resample=resample)


def get_generator(up_types=None, mapping_depth=4, mapping_size=256):

 if up_types is None:
 up_types = ['no_skip', 'no_skip', 'd_affa', 'd_affa', 'd_affa', 'concat']

 x_target = Input(shape=(256, 256, 3))
 z_source = Input(shape=(512,))

 z_id = z_source
 for m in range(np.max([mapping_depth - 1, 0])):
 z_id = Dense(mapping_size)(z_id)
 z_id = LeakyReLU(0.2)(z_id)
 if mapping_depth >= 1:
 z_id = Dense(mapping_size)(z_id)

 x_0 = Conv2D(filters=64, kernel_size=3, strides=1, padding='same')(x_target) # 256

 x_1 = residual_down_block(x_0, 128) # 128

 x_2 = residual_down_block(x_1, 256) # 64

 x_3 = residual_down_block(x_2, 512)

 x_4 = residual_down_block(x_3, 512)

 x_5 = residual_down_block(x_4, 512)

 x_6 = residual_down_block(x_5, 512, resample=False)

 u_5 = residual_up_block([x_6, z_id], 512, resample=False)

 u_4 = make_layer(up_types[0], [x_5, u_5, z_id], 512, resample=True, name='16x16')

 u_3 = make_layer(up_types[1], [x_4, u_4, z_id], 512, resample=True, name='32x32')

 u_2 = make_layer(up_types[2], [x_3, u_3, z_id], 256, resample=True, name='64x64')

 u_1 = make_layer(up_types[3], [x_2, u_2, z_id], 128, resample=True, name='128x128')

 u_0 = make_layer(up_types[4], [x_1, u_1, z_id], 64, resample=True, name='256x256')

 out = make_layer(up_types[5], [x_0, u_0, z_id], 3, resample=False, name='final')

 gen_model = Model([x_target, z_source], out)
 gen_model.summary()

 return gen_model


def get_generator_large(up_types=None, mapping_depth=4, mapping_size=512):

 if up_types is None:
 up_types = ['no_skip', 'no_skip', 'affa', 'affa', 'affa', 'concat']

 x_target = Input(shape=(256, 256, 3))
 z_source = Input(shape=(512,))

 z_id = z_source
 for m in range(np.max([mapping_depth - 1, 0])):
 z_id = Dense(mapping_size)(z_id)
 z_id = LeakyReLU(0.2)(z_id)
 if mapping_depth >= 1:
 z_id = Dense(mapping_size)(z_id)

 x_0 = Conv2D(filters=64, kernel_size=3, strides=1, padding='same')(x_target) # 256

 x_1 = residual_down_block(x_0, 128) # 128

 x_2 = residual_down_block(x_1, 256) # 64

 x_3 = residual_down_block(x_2, 512)

 x_4 = residual_down_block(x_3, 512)

 x_5 = residual_down_block(x_4, 512)

 b_0 = residual_up_block([x_5, z_id], 512, resample=False)

 b_1 = residual_up_block([b_0, z_id], 512, resample=False)

 b_2 = residual_up_block([b_1, z_id], 512, resample=False)

 u_5 = residual_up_block([b_2, z_id], 512, resample=False)

 u_4 = make_layer(up_types[0], [x_5, u_5, z_id], 512, resample=True, name='16x16')

 u_3 = make_layer(up_types[1], [x_4, u_4, z_id], 512, resample=True, name='32x32')

 u_2 = make_layer(up_types[2], [x_3, u_3, z_id], 256, resample=True, name='64x64')

 u_1 = make_layer(up_types[3], [x_2, u_2, z_id], 128, resample=True, name='128x128')

 u_0 = make_layer(up_types[4], [x_1, u_1, z_id], 64, resample=True, name='256x256')

 out = make_layer(up_types[5], [x_0, u_0, z_id], 3, resample=False, name='final')

 gen_model = Model([x_target, z_source], out)
 gen_model.summary()

 return gen_model