Permalink
Browse files

H64, H128, DF, LIAEF128: added pixel loss option.

  • Loading branch information...
iperov committed Feb 11, 2019
1 parent af3dd59 commit f8e63970d230067771019442b0305175ac9de5ed
Showing with 52 additions and 34 deletions.
  1. +11 −3 models/Model_DF/Model.py
  2. +6 −1 models/Model_H128/Model.py
  3. +8 −4 models/Model_H64/Model.py
  4. +11 −3 models/Model_LIAEF128/Model.py
  5. +16 −23 nnlib/nnlib.py
@@ -4,13 +4,21 @@
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 @@ def onInitialize(self, **in_options):
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
@@ -22,6 +22,11 @@ def onInitializeOptions(self, is_first_run, ask_override):
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 @@ def onInitialize(self, **in_options):
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])
@@ -21,7 +21,12 @@ def onInitializeOptions(self, is_first_run, ask_override):
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 @@ def onInitialize(self, **in_options):
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])
@@ -4,14 +4,22 @@
from models import ModelBase
from facelib import FaceType
from samples import *
from utils.console_utils import *

class Model(ModelBase):

encoderH5 = 'encoder.h5'
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 @@ def onInitialize(self, **in_options):
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
@@ -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 @@ class nnlib(object):
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 @@ def __initialize_keras_functions():
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 @@ def __call__(self,y_true, y_pred):
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):

0 comments on commit f8e6397

Please sign in to comment.