diff --git a/models/Model_DF/Model.py b/models/Model_DF/Model.py index ee7dbfa..631e9e9 100644 --- a/models/Model_DF/Model.py +++ b/models/Model_DF/Model.py @@ -4,13 +4,21 @@ from nnlib import nnlib from models import ModelBase from facelib import FaceType from samples import * +from utils.console_utils import * class Model(ModelBase): encoderH5 = 'encoder.h5' decoder_srcH5 = 'decoder_src.h5' decoder_dstH5 = 'decoder_dst.h5' - + + #override + def onInitializeOptions(self, is_first_run, ask_override): + if is_first_run or ask_override: + self.options['pixel_loss'] = self.options['pixel_loss'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", False, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 30-40k epochs to enhance fine details and remove face jitter.") + else: + self.options['pixel_loss'] = self.options.get('pixel_loss', False) + #override def onInitialize(self, **in_options): exec(nnlib.import_all(), locals(), globals()) @@ -29,8 +37,8 @@ class Model(ModelBase): self.autoencoder_src = Model([ae_input_layer,mask_layer], self.decoder_src(self.encoder(ae_input_layer))) self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder_dst(self.encoder(ae_input_layer))) - self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] ) - self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] ) + self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] ) + self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] ) if self.is_training_mode: f = SampleProcessor.TypeFlags diff --git a/models/Model_H128/Model.py b/models/Model_H128/Model.py index 024a3c1..b6f114e 100644 --- a/models/Model_H128/Model.py +++ b/models/Model_H128/Model.py @@ -22,6 +22,11 @@ class Model(ModelBase): self.options.pop ('created_vram_gb') self.options['lighter_ae'] = self.options.get('lighter_ae', default_lighter_ae) + if is_first_run or ask_override: + self.options['pixel_loss'] = self.options['pixel_loss'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", False, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 30-40k epochs to enhance fine details and remove face jitter.") + else: + self.options['pixel_loss'] = self.options.get('pixel_loss', False) + #override def onInitialize(self, **in_options): exec(nnlib.import_all(), locals(), globals()) @@ -44,7 +49,7 @@ class Model(ModelBase): self.ae = Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] ) self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), - loss=[ DSSIMMaskLoss([input_src_mask]), 'mae', DSSIMMaskLoss([input_dst_mask]), 'mae' ] ) + loss=[ DSSIMMSEMaskLoss(input_src_mask, is_mse=self.options['pixel_loss']), 'mae', DSSIMMSEMaskLoss(input_dst_mask, is_mse=self.options['pixel_loss']), 'mae' ] ) self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask]) self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask]) diff --git a/models/Model_H64/Model.py b/models/Model_H64/Model.py index 8bd6789..667cd30 100644 --- a/models/Model_H64/Model.py +++ b/models/Model_H64/Model.py @@ -21,7 +21,12 @@ class Model(ModelBase): if 'created_vram_gb' in self.options.keys(): self.options.pop ('created_vram_gb') self.options['lighter_ae'] = self.options.get('lighter_ae', default_lighter_ae) - + + if is_first_run or ask_override: + self.options['pixel_loss'] = self.options['pixel_loss'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", False, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 30-40k epochs to enhance fine details and remove face jitter.") + else: + self.options['pixel_loss'] = self.options.get('pixel_loss', False) + #override def onInitialize(self, **in_options): exec(nnlib.import_all(), locals(), globals()) @@ -44,9 +49,8 @@ class Model(ModelBase): rec_dst_bgr, rec_dst_mask = self.decoder_dst( self.encoder(input_dst_bgr) ) self.ae = Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] ) - - self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), - loss=[ DSSIMMaskLoss([input_src_mask]), 'mae', DSSIMMaskLoss([input_dst_mask]), 'mae' ] ) + + self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[ DSSIMMSEMaskLoss(input_src_mask, is_mse=self.options['pixel_loss']), 'mae', DSSIMMSEMaskLoss(input_dst_mask, is_mse=self.options['pixel_loss']), 'mae' ] ) self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask]) self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask]) diff --git a/models/Model_LIAEF128/Model.py b/models/Model_LIAEF128/Model.py index e49ea58..ad7a3d8 100644 --- a/models/Model_LIAEF128/Model.py +++ b/models/Model_LIAEF128/Model.py @@ -4,6 +4,7 @@ from nnlib import nnlib from models import ModelBase from facelib import FaceType from samples import * +from utils.console_utils import * class Model(ModelBase): @@ -11,7 +12,14 @@ class Model(ModelBase): decoderH5 = 'decoder.h5' inter_BH5 = 'inter_B.h5' inter_ABH5 = 'inter_AB.h5' - + + #override + def onInitializeOptions(self, is_first_run, ask_override): + if is_first_run or ask_override: + self.options['pixel_loss'] = self.options['pixel_loss'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", False, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 30-40k epochs to enhance fine details and remove face jitter.") + else: + self.options['pixel_loss'] = self.options.get('pixel_loss', False) + #override def onInitialize(self, **in_options): exec(nnlib.import_all(), locals(), globals()) @@ -34,8 +42,8 @@ class Model(ModelBase): self.autoencoder_src = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([AB, AB])) ) self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([B, AB])) ) - self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] ) - self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] ) + self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] ) + self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] ) if self.is_training_mode: f = SampleProcessor.TypeFlags diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index 31dbd8e..9920434 100644 --- a/nnlib/nnlib.py +++ b/nnlib/nnlib.py @@ -37,7 +37,7 @@ class nnlib(object): modelify = None ReflectionPadding2D = None DSSIMLoss = None - DSSIMMaskLoss = None + DSSIMMSEMaskLoss = None PixelShuffler = None SubpixelUpscaler = None AddUniformNoise = None @@ -101,7 +101,7 @@ Adam = keras.optimizers.Adam modelify = nnlib.modelify ReflectionPadding2D = nnlib.ReflectionPadding2D DSSIMLoss = nnlib.DSSIMLoss -DSSIMMaskLoss = nnlib.DSSIMMaskLoss +DSSIMMSEMaskLoss = nnlib.DSSIMMSEMaskLoss PixelShuffler = nnlib.PixelShuffler SubpixelUpscaler = nnlib.SubpixelUpscaler AddUniformNoise = nnlib.AddUniformNoise @@ -417,7 +417,8 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator tf = nnlib.tf keras = nnlib.keras K = keras.backend - + exec (nnlib.code_import_tf, locals(), globals()) + def modelify(model_functor): def func(tensor): return keras.models.Model (tensor, model_functor(tensor)) @@ -451,29 +452,21 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator return (1.0 - tf.image.ssim ((y_true/2+0.5), (y_pred/2+0.5), 1.0)) / 2.0 nnlib.DSSIMLoss = DSSIMLoss - class DSSIMMaskLoss(object): - def __init__(self, mask_list, is_tanh=False): - self.mask_list = mask_list - self.is_tanh = is_tanh + class DSSIMMSEMaskLoss(object): + def __init__(self, mask, is_mse=False): + self.mask = mask + self.is_mse = is_mse def __call__(self,y_true, y_pred): total_loss = None - for mask in self.mask_list: - - if not self.is_tanh: - loss = (1.0 - (tf.image.ssim (y_true*mask, y_pred*mask, 1.0))) / 2.0 - else: - loss = (1.0 - tf.image.ssim ( (y_true/2+0.5)*(mask/2+0.5), (y_pred/2+0.5)*(mask/2+0.5), 1.0)) / 2.0 - - loss = K.cast (loss, K.floatx()) - - if total_loss is None: - total_loss = loss - else: - total_loss += loss - - return total_loss - nnlib.DSSIMMaskLoss = DSSIMMaskLoss + + mask = self.mask + if self.is_mse: + blur_mask = tf_gaussian_blur(max(1, mask.get_shape().as_list()[1] // 32))(mask) + return K.mean ( 100*K.square( y_true*blur_mask - y_pred*blur_mask ) ) + else: + return (1.0 - (tf.image.ssim (y_true*mask, y_pred*mask, 1.0))) / 2.0 + nnlib.DSSIMMSEMaskLoss = DSSIMMSEMaskLoss class PixelShuffler(keras.layers.Layer): def __init__(self, size=(2, 2), data_format=None, **kwargs):