mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-14 02:37:00 -07:00
global refactoring and fixes,
removed support of extracted(aligned) PNG faces. Use old builds to convert from PNG to JPG. fanseg model file in facelib/ is renamed
This commit is contained in:
Colombo
parent
921b464d5b
commit
61472cdaf7
82 changed files with 3838 additions and 3812 deletions
|
|
@ -454,7 +454,7 @@ class ModelBase(object):
|
|||
self.generate_next_samples()
|
||||
|
||||
def finalize(self):
|
||||
nn.tf_close_session()
|
||||
nn.close_session()
|
||||
|
||||
def is_first_run(self):
|
||||
return self.iter == 0
|
||||
|
|
|
|||
|
|
@ -18,21 +18,17 @@ class FANSegModel(ModelBase):
|
|||
device_config = nn.getCurrentDeviceConfig()
|
||||
yn_str = {True:'y',False:'n'}
|
||||
|
||||
#default_resolution = 256
|
||||
#default_face_type = self.options['face_type'] = self.load_or_def_option('face_type', 'f')
|
||||
|
||||
ask_override = self.ask_override()
|
||||
if self.is_first_run() or ask_override:
|
||||
self.ask_autobackup_hour()
|
||||
self.ask_target_iter()
|
||||
self.ask_batch_size(24)
|
||||
|
||||
#if self.is_first_run():
|
||||
#resolution = io.input_int("Resolution", default_resolution, add_info="64-512")
|
||||
#resolution = np.clip ( (resolution // 16) * 16, 64, 512)
|
||||
#self.options['resolution'] = resolution
|
||||
#self.options['face_type'] = io.input_str ("Face type", default_face_type, ['f']).lower()
|
||||
|
||||
default_lr_dropout = self.options['lr_dropout'] = self.load_or_def_option('lr_dropout', False)
|
||||
|
||||
if self.is_first_run() or ask_override:
|
||||
self.options['lr_dropout'] = io.input_bool ("Use learning rate dropout", default_lr_dropout, help_message="When the face is trained enough, you can enable this option to get extra sharpness and reduce subpixel shake for less amount of iterations.")
|
||||
|
||||
#override
|
||||
def on_initialize(self):
|
||||
device_config = nn.getCurrentDeviceConfig()
|
||||
|
|
@ -42,11 +38,7 @@ class FANSegModel(ModelBase):
|
|||
device_config = nn.getCurrentDeviceConfig()
|
||||
devices = device_config.devices
|
||||
|
||||
self.resolution = resolution = 256#self.options['resolution']
|
||||
#self.face_type = {'h' : FaceType.HALF,
|
||||
# 'mf' : FaceType.MID_FULL,
|
||||
# 'f' : FaceType.FULL,
|
||||
# 'wf' : FaceType.WHOLE_FACE}[ self.options['face_type'] ]
|
||||
self.resolution = resolution = 256
|
||||
self.face_type = FaceType.FULL
|
||||
|
||||
place_model_on_cpu = len(devices) == 0
|
||||
|
|
@ -56,13 +48,13 @@ class FANSegModel(ModelBase):
|
|||
mask_shape = nn.get4Dshape(resolution,resolution,1)
|
||||
|
||||
# Initializing model classes
|
||||
self.model = TernausNet(f'{self.model_name}_FANSeg',
|
||||
self.model = TernausNet(f'{self.model_name}_FANSeg_{FaceType.toString(self.face_type)}',
|
||||
resolution,
|
||||
FaceType.toString(self.face_type),
|
||||
load_weights=not self.is_first_run(),
|
||||
weights_file_root=self.get_model_root_path(),
|
||||
training=True,
|
||||
place_model_on_cpu=place_model_on_cpu)
|
||||
place_model_on_cpu=place_model_on_cpu,
|
||||
optimizer=nn.RMSprop(lr=0.0001, lr_dropout=0.3 if self.options['lr_dropout'] else 1.0,name='opt') )
|
||||
|
||||
if self.is_training:
|
||||
# Adjust batch size for multiple GPU
|
||||
|
|
@ -93,15 +85,15 @@ class FANSegModel(ModelBase):
|
|||
gpu_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=gpu_target_t, logits=gpu_pred_logits_t), axis=[1,2,3])
|
||||
gpu_losses += [gpu_loss]
|
||||
|
||||
gpu_loss_gvs += [ nn.tf_gradients ( gpu_loss, self.model.net_weights ) ]
|
||||
gpu_loss_gvs += [ nn.gradients ( gpu_loss, self.model.net_weights ) ]
|
||||
|
||||
|
||||
# Average losses and gradients, and create optimizer update ops
|
||||
with tf.device (models_opt_device):
|
||||
pred = nn.tf_concat(gpu_pred_list, 0)
|
||||
pred = nn.concat(gpu_pred_list, 0)
|
||||
loss = tf.reduce_mean(gpu_losses)
|
||||
|
||||
loss_gv_op = self.model.opt.get_update_op (nn.tf_average_gv_list (gpu_loss_gvs))
|
||||
loss_gv_op = self.model.opt.get_update_op (nn.average_gv_list (gpu_loss_gvs))
|
||||
|
||||
|
||||
# Initializing training and view functions
|
||||
|
|
@ -171,7 +163,7 @@ class FANSegModel(ModelBase):
|
|||
result = []
|
||||
st = []
|
||||
for i in range(n_samples):
|
||||
ar = S[i]*TM[i]+ green_bg*(1-TM[i]), SM[i], S[i]*SM[i] + green_bg*(1-SM[i])
|
||||
ar = S[i]*TM[i] + 0.5*S[i]*(1-TM[i]) + 0.5*green_bg*(1-TM[i]), SM[i], S[i]*SM[i] + green_bg*(1-SM[i])
|
||||
st.append ( np.concatenate ( ar, axis=1) )
|
||||
result += [ ('FANSeg training faces', np.concatenate (st, axis=0 )), ]
|
||||
|
||||
|
|
|
|||
|
|
@ -14,146 +14,11 @@ class QModel(ModelBase):
|
|||
#override
|
||||
def on_initialize(self):
|
||||
device_config = nn.getCurrentDeviceConfig()
|
||||
self.model_data_format = "NCHW" if len(device_config.devices) != 0 and not self.is_debug() else "NHWC"
|
||||
devices = device_config.devices
|
||||
self.model_data_format = "NCHW" if len(devices) != 0 and not self.is_debug() else "NHWC"
|
||||
nn.initialize(data_format=self.model_data_format)
|
||||
tf = nn.tf
|
||||
|
||||
conv_kernel_initializer = nn.initializers.ca()
|
||||
|
||||
class Downscale(nn.ModelBase):
|
||||
def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ):
|
||||
self.in_ch = in_ch
|
||||
self.out_ch = out_ch
|
||||
self.kernel_size = kernel_size
|
||||
self.dilations = dilations
|
||||
self.subpixel = subpixel
|
||||
self.use_activator = use_activator
|
||||
super().__init__(*kwargs)
|
||||
|
||||
def on_build(self, *args, **kwargs ):
|
||||
self.conv1 = nn.Conv2D( self.in_ch,
|
||||
self.out_ch // (4 if self.subpixel else 1),
|
||||
kernel_size=self.kernel_size,
|
||||
strides=1 if self.subpixel else 2,
|
||||
padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer )
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
|
||||
if self.subpixel:
|
||||
x = nn.tf_space_to_depth(x, 2)
|
||||
|
||||
if self.use_activator:
|
||||
x = nn.tf_gelu(x)
|
||||
return x
|
||||
|
||||
def get_out_ch(self):
|
||||
return (self.out_ch // 4) * 4
|
||||
|
||||
class DownscaleBlock(nn.ModelBase):
|
||||
def on_build(self, in_ch, ch, n_downscales, kernel_size, dilations=1, subpixel=True):
|
||||
self.downs = []
|
||||
|
||||
last_ch = in_ch
|
||||
for i in range(n_downscales):
|
||||
cur_ch = ch*( min(2**i, 8) )
|
||||
self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size, dilations=dilations, subpixel=subpixel) )
|
||||
last_ch = self.downs[-1].get_out_ch()
|
||||
|
||||
def forward(self, inp):
|
||||
x = inp
|
||||
for down in self.downs:
|
||||
x = down(x)
|
||||
return x
|
||||
|
||||
class Upscale(nn.ModelBase):
|
||||
def on_build(self, in_ch, out_ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = nn.tf_gelu(x)
|
||||
x = nn.tf_depth_to_space(x, 2)
|
||||
return x
|
||||
|
||||
class ResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.conv1(inp)
|
||||
x = nn.tf_gelu(x)
|
||||
x = self.conv2(x)
|
||||
x = inp + x
|
||||
x = nn.tf_gelu(x)
|
||||
return x
|
||||
|
||||
class Encoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, e_ch):
|
||||
self.down1 = DownscaleBlock(in_ch, e_ch, n_downscales=4, kernel_size=5)
|
||||
def forward(self, inp):
|
||||
return nn.tf_flatten(self.down1(inp))
|
||||
|
||||
class Inter(nn.ModelBase):
|
||||
def __init__(self, in_ch, lowest_dense_res, ae_ch, ae_out_ch, d_ch, **kwargs):
|
||||
self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch, self.d_ch = in_ch, lowest_dense_res, ae_ch, ae_out_ch, d_ch
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def on_build(self):
|
||||
in_ch, lowest_dense_res, ae_ch, ae_out_ch, d_ch = self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch, self.d_ch
|
||||
|
||||
self.dense1 = nn.Dense( in_ch, ae_ch, kernel_initializer=tf.initializers.orthogonal )
|
||||
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch, maxout_features=4, kernel_initializer=tf.initializers.orthogonal )
|
||||
self.upscale1 = Upscale(ae_out_ch, d_ch*8)
|
||||
self.res1 = ResidualBlock(d_ch*8)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.dense1(inp)
|
||||
x = self.dense2(x)
|
||||
x = nn.tf_reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch)
|
||||
x = self.upscale1(x)
|
||||
x = self.res1(x)
|
||||
return x
|
||||
|
||||
def get_out_ch(self):
|
||||
return self.ae_out_ch
|
||||
|
||||
class Decoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, d_ch):
|
||||
self.upscale1 = Upscale(in_ch, d_ch*4)
|
||||
self.res1 = ResidualBlock(d_ch*4)
|
||||
self.upscale2 = Upscale(d_ch*4, d_ch*2)
|
||||
self.res2 = ResidualBlock(d_ch*2)
|
||||
self.upscale3 = Upscale(d_ch*2, d_ch*1)
|
||||
self.res3 = ResidualBlock(d_ch*1)
|
||||
|
||||
self.upscalem1 = Upscale(in_ch, d_ch)
|
||||
self.upscalem2 = Upscale(d_ch, d_ch//2)
|
||||
self.upscalem3 = Upscale(d_ch//2, d_ch//2)
|
||||
|
||||
self.out_conv = nn.Conv2D( d_ch*1, 3, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.out_convm = nn.Conv2D( d_ch//2, 1, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, inp):
|
||||
z = inp
|
||||
x = self.upscale1 (z)
|
||||
x = self.res1 (x)
|
||||
x = self.upscale2 (x)
|
||||
x = self.res2 (x)
|
||||
x = self.upscale3 (x)
|
||||
x = self.res3 (x)
|
||||
|
||||
y = self.upscalem1 (z)
|
||||
y = self.upscalem2 (y)
|
||||
y = self.upscalem3 (y)
|
||||
|
||||
return tf.nn.sigmoid(self.out_conv(x)), \
|
||||
tf.nn.sigmoid(self.out_convm(y))
|
||||
|
||||
device_config = nn.getCurrentDeviceConfig()
|
||||
devices = device_config.devices
|
||||
|
||||
resolution = self.resolution = 96
|
||||
self.face_type = FaceType.FULL
|
||||
ae_dims = 128
|
||||
|
|
@ -169,35 +34,34 @@ class QModel(ModelBase):
|
|||
optimizer_vars_on_cpu = models_opt_device=='/CPU:0'
|
||||
|
||||
input_ch = 3
|
||||
output_ch = 3
|
||||
bgr_shape = nn.get4Dshape(resolution,resolution,input_ch)
|
||||
mask_shape = nn.get4Dshape(resolution,resolution,1)
|
||||
lowest_dense_res = resolution // 16
|
||||
|
||||
self.model_filename_list = []
|
||||
|
||||
|
||||
model_archi = nn.DeepFakeArchi(resolution, mod='quick')
|
||||
|
||||
with tf.device ('/CPU:0'):
|
||||
#Place holders on CPU
|
||||
self.warped_src = tf.placeholder (nn.tf_floatx, bgr_shape)
|
||||
self.warped_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
|
||||
self.warped_src = tf.placeholder (nn.floatx, bgr_shape)
|
||||
self.warped_dst = tf.placeholder (nn.floatx, bgr_shape)
|
||||
|
||||
self.target_src = tf.placeholder (nn.tf_floatx, bgr_shape)
|
||||
self.target_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
|
||||
self.target_src = tf.placeholder (nn.floatx, bgr_shape)
|
||||
self.target_dst = tf.placeholder (nn.floatx, bgr_shape)
|
||||
|
||||
self.target_srcm = tf.placeholder (nn.tf_floatx, mask_shape)
|
||||
self.target_dstm = tf.placeholder (nn.tf_floatx, mask_shape)
|
||||
self.target_srcm = tf.placeholder (nn.floatx, mask_shape)
|
||||
self.target_dstm = tf.placeholder (nn.floatx, mask_shape)
|
||||
|
||||
# Initializing model classes
|
||||
with tf.device (models_opt_device):
|
||||
self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, name='encoder')
|
||||
encoder_out_ch = self.encoder.compute_output_channels ( (nn.tf_floatx, bgr_shape))
|
||||
self.encoder = model_archi.Encoder(in_ch=input_ch, e_ch=e_dims, name='encoder')
|
||||
encoder_out_ch = self.encoder.compute_output_channels ( (nn.floatx, bgr_shape))
|
||||
|
||||
self.inter = Inter (in_ch=encoder_out_ch, lowest_dense_res=lowest_dense_res, ae_ch=ae_dims, ae_out_ch=ae_dims, d_ch=d_dims, name='inter')
|
||||
inter_out_ch = self.inter.compute_output_channels ( (nn.tf_floatx, (None,encoder_out_ch)))
|
||||
self.inter = model_archi.Inter (in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims, d_ch=d_dims, name='inter')
|
||||
inter_out_ch = self.inter.compute_output_channels ( (nn.floatx, (None,encoder_out_ch)))
|
||||
|
||||
self.decoder_src = Decoder(in_ch=inter_out_ch, d_ch=d_dims, name='decoder_src')
|
||||
self.decoder_dst = Decoder(in_ch=inter_out_ch, d_ch=d_dims, name='decoder_dst')
|
||||
self.decoder_src = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, name='decoder_src')
|
||||
self.decoder_dst = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, name='decoder_dst')
|
||||
|
||||
self.model_filename_list += [ [self.encoder, 'encoder.npy' ],
|
||||
[self.inter, 'inter.npy' ],
|
||||
|
|
@ -208,7 +72,7 @@ class QModel(ModelBase):
|
|||
self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()
|
||||
|
||||
# Initialize optimizers
|
||||
self.src_dst_opt = nn.TFRMSpropOptimizer(lr=2e-4, lr_dropout=0.3, name='src_dst_opt')
|
||||
self.src_dst_opt = nn.RMSprop(lr=2e-4, lr_dropout=0.3, name='src_dst_opt')
|
||||
self.src_dst_opt.initialize_variables(self.src_dst_trainable_weights, vars_on_cpu=optimizer_vars_on_cpu )
|
||||
self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]
|
||||
|
||||
|
|
@ -257,8 +121,8 @@ class QModel(ModelBase):
|
|||
gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
|
||||
gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)
|
||||
|
||||
gpu_target_srcm_blur = nn.tf_gaussian_blur(gpu_target_srcm, max(1, resolution // 32) )
|
||||
gpu_target_dstm_blur = nn.tf_gaussian_blur(gpu_target_dstm, max(1, resolution // 32) )
|
||||
gpu_target_srcm_blur = nn.gaussian_blur(gpu_target_srcm, max(1, resolution // 32) )
|
||||
gpu_target_dstm_blur = nn.gaussian_blur(gpu_target_dstm, max(1, resolution // 32) )
|
||||
|
||||
gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur
|
||||
gpu_target_dst_anti_masked = gpu_target_dst*(1.0 - gpu_target_dstm_blur)
|
||||
|
|
@ -272,11 +136,11 @@ class QModel(ModelBase):
|
|||
gpu_psd_target_dst_masked = gpu_pred_src_dst*gpu_target_dstm_blur
|
||||
gpu_psd_target_dst_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_blur)
|
||||
|
||||
gpu_src_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
|
||||
gpu_src_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
|
||||
gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
|
||||
gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )
|
||||
|
||||
gpu_dst_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
|
||||
gpu_dst_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
|
||||
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
|
||||
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )
|
||||
|
||||
|
|
@ -284,21 +148,21 @@ class QModel(ModelBase):
|
|||
gpu_dst_losses += [gpu_dst_loss]
|
||||
|
||||
gpu_G_loss = gpu_src_loss + gpu_dst_loss
|
||||
gpu_src_dst_loss_gvs += [ nn.tf_gradients ( gpu_G_loss, self.src_dst_trainable_weights ) ]
|
||||
gpu_src_dst_loss_gvs += [ nn.gradients ( gpu_G_loss, self.src_dst_trainable_weights ) ]
|
||||
|
||||
|
||||
# Average losses and gradients, and create optimizer update ops
|
||||
with tf.device (models_opt_device):
|
||||
pred_src_src = nn.tf_concat(gpu_pred_src_src_list, 0)
|
||||
pred_dst_dst = nn.tf_concat(gpu_pred_dst_dst_list, 0)
|
||||
pred_src_dst = nn.tf_concat(gpu_pred_src_dst_list, 0)
|
||||
pred_src_srcm = nn.tf_concat(gpu_pred_src_srcm_list, 0)
|
||||
pred_dst_dstm = nn.tf_concat(gpu_pred_dst_dstm_list, 0)
|
||||
pred_src_dstm = nn.tf_concat(gpu_pred_src_dstm_list, 0)
|
||||
pred_src_src = nn.concat(gpu_pred_src_src_list, 0)
|
||||
pred_dst_dst = nn.concat(gpu_pred_dst_dst_list, 0)
|
||||
pred_src_dst = nn.concat(gpu_pred_src_dst_list, 0)
|
||||
pred_src_srcm = nn.concat(gpu_pred_src_srcm_list, 0)
|
||||
pred_dst_dstm = nn.concat(gpu_pred_dst_dstm_list, 0)
|
||||
pred_src_dstm = nn.concat(gpu_pred_src_dstm_list, 0)
|
||||
|
||||
src_loss = nn.tf_average_tensor_list(gpu_src_losses)
|
||||
dst_loss = nn.tf_average_tensor_list(gpu_dst_losses)
|
||||
src_dst_loss_gv = nn.tf_average_gv_list (gpu_src_dst_loss_gvs)
|
||||
src_loss = nn.average_tensor_list(gpu_src_losses)
|
||||
dst_loss = nn.average_tensor_list(gpu_dst_losses)
|
||||
src_dst_loss_gv = nn.average_gv_list (gpu_src_dst_loss_gvs)
|
||||
src_dst_loss_gv_op = self.src_dst_opt.get_update_op (src_dst_loss_gv)
|
||||
|
||||
# Initializing training and view functions
|
||||
|
|
|
|||
|
|
@ -119,13 +119,11 @@ class SAEHDModel(ModelBase):
|
|||
#override
|
||||
def on_initialize(self):
|
||||
device_config = nn.getCurrentDeviceConfig()
|
||||
self.model_data_format = "NCHW" if len(device_config.devices) != 0 and not self.is_debug() else "NHWC"
|
||||
devices = device_config.devices
|
||||
self.model_data_format = "NCHW" if len(devices) != 0 and not self.is_debug() else "NHWC"
|
||||
nn.initialize(data_format=self.model_data_format)
|
||||
tf = nn.tf
|
||||
|
||||
device_config = nn.getCurrentDeviceConfig()
|
||||
devices = device_config.devices
|
||||
|
||||
self.resolution = resolution = self.options['resolution']
|
||||
self.face_type = {'h' : FaceType.HALF,
|
||||
'mf' : FaceType.MID_FULL,
|
||||
|
|
@ -154,41 +152,38 @@ class SAEHDModel(ModelBase):
|
|||
models_opt_device = '/GPU:0' if models_opt_on_gpu and self.is_training else '/CPU:0'
|
||||
optimizer_vars_on_cpu = models_opt_device=='/CPU:0'
|
||||
|
||||
input_ch = 3
|
||||
output_ch = 3
|
||||
input_ch=3
|
||||
bgr_shape = nn.get4Dshape(resolution,resolution,input_ch)
|
||||
mask_shape = nn.get4Dshape(resolution,resolution,1)
|
||||
self.model_filename_list = []
|
||||
|
||||
|
||||
with tf.device ('/CPU:0'):
|
||||
#Place holders on CPU
|
||||
self.warped_src = tf.placeholder (nn.tf_floatx, bgr_shape)
|
||||
self.warped_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
|
||||
self.warped_src = tf.placeholder (nn.floatx, bgr_shape)
|
||||
self.warped_dst = tf.placeholder (nn.floatx, bgr_shape)
|
||||
|
||||
self.target_src = tf.placeholder (nn.tf_floatx, bgr_shape)
|
||||
self.target_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
|
||||
self.target_src = tf.placeholder (nn.floatx, bgr_shape)
|
||||
self.target_dst = tf.placeholder (nn.floatx, bgr_shape)
|
||||
|
||||
self.target_srcm_all = tf.placeholder (nn.tf_floatx, mask_shape)
|
||||
self.target_dstm_all = tf.placeholder (nn.tf_floatx, mask_shape)
|
||||
self.target_srcm_all = tf.placeholder (nn.floatx, mask_shape)
|
||||
self.target_dstm_all = tf.placeholder (nn.floatx, mask_shape)
|
||||
|
||||
# Initializing model classes
|
||||
if archi == 'liaech':
|
||||
lowest_dense_res, Encoder, Inter, Decoder = nn.get_ae_models_chervonij(resolution)
|
||||
model_archi = nn.DeepFakeArchi(resolution, mod='chervonij')
|
||||
else:
|
||||
lowest_dense_res, Encoder, Inter, Decoder = nn.get_ae_models(resolution)
|
||||
|
||||
model_archi = nn.DeepFakeArchi(resolution)
|
||||
|
||||
with tf.device (models_opt_device):
|
||||
if 'df' in archi:
|
||||
self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, is_hd=is_hd, name='encoder')
|
||||
encoder_out_ch = self.encoder.compute_output_channels ( (nn.tf_floatx, bgr_shape))
|
||||
self.encoder = model_archi.Encoder(in_ch=input_ch, e_ch=e_dims, is_hd=is_hd, name='encoder')
|
||||
encoder_out_ch = self.encoder.compute_output_channels ( (nn.floatx, bgr_shape))
|
||||
|
||||
self.inter = Inter (in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims, is_hd=is_hd, name='inter')
|
||||
inter_out_ch = self.inter.compute_output_channels ( (nn.tf_floatx, (None,encoder_out_ch)))
|
||||
self.inter = model_archi.Inter (in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims, is_hd=is_hd, name='inter')
|
||||
inter_out_ch = self.inter.compute_output_channels ( (nn.floatx, (None,encoder_out_ch)))
|
||||
|
||||
self.decoder_src = Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd=is_hd, name='decoder_src')
|
||||
self.decoder_dst = Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd=is_hd, name='decoder_dst')
|
||||
self.decoder_src = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd=is_hd, name='decoder_src')
|
||||
self.decoder_dst = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd=is_hd, name='decoder_dst')
|
||||
|
||||
self.model_filename_list += [ [self.encoder, 'encoder.npy' ],
|
||||
[self.inter, 'inter.npy' ],
|
||||
|
|
@ -197,20 +192,20 @@ class SAEHDModel(ModelBase):
|
|||
|
||||
if self.is_training:
|
||||
if self.options['true_face_power'] != 0:
|
||||
self.code_discriminator = nn.CodeDiscriminator(ae_dims, code_res=lowest_dense_res*2, name='dis' )
|
||||
self.code_discriminator = nn.CodeDiscriminator(ae_dims, code_res=model_archi.Inter.get_code_res()*2, name='dis' )
|
||||
self.model_filename_list += [ [self.code_discriminator, 'code_discriminator.npy'] ]
|
||||
|
||||
elif 'liae' in archi:
|
||||
self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, is_hd=is_hd, name='encoder')
|
||||
encoder_out_ch = self.encoder.compute_output_channels ( (nn.tf_floatx, bgr_shape))
|
||||
self.encoder = model_archi.Encoder(in_ch=input_ch, e_ch=e_dims, is_hd=is_hd, name='encoder')
|
||||
encoder_out_ch = self.encoder.compute_output_channels ( (nn.floatx, bgr_shape))
|
||||
|
||||
self.inter_AB = Inter(in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims*2, is_hd=is_hd, name='inter_AB')
|
||||
self.inter_B = Inter(in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims*2, is_hd=is_hd, name='inter_B')
|
||||
self.inter_AB = model_archi.Inter(in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims*2, is_hd=is_hd, name='inter_AB')
|
||||
self.inter_B = model_archi.Inter(in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims*2, is_hd=is_hd, name='inter_B')
|
||||
|
||||
inter_AB_out_ch = self.inter_AB.compute_output_channels ( (nn.tf_floatx, (None,encoder_out_ch)))
|
||||
inter_B_out_ch = self.inter_B.compute_output_channels ( (nn.tf_floatx, (None,encoder_out_ch)))
|
||||
inter_AB_out_ch = self.inter_AB.compute_output_channels ( (nn.floatx, (None,encoder_out_ch)))
|
||||
inter_B_out_ch = self.inter_B.compute_output_channels ( (nn.floatx, (None,encoder_out_ch)))
|
||||
inters_out_ch = inter_AB_out_ch+inter_B_out_ch
|
||||
self.decoder = Decoder(in_ch=inters_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd=is_hd, name='decoder')
|
||||
self.decoder = model_archi.Decoder(in_ch=inters_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd=is_hd, name='decoder')
|
||||
|
||||
self.model_filename_list += [ [self.encoder, 'encoder.npy'],
|
||||
[self.inter_AB, 'inter_AB.npy'],
|
||||
|
|
@ -219,8 +214,8 @@ class SAEHDModel(ModelBase):
|
|||
|
||||
if self.is_training:
|
||||
if gan_power != 0:
|
||||
self.D_src = nn.PatchDiscriminator(patch_size=resolution//16, in_ch=output_ch, name="D_src")
|
||||
self.D_dst = nn.PatchDiscriminator(patch_size=resolution//16, in_ch=output_ch, name="D_dst")
|
||||
self.D_src = nn.PatchDiscriminator(patch_size=resolution//16, in_ch=input_ch, name="D_src")
|
||||
self.D_dst = nn.PatchDiscriminator(patch_size=resolution//16, in_ch=input_ch, name="D_dst")
|
||||
self.model_filename_list += [ [self.D_src, 'D_src.npy'] ]
|
||||
self.model_filename_list += [ [self.D_dst, 'D_dst.npy'] ]
|
||||
|
||||
|
|
@ -228,7 +223,7 @@ class SAEHDModel(ModelBase):
|
|||
lr=5e-5
|
||||
lr_dropout = 0.3 if self.options['lr_dropout'] and not self.pretrain else 1.0
|
||||
clipnorm = 1.0 if self.options['clipgrad'] else 0.0
|
||||
self.src_dst_opt = nn.TFRMSpropOptimizer(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt')
|
||||
self.src_dst_opt = nn.RMSprop(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt')
|
||||
self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]
|
||||
if 'df' in archi:
|
||||
self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()
|
||||
|
|
@ -238,12 +233,12 @@ class SAEHDModel(ModelBase):
|
|||
self.src_dst_opt.initialize_variables (self.src_dst_trainable_weights, vars_on_cpu=optimizer_vars_on_cpu)
|
||||
|
||||
if self.options['true_face_power'] != 0:
|
||||
self.D_code_opt = nn.TFRMSpropOptimizer(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_code_opt')
|
||||
self.D_code_opt = nn.RMSprop(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_code_opt')
|
||||
self.D_code_opt.initialize_variables ( self.code_discriminator.get_weights(), vars_on_cpu=optimizer_vars_on_cpu)
|
||||
self.model_filename_list += [ (self.D_code_opt, 'D_code_opt.npy') ]
|
||||
|
||||
if gan_power != 0:
|
||||
self.D_src_dst_opt = nn.TFRMSpropOptimizer(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_src_dst_opt')
|
||||
self.D_src_dst_opt = nn.RMSprop(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_src_dst_opt')
|
||||
self.D_src_dst_opt.initialize_variables ( self.D_src.get_weights()+self.D_dst.get_weights(), vars_on_cpu=optimizer_vars_on_cpu)
|
||||
self.model_filename_list += [ (self.D_src_dst_opt, 'D_src_dst_opt.npy') ]
|
||||
|
||||
|
|
@ -316,8 +311,8 @@ class SAEHDModel(ModelBase):
|
|||
gpu_target_srcm_eyes = tf.clip_by_value (gpu_target_srcm_all-1, 0, 1)
|
||||
gpu_target_dstm_eyes = tf.clip_by_value (gpu_target_dstm_all-1, 0, 1)
|
||||
|
||||
gpu_target_srcm_blur = nn.tf_gaussian_blur(gpu_target_srcm, max(1, resolution // 32) )
|
||||
gpu_target_dstm_blur = nn.tf_gaussian_blur(gpu_target_dstm, max(1, resolution // 32) )
|
||||
gpu_target_srcm_blur = nn.gaussian_blur(gpu_target_srcm, max(1, resolution // 32) )
|
||||
gpu_target_dstm_blur = nn.gaussian_blur(gpu_target_dstm, max(1, resolution // 32) )
|
||||
|
||||
gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur
|
||||
gpu_target_dst_anti_masked = gpu_target_dst*(1.0 - gpu_target_dstm_blur)
|
||||
|
|
@ -331,7 +326,7 @@ class SAEHDModel(ModelBase):
|
|||
gpu_psd_target_dst_masked = gpu_pred_src_dst*gpu_target_dstm_blur
|
||||
gpu_psd_target_dst_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_blur)
|
||||
|
||||
gpu_src_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
|
||||
gpu_src_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
|
||||
gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
|
||||
|
||||
if eyes_prio:
|
||||
|
|
@ -341,14 +336,14 @@ class SAEHDModel(ModelBase):
|
|||
|
||||
face_style_power = self.options['face_style_power'] / 100.0
|
||||
if face_style_power != 0 and not self.pretrain:
|
||||
gpu_src_loss += nn.tf_style_loss(gpu_psd_target_dst_masked, gpu_target_dst_masked, gaussian_blur_radius=resolution//16, loss_weight=10000*face_style_power)
|
||||
gpu_src_loss += nn.style_loss(gpu_psd_target_dst_masked, gpu_target_dst_masked, gaussian_blur_radius=resolution//16, loss_weight=10000*face_style_power)
|
||||
|
||||
bg_style_power = self.options['bg_style_power'] / 100.0
|
||||
if bg_style_power != 0 and not self.pretrain:
|
||||
gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*nn.tf_dssim(gpu_psd_target_dst_anti_masked, gpu_target_dst_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
|
||||
gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*nn.dssim(gpu_psd_target_dst_anti_masked, gpu_target_dst_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
|
||||
gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*tf.square( gpu_psd_target_dst_anti_masked - gpu_target_dst_anti_masked), axis=[1,2,3] )
|
||||
|
||||
gpu_dst_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
|
||||
gpu_dst_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
|
||||
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
|
||||
|
||||
if eyes_prio:
|
||||
|
|
@ -376,7 +371,7 @@ class SAEHDModel(ModelBase):
|
|||
gpu_D_code_loss = (DLoss(gpu_src_code_d_ones , gpu_dst_code_d) + \
|
||||
DLoss(gpu_src_code_d_zeros, gpu_src_code_d) ) * 0.5
|
||||
|
||||
gpu_D_code_loss_gvs += [ nn.tf_gradients (gpu_D_code_loss, self.code_discriminator.get_weights() ) ]
|
||||
gpu_D_code_loss_gvs += [ nn.gradients (gpu_D_code_loss, self.code_discriminator.get_weights() ) ]
|
||||
|
||||
if gan_power != 0:
|
||||
gpu_pred_src_src_d = self.D_src(gpu_pred_src_src_masked_opt)
|
||||
|
|
@ -395,32 +390,32 @@ class SAEHDModel(ModelBase):
|
|||
(DLoss(gpu_target_dst_d_ones , gpu_target_dst_d) + \
|
||||
DLoss(gpu_pred_dst_dst_d_zeros, gpu_pred_dst_dst_d) ) * 0.5
|
||||
|
||||
gpu_D_src_dst_loss_gvs += [ nn.tf_gradients (gpu_D_src_dst_loss, self.D_src.get_weights()+self.D_dst.get_weights() ) ]
|
||||
gpu_D_src_dst_loss_gvs += [ nn.gradients (gpu_D_src_dst_loss, self.D_src.get_weights()+self.D_dst.get_weights() ) ]
|
||||
|
||||
gpu_G_loss += gan_power*(DLoss(gpu_pred_src_src_d_ones, gpu_pred_src_src_d) + DLoss(gpu_pred_dst_dst_d_ones, gpu_pred_dst_dst_d))
|
||||
|
||||
|
||||
gpu_G_loss_gvs += [ nn.tf_gradients ( gpu_G_loss, self.src_dst_trainable_weights ) ]
|
||||
gpu_G_loss_gvs += [ nn.gradients ( gpu_G_loss, self.src_dst_trainable_weights ) ]
|
||||
|
||||
|
||||
# Average losses and gradients, and create optimizer update ops
|
||||
with tf.device (models_opt_device):
|
||||
pred_src_src = nn.tf_concat(gpu_pred_src_src_list, 0)
|
||||
pred_dst_dst = nn.tf_concat(gpu_pred_dst_dst_list, 0)
|
||||
pred_src_dst = nn.tf_concat(gpu_pred_src_dst_list, 0)
|
||||
pred_src_srcm = nn.tf_concat(gpu_pred_src_srcm_list, 0)
|
||||
pred_dst_dstm = nn.tf_concat(gpu_pred_dst_dstm_list, 0)
|
||||
pred_src_dstm = nn.tf_concat(gpu_pred_src_dstm_list, 0)
|
||||
pred_src_src = nn.concat(gpu_pred_src_src_list, 0)
|
||||
pred_dst_dst = nn.concat(gpu_pred_dst_dst_list, 0)
|
||||
pred_src_dst = nn.concat(gpu_pred_src_dst_list, 0)
|
||||
pred_src_srcm = nn.concat(gpu_pred_src_srcm_list, 0)
|
||||
pred_dst_dstm = nn.concat(gpu_pred_dst_dstm_list, 0)
|
||||
pred_src_dstm = nn.concat(gpu_pred_src_dstm_list, 0)
|
||||
|
||||
src_loss = tf.concat(gpu_src_losses, 0)
|
||||
dst_loss = tf.concat(gpu_dst_losses, 0)
|
||||
src_dst_loss_gv_op = self.src_dst_opt.get_update_op (nn.tf_average_gv_list (gpu_G_loss_gvs))
|
||||
src_dst_loss_gv_op = self.src_dst_opt.get_update_op (nn.average_gv_list (gpu_G_loss_gvs))
|
||||
|
||||
if self.options['true_face_power'] != 0:
|
||||
D_loss_gv_op = self.D_code_opt.get_update_op (nn.tf_average_gv_list(gpu_D_code_loss_gvs))
|
||||
D_loss_gv_op = self.D_code_opt.get_update_op (nn.average_gv_list(gpu_D_code_loss_gvs))
|
||||
|
||||
if gan_power != 0:
|
||||
src_D_src_dst_loss_gv_op = self.D_src_dst_opt.get_update_op (nn.tf_average_gv_list(gpu_D_src_dst_loss_gvs) )
|
||||
src_D_src_dst_loss_gv_op = self.D_src_dst_opt.get_update_op (nn.average_gv_list(gpu_D_src_dst_loss_gvs) )
|
||||
|
||||
|
||||
# Initializing training and view functions
|
||||
|
|
|
|||
|
|
@ -7,30 +7,29 @@ import numpy as np
|
|||
from core import mathlib
|
||||
from core.interact import interact as io
|
||||
from core.leras import nn
|
||||
from facelib import FaceType, TernausNet
|
||||
from facelib import FaceType, TernausNet, DFLSegNet
|
||||
from models import ModelBase
|
||||
from samplelib import *
|
||||
|
||||
class XSegModel(ModelBase):
|
||||
class SkinSegModel(ModelBase):
|
||||
|
||||
#override
|
||||
def on_initialize_options(self):
|
||||
device_config = nn.getCurrentDeviceConfig()
|
||||
yn_str = {True:'y',False:'n'}
|
||||
|
||||
#default_resolution = 256
|
||||
|
||||
ask_override = self.ask_override()
|
||||
if self.is_first_run() or ask_override:
|
||||
self.ask_autobackup_hour()
|
||||
self.ask_write_preview_history()
|
||||
self.ask_target_iter()
|
||||
self.ask_batch_size(24)
|
||||
|
||||
#if self.is_first_run():
|
||||
#resolution = io.input_int("Resolution", default_resolution, add_info="64-512")
|
||||
#resolution = np.clip ( (resolution // 16) * 16, 64, 512)
|
||||
#self.options['resolution'] = resolution
|
||||
|
||||
self.ask_batch_size(8)
|
||||
|
||||
default_lr_dropout = self.options['lr_dropout'] = self.load_or_def_option('lr_dropout', False)
|
||||
|
||||
if self.is_first_run() or ask_override:
|
||||
self.options['lr_dropout'] = io.input_bool ("Use learning rate dropout", default_lr_dropout, help_message="When the face is trained enough, you can enable this option to get extra sharpness and reduce subpixel shake for less amount of iterations.")
|
||||
|
||||
#override
|
||||
def on_initialize(self):
|
||||
device_config = nn.getCurrentDeviceConfig()
|
||||
|
|
@ -41,22 +40,24 @@ class XSegModel(ModelBase):
|
|||
device_config = nn.getCurrentDeviceConfig()
|
||||
devices = device_config.devices
|
||||
|
||||
self.resolution = resolution = 256#self.options['resolution']
|
||||
|
||||
place_model_on_cpu = True#len(devices) == 0
|
||||
self.resolution = resolution = 256
|
||||
self.face_type = FaceType.WHOLE_FACE
|
||||
|
||||
place_model_on_cpu = True #len(devices) == 0
|
||||
models_opt_device = '/CPU:0' if place_model_on_cpu else '/GPU:0'
|
||||
|
||||
bgr_shape = nn.get4Dshape(resolution,resolution,3)
|
||||
mask_shape = nn.get4Dshape(resolution,resolution,1)
|
||||
|
||||
# Initializing model classes
|
||||
self.model = TernausNet(f'{self.model_name}_SkinSeg',
|
||||
resolution,
|
||||
load_weights=not self.is_first_run(),
|
||||
weights_file_root=self.get_model_root_path(),
|
||||
training=True,
|
||||
place_model_on_cpu=place_model_on_cpu,
|
||||
data_format=nn.data_format)
|
||||
self.model = DFLSegNet(name=f'{self.model_name}_SkinSeg',
|
||||
resolution=resolution,
|
||||
load_weights=not self.is_first_run(),
|
||||
weights_file_root=self.get_model_root_path(),
|
||||
training=True,
|
||||
place_model_on_cpu=place_model_on_cpu,
|
||||
optimizer=nn.RMSprop(lr=0.0001, lr_dropout=0.3 if self.options['lr_dropout'] else 1.0, name='opt'),
|
||||
data_format=nn.data_format)
|
||||
|
||||
if self.is_training:
|
||||
# Adjust batch size for multiple GPU
|
||||
|
|
@ -81,21 +82,21 @@ class XSegModel(ModelBase):
|
|||
gpu_target_t = self.model.target_t [batch_slice,:,:,:]
|
||||
|
||||
# process model tensors
|
||||
gpu_pred_logits_t, gpu_pred_t = self.model.net([gpu_input_t])
|
||||
gpu_pred_logits_t, gpu_pred_t = self.model.flow(gpu_input_t)
|
||||
gpu_pred_list.append(gpu_pred_t)
|
||||
|
||||
gpu_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=gpu_target_t, logits=gpu_pred_logits_t), axis=[1,2,3])
|
||||
gpu_losses += [gpu_loss]
|
||||
|
||||
gpu_loss_gvs += [ nn.tf_gradients ( gpu_loss, self.model.net_weights ) ]
|
||||
gpu_loss_gvs += [ nn.gradients ( gpu_loss, self.model.get_weights() ) ]
|
||||
|
||||
|
||||
# Average losses and gradients, and create optimizer update ops
|
||||
with tf.device (models_opt_device):
|
||||
pred = nn.tf_concat(gpu_pred_list, 0)
|
||||
pred = nn.concat(gpu_pred_list, 0)
|
||||
loss = tf.reduce_mean(gpu_losses)
|
||||
|
||||
loss_gv_op = self.model.opt.get_update_op (nn.tf_average_gv_list (gpu_loss_gvs))
|
||||
loss_gv_op = self.model.opt.get_update_op (nn.average_gv_list (gpu_loss_gvs))
|
||||
|
||||
|
||||
# Initializing training and view functions
|
||||
|
|
@ -114,14 +115,30 @@ class XSegModel(ModelBase):
|
|||
dst_generators_count = cpu_count // 2
|
||||
src_generators_count = int(src_generators_count * 1.5)
|
||||
|
||||
|
||||
src_generator = SampleGeneratorFaceCelebAMaskHQ ( root_path=self.training_data_src_path, debug=self.is_debug(), batch_size=self.get_batch_size(), resolution=256, generators_count=src_generators_count, data_format = nn.data_format)
|
||||
|
||||
dst_generator = SampleGeneratorImage(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
|
||||
"""
|
||||
src_generator = SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True),
|
||||
output_sample_types = [ {'sample_type': SampleProcessor.SampleType.IMAGE, 'warp':False, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'data_format':nn.data_format, 'resolution': resolution} ],
|
||||
output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE, 'warp':True, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR_RANDOM_HSV_SHIFT, 'border_replicate':False, 'face_type':self.face_type, 'motion_blur':(25, 5), 'gaussian_blur':(25,5), 'random_bilinear_resize':(25,75), 'data_format':nn.data_format, 'resolution': resolution},
|
||||
{'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':True, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.NONE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
|
||||
],
|
||||
generators_count=src_generators_count )
|
||||
"""
|
||||
src_generator = SampleGeneratorFaceSkinSegDataset(self.training_data_src_path,
|
||||
debug=self.is_debug(),
|
||||
batch_size=self.get_batch_size(),
|
||||
resolution=resolution,
|
||||
face_type=self.face_type,
|
||||
generators_count=src_generators_count,
|
||||
data_format=nn.data_format)
|
||||
|
||||
dst_generator = SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True),
|
||||
output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE, 'warp':False, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'border_replicate':False, 'face_type':self.face_type, 'motion_blur':(25, 5), 'gaussian_blur':(25,5), 'random_bilinear_resize':(25,75), 'data_format':nn.data_format, 'resolution': resolution},
|
||||
],
|
||||
generators_count=dst_generators_count,
|
||||
raise_on_no_data=False )
|
||||
|
||||
|
||||
if not dst_generator.is_initialized():
|
||||
io.log_info(f"\nTo view the model on unseen faces, place any image faces in {self.training_data_dst_path}.\n")
|
||||
|
||||
|
|
@ -157,9 +174,9 @@ class XSegModel(ModelBase):
|
|||
result = []
|
||||
st = []
|
||||
for i in range(n_samples):
|
||||
ar = I[i]*M[i]+ green_bg*(1-M[i]), IM[i], I[i]*IM[i] + green_bg*(1-IM[i])
|
||||
ar = I[i]*M[i]+0.5*I[i]*(1-M[i])+0.5*green_bg*(1-M[i]), IM[i], I[i]*IM[i] + green_bg*(1-IM[i])
|
||||
st.append ( np.concatenate ( ar, axis=1) )
|
||||
result += [ ('XSeg training faces', np.concatenate (st, axis=0 )), ]
|
||||
result += [ ('SkinSeg training faces', np.concatenate (st, axis=0 )), ]
|
||||
|
||||
if len(dst_samples) != 0:
|
||||
dst_np, = dst_samples
|
||||
|
|
@ -173,8 +190,8 @@ class XSegModel(ModelBase):
|
|||
ar = D[i], DM[i], D[i]*DM[i]+ green_bg*(1-DM[i])
|
||||
st.append ( np.concatenate ( ar, axis=1) )
|
||||
|
||||
result += [ ('XSeg unseen faces', np.concatenate (st, axis=0 )), ]
|
||||
result += [ ('SkinSeg unseen faces', np.concatenate (st, axis=0 )), ]
|
||||
|
||||
return result
|
||||
|
||||
Model = XSegModel
|
||||
Model = SkinSegModel
|
||||
Loading…
Add table
Add a link
Reference in a new issue