mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-19 21:13:20 -07:00
SAE : WARNING, RETRAIN IS REQUIRED !
fixed model sizes from previous update. avoided bug in ML framework(keras) that forces to train the model on random noise. Converter: added blur on the same keys as sharpness Added new model 'TrueFace'. This is a GAN model ported from https://github.com/NVlabs/FUNIT Model produces near zero morphing and high detail face. Model has higher failure rate than other models. Keep src and dst faceset in same lighting conditions.
This commit is contained in:
Colombo
parent
201b762541
commit
dc11ec32be
26 changed files with 1308 additions and 250 deletions
|
|
@ -23,7 +23,7 @@ You can implement your own model. Check examples.
|
|||
class ModelBase(object):
|
||||
|
||||
|
||||
def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, pretraining_data_path=None, debug = False, device_args = None,
|
||||
def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, pretraining_data_path=None, is_training=False, debug = False, device_args = None,
|
||||
ask_enable_autobackup=True,
|
||||
ask_write_preview_history=True,
|
||||
ask_target_iter=True,
|
||||
|
|
@ -56,14 +56,8 @@ class ModelBase(object):
|
|||
self.training_data_dst_path = training_data_dst_path
|
||||
self.pretraining_data_path = pretraining_data_path
|
||||
|
||||
self.src_images_paths = None
|
||||
self.dst_images_paths = None
|
||||
self.src_yaw_images_paths = None
|
||||
self.dst_yaw_images_paths = None
|
||||
self.src_data_generator = None
|
||||
self.dst_data_generator = None
|
||||
self.debug = debug
|
||||
self.is_training_mode = (training_data_src_path is not None and training_data_dst_path is not None)
|
||||
self.is_training_mode = is_training
|
||||
|
||||
self.iter = 0
|
||||
self.options = {}
|
||||
|
|
@ -412,40 +406,60 @@ class ModelBase(object):
|
|||
cv2_imwrite (filepath, img )
|
||||
|
||||
def load_weights_safe(self, model_filename_list, optimizer_filename_list=[]):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
loaded = []
|
||||
not_loaded = []
|
||||
for mf in model_filename_list:
|
||||
model, filename = mf
|
||||
filename = self.get_strpath_storage_for_file(filename)
|
||||
|
||||
if Path(filename).exists():
|
||||
loaded += [ mf ]
|
||||
model.load_weights(filename)
|
||||
|
||||
if issubclass(model.__class__, keras.optimizers.Optimizer):
|
||||
opt = model
|
||||
|
||||
try:
|
||||
with open(filename, "rb") as f:
|
||||
fd = pickle.loads(f.read())
|
||||
|
||||
weights = fd.get('weights', None)
|
||||
if weights is not None:
|
||||
opt.set_weights(weights)
|
||||
|
||||
except Exception as e:
|
||||
print ("Unable to load ", filename)
|
||||
|
||||
else:
|
||||
model.load_weights(filename)
|
||||
else:
|
||||
not_loaded += [ mf ]
|
||||
|
||||
if len(optimizer_filename_list) != 0:
|
||||
opt_filename = self.get_strpath_storage_for_file('opt.h5')
|
||||
if Path(opt_filename).exists():
|
||||
try:
|
||||
with open(opt_filename, "rb") as f:
|
||||
d = pickle.loads(f.read())
|
||||
|
||||
for x in optimizer_filename_list:
|
||||
opt, filename = x
|
||||
if filename in d:
|
||||
weights = d[filename].get('weights', None)
|
||||
if weights:
|
||||
opt.set_weights(weights)
|
||||
print("set ok")
|
||||
except Exception as e:
|
||||
print ("Unable to load ", opt_filename)
|
||||
|
||||
|
||||
return loaded, not_loaded
|
||||
|
||||
def save_weights_safe(self, model_filename_list):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
for model, filename in model_filename_list:
|
||||
filename = self.get_strpath_storage_for_file(filename)
|
||||
model.save_weights( filename + '.tmp' )
|
||||
filename = self.get_strpath_storage_for_file(filename) + '.tmp'
|
||||
|
||||
if issubclass(model.__class__, keras.optimizers.Optimizer):
|
||||
opt = model
|
||||
|
||||
try:
|
||||
fd = {}
|
||||
symbolic_weights = getattr(opt, 'weights')
|
||||
if symbolic_weights:
|
||||
fd['weights'] = self.K.batch_get_value(symbolic_weights)
|
||||
|
||||
with open(filename, 'wb') as f:
|
||||
f.write( pickle.dumps(fd) )
|
||||
except Exception as e:
|
||||
print ("Unable to save ", filename)
|
||||
else:
|
||||
model.save_weights( filename)
|
||||
|
||||
rename_list = model_filename_list
|
||||
|
||||
|
|
|
|||
166
models/Model_DEV_FUNIT/Model.py
Normal file
166
models/Model_DEV_FUNIT/Model.py
Normal file
|
|
@ -0,0 +1,166 @@
|
|||
from functools import partial
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from facelib import FaceType
|
||||
from interact import interact as io
|
||||
from mathlib import get_power_of_two
|
||||
from models import ModelBase
|
||||
from nnlib import nnlib, FUNIT
|
||||
from samplelib import *
|
||||
|
||||
|
||||
|
||||
class FUNITModel(ModelBase):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs,
|
||||
ask_sort_by_yaw=False,
|
||||
ask_random_flip=False,
|
||||
ask_src_scale_mod=False)
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
default_face_type = 'f'
|
||||
if is_first_run:
|
||||
self.options['resolution'] = io.input_int("Resolution ( 128,224 ?:help skip:128) : ", 128, [128,224])
|
||||
else:
|
||||
self.options['resolution'] = self.options.get('resolution', 128)
|
||||
|
||||
if is_first_run:
|
||||
self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="").lower()
|
||||
else:
|
||||
self.options['face_type'] = self.options.get('face_type', default_face_type)
|
||||
|
||||
#override
|
||||
def onInitialize(self, batch_size=-1, **in_options):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
self.set_vram_batch_requirements({4:16})
|
||||
|
||||
resolution = self.options['resolution']
|
||||
face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
|
||||
person_id_max_count = SampleGeneratorFace.get_person_id_max_count(self.training_data_src_path)
|
||||
|
||||
|
||||
self.model = FUNIT( face_type_str=FaceType.toString(face_type),
|
||||
batch_size=self.batch_size,
|
||||
encoder_nf=64,
|
||||
encoder_downs=2,
|
||||
encoder_res_blk=2,
|
||||
class_downs=4,
|
||||
class_nf=64,
|
||||
class_latent=64,
|
||||
mlp_nf=256,
|
||||
mlp_blks=2,
|
||||
dis_nf=64,
|
||||
dis_res_blks=10,
|
||||
num_classes=person_id_max_count,
|
||||
subpixel_decoder=True,
|
||||
initialize_weights=self.is_first_run(),
|
||||
is_training=self.is_training_mode
|
||||
)
|
||||
|
||||
if not self.is_first_run():
|
||||
self.load_weights_safe(self.model.get_model_filename_list())
|
||||
|
||||
if self.is_training_mode:
|
||||
t = SampleProcessor.Types
|
||||
face_type = t.FACE_TYPE_FULL if self.options['face_type'] == 'f' else t.FACE_TYPE_HALF
|
||||
|
||||
output_sample_types=[ {'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':128, 'normalize_tanh':True} ]
|
||||
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True),
|
||||
output_sample_types=output_sample_types, person_id_mode=True ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True),
|
||||
output_sample_types=output_sample_types, person_id_mode=True ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True),
|
||||
output_sample_types=output_sample_types, person_id_mode=True ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True),
|
||||
output_sample_types=output_sample_types, person_id_mode=True ),
|
||||
])
|
||||
|
||||
#override
|
||||
def get_model_filename_list(self):
|
||||
return self.model.get_model_filename_list()
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe(self.model.get_model_filename_list())
|
||||
|
||||
#override
|
||||
def onTrainOneIter(self, generators_samples, generators_list):
|
||||
xa,la = generators_samples[0]
|
||||
xb,lb = generators_samples[1]
|
||||
|
||||
G_loss, D_loss = self.model.train(xa,la,xb,lb)
|
||||
|
||||
return ( ('G_loss', G_loss), ('D_loss', D_loss), )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, generators_samples):
|
||||
xa = generators_samples[0][0]
|
||||
xb = generators_samples[1][0]
|
||||
ta = generators_samples[2][0]
|
||||
tb = generators_samples[3][0]
|
||||
|
||||
view_samples = min(4, xa.shape[0])
|
||||
|
||||
lines_train = []
|
||||
lines_test = []
|
||||
|
||||
for i in range(view_samples):
|
||||
|
||||
s_xa = self.model.get_average_class_code([ xa[i:i+1] ])[0][None,...]
|
||||
s_xb = self.model.get_average_class_code([ xb[i:i+1] ])[0][None,...]
|
||||
|
||||
s_ta = self.model.get_average_class_code([ ta[i:i+1] ])[0][None,...]
|
||||
s_tb = self.model.get_average_class_code([ tb[i:i+1] ])[0][None,...]
|
||||
|
||||
xaxa = self.model.convert ([ xa[i:i+1], s_xa ] )[0][0]
|
||||
xbxb = self.model.convert ([ xb[i:i+1], s_xb ] )[0][0]
|
||||
xaxb = self.model.convert ([ xa[i:i+1], s_xb ] )[0][0]
|
||||
xbxa = self.model.convert ([ xb[i:i+1], s_xa ] )[0][0]
|
||||
|
||||
tata = self.model.convert ([ ta[i:i+1], s_ta ] )[0][0]
|
||||
tbtb = self.model.convert ([ tb[i:i+1], s_tb ] )[0][0]
|
||||
tatb = self.model.convert ([ ta[i:i+1], s_tb ] )[0][0]
|
||||
tbta = self.model.convert ([ tb[i:i+1], s_ta ] )[0][0]
|
||||
|
||||
line_train = [ xa[i], xaxa, xb[i], xbxb, xaxb, xbxa ]
|
||||
line_test = [ ta[i], tata, tb[i], tbtb, tatb, tbta ]
|
||||
|
||||
lines_train += [ np.concatenate([ np.clip(x/2+0.5,0,1) for x in line_train], axis=1) ]
|
||||
lines_test += [ np.concatenate([ np.clip(x/2+0.5,0,1) for x in line_test ], axis=1) ]
|
||||
|
||||
lines_train = np.concatenate ( lines_train, axis=0 )
|
||||
lines_test = np.concatenate ( lines_test, axis=0 )
|
||||
return [ ('TRAIN', lines_train ), ('TEST', lines_test) ]
|
||||
|
||||
def predictor_func (self, face=None, dummy_predict=False):
|
||||
if dummy_predict:
|
||||
self.model.convert ([ np.zeros ( (1, self.options['resolution'], self.options['resolution'], 3), dtype=np.float32 ), self.average_class_code ])
|
||||
else:
|
||||
bgr, = self.model.convert ([ face[np.newaxis,...]*2-1, self.average_class_code ])
|
||||
return bgr[0] / 2 + 0.5
|
||||
|
||||
#override
|
||||
def get_ConverterConfig(self):
|
||||
face_type = FaceType.FULL
|
||||
|
||||
import converters
|
||||
return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), converters.ConverterConfigMasked(face_type=face_type,
|
||||
default_mode = 1,
|
||||
clip_hborder_mask_per=0.0625 if (face_type == FaceType.FULL) else 0,
|
||||
)
|
||||
|
||||
|
||||
Model = FUNITModel
|
||||
1
models/Model_DEV_FUNIT/__init__.py
Normal file
1
models/Model_DEV_FUNIT/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .Model import Model
|
||||
|
|
@ -51,7 +51,7 @@ class SAEModel(ModelBase):
|
|||
default_e_ch_dims = 42
|
||||
default_d_ch_dims = default_e_ch_dims // 2
|
||||
def_ca_weights = False
|
||||
|
||||
|
||||
if is_first_run:
|
||||
self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dims (32-1024 ?:help skip:%d) : " % (default_ae_dims) , default_ae_dims, help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 )
|
||||
self.options['e_ch_dims'] = np.clip ( io.input_int("Encoder dims per channel (21-85 ?:help skip:%d) : " % (default_e_ch_dims) , default_e_ch_dims, help_message="More encoder dims help to recognize more facial features, but require more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 )
|
||||
|
|
@ -133,15 +133,15 @@ class SAEModel(ModelBase):
|
|||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same')(x)))
|
||||
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='valid')(ZeroPadding2D(1)(x))))
|
||||
return func
|
||||
|
||||
def enc_flow(e_dims, ae_dims, lowest_dense_res):
|
||||
def func(x):
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims, kernel_size=5, strides=2, padding='same')(x))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims*2, kernel_size=5, strides=2, padding='same')(x))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims*4, kernel_size=5, strides=2, padding='same')(x))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims*8, kernel_size=5, strides=2, padding='same')(x))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims*2, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims*4, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims*8, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
|
||||
|
||||
x = Dense(ae_dims)(Flatten()(x))
|
||||
x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x)
|
||||
|
|
@ -151,37 +151,37 @@ class SAEModel(ModelBase):
|
|||
return func
|
||||
|
||||
def dec_flow(output_nc, d_ch_dims, add_residual_blocks=True):
|
||||
dims = output_nc * d_ch_dims
|
||||
def ResidualBlock(dim):
|
||||
def func(inp):
|
||||
x = Conv2D(dim, kernel_size=3, padding='same')(inp)
|
||||
x = Conv2D(dim, kernel_size=3, padding='valid')(ZeroPadding2D(1)(inp))
|
||||
x = LeakyReLU(0.2)(x)
|
||||
x = Conv2D(dim, kernel_size=3, padding='same')(x)
|
||||
x = Conv2D(dim, kernel_size=3, padding='valid')(ZeroPadding2D(1)(x))
|
||||
x = Add()([x, inp])
|
||||
x = LeakyReLU(0.2)(x)
|
||||
return x
|
||||
return func
|
||||
|
||||
def func(x):
|
||||
dims = output_nc * d_ch_dims
|
||||
x = upscale(dims*8)(x)
|
||||
|
||||
|
||||
if add_residual_blocks:
|
||||
x = ResidualBlock(dims*8)(x)
|
||||
x = ResidualBlock(dims*8)(x)
|
||||
|
||||
x = upscale(dims*4)(x)
|
||||
|
||||
|
||||
if add_residual_blocks:
|
||||
x = ResidualBlock(dims*4)(x)
|
||||
x = ResidualBlock(dims*4)(x)
|
||||
|
||||
x = upscale(dims*2)(x)
|
||||
|
||||
|
||||
if add_residual_blocks:
|
||||
x = ResidualBlock(dims*2)(x)
|
||||
x = ResidualBlock(dims*2)(x)
|
||||
|
||||
return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return Conv2D(output_nc, kernel_size=5, padding='valid', activation='sigmoid')(ZeroPadding2D(2)(x))
|
||||
return func
|
||||
|
||||
self.encoder = modelify(enc_flow(e_dims, ae_dims, lowest_dense_res)) ( Input(bgr_shape) )
|
||||
|
|
@ -232,20 +232,20 @@ class SAEModel(ModelBase):
|
|||
mask_shape = (resolution, resolution, 1)
|
||||
|
||||
e_dims = output_nc*e_ch_dims
|
||||
d_dims = output_nc*d_ch_dims
|
||||
|
||||
lowest_dense_res = resolution // 16
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same')(x)))
|
||||
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='valid')(ZeroPadding2D(1)(x))))
|
||||
return func
|
||||
|
||||
def enc_flow(e_dims):
|
||||
def func(x):
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims, kernel_size=5, strides=2, padding='same')(x))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims*2, kernel_size=5, strides=2, padding='same')(x))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims*4, kernel_size=5, strides=2, padding='same')(x))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims*8, kernel_size=5, strides=2, padding='same')(x))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims*2, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims*4, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
|
||||
x = LeakyReLU(0.1)(Conv2D(e_dims*8, kernel_size=5, strides=2, padding='valid')(ZeroPadding2D(2)(x)))
|
||||
x = Flatten()(x)
|
||||
return x
|
||||
return func
|
||||
|
|
@ -259,12 +259,13 @@ class SAEModel(ModelBase):
|
|||
return x
|
||||
return func
|
||||
|
||||
def dec_flow(output_nc, d_dims):
|
||||
def dec_flow(output_nc, d_ch_dims, add_residual_blocks=True):
|
||||
d_dims = output_nc*d_ch_dims
|
||||
def ResidualBlock(dim):
|
||||
def func(inp):
|
||||
x = Conv2D(dim, kernel_size=3, padding='same')(inp)
|
||||
x = Conv2D(dim, kernel_size=3, padding='valid')(ZeroPadding2D(1)(inp))
|
||||
x = LeakyReLU(0.2)(x)
|
||||
x = Conv2D(dim, kernel_size=3, padding='same')(x)
|
||||
x = Conv2D(dim, kernel_size=3, padding='valid')(ZeroPadding2D(1)(inp))
|
||||
x = Add()([x, inp])
|
||||
x = LeakyReLU(0.2)(x)
|
||||
return x
|
||||
|
|
@ -272,18 +273,24 @@ class SAEModel(ModelBase):
|
|||
|
||||
def func(x):
|
||||
x = upscale(d_dims*8)(x)
|
||||
x = ResidualBlock(d_dims*8)(x)
|
||||
x = ResidualBlock(d_dims*8)(x)
|
||||
|
||||
if add_residual_blocks:
|
||||
x = ResidualBlock(d_dims*8)(x)
|
||||
x = ResidualBlock(d_dims*8)(x)
|
||||
|
||||
x = upscale(d_dims*4)(x)
|
||||
x = ResidualBlock(d_dims*4)(x)
|
||||
x = ResidualBlock(d_dims*4)(x)
|
||||
|
||||
if add_residual_blocks:
|
||||
x = ResidualBlock(d_dims*4)(x)
|
||||
x = ResidualBlock(d_dims*4)(x)
|
||||
|
||||
x = upscale(d_dims*2)(x)
|
||||
x = ResidualBlock(d_dims*2)(x)
|
||||
x = ResidualBlock(d_dims*2)(x)
|
||||
|
||||
return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
if add_residual_blocks:
|
||||
x = ResidualBlock(d_dims*2)(x)
|
||||
x = ResidualBlock(d_dims*2)(x)
|
||||
|
||||
return Conv2D(output_nc, kernel_size=5, padding='valid', activation='sigmoid')(ZeroPadding2D(2)(x))
|
||||
return func
|
||||
|
||||
self.encoder = modelify(enc_flow(e_dims)) ( Input(bgr_shape) )
|
||||
|
|
@ -293,10 +300,10 @@ class SAEModel(ModelBase):
|
|||
self.inter_AB = modelify(inter_flow(lowest_dense_res, ae_dims)) ( Input(sh) )
|
||||
|
||||
sh = np.array(K.int_shape( self.inter_B.outputs[0] )[1:])*(1,1,2)
|
||||
self.decoder = modelify(dec_flow(output_nc, d_dims)) ( Input(sh) )
|
||||
self.decoder = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) )
|
||||
|
||||
if learn_mask:
|
||||
self.decoderm = modelify(dec_flow(1, d_dims)) ( Input(sh) )
|
||||
self.decoderm = modelify(dec_flow(1, d_ch_dims, add_residual_blocks=False)) ( Input(sh) )
|
||||
|
||||
self.src_dst_trainable_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoder.trainable_weights
|
||||
|
||||
|
|
@ -349,17 +356,17 @@ class SAEModel(ModelBase):
|
|||
loaded, not_loaded = self.load_weights_safe(not_loaded)
|
||||
|
||||
CA_models = []
|
||||
if self.options.get('ca_weights', False):
|
||||
if self.options.get('ca_weights', False):
|
||||
CA_models += [ model for model, _ in not_loaded ]
|
||||
|
||||
|
||||
CA_conv_weights_list = []
|
||||
for model in CA_models:
|
||||
for layer in model.layers:
|
||||
if type(layer) == keras.layers.Conv2D:
|
||||
CA_conv_weights_list += [layer.weights[0]] #- is Conv2D kernel_weights
|
||||
|
||||
|
||||
if len(CA_conv_weights_list) != 0:
|
||||
CAInitializerMP ( CA_conv_weights_list )
|
||||
CAInitializerMP ( CA_conv_weights_list )
|
||||
|
||||
warped_src = self.model.warped_src
|
||||
target_src = Input ( (resolution, resolution, 3) )
|
||||
|
|
@ -501,7 +508,7 @@ class SAEModel(ModelBase):
|
|||
if self.options['learn_mask']:
|
||||
feed = [ warped_src, warped_dst, target_srcm, target_dstm ]
|
||||
src_mask_loss, dst_mask_loss, = self.src_dst_mask_train (feed)
|
||||
|
||||
|
||||
return ( ('src_loss', src_loss), ('dst_loss', dst_loss), )
|
||||
|
||||
#override
|
||||
|
|
|
|||
180
models/Model_TrueFace/Model.py
Normal file
180
models/Model_TrueFace/Model.py
Normal file
|
|
@ -0,0 +1,180 @@
|
|||
import numpy as np
|
||||
|
||||
from facelib import FaceType
|
||||
from interact import interact as io
|
||||
from models import ModelBase
|
||||
from nnlib import nnlib, FUNIT
|
||||
from samplelib import *
|
||||
|
||||
class TrueFaceModel(ModelBase):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs,
|
||||
ask_sort_by_yaw=False,
|
||||
ask_random_flip=False,
|
||||
ask_src_scale_mod=False)
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
default_resolution = 128
|
||||
default_face_type = 'f'
|
||||
|
||||
if is_first_run:
|
||||
resolution = self.options['resolution'] = io.input_int(f"Resolution ( 64-256 ?:help skip:{default_resolution}) : ", default_resolution, help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.")
|
||||
resolution = np.clip (resolution, 64, 256)
|
||||
while np.modf(resolution / 16)[0] != 0.0:
|
||||
resolution -= 1
|
||||
else:
|
||||
self.options['resolution'] = self.options.get('resolution', default_resolution)
|
||||
|
||||
if is_first_run:
|
||||
self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="").lower()
|
||||
else:
|
||||
self.options['face_type'] = self.options.get('face_type', default_face_type)
|
||||
|
||||
#override
|
||||
def onInitialize(self, batch_size=-1, **in_options):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
self.set_vram_batch_requirements({2:1,3:1,4:4,5:8,6:16})
|
||||
|
||||
resolution = self.options['resolution']
|
||||
face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
|
||||
|
||||
self.model = FUNIT( face_type_str=FaceType.toString(face_type),
|
||||
batch_size=self.batch_size,
|
||||
encoder_nf=64,
|
||||
encoder_downs=2,
|
||||
encoder_res_blk=2,
|
||||
class_downs=4,
|
||||
class_nf=64,
|
||||
class_latent=64,
|
||||
mlp_nf=256,
|
||||
mlp_blks=2,
|
||||
dis_nf=64,
|
||||
dis_res_blks=10,
|
||||
num_classes=2,
|
||||
subpixel_decoder=True,
|
||||
initialize_weights=self.is_first_run(),
|
||||
is_training=self.is_training_mode
|
||||
)
|
||||
|
||||
if not self.is_first_run():
|
||||
self.load_weights_safe(self.model.get_model_filename_list())
|
||||
|
||||
t = SampleProcessor.Types
|
||||
face_type = t.FACE_TYPE_FULL if self.options['face_type'] == 'f' else t.FACE_TYPE_HALF
|
||||
if self.is_training_mode:
|
||||
|
||||
output_sample_types=[ {'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh':True},
|
||||
]
|
||||
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True),
|
||||
output_sample_types=output_sample_types ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True),
|
||||
output_sample_types=output_sample_types )
|
||||
])
|
||||
else:
|
||||
generator = SampleGeneratorFace(self.training_data_src_path, batch_size=1,
|
||||
sample_process_options=SampleProcessor.Options(),
|
||||
output_sample_types=[ {'types': (t.IMG_SOURCE, face_type, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh':True} ] )
|
||||
|
||||
io.log_info("Calculating average src face style...")
|
||||
codes = []
|
||||
for i in io.progress_bar_generator(range(generator.get_total_sample_count())):
|
||||
codes += self.model.get_average_class_code( generator.generate_next() )
|
||||
|
||||
self.average_class_code = np.mean ( np.array(codes), axis=0 )[None,...]
|
||||
|
||||
|
||||
#override
|
||||
def get_model_filename_list(self):
|
||||
return self.model.get_model_filename_list()
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe(self.model.get_model_filename_list())
|
||||
|
||||
#override
|
||||
def onTrainOneIter(self, generators_samples, generators_list):
|
||||
bs = self.batch_size
|
||||
lbs = bs // 2
|
||||
hbs = bs - lbs
|
||||
|
||||
src, = generators_samples[0]
|
||||
dst, = generators_samples[1]
|
||||
|
||||
xa = np.concatenate ( [src[0:lbs], dst[0:lbs]], axis=0 )
|
||||
|
||||
la = np.concatenate ( [ np.array ([0]*lbs, np.int32),
|
||||
np.array ([1]*lbs, np.int32) ] )
|
||||
|
||||
xb = np.concatenate ( [src[lbs:], dst[lbs:]], axis=0 )
|
||||
|
||||
lb = np.concatenate ( [ np.array ([0]*hbs, np.int32),
|
||||
np.array ([1]*hbs, np.int32) ] )
|
||||
|
||||
rnd_list = np.arange(lbs*2)
|
||||
np.random.shuffle(rnd_list)
|
||||
xa = xa[rnd_list,...]
|
||||
la = la[rnd_list,...]
|
||||
la = la[...,None]
|
||||
|
||||
rnd_list = np.arange(hbs*2)
|
||||
np.random.shuffle(rnd_list)
|
||||
xb = xb[rnd_list,...]
|
||||
lb = lb[rnd_list,...]
|
||||
lb = lb[...,None]
|
||||
|
||||
G_loss, D_loss = self.model.train(xa,la,xb,lb)
|
||||
|
||||
return ( ('G_loss', G_loss), ('D_loss', D_loss), )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, generators_samples):
|
||||
xa = generators_samples[0][0]
|
||||
xb = generators_samples[1][0]
|
||||
|
||||
view_samples = min(4, xa.shape[0])
|
||||
|
||||
|
||||
s_xa_mean = self.model.get_average_class_code([xa])[0][None,...]
|
||||
s_xb_mean = self.model.get_average_class_code([xb])[0][None,...]
|
||||
|
||||
s_xab_mean = self.model.get_average_class_code([ np.concatenate( [xa,xb], axis=0) ])[0][None,...]
|
||||
|
||||
lines = []
|
||||
|
||||
for i in range(view_samples):
|
||||
xaxa, = self.model.convert ([ xa[i:i+1], s_xa_mean ] )
|
||||
xbxb, = self.model.convert ([ xb[i:i+1], s_xb_mean ] )
|
||||
xbxa, = self.model.convert ([ xb[i:i+1], s_xa_mean ] )
|
||||
|
||||
xa_i,xb_i,xaxa,xbxb,xbxa = [ np.clip(x/2+0.5, 0, 1) for x in [xa[i], xb[i], xaxa[0],xbxb[0],xbxa[0]] ]
|
||||
|
||||
lines += [ np.concatenate( (xa_i, xaxa, xb_i, xbxb, xbxa), axis=1) ]
|
||||
|
||||
r = np.concatenate ( lines, axis=0 )
|
||||
return [ ('TrueFace', r ) ]
|
||||
|
||||
def predictor_func (self, face=None, dummy_predict=False):
|
||||
if dummy_predict:
|
||||
self.model.convert ([ np.zeros ( (1, self.options['resolution'], self.options['resolution'], 3), dtype=np.float32 ), self.average_class_code ])
|
||||
else:
|
||||
bgr, = self.model.convert ([ face[np.newaxis,...]*2-1, self.average_class_code ])
|
||||
return bgr[0] / 2 + 0.5
|
||||
|
||||
#override
|
||||
def get_ConverterConfig(self):
|
||||
face_type = FaceType.FULL
|
||||
|
||||
import converters
|
||||
return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), converters.ConverterConfigMasked(face_type=face_type,
|
||||
default_mode = 1,
|
||||
clip_hborder_mask_per=0.0625 if (face_type == FaceType.FULL) else 0,
|
||||
)
|
||||
|
||||
Model = TrueFaceModel
|
||||
1
models/Model_TrueFace/__init__.py
Normal file
1
models/Model_TrueFace/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .Model import Model
|
||||
Loading…
Add table
Add a link
Reference in a new issue