github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-15 01:23:44 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

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

Browse source
This commit is contained in:
iperov 2019-02-11 12:05:54 +04:00
parent af3dd59f67
commit f8e63970d2
5 changed files with 52 additions and 34 deletions
Download patch file Download diff file Expand all files Collapse all files

12
models/Model_DF/Model.py
View file

@ -4,6 +4,7 @@ from nnlib import nnlib
from models import ModelBase from models import ModelBase
from facelib import FaceType from facelib import FaceType
from samples import * from samples import *
from utils.console_utils import *
class Model(ModelBase): class Model(ModelBase):
@ -11,6 +12,13 @@ class Model(ModelBase):
decoder_srcH5 = 'decoder_src.h5' decoder_srcH5 = 'decoder_src.h5'
decoder_dstH5 = 'decoder_dst.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 #override
def onInitialize(self, **in_options): def onInitialize(self, **in_options):
exec(nnlib.import_all(), locals(), globals()) 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_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_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_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=[DSSIMMaskLoss([mask_layer]), '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: if self.is_training_mode:
f = SampleProcessor.TypeFlags f = SampleProcessor.TypeFlags

7
models/Model_H128/Model.py
View file

@ -22,6 +22,11 @@ class Model(ModelBase):
self.options.pop ('created_vram_gb') self.options.pop ('created_vram_gb')
self.options['lighter_ae'] = self.options.get('lighter_ae', default_lighter_ae) 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 #override
def onInitialize(self, **in_options): def onInitialize(self, **in_options):
exec(nnlib.import_all(), locals(), globals()) 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 = 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), 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.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]) self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask])

8
models/Model_H64/Model.py
View file

@ -22,6 +22,11 @@ class Model(ModelBase):
self.options.pop ('created_vram_gb') self.options.pop ('created_vram_gb')
self.options['lighter_ae'] = self.options.get('lighter_ae', default_lighter_ae) 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 #override
def onInitialize(self, **in_options): def onInitialize(self, **in_options):
exec(nnlib.import_all(), locals(), globals()) exec(nnlib.import_all(), locals(), globals())
@ -45,8 +50,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 = 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), 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' ] )
loss=[ DSSIMMaskLoss([input_src_mask]), 'mae', DSSIMMaskLoss([input_dst_mask]), 'mae' ] )
self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask]) 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]) self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask])

12
models/Model_LIAEF128/Model.py
View file

@ -4,6 +4,7 @@ from nnlib import nnlib
from models import ModelBase from models import ModelBase
from facelib import FaceType from facelib import FaceType
from samples import * from samples import *
from utils.console_utils import *
class Model(ModelBase): class Model(ModelBase):
@ -12,6 +13,13 @@ class Model(ModelBase):
inter_BH5 = 'inter_B.h5' inter_BH5 = 'inter_B.h5'
inter_ABH5 = 'inter_AB.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 #override
def onInitialize(self, **in_options): def onInitialize(self, **in_options):
exec(nnlib.import_all(), locals(), globals()) 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_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_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_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=[DSSIMMaskLoss([mask_layer]), '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: if self.is_training_mode:
f = SampleProcessor.TypeFlags f = SampleProcessor.TypeFlags

35
nnlib/nnlib.py
View file

@ -37,7 +37,7 @@ class nnlib(object):
modelify = None modelify = None
ReflectionPadding2D = None ReflectionPadding2D = None
DSSIMLoss = None DSSIMLoss = None
DSSIMMaskLoss = None DSSIMMSEMaskLoss = None
PixelShuffler = None PixelShuffler = None
SubpixelUpscaler = None SubpixelUpscaler = None
AddUniformNoise = None AddUniformNoise = None
@ -101,7 +101,7 @@ Adam = keras.optimizers.Adam
modelify = nnlib.modelify modelify = nnlib.modelify
ReflectionPadding2D = nnlib.ReflectionPadding2D ReflectionPadding2D = nnlib.ReflectionPadding2D
DSSIMLoss = nnlib.DSSIMLoss DSSIMLoss = nnlib.DSSIMLoss
DSSIMMaskLoss = nnlib.DSSIMMaskLoss DSSIMMSEMaskLoss = nnlib.DSSIMMSEMaskLoss
PixelShuffler = nnlib.PixelShuffler PixelShuffler = nnlib.PixelShuffler
SubpixelUpscaler = nnlib.SubpixelUpscaler SubpixelUpscaler = nnlib.SubpixelUpscaler
AddUniformNoise = nnlib.AddUniformNoise AddUniformNoise = nnlib.AddUniformNoise
@ -417,6 +417,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
tf = nnlib.tf tf = nnlib.tf
keras = nnlib.keras keras = nnlib.keras
K = keras.backend K = keras.backend
exec (nnlib.code_import_tf, locals(), globals())
def modelify(model_functor): def modelify(model_functor):
def func(tensor): def func(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 return (1.0 - tf.image.ssim ((y_true/2+0.5), (y_pred/2+0.5), 1.0)) / 2.0
nnlib.DSSIMLoss = DSSIMLoss nnlib.DSSIMLoss = DSSIMLoss
class DSSIMMaskLoss(object): class DSSIMMSEMaskLoss(object):
def __init__(self, mask_list, is_tanh=False): def __init__(self, mask, is_mse=False):
self.mask_list = mask_list self.mask = mask
self.is_tanh = is_tanh self.is_mse = is_mse
def __call__(self,y_true, y_pred): def __call__(self,y_true, y_pred):
total_loss = None total_loss = None
for mask in self.mask_list:
if not self.is_tanh: mask = self.mask
loss = (1.0 - (tf.image.ssim (y_true*mask, y_pred*mask, 1.0))) / 2.0 if self.is_mse:
else: blur_mask = tf_gaussian_blur(max(1, mask.get_shape().as_list()[1] // 32))(mask)
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 return K.mean ( 100*K.square( y_true*blur_mask - y_pred*blur_mask ) )
else:
loss = K.cast (loss, K.floatx()) return (1.0 - (tf.image.ssim (y_true*mask, y_pred*mask, 1.0))) / 2.0
nnlib.DSSIMMSEMaskLoss = DSSIMMSEMaskLoss
if total_loss is None:
total_loss = loss
else:
total_loss += loss
return total_loss
nnlib.DSSIMMaskLoss = DSSIMMaskLoss
class PixelShuffler(keras.layers.Layer): class PixelShuffler(keras.layers.Layer):
def __init__(self, size=(2, 2), data_format=None, **kwargs): def __init__(self, size=(2, 2), data_format=None, **kwargs):