diff --git a/converters/ConverterMasked.py b/converters/ConverterMasked.py index fd853e9..34e494c 100644 --- a/converters/ConverterMasked.py +++ b/converters/ConverterMasked.py @@ -1,15 +1,18 @@ +import time import traceback -from .Converter import Converter -from facelib import LandmarksProcessor -from facelib import FaceType -from facelib import FANSegmentator + import cv2 import numpy as np + import imagelib +from facelib import FaceType, FANSegmentator, LandmarksProcessor from interact import interact as io from joblib import SubprocessFunctionCaller from utils.pickle_utils import AntiPickler -import time + +from .Converter import Converter + + ''' default_mode = {1:'overlay', 2:'hist-match', @@ -93,10 +96,10 @@ class ConverterMasked(Converter): self.blur_mask_modifier = base_blur_mask_modifier + np.clip ( io.input_int ("Choose blur mask modifier [-200..200] (skip:%d) : " % (default_blur_mask_modifier), default_blur_mask_modifier), -200, 200) self.output_face_scale = np.clip ( 1.0 + io.input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0)*0.01, 0.5, 1.5) - + if self.mode != 'raw': self.color_transfer_mode = io.input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ", None, ['rct','lct']) - + self.super_resolution = io.input_bool("Apply super resolution? (y/n ?:help skip:n) : ", False, help_message="Enhance details by applying DCSCN network.") if self.mode != 'raw': @@ -173,12 +176,12 @@ class ConverterMasked(Converter): prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC) elif self.mask_mode >= 3 and self.mask_mode <= 6: - if self.mask_mode == 3 or self.mask_mode == 5 or self.mask_mode == 6: + if self.mask_mode == 3 or self.mask_mode == 5 or self.mask_mode == 6: prd_face_bgr_256 = cv2.resize (prd_face_bgr, (256,256) ) prd_face_bgr_256_mask = self.fan_seg.extract( prd_face_bgr_256 ) FAN_prd_face_mask_a_0 = cv2.resize (prd_face_bgr_256_mask, (output_size,output_size), cv2.INTER_CUBIC) - if self.mask_mode == 4 or self.mask_mode == 5 or self.mask_mode == 6: + if self.mask_mode == 4 or self.mask_mode == 5 or self.mask_mode == 6: face_256_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, 256, face_type=FaceType.FULL) dst_face_256_bgr = cv2.warpAffine(img_bgr, face_256_mat, (256, 256), flags=cv2.INTER_LANCZOS4 ) dst_face_256_mask = self.fan_seg.extract( dst_face_256_bgr ) @@ -192,7 +195,7 @@ class ConverterMasked(Converter): prd_face_mask_a_0 = FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0 elif self.mask_mode == 6: prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0 - + prd_face_mask_a_0[ prd_face_mask_a_0 < 0.001 ] = 0.0 prd_face_mask_a = prd_face_mask_a_0[...,np.newaxis] @@ -316,7 +319,7 @@ class ConverterMasked(Converter): if self.masked_hist_match: hist_mask_a *= prd_face_mask_a - + white = (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32) hist_match_1 = prd_face_bgr*hist_mask_a + white @@ -326,10 +329,10 @@ class ConverterMasked(Converter): hist_match_2[ hist_match_1 > 1.0 ] = 1.0 prd_face_bgr = imagelib.color_hist_match(hist_match_1, hist_match_2, self.hist_match_threshold ) - + #if self.masked_hist_match: # prd_face_bgr -= white - + if self.mode == 'hist-match-bw': prd_face_bgr = prd_face_bgr.astype(dtype=np.float32) @@ -401,7 +404,7 @@ class ConverterMasked(Converter): if debug: debugs += [ np.clip( cv2.warpAffine( new_out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - + 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 ) @@ -430,4 +433,3 @@ class ConverterMasked(Converter): debugs += [out_img.copy()] return debugs if debug else out_img - diff --git a/models/Model_DEV_FANSEG/Model.py b/models/Model_DEV_FANSEG/Model.py index 6184a43..ac8435d 100644 --- a/models/Model_DEV_FANSEG/Model.py +++ b/models/Model_DEV_FANSEG/Model.py @@ -48,7 +48,7 @@ class Model(ModelBase): SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, sample_process_options=SampleProcessor.Options(random_flip=True), output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution' : self.resolution, 'motion_blur':(25, 1) }, - { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_M, t.FACE_MASK_FULL), 'resolution': self.resolution }, + { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_M), 'resolution': self.resolution }, ]), SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, diff --git a/models/Model_DF/Model.py b/models/Model_DF/Model.py index 879144a..c1725c2 100644 --- a/models/Model_DF/Model.py +++ b/models/Model_DF/Model.py @@ -47,7 +47,7 @@ class Model(ModelBase): t = SampleProcessor.Types output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_BGR), 'resolution':128}, { 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_BGR), 'resolution':128}, - { 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_M, t.FACE_MASK_FULL), 'resolution':128} ] + { 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_M), 'resolution':128} ] self.set_training_data_generators ([ SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, diff --git a/models/Model_H128/Model.py b/models/Model_H128/Model.py index 11041e6..05d6bd0 100644 --- a/models/Model_H128/Model.py +++ b/models/Model_H128/Model.py @@ -57,7 +57,7 @@ class Model(ModelBase): t = SampleProcessor.Types output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_BGR), 'resolution':128}, { 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_BGR), 'resolution':128}, - { 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_M, t.FACE_MASK_FULL), 'resolution':128} ] + { 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_M), 'resolution':128} ] self.set_training_data_generators ([ SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, diff --git a/models/Model_H64/Model.py b/models/Model_H64/Model.py index d37eaa3..3368a3b 100644 --- a/models/Model_H64/Model.py +++ b/models/Model_H64/Model.py @@ -58,7 +58,7 @@ class Model(ModelBase): t = SampleProcessor.Types output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_BGR), 'resolution':64}, { 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_BGR), 'resolution':64}, - { 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_M, t.FACE_MASK_FULL), 'resolution':64} ] + { 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_HALF, t.MODE_M), 'resolution':64} ] self.set_training_data_generators ([ SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, diff --git a/models/Model_LIAEF128/Model.py b/models/Model_LIAEF128/Model.py index 2fe439a..4b4e6bd 100644 --- a/models/Model_LIAEF128/Model.py +++ b/models/Model_LIAEF128/Model.py @@ -52,7 +52,7 @@ class Model(ModelBase): t = SampleProcessor.Types output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_BGR), 'resolution':128}, { 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_BGR), 'resolution':128}, - { 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_M, t.FACE_MASK_FULL), 'resolution':128} ] + { 'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL, t.MODE_M), 'resolution':128} ] self.set_training_data_generators ([ SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index c00fb16..dfe270d 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -59,7 +59,7 @@ class SAEModel(ModelBase): default_e_ch_dims = 42 default_d_ch_dims = default_e_ch_dims // 2 def_ca_weights = False - + if is_first_run: self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dims (32-1024 ?:help skip:%d) : " % (default_ae_dims) , default_ae_dims, help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 ) self.options['e_ch_dims'] = np.clip ( io.input_int("Encoder dims per channel (21-85 ?:help skip:%d) : " % (default_e_ch_dims) , default_e_ch_dims, help_message="More encoder dims help to recognize more facial features, but require more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 ) @@ -87,16 +87,21 @@ class SAEModel(ModelBase): default_bg_style_power = default_bg_style_power if is_first_run else self.options.get('bg_style_power', default_bg_style_power) self.options['bg_style_power'] = np.clip ( io.input_number("Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_bg_style_power), default_bg_style_power, help_message="Learn to transfer image around face. This can make face more like dst. Enabling this option increases the chance of model collapse."), 0.0, 100.0 ) + + default_apply_random_ct = False if is_first_run else self.options.get('apply_random_ct', False) + self.options['apply_random_ct'] = io.input_bool ("Apply random color transfer to src faceset? (y/n, ?:help skip:%s) : " % (yn_str[default_apply_random_ct]), default_apply_random_ct, help_message="Increase variativity of src samples by apply RCT color transfer from random dst samples.") + else: self.options['pixel_loss'] = self.options.get('pixel_loss', False) self.options['face_style_power'] = self.options.get('face_style_power', default_face_style_power) self.options['bg_style_power'] = self.options.get('bg_style_power', default_bg_style_power) - + self.options['apply_random_ct'] = self.options.get('apply_random_ct', False) + if is_first_run: self.options['pretrain'] = io.input_bool ("Pretrain the model? (y/n, ?:help skip:n) : ", False, help_message="Pretrain the model with large amount of various faces. This technique may help to train the fake with overly different face shapes and light conditions of src/dst data. Face will be look more like a morphed. To reduce the morph effect, some model files will be initialized but not be updated after pretrain: LIAE: inter_AB.h5 DF: encoder.h5. The longer you pretrain the model the more morphed face will look. After that, save and run the training again.") else: - self.options['pretrain'] = False - + self.options['pretrain'] = False + #override def onInitialize(self): exec(nnlib.import_all(), locals(), globals()) @@ -110,13 +115,14 @@ class SAEModel(ModelBase): self.pretrain = self.options['pretrain'] = self.options.get('pretrain', False) if not self.pretrain: self.options.pop('pretrain') - + d_residual_blocks = True bgr_shape = (resolution, resolution, 3) mask_shape = (resolution, resolution, 1) self.ms_count = ms_count = 3 if (self.options['multiscale_decoder']) else 1 + apply_random_ct = self.options.get('apply_random_ct', False) masked_training = True warped_src = Input(bgr_shape) @@ -133,8 +139,8 @@ class SAEModel(ModelBase): target_dstm_ar = [ Input ( ( mask_shape[0] // (2**i) ,)*2 + (mask_shape[-1],) ) for i in range(ms_count-1, -1, -1)] common_flow_kwargs = { 'padding': 'zero', - 'norm': 'norm', - 'act':'' } + 'norm': '', + 'act':'' } models_list = [] weights_to_load = [] if 'liae' in self.options['archi']: @@ -149,11 +155,11 @@ class SAEModel(ModelBase): self.decoder = modelify(SAEModel.LIAEDecFlow (bgr_shape[2],ch_dims=d_ch_dims, multiscale_count=self.ms_count, add_residual_blocks=d_residual_blocks, **common_flow_kwargs)) (inter_output_Inputs) models_list += [self.encoder, self.inter_B, self.inter_AB, self.decoder] - + if self.options['learn_mask']: self.decoderm = modelify(SAEModel.LIAEDecFlow (mask_shape[2],ch_dims=d_ch_dims, **common_flow_kwargs)) (inter_output_Inputs) models_list += [self.decoderm] - + if not self.is_first_run(): weights_to_load += [ [self.encoder , 'encoder.h5'], [self.inter_B , 'inter_B.h5'], @@ -191,12 +197,12 @@ class SAEModel(ModelBase): self.decoder_src = modelify(SAEModel.DFDecFlow (bgr_shape[2],ch_dims=d_ch_dims, multiscale_count=self.ms_count, add_residual_blocks=d_residual_blocks, **common_flow_kwargs )) (dec_Inputs) self.decoder_dst = modelify(SAEModel.DFDecFlow (bgr_shape[2],ch_dims=d_ch_dims, multiscale_count=self.ms_count, add_residual_blocks=d_residual_blocks, **common_flow_kwargs )) (dec_Inputs) models_list += [self.encoder, self.decoder_src, self.decoder_dst] - + if self.options['learn_mask']: self.decoder_srcm = modelify(SAEModel.DFDecFlow (mask_shape[2],ch_dims=d_ch_dims, **common_flow_kwargs )) (dec_Inputs) self.decoder_dstm = modelify(SAEModel.DFDecFlow (mask_shape[2],ch_dims=d_ch_dims, **common_flow_kwargs )) (dec_Inputs) models_list += [self.decoder_srcm, self.decoder_dstm] - + if not self.is_first_run(): weights_to_load += [ [self.encoder , 'encoder.h5'], [self.decoder_src, 'decoder_src.h5'], @@ -217,18 +223,18 @@ class SAEModel(ModelBase): pred_src_srcm = self.decoder_srcm(warped_src_code) pred_dst_dstm = self.decoder_dstm(warped_dst_code) pred_src_dstm = self.decoder_srcm(warped_dst_code) - + if self.is_first_run(): if self.options.get('ca_weights',False): conv_weights_list = [] for model in models_list: for layer in model.layers: if type(layer) == keras.layers.Conv2D: - conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights + conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights CAInitializerMP ( conv_weights_list ) else: self.load_weights_safe(weights_to_load) - + pred_src_src, pred_dst_dst, pred_src_dst, = [ [x] if type(x) != list else x for x in [pred_src_src, pred_dst_dst, pred_src_dst, ] ] if self.options['learn_mask']: @@ -264,7 +270,7 @@ class SAEModel(ModelBase): psd_target_dst_masked_ar = [ pred_src_dst_sigm_ar[i]*target_dstm_sigm_ar[i] for i in range(len(pred_src_dst_sigm_ar))] psd_target_dst_anti_masked_ar = [ pred_src_dst_sigm_ar[i]*target_dstm_anti_sigm_ar[i] for i in range(len(pred_src_dst_sigm_ar))] - + if self.is_training_mode: self.src_dst_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999, tf_cpu_mode=self.options['optimizer_mode']-1) self.src_dst_mask_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999, tf_cpu_mode=self.options['optimizer_mode']-1) @@ -328,7 +334,7 @@ class SAEModel(ModelBase): else: self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1] ] ) - + else: if self.options['learn_mask']: self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1], pred_dst_dstm[-1], pred_src_dstm[-1] ]) @@ -345,29 +351,31 @@ class SAEModel(ModelBase): t_mode_bgr = t.MODE_BGR if not self.pretrain else t.MODE_BGR_SHUFFLE - output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t_mode_bgr), 'resolution':resolution} ] - output_sample_types += [ {'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution // (2**i) } for i in range(ms_count)] - output_sample_types += [ {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M, t.FACE_MASK_FULL), 'resolution': resolution // (2**i) } for i in range(ms_count)] - training_data_src_path = self.training_data_src_path training_data_dst_path = self.training_data_dst_path sort_by_yaw = self.sort_by_yaw - + if self.pretrain and self.pretraining_data_path is not None: training_data_src_path = self.pretraining_data_path training_data_dst_path = self.pretraining_data_path sort_by_yaw = False - + self.set_training_data_generators ([ SampleGeneratorFace(training_data_src_path, sort_by_yaw_target_samples_path=training_data_dst_path if sort_by_yaw else None, - debug=self.is_debug(), batch_size=self.batch_size, + random_ct_samples_path=training_data_dst_path if apply_random_ct else None, + debug=self.is_debug(), batch_size=self.batch_size, sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), - output_sample_types=output_sample_types ), + output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t_mode_bgr), 'resolution':resolution, 'apply_ct': apply_random_ct} ] + \ + [ {'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution // (2**i), 'apply_ct': apply_random_ct } for i in range(ms_count)] + \ + [ {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution // (2**i) } for i in range(ms_count)] + ), SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, ), - output_sample_types=output_sample_types ) - ]) + output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t_mode_bgr), 'resolution':resolution} ] + \ + [ {'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution // (2**i)} for i in range(ms_count)] + \ + [ {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution // (2**i) } for i in range(ms_count)]) + ]) #override def onSave(self): @@ -380,23 +388,23 @@ class SAEModel(ModelBase): [self.inter_B, 'inter_B.h5'], [self.decoder, 'decoder.h5'] ] - + if not self.pretrain or self.iter == 0: ar += [ [self.inter_AB, 'inter_AB.h5'], ] - + if self.options['learn_mask']: ar += [ [self.decoderm, 'decoderm.h5'] ] - + elif 'df' in self.options['archi']: if not self.pretrain or self.iter == 0: ar += [ [self.encoder, 'encoder.h5'], ] - + ar += [ [self.decoder_src, 'decoder_src.h5'], [self.decoder_dst, 'decoder_dst.h5'] - ] - + ] + if self.options['learn_mask']: ar += [ [self.decoder_srcm, 'decoder_srcm.h5'], [self.decoder_dstm, 'decoder_dstm.h5'] ] @@ -442,15 +450,15 @@ class SAEModel(ModelBase): for i in range(0, len(test_S)): ar = S[i], SS[i], D[i], DD[i], SD[i] st.append ( np.concatenate ( ar, axis=1) ) - + result += [ ('SAE', np.concatenate (st, axis=0 )), ] - + if self.options['learn_mask']: st_m = [] for i in range(0, len(test_S)): ar = S[i]*test_S_m[i], SS[i], D[i]*test_D_m[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i]) st_m.append ( np.concatenate ( ar, axis=1) ) - + result += [ ('SAE masked', np.concatenate (st_m, axis=0 )), ] return result @@ -458,7 +466,7 @@ class SAEModel(ModelBase): def predictor_func (self, face): if self.options['learn_mask']: bgr, mask_dst_dstm, mask_src_dstm = self.AE_convert ([face[np.newaxis,...]]) - mask = mask_dst_dstm[0] * mask_src_dstm[0] + mask = mask_dst_dstm[0] * mask_src_dstm[0] return bgr[0], mask[...,0] else: bgr, = self.AE_convert ([face[np.newaxis,...]]) @@ -493,13 +501,13 @@ class SAEModel(ModelBase): def NormPass(x): return x - + def Norm(norm=''): if norm == 'bn': return BatchNormalization(axis=-1) else: return NormPass - + def Act(act='', lrelu_alpha=0.1): if act == 'prelu': return PReLU() @@ -549,7 +557,7 @@ class SAEModel(ModelBase): exec (nnlib.import_all(), locals(), globals()) upscale = partial(SAEModel.upscale, **kwargs) downscale = partial(SAEModel.downscale, **kwargs) - + def func(input): dims = K.int_shape(input)[-1]*ch_dims @@ -571,7 +579,7 @@ class SAEModel(ModelBase): def func(input): x = input[0] - x = Dense(ae_dims)(x) + x = Dense(ae_dims)(x) x = Dense(lowest_dense_res * lowest_dense_res * ae_dims*2)(x) x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims*2))(x) x = upscale(ae_dims*2)(x) @@ -635,8 +643,8 @@ class SAEModel(ModelBase): x = downscale(dims*4)(x) x = downscale(dims*8)(x) - x = Dense(ae_dims)(Flatten()(x)) - x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x) + x = Dense(ae_dims)(Flatten()(x)) + x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x) x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x) x = upscale(ae_dims)(x) return x diff --git a/samplelib/SampleGeneratorFace.py b/samplelib/SampleGeneratorFace.py index a781ff7..5837ea9 100644 --- a/samplelib/SampleGeneratorFace.py +++ b/samplelib/SampleGeneratorFace.py @@ -1,11 +1,14 @@ -import traceback -import numpy as np -import cv2 import multiprocessing -from utils import iter_utils -from facelib import LandmarksProcessor +import traceback + +import cv2 +import numpy as np + +from facelib import LandmarksProcessor +from samplelib import (SampleGeneratorBase, SampleLoader, SampleProcessor, + SampleType) +from utils import iter_utils -from samplelib import SampleType, SampleProcessor, SampleLoader, SampleGeneratorBase ''' arg @@ -15,7 +18,7 @@ output_sample_types = [ ] ''' class SampleGeneratorFace(SampleGeneratorBase): - def __init__ (self, samples_path, debug, batch_size, sort_by_yaw=False, sort_by_yaw_target_samples_path=None, sample_process_options=SampleProcessor.Options(), output_sample_types=[], add_sample_idx=False, generators_count=2, generators_random_seed=None, **kwargs): + def __init__ (self, samples_path, debug, batch_size, sort_by_yaw=False, sort_by_yaw_target_samples_path=None, random_ct_samples_path=None, sample_process_options=SampleProcessor.Options(), output_sample_types=[], add_sample_idx=False, generators_count=2, generators_random_seed=None, **kwargs): super().__init__(samples_path, debug, batch_size) self.sample_process_options = sample_process_options self.output_sample_types = output_sample_types @@ -32,15 +35,17 @@ class SampleGeneratorFace(SampleGeneratorBase): raise ValueError("len(generators_random_seed) != generators_count") self.generators_random_seed = generators_random_seed - + samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path) - + + ct_samples = SampleLoader.load (SampleType.FACE, random_ct_samples_path) if random_ct_samples_path is not None else None + if self.debug: self.generators_count = 1 - self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, (0, samples) )] + self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, (0, samples, ct_samples) )] else: self.generators_count = min ( generators_count, len(samples) ) - self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, (i, samples[i::self.generators_count] ) ) for i in range(self.generators_count) ] + self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, (i, samples[i::self.generators_count], ct_samples ) ) for i in range(self.generators_count) ] self.generator_counter = -1 @@ -53,14 +58,16 @@ class SampleGeneratorFace(SampleGeneratorBase): return next(generator) def batch_func(self, param ): - generator_id, samples = param - + generator_id, samples, ct_samples = param + if self.generators_random_seed is not None: np.random.seed ( self.generators_random_seed[generator_id] ) samples_len = len(samples) samples_idxs = [*range(samples_len)] + ct_samples_len = len(ct_samples) if ct_samples is not None else 0 + if len(samples_idxs) == 0: raise ValueError('No training data provided.') @@ -106,7 +113,8 @@ class SampleGeneratorFace(SampleGeneratorBase): if sample is not None: try: - x = SampleProcessor.process (sample, self.sample_process_options, self.output_sample_types, self.debug) + x = SampleProcessor.process (sample, self.sample_process_options, self.output_sample_types, self.debug, + ct_sample=ct_samples[np.random.randint(ct_samples_len)] if ct_samples is not None else None ) except: raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index d1fd6d7..0ab64d6 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -13,7 +13,7 @@ output_sample_types = [ {} opts, ... ] - + opts: 'types' : (S,S,...,S) where S: @@ -23,31 +23,30 @@ opts: 'IMG_TRANSFORMED' 'IMG_LANDMARKS_ARRAY' #currently unused 'IMG_PITCH_YAW_ROLL' - + 'FACE_TYPE_HALF' 'FACE_TYPE_FULL' 'FACE_TYPE_HEAD' #currently unused 'FACE_TYPE_AVATAR' #currently unused - 'FACE_MASK_FULL' - 'FACE_MASK_EYES' #currently unused - 'MODE_BGR' #BGR 'MODE_G' #Grayscale 'MODE_GGG' #3xGrayscale 'MODE_M' #mask only 'MODE_BGR_SHUFFLE' #BGR shuffle - + 'resolution' : N - + 'motion_blur' : (chance_int, range) - chance 0..100 to apply to face (not mask), and range [1..3] where 3 is highest power of motion blur + 'apply_ct' : bool + """ class SampleProcessor(object): class Types(IntEnum): NONE = 0 - + IMG_TYPE_BEGIN = 1 IMG_SOURCE = 1 IMG_WARPED = 2 @@ -57,19 +56,14 @@ class SampleProcessor(object): IMG_PITCH_YAW_ROLL = 6 IMG_PITCH_YAW_ROLL_SIGMOID = 7 IMG_TYPE_END = 10 - + FACE_TYPE_BEGIN = 10 FACE_TYPE_HALF = 10 FACE_TYPE_FULL = 11 FACE_TYPE_HEAD = 12 #currently unused FACE_TYPE_AVATAR = 13 #currently unused FACE_TYPE_END = 20 - - FACE_MASK_BEGIN = 20 - FACE_MASK_FULL = 20 - FACE_MASK_EYES = 21 #currently unused - FACE_MASK_END = 30 - + MODE_BEGIN = 40 MODE_BGR = 40 #BGR MODE_G = 41 #Grayscale @@ -77,7 +71,7 @@ class SampleProcessor(object): MODE_M = 43 #mask only MODE_BGR_SHUFFLE = 44 #BGR shuffle MODE_END = 50 - + class Options(object): def __init__(self, random_flip = True, 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] ): @@ -89,10 +83,12 @@ class SampleProcessor(object): self.ty_range = ty_range @staticmethod - def process (sample, sample_process_options, output_sample_types, debug): + def process (sample, sample_process_options, output_sample_types, debug, ct_sample=None): SPTF = SampleProcessor.Types sample_bgr = sample.load_bgr() + ct_sample_bgr = None + ct_sample_mask = None h,w,c = sample_bgr.shape is_face_sample = sample.landmarks is not None @@ -103,25 +99,26 @@ class SampleProcessor(object): params = imagelib.gen_warp_params(sample_bgr, sample_process_options.random_flip, rotation_range=sample_process_options.rotation_range, scale_range=sample_process_options.scale_range, tx_range=sample_process_options.tx_range, ty_range=sample_process_options.ty_range ) cached_images = collections.defaultdict(dict) - + sample_rnd_seed = np.random.randint(0x80000000) SPTF_FACETYPE_TO_FACETYPE = { SPTF.FACE_TYPE_HALF : FaceType.HALF, SPTF.FACE_TYPE_FULL : FaceType.FULL, SPTF.FACE_TYPE_HEAD : FaceType.HEAD, SPTF.FACE_TYPE_AVATAR : FaceType.AVATAR } - + outputs = [] for opts in output_sample_types: - + resolution = opts.get('resolution', 0) types = opts.get('types', [] ) - + random_sub_res = opts.get('random_sub_res', 0) normalize_std_dev = opts.get('normalize_std_dev', False) normalize_vgg = opts.get('normalize_vgg', False) motion_blur = opts.get('motion_blur', None) - + apply_ct = opts.get('apply_ct', False) + img_type = SPTF.NONE target_face_type = SPTF.NONE face_mask_type = SPTF.NONE @@ -131,11 +128,9 @@ class SampleProcessor(object): img_type = t elif t >= SPTF.FACE_TYPE_BEGIN and t < SPTF.FACE_TYPE_END: target_face_type = t - elif t >= SPTF.FACE_MASK_BEGIN and t < SPTF.FACE_MASK_END: - face_mask_type = t elif t >= SPTF.MODE_BEGIN and t < SPTF.MODE_END: mode_type = t - + if img_type == SPTF.NONE: raise ValueError ('expected IMG_ type') @@ -152,55 +147,52 @@ class SampleProcessor(object): pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll (sample.landmarks) if params['flip']: yaw = -yaw - + if img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: pitch = (pitch+1.0) / 2.0 yaw = (yaw+1.0) / 2.0 roll = (roll+1.0) / 2.0 - - img = (pitch, yaw, roll) + + img = (pitch, yaw, roll) else: if mode_type == SPTF.NONE: raise ValueError ('expected MODE_ type') - - img = cached_images.get(img_type, {}).get(face_mask_type, None) + + img = cached_images.get(img_type, None) if img is None: - + img = sample_bgr + mask = None cur_sample = sample if is_face_sample: if motion_blur is not None: chance, mb_range = motion_blur - chance = np.clip(chance, 0, 100) - + chance = np.clip(chance, 0, 100) + if np.random.randint(100) < chance: mb_range = [3,5,7,9][ : np.clip(mb_range, 0, 3)+1 ] dim = mb_range[ np.random.randint(len(mb_range) ) ] img = imagelib.LinearMotionBlur (img, dim, np.random.randint(180) ) - if face_mask_type == SPTF.FACE_MASK_FULL: - mask = cur_sample.load_fanseg_mask() #using fanseg_mask if exist + mask = cur_sample.load_fanseg_mask() #using fanseg_mask if exist - if mask is None: - mask = LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks) + if mask is None: + mask = LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks) - if cur_sample.ie_polys is not None: - cur_sample.ie_polys.overlay_mask(mask) + if cur_sample.ie_polys is not None: + cur_sample.ie_polys.overlay_mask(mask) - img = np.concatenate( (img, mask ), -1 ) - elif face_mask_type == SPTF.FACE_MASK_EYES: - mask = LandmarksProcessor.get_image_eye_mask (img.shape, cur_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 ) - warp = (img_type==SPTF.IMG_WARPED or img_type==SPTF.IMG_WARPED_TRANSFORMED) transform = (img_type==SPTF.IMG_WARPED_TRANSFORMED or img_type==SPTF.IMG_TRANSFORMED) flip = img_type != SPTF.IMG_WARPED - is_border_replicate = face_mask_type == SPTF.NONE - - img = cached_images[img_type][face_mask_type] = imagelib.warp_by_params (params, img, warp, transform, flip, is_border_replicate) + + img = imagelib.warp_by_params (params, img, warp, transform, flip, True) + if mask is not None: + mask = imagelib.warp_by_params (params, mask, warp, transform, flip, False)[...,np.newaxis] + img = np.concatenate( (img, mask ), -1 ) + + cached_images[img_type] = img if is_face_sample and target_face_type != SPTF.NONE: ft = SPTF_FACETYPE_TO_FACETYPE[target_face_type] @@ -217,9 +209,24 @@ class SampleProcessor(object): start_y = rnd_state.randint(sub_size+1) img = img[start_y:start_y+sub_size,start_x:start_x+sub_size,:] + img = np.clip(img, 0, 1) img_bgr = img[...,0:3] img_mask = img[...,3:4] - + + if apply_ct: + if ct_sample_bgr is None: + ct_sample_bgr = ct_sample.load_bgr() + ct_sample_mask = LandmarksProcessor.get_image_hull_mask (ct_sample_bgr.shape, ct_sample.landmarks) + + ct_sample_bgr_resized = cv2.resize( ct_sample_bgr, (resolution,resolution), cv2.INTER_LINEAR ) + ct_sample_mask_resized = cv2.resize( ct_sample_mask, (resolution,resolution), cv2.INTER_LINEAR )[...,np.newaxis] + + img_bgr = imagelib.reinhard_color_transfer ( np.clip( (img_bgr*255) .astype(np.uint8), 0, 255), + np.clip( (ct_sample_bgr_resized*255).astype(np.uint8), 0, 255), + source_mask=img_mask, target_mask=ct_sample_mask_resized) + + img_bgr = np.clip( img_bgr.astype(np.float32) / 255.0, 0.0, 1.0) + if normalize_std_dev: img_bgr = (img_bgr - img_bgr.mean( (0,1)) ) / img_bgr.std( (0,1) ) elif normalize_vgg: @@ -227,7 +234,7 @@ class SampleProcessor(object): img_bgr[:,:,0] -= 103.939 img_bgr[:,:,1] -= 116.779 img_bgr[:,:,2] -= 123.68 - + if mode_type == SPTF.MODE_BGR: img = img_bgr elif mode_type == SPTF.MODE_BGR_SHUFFLE: @@ -239,8 +246,6 @@ class SampleProcessor(object): elif mode_type == SPTF.MODE_GGG: img = np.concatenate ( ( np.repeat ( np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1), (3,), -1), img_mask), -1) elif mode_type == SPTF.MODE_M and is_face_sample: - if face_mask_type == SPTF.NONE: - raise ValueError ('no face_mask_type defined') img = img_mask if not debug: