Spaces:
Running
Running
| 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 | |