mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-20 21:43:21 -07:00
added new model U-net Face Morpher.
removed AVATAR - useless model was just for demo removed MIAEF128 - use UFM insted removed LIAEF128YAW - use model option sort by yaw on start for any model All models now ask some options on start. Session options (such as target epoch, batch_size, write_preview_history etc) can be overrided by special command arg. Converter now always ask options and no more support to define options via command line. fix bug when ConverterMasked always used not predicted mask. SampleGenerator now always generate samples with replicated border, exclude mask samples. refactorings
This commit is contained in:
iperov
parent
f3782a012b
commit
7b70e7eec1
29 changed files with 673 additions and 1013 deletions
|
|
@ -11,11 +11,12 @@ class ConverterBase(object):
|
|||
#overridable
|
||||
def __init__(self, predictor):
|
||||
self.predictor = predictor
|
||||
|
||||
|
||||
#overridable
|
||||
def get_mode(self):
|
||||
#MODE_FACE calls convert_face
|
||||
#MODE_IMAGE calls convert_image
|
||||
#MODE_IMAGE calls convert_image without landmarks
|
||||
#MODE_IMAGE_WITH_LANDMARKS calls convert_image with landmarks
|
||||
return ConverterBase.MODE_FACE
|
||||
|
||||
#overridable
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ class ConverterImage(ConverterBase):
|
|||
self.predictor ( np.zeros ( (self.predictor_input_size, self.predictor_input_size,3), dtype=np.float32) )
|
||||
|
||||
#override
|
||||
def convert_image (self, img_bgr, debug):
|
||||
def convert_image (self, img_bgr, img_landmarks, debug):
|
||||
img_size = img_bgr.shape[1], img_bgr.shape[0]
|
||||
|
||||
predictor_input_bgr = cv2.resize ( img_bgr, (self.predictor_input_size, self.predictor_input_size), cv2.INTER_LANCZOS4 )
|
||||
|
|
|
|||
|
|
@ -4,47 +4,53 @@ from facelib import FaceType
|
|||
import cv2
|
||||
import numpy as np
|
||||
from utils import image_utils
|
||||
|
||||
from utils.console_utils import *
|
||||
|
||||
class ConverterMasked(ConverterBase):
|
||||
|
||||
#override
|
||||
def __init__(self, predictor,
|
||||
predictor_input_size=0,
|
||||
output_size=0,
|
||||
face_type=FaceType.FULL,
|
||||
clip_border_mask_per = 0,
|
||||
masked_hist_match = True,
|
||||
hist_match_threshold = 255,
|
||||
mode='seamless',
|
||||
use_predicted_mask = True,
|
||||
erode_mask_modifier=0,
|
||||
blur_mask_modifier=0,
|
||||
seamless_erode_mask_modifier=0,
|
||||
output_face_scale_modifier=0.0,
|
||||
transfercolor=False,
|
||||
final_image_color_degrade_power=0,
|
||||
alpha=False,
|
||||
face_type=FaceType.FULL,
|
||||
base_erode_mask_modifier = 0,
|
||||
base_blur_mask_modifier = 0,
|
||||
|
||||
**in_options):
|
||||
|
||||
super().__init__(predictor)
|
||||
|
||||
self.predictor_input_size = predictor_input_size
|
||||
self.output_size = output_size
|
||||
self.face_type = face_type
|
||||
self.use_predicted_mask = use_predicted_mask
|
||||
self.clip_border_mask_per = clip_border_mask_per
|
||||
self.masked_hist_match = masked_hist_match
|
||||
self.hist_match_threshold = hist_match_threshold
|
||||
self.mode = mode
|
||||
self.erode_mask_modifier = erode_mask_modifier
|
||||
self.blur_mask_modifier = blur_mask_modifier
|
||||
self.seamless_erode_mask_modifier = seamless_erode_mask_modifier
|
||||
self.output_face_scale = np.clip(1.0 + output_face_scale_modifier*0.01, 0.5, 1.5)
|
||||
self.transfercolor = transfercolor
|
||||
self.face_type = face_type
|
||||
self.TFLabConverter = None
|
||||
self.final_image_color_degrade_power = np.clip (final_image_color_degrade_power, 0, 100)
|
||||
self.alpha = alpha
|
||||
|
||||
|
||||
mode = input_int ("Choose mode: (1) overlay, (2) hist match, (3) hist match bw, (4) seamless (default), (5) seamless hist match : ", 4)
|
||||
self.mode = {1:'overlay',
|
||||
2:'hist-match',
|
||||
3:'hist-match-bw',
|
||||
4:'seamless',
|
||||
5:'seamless-hist-match'}.get (mode, 'seamless')
|
||||
|
||||
if self.mode == 'hist-match' or self.mode == 'hist-match-bw':
|
||||
self.masked_hist_match = input_bool("Masked hist match? (y/n skip:y) : ", True)
|
||||
|
||||
if self.mode == 'hist-match' or self.mode == 'hist-match-bw' or self.mode == 'seamless-hist-match':
|
||||
self.hist_match_threshold = np.clip ( input_int("Hist match threshold [0..255] (skip:255) : ", 255), 0, 255)
|
||||
|
||||
self.use_predicted_mask = input_bool("Use predicted mask? (y/n skip:y) : ", True)
|
||||
self.erode_mask_modifier = base_erode_mask_modifier + np.clip ( input_int ("Choose erode mask modifier [-200..200] (skip:0) : ", 0), -200, 200)
|
||||
self.blur_mask_modifier = base_blur_mask_modifier + np.clip ( input_int ("Choose blur mask modifier [-200..200] (skip:0) : ", 0), -200, 200)
|
||||
|
||||
self.seamless_erode_mask_modifier = 0
|
||||
if self.mode == 'seamless' or self.mode == 'seamless-hist-match':
|
||||
self.seamless_erode_mask_modifier = np.clip ( input_int ("Choose seamless erode mask modifier [-100..100] (skip:0) : ", 0), -100, 100)
|
||||
|
||||
self.output_face_scale = np.clip ( 1.0 + input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0)*0.01, 0.5, 1.5)
|
||||
self.transfercolor = input_bool("Transfer color from dst face to converted final face? (y/n skip:n) : ", False)
|
||||
self.final_image_color_degrade_power = np.clip ( input_int ("Degrade color power of final image [0..100] (skip:0) : ", 0), 0, 100)
|
||||
self.alpha = input_bool("Export png with alpha channel? (y/n skip:n) : ", False)
|
||||
print ("")
|
||||
|
||||
#override
|
||||
def get_mode(self):
|
||||
return ConverterBase.MODE_FACE
|
||||
|
|
@ -79,7 +85,7 @@ class ConverterMasked(ConverterBase):
|
|||
|
||||
if not self.use_predicted_mask:
|
||||
prd_face_mask_a_0 = predictor_input_mask_a_0
|
||||
|
||||
|
||||
prd_face_mask_a_0[ prd_face_mask_a_0 < 0.001 ] = 0.0
|
||||
|
||||
prd_face_mask_a = np.expand_dims (prd_face_mask_a_0, axis=-1)
|
||||
|
|
@ -145,16 +151,6 @@ class ConverterMasked(ConverterBase):
|
|||
print ("blur_size = %d" % (blur) )
|
||||
|
||||
img_mask_blurry_aaa = np.clip( img_mask_blurry_aaa, 0, 1.0 )
|
||||
|
||||
#if self.clip_border_mask_per > 0:
|
||||
# prd_border_rect_mask_a = np.ones ( prd_face_mask_a.shape, dtype=prd_face_mask_a.dtype)
|
||||
# prd_border_size = int ( prd_border_rect_mask_a.shape[1] * self.clip_border_mask_per )
|
||||
#
|
||||
# prd_border_rect_mask_a[0:prd_border_size,:,:] = 0
|
||||
# prd_border_rect_mask_a[-prd_border_size:,:,:] = 0
|
||||
# prd_border_rect_mask_a[:,0:prd_border_size,:] = 0
|
||||
# prd_border_rect_mask_a[:,-prd_border_size:,:] = 0
|
||||
# prd_border_rect_mask_a = np.expand_dims(cv2.blur(prd_border_rect_mask_a, (prd_border_size, prd_border_size) ),-1)
|
||||
|
||||
if self.mode == 'hist-match-bw':
|
||||
prd_face_bgr = cv2.cvtColor(prd_face_bgr, cv2.COLOR_BGR2GRAY)
|
||||
|
|
@ -174,22 +170,21 @@ class ConverterMasked(ConverterBase):
|
|||
|
||||
hist_match_2 = dst_face_bgr*hist_mask_a + (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=prd_face_bgr.dtype)
|
||||
hist_match_2[ hist_match_1 > 1.0 ] = 1.0
|
||||
|
||||
new_prd_face_bgr = image_utils.color_hist_match(hist_match_1, hist_match_2, self.hist_match_threshold )
|
||||
|
||||
prd_face_bgr = new_prd_face_bgr
|
||||
prd_face_bgr = image_utils.color_hist_match(hist_match_1, hist_match_2, self.hist_match_threshold )
|
||||
|
||||
if self.mode == 'hist-match-bw':
|
||||
prd_face_bgr = prd_face_bgr.astype(np.float32)
|
||||
|
||||
|
||||
|
||||
out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT )
|
||||
|
||||
|
||||
if debug:
|
||||
debugs += [out_img.copy()]
|
||||
debugs += [img_mask_blurry_aaa.copy()]
|
||||
|
||||
if self.mode == 'overlay':
|
||||
pass
|
||||
|
||||
if self.mode == 'seamless' or self.mode == 'seamless-hist-match':
|
||||
out_img = np.clip( img_bgr*(1-img_face_mask_aaa) + (out_img*img_face_mask_aaa) , 0, 1.0 )
|
||||
if debug:
|
||||
|
|
@ -200,14 +195,7 @@ class ConverterMasked(ConverterBase):
|
|||
|
||||
if debug:
|
||||
debugs += [out_img.copy()]
|
||||
|
||||
#if self.clip_border_mask_per > 0:
|
||||
# img_prd_border_rect_mask_a = cv2.warpAffine( prd_border_rect_mask_a, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT )
|
||||
# img_prd_border_rect_mask_a = np.expand_dims (img_prd_border_rect_mask_a, -1)
|
||||
#
|
||||
# out_img = out_img * img_prd_border_rect_mask_a + img_bgr * (1.0 - img_prd_border_rect_mask_a)
|
||||
# img_mask_blurry_aaa *= img_prd_border_rect_mask_a
|
||||
|
||||
|
||||
out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (out_img*img_mask_blurry_aaa) , 0, 1.0 )
|
||||
|
||||
if self.mode == 'seamless-hist-match':
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from pathlib import Path
|
|||
from utils import Path_utils
|
||||
from utils import std_utils
|
||||
from utils import image_utils
|
||||
from utils.console_utils import *
|
||||
import numpy as np
|
||||
import cv2
|
||||
from samples import SampleGeneratorBase
|
||||
|
|
@ -18,8 +19,11 @@ class ModelBase(object):
|
|||
|
||||
#DONT OVERRIDE
|
||||
def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None,
|
||||
batch_size=0,
|
||||
write_preview_history = False,
|
||||
ask_for_session_options=False,
|
||||
session_write_preview_history = None,
|
||||
session_target_epoch=0,
|
||||
session_batch_size=0,
|
||||
|
||||
debug = False, **in_options
|
||||
):
|
||||
print ("Loading model...")
|
||||
|
|
@ -35,56 +39,94 @@ class ModelBase(object):
|
|||
self.dst_yaw_images_paths = None
|
||||
self.src_data_generator = None
|
||||
self.dst_data_generator = None
|
||||
self.is_training_mode = (training_data_src_path is not None and training_data_dst_path is not None)
|
||||
self.batch_size = batch_size
|
||||
self.write_preview_history = write_preview_history
|
||||
self.debug = debug
|
||||
self.is_training_mode = (training_data_src_path is not None and training_data_dst_path is not None)
|
||||
|
||||
self.supress_std_once = ('TF_SUPPRESS_STD' in os.environ.keys() and os.environ['TF_SUPPRESS_STD'] == '1')
|
||||
|
||||
self.epoch = 0
|
||||
self.options = {}
|
||||
self.loss_history = []
|
||||
self.sample_for_preview = None
|
||||
if self.model_data_path.exists():
|
||||
model_data = pickle.loads ( self.model_data_path.read_bytes() )
|
||||
self.epoch = model_data['epoch']
|
||||
self.options = model_data['options']
|
||||
self.loss_history = model_data['loss_history'] if 'loss_history' in model_data.keys() else []
|
||||
self.sample_for_preview = model_data['sample_for_preview'] if 'sample_for_preview' in model_data.keys() else None
|
||||
else:
|
||||
self.epoch = 0
|
||||
self.options = {}
|
||||
self.loss_history = []
|
||||
self.sample_for_preview = None
|
||||
if self.epoch != 0:
|
||||
self.options = model_data['options']
|
||||
self.loss_history = model_data['loss_history'] if 'loss_history' in model_data.keys() else []
|
||||
self.sample_for_preview = model_data['sample_for_preview'] if 'sample_for_preview' in model_data.keys() else None
|
||||
|
||||
if self.write_preview_history:
|
||||
self.preview_history_path = self.model_path / ( '%s_history' % (self.get_model_name()) )
|
||||
|
||||
if not self.preview_history_path.exists():
|
||||
self.preview_history_path.mkdir(exist_ok=True)
|
||||
else:
|
||||
if self.epoch == 0:
|
||||
for filename in Path_utils.get_image_paths(self.preview_history_path):
|
||||
Path(filename).unlink()
|
||||
|
||||
self.device_config = nnlib.DeviceConfig(allow_growth=False, **in_options)
|
||||
|
||||
if self.epoch == 0:
|
||||
#first run
|
||||
self.options['created_vram_gb'] = self.device_config.gpu_total_vram_gb
|
||||
self.created_vram_gb = self.device_config.gpu_total_vram_gb
|
||||
print ("\nModel first run. Enter model options as default for each run.")
|
||||
self.options['write_preview_history'] = input_bool("Write preview history? (y/n skip:n) : ", False)
|
||||
self.options['target_epoch'] = max(0, input_int("Target epoch (skip:unlimited) : ", 0))
|
||||
self.options['batch_size'] = max(0, input_int("Batch_size (skip:model choice) : ", 0))
|
||||
self.options['sort_by_yaw'] = input_bool("Feed faces to network sorted by yaw? (y/n skip:n) : ", False)
|
||||
|
||||
#self.options['use_fp16'] = use_fp16 = input_bool("Use float16? (y/n skip:n) : ", False)
|
||||
else:
|
||||
#not first run
|
||||
if 'created_vram_gb' in self.options.keys():
|
||||
self.created_vram_gb = self.options['created_vram_gb']
|
||||
else:
|
||||
self.options['created_vram_gb'] = self.device_config.gpu_total_vram_gb
|
||||
self.created_vram_gb = self.device_config.gpu_total_vram_gb
|
||||
self.options['write_preview_history'] = self.options.get('write_preview_history', False)
|
||||
self.options['target_epoch'] = self.options.get('target_epoch', 0)
|
||||
self.options['batch_size'] = self.options.get('batch_size', 0)
|
||||
self.options['sort_by_yaw'] = self.options.get('sort_by_yaw', False)
|
||||
#self.options['use_fp16'] = use_fp16 = self.options['use_fp16'] if 'use_fp16' in self.options.keys() else False
|
||||
|
||||
use_fp16 = False #currently models fails with fp16
|
||||
|
||||
if ask_for_session_options:
|
||||
print ("Override options for current session:")
|
||||
session_write_preview_history = input_bool("Write preview history? (y/n skip:default) : ", None )
|
||||
session_target_epoch = input_int("Target epoch (skip:default) : ", 0)
|
||||
session_batch_size = input_int("Batch_size (skip:default) : ", 0)
|
||||
|
||||
if self.options['write_preview_history']:
|
||||
if session_write_preview_history is None:
|
||||
session_write_preview_history = self.options['write_preview_history']
|
||||
else:
|
||||
self.options.pop('write_preview_history')
|
||||
|
||||
if self.options['target_epoch'] != 0:
|
||||
if session_target_epoch == 0:
|
||||
session_target_epoch = self.options['target_epoch']
|
||||
else:
|
||||
self.options.pop('target_epoch')
|
||||
|
||||
if self.options['batch_size'] != 0:
|
||||
if session_batch_size == 0:
|
||||
session_batch_size = self.options['batch_size']
|
||||
else:
|
||||
self.options.pop('batch_size')
|
||||
|
||||
self.sort_by_yaw = self.options['sort_by_yaw']
|
||||
if not self.sort_by_yaw:
|
||||
self.options.pop('sort_by_yaw')
|
||||
|
||||
self.write_preview_history = session_write_preview_history
|
||||
self.target_epoch = session_target_epoch
|
||||
self.batch_size = session_batch_size
|
||||
|
||||
self.device_config = nnlib.DeviceConfig(allow_growth=False, use_fp16=use_fp16, **in_options)
|
||||
|
||||
self.created_vram_gb = self.options['created_vram_gb'] if 'created_vram_gb' in self.options.keys() else self.device_config.gpu_total_vram_gb
|
||||
|
||||
self.onInitializeOptions(self.epoch == 0, ask_for_session_options)
|
||||
nnlib.import_all (self.device_config)
|
||||
|
||||
self.onInitialize(**in_options)
|
||||
|
||||
if self.debug or self.batch_size == 0:
|
||||
self.batch_size = 1
|
||||
|
||||
if self.is_training_mode:
|
||||
if self.write_preview_history:
|
||||
self.preview_history_path = self.model_path / ( '%s_history' % (self.get_model_name()) )
|
||||
|
||||
if not self.preview_history_path.exists():
|
||||
self.preview_history_path.mkdir(exist_ok=True)
|
||||
else:
|
||||
if self.epoch == 0:
|
||||
for filename in Path_utils.get_image_paths(self.preview_history_path):
|
||||
Path(filename).unlink()
|
||||
|
||||
if self.generator_list is None:
|
||||
raise Exception( 'You didnt set_training_data_generators()')
|
||||
else:
|
||||
|
|
@ -100,11 +142,18 @@ class ModelBase(object):
|
|||
print ("==")
|
||||
print ("== Current epoch: " + str(self.epoch) )
|
||||
print ("==")
|
||||
print ("== Options:")
|
||||
print ("== |== batch_size : %s " % (self.batch_size) )
|
||||
print ("== |== multi_gpu : %s " % (self.device_config.multi_gpu) )
|
||||
print ("== Model options:")
|
||||
for key in self.options.keys():
|
||||
print ("== |== %s : %s" % (key, self.options[key]) )
|
||||
print ("== Session options:")
|
||||
if self.write_preview_history:
|
||||
print ("== |== write_preview_history : True ")
|
||||
if self.target_epoch != 0:
|
||||
print ("== |== target_epoch : %s " % (self.target_epoch) )
|
||||
print ("== |== batch_size : %s " % (self.batch_size) )
|
||||
if self.device_config.multi_gpu:
|
||||
print ("== |== multi_gpu : True ")
|
||||
|
||||
|
||||
print ("== Running on:")
|
||||
if self.device_config.cpu_only:
|
||||
|
|
@ -122,6 +171,10 @@ class ModelBase(object):
|
|||
|
||||
print ("=========================")
|
||||
|
||||
#overridable
|
||||
def onInitializeOptions(self, is_first_run, ask_for_session_options):
|
||||
pass
|
||||
|
||||
#overridable
|
||||
def onInitialize(self, **in_options):
|
||||
'''
|
||||
|
|
@ -161,6 +214,12 @@ class ModelBase(object):
|
|||
from .ConverterBase import ConverterBase
|
||||
return ConverterBase(self, **in_options)
|
||||
|
||||
def get_target_epoch(self):
|
||||
return self.target_epoch
|
||||
|
||||
def is_reached_epoch_goal(self):
|
||||
return self.target_epoch != 0 and self.epoch >= self.target_epoch
|
||||
|
||||
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
|
||||
|
|
@ -305,9 +364,6 @@ class ModelBase(object):
|
|||
if self.batch_size == 0:
|
||||
self.batch_size = 2
|
||||
else:
|
||||
if self.device_config.gpu_total_vram_gb < keys[0]:
|
||||
raise Exception ('Sorry, this model works only on %dGB+ GPU' % ( keys[0] ) )
|
||||
|
||||
if self.batch_size == 0:
|
||||
for x in keys:
|
||||
if self.device_config.gpu_total_vram_gb <= x:
|
||||
|
|
|
|||
|
|
@ -1,251 +0,0 @@
|
|||
import numpy as np
|
||||
import cv2
|
||||
from models import ModelBase
|
||||
from samples import *
|
||||
from nnlib import nnlib
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
encoder64H5 = 'encoder64.h5'
|
||||
decoder64_srcH5 = 'decoder64_src.h5'
|
||||
decoder64_dstH5 = 'decoder64_dst.h5'
|
||||
encoder256H5 = 'encoder256.h5'
|
||||
decoder256H5 = 'decoder256.h5'
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
|
||||
self.set_vram_batch_requirements( {3.5:8,4:8,5:12,6:16,7:24,8:32,9:48} )
|
||||
if self.batch_size < 4:
|
||||
self.batch_size = 4
|
||||
|
||||
img_shape64, img_shape256, self.encoder64, self.decoder64_src, self.decoder64_dst, self.encoder256, self.decoder256 = self.Build()
|
||||
|
||||
if not self.is_first_run():
|
||||
self.encoder64.load_weights (self.get_strpath_storage_for_file(self.encoder64H5))
|
||||
self.decoder64_src.load_weights (self.get_strpath_storage_for_file(self.decoder64_srcH5))
|
||||
self.decoder64_dst.load_weights (self.get_strpath_storage_for_file(self.decoder64_dstH5))
|
||||
self.encoder256.load_weights (self.get_strpath_storage_for_file(self.encoder256H5))
|
||||
self.decoder256.load_weights (self.get_strpath_storage_for_file(self.decoder256H5))
|
||||
|
||||
#if self.is_training_mode:
|
||||
# self.encoder64, self.decoder64_src, self.decoder64_dst, self.encoder256, self.decoder256 = self.to_multi_gpu_model_if_possible ( [self.encoder64, self.decoder64_src, self.decoder64_dst, self.encoder256, self.decoder256] )
|
||||
|
||||
input_A_warped64 = Input(img_shape64)
|
||||
input_B_warped64 = Input(img_shape64)
|
||||
A_rec64 = self.decoder64_src(self.encoder64(input_A_warped64))
|
||||
B_rec64 = self.decoder64_dst(self.encoder64(input_B_warped64))
|
||||
self.ae64 = Model([input_A_warped64, input_B_warped64], [A_rec64, B_rec64] )
|
||||
|
||||
if self.is_training_mode:
|
||||
self.ae64, = self.to_multi_gpu_model_if_possible ( [self.ae64,] )
|
||||
|
||||
self.ae64.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[DSSIMLoss(), DSSIMLoss()] )
|
||||
|
||||
self.A64_view = K.function ([input_A_warped64], [A_rec64])
|
||||
self.B64_view = K.function ([input_B_warped64], [B_rec64])
|
||||
|
||||
input_A_warped64 = Input(img_shape64)
|
||||
input_A_target256 = Input(img_shape256)
|
||||
A_rec256 = self.decoder256( self.encoder256(input_A_warped64) )
|
||||
|
||||
input_B_warped64 = Input(img_shape64)
|
||||
BA_rec64 = self.decoder64_src( self.encoder64(input_B_warped64) )
|
||||
BA_rec256 = self.decoder256( self.encoder256(BA_rec64) )
|
||||
|
||||
self.ae256 = Model([input_A_warped64], [A_rec256] )
|
||||
|
||||
if self.is_training_mode:
|
||||
self.ae256, = self.to_multi_gpu_model_if_possible ( [self.ae256,] )
|
||||
|
||||
self.ae256.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[DSSIMLoss()])
|
||||
|
||||
self.A256_view = K.function ([input_A_warped64], [A_rec256])
|
||||
self.BA256_view = K.function ([input_B_warped64], [BA_rec256])
|
||||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[
|
||||
[f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 256],
|
||||
[f.SOURCE | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.SOURCE | f.FACE_ALIGN_HALF | f.MODE_BGR, 256] ] ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[
|
||||
[f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.SOURCE | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.SOURCE | f.FACE_ALIGN_HALF | f.MODE_BGR, 256] ] )
|
||||
])
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( [[self.encoder64, self.get_strpath_storage_for_file(self.encoder64H5)],
|
||||
[self.decoder64_src, self.get_strpath_storage_for_file(self.decoder64_srcH5)],
|
||||
[self.decoder64_dst, self.get_strpath_storage_for_file(self.decoder64_dstH5)],
|
||||
[self.encoder256, self.get_strpath_storage_for_file(self.encoder256H5)],
|
||||
[self.decoder256, self.get_strpath_storage_for_file(self.decoder256H5)],
|
||||
] )
|
||||
|
||||
#override
|
||||
def onTrainOneEpoch(self, sample):
|
||||
warped_src64, target_src64, target_src256, target_src_source64, target_src_source256 = sample[0]
|
||||
warped_dst64, target_dst64, target_dst_source64, target_dst_source256 = sample[1]
|
||||
|
||||
loss64, loss_src64, loss_dst64 = self.ae64.train_on_batch ([warped_src64, warped_dst64], [target_src64, target_dst64])
|
||||
|
||||
loss256 = self.ae256.train_on_batch ([warped_src64], [target_src256])
|
||||
|
||||
return ( ('loss64', loss64 ), ('loss256', loss256), )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
sample_src64_source = sample[0][3][0:4]
|
||||
sample_src256_source = sample[0][4][0:4]
|
||||
|
||||
sample_dst64_source = sample[1][2][0:4]
|
||||
sample_dst256_source = sample[1][3][0:4]
|
||||
|
||||
SRC64, = self.A64_view ([sample_src64_source])
|
||||
DST64, = self.B64_view ([sample_dst64_source])
|
||||
SRCDST64, = self.A64_view ([sample_dst64_source])
|
||||
DSTSRC64, = self.B64_view ([sample_src64_source])
|
||||
|
||||
SRC_x1_256, = self.A256_view ([sample_src64_source])
|
||||
DST_x2_256, = self.BA256_view ([sample_dst64_source])
|
||||
|
||||
b1 = np.concatenate ( (
|
||||
np.concatenate ( (sample_src64_source[0], SRC64[0], sample_src64_source[1], SRC64[1], ), axis=1),
|
||||
np.concatenate ( (sample_src64_source[1], SRC64[1], sample_src64_source[3], SRC64[3], ), axis=1),
|
||||
np.concatenate ( (sample_dst64_source[0], DST64[0], sample_dst64_source[1], DST64[1], ), axis=1),
|
||||
np.concatenate ( (sample_dst64_source[2], DST64[2], sample_dst64_source[3], DST64[3], ), axis=1),
|
||||
), axis=0 )
|
||||
|
||||
b2 = np.concatenate ( (
|
||||
np.concatenate ( (sample_src64_source[0], DSTSRC64[0], sample_src64_source[1], DSTSRC64[1], ), axis=1),
|
||||
np.concatenate ( (sample_src64_source[2], DSTSRC64[2], sample_src64_source[3], DSTSRC64[3], ), axis=1),
|
||||
np.concatenate ( (sample_dst64_source[0], SRCDST64[0], sample_dst64_source[1], SRCDST64[1], ), axis=1),
|
||||
np.concatenate ( (sample_dst64_source[2], SRCDST64[2], sample_dst64_source[3], SRCDST64[3], ), axis=1),
|
||||
|
||||
), axis=0 )
|
||||
|
||||
result = np.concatenate ( ( np.concatenate ( (b1, sample_src256_source[0], SRC_x1_256[0] ), axis=1 ),
|
||||
np.concatenate ( (b2, sample_dst256_source[0], DST_x2_256[0] ), axis=1 ),
|
||||
), axis = 0 )
|
||||
|
||||
return [ ('AVATAR', result ) ]
|
||||
|
||||
def predictor_func (self, img):
|
||||
x, = self.BA256_view ([ np.expand_dims(img, 0) ])[0]
|
||||
return x
|
||||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
return ConverterAvatar(self.predictor_func, predictor_input_size=64, output_size=256, **in_options)
|
||||
|
||||
def Build(self):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
img_shape64 = (64,64,3)
|
||||
img_shape256 = (256,256,3)
|
||||
|
||||
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):
|
||||
x = _input
|
||||
x = Conv2D(90, kernel_size=5, strides=1, padding='same')(x)
|
||||
x = Conv2D(90, kernel_size=5, strides=1, padding='same')(x)
|
||||
x = MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
|
||||
|
||||
x = Conv2D(180, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = Conv2D(180, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
|
||||
|
||||
x = Conv2D(360, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = Conv2D(360, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
|
||||
|
||||
x = Dense (1024)(x)
|
||||
x = LeakyReLU(0.1)(x)
|
||||
x = Dropout(0.5)(x)
|
||||
|
||||
x = Dense (1024)(x)
|
||||
x = LeakyReLU(0.1)(x)
|
||||
x = Dropout(0.5)(x)
|
||||
x = Flatten()(x)
|
||||
x = Dense (64)(x)
|
||||
|
||||
return keras.models.Model (_input, x)
|
||||
|
||||
encoder256 = Encoder( Input (img_shape64) )
|
||||
encoder64 = Encoder( Input (img_shape64) )
|
||||
|
||||
def decoder256(encoder):
|
||||
decoder_input = Input ( K.int_shape(encoder.outputs[0])[1:] )
|
||||
x = decoder_input
|
||||
x = Dense(16 * 16 * 720)(x)
|
||||
x = Reshape ( (16, 16, 720) )(x)
|
||||
x = upscale(720)(x)
|
||||
x = upscale(360)(x)
|
||||
x = upscale(180)(x)
|
||||
x = upscale(90)(x)
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return keras.models.Model(decoder_input, x)
|
||||
|
||||
def decoder64(encoder):
|
||||
decoder_input = Input ( K.int_shape(encoder.outputs[0])[1:] )
|
||||
x = decoder_input
|
||||
x = Dense(8 * 8 * 720)(x)
|
||||
x = Reshape ( (8, 8, 720) )(x)
|
||||
x = upscale(360)(x)
|
||||
x = upscale(180)(x)
|
||||
x = upscale(90)(x)
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return Model(decoder_input, x)
|
||||
|
||||
return img_shape64, img_shape256, encoder64, decoder64(encoder64), decoder64(encoder64), encoder256, decoder256(encoder256)
|
||||
|
||||
from models import ConverterBase
|
||||
from facelib import FaceType
|
||||
from facelib import LandmarksProcessor
|
||||
class ConverterAvatar(ConverterBase):
|
||||
|
||||
#override
|
||||
def __init__(self, predictor,
|
||||
predictor_input_size=0,
|
||||
output_size=0,
|
||||
**in_options):
|
||||
|
||||
super().__init__(predictor)
|
||||
|
||||
self.predictor_input_size = predictor_input_size
|
||||
self.output_size = output_size
|
||||
|
||||
#override
|
||||
def get_mode(self):
|
||||
return ConverterBase.MODE_IMAGE_WITH_LANDMARKS
|
||||
|
||||
#override
|
||||
def dummy_predict(self):
|
||||
self.predictor ( np.zeros ( (self.predictor_input_size, self.predictor_input_size,3), dtype=np.float32) )
|
||||
|
||||
#override
|
||||
def convert_image (self, img_bgr, img_face_landmarks, debug):
|
||||
img_size = img_bgr.shape[1], img_bgr.shape[0]
|
||||
|
||||
face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, self.predictor_input_size, face_type=FaceType.HALF )
|
||||
predictor_input_bgr = cv2.warpAffine( img_bgr, face_mat, (self.predictor_input_size, self.predictor_input_size), flags=cv2.INTER_LANCZOS4 )
|
||||
|
||||
predicted_bgr = self.predictor ( predictor_input_bgr )
|
||||
|
||||
output = cv2.resize ( predicted_bgr, (self.output_size, self.output_size), cv2.INTER_LANCZOS4 )
|
||||
if debug:
|
||||
return (img_bgr,output,)
|
||||
return output
|
||||
|
|
@ -38,7 +38,8 @@ class Model(ModelBase):
|
|||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
|
||||
debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
|
||||
|
|
@ -107,16 +108,14 @@ class Model(ModelBase):
|
|||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterMasked
|
||||
|
||||
if 'erode_mask_modifier' not in in_options.keys():
|
||||
in_options['erode_mask_modifier'] = 0
|
||||
in_options['erode_mask_modifier'] += 30
|
||||
|
||||
if 'blur_mask_modifier' not in in_options.keys():
|
||||
in_options['blur_mask_modifier'] = 0
|
||||
|
||||
return ConverterMasked(self.predictor_func, predictor_input_size=128, output_size=128, face_type=FaceType.FULL, clip_border_mask_per=0.046875, **in_options)
|
||||
from models import ConverterMasked
|
||||
return ConverterMasked(self.predictor_func,
|
||||
predictor_input_size=128,
|
||||
output_size=128,
|
||||
face_type=FaceType.FULL,
|
||||
base_erode_mask_modifier=30,
|
||||
base_blur_mask_modifier=100,
|
||||
**in_options)
|
||||
|
||||
def Build(self, input_layer):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
|
|
|||
|
|
@ -44,7 +44,8 @@ class Model(ModelBase):
|
|||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
|
||||
debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
|
||||
|
|
@ -112,16 +113,13 @@ class Model(ModelBase):
|
|||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterMasked
|
||||
|
||||
if 'erode_mask_modifier' not in in_options.keys():
|
||||
in_options['erode_mask_modifier'] = 0
|
||||
in_options['erode_mask_modifier'] += 100
|
||||
|
||||
if 'blur_mask_modifier' not in in_options.keys():
|
||||
in_options['blur_mask_modifier'] = 0
|
||||
in_options['blur_mask_modifier'] += 100
|
||||
|
||||
return ConverterMasked(self.predictor_func, predictor_input_size=128, output_size=128, face_type=FaceType.HALF, **in_options)
|
||||
return ConverterMasked(self.predictor_func,
|
||||
predictor_input_size=128,
|
||||
output_size=128,
|
||||
face_type=FaceType.HALF,
|
||||
base_erode_mask_modifier=100,
|
||||
base_blur_mask_modifier=100,
|
||||
**in_options)
|
||||
|
||||
def Build(self, created_vram_gb):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from nnlib import nnlib
|
|||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
from utils.console_utils import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
|
|
@ -44,7 +45,8 @@ class Model(ModelBase):
|
|||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
|
||||
debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 64] ] ),
|
||||
|
|
@ -66,7 +68,6 @@ class Model(ModelBase):
|
|||
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) )
|
||||
|
|
@ -114,16 +115,13 @@ class Model(ModelBase):
|
|||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterMasked
|
||||
|
||||
if 'erode_mask_modifier' not in in_options.keys():
|
||||
in_options['erode_mask_modifier'] = 0
|
||||
in_options['erode_mask_modifier'] += 100
|
||||
|
||||
if 'blur_mask_modifier' not in in_options.keys():
|
||||
in_options['blur_mask_modifier'] = 0
|
||||
in_options['blur_mask_modifier'] += 100
|
||||
|
||||
return ConverterMasked(self.predictor_func, predictor_input_size=64, output_size=64, face_type=FaceType.HALF, **in_options)
|
||||
return ConverterMasked(self.predictor_func,
|
||||
predictor_input_size=64,
|
||||
output_size=64,
|
||||
face_type=FaceType.HALF,
|
||||
base_erode_mask_modifier=100,
|
||||
base_blur_mask_modifier=100,
|
||||
**in_options)
|
||||
|
||||
def Build(self, created_vram_gb):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
|
|
|||
|
|
@ -42,8 +42,11 @@ class Model(ModelBase):
|
|||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
self.set_training_data_generators ([
|
||||
|
||||
|
||||
SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
|
||||
debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
|
||||
|
|
@ -115,15 +118,13 @@ class Model(ModelBase):
|
|||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterMasked
|
||||
|
||||
if 'erode_mask_modifier' not in in_options.keys():
|
||||
in_options['erode_mask_modifier'] = 0
|
||||
in_options['erode_mask_modifier'] += 30
|
||||
|
||||
if 'blur_mask_modifier' not in in_options.keys():
|
||||
in_options['blur_mask_modifier'] = 0
|
||||
|
||||
return ConverterMasked(self.predictor_func, predictor_input_size=128, output_size=128, face_type=FaceType.FULL, clip_border_mask_per=0.046875, **in_options)
|
||||
return ConverterMasked(self.predictor_func,
|
||||
predictor_input_size=128,
|
||||
output_size=128,
|
||||
face_type=FaceType.FULL,
|
||||
base_erode_mask_modifier=30,
|
||||
base_blur_mask_modifier=0,
|
||||
**in_options)
|
||||
|
||||
def Build(self, input_layer):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
|
|
|||
|
|
@ -1,175 +0,0 @@
|
|||
import numpy as np
|
||||
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
encoderH5 = 'encoder.h5'
|
||||
decoderH5 = 'decoder.h5'
|
||||
inter_BH5 = 'inter_B.h5'
|
||||
inter_ABH5 = 'inter_AB.h5'
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {4.5:4,5:4,6:8,7:12,8:16,9:20,10:24,11:24,12:32,13:48} )
|
||||
|
||||
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():
|
||||
self.encoder.load_weights (self.get_strpath_storage_for_file(self.encoderH5))
|
||||
self.decoder.load_weights (self.get_strpath_storage_for_file(self.decoderH5))
|
||||
self.inter_B.load_weights (self.get_strpath_storage_for_file(self.inter_BH5))
|
||||
self.inter_AB.load_weights (self.get_strpath_storage_for_file(self.inter_ABH5))
|
||||
|
||||
code = self.encoder(ae_input_layer)
|
||||
AB = self.inter_AB(code)
|
||||
B = self.inter_B(code)
|
||||
self.autoencoder_src = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([AB, AB])) )
|
||||
self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([B, AB])) )
|
||||
|
||||
if self.is_training_mode:
|
||||
self.autoencoder_src, self.autoencoder_dst = self.to_multi_gpu_model_if_possible ( [self.autoencoder_src, self.autoencoder_dst] )
|
||||
|
||||
self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] )
|
||||
])
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( [[self.encoder, self.get_strpath_storage_for_file(self.encoderH5)],
|
||||
[self.decoder, self.get_strpath_storage_for_file(self.decoderH5)],
|
||||
[self.inter_B, self.get_strpath_storage_for_file(self.inter_BH5)],
|
||||
[self.inter_AB, self.get_strpath_storage_for_file(self.inter_ABH5)]] )
|
||||
|
||||
#override
|
||||
def onTrainOneEpoch(self, sample):
|
||||
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 [ ('LIAEF128YAW', np.concatenate ( st, axis=0 ) ) ]
|
||||
|
||||
def predictor_func (self, face):
|
||||
|
||||
face_128_bgr = face[...,0:3]
|
||||
face_128_mask = np.expand_dims(face[...,3],-1)
|
||||
|
||||
x, mx = self.autoencoder_src.predict ( [ np.expand_dims(face_128_bgr,0), np.expand_dims(face_128_mask,0) ] )
|
||||
x, mx = x[0], mx[0]
|
||||
|
||||
return np.concatenate ( (x,mx), -1 )
|
||||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterMasked
|
||||
|
||||
if 'erode_mask_modifier' not in in_options.keys():
|
||||
in_options['erode_mask_modifier'] = 0
|
||||
in_options['erode_mask_modifier'] += 30
|
||||
|
||||
if 'blur_mask_modifier' not in in_options.keys():
|
||||
in_options['blur_mask_modifier'] = 0
|
||||
|
||||
return ConverterMasked(self.predictor_func, predictor_input_size=128, output_size=128, face_type=FaceType.FULL, clip_border_mask_per=0.046875, **in_options)
|
||||
|
||||
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,225 +0,0 @@
|
|||
import numpy as np
|
||||
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
encoderH5 = 'encoder.h5'
|
||||
decoderMaskH5 = 'decoderMask.h5'
|
||||
decoderCommonAH5 = 'decoderCommonA.h5'
|
||||
decoderCommonBH5 = 'decoderCommonB.h5'
|
||||
decoderRGBH5 = 'decoderRGB.h5'
|
||||
decoderBWH5 = 'decoderBW.h5'
|
||||
inter_BH5 = 'inter_B.h5'
|
||||
inter_AH5 = 'inter_A.h5'
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {4.5:4,5:4,6:8,7:12,8:16,9:20,10:24,11:24,12:32,13:48} )
|
||||
|
||||
ae_input_layer = Input(shape=(128, 128, 3))
|
||||
mask_layer = Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder, self.decoderMask, self.decoderCommonA, self.decoderCommonB, self.decoderRGB, \
|
||||
self.decoderBW, self.inter_A, self.inter_B = self.Build(ae_input_layer)
|
||||
|
||||
if not self.is_first_run():
|
||||
self.encoder.load_weights (self.get_strpath_storage_for_file(self.encoderH5))
|
||||
self.decoderMask.load_weights (self.get_strpath_storage_for_file(self.decoderMaskH5))
|
||||
self.decoderCommonA.load_weights (self.get_strpath_storage_for_file(self.decoderCommonAH5))
|
||||
self.decoderCommonB.load_weights (self.get_strpath_storage_for_file(self.decoderCommonBH5))
|
||||
self.decoderRGB.load_weights (self.get_strpath_storage_for_file(self.decoderRGBH5))
|
||||
self.decoderBW.load_weights (self.get_strpath_storage_for_file(self.decoderBWH5))
|
||||
self.inter_A.load_weights (self.get_strpath_storage_for_file(self.inter_AH5))
|
||||
self.inter_B.load_weights (self.get_strpath_storage_for_file(self.inter_BH5))
|
||||
|
||||
code = self.encoder(ae_input_layer)
|
||||
A = self.inter_A(code)
|
||||
B = self.inter_B(code)
|
||||
|
||||
inter_A_A = Concatenate()([A, A])
|
||||
inter_B_A = Concatenate()([B, A])
|
||||
|
||||
x1,m1 = self.decoderCommonA (inter_A_A)
|
||||
x2,m2 = self.decoderCommonA (inter_A_A)
|
||||
self.autoencoder_src = Model([ae_input_layer,mask_layer],
|
||||
[ self.decoderBW (Concatenate()([x1,x2]) ),
|
||||
self.decoderMask(Concatenate()([m1,m2]) )
|
||||
])
|
||||
|
||||
x1,m1 = self.decoderCommonA (inter_A_A)
|
||||
x2,m2 = self.decoderCommonB (inter_A_A)
|
||||
self.autoencoder_src_RGB = Model([ae_input_layer,mask_layer],
|
||||
[ self.decoderRGB (Concatenate()([x1,x2]) ),
|
||||
self.decoderMask (Concatenate()([m1,m2]) )
|
||||
])
|
||||
|
||||
x1,m1 = self.decoderCommonA (inter_B_A)
|
||||
x2,m2 = self.decoderCommonB (inter_B_A)
|
||||
self.autoencoder_dst = Model([ae_input_layer,mask_layer],
|
||||
[ self.decoderRGB (Concatenate()([x1,x2]) ),
|
||||
self.decoderMask (Concatenate()([m1,m2]) )
|
||||
])
|
||||
|
||||
if self.is_training_mode:
|
||||
self.autoencoder_src, self.autoencoder_dst = self.to_multi_gpu_model_if_possible ( [self.autoencoder_src, self.autoencoder_dst] )
|
||||
|
||||
self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_GGG, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_G , 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_GGG, 128] ] ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
|
||||
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128]] )
|
||||
])
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( [[self.encoder, self.get_strpath_storage_for_file(self.encoderH5)],
|
||||
[self.decoderMask, self.get_strpath_storage_for_file(self.decoderMaskH5)],
|
||||
[self.decoderCommonA, self.get_strpath_storage_for_file(self.decoderCommonAH5)],
|
||||
[self.decoderCommonB, self.get_strpath_storage_for_file(self.decoderCommonBH5)],
|
||||
[self.decoderRGB, self.get_strpath_storage_for_file(self.decoderRGBH5)],
|
||||
[self.decoderBW, self.get_strpath_storage_for_file(self.decoderBWH5)],
|
||||
[self.inter_A, self.get_strpath_storage_for_file(self.inter_AH5)],
|
||||
[self.inter_B, self.get_strpath_storage_for_file(self.inter_BH5)]] )
|
||||
|
||||
|
||||
#override
|
||||
def onTrainOneEpoch(self, sample):
|
||||
warped_src, target_src, target_src_mask, target_src_GGG = 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][3][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_RGB.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 ( (
|
||||
np.repeat (np.expand_dims (test_A[i,:,:,0],-1), (3,), -1) ,
|
||||
np.repeat (AA[i], (3,), -1),
|
||||
#mAA[i],
|
||||
test_B[i,:,:,0:3],
|
||||
BB[i],
|
||||
#mBB[i],
|
||||
AB[i],
|
||||
#mAB[i]
|
||||
), axis=1) )
|
||||
|
||||
return [ ('MIAEF128', np.concatenate ( st, axis=0 ) ) ]
|
||||
|
||||
def predictor_func (self, face):
|
||||
face_128_bgr = face[...,0:3]
|
||||
face_128_mask = np.expand_dims(face[...,-1],-1)
|
||||
|
||||
x, mx = self.autoencoder_src_RGB.predict ( [ np.expand_dims(face_128_bgr,0), np.expand_dims(face_128_mask,0) ] )
|
||||
x, mx = x[0], mx[0]
|
||||
|
||||
return np.concatenate ( (x,mx), -1 )
|
||||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterMasked
|
||||
|
||||
if 'erode_mask_modifier' not in in_options.keys():
|
||||
in_options['erode_mask_modifier'] = 0
|
||||
in_options['erode_mask_modifier'] += 30
|
||||
|
||||
if 'blur_mask_modifier' not in in_options.keys():
|
||||
in_options['blur_mask_modifier'] = 0
|
||||
|
||||
return ConverterMasked(self.predictor_func, predictor_input_size=128, output_size=128, face_type=FaceType.FULL, clip_border_mask_per=0.046875, **in_options)
|
||||
|
||||
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 DecoderCommon():
|
||||
input_ = Input(shape=(16, 16, 1024))
|
||||
x = input_
|
||||
x = upscale(512)(x)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
|
||||
y = input_
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
y = upscale(64)(y)
|
||||
|
||||
return Model(input_, [x,y])
|
||||
|
||||
def DecoderRGB():
|
||||
input_ = Input(shape=(128, 128, 256))
|
||||
x = input_
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return Model(input_, [x])
|
||||
|
||||
def DecoderBW():
|
||||
input_ = Input(shape=(128, 128, 256))
|
||||
x = input_
|
||||
x = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return Model(input_, [x])
|
||||
|
||||
def DecoderMask():
|
||||
input_ = Input(shape=(128, 128, 128))
|
||||
y = input_
|
||||
y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
return Model(input_, [y])
|
||||
|
||||
return Encoder(), DecoderMask(), DecoderCommon(), DecoderCommon(), DecoderRGB(), DecoderBW(), Intermediate(), Intermediate()
|
||||
|
|
@ -1 +0,0 @@
|
|||
from .Model import Model
|
||||
|
|
@ -25,7 +25,7 @@ class Model(ModelBase):
|
|||
if self.epoch == 0:
|
||||
#first run
|
||||
|
||||
print ("\nModel first run. Enter options.")
|
||||
|
||||
|
||||
try:
|
||||
created_resolution = int ( input ("Resolution (default:64, valid: 64,128,256) : ") )
|
||||
|
|
@ -68,9 +68,9 @@ class Model(ModelBase):
|
|||
|
||||
self.set_batch_size(created_batch_size)
|
||||
|
||||
use_batch_norm = created_batch_size > 1
|
||||
self.GA = modelify(ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=False))(Input(bgr_shape))
|
||||
self.GB = modelify(ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=False))(Input(bgr_shape))
|
||||
use_batch_norm = False #created_batch_size > 1
|
||||
self.GA = modelify(ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))(Input(bgr_shape))
|
||||
self.GB = modelify(ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))(Input(bgr_shape))
|
||||
#self.GA = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
|
||||
#self.GB = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
|
||||
|
||||
|
|
@ -211,7 +211,7 @@ class Model(ModelBase):
|
|||
loss_G, = self.G_train ( feed )
|
||||
loss_DA, = self.DA_train( feed )
|
||||
loss_DB, = self.DB_train( feed )
|
||||
|
||||
#return ( ('G', loss_G), )
|
||||
return ( ('G', loss_G), ('DA', loss_DA), ('DB', loss_DB) )
|
||||
|
||||
#override
|
||||
|
|
@ -242,7 +242,9 @@ class Model(ModelBase):
|
|||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterImage
|
||||
|
||||
return ConverterImage(self.predictor_func, predictor_input_size=self.options['created_resolution'], output_size=self.options['created_resolution'], **in_options)
|
||||
from models import ConverterImage
|
||||
return ConverterImage(self.predictor_func,
|
||||
predictor_input_size=self.options['created_resolution'],
|
||||
output_size=self.options['created_resolution'],
|
||||
**in_options)
|
||||
|
||||
|
|
|
|||
298
models/Model_UFM/Model.py
Normal file
298
models/Model_UFM/Model.py
Normal file
|
|
@ -0,0 +1,298 @@
|
|||
import numpy as np
|
||||
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
from utils.console_utils import *
|
||||
|
||||
#U-net Face Morpher
|
||||
class UFMModel(ModelBase):
|
||||
|
||||
encoderH5 = 'encoder.h5'
|
||||
decoder_srcH5 = 'decoder_src.h5'
|
||||
decoder_dstH5 = 'decoder_dst.h5'
|
||||
decoder_srcmH5 = 'decoder_srcm.h5'
|
||||
decoder_dstmH5 = 'decoder_dstm.h5'
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_for_session_options):
|
||||
default_resolution = 128
|
||||
default_filters = 64
|
||||
default_match_style = True
|
||||
default_face_type = 'f'
|
||||
|
||||
if is_first_run:
|
||||
#first run
|
||||
self.options['resolution'] = input_int("Resolution (valid: 64,128,256, skip:128) : ", default_resolution, [64,128,256])
|
||||
self.options['filters'] = np.clip ( input_int("Number of U-net filters (valid: 32-128, skip:64) : ", default_filters), 32, 128 )
|
||||
self.options['match_style'] = input_bool ("Match style? (y/n skip:y) : ", default_match_style)
|
||||
self.options['face_type'] = input_str ("Half or Full face? (h/f, skip:f) : ", default_face_type, ['h','f'])
|
||||
|
||||
else:
|
||||
#not first run
|
||||
self.options['resolution'] = self.options.get('resolution', default_resolution)
|
||||
self.options['filters'] = self.options.get('filters', default_filters)
|
||||
self.options['match_style'] = self.options.get('match_style', default_match_style)
|
||||
self.options['face_type'] = self.options.get('face_type', default_face_type)
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
|
||||
self.set_vram_batch_requirements({2:1,3:2,4:6,5:8,6:16,7:24,8:32})
|
||||
|
||||
resolution = self.options['resolution']
|
||||
bgr_shape = (resolution, resolution, 3)
|
||||
mask_shape = (resolution, resolution, 1)
|
||||
|
||||
filters = self.options['filters']
|
||||
|
||||
if resolution == 64:
|
||||
lowest_dense = 512
|
||||
elif resolution == 128:
|
||||
lowest_dense = 512
|
||||
elif resolution == 256:
|
||||
lowest_dense = 256
|
||||
|
||||
self.encoder = modelify(UFMModel.EncFlow (ngf=filters, lowest_dense=lowest_dense)) (Input(bgr_shape))
|
||||
|
||||
dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ]
|
||||
|
||||
self.decoder_src = modelify(UFMModel.DecFlow (bgr_shape[2], ngf=filters)) (dec_Inputs)
|
||||
self.decoder_dst = modelify(UFMModel.DecFlow (bgr_shape[2], ngf=filters)) (dec_Inputs)
|
||||
|
||||
self.decoder_srcm = modelify(UFMModel.DecFlow (mask_shape[2], ngf=filters//2)) (dec_Inputs)
|
||||
self.decoder_dstm = modelify(UFMModel.DecFlow (mask_shape[2], ngf=filters//2)) (dec_Inputs)
|
||||
|
||||
if not self.is_first_run():
|
||||
self.encoder.load_weights (self.get_strpath_storage_for_file(self.encoderH5))
|
||||
self.decoder_src.load_weights (self.get_strpath_storage_for_file(self.decoder_srcH5))
|
||||
self.decoder_dst.load_weights (self.get_strpath_storage_for_file(self.decoder_dstH5))
|
||||
self.decoder_srcm.load_weights (self.get_strpath_storage_for_file(self.decoder_srcmH5))
|
||||
self.decoder_dstm.load_weights (self.get_strpath_storage_for_file(self.decoder_dstmH5))
|
||||
|
||||
warped_src = Input(bgr_shape)
|
||||
target_src = Input(bgr_shape)
|
||||
target_srcm = Input(mask_shape)
|
||||
|
||||
warped_src_code = self.encoder (warped_src)
|
||||
pred_src_src = self.decoder_src(warped_src_code)
|
||||
pred_src_srcm = self.decoder_srcm(warped_src_code)
|
||||
|
||||
warped_dst = Input(bgr_shape)
|
||||
target_dst = Input(bgr_shape)
|
||||
target_dstm = Input(mask_shape)
|
||||
|
||||
warped_dst_code = self.encoder (warped_dst)
|
||||
pred_dst_dst = self.decoder_dst(warped_dst_code)
|
||||
pred_dst_dstm = self.decoder_dstm(warped_dst_code)
|
||||
|
||||
pred_src_dst = self.decoder_src(warped_dst_code)
|
||||
pred_src_dstm = self.decoder_srcm(warped_dst_code)
|
||||
|
||||
target_srcm_blurred = tf_gaussian_blur(4.0)(target_srcm)
|
||||
target_srcm_sigm = target_srcm_blurred / 2.0 + 0.5
|
||||
target_srcm_anti_sigm = 1.0 - target_srcm_sigm
|
||||
|
||||
target_dstm_blurred = tf_gaussian_blur(4.0)(target_dstm)
|
||||
target_dstm_sigm = target_dstm_blurred / 2.0 + 0.5
|
||||
target_dstm_anti_sigm = 1.0 - target_dstm_sigm
|
||||
|
||||
target_src_sigm = target_src+1
|
||||
target_dst_sigm = target_dst+1
|
||||
|
||||
pred_src_src_sigm = pred_src_src+1
|
||||
pred_dst_dst_sigm = pred_dst_dst+1
|
||||
pred_src_dst_sigm = pred_src_dst+1
|
||||
|
||||
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
|
||||
|
||||
pred_src_src_masked = pred_src_src_sigm * target_srcm_sigm
|
||||
pred_dst_dst_masked = pred_dst_dst_sigm * target_dstm_sigm
|
||||
|
||||
pred_src_dst_target_dst_masked = pred_src_dst_sigm * target_dstm_sigm
|
||||
pred_src_dst_target_dst_anti_masked = pred_src_dst_sigm * target_dstm_anti_sigm
|
||||
|
||||
|
||||
src_loss = K.mean( 100*K.square(tf_dssim(2.0)( target_src_masked, pred_src_src_masked )) )
|
||||
if self.options['match_style']:
|
||||
src_loss += tf_style_loss(gaussian_blur_radius=resolution // 8, loss_weight=0.015)(pred_src_dst_target_dst_masked, target_dst_masked)
|
||||
src_loss += 0.05 * K.mean( tf_dssim(2.0)( pred_src_dst_target_dst_anti_masked, target_dst_anti_masked ))
|
||||
|
||||
self.src_train = K.function ([warped_src, target_src, target_srcm, warped_dst, target_dst, target_dstm ],[src_loss],
|
||||
Adam(lr=5e-5, beta_1=0.5, beta_2=0.999).get_updates(src_loss, self.encoder.trainable_weights + self.decoder_src.trainable_weights) )
|
||||
|
||||
dst_loss = K.mean( 100*K.square(tf_dssim(2.0)( target_dst_masked, pred_dst_dst_masked )) )
|
||||
self.dst_train = K.function ([warped_dst, target_dst, target_dstm],[dst_loss],
|
||||
Adam(lr=5e-5, beta_1=0.5, beta_2=0.999).get_updates(dst_loss, self.encoder.trainable_weights + self.decoder_dst.trainable_weights) )
|
||||
|
||||
|
||||
src_mask_loss = K.mean(K.square(target_srcm-pred_src_srcm))
|
||||
self.src_mask_train = K.function ([warped_src, target_srcm],[src_mask_loss],
|
||||
Adam(lr=5e-5, beta_1=0.5, beta_2=0.999).get_updates(src_mask_loss, self.encoder.trainable_weights + self.decoder_srcm.trainable_weights) )
|
||||
|
||||
dst_mask_loss = K.mean(K.square(target_dstm-pred_dst_dstm))
|
||||
self.dst_mask_train = K.function ([warped_dst, target_dstm],[dst_mask_loss],
|
||||
Adam(lr=5e-5, beta_1=0.5, beta_2=0.999).get_updates(dst_mask_loss, self.encoder.trainable_weights + self.decoder_dstm.trainable_weights) )
|
||||
|
||||
self.AE_view = K.function ([warped_src, warped_dst],[pred_src_src, pred_src_srcm, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm])
|
||||
self.AE_convert = K.function ([warped_dst],[pred_src_dst, pred_src_dstm])
|
||||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
|
||||
face_type = f.FACE_ALIGN_FULL if self.options['face_type'] == 'f' else f.FACE_ALIGN_HALF
|
||||
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
|
||||
debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(normalize_tanh = True),
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | face_type | f.MODE_BGR, resolution],
|
||||
[f.TRANSFORMED | face_type | f.MODE_BGR, resolution],
|
||||
[f.TRANSFORMED | face_type | f.MODE_M | f.FACE_MASK_FULL, resolution] ] ),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(normalize_tanh = True),
|
||||
output_sample_types=[ [f.WARPED_TRANSFORMED | face_type | f.MODE_BGR, resolution],
|
||||
[f.TRANSFORMED | face_type | f.MODE_BGR, resolution],
|
||||
[f.TRANSFORMED | face_type | f.MODE_M | f.FACE_MASK_FULL, resolution] ] )
|
||||
])
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( [[self.encoder, self.get_strpath_storage_for_file(self.encoderH5)],
|
||||
[self.decoder_src, self.get_strpath_storage_for_file(self.decoder_srcH5)],
|
||||
[self.decoder_dst, self.get_strpath_storage_for_file(self.decoder_dstH5)],
|
||||
[self.decoder_srcm, self.get_strpath_storage_for_file(self.decoder_srcmH5)],
|
||||
[self.decoder_dstm, self.get_strpath_storage_for_file(self.decoder_dstmH5)]
|
||||
] )
|
||||
|
||||
#override
|
||||
def onTrainOneEpoch(self, sample):
|
||||
warped_src, target_src, target_src_mask = sample[0]
|
||||
warped_dst, target_dst, target_dst_mask = sample[1]
|
||||
|
||||
src_loss, = self.src_train ([warped_src, target_src, target_src_mask, warped_dst, target_dst, target_dst_mask])
|
||||
dst_loss, = self.dst_train ([warped_dst, target_dst, target_dst_mask])
|
||||
|
||||
src_mask_loss, = self.src_mask_train ([warped_src, target_src_mask])
|
||||
dst_mask_loss, = self.dst_mask_train ([warped_dst, target_dst_mask])
|
||||
|
||||
return ( ('src_loss', src_loss), ('dst_loss', dst_loss) )
|
||||
|
||||
|
||||
#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]
|
||||
|
||||
S = test_A
|
||||
D = test_B
|
||||
|
||||
SS, SM, DD, DM, SD, SDM = self.AE_view ([test_A, test_B])
|
||||
S, D, SS, SM, DD, DM, SD, SDM = [ x / 2 + 0.5 for x in [S, D, SS, SM, DD, DM, SD, SDM] ]
|
||||
|
||||
SM, DM, SDM = [ np.repeat (x, (3,), -1) for x in [SM, DM, SDM] ]
|
||||
|
||||
st = []
|
||||
for i in range(0, len(test_A)):
|
||||
st.append ( np.concatenate ( (
|
||||
S[i], SS[i], #SM[i],
|
||||
D[i], DD[i], #DM[i],
|
||||
SD[i], #SDM[i]
|
||||
), axis=1) )
|
||||
|
||||
return [ ('U-net Face Morpher', np.concatenate ( st, axis=0 ) ) ]
|
||||
|
||||
def predictor_func (self, face):
|
||||
|
||||
face = face * 2.0 - 1.0
|
||||
|
||||
face_128_bgr = face[...,0:3]
|
||||
|
||||
x, mx = [ (x[0] + 1.0) / 2.0 for x in self.AE_convert ( [ np.expand_dims(face_128_bgr,0) ] ) ]
|
||||
|
||||
if self.options['match_style']:
|
||||
res = self.options['resolution']
|
||||
s = int( res * 0.96875 )
|
||||
mx = np.pad ( np.ones ( (s,s) ), (res-s) // 2 , mode='constant')
|
||||
mx = np.expand_dims(mx, -1)
|
||||
|
||||
return np.concatenate ( (x,mx), -1 )
|
||||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterMasked
|
||||
|
||||
if self.options['match_style']:
|
||||
base_erode_mask_modifier = 50
|
||||
base_blur_mask_modifier = 50
|
||||
else:
|
||||
base_erode_mask_modifier = 30 if self.options['face_type'] == 'f' else 100
|
||||
base_blur_mask_modifier = 0 if self.options['face_type'] == 'f' else 100
|
||||
|
||||
face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
|
||||
|
||||
return ConverterMasked(self.predictor_func,
|
||||
predictor_input_size=self.options['resolution'],
|
||||
output_size=self.options['resolution'],
|
||||
face_type=face_type,
|
||||
base_erode_mask_modifier=base_erode_mask_modifier,
|
||||
base_blur_mask_modifier=base_blur_mask_modifier,
|
||||
**in_options)
|
||||
|
||||
@staticmethod
|
||||
def EncFlow(ngf=64, num_downs=4, lowest_dense=512):
|
||||
exec (nnlib.import_all(), locals(), globals())
|
||||
|
||||
use_bias = True
|
||||
def XNormalization(x):
|
||||
return InstanceNormalization (axis=3, gamma_initializer=RandomNormal(1., 0.02))(x)
|
||||
|
||||
def Conv2D (filters, kernel_size, strides=(1, 1), padding='valid', data_format=None, dilation_rate=(1, 1), activation=None, use_bias=use_bias, kernel_initializer=RandomNormal(0, 0.02), bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None):
|
||||
return keras.layers.Conv2D( filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint )
|
||||
|
||||
def func(input):
|
||||
x = input
|
||||
|
||||
result = []
|
||||
for i in range(num_downs):
|
||||
x = LeakyReLU(0.1)(XNormalization(Conv2D( min(ngf* (2**i), ngf*8) , 5, 2, 'same')(x)))
|
||||
|
||||
if i == num_downs-1:
|
||||
x_shape = K.int_shape(x)[1:]
|
||||
x = Reshape(x_shape)(Dense( np.prod(x_shape) )(Dense(lowest_dense)(Flatten()(x))))
|
||||
result += [x]
|
||||
|
||||
return result
|
||||
return func
|
||||
|
||||
@staticmethod
|
||||
def DecFlow(output_nc, ngf=64, activation='tanh'):
|
||||
exec (nnlib.import_all(), locals(), globals())
|
||||
|
||||
use_bias = True
|
||||
def XNormalization(x):
|
||||
return InstanceNormalization (axis=3, gamma_initializer=RandomNormal(1., 0.02))(x)
|
||||
|
||||
def Conv2D (filters, kernel_size, strides=(1, 1), padding='valid', data_format=None, dilation_rate=(1, 1), activation=None, use_bias=use_bias, kernel_initializer=RandomNormal(0, 0.02), bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None):
|
||||
return keras.layers.Conv2D( filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint )
|
||||
|
||||
def func(input):
|
||||
input_len = len(input)
|
||||
|
||||
x = input[input_len-1]
|
||||
for i in range(input_len-1, -1, -1):
|
||||
x = SubpixelUpscaler()( LeakyReLU(0.1)(XNormalization(Conv2D( min(ngf* (2**i) *4, ngf*8 *4 ), 3, 1, 'same')(x))) )
|
||||
if i != 0:
|
||||
x = Concatenate(axis=3)([ input[i-1] , x])
|
||||
|
||||
return Conv2D(output_nc, 3, 1, 'same', activation=activation)(x)
|
||||
return func
|
||||
|
||||
Model = UFMModel
|
||||
Loading…
Add table
Add a link
Reference in a new issue