mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-20 05:23:22 -07:00
DFL-2.0 initial branch commit
This commit is contained in:
Colombo
parent
52a67a61b3
commit
38b85108b3
154 changed files with 5251 additions and 9414 deletions
|
|
@ -1,162 +1,179 @@
|
|||
import colorsys
|
||||
import inspect
|
||||
import json
|
||||
import operator
|
||||
import os
|
||||
import pickle
|
||||
import shutil
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
import imagelib
|
||||
from interact import interact as io
|
||||
from nnlib import nnlib
|
||||
from core import imagelib
|
||||
from core.interact import interact as io
|
||||
from core.leras import nn
|
||||
from samplelib import SampleGeneratorBase
|
||||
from utils import Path_utils, std_utils
|
||||
from utils.cv2_utils import *
|
||||
from core import pathex
|
||||
from core.cv2ex import *
|
||||
|
||||
|
||||
'''
|
||||
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, is_training=False, debug = False, no_preview=False, device_args = None,
|
||||
ask_enable_autobackup=True,
|
||||
ask_write_preview_history=True,
|
||||
ask_target_iter=True,
|
||||
ask_batch_size=True,
|
||||
ask_random_flip=True, **kwargs):
|
||||
|
||||
device_args['force_gpu_idx'] = device_args.get('force_gpu_idx',-1)
|
||||
device_args['cpu_only'] = True if debug else device_args.get('cpu_only',False)
|
||||
|
||||
if device_args['force_gpu_idx'] == -1 and not device_args['cpu_only']:
|
||||
idxs_names_list = nnlib.device.getValidDevicesIdxsWithNamesList()
|
||||
if len(idxs_names_list) > 1:
|
||||
io.log_info ("You have multi GPUs in a system: ")
|
||||
for idx, name in idxs_names_list:
|
||||
io.log_info ("[%d] : %s" % (idx, name) )
|
||||
|
||||
device_args['force_gpu_idx'] = io.input_int("Which GPU idx to choose? ( skip: best GPU ) : ", -1, [ x[0] for x in idxs_names_list] )
|
||||
self.device_args = device_args
|
||||
|
||||
self.device_config = nnlib.DeviceConfig(allow_growth=True, **self.device_args)
|
||||
|
||||
io.log_info ("Loading model...")
|
||||
|
||||
self.model_path = model_path
|
||||
self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') )
|
||||
|
||||
def __init__(self, is_training=False,
|
||||
saved_models_path=None,
|
||||
training_data_src_path=None,
|
||||
training_data_dst_path=None,
|
||||
pretraining_data_path=None,
|
||||
pretrained_model_path=None,
|
||||
no_preview=False,
|
||||
force_model_name=None,
|
||||
force_gpu_idxs=None,
|
||||
cpu_only=False,
|
||||
debug=False,
|
||||
**kwargs):
|
||||
self.is_training = is_training
|
||||
self.saved_models_path = saved_models_path
|
||||
self.training_data_src_path = training_data_src_path
|
||||
self.training_data_dst_path = training_data_dst_path
|
||||
self.pretraining_data_path = pretraining_data_path
|
||||
|
||||
self.debug = debug
|
||||
self.pretrained_model_path = pretrained_model_path
|
||||
self.no_preview = no_preview
|
||||
self.is_training_mode = is_training
|
||||
self.debug = debug
|
||||
|
||||
self.model_class_name = model_class_name = Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1]
|
||||
|
||||
if force_model_name is not None:
|
||||
self.model_name = force_model_name
|
||||
else:
|
||||
while True:
|
||||
# gather all model dat files
|
||||
saved_models_names = []
|
||||
for filepath in pathex.get_file_paths(saved_models_path):
|
||||
filepath_name = filepath.name
|
||||
if filepath_name.endswith(f'{model_class_name}_data.dat'):
|
||||
saved_models_names += [ (filepath_name.split('_')[0], os.path.getmtime(filepath)) ]
|
||||
|
||||
# sort by modified datetime
|
||||
saved_models_names = sorted(saved_models_names, key=operator.itemgetter(1), reverse=True )
|
||||
saved_models_names = [ x[0] for x in saved_models_names ]
|
||||
|
||||
if len(saved_models_names) != 0:
|
||||
io.log_info ("Choose one of saved models, or enter a name to create a new model.")
|
||||
io.log_info ("[r] : rename")
|
||||
io.log_info ("[d] : delete")
|
||||
io.log_info ("")
|
||||
for i, model_name in enumerate(saved_models_names):
|
||||
s = f"[{i}] : {model_name} "
|
||||
if i == 0:
|
||||
s += "- latest"
|
||||
io.log_info (s)
|
||||
|
||||
inp = io.input_str(f"", "0", show_default_value=False )
|
||||
model_idx = -1
|
||||
try:
|
||||
model_idx = np.clip ( int(inp), 0, len(saved_models_names)-1 )
|
||||
except:
|
||||
pass
|
||||
|
||||
if model_idx == -1:
|
||||
if len(inp) == 1:
|
||||
is_rename = inp[0] == 'r'
|
||||
is_delete = inp[0] == 'd'
|
||||
|
||||
if is_rename or is_delete:
|
||||
if len(saved_models_names) != 0:
|
||||
|
||||
if is_rename:
|
||||
name = io.input_str(f"Enter the name of the model you want to rename")
|
||||
elif is_delete:
|
||||
name = io.input_str(f"Enter the name of the model you want to delete")
|
||||
|
||||
if name in saved_models_names:
|
||||
|
||||
if is_rename:
|
||||
new_model_name = io.input_str(f"Enter new name of the model")
|
||||
|
||||
for filepath in pathex.get_file_paths(saved_models_path):
|
||||
filepath_name = filepath.name
|
||||
|
||||
model_filename, remain_filename = filepath_name.split('_', 1)
|
||||
if model_filename == name:
|
||||
|
||||
if is_rename:
|
||||
new_filepath = filepath.parent / ( new_model_name + '_' + remain_filename )
|
||||
filepath.rename (new_filepath)
|
||||
elif is_delete:
|
||||
filepath.unlink()
|
||||
continue
|
||||
|
||||
self.model_name = inp
|
||||
else:
|
||||
self.model_name = saved_models_names[model_idx]
|
||||
|
||||
else:
|
||||
self.model_name = io.input_str(f"No saved models found. Enter a name of a new model", "noname")
|
||||
|
||||
break
|
||||
|
||||
self.model_name = self.model_name + '_' + self.model_class_name
|
||||
|
||||
self.iter = 0
|
||||
self.options = {}
|
||||
self.loss_history = []
|
||||
self.sample_for_preview = None
|
||||
self.choosed_gpu_indexes = None
|
||||
|
||||
model_data = {}
|
||||
self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') )
|
||||
if self.model_data_path.exists():
|
||||
io.log_info (f"Loading {self.model_name} model...")
|
||||
model_data = pickle.loads ( self.model_data_path.read_bytes() )
|
||||
self.iter = max( model_data.get('iter',0), model_data.get('epoch',0) )
|
||||
if 'epoch' in self.options:
|
||||
self.options.pop('epoch')
|
||||
self.iter = model_data.get('iter',0)
|
||||
if self.iter != 0:
|
||||
self.options = model_data['options']
|
||||
self.loss_history = model_data.get('loss_history', [])
|
||||
self.sample_for_preview = model_data.get('sample_for_preview', None)
|
||||
self.choosed_gpu_indexes = model_data.get('choosed_gpu_indexes', None)
|
||||
|
||||
ask_override = self.is_training_mode and self.iter != 0 and io.input_in_time ("Press enter in 2 seconds to override model settings.", 5 if io.is_colab() else 2 )
|
||||
|
||||
yn_str = {True:'y',False:'n'}
|
||||
|
||||
if self.iter == 0:
|
||||
if self.is_first_run():
|
||||
io.log_info ("\nModel first run.")
|
||||
|
||||
if ask_enable_autobackup and (self.iter == 0 or ask_override):
|
||||
default_autobackup = False if self.iter == 0 else self.options.get('autobackup',False)
|
||||
self.options['autobackup'] = io.input_bool("Enable autobackup? (y/n ?:help skip:%s) : " % (yn_str[default_autobackup]) , default_autobackup, help_message="Autobackup model files with preview every hour for last 15 hours. Latest backup located in model/<>_autobackups/01")
|
||||
else:
|
||||
self.options['autobackup'] = self.options.get('autobackup', False)
|
||||
self.device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(suggest_best_multi_gpu=True)) \
|
||||
if not cpu_only else nn.DeviceConfig.CPU()
|
||||
|
||||
if ask_write_preview_history and (self.iter == 0 or ask_override):
|
||||
default_write_preview_history = False if self.iter == 0 else self.options.get('write_preview_history',False)
|
||||
self.options['write_preview_history'] = io.input_bool("Write preview history? (y/n ?:help skip:%s) : " % (yn_str[default_write_preview_history]) , default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.")
|
||||
else:
|
||||
self.options['write_preview_history'] = self.options.get('write_preview_history', False)
|
||||
nn.initialize(self.device_config)
|
||||
|
||||
if (self.iter == 0 or ask_override) and self.options['write_preview_history'] and io.is_support_windows():
|
||||
choose_preview_history = io.input_bool("Choose image for the preview history? (y/n skip:%s) : " % (yn_str[False]) , False)
|
||||
elif (self.iter == 0 or ask_override) and self.options['write_preview_history'] and io.is_colab():
|
||||
choose_preview_history = io.input_bool("Randomly choose new image for preview history? (y/n ?:help skip:%s) : " % (yn_str[False]), False, help_message="Preview image history will stay stuck with old faces if you reuse the same model on different celebs. Choose no unless you are changing src/dst to a new person")
|
||||
else:
|
||||
choose_preview_history = False
|
||||
####
|
||||
self.default_options_path = saved_models_path / f'{self.model_class_name}_default_options.dat'
|
||||
self.default_options = {}
|
||||
if self.default_options_path.exists():
|
||||
try:
|
||||
self.default_options = pickle.loads ( self.default_options_path.read_bytes() )
|
||||
except:
|
||||
pass
|
||||
|
||||
if ask_target_iter:
|
||||
if (self.iter == 0 or ask_override):
|
||||
self.options['target_iter'] = max(0, io.input_int("Target iteration (skip:unlimited/default) : ", 0))
|
||||
else:
|
||||
self.options['target_iter'] = max(model_data.get('target_iter',0), self.options.get('target_epoch',0))
|
||||
if 'target_epoch' in self.options:
|
||||
self.options.pop('target_epoch')
|
||||
self.choose_preview_history = False
|
||||
self.batch_size = self.load_or_def_option('batch_size', 1)
|
||||
#####
|
||||
|
||||
if ask_batch_size and (self.iter == 0 or ask_override):
|
||||
default_batch_size = 0 if self.iter == 0 else self.options.get('batch_size',0)
|
||||
self.batch_size = max(0, io.input_int("Batch_size (?:help skip:%d) : " % (default_batch_size), default_batch_size, help_message="Larger batch size is better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually."))
|
||||
else:
|
||||
self.batch_size = self.options.get('batch_size', 0)
|
||||
|
||||
if ask_random_flip:
|
||||
default_random_flip = self.options.get('random_flip', True)
|
||||
if (self.iter == 0 or ask_override):
|
||||
self.options['random_flip'] = io.input_bool(f"Flip faces randomly? (y/n ?:help skip:{yn_str[default_random_flip]}) : ", default_random_flip, help_message="Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset.")
|
||||
else:
|
||||
self.options['random_flip'] = self.options.get('random_flip', default_random_flip)
|
||||
self.on_initialize_options()
|
||||
if self.is_first_run():
|
||||
# save as default options only for first run model initialize
|
||||
self.default_options_path.write_bytes( pickle.dumps (self.options) )
|
||||
|
||||
self.autobackup = self.options.get('autobackup', False)
|
||||
if not self.autobackup and 'autobackup' in self.options:
|
||||
self.options.pop('autobackup')
|
||||
|
||||
self.write_preview_history = self.options.get('write_preview_history', False)
|
||||
if not self.write_preview_history and 'write_preview_history' in self.options:
|
||||
self.options.pop('write_preview_history')
|
||||
|
||||
self.target_iter = self.options.get('target_iter',0)
|
||||
if self.target_iter == 0 and 'target_iter' in self.options:
|
||||
self.options.pop('target_iter')
|
||||
|
||||
#self.batch_size = self.options.get('batch_size',0)
|
||||
self.sort_by_yaw = self.options.get('sort_by_yaw',False)
|
||||
self.random_flip = self.options.get('random_flip',True)
|
||||
|
||||
self.onInitializeOptions(self.iter == 0, ask_override)
|
||||
|
||||
nnlib.import_all(self.device_config)
|
||||
self.keras = nnlib.keras
|
||||
self.K = nnlib.keras.backend
|
||||
|
||||
self.onInitialize()
|
||||
|
||||
self.on_initialize()
|
||||
self.options['batch_size'] = self.batch_size
|
||||
|
||||
if self.debug or self.batch_size == 0:
|
||||
self.batch_size = 1
|
||||
|
||||
if self.is_training_mode:
|
||||
if self.device_args['force_gpu_idx'] == -1:
|
||||
self.preview_history_path = self.model_path / ( '%s_history' % (self.get_model_name()) )
|
||||
self.autobackups_path = self.model_path / ( '%s_autobackups' % (self.get_model_name()) )
|
||||
else:
|
||||
self.preview_history_path = self.model_path / ( '%d_%s_history' % (self.device_args['force_gpu_idx'], self.get_model_name()) )
|
||||
self.autobackups_path = self.model_path / ( '%d_%s_autobackups' % (self.device_args['force_gpu_idx'], self.get_model_name()) )
|
||||
|
||||
if self.is_training:
|
||||
self.preview_history_path = self.saved_models_path / ( f'{self.get_model_name()}_history' )
|
||||
self.autobackups_path = self.saved_models_path / ( f'{self.get_model_name()}_autobackups' )
|
||||
|
||||
if self.autobackup:
|
||||
self.autobackup_current_hour = time.localtime().tm_hour
|
||||
|
|
@ -169,7 +186,7 @@ class ModelBase(object):
|
|||
self.preview_history_path.mkdir(exist_ok=True)
|
||||
else:
|
||||
if self.iter == 0:
|
||||
for filename in Path_utils.get_image_paths(self.preview_history_path):
|
||||
for filename in pathex.get_image_paths(self.preview_history_path):
|
||||
Path(filename).unlink()
|
||||
|
||||
if self.generator_list is None:
|
||||
|
|
@ -179,15 +196,15 @@ class ModelBase(object):
|
|||
if not isinstance(generator, SampleGeneratorBase):
|
||||
raise ValueError('training data generator is not subclass of SampleGeneratorBase')
|
||||
|
||||
if self.sample_for_preview is None or choose_preview_history:
|
||||
if choose_preview_history and io.is_support_windows():
|
||||
if self.sample_for_preview is None or self.choose_preview_history:
|
||||
if self.choose_preview_history and io.is_support_windows():
|
||||
io.log_info ("Choose image for the preview history. [p] - next. [enter] - confirm.")
|
||||
wnd_name = "[p] - next. [enter] - confirm."
|
||||
io.named_window(wnd_name)
|
||||
io.capture_keys(wnd_name)
|
||||
choosed = False
|
||||
while not choosed:
|
||||
self.sample_for_preview = self.generate_next_sample()
|
||||
self.sample_for_preview = self.generate_next_samples()
|
||||
preview = self.get_static_preview()
|
||||
io.show_image( wnd_name, (preview*255).astype(np.uint8) )
|
||||
|
||||
|
|
@ -207,73 +224,66 @@ class ModelBase(object):
|
|||
|
||||
io.destroy_window(wnd_name)
|
||||
else:
|
||||
self.sample_for_preview = self.generate_next_sample()
|
||||
self.sample_for_preview = self.generate_next_samples()
|
||||
|
||||
try:
|
||||
self.get_static_preview()
|
||||
except:
|
||||
self.sample_for_preview = self.generate_next_sample()
|
||||
self.sample_for_preview = self.generate_next_samples()
|
||||
|
||||
self.last_sample = self.sample_for_preview
|
||||
|
||||
###Generate text summary of model hyperparameters
|
||||
#Find the longest key name and value string. Used as column widths.
|
||||
width_name = max([len(k) for k in self.options.keys()] + [17]) + 1 # Single space buffer to left edge. Minimum of 17, the length of the longest static string used "Current iteration"
|
||||
width_value = max([len(str(x)) for x in self.options.values()] + [len(str(self.iter)), len(self.get_model_name())]) + 1 # Single space buffer to right edge
|
||||
if not self.device_config.cpu_only: #Check length of GPU names
|
||||
width_value = max([len(nnlib.device.getDeviceName(idx))+1 for idx in self.device_config.gpu_idxs] + [width_value])
|
||||
width_total = width_name + width_value + 2 #Plus 2 for ": "
|
||||
io.log_info( self.get_summary_text() )
|
||||
|
||||
model_summary_text = []
|
||||
model_summary_text += [f'=={" Model Summary ":=^{width_total}}=='] # Model/status summary
|
||||
model_summary_text += [f'=={" "*width_total}==']
|
||||
model_summary_text += [f'=={"Model name": >{width_name}}: {self.get_model_name(): <{width_value}}=='] # Name
|
||||
model_summary_text += [f'=={" "*width_total}==']
|
||||
model_summary_text += [f'=={"Current iteration": >{width_name}}: {str(self.iter): <{width_value}}=='] # Iter
|
||||
model_summary_text += [f'=={" "*width_total}==']
|
||||
def load_or_def_option(self, name, def_value):
|
||||
options_val = self.options.get(name, None)
|
||||
if options_val is not None:
|
||||
return options_val
|
||||
|
||||
model_summary_text += [f'=={" Model Options ":-^{width_total}}=='] # Model options
|
||||
model_summary_text += [f'=={" "*width_total}==']
|
||||
for key in self.options.keys():
|
||||
model_summary_text += [f'=={key: >{width_name}}: {str(self.options[key]): <{width_value}}=='] # self.options key/value pairs
|
||||
model_summary_text += [f'=={" "*width_total}==']
|
||||
def_opt_val = self.default_options.get(name, None)
|
||||
if def_opt_val is not None:
|
||||
return def_opt_val
|
||||
|
||||
model_summary_text += [f'=={" Running On ":-^{width_total}}=='] # Training hardware info
|
||||
model_summary_text += [f'=={" "*width_total}==']
|
||||
if self.device_config.multi_gpu:
|
||||
model_summary_text += [f'=={"Using multi_gpu": >{width_name}}: {"True": <{width_value}}=='] # multi_gpu
|
||||
model_summary_text += [f'=={" "*width_total}==']
|
||||
if self.device_config.cpu_only:
|
||||
model_summary_text += [f'=={"Using device": >{width_name}}: {"CPU": <{width_value}}=='] # cpu_only
|
||||
else:
|
||||
for idx in self.device_config.gpu_idxs:
|
||||
model_summary_text += [f'=={"Device index": >{width_name}}: {idx: <{width_value}}=='] # GPU hardware device index
|
||||
model_summary_text += [f'=={"Name": >{width_name}}: {nnlib.device.getDeviceName(idx): <{width_value}}=='] # GPU name
|
||||
vram_str = f'{nnlib.device.getDeviceVRAMTotalGb(idx):.2f}GB' # GPU VRAM - Formated as #.## (or ##.##)
|
||||
model_summary_text += [f'=={"VRAM": >{width_name}}: {vram_str: <{width_value}}==']
|
||||
model_summary_text += [f'=={" "*width_total}==']
|
||||
model_summary_text += [f'=={"="*width_total}==']
|
||||
return def_value
|
||||
|
||||
if not self.device_config.cpu_only and self.device_config.gpu_vram_gb[0] <= 2: # Low VRAM warning
|
||||
model_summary_text += ["/!\\"]
|
||||
model_summary_text += ["/!\\ WARNING:"]
|
||||
model_summary_text += ["/!\\ You are using a GPU with 2GB or less VRAM. This may significantly reduce the quality of your result!"]
|
||||
model_summary_text += ["/!\\ If training does not start, close all programs and try again."]
|
||||
model_summary_text += ["/!\\ Also you can disable Windows Aero Desktop to increase available VRAM."]
|
||||
model_summary_text += ["/!\\"]
|
||||
def ask_override(self):
|
||||
return self.is_training and self.iter != 0 and io.input_in_time ("Press enter in 2 seconds to override model settings.", 5 if io.is_colab() else 2 )
|
||||
|
||||
def ask_enable_autobackup(self):
|
||||
default_autobackup = self.options['autobackup'] = self.load_or_def_option('autobackup', False)
|
||||
self.options['autobackup'] = io.input_bool(f"Enable autobackup", default_autobackup, help_message="Autobackup model files with preview every hour for last 15 hours. Latest backup located in model/<>_autobackups/01")
|
||||
|
||||
def ask_write_preview_history(self):
|
||||
default_write_preview_history = self.load_or_def_option('write_preview_history', False)
|
||||
self.options['write_preview_history'] = io.input_bool(f"Write preview history", default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.")
|
||||
|
||||
if self.options['write_preview_history']:
|
||||
if io.is_support_windows():
|
||||
self.choose_preview_history = io.input_bool("Choose image for the preview history", False)
|
||||
elif io.is_colab():
|
||||
self.choose_preview_history = io.input_bool("Randomly choose new image for preview history", False, help_message="Preview image history will stay stuck with old faces if you reuse the same model on different celebs. Choose no unless you are changing src/dst to a new person")
|
||||
|
||||
def ask_target_iter(self):
|
||||
default_target_iter = self.load_or_def_option('target_iter', 0)
|
||||
self.options['target_iter'] = max(0, io.input_int("Target iteration", default_target_iter))
|
||||
|
||||
def ask_random_flip(self):
|
||||
default_random_flip = self.load_or_def_option('random_flip', True)
|
||||
self.options['random_flip'] = io.input_bool("Flip faces randomly", default_random_flip, help_message="Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset.")
|
||||
|
||||
def ask_batch_size(self, suggest_batch_size=None):
|
||||
default_batch_size = self.load_or_def_option('batch_size', suggest_batch_size or self.batch_size)
|
||||
self.batch_size = max(0, io.input_int("Batch_size", default_batch_size, help_message="Larger batch size is better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually."))
|
||||
|
||||
model_summary_text = "\n".join (model_summary_text)
|
||||
self.model_summary_text = model_summary_text
|
||||
io.log_info(model_summary_text)
|
||||
|
||||
#overridable
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
def on_initialize_options(self):
|
||||
pass
|
||||
|
||||
#overridable
|
||||
def onInitialize(self):
|
||||
def on_initialize(self):
|
||||
'''
|
||||
initialize your keras models
|
||||
initialize your models
|
||||
|
||||
store and retrieve your model options in self.options['']
|
||||
|
||||
|
|
@ -283,12 +293,12 @@ class ModelBase(object):
|
|||
|
||||
#overridable
|
||||
def onSave(self):
|
||||
#save your keras models here
|
||||
#save your models here
|
||||
pass
|
||||
|
||||
#overridable
|
||||
def onTrainOneIter(self, sample, generator_list):
|
||||
#train your keras models here
|
||||
#train your models here
|
||||
|
||||
#return array of losses
|
||||
return ( ('loss_src', 0), ('loss_dst', 0) )
|
||||
|
|
@ -301,42 +311,26 @@ class ModelBase(object):
|
|||
|
||||
#overridable if you want model name differs from folder name
|
||||
def get_model_name(self):
|
||||
return Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1]
|
||||
return self.model_name
|
||||
|
||||
#overridable , return [ [model, filename],... ] list
|
||||
def get_model_filename_list(self):
|
||||
return []
|
||||
|
||||
#overridable
|
||||
def get_ConverterConfig(self):
|
||||
#return predictor_func, predictor_input_shape, ConverterConfig() for the model
|
||||
def get_MergerConfig(self):
|
||||
#return predictor_func, predictor_input_shape, MergerConfig() for the model
|
||||
raise NotImplementedError
|
||||
|
||||
def get_pretraining_data_path(self):
|
||||
return self.pretraining_data_path
|
||||
|
||||
def get_target_iter(self):
|
||||
return self.target_iter
|
||||
|
||||
def is_reached_iter_goal(self):
|
||||
return self.target_iter != 0 and self.iter >= self.target_iter
|
||||
|
||||
#multi gpu in keras actually is fake and doesn't work for training https://github.com/keras-team/keras/issues/11976
|
||||
#def to_multi_gpu_model_if_possible (self, models_list):
|
||||
# if len(self.device_config.gpu_idxs) > 1:
|
||||
# #make batch_size to divide on GPU count without remainder
|
||||
# self.batch_size = int( self.batch_size / len(self.device_config.gpu_idxs) )
|
||||
# if self.batch_size == 0:
|
||||
# self.batch_size = 1
|
||||
# self.batch_size *= len(self.device_config.gpu_idxs)
|
||||
#
|
||||
# result = []
|
||||
# for model in models_list:
|
||||
# for i in range( len(model.output_names) ):
|
||||
# model.output_names = 'output_%d' % (i)
|
||||
# result += [ nnlib.keras.utils.multi_gpu_model( model, self.device_config.gpu_idxs ) ]
|
||||
#
|
||||
# return result
|
||||
# else:
|
||||
# return models_list
|
||||
|
||||
def get_previews(self):
|
||||
return self.onGetPreview ( self.last_sample )
|
||||
|
||||
|
|
@ -345,21 +339,23 @@ class ModelBase(object):
|
|||
|
||||
def save(self):
|
||||
summary_path = self.get_strpath_storage_for_file('summary.txt')
|
||||
Path( summary_path ).write_text(self.model_summary_text)
|
||||
Path( summary_path ).write_text( self.get_summary_text() )
|
||||
|
||||
self.onSave()
|
||||
|
||||
model_data = {
|
||||
'iter': self.iter,
|
||||
'options': self.options,
|
||||
'loss_history': self.loss_history,
|
||||
'sample_for_preview' : self.sample_for_preview
|
||||
'sample_for_preview' : self.sample_for_preview,
|
||||
'choosed_gpu_indexes' : self.choosed_gpu_indexes,
|
||||
}
|
||||
self.model_data_path.write_bytes( pickle.dumps(model_data) )
|
||||
|
||||
bckp_filename_list = [ self.get_strpath_storage_for_file(filename) for _, filename in self.get_model_filename_list() ]
|
||||
bckp_filename_list += [ str(summary_path), str(self.model_data_path) ]
|
||||
pathex.write_bytes_safe (self.model_data_path, pickle.dumps(model_data) )
|
||||
|
||||
if self.autobackup:
|
||||
bckp_filename_list = [ self.get_strpath_storage_for_file(filename) for _, filename in self.get_model_filename_list() ]
|
||||
bckp_filename_list += [ str(summary_path), str(self.model_data_path) ]
|
||||
|
||||
current_hour = time.localtime().tm_hour
|
||||
if self.autobackup_current_hour != current_hour:
|
||||
self.autobackup_current_hour = current_hour
|
||||
|
|
@ -373,10 +369,10 @@ class ModelBase(object):
|
|||
|
||||
if idx_backup_path.exists():
|
||||
if i == 15:
|
||||
Path_utils.delete_all_files(idx_backup_path)
|
||||
pathex.delete_all_files(idx_backup_path)
|
||||
else:
|
||||
next_idx_packup_path.mkdir(exist_ok=True)
|
||||
Path_utils.move_all_files (idx_backup_path, next_idx_packup_path)
|
||||
pathex.move_all_files (idx_backup_path, next_idx_packup_path)
|
||||
|
||||
if i == 1:
|
||||
idx_backup_path.mkdir(exist_ok=True)
|
||||
|
|
@ -394,97 +390,6 @@ class ModelBase(object):
|
|||
img = (np.concatenate ( [preview_lh, preview], axis=0 ) * 255).astype(np.uint8)
|
||||
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 ]
|
||||
|
||||
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 ]
|
||||
|
||||
|
||||
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) + '.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
|
||||
|
||||
"""
|
||||
#unused
|
||||
, optimizer_filename_list=[]
|
||||
if len(optimizer_filename_list) != 0:
|
||||
opt_filename = self.get_strpath_storage_for_file('opt.h5')
|
||||
|
||||
try:
|
||||
d = {}
|
||||
for opt, filename in optimizer_filename_list:
|
||||
fd = {}
|
||||
symbolic_weights = getattr(opt, 'weights')
|
||||
if symbolic_weights:
|
||||
fd['weights'] = self.K.batch_get_value(symbolic_weights)
|
||||
|
||||
d[filename] = fd
|
||||
|
||||
with open(opt_filename+'.tmp', 'wb') as f:
|
||||
f.write( pickle.dumps(d) )
|
||||
|
||||
rename_list += [('', 'opt.h5')]
|
||||
except Exception as e:
|
||||
print ("Unable to save ", opt_filename)
|
||||
"""
|
||||
|
||||
for _, filename in rename_list:
|
||||
filename = self.get_strpath_storage_for_file(filename)
|
||||
source_filename = Path(filename+'.tmp')
|
||||
if source_filename.exists():
|
||||
target_filename = Path(filename)
|
||||
if target_filename.exists():
|
||||
target_filename.unlink()
|
||||
source_filename.rename ( str(target_filename) )
|
||||
|
||||
def debug_one_iter(self):
|
||||
images = []
|
||||
for generator in self.generator_list:
|
||||
|
|
@ -494,19 +399,15 @@ class ModelBase(object):
|
|||
|
||||
return imagelib.equalize_and_stack_square (images)
|
||||
|
||||
def generate_next_sample(self):
|
||||
return [ generator.generate_next() for generator in self.generator_list]
|
||||
|
||||
#overridable
|
||||
def on_success_train_one_iter(self):
|
||||
pass
|
||||
def generate_next_samples(self):
|
||||
self.last_sample = sample = [ generator.generate_next() for generator in self.generator_list]
|
||||
return sample
|
||||
|
||||
def train_one_iter(self):
|
||||
sample = self.generate_next_sample()
|
||||
|
||||
iter_time = time.time()
|
||||
losses = self.onTrainOneIter(sample, self.generator_list)
|
||||
losses = self.onTrainOneIter()
|
||||
iter_time = time.time() - iter_time
|
||||
self.last_sample = sample
|
||||
|
||||
self.loss_history.append ( [float(loss[1]) for loss in losses] )
|
||||
|
||||
|
|
@ -527,17 +428,15 @@ class ModelBase(object):
|
|||
img = (np.concatenate ( [preview_lh, preview], axis=0 ) * 255).astype(np.uint8)
|
||||
cv2_imwrite (filepath, img )
|
||||
|
||||
self.on_success_train_one_iter()
|
||||
|
||||
self.iter += 1
|
||||
|
||||
return self.iter, iter_time
|
||||
|
||||
def pass_one_iter(self):
|
||||
self.last_sample = self.generate_next_sample()
|
||||
self.generate_next_samples()
|
||||
|
||||
def finalize(self):
|
||||
nnlib.finalize_all()
|
||||
nn.tf_close_session()
|
||||
|
||||
def is_first_run(self):
|
||||
return self.iter == 0
|
||||
|
|
@ -554,6 +453,10 @@ class ModelBase(object):
|
|||
def get_iter(self):
|
||||
return self.iter
|
||||
|
||||
def set_iter(self, iter):
|
||||
self.iter = iter
|
||||
self.loss_history = self.loss_history[:iter]
|
||||
|
||||
def get_loss_history(self):
|
||||
return self.loss_history
|
||||
|
||||
|
|
@ -564,30 +467,48 @@ class ModelBase(object):
|
|||
return self.generator_list
|
||||
|
||||
def get_model_root_path(self):
|
||||
return self.model_path
|
||||
return self.saved_models_path
|
||||
|
||||
def get_strpath_storage_for_file(self, filename):
|
||||
if self.device_args['force_gpu_idx'] == -1:
|
||||
return str( self.model_path / ( self.get_model_name() + '_' + filename) )
|
||||
else:
|
||||
return str( self.model_path / ( str(self.device_args['force_gpu_idx']) + '_' + self.get_model_name() + '_' + filename) )
|
||||
return str( self.saved_models_path / ( self.get_model_name() + '_' + filename) )
|
||||
|
||||
def set_vram_batch_requirements (self, d):
|
||||
#example d = {2:2,3:4,4:8,5:16,6:32,7:32,8:32,9:48}
|
||||
keys = [x for x in d.keys()]
|
||||
def get_summary_text(self):
|
||||
###Generate text summary of model hyperparameters
|
||||
#Find the longest key name and value string. Used as column widths.
|
||||
width_name = max([len(k) for k in self.options.keys()] + [17]) + 1 # Single space buffer to left edge. Minimum of 17, the length of the longest static string used "Current iteration"
|
||||
width_value = max([len(str(x)) for x in self.options.values()] + [len(str(self.get_iter())), len(self.get_model_name())]) + 1 # Single space buffer to right edge
|
||||
if not self.device_config.cpu_only: #Check length of GPU names
|
||||
width_value = max([len(device.name)+1 for device in self.device_config.devices] + [width_value])
|
||||
width_total = width_name + width_value + 2 #Plus 2 for ": "
|
||||
|
||||
summary_text = []
|
||||
summary_text += [f'=={" Model Summary ":=^{width_total}}=='] # Model/status summary
|
||||
summary_text += [f'=={" "*width_total}==']
|
||||
summary_text += [f'=={"Model name": >{width_name}}: {self.get_model_name(): <{width_value}}=='] # Name
|
||||
summary_text += [f'=={" "*width_total}==']
|
||||
summary_text += [f'=={"Current iteration": >{width_name}}: {str(self.get_iter()): <{width_value}}=='] # Iter
|
||||
summary_text += [f'=={" "*width_total}==']
|
||||
|
||||
summary_text += [f'=={" Model Options ":-^{width_total}}=='] # Model options
|
||||
summary_text += [f'=={" "*width_total}==']
|
||||
for key in self.options.keys():
|
||||
summary_text += [f'=={key: >{width_name}}: {str(self.options[key]): <{width_value}}=='] # self.options key/value pairs
|
||||
summary_text += [f'=={" "*width_total}==']
|
||||
|
||||
summary_text += [f'=={" Running On ":-^{width_total}}=='] # Training hardware info
|
||||
summary_text += [f'=={" "*width_total}==']
|
||||
if self.device_config.cpu_only:
|
||||
if self.batch_size == 0:
|
||||
self.batch_size = 2
|
||||
summary_text += [f'=={"Using device": >{width_name}}: {"CPU": <{width_value}}=='] # cpu_only
|
||||
else:
|
||||
if self.batch_size == 0:
|
||||
for x in keys:
|
||||
if self.device_config.gpu_vram_gb[0] <= x:
|
||||
self.batch_size = d[x]
|
||||
break
|
||||
|
||||
if self.batch_size == 0:
|
||||
self.batch_size = d[ keys[-1] ]
|
||||
for device in self.device_config.devices:
|
||||
summary_text += [f'=={"Device index": >{width_name}}: {device.index: <{width_value}}=='] # GPU hardware device index
|
||||
summary_text += [f'=={"Name": >{width_name}}: {device.name: <{width_value}}=='] # GPU name
|
||||
vram_str = f'{device.total_mem_gb:.2f}GB' # GPU VRAM - Formated as #.## (or ##.##)
|
||||
summary_text += [f'=={"VRAM": >{width_name}}: {vram_str: <{width_value}}==']
|
||||
summary_text += [f'=={" "*width_total}==']
|
||||
summary_text += [f'=={"="*width_total}==']
|
||||
summary_text = "\n".join (summary_text)
|
||||
return summary_text
|
||||
|
||||
@staticmethod
|
||||
def get_loss_history_preview(loss_history, iter, w, c):
|
||||
|
|
|
|||
|
|
@ -1,490 +0,0 @@
|
|||
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
|
||||
from samplelib import *
|
||||
|
||||
from facelib import PoseEstimator
|
||||
|
||||
class AVATARModel(ModelBase):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs,
|
||||
ask_random_flip=False)
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run:
|
||||
#avatar_type = io.input_int("Avatar type ( 0:source, 1:head, 2:full_face ?:help skip:1) : ", 1, [0,1,2],
|
||||
# help_message="Training target for the model. Source is direct untouched images. Full_face or head are centered nose unaligned faces.")
|
||||
#avatar_type = {0:'source',
|
||||
# 1:'head',
|
||||
# 2:'full_face'}[avatar_type]
|
||||
|
||||
self.options['avatar_type'] = 'head'
|
||||
else:
|
||||
self.options['avatar_type'] = self.options.get('avatar_type', 'head')
|
||||
|
||||
if is_first_run or ask_override:
|
||||
def_stage = self.options.get('stage', 1)
|
||||
self.options['stage'] = io.input_int("Stage (0, 1, 2 ?:help skip:%d) : " % def_stage, def_stage, [0,1,2], help_message="Train first stage, then second. Tune batch size to maximum possible for both stages.")
|
||||
else:
|
||||
self.options['stage'] = self.options.get('stage', 1)
|
||||
|
||||
#override
|
||||
def onInitialize(self, batch_size=-1, **in_options):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
self.set_vram_batch_requirements({6:4})
|
||||
|
||||
resolution = self.resolution = 224
|
||||
avatar_type = self.options['avatar_type']
|
||||
stage = self.stage = self.options['stage']
|
||||
df_res = self.df_res = 128
|
||||
df_bgr_shape = (df_res, df_res, 3)
|
||||
df_mask_shape = (df_res, df_res, 1)
|
||||
res_bgr_shape = (resolution, resolution, 3)
|
||||
res_bgr_t_shape = (resolution, resolution, 9)
|
||||
|
||||
self.enc = modelify(AVATARModel.EncFlow())( [Input(df_bgr_shape),] )
|
||||
|
||||
self.decA64 = modelify(AVATARModel.DecFlow()) ( [ Input(K.int_shape(self.enc.outputs[0])[1:]) ] )
|
||||
self.decB64 = modelify(AVATARModel.DecFlow()) ( [ Input(K.int_shape(self.enc.outputs[0])[1:]) ] )
|
||||
self.D = modelify(AVATARModel.Discriminator() ) (Input(df_bgr_shape))
|
||||
self.C = modelify(AVATARModel.ResNet (9, n_blocks=6, ngf=128, use_dropout=False))( Input(res_bgr_t_shape))
|
||||
|
||||
self.CA_conv_weights_list = []
|
||||
if self.is_first_run():
|
||||
for model, _ in self.get_model_filename_list():
|
||||
for layer in model.layers:
|
||||
if type(layer) == keras.layers.Conv2D:
|
||||
self.CA_conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights
|
||||
|
||||
if not self.is_first_run():
|
||||
self.load_weights_safe( self.get_model_filename_list() )
|
||||
|
||||
def DLoss(labels,logits):
|
||||
return K.mean(K.binary_crossentropy(labels,logits))
|
||||
|
||||
warped_A64 = Input(df_bgr_shape)
|
||||
real_A64 = Input(df_bgr_shape)
|
||||
real_A64m = Input(df_mask_shape)
|
||||
|
||||
real_B64_t0 = Input(df_bgr_shape)
|
||||
real_B64_t1 = Input(df_bgr_shape)
|
||||
real_B64_t2 = Input(df_bgr_shape)
|
||||
|
||||
real_A64_t0 = Input(df_bgr_shape)
|
||||
real_A64m_t0 = Input(df_mask_shape)
|
||||
real_A_t0 = Input(res_bgr_shape)
|
||||
real_A64_t1 = Input(df_bgr_shape)
|
||||
real_A64m_t1 = Input(df_mask_shape)
|
||||
real_A_t1 = Input(res_bgr_shape)
|
||||
real_A64_t2 = Input(df_bgr_shape)
|
||||
real_A64m_t2 = Input(df_mask_shape)
|
||||
real_A_t2 = Input(res_bgr_shape)
|
||||
|
||||
warped_B64 = Input(df_bgr_shape)
|
||||
real_B64 = Input(df_bgr_shape)
|
||||
real_B64m = Input(df_mask_shape)
|
||||
|
||||
warped_A_code = self.enc (warped_A64)
|
||||
warped_B_code = self.enc (warped_B64)
|
||||
|
||||
rec_A64 = self.decA64(warped_A_code)
|
||||
rec_B64 = self.decB64(warped_B_code)
|
||||
rec_AB64 = self.decA64(warped_B_code)
|
||||
|
||||
def Lambda_grey_mask (x,m):
|
||||
return Lambda (lambda x: x[0]*m+(1-m)*0.5, output_shape= K.int_shape(x)[1:3] + (3,)) ([x, m])
|
||||
|
||||
def Lambda_gray_pad(x):
|
||||
a = np.ones((resolution,resolution,3))*0.5
|
||||
pad = ( resolution - df_res ) // 2
|
||||
a[pad:-pad:,pad:-pad:,:] = 0
|
||||
|
||||
return Lambda ( lambda x: K.spatial_2d_padding(x, padding=((pad, pad), (pad, pad)) ) + K.constant(a, dtype=K.floatx() ),
|
||||
output_shape=(resolution,resolution,3) ) (x)
|
||||
|
||||
def Lambda_concat ( x ):
|
||||
c = sum ( [ K.int_shape(l)[-1] for l in x ] )
|
||||
return Lambda ( lambda x: K.concatenate (x, axis=-1), output_shape=K.int_shape(x[0])[1:3] + (c,) ) (x)
|
||||
|
||||
def Lambda_Cto3t(x):
|
||||
return Lambda ( lambda x: x[...,0:3], output_shape= K.int_shape(x)[1:3] + (3,) ) (x), \
|
||||
Lambda ( lambda x: x[...,3:6], output_shape= K.int_shape(x)[1:3] + (3,) ) (x), \
|
||||
Lambda ( lambda x: x[...,6:9], output_shape= K.int_shape(x)[1:3] + (3,) ) (x)
|
||||
|
||||
real_A64_d = self.D( Lambda_grey_mask(real_A64, real_A64m) )
|
||||
|
||||
real_A64_d_ones = K.ones_like(real_A64_d)
|
||||
fake_A64_d = self.D(rec_AB64)
|
||||
fake_A64_d_ones = K.ones_like(fake_A64_d)
|
||||
fake_A64_d_zeros = K.zeros_like(fake_A64_d)
|
||||
|
||||
rec_AB_t0 = Lambda_gray_pad( self.decA64 (self.enc (real_B64_t0)) )
|
||||
rec_AB_t1 = Lambda_gray_pad( self.decA64 (self.enc (real_B64_t1)) )
|
||||
rec_AB_t2 = Lambda_gray_pad( self.decA64 (self.enc (real_B64_t2)) )
|
||||
|
||||
C_in_A_t0 = Lambda_gray_pad( Lambda_grey_mask (real_A64_t0, real_A64m_t0) )
|
||||
C_in_A_t1 = Lambda_gray_pad( Lambda_grey_mask (real_A64_t1, real_A64m_t1) )
|
||||
C_in_A_t2 = Lambda_gray_pad( Lambda_grey_mask (real_A64_t2, real_A64m_t2) )
|
||||
|
||||
rec_C_A_t0, rec_C_A_t1, rec_C_A_t2 = Lambda_Cto3t ( self.C ( Lambda_concat ( [C_in_A_t0, C_in_A_t1, C_in_A_t2]) ) )
|
||||
rec_C_AB_t0, rec_C_AB_t1, rec_C_AB_t2 = Lambda_Cto3t( self.C ( Lambda_concat ( [rec_AB_t0, rec_AB_t1, rec_AB_t2]) ) )
|
||||
|
||||
#real_A_t012_d = self.CD ( K.concatenate ( [real_A_t0, real_A_t1,real_A_t2], axis=-1) )
|
||||
#real_A_t012_d_ones = K.ones_like(real_A_t012_d)
|
||||
#rec_C_AB_t012_d = self.CD ( K.concatenate ( [rec_C_AB_t0,rec_C_AB_t1, rec_C_AB_t2], axis=-1) )
|
||||
#rec_C_AB_t012_d_ones = K.ones_like(rec_C_AB_t012_d)
|
||||
#rec_C_AB_t012_d_zeros = K.zeros_like(rec_C_AB_t012_d)
|
||||
|
||||
self.G64_view = K.function([warped_A64, warped_B64],[rec_A64, rec_B64, rec_AB64])
|
||||
self.G_view = K.function([real_A64_t0, real_A64m_t0, real_A64_t1, real_A64m_t1, real_A64_t2, real_A64m_t2, real_B64_t0, real_B64_t1, real_B64_t2], [rec_C_A_t0, rec_C_A_t1, rec_C_A_t2, rec_C_AB_t0, rec_C_AB_t1, rec_C_AB_t2])
|
||||
|
||||
if self.is_training_mode:
|
||||
loss_AB64 = K.mean(10 * dssim(kernel_size=int(df_res/11.6),max_value=1.0) ( rec_A64, real_A64*real_A64m + (1-real_A64m)*0.5) ) + \
|
||||
K.mean(10 * dssim(kernel_size=int(df_res/11.6),max_value=1.0) ( rec_B64, real_B64*real_B64m + (1-real_B64m)*0.5) ) + 0.1*DLoss(fake_A64_d_ones, fake_A64_d )
|
||||
|
||||
weights_AB64 = self.enc.trainable_weights + self.decA64.trainable_weights + self.decB64.trainable_weights
|
||||
|
||||
loss_C = K.mean( 10 * dssim(kernel_size=int(resolution/11.6),max_value=1.0) ( real_A_t0, rec_C_A_t0 ) ) + \
|
||||
K.mean( 10 * dssim(kernel_size=int(resolution/11.6),max_value=1.0) ( real_A_t1, rec_C_A_t1 ) ) + \
|
||||
K.mean( 10 * dssim(kernel_size=int(resolution/11.6),max_value=1.0) ( real_A_t2, rec_C_A_t2 ) )
|
||||
#0.1*DLoss(rec_C_AB_t012_d_ones, rec_C_AB_t012_d )
|
||||
|
||||
weights_C = self.C.trainable_weights
|
||||
|
||||
loss_D = (DLoss(real_A64_d_ones, real_A64_d ) + \
|
||||
DLoss(fake_A64_d_zeros, fake_A64_d ) ) * 0.5
|
||||
|
||||
#loss_CD = ( DLoss(real_A_t012_d_ones, real_A_t012_d) + \
|
||||
# DLoss(rec_C_AB_t012_d_zeros, rec_C_AB_t012_d) ) * 0.5
|
||||
#
|
||||
#weights_CD = self.CD.trainable_weights
|
||||
|
||||
def opt(lr=5e-5):
|
||||
return Adam(lr=lr, beta_1=0.5, beta_2=0.999, tf_cpu_mode=2 if 'tensorflow' in self.device_config.backend else 0 )
|
||||
|
||||
self.AB64_train = K.function ([warped_A64, real_A64, real_A64m, warped_B64, real_B64, real_B64m], [loss_AB64], opt().get_updates(loss_AB64, weights_AB64) )
|
||||
self.C_train = K.function ([real_A64_t0, real_A64m_t0, real_A_t0,
|
||||
real_A64_t1, real_A64m_t1, real_A_t1,
|
||||
real_A64_t2, real_A64m_t2, real_A_t2,
|
||||
real_B64_t0, real_B64_t1, real_B64_t2],[ loss_C ], opt().get_updates(loss_C, weights_C) )
|
||||
|
||||
self.D_train = K.function ([warped_A64, real_A64, real_A64m, warped_B64, real_B64, real_B64m],[loss_D], opt().get_updates(loss_D, self.D.trainable_weights) )
|
||||
|
||||
|
||||
#self.CD_train = K.function ([real_A64_t0, real_A64m_t0, real_A_t0,
|
||||
# real_A64_t1, real_A64m_t1, real_A_t1,
|
||||
# real_A64_t2, real_A64m_t2, real_A_t2,
|
||||
# real_B64_t0, real_B64_t1, real_B64_t2 ],[ loss_CD ], opt().get_updates(loss_CD, weights_CD) )
|
||||
|
||||
###########
|
||||
t = SampleProcessor.Types
|
||||
|
||||
training_target = {'source' : t.NONE,
|
||||
'full_face' : t.FACE_TYPE_FULL_NO_ALIGN,
|
||||
'head' : t.FACE_TYPE_HEAD_NO_ALIGN}[avatar_type]
|
||||
|
||||
generators = [
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=False),
|
||||
output_sample_types=[ {'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_BGR), 'resolution':df_res},
|
||||
{'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_BGR), 'resolution':df_res},
|
||||
{'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_M), 'resolution':df_res}
|
||||
] ),
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=False),
|
||||
output_sample_types=[ {'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_BGR), 'resolution':df_res},
|
||||
{'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_BGR), 'resolution':df_res},
|
||||
{'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_M), 'resolution':df_res}
|
||||
] ),
|
||||
|
||||
SampleGeneratorFaceTemporal(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
temporal_image_count=3,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=False),
|
||||
output_sample_types=[{'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_BGR), 'resolution':df_res},#IMG_WARPED_TRANSFORMED
|
||||
{'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_M), 'resolution':df_res},
|
||||
{'types': (t.IMG_SOURCE, training_target, t.MODE_BGR), 'resolution':resolution},
|
||||
] ),
|
||||
|
||||
SampleGeneratorFaceTemporal(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
temporal_image_count=3,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=False),
|
||||
output_sample_types=[{'types': (t.IMG_SOURCE, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_BGR), 'resolution':df_res},
|
||||
{'types': (t.IMG_SOURCE, t.NONE, t.MODE_BGR), 'resolution':resolution},
|
||||
] ),
|
||||
]
|
||||
|
||||
if self.stage == 1:
|
||||
generators[2].set_active(False)
|
||||
generators[3].set_active(False)
|
||||
elif self.stage == 2:
|
||||
generators[0].set_active(False)
|
||||
generators[1].set_active(False)
|
||||
|
||||
self.set_training_data_generators (generators)
|
||||
else:
|
||||
self.G_convert = K.function([real_B64_t0, real_B64_t1, real_B64_t2],[rec_C_AB_t1])
|
||||
|
||||
#override , return [ [model, filename],... ] list
|
||||
def get_model_filename_list(self):
|
||||
return [ [self.enc, 'enc.h5'],
|
||||
[self.decA64, 'decA64.h5'],
|
||||
[self.decB64, 'decB64.h5'],
|
||||
[self.C, 'C.h5'],
|
||||
[self.D, 'D.h5'],
|
||||
#[self.CD, 'CD.h5'],
|
||||
]
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( self.get_model_filename_list() )
|
||||
|
||||
#override
|
||||
def on_success_train_one_iter(self):
|
||||
if len(self.CA_conv_weights_list) != 0:
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
CAInitializerMP ( self.CA_conv_weights_list )
|
||||
self.CA_conv_weights_list = []
|
||||
|
||||
#override
|
||||
def onTrainOneIter(self, generators_samples, generators_list):
|
||||
warped_src64, src64, src64m = generators_samples[0]
|
||||
warped_dst64, dst64, dst64m = generators_samples[1]
|
||||
|
||||
real_A64_t0, real_A64m_t0, real_A_t0, real_A64_t1, real_A64m_t1, real_A_t1, real_A64_t2, real_A64m_t2, real_A_t2 = generators_samples[2]
|
||||
real_B64_t0, _, real_B64_t1, _, real_B64_t2, _ = generators_samples[3]
|
||||
|
||||
if self.stage == 0 or self.stage == 1:
|
||||
loss, = self.AB64_train ( [warped_src64, src64, src64m, warped_dst64, dst64, dst64m] )
|
||||
loss_D, = self.D_train ( [warped_src64, src64, src64m, warped_dst64, dst64, dst64m] )
|
||||
if self.stage != 0:
|
||||
loss_C = loss_CD = 0
|
||||
|
||||
if self.stage == 0 or self.stage == 2:
|
||||
loss_C1, = self.C_train ( [real_A64_t0, real_A64m_t0, real_A_t0,
|
||||
real_A64_t1, real_A64m_t1, real_A_t1,
|
||||
real_A64_t2, real_A64m_t2, real_A_t2,
|
||||
real_B64_t0, real_B64_t1, real_B64_t2] )
|
||||
|
||||
loss_C2, = self.C_train ( [real_A64_t2, real_A64m_t2, real_A_t2,
|
||||
real_A64_t1, real_A64m_t1, real_A_t1,
|
||||
real_A64_t0, real_A64m_t0, real_A_t0,
|
||||
real_B64_t0, real_B64_t1, real_B64_t2] )
|
||||
|
||||
#loss_CD1, = self.CD_train ( [real_A64_t0, real_A64m_t0, real_A_t0,
|
||||
# real_A64_t1, real_A64m_t1, real_A_t1,
|
||||
# real_A64_t2, real_A64m_t2, real_A_t2,
|
||||
# real_B64_t0, real_B64_t1, real_B64_t2] )
|
||||
#
|
||||
#loss_CD2, = self.CD_train ( [real_A64_t2, real_A64m_t2, real_A_t2,
|
||||
# real_A64_t1, real_A64m_t1, real_A_t1,
|
||||
# real_A64_t0, real_A64m_t0, real_A_t0,
|
||||
# real_B64_t0, real_B64_t1, real_B64_t2] )
|
||||
|
||||
loss_C = (loss_C1 + loss_C2) / 2
|
||||
#loss_CD = (loss_CD1 + loss_CD2) / 2
|
||||
if self.stage != 0:
|
||||
loss = loss_D = 0
|
||||
|
||||
return ( ('loss', loss), ('D', loss_D), ('C', loss_C), ) #('CD', loss_CD) )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
test_A064w = sample[0][0][0:4]
|
||||
test_A064r = sample[0][1][0:4]
|
||||
test_A064m = sample[0][2][0:4]
|
||||
|
||||
test_B064w = sample[1][0][0:4]
|
||||
test_B064r = sample[1][1][0:4]
|
||||
test_B064m = sample[1][2][0:4]
|
||||
|
||||
t_src64_0 = sample[2][0][0:4]
|
||||
t_src64m_0 = sample[2][1][0:4]
|
||||
t_src_0 = sample[2][2][0:4]
|
||||
t_src64_1 = sample[2][3][0:4]
|
||||
t_src64m_1 = sample[2][4][0:4]
|
||||
t_src_1 = sample[2][5][0:4]
|
||||
t_src64_2 = sample[2][6][0:4]
|
||||
t_src64m_2 = sample[2][7][0:4]
|
||||
t_src_2 = sample[2][8][0:4]
|
||||
|
||||
t_dst64_0 = sample[3][0][0:4]
|
||||
t_dst_0 = sample[3][1][0:4]
|
||||
t_dst64_1 = sample[3][2][0:4]
|
||||
t_dst_1 = sample[3][3][0:4]
|
||||
t_dst64_2 = sample[3][4][0:4]
|
||||
t_dst_2 = sample[3][5][0:4]
|
||||
|
||||
G64_view_result = self.G64_view ([test_A064r, test_B064r])
|
||||
test_A064r, test_B064r, rec_A64, rec_B64, rec_AB64 = [ x[0] for x in ([test_A064r, test_B064r] + G64_view_result) ]
|
||||
|
||||
sample64x4 = np.concatenate ([ np.concatenate ( [rec_B64, rec_A64], axis=1 ),
|
||||
np.concatenate ( [test_B064r, rec_AB64], axis=1) ], axis=0 )
|
||||
|
||||
sample64x4 = cv2.resize (sample64x4, (self.resolution, self.resolution) )
|
||||
|
||||
G_view_result = self.G_view([t_src64_0, t_src64m_0, t_src64_1, t_src64m_1, t_src64_2, t_src64m_2, t_dst64_0, t_dst64_1, t_dst64_2 ])
|
||||
|
||||
t_dst_0, t_dst_1, t_dst_2, rec_C_A_t0, rec_C_A_t1, rec_C_A_t2, rec_C_AB_t0, rec_C_AB_t1, rec_C_AB_t2 = [ x[0] for x in ([t_dst_0, t_dst_1, t_dst_2, ] + G_view_result) ]
|
||||
|
||||
c1 = np.concatenate ( (sample64x4, rec_C_A_t0, t_dst_0, rec_C_AB_t0 ), axis=1 )
|
||||
c2 = np.concatenate ( (sample64x4, rec_C_A_t1, t_dst_1, rec_C_AB_t1 ), axis=1 )
|
||||
c3 = np.concatenate ( (sample64x4, rec_C_A_t2, t_dst_2, rec_C_AB_t2 ), axis=1 )
|
||||
|
||||
r = np.concatenate ( [c1,c2,c3], axis=0 )
|
||||
|
||||
return [ ('AVATAR', r ) ]
|
||||
|
||||
def predictor_func (self, prev_imgs=None, img=None, next_imgs=None, dummy_predict=False):
|
||||
if dummy_predict:
|
||||
z = np.zeros ( (1, self.df_res, self.df_res, 3), dtype=np.float32 )
|
||||
self.G_convert ([z,z,z])
|
||||
else:
|
||||
feed = [ prev_imgs[-1][np.newaxis,...], img[np.newaxis,...], next_imgs[0][np.newaxis,...] ]
|
||||
x = self.G_convert (feed)[0]
|
||||
return np.clip ( x[0], 0, 1)
|
||||
|
||||
#override
|
||||
def get_ConverterConfig(self):
|
||||
import converters
|
||||
return self.predictor_func, (self.df_res, self.df_res, 3), converters.ConverterConfigFaceAvatar(temporal_face_count=1)
|
||||
|
||||
@staticmethod
|
||||
def Discriminator(ndf=128):
|
||||
exec (nnlib.import_all(), locals(), globals())
|
||||
|
||||
def func(input):
|
||||
b,h,w,c = K.int_shape(input)
|
||||
|
||||
x = input
|
||||
|
||||
x = Conv2D( ndf, 4, strides=2, padding='valid')( ZeroPadding2D(1)(x) )
|
||||
x = LeakyReLU(0.2)(x)
|
||||
|
||||
x = Conv2D( ndf*2, 4, strides=2, padding='valid')( ZeroPadding2D(1)(x) )
|
||||
x = InstanceNormalization (axis=-1)(x)
|
||||
x = LeakyReLU(0.2)(x)
|
||||
|
||||
x = Conv2D( ndf*4, 4, strides=2, padding='valid')( ZeroPadding2D(1)(x) )
|
||||
x = InstanceNormalization (axis=-1)(x)
|
||||
x = LeakyReLU(0.2)(x)
|
||||
|
||||
x = Conv2D( ndf*8, 4, strides=2, padding='valid')( ZeroPadding2D(1)(x) )
|
||||
x = InstanceNormalization (axis=-1)(x)
|
||||
x = LeakyReLU(0.2)(x)
|
||||
|
||||
return Conv2D( 1, 4, strides=1, padding='valid', activation='sigmoid')( ZeroPadding2D(3)(x) )
|
||||
return func
|
||||
|
||||
@staticmethod
|
||||
def EncFlow():
|
||||
exec (nnlib.import_all(), locals(), globals())
|
||||
|
||||
def downscale (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)( Conv2D(dim, 5, strides=2, padding='same')(x))
|
||||
return func
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
|
||||
def func(input):
|
||||
x, = input
|
||||
b,h,w,c = K.int_shape(x)
|
||||
|
||||
dim_res = w // 16
|
||||
|
||||
x = downscale(64)(x)
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
|
||||
x = Dense(512)(Flatten()(x))
|
||||
x = Dense(dim_res * dim_res * 512)(x)
|
||||
x = Reshape((dim_res, dim_res, 512))(x)
|
||||
x = upscale(512)(x)
|
||||
return x
|
||||
|
||||
return func
|
||||
|
||||
@staticmethod
|
||||
def DecFlow(output_nc=3, **kwargs):
|
||||
exec (nnlib.import_all(), locals(), globals())
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
def to_bgr (output_nc, **kwargs):
|
||||
def func(x):
|
||||
return Conv2D(output_nc, kernel_size=5, strides=1, padding='same', activation='sigmoid')(x)
|
||||
return func
|
||||
|
||||
def func(input):
|
||||
x = input[0]
|
||||
|
||||
x = upscale(512)(x)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
return to_bgr(output_nc) (x)
|
||||
|
||||
return func
|
||||
|
||||
@staticmethod
|
||||
def ResNet(output_nc, ngf=64, n_blocks=6, use_dropout=False):
|
||||
exec (nnlib.import_all(), locals(), globals())
|
||||
|
||||
def func(input):
|
||||
def ResnetBlock(dim, use_dropout=False):
|
||||
def func(input):
|
||||
x = input
|
||||
|
||||
x = Conv2D(dim, 3, strides=1, padding='same')(x)
|
||||
x = InstanceNormalization (axis=-1)(x)
|
||||
x = ReLU()(x)
|
||||
|
||||
if use_dropout:
|
||||
x = Dropout(0.5)(x)
|
||||
|
||||
x = Conv2D(dim, 3, strides=1, padding='same')(x)
|
||||
x = InstanceNormalization (axis=-1)(x)
|
||||
x = ReLU()(x)
|
||||
return Add()([x,input])
|
||||
return func
|
||||
|
||||
x = input
|
||||
|
||||
x = ReLU()(InstanceNormalization (axis=-1)(Conv2D(ngf, 7, strides=1, padding='same')(x)))
|
||||
|
||||
x = ReLU()(InstanceNormalization (axis=-1)(Conv2D(ngf*2, 3, strides=2, padding='same')(x)))
|
||||
x = ReLU()(InstanceNormalization (axis=-1)(Conv2D(ngf*4, 3, strides=2, padding='same')(x)))
|
||||
|
||||
x = ReLU()(InstanceNormalization (axis=-1)(Conv2D(ngf*4, 3, strides=2, padding='same')(x)))
|
||||
|
||||
for i in range(n_blocks):
|
||||
x = ResnetBlock(ngf*4, use_dropout=use_dropout)(x)
|
||||
|
||||
x = ReLU()(InstanceNormalization (axis=-1)(Conv2DTranspose(ngf*4, 3, strides=2, padding='same')(x)))
|
||||
|
||||
x = ReLU()(InstanceNormalization (axis=-1)(Conv2DTranspose(ngf*2, 3, strides=2, padding='same')(x)))
|
||||
x = ReLU()(InstanceNormalization (axis=-1)(Conv2DTranspose(ngf , 3, strides=2, padding='same')(x)))
|
||||
|
||||
x = Conv2D(output_nc, 7, strides=1, activation='sigmoid', padding='same')(x)
|
||||
|
||||
return x
|
||||
|
||||
return func
|
||||
|
||||
Model = AVATARModel
|
||||
|
|
@ -1 +0,0 @@
|
|||
from .Model import Model
|
||||
|
|
@ -1,103 +0,0 @@
|
|||
import numpy as np
|
||||
|
||||
from nnlib import nnlib, TernausNet
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samplelib import *
|
||||
from interact import interact as io
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs,
|
||||
ask_enable_autobackup=False,
|
||||
ask_write_preview_history=False,
|
||||
ask_target_iter=False,
|
||||
ask_random_flip=False)
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
default_face_type = 'f'
|
||||
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):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {1.5:4, 11:48} )
|
||||
|
||||
self.resolution = 256
|
||||
self.face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
|
||||
|
||||
model_name = 'FANSeg'
|
||||
self.fan_seg = TernausNet(model_name, self.resolution,
|
||||
FaceType.toString(self.face_type),
|
||||
load_weights=not self.is_first_run(),
|
||||
weights_file_root=self.get_model_root_path(),
|
||||
training=True)
|
||||
|
||||
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
|
||||
|
||||
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=[ { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR_RANDOM_HSV_SHIFT), 'resolution' : self.resolution, 'motion_blur':(25, 5), 'gaussian_blur':(25,5), 'border_replicate':False},
|
||||
{ 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_M), 'resolution': self.resolution },
|
||||
]),
|
||||
|
||||
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=[ { 'types': (t.IMG_TRANSFORMED , face_type, t.MODE_BGR_RANDOM_HSV_SHIFT), 'resolution' : self.resolution},
|
||||
])
|
||||
])
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.fan_seg.save_weights()
|
||||
|
||||
#override
|
||||
def onTrainOneIter(self, generators_samples, generators_list):
|
||||
target_src, target_src_mask = generators_samples[0]
|
||||
|
||||
loss = self.fan_seg.train( target_src, target_src_mask )
|
||||
|
||||
return ( ('loss', loss), )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
test_A = sample[0][0][0:4] #first 4 samples
|
||||
test_Am = sample[0][1][0:4] #first 4 samples
|
||||
test_B = sample[1][0][0:4] #first 4 samples
|
||||
|
||||
|
||||
mAA = self.fan_seg.extract(test_A)
|
||||
mBB = self.fan_seg.extract(test_B)
|
||||
|
||||
test_Am = np.repeat ( test_Am, (3,), -1)
|
||||
mAA = np.repeat ( mAA, (3,), -1)
|
||||
mBB = np.repeat ( mBB, (3,), -1)
|
||||
|
||||
st = []
|
||||
for i in range(0, len(test_A)):
|
||||
st.append ( np.concatenate ( (
|
||||
test_A[i,:,:,0:3],
|
||||
test_Am[i],
|
||||
mAA[i],
|
||||
test_A[i,:,:,0:3]*mAA[i],
|
||||
), axis=1) )
|
||||
|
||||
st2 = []
|
||||
for i in range(0, len(test_B)):
|
||||
st2.append ( np.concatenate ( (
|
||||
test_B[i,:,:,0:3],
|
||||
mBB[i],
|
||||
test_B[i,:,:,0:3]*mBB[i],
|
||||
), axis=1) )
|
||||
|
||||
return [ ('training data', np.concatenate ( st, axis=0 ) ),
|
||||
('evaluating data', np.concatenate ( st2, axis=0 ) ),
|
||||
]
|
||||
|
|
@ -1 +0,0 @@
|
|||
from .Model import Model
|
||||
|
|
@ -1,178 +0,0 @@
|
|||
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_random_flip=False)
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
|
||||
default_resolution = 64
|
||||
if is_first_run:
|
||||
self.options['resolution'] = io.input_int(f"Resolution ( 64,96,128,224 ?:help skip:{default_resolution}) : ", default_resolution, [64,96,128,224])
|
||||
else:
|
||||
self.options['resolution'] = self.options.get('resolution', default_resolution)
|
||||
|
||||
default_face_type = 'mf'
|
||||
if is_first_run:
|
||||
self.options['face_type'] = io.input_str (f"Half or Full face? (h/mf/f, ?:help skip:{default_face_type}) : ", default_face_type, ['h','mf','f'], help_message="").lower()
|
||||
else:
|
||||
self.options['face_type'] = self.options.get('face_type', default_face_type)
|
||||
|
||||
if (is_first_run or ask_override) and 'tensorflow' in self.device_config.backend:
|
||||
def_optimizer_mode = self.options.get('optimizer_mode', 1)
|
||||
self.options['optimizer_mode'] = io.input_int ("Optimizer mode? ( 1,2,3 ?:help skip:%d) : " % (def_optimizer_mode), def_optimizer_mode, help_message="1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power.")
|
||||
else:
|
||||
self.options['optimizer_mode'] = self.options.get('optimizer_mode', 1)
|
||||
|
||||
#override
|
||||
def onInitialize(self, batch_size=-1, **in_options):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
self.set_vram_batch_requirements({4:16,11:24})
|
||||
|
||||
resolution = self.options['resolution']
|
||||
face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
|
||||
person_id_max_count = SampleGeneratorFacePerson.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_blks=2,
|
||||
dis_nf=64,
|
||||
dis_res_blks=8,#10
|
||||
num_classes=person_id_max_count,
|
||||
subpixel_decoder=True,
|
||||
initialize_weights=self.is_first_run(),
|
||||
is_training=self.is_training_mode,
|
||||
tf_cpu_mode=self.options['optimizer_mode']-1
|
||||
)
|
||||
|
||||
if not self.is_first_run():
|
||||
self.load_weights_safe(self.model.get_model_filename_list())
|
||||
|
||||
if self.is_training_mode:
|
||||
t = SampleProcessor.Types
|
||||
if self.options['face_type'] == 'h':
|
||||
face_type = t.FACE_TYPE_HALF
|
||||
elif self.options['face_type'] == 'mf':
|
||||
face_type = t.FACE_TYPE_MID_FULL
|
||||
elif self.options['face_type'] == 'f':
|
||||
face_type = t.FACE_TYPE_FULL
|
||||
|
||||
output_sample_types=[ {'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh':True} ]
|
||||
output_sample_types1=[ {'types': (t.IMG_SOURCE, face_type, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh':True} ]
|
||||
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFacePerson(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True, rotation_range=[0,0] ),
|
||||
output_sample_types=output_sample_types, person_id_mode=1, ),
|
||||
|
||||
SampleGeneratorFacePerson(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True, rotation_range=[0,0] ),
|
||||
output_sample_types=output_sample_types, person_id_mode=1, ),
|
||||
|
||||
SampleGeneratorFacePerson(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True, rotation_range=[0,0]),
|
||||
output_sample_types=output_sample_types1, person_id_mode=1, ),
|
||||
|
||||
SampleGeneratorFacePerson(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True, rotation_range=[0,0]),
|
||||
output_sample_types=output_sample_types1, person_id_mode=1, ),
|
||||
])
|
||||
|
||||
#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 +0,0 @@
|
|||
from .Model import Model
|
||||
|
|
@ -1,120 +0,0 @@
|
|||
import numpy as np
|
||||
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from facelib import PoseEstimator
|
||||
from samplelib import *
|
||||
from interact import interact as io
|
||||
import imagelib
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs,
|
||||
ask_enable_autobackup=False,
|
||||
ask_write_preview_history=False,
|
||||
ask_target_iter=False,
|
||||
ask_random_flip=False)
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
yn_str = {True:'y',False:'n'}
|
||||
|
||||
default_face_type = 'f'
|
||||
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="Half face has better resolution, but covers less area of cheeks.").lower()
|
||||
else:
|
||||
self.options['face_type'] = self.options.get('face_type', default_face_type)
|
||||
|
||||
def_train_bgr = self.options.get('train_bgr', True)
|
||||
if is_first_run or ask_override:
|
||||
self.options['train_bgr'] = io.input_bool ("Train bgr? (y/n, ?:help skip: %s) : " % (yn_str[def_train_bgr]), def_train_bgr)
|
||||
else:
|
||||
self.options['train_bgr'] = self.options.get('train_bgr', def_train_bgr)
|
||||
|
||||
#override
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {4:64} )
|
||||
|
||||
self.resolution = 128
|
||||
self.face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
|
||||
|
||||
|
||||
self.pose_est = PoseEstimator(self.resolution,
|
||||
FaceType.toString(self.face_type),
|
||||
load_weights=not self.is_first_run(),
|
||||
weights_file_root=self.get_model_root_path(),
|
||||
training=True)
|
||||
|
||||
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
|
||||
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, generators_count=4,
|
||||
sample_process_options=SampleProcessor.Options( rotation_range=[0,0] ), #random_flip=True,
|
||||
output_sample_types=[ {'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution':self.resolution, 'motion_blur':(25, 1) },
|
||||
{'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution':self.resolution },
|
||||
{'types': (t.IMG_PITCH_YAW_ROLL,)}
|
||||
]),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, generators_count=4,
|
||||
sample_process_options=SampleProcessor.Options( rotation_range=[0,0] ), #random_flip=True,
|
||||
output_sample_types=[ {'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':self.resolution },
|
||||
{'types': (t.IMG_PITCH_YAW_ROLL,)}
|
||||
])
|
||||
])
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.pose_est.save_weights()
|
||||
|
||||
#override
|
||||
def onTrainOneIter(self, generators_samples, generators_list):
|
||||
target_srcw, target_src, pitch_yaw_roll = generators_samples[0]
|
||||
|
||||
bgr_loss, pyr_loss = self.pose_est.train_on_batch( target_srcw, target_src, pitch_yaw_roll, skip_bgr_train=not self.options['train_bgr'] )
|
||||
|
||||
return ( ('bgr_loss', bgr_loss), ('pyr_loss', pyr_loss), )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, generators_samples):
|
||||
test_src = generators_samples[0][1][0:4] #first 4 samples
|
||||
test_pyr_src = generators_samples[0][2][0:4]
|
||||
test_dst = generators_samples[1][0][0:4]
|
||||
test_pyr_dst = generators_samples[1][1][0:4]
|
||||
|
||||
h,w,c = self.resolution,self.resolution,3
|
||||
h_line = 13
|
||||
|
||||
result = []
|
||||
for name, img, pyr in [ ['training data', test_src, test_pyr_src], \
|
||||
['evaluating data',test_dst, test_pyr_dst] ]:
|
||||
bgr_pred, pyr_pred = self.pose_est.extract(img)
|
||||
|
||||
hor_imgs = []
|
||||
for i in range(len(img)):
|
||||
img_info = np.ones ( (h,w,c) ) * 0.1
|
||||
|
||||
i_pyr = pyr[i]
|
||||
i_pyr_pred = pyr_pred[i]
|
||||
lines = ["%.4f %.4f %.4f" % (i_pyr[0],i_pyr[1],i_pyr[2]),
|
||||
"%.4f %.4f %.4f" % (i_pyr_pred[0],i_pyr_pred[1],i_pyr_pred[2]) ]
|
||||
|
||||
lines_count = len(lines)
|
||||
for ln in range(lines_count):
|
||||
img_info[ ln*h_line:(ln+1)*h_line, 0:w] += \
|
||||
imagelib.get_text_image ( (h_line,w,c), lines[ln], color=[0.8]*c )
|
||||
|
||||
hor_imgs.append ( np.concatenate ( (
|
||||
img[i,:,:,0:3],
|
||||
bgr_pred[i],
|
||||
img_info
|
||||
), axis=1) )
|
||||
|
||||
|
||||
result += [ (name, np.concatenate (hor_imgs, axis=0)) ]
|
||||
|
||||
return result
|
||||
|
|
@ -1 +0,0 @@
|
|||
from .Model import Model
|
||||
|
|
@ -1,169 +0,0 @@
|
|||
import numpy as np
|
||||
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samplelib import *
|
||||
from interact import interact as io
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run or ask_override:
|
||||
def_pixel_loss = self.options.get('pixel_loss', False)
|
||||
self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.")
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
||||
#override
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {4.5:4} )
|
||||
|
||||
ae_input_layer = Input(shape=(128, 128, 3))
|
||||
mask_layer = Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder, self.decoder_src, self.decoder_dst = self.Build(ae_input_layer)
|
||||
|
||||
if not self.is_first_run():
|
||||
weights_to_load = [ [self.encoder , 'encoder.h5'],
|
||||
[self.decoder_src, 'decoder_src.h5'],
|
||||
[self.decoder_dst, 'decoder_dst.h5']
|
||||
]
|
||||
self.load_weights_safe(weights_to_load)
|
||||
|
||||
rec_src = self.decoder_src(self.encoder(ae_input_layer))
|
||||
rec_dst = self.decoder_dst(self.encoder(ae_input_layer))
|
||||
self.autoencoder_src = Model([ae_input_layer,mask_layer], rec_src)
|
||||
self.autoencoder_dst = Model([ae_input_layer,mask_layer], rec_dst)
|
||||
|
||||
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'] )
|
||||
|
||||
self.convert = K.function([ae_input_layer], rec_src)
|
||||
|
||||
if self.is_training_mode:
|
||||
t = SampleProcessor.Types
|
||||
output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_BGR), 'resolution':128},
|
||||
{ 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_BGR), 'resolution':128},
|
||||
{ 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_M), 'resolution':128} ]
|
||||
|
||||
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=self.random_flip, scale_range=np.array([-0.05, 0.05]) ),
|
||||
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=self.random_flip),
|
||||
output_sample_types=output_sample_types)
|
||||
])
|
||||
|
||||
#override
|
||||
def get_model_filename_list(self):
|
||||
return [[self.encoder, 'encoder.h5'],
|
||||
[self.decoder_src, 'decoder_src.h5'],
|
||||
[self.decoder_dst, 'decoder_dst.h5']]
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( self.get_model_filename_list() )
|
||||
|
||||
#override
|
||||
def onTrainOneIter(self, sample, generators_list):
|
||||
warped_src, target_src, target_src_mask = sample[0]
|
||||
warped_dst, target_dst, target_dst_mask = sample[1]
|
||||
|
||||
loss_src = self.autoencoder_src.train_on_batch( [warped_src, target_src_mask], [target_src, target_src_mask] )
|
||||
loss_dst = self.autoencoder_dst.train_on_batch( [warped_dst, target_dst_mask], [target_dst, target_dst_mask] )
|
||||
|
||||
return ( ('loss_src', loss_src[0]), ('loss_dst', loss_dst[0]) )
|
||||
|
||||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
test_A = sample[0][1][0:4] #first 4 samples
|
||||
test_A_m = sample[0][2][0:4] #first 4 samples
|
||||
test_B = sample[1][1][0:4]
|
||||
test_B_m = sample[1][2][0:4]
|
||||
|
||||
AA, mAA = self.autoencoder_src.predict([test_A, test_A_m])
|
||||
AB, mAB = self.autoencoder_src.predict([test_B, test_B_m])
|
||||
BB, mBB = self.autoencoder_dst.predict([test_B, test_B_m])
|
||||
|
||||
mAA = np.repeat ( mAA, (3,), -1)
|
||||
mAB = np.repeat ( mAB, (3,), -1)
|
||||
mBB = np.repeat ( mBB, (3,), -1)
|
||||
|
||||
st = []
|
||||
for i in range(0, len(test_A)):
|
||||
st.append ( np.concatenate ( (
|
||||
test_A[i,:,:,0:3],
|
||||
AA[i],
|
||||
#mAA[i],
|
||||
test_B[i,:,:,0:3],
|
||||
BB[i],
|
||||
#mBB[i],
|
||||
AB[i],
|
||||
#mAB[i]
|
||||
), axis=1) )
|
||||
|
||||
return [ ('DF', np.concatenate ( st, axis=0 ) ) ]
|
||||
|
||||
def predictor_func (self, face=None, dummy_predict=False):
|
||||
if dummy_predict:
|
||||
self.convert ([ np.zeros ( (1, 128, 128, 3), dtype=np.float32 ) ])
|
||||
else:
|
||||
x, mx = self.convert ( [ face[np.newaxis,...] ] )
|
||||
return x[0], mx[0][...,0]
|
||||
|
||||
#override
|
||||
def get_ConverterConfig(self):
|
||||
import converters
|
||||
return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.FULL, default_mode='seamless')
|
||||
|
||||
def Build(self, input_layer):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
def downscale (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
|
||||
return func
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
def Encoder(input_layer):
|
||||
x = input_layer
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(1024)(x)
|
||||
|
||||
x = Dense(512)(Flatten()(x))
|
||||
x = Dense(8 * 8 * 512)(x)
|
||||
x = Reshape((8, 8, 512))(x)
|
||||
x = upscale(512)(x)
|
||||
|
||||
return Model(input_layer, x)
|
||||
|
||||
def Decoder():
|
||||
input_ = Input(shape=(16, 16, 512))
|
||||
x = input_
|
||||
x = upscale(512)(x)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(512)(y)
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
|
||||
return Model(input_, [x,y])
|
||||
|
||||
return Encoder(input_layer), Decoder(), Decoder()
|
||||
|
|
@ -1 +0,0 @@
|
|||
from .Model import Model
|
||||
|
|
@ -1,203 +0,0 @@
|
|||
import numpy as np
|
||||
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samplelib import *
|
||||
from interact import interact as io
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run:
|
||||
self.options['lighter_ae'] = io.input_bool ("Use lightweight autoencoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight autoencoder is faster, requires less VRAM, sacrificing overall quality. If your GPU VRAM <= 4, you should to choose this option.")
|
||||
else:
|
||||
default_lighter_ae = self.options.get('created_vram_gb', 99) <= 4 #temporally support old models, deprecate in future
|
||||
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:
|
||||
def_pixel_loss = self.options.get('pixel_loss', False)
|
||||
self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.")
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
||||
#override
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {2.5:4} )
|
||||
|
||||
bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build( self.options['lighter_ae'] )
|
||||
if not self.is_first_run():
|
||||
weights_to_load = [ [self.encoder , 'encoder.h5'],
|
||||
[self.decoder_src, 'decoder_src.h5'],
|
||||
[self.decoder_dst, 'decoder_dst.h5']
|
||||
]
|
||||
self.load_weights_safe(weights_to_load)
|
||||
|
||||
input_src_bgr = Input(bgr_shape)
|
||||
input_src_mask = Input(mask_shape)
|
||||
input_dst_bgr = Input(bgr_shape)
|
||||
input_dst_mask = Input(mask_shape)
|
||||
|
||||
rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) )
|
||||
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=[ 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])
|
||||
|
||||
if self.is_training_mode:
|
||||
t = SampleProcessor.Types
|
||||
output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_BGR), 'resolution':128},
|
||||
{ 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_BGR), 'resolution':128},
|
||||
{ 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_M), 'resolution':128} ]
|
||||
|
||||
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=self.random_flip, scale_range=np.array([-0.05, 0.05]) ),
|
||||
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=self.random_flip),
|
||||
output_sample_types=output_sample_types )
|
||||
])
|
||||
|
||||
#override
|
||||
def get_model_filename_list(self):
|
||||
return [[self.encoder, 'encoder.h5'],
|
||||
[self.decoder_src, 'decoder_src.h5'],
|
||||
[self.decoder_dst, 'decoder_dst.h5']]
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( self.get_model_filename_list() )
|
||||
|
||||
#override
|
||||
def onTrainOneIter(self, sample, generators_list):
|
||||
warped_src, target_src, target_src_mask = sample[0]
|
||||
warped_dst, target_dst, target_dst_mask = sample[1]
|
||||
|
||||
total, loss_src_bgr, loss_src_mask, loss_dst_bgr, loss_dst_mask = self.ae.train_on_batch( [warped_src, target_src_mask, warped_dst, target_dst_mask], [target_src, target_src_mask, target_dst, target_dst_mask] )
|
||||
|
||||
return ( ('loss_src', loss_src_bgr), ('loss_dst', loss_dst_bgr) )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
test_A = sample[0][1][0:4] #first 4 samples
|
||||
test_A_m = sample[0][2][0:4] #first 4 samples
|
||||
test_B = sample[1][1][0:4]
|
||||
test_B_m = sample[1][2][0:4]
|
||||
|
||||
AA, mAA = self.src_view([test_A])
|
||||
AB, mAB = self.src_view([test_B])
|
||||
BB, mBB = self.dst_view([test_B])
|
||||
|
||||
mAA = np.repeat ( mAA, (3,), -1)
|
||||
mAB = np.repeat ( mAB, (3,), -1)
|
||||
mBB = np.repeat ( mBB, (3,), -1)
|
||||
|
||||
st = []
|
||||
for i in range(0, len(test_A)):
|
||||
st.append ( np.concatenate ( (
|
||||
test_A[i,:,:,0:3],
|
||||
AA[i],
|
||||
#mAA[i],
|
||||
test_B[i,:,:,0:3],
|
||||
BB[i],
|
||||
#mBB[i],
|
||||
AB[i],
|
||||
#mAB[i]
|
||||
), axis=1) )
|
||||
|
||||
return [ ('H128', np.concatenate ( st, axis=0 ) ) ]
|
||||
|
||||
def predictor_func (self, face=None, dummy_predict=False):
|
||||
if dummy_predict:
|
||||
self.src_view ([ np.zeros ( (1, 128, 128, 3), dtype=np.float32 ) ])
|
||||
else:
|
||||
x, mx = self.src_view ( [ face[np.newaxis,...] ] )
|
||||
return x[0], mx[0][...,0]
|
||||
|
||||
#override
|
||||
def get_ConverterConfig(self):
|
||||
import converters
|
||||
return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.HALF, default_mode='seamless')
|
||||
|
||||
def Build(self, lighter_ae):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
bgr_shape = (128, 128, 3)
|
||||
mask_shape = (128, 128, 1)
|
||||
|
||||
def downscale (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
|
||||
return func
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
def Encoder(input_shape):
|
||||
input_layer = Input(input_shape)
|
||||
x = input_layer
|
||||
if not lighter_ae:
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(1024)(x)
|
||||
x = Dense(512)(Flatten()(x))
|
||||
x = Dense(8 * 8 * 512)(x)
|
||||
x = Reshape((8, 8, 512))(x)
|
||||
x = upscale(512)(x)
|
||||
else:
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(1024)(x)
|
||||
x = Dense(256)(Flatten()(x))
|
||||
x = Dense(8 * 8 * 256)(x)
|
||||
x = Reshape((8, 8, 256))(x)
|
||||
x = upscale(256)(x)
|
||||
|
||||
return Model(input_layer, x)
|
||||
|
||||
def Decoder():
|
||||
if not lighter_ae:
|
||||
input_ = Input(shape=(16, 16, 512))
|
||||
x = input_
|
||||
x = upscale(512)(x)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(512)(y)
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
else:
|
||||
input_ = Input(shape=(16, 16, 256))
|
||||
x = input_
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
x = upscale(64)(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
y = upscale(64)(y)
|
||||
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
|
||||
|
||||
return Model(input_, [x,y])
|
||||
|
||||
return bgr_shape, mask_shape, Encoder(bgr_shape), Decoder(), Decoder()
|
||||
|
|
@ -1 +0,0 @@
|
|||
from .Model import Model
|
||||
|
|
@ -1,200 +0,0 @@
|
|||
import numpy as np
|
||||
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samplelib import *
|
||||
from interact import interact as io
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run:
|
||||
self.options['lighter_ae'] = io.input_bool ("Use lightweight autoencoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight autoencoder is faster, requires less VRAM, sacrificing overall quality. If your GPU VRAM <= 4, you should to choose this option.")
|
||||
else:
|
||||
default_lighter_ae = self.options.get('created_vram_gb', 99) <= 4 #temporally support old models, deprecate in future
|
||||
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:
|
||||
def_pixel_loss = self.options.get('pixel_loss', False)
|
||||
self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.")
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
||||
#override
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {1.5:4} )
|
||||
|
||||
|
||||
bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build(self.options['lighter_ae'])
|
||||
|
||||
if not self.is_first_run():
|
||||
weights_to_load = [ [self.encoder , 'encoder.h5'],
|
||||
[self.decoder_src, 'decoder_src.h5'],
|
||||
[self.decoder_dst, 'decoder_dst.h5']
|
||||
]
|
||||
self.load_weights_safe(weights_to_load)
|
||||
|
||||
input_src_bgr = Input(bgr_shape)
|
||||
input_src_mask = Input(mask_shape)
|
||||
input_dst_bgr = Input(bgr_shape)
|
||||
input_dst_mask = Input(mask_shape)
|
||||
|
||||
rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) )
|
||||
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=[ 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])
|
||||
|
||||
if self.is_training_mode:
|
||||
t = SampleProcessor.Types
|
||||
output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_BGR), 'resolution':64},
|
||||
{ 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_BGR), 'resolution':64},
|
||||
{ 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_M), 'resolution':64} ]
|
||||
|
||||
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=self.random_flip, scale_range=np.array([-0.05, 0.05]) ),
|
||||
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=self.random_flip),
|
||||
output_sample_types=output_sample_types)
|
||||
])
|
||||
|
||||
#override
|
||||
def get_model_filename_list(self):
|
||||
return [[self.encoder, 'encoder.h5'],
|
||||
[self.decoder_src, 'decoder_src.h5'],
|
||||
[self.decoder_dst, 'decoder_dst.h5']]
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( self.get_model_filename_list() )
|
||||
|
||||
#override
|
||||
def onTrainOneIter(self, sample, generators_list):
|
||||
warped_src, target_src, target_src_full_mask = sample[0]
|
||||
warped_dst, target_dst, target_dst_full_mask = sample[1]
|
||||
|
||||
total, loss_src_bgr, loss_src_mask, loss_dst_bgr, loss_dst_mask = self.ae.train_on_batch( [warped_src, target_src_full_mask, warped_dst, target_dst_full_mask], [target_src, target_src_full_mask, target_dst, target_dst_full_mask] )
|
||||
|
||||
return ( ('loss_src', loss_src_bgr), ('loss_dst', loss_dst_bgr) )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
test_A = sample[0][1][0:4] #first 4 samples
|
||||
test_A_m = sample[0][2][0:4]
|
||||
test_B = sample[1][1][0:4]
|
||||
test_B_m = sample[1][2][0:4]
|
||||
|
||||
AA, mAA = self.src_view([test_A])
|
||||
AB, mAB = self.src_view([test_B])
|
||||
BB, mBB = self.dst_view([test_B])
|
||||
|
||||
mAA = np.repeat ( mAA, (3,), -1)
|
||||
mAB = np.repeat ( mAB, (3,), -1)
|
||||
mBB = np.repeat ( mBB, (3,), -1)
|
||||
|
||||
st = []
|
||||
for i in range(0, len(test_A)):
|
||||
st.append ( np.concatenate ( (
|
||||
test_A[i,:,:,0:3],
|
||||
AA[i],
|
||||
#mAA[i],
|
||||
test_B[i,:,:,0:3],
|
||||
BB[i],
|
||||
#mBB[i],
|
||||
AB[i],
|
||||
#mAB[i]
|
||||
), axis=1) )
|
||||
|
||||
return [ ('H64', np.concatenate ( st, axis=0 ) ) ]
|
||||
|
||||
def predictor_func (self, face=None, dummy_predict=False):
|
||||
if dummy_predict:
|
||||
self.src_view ([ np.zeros ( (1, 64, 64, 3), dtype=np.float32 ) ])
|
||||
else:
|
||||
x, mx = self.src_view ( [ face[np.newaxis,...] ] )
|
||||
return x[0], mx[0][...,0]
|
||||
|
||||
#override
|
||||
def get_ConverterConfig(self):
|
||||
import converters
|
||||
return self.predictor_func, (64,64,3), converters.ConverterConfigMasked(face_type=FaceType.HALF, default_mode='seamless')
|
||||
|
||||
def Build(self, lighter_ae):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
bgr_shape = (64, 64, 3)
|
||||
mask_shape = (64, 64, 1)
|
||||
|
||||
def downscale (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
|
||||
return func
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
def Encoder(input_shape):
|
||||
input_layer = Input(input_shape)
|
||||
x = input_layer
|
||||
if not lighter_ae:
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(1024)(x)
|
||||
x = Dense(1024)(Flatten()(x))
|
||||
x = Dense(4 * 4 * 1024)(x)
|
||||
x = Reshape((4, 4, 1024))(x)
|
||||
x = upscale(512)(x)
|
||||
else:
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(768)(x)
|
||||
x = Dense(512)(Flatten()(x))
|
||||
x = Dense(4 * 4 * 512)(x)
|
||||
x = Reshape((4, 4, 512))(x)
|
||||
x = upscale(256)(x)
|
||||
return Model(input_layer, x)
|
||||
|
||||
def Decoder():
|
||||
if not lighter_ae:
|
||||
input_ = Input(shape=(8, 8, 512))
|
||||
x = input_
|
||||
|
||||
x = upscale(512)(x)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
|
||||
else:
|
||||
input_ = Input(shape=(8, 8, 256))
|
||||
|
||||
x = input_
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
x = upscale(64)(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
y = upscale(64)(y)
|
||||
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
|
||||
return Model(input_, [x,y])
|
||||
|
||||
return bgr_shape, mask_shape, Encoder(bgr_shape), Decoder(), Decoder()
|
||||
|
|
@ -1 +0,0 @@
|
|||
from .Model import Model
|
||||
|
|
@ -1,178 +0,0 @@
|
|||
import numpy as np
|
||||
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samplelib import *
|
||||
from interact import interact as io
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run or ask_override:
|
||||
def_pixel_loss = self.options.get('pixel_loss', False)
|
||||
self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.")
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
||||
#override
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {4.5:4} )
|
||||
|
||||
ae_input_layer = Input(shape=(128, 128, 3))
|
||||
mask_layer = Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder, self.decoder, self.inter_B, self.inter_AB = self.Build(ae_input_layer)
|
||||
|
||||
if not self.is_first_run():
|
||||
weights_to_load = [ [self.encoder, 'encoder.h5'],
|
||||
[self.decoder, 'decoder.h5'],
|
||||
[self.inter_B, 'inter_B.h5'],
|
||||
[self.inter_AB, 'inter_AB.h5']
|
||||
]
|
||||
self.load_weights_safe(weights_to_load)
|
||||
|
||||
code = self.encoder(ae_input_layer)
|
||||
AB = self.inter_AB(code)
|
||||
B = self.inter_B(code)
|
||||
rec_src = self.decoder(Concatenate()([AB, AB]))
|
||||
rec_dst = self.decoder(Concatenate()([B, AB]))
|
||||
self.autoencoder_src = Model([ae_input_layer,mask_layer], rec_src )
|
||||
self.autoencoder_dst = Model([ae_input_layer,mask_layer], rec_dst )
|
||||
|
||||
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'] )
|
||||
|
||||
self.convert = K.function([ae_input_layer],rec_src)
|
||||
|
||||
|
||||
if self.is_training_mode:
|
||||
t = SampleProcessor.Types
|
||||
output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_BGR), 'resolution':128},
|
||||
{ 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_BGR), 'resolution':128},
|
||||
{ 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_M), 'resolution':128} ]
|
||||
|
||||
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=self.random_flip, scale_range=np.array([-0.05, 0.05]) ),
|
||||
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=self.random_flip),
|
||||
output_sample_types=output_sample_types)
|
||||
])
|
||||
|
||||
#override
|
||||
def get_model_filename_list(self):
|
||||
return [[self.encoder, 'encoder.h5'],
|
||||
[self.decoder, 'decoder.h5'],
|
||||
[self.inter_B, 'inter_B.h5'],
|
||||
[self.inter_AB, 'inter_AB.h5']]
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( self.get_model_filename_list() )
|
||||
|
||||
#override
|
||||
def onTrainOneIter(self, sample, generators_list):
|
||||
warped_src, target_src, target_src_mask = sample[0]
|
||||
warped_dst, target_dst, target_dst_mask = sample[1]
|
||||
|
||||
loss_src = self.autoencoder_src.train_on_batch( [warped_src, target_src_mask], [target_src, target_src_mask] )
|
||||
loss_dst = self.autoencoder_dst.train_on_batch( [warped_dst, target_dst_mask], [target_dst, target_dst_mask] )
|
||||
|
||||
return ( ('loss_src', loss_src[0]), ('loss_dst', loss_dst[0]) )
|
||||
|
||||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
test_A = sample[0][1][0:4] #first 4 samples
|
||||
test_A_m = sample[0][2][0:4] #first 4 samples
|
||||
test_B = sample[1][1][0:4]
|
||||
test_B_m = sample[1][2][0:4]
|
||||
|
||||
AA, mAA = self.autoencoder_src.predict([test_A, test_A_m])
|
||||
AB, mAB = self.autoencoder_src.predict([test_B, test_B_m])
|
||||
BB, mBB = self.autoencoder_dst.predict([test_B, test_B_m])
|
||||
|
||||
mAA = np.repeat ( mAA, (3,), -1)
|
||||
mAB = np.repeat ( mAB, (3,), -1)
|
||||
mBB = np.repeat ( mBB, (3,), -1)
|
||||
|
||||
st = []
|
||||
for i in range(0, len(test_A)):
|
||||
st.append ( np.concatenate ( (
|
||||
test_A[i,:,:,0:3],
|
||||
AA[i],
|
||||
#mAA[i],
|
||||
test_B[i,:,:,0:3],
|
||||
BB[i],
|
||||
#mBB[i],
|
||||
AB[i],
|
||||
#mAB[i]
|
||||
), axis=1) )
|
||||
|
||||
return [ ('LIAEF128', np.concatenate ( st, axis=0 ) ) ]
|
||||
|
||||
def predictor_func (self, face=None, dummy_predict=False):
|
||||
if dummy_predict:
|
||||
self.convert ([ np.zeros ( (1, 128, 128, 3), dtype=np.float32 ) ])
|
||||
else:
|
||||
x, mx = self.convert ( [ face[np.newaxis,...] ] )
|
||||
return x[0], mx[0][...,0]
|
||||
|
||||
#override
|
||||
def get_ConverterConfig(self):
|
||||
import converters
|
||||
return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.FULL, default_mode='seamless')
|
||||
|
||||
def Build(self, input_layer):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
def downscale (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
|
||||
return func
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
def Encoder():
|
||||
x = input_layer
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(1024)(x)
|
||||
x = Flatten()(x)
|
||||
return Model(input_layer, x)
|
||||
|
||||
def Intermediate():
|
||||
input_layer = Input(shape=(None, 8 * 8 * 1024))
|
||||
x = input_layer
|
||||
x = Dense(256)(x)
|
||||
x = Dense(8 * 8 * 512)(x)
|
||||
x = Reshape((8, 8, 512))(x)
|
||||
x = upscale(512)(x)
|
||||
return Model(input_layer, x)
|
||||
|
||||
def Decoder():
|
||||
input_ = Input(shape=(16, 16, 1024))
|
||||
x = input_
|
||||
x = upscale(512)(x)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(512)(y)
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid' )(y)
|
||||
|
||||
return Model(input_, [x,y])
|
||||
|
||||
return Encoder(), Decoder(), Intermediate(), Intermediate()
|
||||
|
|
@ -1 +0,0 @@
|
|||
from .Model import Model
|
||||
|
|
@ -1,261 +1,503 @@
|
|||
import multiprocessing
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
|
||||
import mathlib
|
||||
from core import mathlib
|
||||
from core.interact import interact as io
|
||||
from core.leras import nn
|
||||
from facelib import FaceType
|
||||
from interact import interact as io
|
||||
from models import ModelBase
|
||||
from nnlib import nnlib
|
||||
from samplelib import *
|
||||
|
||||
|
||||
class Quick96Model(ModelBase):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs,
|
||||
ask_enable_autobackup=False,
|
||||
ask_write_preview_history=False,
|
||||
ask_target_iter=True,
|
||||
ask_batch_size=False,
|
||||
ask_random_flip=False)
|
||||
|
||||
class QModel(ModelBase):
|
||||
#override
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements({1.5:2,2:4})
|
||||
def on_initialize(self):
|
||||
nn.initialize()
|
||||
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 = tf.nn.space_to_depth(x, 2)
|
||||
|
||||
if self.use_activator:
|
||||
x = tf.nn.leaky_relu(x, 0.2)
|
||||
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 = tf.nn.leaky_relu(x, 0.2)
|
||||
x = tf.nn.depth_to_space(x, 2)
|
||||
return x
|
||||
|
||||
class UpdownResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, inner_ch, kernel_size=3 ):
|
||||
self.up = Upscale (ch, inner_ch, kernel_size=kernel_size)
|
||||
self.res = ResidualBlock (inner_ch, kernel_size=kernel_size)
|
||||
self.down = Downscale (inner_ch, ch, kernel_size=kernel_size, use_activator=False)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.up(inp)
|
||||
x = upx = self.res(x)
|
||||
x = self.down(x)
|
||||
x = x + inp
|
||||
x = tf.nn.leaky_relu(x, 0.2)
|
||||
return x, upx
|
||||
|
||||
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 = tf.nn.leaky_relu(x, 0.2)
|
||||
x = self.conv2(x)
|
||||
x = inp + x
|
||||
x = tf.nn.leaky_relu(x, 0.2)
|
||||
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=2, 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 = tf.reshape (x, (-1, 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 = UpdownResidualBlock(d_ch*4, d_ch*2)
|
||||
self.upscale2 = Upscale(d_ch*4, d_ch*2)
|
||||
self.res2 = UpdownResidualBlock(d_ch*2, d_ch)
|
||||
self.upscale3 = Upscale(d_ch*2, d_ch*1)
|
||||
self.res3 = UpdownResidualBlock(d_ch, d_ch//2)
|
||||
|
||||
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, upx = self.res1(x)
|
||||
|
||||
x = self.upscale2(x)
|
||||
x = tf.nn.leaky_relu(x + upx, 0.2)
|
||||
x, upx = self.res2(x)
|
||||
|
||||
x = self.upscale3(x)
|
||||
x = tf.nn.leaky_relu(x + upx, 0.2)
|
||||
x, upx = self.res3(x)
|
||||
|
||||
"""
|
||||
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
|
||||
ae_dims = 128
|
||||
e_dims = 128
|
||||
d_dims = 64
|
||||
self.pretrain = True
|
||||
self.pretrain_just_disabled = False
|
||||
|
||||
class CommonModel(object):
|
||||
def downscale (self, dim, kernel_size=5, dilation_rate=1):
|
||||
def func(x):
|
||||
return SubpixelDownscaler()(ELU()(Conv2D(dim // 4, kernel_size=kernel_size, strides=1, dilation_rate=dilation_rate, padding='same')(x)))
|
||||
return func
|
||||
masked_training = True
|
||||
|
||||
def upscale (self, dim, size=(2,2)):
|
||||
def func(x):
|
||||
return SubpixelUpscaler(size=size)(ELU()(Conv2D(dim * np.prod(size) , kernel_size=3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
models_opt_on_gpu = len(devices) == 1 and devices[0].total_mem_gb >= 4
|
||||
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'
|
||||
|
||||
def ResidualBlock(self, dim):
|
||||
def func(inp):
|
||||
x = Conv2D(dim, kernel_size=3, padding='same')(inp)
|
||||
x = LeakyReLU(0.2)(x)
|
||||
x = Conv2D(dim, kernel_size=3, padding='same')(x)
|
||||
x = Add()([x, inp])
|
||||
x = LeakyReLU(0.2)(x)
|
||||
return x
|
||||
return func
|
||||
input_nc = 3
|
||||
output_nc = 3
|
||||
bgr_shape = (resolution, resolution, output_nc)
|
||||
mask_shape = (resolution, resolution, 1)
|
||||
lowest_dense_res = resolution // 16
|
||||
|
||||
class QModel(CommonModel):
|
||||
def __init__(self, resolution, ae_dims, e_dims, d_dims):
|
||||
super().__init__()
|
||||
bgr_shape = (resolution, resolution, 3)
|
||||
mask_shape = (resolution, resolution, 1)
|
||||
lowest_dense_res = resolution // 16
|
||||
self.model_filename_list = []
|
||||
|
||||
def enc_flow():
|
||||
def func(inp):
|
||||
x = self.downscale(e_dims, 3, 1 )(inp)
|
||||
x = self.downscale(e_dims*2, 3, 1 )(x)
|
||||
x = self.downscale(e_dims*4, 3, 1 )(x)
|
||||
x0 = self.downscale(e_dims*8, 3, 1 )(x)
|
||||
|
||||
x = self.downscale(e_dims, 3, 2 )(inp)
|
||||
x = self.downscale(e_dims*2, 3, 2 )(x)
|
||||
x = self.downscale(e_dims*4, 3, 2 )(x)
|
||||
x1 = self.downscale(e_dims*8, 3, 2 )(x)
|
||||
|
||||
x = Concatenate()([x0,x1])
|
||||
|
||||
x = DenseMaxout(ae_dims, kernel_initializer='orthogonal')(Flatten()(x))
|
||||
x = DenseMaxout(lowest_dense_res * lowest_dense_res * ae_dims, kernel_initializer='orthogonal')(x)
|
||||
x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x)
|
||||
|
||||
x = self.ResidualBlock(ae_dims)(x)
|
||||
x = self.upscale(d_dims*8)(x)
|
||||
x = self.ResidualBlock(d_dims*8)(x)
|
||||
return x
|
||||
return func
|
||||
|
||||
def dec_flow():
|
||||
def func(inp):
|
||||
x = self.upscale(d_dims*4)(inp)
|
||||
x = self.ResidualBlock(d_dims*4)(x)
|
||||
x = self.upscale(d_dims*2)(x)
|
||||
x = self.ResidualBlock(d_dims*2)(x)
|
||||
x = self.upscale(d_dims)(x)
|
||||
x = self.ResidualBlock(d_dims)(x)
|
||||
|
||||
y = self.upscale(d_dims)(inp)
|
||||
y = self.upscale(d_dims//2)(y)
|
||||
y = self.upscale(d_dims//4)(y)
|
||||
|
||||
return Conv2D(3, kernel_size=1, padding='same', activation='tanh')(x), \
|
||||
Conv2D(1, kernel_size=1, padding='same', activation='sigmoid')(y)
|
||||
with tf.device ('/CPU:0'):
|
||||
#Place holders on CPU
|
||||
self.warped_src = tf.placeholder (tf.float32, (None,)+bgr_shape)
|
||||
self.warped_dst = tf.placeholder (tf.float32, (None,)+bgr_shape)
|
||||
|
||||
return func
|
||||
self.target_src = tf.placeholder (tf.float32, (None,)+bgr_shape)
|
||||
self.target_dst = tf.placeholder (tf.float32, (None,)+bgr_shape)
|
||||
|
||||
self.encoder = modelify(enc_flow()) ( Input(bgr_shape) )
|
||||
self.target_srcm = tf.placeholder (tf.float32, (None,)+mask_shape)
|
||||
self.target_dstm = tf.placeholder (tf.float32, (None,)+mask_shape)
|
||||
|
||||
sh = K.int_shape( self.encoder.outputs[0] )[1:]
|
||||
self.decoder_src = modelify(dec_flow()) ( Input(sh) )
|
||||
self.decoder_dst = modelify(dec_flow()) ( Input(sh) )
|
||||
# Initializing model classes
|
||||
with tf.device (models_opt_device):
|
||||
self.encoder = Encoder(in_ch=input_nc, e_ch=e_dims, name='encoder')
|
||||
encoder_out_ch = self.encoder.compute_output_shape ( (tf.float32, (None,resolution,resolution,input_nc)))[-1]
|
||||
|
||||
self.src_trainable_weights = self.encoder.trainable_weights + self.decoder_src.trainable_weights
|
||||
self.dst_trainable_weights = self.encoder.trainable_weights + self.decoder_dst.trainable_weights
|
||||
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_shape ( (tf.float32, (None,encoder_out_ch)))[-1]
|
||||
|
||||
self.warped_src, self.warped_dst = Input(bgr_shape), Input(bgr_shape)
|
||||
self.target_src, self.target_dst = Input(bgr_shape), Input(bgr_shape)
|
||||
self.target_srcm, self.target_dstm = Input(mask_shape), Input(mask_shape)
|
||||
|
||||
self.src_code = self.encoder(self.warped_src)
|
||||
self.dst_code = self.encoder(self.warped_dst)
|
||||
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.pred_src_src, self.pred_src_srcm = self.decoder_src(self.src_code)
|
||||
self.pred_dst_dst, self.pred_dst_dstm = self.decoder_dst(self.dst_code)
|
||||
self.pred_src_dst, self.pred_src_dstm = self.decoder_src(self.dst_code)
|
||||
self.model_filename_list += [ [self.encoder, 'encoder.npy' ],
|
||||
[self.inter, 'inter.npy' ],
|
||||
[self.decoder_src, 'decoder_src.npy'],
|
||||
[self.decoder_dst, 'decoder_dst.npy'] ]
|
||||
|
||||
def get_model_filename_list(self, exclude_for_pretrain=False):
|
||||
ar = []
|
||||
if not exclude_for_pretrain:
|
||||
ar += [ [self.encoder, 'encoder.h5'] ]
|
||||
ar += [ [self.decoder_src, 'decoder_src.h5'],
|
||||
[self.decoder_dst, 'decoder_dst.h5'] ]
|
||||
|
||||
return ar
|
||||
if self.is_training:
|
||||
self.src_dst_trainable_weights = self.encoder.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()
|
||||
|
||||
self.model = QModel (resolution, 128, 64, 64)
|
||||
# Initialize optimizers
|
||||
self.src_dst_opt = nn.TFRMSpropOptimizer(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') ]
|
||||
|
||||
loaded, not_loaded = [], self.model.get_model_filename_list()
|
||||
if not self.is_first_run():
|
||||
loaded, not_loaded = self.load_weights_safe(not_loaded)
|
||||
if self.is_training:
|
||||
# Adjust batch size for multiple GPU
|
||||
gpu_count = max(1, len(devices) )
|
||||
bs_per_gpu = max(1, 4 // gpu_count)
|
||||
self.set_batch_size( gpu_count*bs_per_gpu)
|
||||
|
||||
CA_models = [ model for model, _ in not_loaded ]
|
||||
|
||||
self.CA_conv_weights_list = []
|
||||
for model in CA_models:
|
||||
for layer in model.layers:
|
||||
if type(layer) == keras.layers.Conv2D:
|
||||
self.CA_conv_weights_list += [layer.weights[0]] #- is Conv2D kernel_weights
|
||||
|
||||
if self.is_training_mode:
|
||||
lr_dropout = 0.3 if nnlib.device.backend != 'plaidML' else 0.0
|
||||
self.src_dst_opt = RMSprop(lr=2e-4, lr_dropout=lr_dropout)
|
||||
self.src_dst_mask_opt = RMSprop(lr=2e-4, lr_dropout=lr_dropout)
|
||||
|
||||
target_src_masked = self.model.target_src*self.model.target_srcm
|
||||
target_dst_masked = self.model.target_dst*self.model.target_dstm
|
||||
|
||||
pred_src_src_masked = self.model.pred_src_src*self.model.target_srcm
|
||||
pred_dst_dst_masked = self.model.pred_dst_dst*self.model.target_dstm
|
||||
# Compute losses per GPU
|
||||
gpu_pred_src_src_list = []
|
||||
gpu_pred_dst_dst_list = []
|
||||
gpu_pred_src_dst_list = []
|
||||
gpu_pred_src_srcm_list = []
|
||||
gpu_pred_dst_dstm_list = []
|
||||
gpu_pred_src_dstm_list = []
|
||||
|
||||
src_loss = K.mean ( 10*dssim(kernel_size=int(resolution/11.6),max_value=2.0)( target_src_masked+1, pred_src_src_masked+1) )
|
||||
src_loss += K.mean ( 10*K.square( target_src_masked - pred_src_src_masked ) )
|
||||
src_loss += K.mean(K.square(self.model.target_srcm-self.model.pred_src_srcm))
|
||||
gpu_src_losses = []
|
||||
gpu_dst_losses = []
|
||||
gpu_src_dst_loss_gvs = []
|
||||
|
||||
dst_loss = K.mean( 10*dssim(kernel_size=int(resolution/11.6),max_value=2.0)(target_dst_masked+1, pred_dst_dst_masked+1) )
|
||||
dst_loss += K.mean( 10*K.square( target_dst_masked - pred_dst_dst_masked ) )
|
||||
dst_loss += K.mean(K.square(self.model.target_dstm-self.model.pred_dst_dstm))
|
||||
for gpu_id in range(gpu_count):
|
||||
with tf.device( f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
|
||||
batch_slice = slice( gpu_id*bs_per_gpu, (gpu_id+1)*bs_per_gpu )
|
||||
with tf.device(f'/CPU:0'):
|
||||
# slice on CPU, otherwise all batch data will be transfered to GPU first
|
||||
gpu_warped_src = self.warped_src [batch_slice,:,:,:]
|
||||
gpu_warped_dst = self.warped_dst [batch_slice,:,:,:]
|
||||
gpu_target_src = self.target_src [batch_slice,:,:,:]
|
||||
gpu_target_dst = self.target_dst [batch_slice,:,:,:]
|
||||
gpu_target_srcm = self.target_srcm[batch_slice,:,:,:]
|
||||
gpu_target_dstm = self.target_dstm[batch_slice,:,:,:]
|
||||
|
||||
self.src_train = K.function ([self.model.warped_src, self.model.target_src, self.model.target_srcm], [src_loss], self.src_dst_opt.get_updates( src_loss, self.model.src_trainable_weights) )
|
||||
self.dst_train = K.function ([self.model.warped_dst, self.model.target_dst, self.model.target_dstm], [dst_loss], self.src_dst_opt.get_updates( dst_loss, self.model.dst_trainable_weights) )
|
||||
self.AE_view = K.function ([self.model.warped_src, self.model.warped_dst], [self.model.pred_src_src, self.model.pred_dst_dst, self.model.pred_dst_dstm, self.model.pred_src_dst, self.model.pred_src_dstm])
|
||||
# process model tensors
|
||||
gpu_src_code = self.inter(self.encoder(gpu_warped_src))
|
||||
gpu_dst_code = self.inter(self.encoder(gpu_warped_dst))
|
||||
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(gpu_src_code)
|
||||
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
|
||||
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
|
||||
|
||||
gpu_pred_src_src_list.append(gpu_pred_src_src)
|
||||
gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
|
||||
gpu_pred_src_dst_list.append(gpu_pred_src_dst)
|
||||
|
||||
gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
|
||||
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_dst_masked = gpu_target_dst*gpu_target_dstm_blur
|
||||
gpu_target_dst_anti_masked = gpu_target_dst*(1.0 - gpu_target_dstm_blur)
|
||||
|
||||
gpu_target_srcmasked_opt = gpu_target_src*gpu_target_srcm_blur if masked_training else gpu_target_src
|
||||
gpu_target_dst_masked_opt = gpu_target_dst_masked if masked_training else gpu_target_dst
|
||||
|
||||
gpu_pred_src_src_masked_opt = gpu_pred_src_src*gpu_target_srcm_blur if masked_training else gpu_pred_src_src
|
||||
gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst*gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst
|
||||
|
||||
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_srcmasked_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_srcmasked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
|
||||
gpu_src_loss += tf.reduce_mean ( 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*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
|
||||
gpu_dst_loss += tf.reduce_mean ( tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )
|
||||
|
||||
gpu_src_losses += [gpu_src_loss]
|
||||
gpu_dst_losses += [gpu_dst_loss]
|
||||
|
||||
gpu_src_dst_loss = gpu_src_loss + gpu_dst_loss
|
||||
gpu_src_dst_loss_gvs += [ nn.tf_gradients ( gpu_src_dst_loss, self.src_dst_trainable_weights ) ]
|
||||
|
||||
|
||||
# Average losses and gradients, and create optimizer update ops
|
||||
with tf.device (models_opt_device):
|
||||
if gpu_count == 1:
|
||||
pred_src_src = gpu_pred_src_src_list[0]
|
||||
pred_dst_dst = gpu_pred_dst_dst_list[0]
|
||||
pred_src_dst = gpu_pred_src_dst_list[0]
|
||||
pred_src_srcm = gpu_pred_src_srcm_list[0]
|
||||
pred_dst_dstm = gpu_pred_dst_dstm_list[0]
|
||||
pred_src_dstm = gpu_pred_src_dstm_list[0]
|
||||
|
||||
src_loss = gpu_src_losses[0]
|
||||
dst_loss = gpu_dst_losses[0]
|
||||
src_dst_loss_gv = gpu_src_dst_loss_gvs[0]
|
||||
else:
|
||||
pred_src_src = tf.concat(gpu_pred_src_src_list, 0)
|
||||
pred_dst_dst = tf.concat(gpu_pred_dst_dst_list, 0)
|
||||
pred_src_dst = tf.concat(gpu_pred_src_dst_list, 0)
|
||||
pred_src_srcm = tf.concat(gpu_pred_src_srcm_list, 0)
|
||||
pred_dst_dstm = tf.concat(gpu_pred_dst_dstm_list, 0)
|
||||
pred_src_dstm = tf.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_dst_loss_gv_op = self.src_dst_opt.get_update_op (src_dst_loss_gv)
|
||||
|
||||
# Initializing training and view functions
|
||||
def src_dst_train(warped_src, target_src, target_srcm, \
|
||||
warped_dst, target_dst, target_dstm):
|
||||
s, d, _ = nn.tf_sess.run ( [ src_loss, dst_loss, src_dst_loss_gv_op],
|
||||
feed_dict={self.warped_src :warped_src,
|
||||
self.target_src :target_src,
|
||||
self.target_srcm:target_srcm,
|
||||
self.warped_dst :warped_dst,
|
||||
self.target_dst :target_dst,
|
||||
self.target_dstm:target_dstm,
|
||||
})
|
||||
s = np.mean(s)
|
||||
d = np.mean(d)
|
||||
return s, d
|
||||
self.src_dst_train = src_dst_train
|
||||
|
||||
def AE_view(warped_src, warped_dst):
|
||||
return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm],
|
||||
feed_dict={self.warped_src:warped_src,
|
||||
self.warped_dst:warped_dst})
|
||||
|
||||
self.AE_view = AE_view
|
||||
else:
|
||||
self.AE_convert = K.function ([self.model.warped_dst],[ self.model.pred_src_dst, self.model.pred_dst_dstm, self.model.pred_src_dstm ])
|
||||
# Initializing merge function
|
||||
with tf.device( f'/GPU:0' if len(devices) != 0 else f'/CPU:0'):
|
||||
gpu_dst_code = self.inter(self.encoder(self.warped_dst))
|
||||
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
|
||||
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
|
||||
|
||||
if self.is_training_mode:
|
||||
def AE_merge( warped_dst):
|
||||
return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst})
|
||||
|
||||
self.AE_merge = AE_merge
|
||||
|
||||
|
||||
|
||||
|
||||
# Loading/initializing all models/optimizers weights
|
||||
for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
|
||||
do_init = self.is_first_run()
|
||||
|
||||
if self.pretrain_just_disabled:
|
||||
if model == self.inter:
|
||||
do_init = True
|
||||
|
||||
if not do_init:
|
||||
do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) )
|
||||
|
||||
if do_init and self.pretrained_model_path is not None:
|
||||
pretrained_filepath = self.pretrained_model_path / filename
|
||||
if pretrained_filepath.exists():
|
||||
do_init = not model.load_weights(pretrained_filepath)
|
||||
|
||||
if do_init:
|
||||
model.init_weights()
|
||||
|
||||
# initializing sample generators
|
||||
|
||||
if self.is_training:
|
||||
t = SampleProcessor.Types
|
||||
face_type = t.FACE_TYPE_FULL
|
||||
|
||||
training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path()
|
||||
training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path()
|
||||
|
||||
cpu_count = multiprocessing.cpu_count()
|
||||
|
||||
src_generators_count = cpu_count // 2
|
||||
dst_generators_count = cpu_count - src_generators_count
|
||||
|
||||
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=False, scale_range=np.array([-0.05, 0.05]) ),
|
||||
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_BGR), 'resolution': resolution, 'normalize_tanh':True },
|
||||
{'types' : (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_BGR), 'resolution': resolution, 'normalize_tanh':True },
|
||||
{'types' : (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_M), 'resolution': resolution } ]
|
||||
),
|
||||
SampleGeneratorFace(training_data_src_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False),
|
||||
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':resolution, },
|
||||
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution': resolution, },
|
||||
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution } ],
|
||||
generators_count=src_generators_count ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=False, ),
|
||||
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_BGR), 'resolution': resolution, 'normalize_tanh':True },
|
||||
{'types' : (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_BGR), 'resolution': resolution, 'normalize_tanh':True },
|
||||
{'types' : (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_M), 'resolution': resolution} ])
|
||||
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False),
|
||||
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':resolution},
|
||||
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution': resolution},
|
||||
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution} ],
|
||||
generators_count=dst_generators_count )
|
||||
])
|
||||
self.counter = 0
|
||||
|
||||
|
||||
self.last_samples = None
|
||||
|
||||
#override
|
||||
def get_model_filename_list(self):
|
||||
return self.model.get_model_filename_list ()
|
||||
return self.model_filename_list
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( self.get_model_filename_list() )
|
||||
for model, filename in io.progress_bar_generator(self.get_model_filename_list(), "Saving", leave=False):
|
||||
model.save_weights ( self.get_strpath_storage_for_file(filename) )
|
||||
|
||||
|
||||
#override
|
||||
def on_success_train_one_iter(self):
|
||||
if len(self.CA_conv_weights_list) != 0:
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
CAInitializerMP ( self.CA_conv_weights_list )
|
||||
self.CA_conv_weights_list = []
|
||||
|
||||
#override
|
||||
def onTrainOneIter(self, generators_samples, generators_list):
|
||||
warped_src, target_src, target_srcm = generators_samples[0]
|
||||
warped_dst, target_dst, target_dstm = generators_samples[1]
|
||||
|
||||
self.counter += 1
|
||||
if self.counter % 3 == 0:
|
||||
src_loss, = self.src_train ([warped_src, target_src, target_srcm])
|
||||
dst_loss, = self.dst_train ([warped_dst, target_dst, target_dstm])
|
||||
def onTrainOneIter(self):
|
||||
if self.get_iter() % 3 == 0 and self.last_samples is not None:
|
||||
( (warped_src, target_src, target_srcm), \
|
||||
(warped_dst, target_dst, target_dstm) ) = self.last_samples
|
||||
src_loss, dst_loss = self.src_dst_train (target_src, target_src, target_srcm,
|
||||
target_dst, target_dst, target_dstm)
|
||||
else:
|
||||
src_loss, = self.src_train ([target_src, target_src, target_srcm])
|
||||
dst_loss, = self.dst_train ([target_dst, target_dst, target_dstm])
|
||||
samples = self.last_samples = self.generate_next_samples()
|
||||
( (warped_src, target_src, target_srcm), \
|
||||
(warped_dst, target_dst, target_dstm) ) = samples
|
||||
|
||||
src_loss, dst_loss = self.src_dst_train (warped_src, target_src, target_srcm,
|
||||
warped_dst, target_dst, target_dstm)
|
||||
|
||||
return ( ('src_loss', src_loss), ('dst_loss', dst_loss), )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
test_S = sample[0][1][0:4] #first 4 samples
|
||||
test_S_m = sample[0][2][0:4] #first 4 samples
|
||||
test_D = sample[1][1][0:4]
|
||||
test_D_m = sample[1][2][0:4]
|
||||
def onGetPreview(self, samples):
|
||||
n_samples = min(4, self.get_batch_size() )
|
||||
|
||||
S, D, SS, DD, DDM, SD, SDM = [test_S,test_D] + self.AE_view ([test_S, test_D])
|
||||
S, D, SS, DD, SD, = [ np.clip(x/2+0.5, 0.0, 1.0) for x in [S, D, SS, DD, SD] ]
|
||||
DDM, SDM, = [ np.clip( np.repeat (x, (3,), -1), 0, 1) for x in [DDM, SDM] ]
|
||||
( (warped_src, target_src, target_srcm),
|
||||
(warped_dst, target_dst, target_dstm) ) = \
|
||||
[ [sample[0:n_samples] for sample in sample_list ]
|
||||
for sample_list in samples ]
|
||||
|
||||
S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ]
|
||||
DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ]
|
||||
|
||||
result = []
|
||||
st = []
|
||||
for i in range(len(test_S)):
|
||||
for i in range(n_samples):
|
||||
ar = S[i], SS[i], D[i], DD[i], SD[i]
|
||||
st.append ( np.concatenate ( ar, axis=1) )
|
||||
|
||||
result += [ ('Quick96', np.concatenate (st, axis=0 )), ]
|
||||
|
||||
|
||||
st_m = []
|
||||
for i in range(len(test_S)):
|
||||
ar = S[i]*test_S_m[i], SS[i], D[i]*test_D_m[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i])
|
||||
for i in range(n_samples):
|
||||
ar = S[i]*target_srcm[i], SS[i], D[i]*target_dstm[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i])
|
||||
st_m.append ( np.concatenate ( ar, axis=1) )
|
||||
|
||||
result += [ ('Quick96 masked', np.concatenate (st_m, axis=0 )), ]
|
||||
|
||||
return result
|
||||
|
||||
def predictor_func (self, face=None, dummy_predict=False):
|
||||
if dummy_predict:
|
||||
self.AE_convert ([ np.zeros ( (1, self.resolution, self.resolution, 3), dtype=np.float32 ) ])
|
||||
else:
|
||||
face = face * 2 - 1
|
||||
bgr, mask_dst_dstm, mask_src_dstm = self.AE_convert ([face[np.newaxis,...]])
|
||||
bgr = bgr /2 + 0.5
|
||||
mask = mask_dst_dstm[0] * mask_src_dstm[0]
|
||||
return bgr[0], mask[...,0]
|
||||
def predictor_func (self, face=None):
|
||||
|
||||
bgr, mask_dst_dstm, mask_src_dstm = self.AE_merge (face[np.newaxis,...])
|
||||
mask = mask_dst_dstm[0] * mask_src_dstm[0]
|
||||
return bgr[0], mask[...,0]
|
||||
|
||||
#override
|
||||
def get_ConverterConfig(self):
|
||||
import converters
|
||||
return self.predictor_func, (self.resolution, self.resolution, 3), converters.ConverterConfigMasked(face_type=FaceType.FULL,
|
||||
default_mode='seamless', clip_hborder_mask_per=0.0625)
|
||||
def get_MergerConfig(self):
|
||||
face_type = FaceType.FULL
|
||||
|
||||
Model = Quick96Model
|
||||
import merger
|
||||
return self.predictor_func, (self.resolution, self.resolution, 3), merger.MergerConfigMasked(face_type=face_type,
|
||||
default_mode = 'overlay',
|
||||
clip_hborder_mask_per=0.0625 if (face_type != FaceType.HALF) else 0,
|
||||
)
|
||||
|
||||
Model = QModel
|
||||
|
|
|
|||
|
|
@ -1,568 +0,0 @@
|
|||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
|
||||
import mathlib
|
||||
from facelib import FaceType
|
||||
from interact import interact as io
|
||||
from models import ModelBase
|
||||
from nnlib import nnlib
|
||||
from samplelib import *
|
||||
|
||||
|
||||
#SAE - Styled AutoEncoder
|
||||
class SAEModel(ModelBase):
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
yn_str = {True:'y',False:'n'}
|
||||
|
||||
default_resolution = 128
|
||||
default_archi = 'df'
|
||||
default_face_type = 'f'
|
||||
|
||||
|
||||
if is_first_run:
|
||||
resolution = io.input_int("Resolution ( 64-256 ?:help skip:128) : ", 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
|
||||
self.options['resolution'] = resolution
|
||||
|
||||
self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="Half face has better resolution, but covers less area of cheeks.").lower()
|
||||
else:
|
||||
self.options['resolution'] = self.options.get('resolution', default_resolution)
|
||||
self.options['face_type'] = self.options.get('face_type', default_face_type)
|
||||
|
||||
default_learn_mask = self.options.get('learn_mask', True)
|
||||
if is_first_run or ask_override:
|
||||
self.options['learn_mask'] = io.input_bool ( f"Learn mask? (y/n, ?:help skip:{yn_str[default_learn_mask]} ) : " , default_learn_mask, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted. Model with style values can be learned without mask and produce same quality result.")
|
||||
else:
|
||||
self.options['learn_mask'] = self.options.get('learn_mask', default_learn_mask)
|
||||
|
||||
if (is_first_run or ask_override) and 'tensorflow' in self.device_config.backend:
|
||||
def_optimizer_mode = self.options.get('optimizer_mode', 1)
|
||||
self.options['optimizer_mode'] = io.input_int ("Optimizer mode? ( 1,2,3 ?:help skip:%d) : " % (def_optimizer_mode), def_optimizer_mode, help_message="1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power.")
|
||||
else:
|
||||
self.options['optimizer_mode'] = self.options.get('optimizer_mode', 1)
|
||||
|
||||
if is_first_run:
|
||||
self.options['archi'] = io.input_str ("AE architecture (df, liae ?:help skip:%s) : " % (default_archi) , default_archi, ['df','liae'], help_message="'df' keeps faces more natural. 'liae' can fix overly different face shapes.").lower() #-s version is slower, but has decreased change to collapse.
|
||||
else:
|
||||
self.options['archi'] = self.options.get('archi', default_archi)
|
||||
|
||||
default_ae_dims = 256 if 'liae' in self.options['archi'] else 512
|
||||
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 )
|
||||
default_d_ch_dims = self.options['e_ch_dims'] // 2
|
||||
self.options['d_ch_dims'] = np.clip ( io.input_int("Decoder dims per channel (10-85 ?:help skip:%d) : " % (default_d_ch_dims) , default_d_ch_dims, help_message="More decoder dims help to get better details, but require more VRAM. You can fine-tune model size to fit your GPU." ), 10, 85 )
|
||||
self.options['ca_weights'] = io.input_bool (f"Use CA weights? (y/n, ?:help skip:{yn_str[def_ca_weights]} ) : ", def_ca_weights, help_message="Initialize network with 'Convolution Aware' weights. This may help to achieve a higher accuracy model, but consumes a time at first run.")
|
||||
else:
|
||||
self.options['ae_dims'] = self.options.get('ae_dims', default_ae_dims)
|
||||
self.options['e_ch_dims'] = self.options.get('e_ch_dims', default_e_ch_dims)
|
||||
self.options['d_ch_dims'] = self.options.get('d_ch_dims', default_d_ch_dims)
|
||||
self.options['ca_weights'] = self.options.get('ca_weights', def_ca_weights)
|
||||
|
||||
default_face_style_power = 0.0
|
||||
default_bg_style_power = 0.0
|
||||
if is_first_run or ask_override:
|
||||
def_pixel_loss = self.options.get('pixel_loss', False)
|
||||
self.options['pixel_loss'] = io.input_bool (f"Use pixel loss? (y/n, ?:help skip:{yn_str[def_pixel_loss]} ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time. Enabling this option too early increases the chance of model collapse.")
|
||||
|
||||
default_face_style_power = default_face_style_power if is_first_run else self.options.get('face_style_power', default_face_style_power)
|
||||
self.options['face_style_power'] = np.clip ( io.input_number("Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_face_style_power), default_face_style_power,
|
||||
help_message="Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes. Enabling this option increases the chance of model collapse."), 0.0, 100.0 )
|
||||
|
||||
default_bg_style_power = default_bg_style_power if is_first_run else self.options.get('bg_style_power', default_bg_style_power)
|
||||
self.options['bg_style_power'] = np.clip ( io.input_number("Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_bg_style_power), default_bg_style_power,
|
||||
help_message="Learn to transfer image around face. This can make face more like dst. Enabling this option increases the chance of model collapse."), 0.0, 100.0 )
|
||||
|
||||
default_ct_mode = self.options.get('ct_mode', 'none')
|
||||
self.options['ct_mode'] = io.input_str (f"Color transfer mode apply to src faceset. ( none/rct/lct/mkl/idt/sot, ?:help skip:{default_ct_mode}) : ", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.")
|
||||
|
||||
if nnlib.device.backend != 'plaidML': # todo https://github.com/plaidml/plaidml/issues/301
|
||||
default_clipgrad = False if is_first_run else self.options.get('clipgrad', False)
|
||||
self.options['clipgrad'] = io.input_bool (f"Enable gradient clipping? (y/n, ?:help skip:{yn_str[default_clipgrad]}) : ", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.")
|
||||
else:
|
||||
self.options['clipgrad'] = False
|
||||
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
self.options['face_style_power'] = self.options.get('face_style_power', default_face_style_power)
|
||||
self.options['bg_style_power'] = self.options.get('bg_style_power', default_bg_style_power)
|
||||
self.options['ct_mode'] = self.options.get('ct_mode', 'none')
|
||||
self.options['clipgrad'] = self.options.get('clipgrad', False)
|
||||
|
||||
if is_first_run:
|
||||
self.options['pretrain'] = io.input_bool ("Pretrain the model? (y/n, ?:help skip:n) : ", False, help_message="Pretrain the model with large amount of various faces. This technique may help to train the fake with overly different face shapes and light conditions of src/dst data. Face will be look more like a morphed. To reduce the morph effect, some model files will be initialized but not be updated after pretrain: LIAE: inter_AB.h5 DF: encoder.h5. The longer you pretrain the model the more morphed face will look. After that, save and run the training again.")
|
||||
else:
|
||||
self.options['pretrain'] = False
|
||||
|
||||
#override
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements({1.5:4})
|
||||
|
||||
resolution = self.options['resolution']
|
||||
learn_mask = self.options['learn_mask']
|
||||
|
||||
ae_dims = self.options['ae_dims']
|
||||
e_ch_dims = self.options['e_ch_dims']
|
||||
d_ch_dims = self.options['d_ch_dims']
|
||||
self.pretrain = self.options['pretrain'] = self.options.get('pretrain', False)
|
||||
if not self.pretrain:
|
||||
self.options.pop('pretrain')
|
||||
|
||||
bgr_shape = (resolution, resolution, 3)
|
||||
mask_shape = (resolution, resolution, 1)
|
||||
|
||||
masked_training = True
|
||||
|
||||
class SAEDFModel(object):
|
||||
def __init__(self, resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask):
|
||||
super().__init__()
|
||||
self.learn_mask = learn_mask
|
||||
|
||||
output_nc = 3
|
||||
bgr_shape = (resolution, resolution, output_nc)
|
||||
mask_shape = (resolution, resolution, 1)
|
||||
lowest_dense_res = resolution // 16
|
||||
e_dims = output_nc*e_ch_dims
|
||||
|
||||
def upscale (dim):
|
||||
def func(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='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)
|
||||
x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x)
|
||||
x = upscale(ae_dims)(x)
|
||||
return x
|
||||
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='valid')(ZeroPadding2D(1)(inp))
|
||||
x = LeakyReLU(0.2)(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):
|
||||
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='valid', activation='sigmoid')(ZeroPadding2D(2)(x))
|
||||
return func
|
||||
|
||||
self.encoder = modelify(enc_flow(e_dims, ae_dims, lowest_dense_res)) ( Input(bgr_shape) )
|
||||
|
||||
sh = K.int_shape( self.encoder.outputs[0] )[1:]
|
||||
self.decoder_src = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) )
|
||||
self.decoder_dst = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) )
|
||||
|
||||
if learn_mask:
|
||||
self.decoder_srcm = modelify(dec_flow(1, d_ch_dims, add_residual_blocks=False)) ( Input(sh) )
|
||||
self.decoder_dstm = modelify(dec_flow(1, d_ch_dims, add_residual_blocks=False)) ( Input(sh) )
|
||||
|
||||
self.src_dst_trainable_weights = self.encoder.trainable_weights + self.decoder_src.trainable_weights + self.decoder_dst.trainable_weights
|
||||
|
||||
if learn_mask:
|
||||
self.src_dst_mask_trainable_weights = self.encoder.trainable_weights + self.decoder_srcm.trainable_weights + self.decoder_dstm.trainable_weights
|
||||
|
||||
self.warped_src, self.warped_dst = Input(bgr_shape), Input(bgr_shape)
|
||||
src_code, dst_code = self.encoder(self.warped_src), self.encoder(self.warped_dst)
|
||||
|
||||
self.pred_src_src = self.decoder_src(src_code)
|
||||
self.pred_dst_dst = self.decoder_dst(dst_code)
|
||||
self.pred_src_dst = self.decoder_src(dst_code)
|
||||
|
||||
if learn_mask:
|
||||
self.pred_src_srcm = self.decoder_srcm(src_code)
|
||||
self.pred_dst_dstm = self.decoder_dstm(dst_code)
|
||||
self.pred_src_dstm = self.decoder_srcm(dst_code)
|
||||
|
||||
def get_model_filename_list(self, exclude_for_pretrain=False):
|
||||
ar = []
|
||||
if not exclude_for_pretrain:
|
||||
ar += [ [self.encoder, 'encoder.h5'] ]
|
||||
ar += [ [self.decoder_src, 'decoder_src.h5'],
|
||||
[self.decoder_dst, 'decoder_dst.h5'] ]
|
||||
if self.learn_mask:
|
||||
ar += [ [self.decoder_srcm, 'decoder_srcm.h5'],
|
||||
[self.decoder_dstm, 'decoder_dstm.h5'] ]
|
||||
return ar
|
||||
|
||||
class SAELIAEModel(object):
|
||||
def __init__(self, resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask):
|
||||
super().__init__()
|
||||
self.learn_mask = learn_mask
|
||||
|
||||
output_nc = 3
|
||||
bgr_shape = (resolution, resolution, output_nc)
|
||||
mask_shape = (resolution, resolution, 1)
|
||||
|
||||
e_dims = output_nc*e_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='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='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
|
||||
|
||||
def inter_flow(lowest_dense_res, ae_dims):
|
||||
def func(x):
|
||||
x = Dense(ae_dims)(x)
|
||||
x = Dense(lowest_dense_res * lowest_dense_res * ae_dims*2)(x)
|
||||
x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims*2))(x)
|
||||
x = upscale(ae_dims*2)(x)
|
||||
return x
|
||||
return func
|
||||
|
||||
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='valid')(ZeroPadding2D(1)(inp))
|
||||
x = LeakyReLU(0.2)(x)
|
||||
x = Conv2D(dim, kernel_size=3, padding='valid')(ZeroPadding2D(1)(inp))
|
||||
x = Add()([x, inp])
|
||||
x = LeakyReLU(0.2)(x)
|
||||
return x
|
||||
return func
|
||||
|
||||
def func(x):
|
||||
x = upscale(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)
|
||||
|
||||
if add_residual_blocks:
|
||||
x = ResidualBlock(d_dims*4)(x)
|
||||
x = ResidualBlock(d_dims*4)(x)
|
||||
|
||||
x = upscale(d_dims*2)(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) )
|
||||
|
||||
sh = K.int_shape( self.encoder.outputs[0] )[1:]
|
||||
self.inter_B = modelify(inter_flow(lowest_dense_res, ae_dims)) ( Input(sh) )
|
||||
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_ch_dims)) ( Input(sh) )
|
||||
|
||||
if learn_mask:
|
||||
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
|
||||
|
||||
if learn_mask:
|
||||
self.src_dst_mask_trainable_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoderm.trainable_weights
|
||||
|
||||
self.warped_src, self.warped_dst = Input(bgr_shape), Input(bgr_shape)
|
||||
|
||||
warped_src_code = self.encoder (self.warped_src)
|
||||
warped_src_inter_AB_code = self.inter_AB (warped_src_code)
|
||||
warped_src_inter_code = Concatenate()([warped_src_inter_AB_code,warped_src_inter_AB_code])
|
||||
|
||||
warped_dst_code = self.encoder (self.warped_dst)
|
||||
warped_dst_inter_B_code = self.inter_B (warped_dst_code)
|
||||
warped_dst_inter_AB_code = self.inter_AB (warped_dst_code)
|
||||
warped_dst_inter_code = Concatenate()([warped_dst_inter_B_code,warped_dst_inter_AB_code])
|
||||
|
||||
warped_src_dst_inter_code = Concatenate()([warped_dst_inter_AB_code,warped_dst_inter_AB_code])
|
||||
|
||||
self.pred_src_src = self.decoder(warped_src_inter_code)
|
||||
self.pred_dst_dst = self.decoder(warped_dst_inter_code)
|
||||
self.pred_src_dst = self.decoder(warped_src_dst_inter_code)
|
||||
|
||||
if learn_mask:
|
||||
self.pred_src_srcm = self.decoderm(warped_src_inter_code)
|
||||
self.pred_dst_dstm = self.decoderm(warped_dst_inter_code)
|
||||
self.pred_src_dstm = self.decoderm(warped_src_dst_inter_code)
|
||||
|
||||
def get_model_filename_list(self, exclude_for_pretrain=False):
|
||||
ar = [ [self.encoder, 'encoder.h5'],
|
||||
[self.inter_B, 'inter_B.h5'] ]
|
||||
|
||||
if not exclude_for_pretrain:
|
||||
ar += [ [self.inter_AB, 'inter_AB.h5'] ]
|
||||
|
||||
ar += [ [self.decoder, 'decoder.h5'] ]
|
||||
|
||||
if self.learn_mask:
|
||||
ar += [ [self.decoderm, 'decoderm.h5'] ]
|
||||
|
||||
return ar
|
||||
|
||||
if 'df' in self.options['archi']:
|
||||
self.model = SAEDFModel (resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask)
|
||||
elif 'liae' in self.options['archi']:
|
||||
self.model = SAELIAEModel (resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask)
|
||||
|
||||
loaded, not_loaded = [], self.model.get_model_filename_list()
|
||||
if not self.is_first_run():
|
||||
loaded, not_loaded = self.load_weights_safe(not_loaded)
|
||||
|
||||
CA_models = []
|
||||
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 )
|
||||
|
||||
warped_src = self.model.warped_src
|
||||
target_src = Input ( (resolution, resolution, 3) )
|
||||
target_srcm = Input ( (resolution, resolution, 1) )
|
||||
|
||||
warped_dst = self.model.warped_dst
|
||||
target_dst = Input ( (resolution, resolution, 3) )
|
||||
target_dstm = Input ( (resolution, resolution, 1) )
|
||||
|
||||
target_src_sigm = target_src
|
||||
target_dst_sigm = target_dst
|
||||
|
||||
target_srcm_sigm = gaussian_blur( max(1, K.int_shape(target_srcm)[1] // 32) )(target_srcm)
|
||||
target_dstm_sigm = gaussian_blur( max(1, K.int_shape(target_dstm)[1] // 32) )(target_dstm)
|
||||
target_dstm_anti_sigm = 1.0 - target_dstm_sigm
|
||||
|
||||
target_src_masked = target_src_sigm*target_srcm_sigm
|
||||
target_dst_masked = target_dst_sigm*target_dstm_sigm
|
||||
target_dst_anti_masked = target_dst_sigm*target_dstm_anti_sigm
|
||||
|
||||
target_src_masked_opt = target_src_masked if masked_training else target_src_sigm
|
||||
target_dst_masked_opt = target_dst_masked if masked_training else target_dst_sigm
|
||||
|
||||
pred_src_src = self.model.pred_src_src
|
||||
pred_dst_dst = self.model.pred_dst_dst
|
||||
pred_src_dst = self.model.pred_src_dst
|
||||
if learn_mask:
|
||||
pred_src_srcm = self.model.pred_src_srcm
|
||||
pred_dst_dstm = self.model.pred_dst_dstm
|
||||
pred_src_dstm = self.model.pred_src_dstm
|
||||
|
||||
pred_src_src_sigm = self.model.pred_src_src
|
||||
pred_dst_dst_sigm = self.model.pred_dst_dst
|
||||
pred_src_dst_sigm = self.model.pred_src_dst
|
||||
|
||||
pred_src_src_masked = pred_src_src_sigm*target_srcm_sigm
|
||||
pred_dst_dst_masked = pred_dst_dst_sigm*target_dstm_sigm
|
||||
|
||||
pred_src_src_masked_opt = pred_src_src_masked if masked_training else pred_src_src_sigm
|
||||
pred_dst_dst_masked_opt = pred_dst_dst_masked if masked_training else pred_dst_dst_sigm
|
||||
|
||||
psd_target_dst_masked = pred_src_dst_sigm*target_dstm_sigm
|
||||
psd_target_dst_anti_masked = pred_src_dst_sigm*target_dstm_anti_sigm
|
||||
|
||||
if self.is_training_mode:
|
||||
self.src_dst_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1)
|
||||
self.src_dst_mask_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1)
|
||||
|
||||
if not self.options['pixel_loss']:
|
||||
src_loss = K.mean ( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_src_masked_opt, pred_src_src_masked_opt) )
|
||||
else:
|
||||
src_loss = K.mean ( 50*K.square( target_src_masked_opt - pred_src_src_masked_opt ) )
|
||||
|
||||
face_style_power = self.options['face_style_power'] / 100.0
|
||||
if face_style_power != 0:
|
||||
src_loss += style_loss(gaussian_blur_radius=resolution//16, loss_weight=face_style_power, wnd_size=0)( psd_target_dst_masked, target_dst_masked )
|
||||
|
||||
bg_style_power = self.options['bg_style_power'] / 100.0
|
||||
if bg_style_power != 0:
|
||||
if not self.options['pixel_loss']:
|
||||
src_loss += K.mean( (10*bg_style_power)*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( psd_target_dst_anti_masked, target_dst_anti_masked ))
|
||||
else:
|
||||
src_loss += K.mean( (50*bg_style_power)*K.square( psd_target_dst_anti_masked - target_dst_anti_masked ))
|
||||
|
||||
if not self.options['pixel_loss']:
|
||||
dst_loss = K.mean( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)(target_dst_masked_opt, pred_dst_dst_masked_opt) )
|
||||
else:
|
||||
dst_loss = K.mean( 50*K.square( target_dst_masked_opt - pred_dst_dst_masked_opt ) )
|
||||
|
||||
self.src_dst_train = K.function ([warped_src, warped_dst, target_src, target_srcm, target_dst, target_dstm],[src_loss,dst_loss], self.src_dst_opt.get_updates(src_loss+dst_loss, self.model.src_dst_trainable_weights) )
|
||||
|
||||
if self.options['learn_mask']:
|
||||
src_mask_loss = K.mean(K.square(target_srcm-pred_src_srcm))
|
||||
dst_mask_loss = K.mean(K.square(target_dstm-pred_dst_dstm))
|
||||
self.src_dst_mask_train = K.function ([warped_src, warped_dst, target_srcm, target_dstm],[src_mask_loss, dst_mask_loss], self.src_dst_mask_opt.get_updates(src_mask_loss+dst_mask_loss, self.model.src_dst_mask_trainable_weights ) )
|
||||
|
||||
if self.options['learn_mask']:
|
||||
self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm])
|
||||
else:
|
||||
self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src, pred_dst_dst, pred_src_dst ])
|
||||
|
||||
else:
|
||||
if self.options['learn_mask']:
|
||||
self.AE_convert = K.function ([warped_dst],[ pred_src_dst, pred_dst_dstm, pred_src_dstm ])
|
||||
else:
|
||||
self.AE_convert = K.function ([warped_dst],[ pred_src_dst ])
|
||||
|
||||
|
||||
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
|
||||
|
||||
t_mode_bgr = t.MODE_BGR if not self.pretrain else t.MODE_BGR_SHUFFLE
|
||||
|
||||
training_data_src_path = self.training_data_src_path
|
||||
training_data_dst_path = self.training_data_dst_path
|
||||
|
||||
if self.pretrain and self.pretraining_data_path is not None:
|
||||
training_data_src_path = self.pretraining_data_path
|
||||
training_data_dst_path = self.pretraining_data_path
|
||||
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(training_data_src_path, random_ct_samples_path=training_data_dst_path if self.options['ct_mode'] != 'none' else None,
|
||||
debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05]) ),
|
||||
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t_mode_bgr), 'resolution':resolution, 'ct_mode': self.options['ct_mode'] },
|
||||
{'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution, 'ct_mode': self.options['ct_mode'] },
|
||||
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution } ]
|
||||
),
|
||||
|
||||
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, ),
|
||||
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t_mode_bgr), 'resolution':resolution},
|
||||
{'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution},
|
||||
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution} ])
|
||||
])
|
||||
|
||||
#override
|
||||
def get_model_filename_list(self):
|
||||
ar = self.model.get_model_filename_list ( exclude_for_pretrain=(self.pretrain and self.iter != 0) )
|
||||
return ar
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( self.get_model_filename_list() )
|
||||
|
||||
#override
|
||||
def onTrainOneIter(self, generators_samples, generators_list):
|
||||
warped_src, target_src, target_srcm = generators_samples[0]
|
||||
warped_dst, target_dst, target_dstm = generators_samples[1]
|
||||
|
||||
feed = [warped_src, warped_dst, target_src, target_srcm, target_dst, target_dstm]
|
||||
|
||||
src_loss, dst_loss, = self.src_dst_train (feed)
|
||||
|
||||
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
|
||||
def onGetPreview(self, sample):
|
||||
test_S = sample[0][1][0:4] #first 4 samples
|
||||
test_S_m = sample[0][2][0:4] #first 4 samples
|
||||
test_D = sample[1][1][0:4]
|
||||
test_D_m = sample[1][2][0:4]
|
||||
|
||||
if self.options['learn_mask']:
|
||||
S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ]
|
||||
DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ]
|
||||
else:
|
||||
S, D, SS, DD, SD, = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ]
|
||||
|
||||
result = []
|
||||
st = []
|
||||
for i in range(len(test_S)):
|
||||
ar = S[i], SS[i], D[i], DD[i], SD[i]
|
||||
|
||||
st.append ( np.concatenate ( ar, axis=1) )
|
||||
|
||||
result += [ ('SAE', np.concatenate (st, axis=0 )), ]
|
||||
|
||||
if self.options['learn_mask']:
|
||||
st_m = []
|
||||
for i in range(len(test_S)):
|
||||
ar = S[i]*test_S_m[i], SS[i], D[i]*test_D_m[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i])
|
||||
st_m.append ( np.concatenate ( ar, axis=1) )
|
||||
|
||||
result += [ ('SAE masked', np.concatenate (st_m, axis=0 )), ]
|
||||
|
||||
return result
|
||||
|
||||
def predictor_func (self, face=None, dummy_predict=False):
|
||||
if dummy_predict:
|
||||
self.AE_convert ([ np.zeros ( (1, self.options['resolution'], self.options['resolution'], 3), dtype=np.float32 ) ])
|
||||
else:
|
||||
if self.options['learn_mask']:
|
||||
bgr, mask_dst_dstm, mask_src_dstm = self.AE_convert ([face[np.newaxis,...]])
|
||||
mask = mask_dst_dstm[0] * mask_src_dstm[0]
|
||||
return bgr[0], mask[...,0]
|
||||
else:
|
||||
bgr, = self.AE_convert ([face[np.newaxis,...]])
|
||||
return bgr[0]
|
||||
|
||||
#override
|
||||
def get_ConverterConfig(self):
|
||||
face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
|
||||
|
||||
import converters
|
||||
return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), converters.ConverterConfigMasked(face_type=face_type,
|
||||
default_mode = 'overlay' if self.options['ct_mode'] != 'none' or self.options['face_style_power'] or self.options['bg_style_power'] else 'seamless',
|
||||
clip_hborder_mask_per=0.0625 if (self.options['face_type'] == 'f') else 0,
|
||||
)
|
||||
|
||||
Model = SAEModel
|
||||
|
|
@ -1 +0,0 @@
|
|||
from .Model import Model
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -1,5 +1,5 @@
|
|||
from .ModelBase import ModelBase
|
||||
|
||||
def import_model(name):
|
||||
module = __import__('Model_'+name, globals(), locals(), [], 1)
|
||||
def import_model(model_class_name):
|
||||
module = __import__('Model_'+model_class_name, globals(), locals(), [], 1)
|
||||
return getattr(module, 'Model')
|
||||
|
|
|
|||
Binary file not shown.
Loading…
Add table
Add a link
Reference in a new issue