|
|
|
@@ -124,25 +124,26 @@ class SAEModel(ModelBase):
|
|
|
|
|
target_dst_ar = [ Input ( ( bgr_shape[0] // (2**i) ,)*2 + (bgr_shape[-1],) ) for i in range(ms_count-1, -1, -1)]
|
|
|
|
|
target_dstm_ar = [ Input ( ( mask_shape[0] // (2**i) ,)*2 + (mask_shape[-1],) ) for i in range(ms_count-1, -1, -1)]
|
|
|
|
|
|
|
|
|
|
use_bn = True
|
|
|
|
|
|
|
|
|
|
models_list = []
|
|
|
|
|
weights_to_load = []
|
|
|
|
|
if self.options['archi'] == 'liae':
|
|
|
|
|
self.encoder = modelify(SAEModel.LIAEEncFlow(resolution, self.options['lighter_encoder'], ed_ch_dims=ed_ch_dims) ) (Input(bgr_shape))
|
|
|
|
|
self.encoder = modelify(SAEModel.LIAEEncFlow(resolution, self.options['lighter_encoder'], ed_ch_dims=ed_ch_dims, use_bn=use_bn) ) (Input(bgr_shape))
|
|
|
|
|
|
|
|
|
|
enc_output_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ]
|
|
|
|
|
|
|
|
|
|
self.inter_B = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims)) (enc_output_Inputs)
|
|
|
|
|
self.inter_AB = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims)) (enc_output_Inputs)
|
|
|
|
|
self.inter_B = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims, use_bn=use_bn)) (enc_output_Inputs)
|
|
|
|
|
self.inter_AB = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims, use_bn=use_bn)) (enc_output_Inputs)
|
|
|
|
|
|
|
|
|
|
inter_output_Inputs = [ Input( np.array(K.int_shape(x)[1:])*(1,1,2) ) for x in self.inter_B.outputs ]
|
|
|
|
|
|
|
|
|
|
self.decoder = modelify(SAEModel.LIAEDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count )) (inter_output_Inputs)
|
|
|
|
|
self.decoder = modelify(SAEModel.LIAEDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count, use_bn=use_bn )) (inter_output_Inputs)
|
|
|
|
|
|
|
|
|
|
models_list += [self.encoder, self.inter_B, self.inter_AB, self.decoder]
|
|
|
|
|
|
|
|
|
|
if self.options['learn_mask']:
|
|
|
|
|
self.decoderm = modelify(SAEModel.LIAEDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (inter_output_Inputs)
|
|
|
|
|
self.decoderm = modelify(SAEModel.LIAEDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5), use_bn=use_bn )) (inter_output_Inputs)
|
|
|
|
|
models_list += [self.decoderm]
|
|
|
|
|
|
|
|
|
|
if not self.is_first_run():
|
|
|
|
@@ -495,6 +496,10 @@ class SAEModel(ModelBase):
|
|
|
|
|
def initialize_nn_functions():
|
|
|
|
|
exec (nnlib.import_all(), locals(), globals())
|
|
|
|
|
|
|
|
|
|
def BatchNorm():
|
|
|
|
|
return BatchNormalization(axis=-1, gamma_initializer=RandomNormal(1., 0.02) )
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ResidualBlock(object):
|
|
|
|
|
def __init__(self, filters, kernel_size=3, padding='same', use_reflection_padding=False):
|
|
|
|
|
self.filters = filters
|
|
|
|
@@ -521,21 +526,30 @@ class SAEModel(ModelBase):
|
|
|
|
|
return var_x
|
|
|
|
|
SAEModel.ResidualBlock = ResidualBlock
|
|
|
|
|
|
|
|
|
|
def downscale (dim):
|
|
|
|
|
def downscale (dim, use_bn=False):
|
|
|
|
|
def func(x):
|
|
|
|
|
return LeakyReLU(0.1)(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x))
|
|
|
|
|
if use_bn:
|
|
|
|
|
return LeakyReLU(0.1)(BatchNorm()(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02), use_bias=False)(x)))
|
|
|
|
|
else:
|
|
|
|
|
return LeakyReLU(0.1)(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x))
|
|
|
|
|
return func
|
|
|
|
|
SAEModel.downscale = downscale
|
|
|
|
|
|
|
|
|
|
def downscale_sep (dim):
|
|
|
|
|
def downscale_sep (dim, use_bn=False):
|
|
|
|
|
def func(x):
|
|
|
|
|
return LeakyReLU(0.1)(SeparableConv2D(dim, kernel_size=5, strides=2, padding='same', depthwise_initializer=RandomNormal(0, 0.02), pointwise_initializer=RandomNormal(0, 0.02) )(x))
|
|
|
|
|
if use_bn:
|
|
|
|
|
return LeakyReLU(0.1)(BatchNorm()(SeparableConv2D(dim, kernel_size=5, strides=2, padding='same', depthwise_initializer=RandomNormal(0, 0.02), pointwise_initializer=RandomNormal(0, 0.02), use_bias=False )(x)))
|
|
|
|
|
else:
|
|
|
|
|
return LeakyReLU(0.1)(SeparableConv2D(dim, kernel_size=5, strides=2, padding='same', depthwise_initializer=RandomNormal(0, 0.02), pointwise_initializer=RandomNormal(0, 0.02) )(x))
|
|
|
|
|
return func
|
|
|
|
|
SAEModel.downscale_sep = downscale_sep
|
|
|
|
|
|
|
|
|
|
def upscale (dim):
|
|
|
|
|
def upscale (dim, use_bn=False):
|
|
|
|
|
def func(x):
|
|
|
|
|
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0, 0.02) )(x)))
|
|
|
|
|
if use_bn:
|
|
|
|
|
return SubpixelUpscaler()(LeakyReLU(0.1)(BatchNorm()(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0,0.02), use_bias=False )(x))))
|
|
|
|
|
else:
|
|
|
|
|
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0, 0.02) )(x)))
|
|
|
|
|
return func
|
|
|
|
|
SAEModel.upscale = upscale
|
|
|
|
|
|
|
|
|
@@ -547,7 +561,7 @@ class SAEModel(ModelBase):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def LIAEEncFlow(resolution, light_enc, ed_ch_dims=42):
|
|
|
|
|
def LIAEEncFlow(resolution, light_enc, ed_ch_dims=42, use_bn=False):
|
|
|
|
|
exec (nnlib.import_all(), locals(), globals())
|
|
|
|
|
upscale = SAEModel.upscale
|
|
|
|
|
downscale = SAEModel.downscale
|
|
|
|
@@ -559,20 +573,20 @@ class SAEModel(ModelBase):
|
|
|
|
|
x = input
|
|
|
|
|
x = downscale(ed_dims)(x)
|
|
|
|
|
if not light_enc:
|
|
|
|
|
x = downscale(ed_dims*2)(x)
|
|
|
|
|
x = downscale(ed_dims*4)(x)
|
|
|
|
|
x = downscale(ed_dims*8)(x)
|
|
|
|
|
x = downscale(ed_dims*2, use_bn=use_bn)(x)
|
|
|
|
|
x = downscale(ed_dims*4, use_bn=use_bn)(x)
|
|
|
|
|
x = downscale(ed_dims*8, use_bn=use_bn)(x)
|
|
|
|
|
else:
|
|
|
|
|
x = downscale_sep(ed_dims*2)(x)
|
|
|
|
|
x = downscale(ed_dims*4)(x)
|
|
|
|
|
x = downscale_sep(ed_dims*8)(x)
|
|
|
|
|
x = downscale_sep(ed_dims*2, use_bn=use_bn)(x)
|
|
|
|
|
x = downscale(ed_dims*4, use_bn=use_bn)(x)
|
|
|
|
|
x = downscale_sep(ed_dims*8, use_bn=use_bn)(x)
|
|
|
|
|
|
|
|
|
|
x = Flatten()(x)
|
|
|
|
|
return x
|
|
|
|
|
return func
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def LIAEInterFlow(resolution, ae_dims=256):
|
|
|
|
|
def LIAEInterFlow(resolution, ae_dims=256, use_bn=False):
|
|
|
|
|
exec (nnlib.import_all(), locals(), globals())
|
|
|
|
|
upscale = SAEModel.upscale
|
|
|
|
|
lowest_dense_res=resolution // 16
|
|
|
|
@@ -582,12 +596,12 @@ class SAEModel(ModelBase):
|
|
|
|
|
x = Dense(ae_dims)(x)
|
|
|
|
|
x = Dense(lowest_dense_res * lowest_dense_res * ae_dims*2)(x)
|
|
|
|
|
x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims*2))(x)
|
|
|
|
|
x = upscale(ae_dims*2)(x)
|
|
|
|
|
x = upscale(ae_dims*2, use_bn=use_bn)(x)
|
|
|
|
|
return x
|
|
|
|
|
return func
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def LIAEDecFlow(output_nc,ed_ch_dims=21, multiscale_count=1):
|
|
|
|
|
def LIAEDecFlow(output_nc,ed_ch_dims=21, multiscale_count=1, use_bn=False):
|
|
|
|
|
exec (nnlib.import_all(), locals(), globals())
|
|
|
|
|
upscale = SAEModel.upscale
|
|
|
|
|
to_bgr = SAEModel.to_bgr
|
|
|
|
@@ -597,17 +611,17 @@ class SAEModel(ModelBase):
|
|
|
|
|
x = input[0]
|
|
|
|
|
|
|
|
|
|
outputs = []
|
|
|
|
|
x1 = upscale(ed_dims*8)( x )
|
|
|
|
|
x1 = upscale(ed_dims*8, use_bn=use_bn)( x )
|
|
|
|
|
|
|
|
|
|
if multiscale_count >= 3:
|
|
|
|
|
outputs += [ to_bgr(output_nc) ( x1 ) ]
|
|
|
|
|
|
|
|
|
|
x2 = upscale(ed_dims*4)( x1 )
|
|
|
|
|
x2 = upscale(ed_dims*4, use_bn=use_bn)( x1 )
|
|
|
|
|
|
|
|
|
|
if multiscale_count >= 2:
|
|
|
|
|
outputs += [ to_bgr(output_nc) ( x2 ) ]
|
|
|
|
|
|
|
|
|
|
x3 = upscale(ed_dims*2)( x2 )
|
|
|
|
|
x3 = upscale(ed_dims*2, use_bn=use_bn)( x2 )
|
|
|
|
|
|
|
|
|
|
outputs += [ to_bgr(output_nc) ( x3 ) ]
|
|
|
|
|
|
|
|
|
|