mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-20 05:23:22 -07:00
initial
This commit is contained in:
iperov
parent
73de93b4f1
commit
6bd5a44264
71 changed files with 8448 additions and 0 deletions
50
models/BaseTypes.py
Normal file
50
models/BaseTypes.py
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
from enum import IntEnum
|
||||
import cv2
|
||||
import numpy as np
|
||||
from random import randint
|
||||
from facelib import FaceType
|
||||
|
||||
|
||||
class TrainingDataType(IntEnum):
|
||||
IMAGE = 0 #raw image
|
||||
|
||||
FACE_BEGIN = 1
|
||||
FACE = 1 #aligned face unsorted
|
||||
FACE_YAW_SORTED = 2 #sorted by yaw
|
||||
FACE_YAW_SORTED_AS_TARGET = 3 #sorted by yaw and included only yaws which exist in TARGET also automatic mirrored
|
||||
FACE_END = 3
|
||||
|
||||
QTY = 4
|
||||
|
||||
|
||||
class TrainingDataSample(object):
|
||||
|
||||
def __init__(self, filename=None, face_type=None, shape=None, landmarks=None, yaw=None, mirror=None, nearest_target_list=None):
|
||||
self.filename = filename
|
||||
self.face_type = face_type
|
||||
self.shape = shape
|
||||
self.landmarks = np.array(landmarks) if landmarks is not None else None
|
||||
self.yaw = yaw
|
||||
self.mirror = mirror
|
||||
self.nearest_target_list = nearest_target_list
|
||||
|
||||
def copy_and_set(self, filename=None, face_type=None, shape=None, landmarks=None, yaw=None, mirror=None, nearest_target_list=None):
|
||||
return TrainingDataSample(
|
||||
filename=filename if filename is not None else self.filename,
|
||||
face_type=face_type if face_type is not None else self.face_type,
|
||||
shape=shape if shape is not None else self.shape,
|
||||
landmarks=landmarks if landmarks is not None else self.landmarks.copy(),
|
||||
yaw=yaw if yaw is not None else self.yaw,
|
||||
mirror=mirror if mirror is not None else self.mirror,
|
||||
nearest_target_list=nearest_target_list if nearest_target_list is not None else self.nearest_target_list)
|
||||
|
||||
def load_bgr(self):
|
||||
img = cv2.imread (self.filename).astype(np.float32) / 255.0
|
||||
if self.mirror:
|
||||
img = img[:,::-1].copy()
|
||||
return img
|
||||
|
||||
def get_random_nearest_target_sample(self):
|
||||
if self.nearest_target_list is None:
|
||||
return None
|
||||
return self.nearest_target_list[randint (0, len(self.nearest_target_list)-1)]
|
||||
44
models/ConverterBase.py
Normal file
44
models/ConverterBase.py
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
import copy
|
||||
'''
|
||||
You can implement your own Converter, check example ConverterMasked.py
|
||||
'''
|
||||
|
||||
class ConverterBase(object):
|
||||
MODE_FACE = 0
|
||||
MODE_IMAGE = 1
|
||||
|
||||
#overridable
|
||||
def __init__(self, predictor):
|
||||
self.predictor = predictor
|
||||
|
||||
#overridable
|
||||
def get_mode(self):
|
||||
#MODE_FACE calls convert_face
|
||||
#MODE_IMAGE calls convert_image
|
||||
return ConverterBase.MODE_FACE
|
||||
|
||||
#overridable
|
||||
def convert_face (self, img_bgr, img_face_landmarks, debug):
|
||||
#return float32 image
|
||||
#if debug , return tuple ( images of any size and channels, ...)
|
||||
return image
|
||||
|
||||
#overridable
|
||||
def convert_image (self, img_bgr, img_landmarks, debug):
|
||||
#img_landmarks not None, if input image is png with embedded data
|
||||
#return float32 image
|
||||
#if debug , return tuple ( images of any size and channels, ...)
|
||||
return image
|
||||
|
||||
#overridable
|
||||
def dummy_predict(self):
|
||||
#do dummy predict here
|
||||
pass
|
||||
|
||||
def copy(self):
|
||||
return copy.copy(self)
|
||||
|
||||
def copy_and_set_predictor(self, predictor):
|
||||
result = self.copy()
|
||||
result.predictor = predictor
|
||||
return result
|
||||
46
models/ConverterImage.py
Normal file
46
models/ConverterImage.py
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
from models import ConverterBase
|
||||
from facelib import LandmarksProcessor
|
||||
from facelib import FaceType
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from utils import image_utils
|
||||
|
||||
'''
|
||||
predictor:
|
||||
input: [predictor_input_size, predictor_input_size, BGR]
|
||||
output: [predictor_input_size, predictor_input_size, BGR]
|
||||
'''
|
||||
|
||||
class ConverterImage(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
|
||||
|
||||
#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_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 )
|
||||
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
|
||||
194
models/ConverterMasked.py
Normal file
194
models/ConverterMasked.py
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
from models import ConverterBase
|
||||
from facelib import LandmarksProcessor
|
||||
from facelib import FaceType
|
||||
import cv2
|
||||
import numpy as np
|
||||
from utils import image_utils
|
||||
|
||||
'''
|
||||
predictor:
|
||||
input: [predictor_input_size, predictor_input_size, BGRA]
|
||||
output: [predictor_input_size, predictor_input_size, BGRA]
|
||||
'''
|
||||
|
||||
class ConverterMasked(ConverterBase):
|
||||
|
||||
#override
|
||||
def __init__(self, predictor,
|
||||
predictor_input_size=0,
|
||||
output_size=0,
|
||||
face_type=FaceType.FULL,
|
||||
erode_mask = True,
|
||||
blur_mask = True,
|
||||
clip_border_mask_per = 0,
|
||||
masked_hist_match = False,
|
||||
mode='seamless',
|
||||
erode_mask_modifier=0,
|
||||
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.erode_mask = erode_mask
|
||||
self.blur_mask = blur_mask
|
||||
self.clip_border_mask_per = clip_border_mask_per
|
||||
self.masked_hist_match = masked_hist_match
|
||||
self.mode = mode
|
||||
self.erode_mask_modifier = erode_mask_modifier
|
||||
self.blur_mask_modifier = blur_mask_modifier
|
||||
|
||||
if self.erode_mask_modifier != 0 and not self.erode_mask:
|
||||
print ("Erode mask modifier not used in this model.")
|
||||
|
||||
if self.blur_mask_modifier != 0 and not self.blur_mask:
|
||||
print ("Blur modifier not used in this model.")
|
||||
|
||||
#override
|
||||
def get_mode(self):
|
||||
return ConverterBase.MODE_FACE
|
||||
|
||||
#override
|
||||
def dummy_predict(self):
|
||||
self.predictor ( np.zeros ( (self.predictor_input_size,self.predictor_input_size,4), dtype=np.float32 ) )
|
||||
|
||||
#override
|
||||
def convert_face (self, img_bgr, img_face_landmarks, debug):
|
||||
if debug:
|
||||
debugs = [img_bgr.copy()]
|
||||
|
||||
img_size = img_bgr.shape[1], img_bgr.shape[0]
|
||||
|
||||
img_face_mask_a = LandmarksProcessor.get_image_hull_mask (img_bgr, img_face_landmarks)
|
||||
|
||||
face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, self.output_size, face_type=self.face_type)
|
||||
dst_face_bgr = cv2.warpAffine( img_bgr , face_mat, (self.output_size, self.output_size), flags=cv2.INTER_LANCZOS4 )
|
||||
dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (self.output_size, self.output_size), flags=cv2.INTER_LANCZOS4 )
|
||||
|
||||
predictor_input_bgr = cv2.resize (dst_face_bgr, (self.predictor_input_size,self.predictor_input_size))
|
||||
predictor_input_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size))
|
||||
predictor_input_mask_a = np.expand_dims (predictor_input_mask_a_0, -1)
|
||||
|
||||
predicted_bgra = self.predictor ( np.concatenate( (predictor_input_bgr, predictor_input_mask_a), -1) )
|
||||
|
||||
prd_face_bgr = np.clip (predicted_bgra[:,:,0:3], 0, 1.0 )
|
||||
prd_face_mask_a_0 = np.clip (predicted_bgra[:,:,3], 0.0, 1.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)
|
||||
prd_face_mask_aaa = np.repeat (prd_face_mask_a, (3,), axis=-1)
|
||||
|
||||
img_prd_face_mask_aaa = cv2.warpAffine( prd_face_mask_aaa, face_mat, img_size, np.zeros(img_bgr.shape, dtype=float), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4 )
|
||||
img_prd_face_mask_aaa = np.clip (img_prd_face_mask_aaa, 0.0, 1.0)
|
||||
|
||||
img_face_mask_aaa = img_prd_face_mask_aaa
|
||||
|
||||
if debug:
|
||||
debugs += [img_face_mask_aaa.copy()]
|
||||
|
||||
img_face_mask_aaa [ img_face_mask_aaa <= 0.1 ] = 0.0
|
||||
|
||||
img_face_mask_flatten_aaa = img_face_mask_aaa.copy()
|
||||
img_face_mask_flatten_aaa[img_face_mask_flatten_aaa > 0.9] = 1.0
|
||||
|
||||
maxregion = np.argwhere(img_face_mask_flatten_aaa==1.0)
|
||||
|
||||
out_img = img_bgr.copy()
|
||||
if maxregion.size != 0:
|
||||
miny,minx = maxregion.min(axis=0)[:2]
|
||||
maxy,maxx = maxregion.max(axis=0)[:2]
|
||||
lenx = maxx - minx
|
||||
leny = maxy - miny
|
||||
masky = int(minx+(lenx//2))
|
||||
maskx = int(miny+(leny//2))
|
||||
lowest_len = min (lenx, leny)
|
||||
|
||||
if debug:
|
||||
print ("lowest_len = %f" % (lowest_len) )
|
||||
|
||||
ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*self.erode_mask_modifier )
|
||||
blur = int( lowest_len * 0.10 * 0.01*self.blur_mask_modifier )
|
||||
|
||||
if debug:
|
||||
print ("ero = %d, blur = %d" % (ero, blur) )
|
||||
|
||||
img_mask_blurry_aaa = img_face_mask_aaa
|
||||
if self.erode_mask:
|
||||
if ero > 0:
|
||||
img_mask_blurry_aaa = cv2.erode(img_mask_blurry_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 )
|
||||
elif ero < 0:
|
||||
img_mask_blurry_aaa = cv2.dilate(img_mask_blurry_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(-ero,-ero)), iterations = 1 )
|
||||
|
||||
if self.blur_mask and blur > 0:
|
||||
img_mask_blurry_aaa = cv2.blur(img_mask_blurry_aaa, (blur, 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)
|
||||
prd_face_bgr = np.repeat( np.expand_dims (prd_face_bgr, -1), (3,), -1 )
|
||||
|
||||
if self.mode == 'hist-match' or self.mode == 'hist-match-bw':
|
||||
if debug:
|
||||
debugs += [ cv2.warpAffine( prd_face_bgr, face_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) ]
|
||||
|
||||
hist_mask_a = np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=prd_face_bgr.dtype)
|
||||
|
||||
if self.masked_hist_match:
|
||||
hist_mask_a *= prd_face_mask_a
|
||||
|
||||
new_prd_face_bgr = image_utils.color_hist_match(prd_face_bgr*hist_mask_a, dst_face_bgr*hist_mask_a )
|
||||
|
||||
prd_face_bgr = new_prd_face_bgr
|
||||
|
||||
if self.mode == 'hist-match-bw':
|
||||
prd_face_bgr = prd_face_bgr.astype(np.float32)
|
||||
|
||||
out_img = cv2.warpAffine( prd_face_bgr, face_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 == '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:
|
||||
debugs += [out_img.copy()]
|
||||
out_img = cv2.seamlessClone( (out_img*255).astype(np.uint8), (img_bgr*255).astype(np.uint8), (img_face_mask_flatten_aaa*255).astype(np.uint8), (masky,maskx) , cv2.NORMAL_CLONE )
|
||||
out_img = out_img.astype(np.float32) / 255.0
|
||||
|
||||
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_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':
|
||||
out_face_bgr = cv2.warpAffine( out_img, face_mat, (self.output_size, self.output_size) )
|
||||
new_out_face_bgr = image_utils.color_hist_match(out_face_bgr, dst_face_bgr )
|
||||
new_out = cv2.warpAffine( new_out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT )
|
||||
out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (new_out*img_mask_blurry_aaa) , 0, 1.0 )
|
||||
|
||||
if debug:
|
||||
debugs += [out_img.copy()]
|
||||
|
||||
return debugs if debug else out_img
|
||||
|
||||
332
models/ModelBase.py
Normal file
332
models/ModelBase.py
Normal file
|
|
@ -0,0 +1,332 @@
|
|||
import os
|
||||
import time
|
||||
import inspect
|
||||
import operator
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
from utils import Path_utils
|
||||
from utils import std_utils
|
||||
from utils import image_utils
|
||||
import numpy as np
|
||||
import cv2
|
||||
import gpufmkmgr
|
||||
from .TrainingDataGeneratorBase import TrainingDataGeneratorBase
|
||||
|
||||
'''
|
||||
You can implement your own model. Check examples.
|
||||
'''
|
||||
class ModelBase(object):
|
||||
|
||||
#DONT OVERRIDE
|
||||
def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None,
|
||||
batch_size=0,
|
||||
multi_gpu = False,
|
||||
choose_worst_gpu = False,
|
||||
force_best_gpu_idx = -1,
|
||||
force_gpu_idxs = None,
|
||||
write_preview_history = False,
|
||||
debug = False, **in_options
|
||||
):
|
||||
print ("Loading model...")
|
||||
self.model_path = model_path
|
||||
self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') )
|
||||
|
||||
self.training_data_src_path = training_data_src_path
|
||||
self.training_data_dst_path = training_data_dst_path
|
||||
|
||||
self.src_images_paths = None
|
||||
self.dst_images_paths = None
|
||||
self.src_yaw_images_paths = None
|
||||
self.dst_yaw_images_paths = None
|
||||
self.src_data_generator = None
|
||||
self.dst_data_generator = None
|
||||
self.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.supress_std_once = ('TF_SUPPRESS_STD' in os.environ.keys() and os.environ['TF_SUPPRESS_STD'] == '1')
|
||||
|
||||
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.generator_dict_states = model_data['generator_dict_states'] if 'generator_dict_states' in model_data.keys() else None
|
||||
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.generator_dict_states = None
|
||||
self.sample_for_preview = 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.multi_gpu = multi_gpu
|
||||
|
||||
gpu_idx = force_best_gpu_idx if (force_best_gpu_idx >= 0 and gpufmkmgr.isValidDeviceIdx(force_best_gpu_idx)) else gpufmkmgr.getBestDeviceIdx() if not choose_worst_gpu else gpufmkmgr.getWorstDeviceIdx()
|
||||
gpu_total_vram_gb = gpufmkmgr.getDeviceVRAMTotalGb (gpu_idx)
|
||||
is_gpu_low_mem = (gpu_total_vram_gb < 4)
|
||||
|
||||
self.gpu_total_vram_gb = gpu_total_vram_gb
|
||||
|
||||
if self.epoch == 0:
|
||||
#first run
|
||||
self.options['created_vram_gb'] = gpu_total_vram_gb
|
||||
self.created_vram_gb = gpu_total_vram_gb
|
||||
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'] = gpu_total_vram_gb
|
||||
self.created_vram_gb = gpu_total_vram_gb
|
||||
|
||||
if force_gpu_idxs is not None:
|
||||
self.gpu_idxs = [ int(x) for x in force_gpu_idxs.split(',') ]
|
||||
else:
|
||||
if self.multi_gpu:
|
||||
self.gpu_idxs = gpufmkmgr.getDeviceIdxsEqualModel( gpu_idx )
|
||||
if len(self.gpu_idxs) <= 1:
|
||||
self.multi_gpu = False
|
||||
else:
|
||||
self.gpu_idxs = [gpu_idx]
|
||||
|
||||
self.tf = gpufmkmgr.import_tf(self.gpu_idxs,allow_growth=False)
|
||||
self.tf_sess = gpufmkmgr.get_tf_session()
|
||||
self.keras = gpufmkmgr.import_keras()
|
||||
self.keras_contrib = gpufmkmgr.import_keras_contrib()
|
||||
|
||||
self.onInitialize(**in_options)
|
||||
|
||||
if self.debug or self.batch_size == 0:
|
||||
self.batch_size = 1
|
||||
|
||||
if self.is_training_mode:
|
||||
if self.generator_list is None:
|
||||
raise Exception( 'You didnt set_training_data_generators()')
|
||||
else:
|
||||
for i, generator in enumerate(self.generator_list):
|
||||
if not isinstance(generator, TrainingDataGeneratorBase):
|
||||
raise Exception('training data generator is not subclass of TrainingDataGeneratorBase')
|
||||
|
||||
if self.generator_dict_states is not None and i < len(self.generator_dict_states):
|
||||
generator.set_dict_state ( self.generator_dict_states[i] )
|
||||
|
||||
if self.sample_for_preview is None:
|
||||
self.sample_for_preview = self.generate_next_sample()
|
||||
|
||||
print ("===== Model summary =====")
|
||||
print ("== Model name: " + self.get_model_name())
|
||||
print ("==")
|
||||
print ("== Current epoch: " + str(self.epoch) )
|
||||
print ("==")
|
||||
print ("== Options:")
|
||||
print ("== |== batch_size : %s " % (self.batch_size) )
|
||||
print ("== |== multi_gpu : %s " % (self.multi_gpu) )
|
||||
for key in self.options.keys():
|
||||
print ("== |== %s : %s" % (key, self.options[key]) )
|
||||
|
||||
print ("== Running on:")
|
||||
for idx in self.gpu_idxs:
|
||||
print ("== |== [%d : %s]" % (idx, gpufmkmgr.getDeviceName(idx)) )
|
||||
|
||||
if self.gpu_total_vram_gb == 2:
|
||||
print ("==")
|
||||
print ("== WARNING: You are using 2GB GPU. Result quality may be significantly decreased.")
|
||||
print ("== If training does not start, close all programs and try again.")
|
||||
print ("== Also you can disable Windows Aero Desktop to get extra free VRAM.")
|
||||
print ("==")
|
||||
|
||||
print ("=========================")
|
||||
|
||||
#overridable
|
||||
def onInitialize(self, **in_options):
|
||||
'''
|
||||
initialize your keras models
|
||||
|
||||
store and retrieve your model options in self.options['']
|
||||
|
||||
check example
|
||||
'''
|
||||
pass
|
||||
|
||||
#overridable
|
||||
def onSave(self):
|
||||
#save your keras models here
|
||||
pass
|
||||
|
||||
#overridable
|
||||
def onTrainOneEpoch(self, sample):
|
||||
#train your keras models here
|
||||
|
||||
#return array of losses
|
||||
return ( ('loss_src', 0), ('loss_dst', 0) )
|
||||
|
||||
#overridable
|
||||
def onGetPreview(self, sample):
|
||||
#you can return multiple previews
|
||||
#return [ ('preview_name',preview_rgb), ... ]
|
||||
return []
|
||||
|
||||
#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]
|
||||
|
||||
#overridable
|
||||
def get_converter(self, **in_options):
|
||||
#return existing or your own converter which derived from base
|
||||
from .ConverterBase import ConverterBase
|
||||
return ConverterBase(self, **in_options)
|
||||
|
||||
def to_multi_gpu_model_if_possible (self, models_list):
|
||||
if len(self.gpu_idxs) > 1:
|
||||
#make batch_size to divide on GPU count without remainder
|
||||
self.batch_size = int( self.batch_size / len(self.gpu_idxs) )
|
||||
if self.batch_size == 0:
|
||||
self.batch_size = 1
|
||||
self.batch_size *= len(self.gpu_idxs)
|
||||
|
||||
result = []
|
||||
for model in models_list:
|
||||
for i in range( len(model.output_names) ):
|
||||
model.output_names = 'output_%d' % (i)
|
||||
result += [ self.keras.utils.multi_gpu_model( model, self.gpu_idxs ) ]
|
||||
|
||||
return result
|
||||
else:
|
||||
return models_list
|
||||
|
||||
def get_previews(self):
|
||||
return self.onGetPreview ( self.last_sample )
|
||||
|
||||
def get_static_preview(self):
|
||||
return self.onGetPreview (self.sample_for_preview)[0][1] #first preview, and bgr
|
||||
|
||||
def save(self):
|
||||
print ("Saving...")
|
||||
|
||||
if self.supress_std_once:
|
||||
supressor = std_utils.suppress_stdout_stderr()
|
||||
supressor.__enter__()
|
||||
|
||||
self.onSave()
|
||||
|
||||
if self.supress_std_once:
|
||||
supressor.__exit__()
|
||||
|
||||
model_data = {
|
||||
'epoch': self.epoch,
|
||||
'options': self.options,
|
||||
'loss_history': self.loss_history,
|
||||
'generator_dict_states' : [generator.get_dict_state() for generator in self.generator_list],
|
||||
'sample_for_preview' : self.sample_for_preview
|
||||
}
|
||||
self.model_data_path.write_bytes( pickle.dumps(model_data) )
|
||||
|
||||
def save_weights_safe(self, model_filename_list):
|
||||
for model, filename in model_filename_list:
|
||||
model.save_weights( filename + '.tmp' )
|
||||
|
||||
for model, filename in model_filename_list:
|
||||
source_filename = Path(filename+'.tmp')
|
||||
target_filename = Path(filename)
|
||||
if target_filename.exists():
|
||||
target_filename.unlink()
|
||||
|
||||
source_filename.rename ( str(target_filename) )
|
||||
|
||||
def debug_one_epoch(self):
|
||||
images = []
|
||||
for generator in self.generator_list:
|
||||
for i,batch in enumerate(next(generator)):
|
||||
images.append( batch[0] )
|
||||
|
||||
return image_utils.equalize_and_stack_square (images)
|
||||
|
||||
def generate_next_sample(self):
|
||||
return [next(generator) for generator in self.generator_list]
|
||||
|
||||
def train_one_epoch(self):
|
||||
if self.supress_std_once:
|
||||
supressor = std_utils.suppress_stdout_stderr()
|
||||
supressor.__enter__()
|
||||
|
||||
self.last_sample = self.generate_next_sample()
|
||||
|
||||
epoch_time = time.time()
|
||||
|
||||
losses = self.onTrainOneEpoch(self.last_sample)
|
||||
|
||||
epoch_time = time.time() - epoch_time
|
||||
|
||||
self.loss_history.append ( [float(loss[1]) for loss in losses] )
|
||||
|
||||
if self.supress_std_once:
|
||||
supressor.__exit__()
|
||||
self.supress_std_once = False
|
||||
|
||||
if self.write_preview_history:
|
||||
if self.epoch % 10 == 0:
|
||||
img = (self.get_static_preview() * 255).astype(np.uint8)
|
||||
cv2.imwrite ( str (self.preview_history_path / ('%.6d.jpg' %( self.epoch) )), img )
|
||||
|
||||
self.epoch += 1
|
||||
|
||||
#............."Saving...
|
||||
loss_string = "Training [#{0:06d}][{1:04d}ms]".format ( self.epoch, int(epoch_time*1000) % 10000 )
|
||||
for (loss_name, loss_value) in losses:
|
||||
loss_string += " %s:%.3f" % (loss_name, loss_value)
|
||||
|
||||
return loss_string
|
||||
|
||||
def pass_one_epoch(self):
|
||||
self.last_sample = self.generate_next_sample()
|
||||
|
||||
def finalize(self):
|
||||
gpufmkmgr.finalize_keras()
|
||||
|
||||
def is_first_run(self):
|
||||
return self.epoch == 0
|
||||
|
||||
def is_debug(self):
|
||||
return self.debug
|
||||
|
||||
def get_epoch(self):
|
||||
return self.epoch
|
||||
|
||||
def get_loss_history(self):
|
||||
return self.loss_history
|
||||
|
||||
def set_training_data_generators (self, generator_list):
|
||||
self.generator_list = generator_list
|
||||
|
||||
def get_training_data_generators (self):
|
||||
return self.generator_list
|
||||
|
||||
def get_strpath_storage_for_file(self, filename):
|
||||
return str( self.model_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()]
|
||||
|
||||
if self.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.gpu_total_vram_gb <= x:
|
||||
self.batch_size = d[x]
|
||||
break
|
||||
|
||||
if self.batch_size == 0:
|
||||
self.batch_size = d[ keys[-1] ]
|
||||
223
models/Model_AVATAR/Model.py
Normal file
223
models/Model_AVATAR/Model.py
Normal file
|
|
@ -0,0 +1,223 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
from nnlib import tf_dssim
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
encoder64H5 = 'encoder64.h5'
|
||||
decoder64_srcH5 = 'decoder64_src.h5'
|
||||
decoder64_dstH5 = 'decoder64_dst.h5'
|
||||
encoder128H5 = 'encoder128.h5'
|
||||
decoder128_srcH5 = 'decoder128_src.h5'
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
tf = self.tf
|
||||
keras = self.keras
|
||||
K = keras.backend
|
||||
|
||||
self.set_vram_batch_requirements( {4:8,5:16,6:20,7:24,8:32,9:48} )
|
||||
|
||||
self.encoder64, self.decoder64_src, self.decoder64_dst, self.encoder128, self.decoder128_src = self.BuildAE()
|
||||
img_shape64 = (64,64,1)
|
||||
img_shape128 = (256,256,3)
|
||||
|
||||
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.encoder128.load_weights (self.get_strpath_storage_for_file(self.encoder128H5))
|
||||
self.decoder128_src.load_weights (self.get_strpath_storage_for_file(self.decoder128_srcH5))
|
||||
|
||||
if self.is_training_mode:
|
||||
self.encoder64, self.decoder64_src, self.decoder64_dst, self.encoder128, self.decoder128_src = self.to_multi_gpu_model_if_possible ( [self.encoder64, self.decoder64_src, self.decoder64_dst, self.encoder128, self.decoder128_src] )
|
||||
|
||||
input_src_64 = keras.layers.Input(img_shape64)
|
||||
input_src_target64 = keras.layers.Input(img_shape64)
|
||||
input_src_target128 = keras.layers.Input(img_shape128)
|
||||
input_dst_64 = keras.layers.Input(img_shape64)
|
||||
input_dst_target64 = keras.layers.Input(img_shape64)
|
||||
|
||||
src_code64 = self.encoder64(input_src_64)
|
||||
dst_code64 = self.encoder64(input_dst_64)
|
||||
|
||||
rec_src64 = self.decoder64_src(src_code64)
|
||||
rec_dst64 = self.decoder64_dst(dst_code64)
|
||||
|
||||
src64_loss = tf_dssim(tf, input_src_target64, rec_src64)
|
||||
dst64_loss = tf_dssim(tf, input_dst_target64, rec_dst64)
|
||||
total64_loss = src64_loss + dst64_loss
|
||||
|
||||
self.ed64_train = K.function ([input_src_64, input_src_target64, input_dst_64, input_dst_target64],[K.mean(total64_loss)],
|
||||
self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999).get_updates(total64_loss, self.encoder64.trainable_weights + self.decoder64_src.trainable_weights + self.decoder64_dst.trainable_weights)
|
||||
)
|
||||
|
||||
src_code128 = self.encoder128(input_src_64)
|
||||
rec_src128 = self.decoder128_src(src_code128)
|
||||
src128_loss = tf_dssim(tf, input_src_target128, rec_src128)
|
||||
|
||||
self.ed128_train = K.function ([input_src_64, input_src_target128],[K.mean(src128_loss)],
|
||||
self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999).get_updates(src128_loss, self.encoder128.trainable_weights + self.decoder128_src.trainable_weights)
|
||||
)
|
||||
|
||||
src_code128 = self.encoder128(rec_src64)
|
||||
rec_src128 = self.decoder128_src(src_code128)
|
||||
|
||||
self.src128_view = K.function ([input_src_64], [rec_src128])
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[
|
||||
[f.WARPED_TRANSFORMED | f.HALF_FACE | f.MODE_G, 64],
|
||||
[f.TRANSFORMED | f.HALF_FACE | f.MODE_G, 64],
|
||||
[f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 256],
|
||||
[f.SOURCE | f.HALF_FACE | f.MODE_G, 64],
|
||||
[f.SOURCE | f.HALF_FACE | f.MODE_GGG, 256] ] ),
|
||||
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[
|
||||
[f.WARPED_TRANSFORMED | f.HALF_FACE | f.MODE_G, 64],
|
||||
[f.TRANSFORMED | f.HALF_FACE | f.MODE_G, 64],
|
||||
[f.SOURCE | f.HALF_FACE | f.MODE_G, 64],
|
||||
[f.SOURCE | f.HALF_FACE | f.MODE_GGG, 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.encoder128, self.get_strpath_storage_for_file(self.encoder128H5)],
|
||||
[self.decoder128_src, self.get_strpath_storage_for_file(self.decoder128_srcH5)],
|
||||
] )
|
||||
|
||||
#override
|
||||
def onTrainOneEpoch(self, sample):
|
||||
warped_src64, target_src64, target_src128, target_src_source64_G, target_src_source128_GGG = sample[0]
|
||||
warped_dst64, target_dst64, target_dst_source64_G, target_dst_source128_GGG = sample[1]
|
||||
|
||||
loss64, = self.ed64_train ([warped_src64, target_src64, warped_dst64, target_dst64])
|
||||
loss256, = self.ed128_train ([warped_src64, target_src128])
|
||||
|
||||
return ( ('loss64', loss64), ('loss256', loss256) )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
n_samples = 4
|
||||
test_B = sample[1][2][0:n_samples]
|
||||
test_B128 = sample[1][3][0:n_samples]
|
||||
|
||||
BB, = self.src128_view ([test_B])
|
||||
|
||||
st = []
|
||||
for i in range(n_samples // 2):
|
||||
st.append ( np.concatenate ( (
|
||||
test_B128[i*2+0], BB[i*2+0], test_B128[i*2+1], BB[i*2+1],
|
||||
), axis=1) )
|
||||
return [ ('AVATAR', np.concatenate ( st, axis=0 ) ) ]
|
||||
|
||||
def predictor_func (self, img):
|
||||
x, = self.src128_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 BuildAE(self):
|
||||
keras, K = self.keras, self.keras.backend
|
||||
|
||||
def Encoder(_input):
|
||||
x = keras.layers.convolutional.Conv2D(90, kernel_size=5, strides=1, padding='same')(_input)
|
||||
x = keras.layers.convolutional.Conv2D(90, kernel_size=5, strides=1, padding='same')(x)
|
||||
x = keras.layers.MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
|
||||
|
||||
x = keras.layers.convolutional.Conv2D(180, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = keras.layers.convolutional.Conv2D(180, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = keras.layers.MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
|
||||
|
||||
x = keras.layers.convolutional.Conv2D(360, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = keras.layers.convolutional.Conv2D(360, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = keras.layers.MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
|
||||
|
||||
x = keras.layers.Dense (1024)(x)
|
||||
x = keras.layers.advanced_activations.LeakyReLU(0.1)(x)
|
||||
x = keras.layers.Dropout(0.5)(x)
|
||||
|
||||
x = keras.layers.Dense (1024)(x)
|
||||
x = keras.layers.advanced_activations.LeakyReLU(0.1)(x)
|
||||
x = keras.layers.Dropout(0.5)(x)
|
||||
x = keras.layers.Flatten()(x)
|
||||
x = keras.layers.Dense (64)(x)
|
||||
return keras.models.Model (_input, x)
|
||||
|
||||
encoder128 = Encoder( keras.layers.Input ( (64, 64, 1) ) )
|
||||
encoder64 = Encoder( keras.layers.Input ( (64, 64, 1) ) )
|
||||
|
||||
def decoder128_3(encoder):
|
||||
decoder_input = keras.layers.Input ( K.int_shape(encoder.outputs[0])[1:] )
|
||||
x = decoder_input
|
||||
x = self.keras.layers.Dense(16 * 16 * 720)(x)
|
||||
x = keras.layers.Reshape ( (16, 16, 720) )(x)
|
||||
x = upscale(keras, x, 720)
|
||||
x = upscale(keras, x, 360)
|
||||
x = upscale(keras, x, 180)
|
||||
x = upscale(keras, x, 90)
|
||||
x = keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return keras.models.Model(decoder_input, x)
|
||||
|
||||
def decoder64_1(encoder):
|
||||
decoder_input = keras.layers.Input ( K.int_shape(encoder.outputs[0])[1:] )
|
||||
x = decoder_input
|
||||
x = self.keras.layers.Dense(8 * 8 * 720)(x)
|
||||
x = keras.layers.Reshape ( (8,8,720) )(x)
|
||||
x = upscale(keras, x, 360)
|
||||
x = upscale(keras, x, 180)
|
||||
x = upscale(keras, x, 90)
|
||||
x = keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return keras.models.Model(decoder_input, x)
|
||||
|
||||
return encoder64, decoder64_1(encoder64), decoder64_1(encoder64), encoder128, decoder128_3(encoder128)
|
||||
|
||||
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
|
||||
|
||||
#override
|
||||
def dummy_predict(self):
|
||||
self.predictor ( np.zeros ( (self.predictor_input_size, self.predictor_input_size,1), 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 )
|
||||
predictor_input_g = np.expand_dims(cv2.cvtColor(predictor_input_bgr, cv2.COLOR_BGR2GRAY),-1)
|
||||
|
||||
predicted_bgr = self.predictor ( predictor_input_g )
|
||||
|
||||
output = cv2.resize ( predicted_bgr, (self.output_size, self.output_size), cv2.INTER_LANCZOS4 )
|
||||
if debug:
|
||||
return (img_bgr,output,)
|
||||
return output
|
||||
1
models/Model_AVATAR/__init__.py
Normal file
1
models/Model_AVATAR/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .Model import Model
|
||||
153
models/Model_DF/Model.py
Normal file
153
models/Model_DF/Model.py
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from facelib import FaceType
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
encoderH5 = 'encoder.h5'
|
||||
decoder_srcH5 = 'decoder_src.h5'
|
||||
decoder_dstH5 = 'decoder_dst.h5'
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
self.set_vram_batch_requirements( {5:16,6:16,7:16,8:24,9:24,10:32,11:32,12:32,13:48} )
|
||||
|
||||
ae_input_layer = self.keras.layers.Input(shape=(128, 128, 3))
|
||||
mask_layer = self.keras.layers.Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder = self.Encoder(ae_input_layer)
|
||||
self.decoder_src = self.Decoder()
|
||||
self.decoder_dst = self.Decoder()
|
||||
|
||||
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.autoencoder_src = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder_src(self.encoder(ae_input_layer)))
|
||||
self.autoencoder_dst = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder_dst(self.encoder(ae_input_layer)))
|
||||
|
||||
if self.is_training_mode:
|
||||
self.autoencoder_src, self.autoencoder_dst = self.to_multi_gpu_model_if_possible ( [self.autoencoder_src, self.autoencoder_dst] )
|
||||
|
||||
optimizer = self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
|
||||
dssimloss = DSSIMMaskLossClass(self.tf)([mask_layer])
|
||||
self.autoencoder_src.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True ),
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True )
|
||||
])
|
||||
#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)]] )
|
||||
|
||||
#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 [ ('DF', 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 'masked_hist_match' not in in_options.keys() or in_options['masked_hist_match'] is None:
|
||||
in_options['masked_hist_match'] = True
|
||||
|
||||
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 Encoder(self, input_layer):
|
||||
x = input_layer
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
|
||||
x = self.keras.layers.Dense(512)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(8 * 8 * 512)(x)
|
||||
x = self.keras.layers.Reshape((8, 8, 512))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def Decoder(self):
|
||||
input_ = self.keras.layers.Input(shape=(16, 16, 512))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 512)
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
1
models/Model_DF/__init__.py
Normal file
1
models/Model_DF/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .Model import Model
|
||||
174
models/Model_H128/Model.py
Normal file
174
models/Model_H128/Model.py
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from facelib import FaceType
|
||||
|
||||
import cv2
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
encoderH5 = 'encoder.h5'
|
||||
decoder_srcH5 = 'decoder_src.h5'
|
||||
decoder_dstH5 = 'decoder_dst.h5'
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
self.set_vram_batch_requirements( {3:2,4:2,4:4,5:8,6:8,7:16,8:16,9:24,10:24,11:32,12:32,13:48} )
|
||||
|
||||
ae_input_layer = self.keras.layers.Input(shape=(128, 128, 3))
|
||||
mask_layer = self.keras.layers.Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder = self.Encoder(ae_input_layer, self.created_vram_gb)
|
||||
self.decoder_src = self.Decoder(self.created_vram_gb)
|
||||
self.decoder_dst = self.Decoder(self.created_vram_gb)
|
||||
|
||||
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.autoencoder_src = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder_src(self.encoder(ae_input_layer)))
|
||||
self.autoencoder_dst = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder_dst(self.encoder(ae_input_layer)))
|
||||
|
||||
if self.is_training_mode:
|
||||
self.autoencoder_src, self.autoencoder_dst = self.to_multi_gpu_model_if_possible ( [self.autoencoder_src, self.autoencoder_dst] )
|
||||
|
||||
optimizer = self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
|
||||
dssimloss = DSSIMMaskLossClass(self.tf)([mask_layer])
|
||||
self.autoencoder_src.compile(optimizer=optimizer, loss=[dssimloss, 'mae'])
|
||||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mae'])
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.HALF_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True ),
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.HALF_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True )
|
||||
])
|
||||
|
||||
#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)]])
|
||||
|
||||
#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 [ ('H128', 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 'masked_hist_match' not in in_options.keys() or in_options['masked_hist_match'] is None:
|
||||
in_options['masked_hist_match'] = True
|
||||
|
||||
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)
|
||||
|
||||
def Encoder(self, input_layer, created_vram_gb):
|
||||
x = input_layer
|
||||
|
||||
if created_vram_gb >= 5:
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Dense(512)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(8 * 8 * 512)(x)
|
||||
x = self.keras.layers.Reshape((8, 8, 512))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
else:
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Dense(256)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(8 * 8 * 256)(x)
|
||||
x = self.keras.layers.Reshape((8, 8, 256))(x)
|
||||
x = upscale(self.keras, x, 256)
|
||||
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def Decoder(self, created_vram_gb):
|
||||
if created_vram_gb >= 5:
|
||||
input_ = self.keras.layers.Input(shape=(16, 16, 512))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 512)
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
else:
|
||||
input_ = self.keras.layers.Input(shape=(16, 16, 256))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
x = upscale(self.keras, x, 64)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = upscale(self.keras, y, 64)
|
||||
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
1
models/Model_H128/__init__.py
Normal file
1
models/Model_H128/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .Model import Model
|
||||
167
models/Model_H64/Model.py
Normal file
167
models/Model_H64/Model.py
Normal file
|
|
@ -0,0 +1,167 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from facelib import FaceType
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
encoderH5 = 'encoder.h5'
|
||||
decoder_srcH5 = 'decoder_src.h5'
|
||||
decoder_dstH5 = 'decoder_dst.h5'
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
self.set_vram_batch_requirements( {2:2,3:4,4:8,5:16,6:32,7:32,8:32,9:48} )
|
||||
|
||||
ae_input_layer = self.keras.layers.Input(shape=(64, 64, 3))
|
||||
mask_layer = self.keras.layers.Input(shape=(64, 64, 1)) #same as output
|
||||
|
||||
self.encoder = self.Encoder(ae_input_layer, self.created_vram_gb)
|
||||
self.decoder_src = self.Decoder(self.created_vram_gb)
|
||||
self.decoder_dst = self.Decoder(self.created_vram_gb)
|
||||
|
||||
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.autoencoder_src = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder_src(self.encoder(ae_input_layer)))
|
||||
self.autoencoder_dst = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder_dst(self.encoder(ae_input_layer)))
|
||||
|
||||
if self.is_training_mode:
|
||||
self.autoencoder_src, self.autoencoder_dst = self.to_multi_gpu_model_if_possible ( [self.autoencoder_src, self.autoencoder_dst] )
|
||||
|
||||
optimizer = self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
|
||||
dssimloss = DSSIMMaskLossClass(self.tf)([mask_layer])
|
||||
self.autoencoder_src.compile(optimizer=optimizer, loss=[dssimloss, 'mae'])
|
||||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mae'])
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 64], [f.TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 64], [f.TRANSFORMED | f.HALF_FACE | f.MODE_M | f.MASK_FULL, 64] ], random_flip=True ),
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 64], [f.TRANSFORMED | f.HALF_FACE | f.MODE_BGR, 64], [f.TRANSFORMED | f.HALF_FACE | f.MODE_M | f.MASK_FULL, 64] ], random_flip=True )
|
||||
])
|
||||
|
||||
#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)]] )
|
||||
|
||||
#override
|
||||
def onTrainOneEpoch(self, sample):
|
||||
warped_src, target_src, target_src_full_mask = sample[0]
|
||||
warped_dst, target_dst, target_dst_full_mask = sample[1]
|
||||
|
||||
loss_src = self.autoencoder_src.train_on_batch( [warped_src, target_src_full_mask], [target_src, target_src_full_mask] )
|
||||
loss_dst = self.autoencoder_dst.train_on_batch( [warped_dst, target_dst_full_mask], [target_dst, target_dst_full_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]
|
||||
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 [ ('H64', np.concatenate ( st, axis=0 ) ) ]
|
||||
|
||||
def predictor_func (self, face):
|
||||
|
||||
face_64_bgr = face[...,0:3]
|
||||
face_64_mask = np.expand_dims(face[...,3],-1)
|
||||
|
||||
x, mx = self.autoencoder_src.predict ( [ np.expand_dims(face_64_bgr,0), np.expand_dims(face_64_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 'masked_hist_match' not in in_options.keys() or in_options['masked_hist_match'] is None:
|
||||
in_options['masked_hist_match'] = True
|
||||
|
||||
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)
|
||||
|
||||
def Encoder(self, input_layer, created_vram_gb):
|
||||
x = input_layer
|
||||
if created_vram_gb >= 4:
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Dense(1024)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(4 * 4 * 1024)(x)
|
||||
x = self.keras.layers.Reshape((4, 4, 1024))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
else:
|
||||
x = conv(self.keras, x, 128 )
|
||||
x = conv(self.keras, x, 256 )
|
||||
x = conv(self.keras, x, 512 )
|
||||
x = conv(self.keras, x, 768 )
|
||||
x = self.keras.layers.Dense(512)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(4 * 4 * 512)(x)
|
||||
x = self.keras.layers.Reshape((4, 4, 512))(x)
|
||||
x = upscale(self.keras, x, 256)
|
||||
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def Decoder(self, created_vram_gb):
|
||||
if created_vram_gb >= 4:
|
||||
input_ = self.keras.layers.Input(shape=(8, 8, 512))
|
||||
else:
|
||||
input_ = self.keras.layers.Input(shape=(8, 8, 256))
|
||||
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
x = upscale(self.keras, x, 64)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = upscale(self.keras, y, 64)
|
||||
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
1
models/Model_H64/__init__.py
Normal file
1
models/Model_H64/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .Model import Model
|
||||
164
models/Model_LIAEF128/Model.py
Normal file
164
models/Model_LIAEF128/Model.py
Normal file
|
|
@ -0,0 +1,164 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from facelib import FaceType
|
||||
|
||||
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):
|
||||
self.set_vram_batch_requirements( {5:4,6:8,7:12,8:16,9:20,10:24,11:24,12:32,13:48} )
|
||||
|
||||
ae_input_layer = self.keras.layers.Input(shape=(128, 128, 3))
|
||||
mask_layer = self.keras.layers.Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder = self.Encoder(ae_input_layer)
|
||||
self.decoder = self.Decoder()
|
||||
self.inter_B = self.Intermediate ()
|
||||
self.inter_AB = self.Intermediate ()
|
||||
|
||||
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 = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder(self.keras.layers.Concatenate()([AB, AB])) )
|
||||
self.autoencoder_dst = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder(self.keras.layers.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] )
|
||||
|
||||
optimizer = self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
|
||||
dssimloss = DSSIMMaskLossClass(self.tf)([mask_layer])
|
||||
self.autoencoder_src.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True ),
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True )
|
||||
])
|
||||
|
||||
#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 [ ('LIAEF128', 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 'masked_hist_match' not in in_options.keys() or in_options['masked_hist_match'] is None:
|
||||
in_options['masked_hist_match'] = True
|
||||
|
||||
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 Encoder(self, input_layer,):
|
||||
x = input_layer
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Flatten()(x)
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def Intermediate(self):
|
||||
input_layer = self.keras.layers.Input(shape=(None, 8 * 8 * 1024))
|
||||
x = input_layer
|
||||
x = self.keras.layers.Dense(256)(x)
|
||||
x = self.keras.layers.Dense(8 * 8 * 512)(x)
|
||||
x = self.keras.layers.Reshape((8, 8, 512))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def Decoder(self):
|
||||
input_ = self.keras.layers.Input(shape=(16, 16, 1024))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 512)
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid' )(y)
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
1
models/Model_LIAEF128/__init__.py
Normal file
1
models/Model_LIAEF128/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .Model import Model
|
||||
164
models/Model_LIAEF128YAW/Model.py
Normal file
164
models/Model_LIAEF128YAW/Model.py
Normal file
|
|
@ -0,0 +1,164 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from facelib import FaceType
|
||||
|
||||
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):
|
||||
self.set_vram_batch_requirements( {5:4,6:8,7:12,8:16,9:20,10:24,11:24,12:32,13:48} )
|
||||
|
||||
ae_input_layer = self.keras.layers.Input(shape=(128, 128, 3))
|
||||
mask_layer = self.keras.layers.Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder = self.Encoder(ae_input_layer)
|
||||
self.decoder = self.Decoder()
|
||||
self.inter_B = self.Intermediate ()
|
||||
self.inter_AB = self.Intermediate ()
|
||||
|
||||
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 = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder(self.keras.layers.Concatenate()([AB, AB])) )
|
||||
self.autoencoder_dst = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder(self.keras.layers.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] )
|
||||
|
||||
optimizer = self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
|
||||
dssimloss = DSSIMMaskLossClass(self.tf)([mask_layer])
|
||||
self.autoencoder_src.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE_YAW_SORTED_AS_TARGET, self.training_data_src_path, target_training_data_path=self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True ),
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128] ], random_flip=True )
|
||||
])
|
||||
|
||||
#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 'masked_hist_match' not in in_options.keys() or in_options['masked_hist_match'] is None:
|
||||
in_options['masked_hist_match'] = True
|
||||
|
||||
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 Encoder(self, input_layer,):
|
||||
x = input_layer
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Flatten()(x)
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def Intermediate(self):
|
||||
input_layer = self.keras.layers.Input(shape=(None, 8 * 8 * 1024))
|
||||
x = input_layer
|
||||
x = self.keras.layers.Dense(256)(x)
|
||||
x = self.keras.layers.Dense(8 * 8 * 512)(x)
|
||||
x = self.keras.layers.Reshape((8, 8, 512))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def Decoder(self):
|
||||
input_ = self.keras.layers.Input(shape=(16, 16, 1024))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 512)
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid' )(y)
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
1
models/Model_LIAEF128YAW/__init__.py
Normal file
1
models/Model_LIAEF128YAW/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .Model import Model
|
||||
217
models/Model_MIAEF128/Model.py
Normal file
217
models/Model_MIAEF128/Model.py
Normal file
|
|
@ -0,0 +1,217 @@
|
|||
from models import ModelBase
|
||||
from models import TrainingDataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from facelib import FaceType
|
||||
|
||||
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):
|
||||
self.set_vram_batch_requirements( {5:4,6:8,7:12,8:16,9:20,10:24,11:24,12:32,13:48} )
|
||||
|
||||
ae_input_layer = self.keras.layers.Input(shape=(128, 128, 3))
|
||||
mask_layer = self.keras.layers.Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder = self.Encoder(ae_input_layer)
|
||||
self.decoderMask = self.DecoderMask()
|
||||
self.decoderCommonA = self.DecoderCommon()
|
||||
self.decoderCommonB = self.DecoderCommon()
|
||||
self.decoderRGB = self.DecoderRGB()
|
||||
self.decoderBW = self.DecoderBW()
|
||||
self.inter_A = self.Intermediate ()
|
||||
self.inter_B = self.Intermediate ()
|
||||
|
||||
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 = self.keras.layers.Concatenate()([A, A])
|
||||
inter_B_A = self.keras.layers.Concatenate()([B, A])
|
||||
|
||||
x1,m1 = self.decoderCommonA (inter_A_A)
|
||||
x2,m2 = self.decoderCommonA (inter_A_A)
|
||||
self.autoencoder_src = self.keras.models.Model([ae_input_layer,mask_layer],
|
||||
[ self.decoderBW (self.keras.layers.Concatenate()([x1,x2]) ),
|
||||
self.decoderMask(self.keras.layers.Concatenate()([m1,m2]) )
|
||||
])
|
||||
|
||||
x1,m1 = self.decoderCommonA (inter_A_A)
|
||||
x2,m2 = self.decoderCommonB (inter_A_A)
|
||||
self.autoencoder_src_RGB = self.keras.models.Model([ae_input_layer,mask_layer],
|
||||
[ self.decoderRGB (self.keras.layers.Concatenate()([x1,x2]) ),
|
||||
self.decoderMask (self.keras.layers.Concatenate()([m1,m2]) )
|
||||
])
|
||||
|
||||
x1,m1 = self.decoderCommonA (inter_B_A)
|
||||
x2,m2 = self.decoderCommonB (inter_B_A)
|
||||
self.autoencoder_dst = self.keras.models.Model([ae_input_layer,mask_layer],
|
||||
[ self.decoderRGB (self.keras.layers.Concatenate()([x1,x2]) ),
|
||||
self.decoderMask (self.keras.layers.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] )
|
||||
|
||||
optimizer = self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
|
||||
dssimloss = DSSIMMaskLossClass(self.tf)([mask_layer])
|
||||
self.autoencoder_src.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
from models import TrainingDataGenerator
|
||||
f = TrainingDataGenerator.SampleTypeFlags
|
||||
self.set_training_data_generators ([
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_GGG, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_G , 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_GGG, 128] ], random_flip=True ),
|
||||
TrainingDataGenerator(TrainingDataType.FACE, self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, output_sample_types=[ [f.WARPED_TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_BGR, 128], [f.TRANSFORMED | f.FULL_FACE | f.MODE_M | f.MASK_FULL, 128]], random_flip=True )
|
||||
])
|
||||
#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 'masked_hist_match' not in in_options.keys() or in_options['masked_hist_match'] is None:
|
||||
in_options['masked_hist_match'] = False
|
||||
|
||||
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 Encoder(self, input_layer,):
|
||||
x = input_layer
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Flatten()(x)
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def Intermediate(self):
|
||||
input_layer = self.keras.layers.Input(shape=(None, 8 * 8 * 1024))
|
||||
x = input_layer
|
||||
x = self.keras.layers.Dense(256)(x)
|
||||
x = self.keras.layers.Dense(8 * 8 * 512)(x)
|
||||
x = self.keras.layers.Reshape((8, 8, 512))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def DecoderCommon(self):
|
||||
input_ = self.keras.layers.Input(shape=(16, 16, 1024))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
|
||||
y = input_
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = upscale(self.keras, y, 64)
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
|
||||
def DecoderRGB(self):
|
||||
input_ = self.keras.layers.Input(shape=(128, 128, 256))
|
||||
x = input_
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return self.keras.models.Model(input_, [x])
|
||||
|
||||
def DecoderBW(self):
|
||||
input_ = self.keras.layers.Input(shape=(128, 128, 256))
|
||||
x = input_
|
||||
x = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return self.keras.models.Model(input_, [x])
|
||||
|
||||
def DecoderMask(self):
|
||||
input_ = self.keras.layers.Input(shape=(128, 128, 128))
|
||||
y = input_
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
return self.keras.models.Model(input_, [y])
|
||||
1
models/Model_MIAEF128/__init__.py
Normal file
1
models/Model_MIAEF128/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .Model import Model
|
||||
149
models/TrainingDataGenerator.py
Normal file
149
models/TrainingDataGenerator.py
Normal file
|
|
@ -0,0 +1,149 @@
|
|||
from facelib import FaceType
|
||||
from facelib import LandmarksProcessor
|
||||
import cv2
|
||||
import numpy as np
|
||||
from models import TrainingDataGeneratorBase
|
||||
from utils import image_utils
|
||||
from utils import random_utils
|
||||
from enum import IntEnum
|
||||
from models import TrainingDataType
|
||||
|
||||
class TrainingDataGenerator(TrainingDataGeneratorBase):
|
||||
class SampleTypeFlags(IntEnum):
|
||||
SOURCE = 0x000001,
|
||||
WARPED = 0x000002,
|
||||
WARPED_TRANSFORMED = 0x000004,
|
||||
TRANSFORMED = 0x000008,
|
||||
|
||||
HALF_FACE = 0x000010,
|
||||
FULL_FACE = 0x000020,
|
||||
HEAD_FACE = 0x000040,
|
||||
AVATAR_FACE = 0x000080,
|
||||
MARK_ONLY_FACE = 0x000100,
|
||||
|
||||
MODE_BGR = 0x001000, #BGR
|
||||
MODE_G = 0x002000, #Grayscale
|
||||
MODE_GGG = 0x004000, #3xGrayscale
|
||||
MODE_M = 0x008000, #mask only
|
||||
MODE_BGR_SHUFFLE = 0x010000, #BGR shuffle
|
||||
|
||||
MASK_FULL = 0x100000,
|
||||
MASK_EYES = 0x200000,
|
||||
|
||||
#overrided
|
||||
def onInitialize(self, random_flip=False, normalize_tanh=False, rotation_range=[-10,10], scale_range=[-0.05, 0.05], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05], output_sample_types=[], **kwargs):
|
||||
self.random_flip = random_flip
|
||||
self.normalize_tanh = normalize_tanh
|
||||
self.output_sample_types = output_sample_types
|
||||
self.rotation_range = rotation_range
|
||||
self.scale_range = scale_range
|
||||
self.tx_range = tx_range
|
||||
self.ty_range = ty_range
|
||||
|
||||
#overrided
|
||||
def onProcessSample(self, sample, debug):
|
||||
source = sample.load_bgr()
|
||||
h,w,c = source.shape
|
||||
|
||||
is_face_sample = self.trainingdatatype >= TrainingDataType.FACE_BEGIN and self.trainingdatatype <= TrainingDataType.FACE_END
|
||||
|
||||
if debug and is_face_sample:
|
||||
LandmarksProcessor.draw_landmarks (source, sample.landmarks, (0, 1, 0))
|
||||
|
||||
params = image_utils.gen_warp_params(source, self.random_flip, rotation_range=self.rotation_range, scale_range=self.scale_range, tx_range=self.tx_range, ty_range=self.ty_range )
|
||||
|
||||
images = [[None]*3 for _ in range(4)]
|
||||
|
||||
outputs = []
|
||||
for t,size in self.output_sample_types:
|
||||
if t & self.SampleTypeFlags.SOURCE != 0:
|
||||
img_type = 0
|
||||
elif t & self.SampleTypeFlags.WARPED != 0:
|
||||
img_type = 1
|
||||
elif t & self.SampleTypeFlags.WARPED_TRANSFORMED != 0:
|
||||
img_type = 2
|
||||
elif t & self.SampleTypeFlags.TRANSFORMED != 0:
|
||||
img_type = 3
|
||||
else:
|
||||
raise ValueError ('expected SampleTypeFlags type')
|
||||
|
||||
mask_type = 0
|
||||
if t & self.SampleTypeFlags.MASK_FULL != 0:
|
||||
mask_type = 1
|
||||
elif t & self.SampleTypeFlags.MASK_EYES != 0:
|
||||
mask_type = 2
|
||||
|
||||
if images[img_type][mask_type] is None:
|
||||
img = source
|
||||
if is_face_sample:
|
||||
if mask_type == 1:
|
||||
img = np.concatenate( (img, LandmarksProcessor.get_image_hull_mask (source, sample.landmarks) ), -1 )
|
||||
elif mask_type == 2:
|
||||
mask = LandmarksProcessor.get_image_eye_mask (source, sample.landmarks)
|
||||
mask = np.expand_dims (cv2.blur (mask, ( w // 32, w // 32 ) ), -1)
|
||||
mask[mask > 0.0] = 1.0
|
||||
img = np.concatenate( (img, mask ), -1 )
|
||||
|
||||
images[img_type][mask_type] = image_utils.warp_by_params (params, img, (img_type==1 or img_type==2), (img_type==2 or img_type==3), img_type != 0)
|
||||
|
||||
img = images[img_type][mask_type]
|
||||
|
||||
target_face_type = -1
|
||||
if t & self.SampleTypeFlags.HALF_FACE != 0:
|
||||
target_face_type = FaceType.HALF
|
||||
elif t & self.SampleTypeFlags.FULL_FACE != 0:
|
||||
target_face_type = FaceType.FULL
|
||||
elif t & self.SampleTypeFlags.HEAD_FACE != 0:
|
||||
target_face_type = FaceType.HEAD
|
||||
elif t & self.SampleTypeFlags.AVATAR_FACE != 0:
|
||||
target_face_type = FaceType.AVATAR
|
||||
elif t & self.SampleTypeFlags.MARK_ONLY_FACE != 0:
|
||||
target_face_type = FaceType.MARK_ONLY
|
||||
|
||||
if is_face_sample and target_face_type != -1 and target_face_type != FaceType.MARK_ONLY:
|
||||
if target_face_type > sample.face_type:
|
||||
raise Exception ('sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, target_face_type) )
|
||||
|
||||
img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, size, target_face_type), (size,size), flags=cv2.INTER_LANCZOS4 )
|
||||
else:
|
||||
img = cv2.resize( img, (size,size), cv2.INTER_LANCZOS4 )
|
||||
|
||||
img_bgr = img[...,0:3]
|
||||
img_mask = img[...,3:4]
|
||||
|
||||
if t & self.SampleTypeFlags.MODE_BGR != 0:
|
||||
img = img
|
||||
elif t & self.SampleTypeFlags.MODE_BGR_SHUFFLE != 0:
|
||||
img_bgr = np.take (img_bgr, np.random.permutation(img_bgr.shape[-1]), axis=-1)
|
||||
img = np.concatenate ( (img_bgr,img_mask) , -1 )
|
||||
elif t & self.SampleTypeFlags.MODE_G != 0:
|
||||
img = np.concatenate ( (np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1),img_mask) , -1 )
|
||||
elif t & self.SampleTypeFlags.MODE_GGG != 0:
|
||||
img = np.concatenate ( ( np.repeat ( np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1), (3,), -1), img_mask), -1)
|
||||
elif is_face_sample and t & self.SampleTypeFlags.MODE_M != 0:
|
||||
if mask_type== 0:
|
||||
raise ValueError ('no mask mode defined')
|
||||
img = img_mask
|
||||
else:
|
||||
raise ValueError ('expected SampleTypeFlags mode')
|
||||
|
||||
if not debug and self.normalize_tanh:
|
||||
img = img * 2.0 - 1.0
|
||||
|
||||
outputs.append ( img )
|
||||
|
||||
if debug:
|
||||
result = ()
|
||||
|
||||
for output in outputs:
|
||||
if output.shape[2] < 4:
|
||||
result += (output,)
|
||||
elif output.shape[2] == 4:
|
||||
result += (output[...,0:3]*output[...,3:4],)
|
||||
|
||||
return result
|
||||
else:
|
||||
return outputs
|
||||
|
||||
|
||||
|
||||
245
models/TrainingDataGeneratorBase.py
Normal file
245
models/TrainingDataGeneratorBase.py
Normal file
|
|
@ -0,0 +1,245 @@
|
|||
import traceback
|
||||
import random
|
||||
from pathlib import Path
|
||||
from tqdm import tqdm
|
||||
import numpy as np
|
||||
import cv2
|
||||
from utils.AlignedPNG import AlignedPNG
|
||||
from utils import iter_utils
|
||||
from utils import Path_utils
|
||||
from .BaseTypes import TrainingDataType
|
||||
from .BaseTypes import TrainingDataSample
|
||||
from facelib import FaceType
|
||||
from facelib import LandmarksProcessor
|
||||
|
||||
'''
|
||||
You can implement your own TrainingDataGenerator
|
||||
'''
|
||||
class TrainingDataGeneratorBase(object):
|
||||
cache = dict()
|
||||
|
||||
#DONT OVERRIDE
|
||||
#use YourOwnTrainingDataGenerator (..., your_opt=1)
|
||||
#and then this opt will be passed in YourOwnTrainingDataGenerator.onInitialize ( your_opt )
|
||||
def __init__ (self, trainingdatatype, training_data_path, target_training_data_path=None, debug=False, batch_size=1, **kwargs):
|
||||
if not isinstance(trainingdatatype, TrainingDataType):
|
||||
raise Exception('TrainingDataGeneratorBase() trainingdatatype is not TrainingDataType')
|
||||
|
||||
if training_data_path is None:
|
||||
raise Exception('training_data_path is None')
|
||||
|
||||
self.training_data_path = Path(training_data_path)
|
||||
self.target_training_data_path = Path(target_training_data_path) if target_training_data_path is not None else None
|
||||
|
||||
self.debug = debug
|
||||
self.batch_size = 1 if self.debug else batch_size
|
||||
self.trainingdatatype = trainingdatatype
|
||||
self.data = TrainingDataGeneratorBase.load (trainingdatatype, self.training_data_path, self.target_training_data_path)
|
||||
|
||||
if self.debug:
|
||||
self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, self.data)]
|
||||
else:
|
||||
if len(self.data) > 1:
|
||||
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, self.data[0::2] ),
|
||||
iter_utils.SubprocessGenerator ( self.batch_func, self.data[1::2] )]
|
||||
else:
|
||||
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, self.data )]
|
||||
|
||||
self.generator_counter = -1
|
||||
self.onInitialize(**kwargs)
|
||||
|
||||
#overridable
|
||||
def onInitialize(self, **kwargs):
|
||||
#your TrainingDataGenerator initialization here
|
||||
pass
|
||||
|
||||
#overridable
|
||||
def onProcessSample(self, sample, debug):
|
||||
#process sample and return tuple of images for your model in onTrainOneEpoch
|
||||
return ( np.zeros( (64,64,4), dtype=np.float32 ), )
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
self.generator_counter += 1
|
||||
generator = self.generators[self.generator_counter % len(self.generators) ]
|
||||
x = next(generator)
|
||||
return x
|
||||
|
||||
def batch_func(self, data):
|
||||
data_len = len(data)
|
||||
if data_len == 0:
|
||||
raise ValueError('No training data provided.')
|
||||
|
||||
if self.trainingdatatype == TrainingDataType.FACE_YAW_SORTED or self.trainingdatatype == TrainingDataType.FACE_YAW_SORTED_AS_TARGET:
|
||||
if all ( [ x == None for x in data] ):
|
||||
raise ValueError('Not enough training data. Gather more faces!')
|
||||
|
||||
if self.trainingdatatype == TrainingDataType.IMAGE or self.trainingdatatype == TrainingDataType.FACE:
|
||||
shuffle_idxs = []
|
||||
elif self.trainingdatatype == TrainingDataType.FACE_YAW_SORTED or self.trainingdatatype == TrainingDataType.FACE_YAW_SORTED_AS_TARGET:
|
||||
shuffle_idxs = []
|
||||
shuffle_idxs_2D = [[]]*data_len
|
||||
|
||||
while True:
|
||||
|
||||
batches = None
|
||||
for n_batch in range(0, self.batch_size):
|
||||
while True:
|
||||
sample = None
|
||||
|
||||
if self.trainingdatatype == TrainingDataType.IMAGE or self.trainingdatatype == TrainingDataType.FACE:
|
||||
if len(shuffle_idxs) == 0:
|
||||
shuffle_idxs = [ i for i in range(0, data_len) ]
|
||||
random.shuffle(shuffle_idxs)
|
||||
idx = shuffle_idxs.pop()
|
||||
sample = data[ idx ]
|
||||
elif self.trainingdatatype == TrainingDataType.FACE_YAW_SORTED or self.trainingdatatype == TrainingDataType.FACE_YAW_SORTED_AS_TARGET:
|
||||
if len(shuffle_idxs) == 0:
|
||||
shuffle_idxs = [ i for i in range(0, data_len) ]
|
||||
random.shuffle(shuffle_idxs)
|
||||
|
||||
idx = shuffle_idxs.pop()
|
||||
if data[idx] != None:
|
||||
if len(shuffle_idxs_2D[idx]) == 0:
|
||||
shuffle_idxs_2D[idx] = [ i for i in range(0, len(data[idx])) ]
|
||||
random.shuffle(shuffle_idxs_2D[idx])
|
||||
|
||||
idx2 = shuffle_idxs_2D[idx].pop()
|
||||
sample = data[idx][idx2]
|
||||
|
||||
if sample is not None:
|
||||
try:
|
||||
x = self.onProcessSample (sample, self.debug)
|
||||
except:
|
||||
raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) )
|
||||
|
||||
if type(x) != tuple and type(x) != list:
|
||||
raise Exception('TrainingDataGenerator.onProcessSample() returns NOT tuple/list')
|
||||
|
||||
x_len = len(x)
|
||||
if batches is None:
|
||||
batches = [ [] for _ in range(0,x_len) ]
|
||||
|
||||
for i in range(0,x_len):
|
||||
batches[i].append ( x[i] )
|
||||
|
||||
break
|
||||
|
||||
yield [ np.array(batch) for batch in batches]
|
||||
|
||||
def get_dict_state(self):
|
||||
return {}
|
||||
|
||||
def set_dict_state(self, state):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def load(trainingdatatype, training_data_path, target_training_data_path=None):
|
||||
cache = TrainingDataGeneratorBase.cache
|
||||
|
||||
if str(training_data_path) not in cache.keys():
|
||||
cache[str(training_data_path)] = [None]*TrainingDataType.QTY
|
||||
|
||||
if target_training_data_path is not None and str(target_training_data_path) not in cache.keys():
|
||||
cache[str(target_training_data_path)] = [None]*TrainingDataType.QTY
|
||||
|
||||
datas = cache[str(training_data_path)]
|
||||
|
||||
if trainingdatatype == TrainingDataType.IMAGE:
|
||||
if datas[trainingdatatype] is None:
|
||||
datas[trainingdatatype] = [ TrainingDataSample(filename=filename) for filename in tqdm( Path_utils.get_image_paths(training_data_path), desc="Loading" ) ]
|
||||
|
||||
elif trainingdatatype == TrainingDataType.FACE:
|
||||
if datas[trainingdatatype] is None:
|
||||
datas[trainingdatatype] = X_LOAD( [ TrainingDataSample(filename=filename) for filename in Path_utils.get_image_paths(training_data_path) ] )
|
||||
|
||||
elif trainingdatatype == TrainingDataType.FACE_YAW_SORTED:
|
||||
if datas[trainingdatatype] is None:
|
||||
datas[trainingdatatype] = X_YAW_SORTED( TrainingDataGeneratorBase.load(TrainingDataType.FACE, training_data_path) )
|
||||
|
||||
elif trainingdatatype == TrainingDataType.FACE_YAW_SORTED_AS_TARGET:
|
||||
if datas[trainingdatatype] is None:
|
||||
if target_training_data_path is None:
|
||||
raise Exception('target_training_data_path is None for FACE_YAW_SORTED_AS_TARGET')
|
||||
datas[trainingdatatype] = X_YAW_AS_Y_SORTED( TrainingDataGeneratorBase.load(TrainingDataType.FACE_YAW_SORTED, training_data_path), TrainingDataGeneratorBase.load(TrainingDataType.FACE_YAW_SORTED, target_training_data_path) )
|
||||
|
||||
return datas[trainingdatatype]
|
||||
|
||||
def X_LOAD ( RAWS ):
|
||||
sample_list = []
|
||||
|
||||
for s in tqdm( RAWS, desc="Loading" ):
|
||||
|
||||
s_filename_path = Path(s.filename)
|
||||
if s_filename_path.suffix != '.png':
|
||||
print ("%s is not a png file required for training" % (s_filename_path.name) )
|
||||
continue
|
||||
|
||||
a_png = AlignedPNG.load ( str(s_filename_path) )
|
||||
if a_png is None:
|
||||
print ("%s failed to load" % (s_filename_path.name) )
|
||||
continue
|
||||
|
||||
d = a_png.getFaceswapDictData()
|
||||
if d is None or d['landmarks'] is None or d['yaw_value'] is None:
|
||||
print ("%s - no embedded faceswap info found required for training" % (s_filename_path.name) )
|
||||
continue
|
||||
|
||||
face_type = d['face_type'] if 'face_type' in d.keys() else 'full_face'
|
||||
face_type = FaceType.fromString (face_type)
|
||||
sample_list.append( s.copy_and_set(face_type=face_type, shape=a_png.get_shape(), landmarks=d['landmarks'], yaw=d['yaw_value']) )
|
||||
|
||||
return sample_list
|
||||
|
||||
def X_YAW_SORTED( YAW_RAWS ):
|
||||
|
||||
lowest_yaw, highest_yaw = -32, +32
|
||||
gradations = 64
|
||||
diff_rot_per_grad = abs(highest_yaw-lowest_yaw) / gradations
|
||||
|
||||
yaws_sample_list = [None]*gradations
|
||||
|
||||
for i in tqdm( range(0, gradations), desc="Sorting" ):
|
||||
yaw = lowest_yaw + i*diff_rot_per_grad
|
||||
next_yaw = lowest_yaw + (i+1)*diff_rot_per_grad
|
||||
|
||||
yaw_samples = []
|
||||
for s in YAW_RAWS:
|
||||
s_yaw = s.yaw
|
||||
if (i == 0 and s_yaw < next_yaw) or \
|
||||
(i < gradations-1 and s_yaw >= yaw and s_yaw < next_yaw) or \
|
||||
(i == gradations-1 and s_yaw >= yaw):
|
||||
yaw_samples.append ( s )
|
||||
|
||||
if len(yaw_samples) > 0:
|
||||
yaws_sample_list[i] = yaw_samples
|
||||
|
||||
return yaws_sample_list
|
||||
|
||||
def X_YAW_AS_Y_SORTED (s, t):
|
||||
l = len(s)
|
||||
if l != len(t):
|
||||
raise Exception('X_YAW_AS_Y_SORTED() s_len != t_len')
|
||||
b = l // 2
|
||||
|
||||
s_idxs = np.argwhere ( np.array ( [ 1 if x != None else 0 for x in s] ) == 1 )[:,0]
|
||||
t_idxs = np.argwhere ( np.array ( [ 1 if x != None else 0 for x in t] ) == 1 )[:,0]
|
||||
|
||||
new_s = [None]*l
|
||||
|
||||
for t_idx in t_idxs:
|
||||
search_idxs = []
|
||||
for i in range(0,l):
|
||||
search_idxs += [t_idx - i, (l-t_idx-1) - i, t_idx + i, (l-t_idx-1) + i]
|
||||
|
||||
for search_idx in search_idxs:
|
||||
if search_idx in s_idxs:
|
||||
mirrored = ( t_idx != search_idx and ((t_idx < b and search_idx >= b) or (search_idx < b and t_idx >= b)) )
|
||||
new_s[t_idx] = [ sample.copy_and_set(mirror=True, yaw=-sample.yaw, landmarks=LandmarksProcessor.mirror_landmarks (sample.landmarks, sample.shape[1] ))
|
||||
for sample in s[search_idx]
|
||||
] if mirrored else s[search_idx]
|
||||
break
|
||||
|
||||
return new_s
|
||||
13
models/__init__.py
Normal file
13
models/__init__.py
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
from .BaseTypes import TrainingDataType
|
||||
from .BaseTypes import TrainingDataSample
|
||||
|
||||
from .ModelBase import ModelBase
|
||||
from .ConverterBase import ConverterBase
|
||||
from .ConverterMasked import ConverterMasked
|
||||
from .ConverterImage import ConverterImage
|
||||
from .TrainingDataGeneratorBase import TrainingDataGeneratorBase
|
||||
from .TrainingDataGenerator import TrainingDataGenerator
|
||||
|
||||
def import_model(name):
|
||||
module = __import__('Model_'+name, globals(), locals(), [], 1)
|
||||
return getattr(module, 'Model')
|
||||
Loading…
Add table
Add a link
Reference in a new issue