diff --git a/converters/ConvertAvatar.py b/converters/ConvertAvatar.py index 08f8394..e5264a0 100644 --- a/converters/ConvertAvatar.py +++ b/converters/ConvertAvatar.py @@ -14,8 +14,8 @@ def process_frame_info(frame_info, inp_sh): img = cv2.warpAffine( img, img_mat, inp_sh[0:2], borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC ) return img -def ConvertFaceAvatar (cfg, prev_temporal_frame_infos, frame_info, next_temporal_frame_infos): - inp_sh = cfg.predictor_input_shape +def ConvertFaceAvatar (predictor_func, predictor_input_shape, cfg, prev_temporal_frame_infos, frame_info, next_temporal_frame_infos): + inp_sh = predictor_input_shape prev_imgs=[] next_imgs=[] @@ -24,7 +24,7 @@ def ConvertFaceAvatar (cfg, prev_temporal_frame_infos, frame_info, next_temporal next_imgs.append( process_frame_info(next_temporal_frame_infos[i], inp_sh) ) img = process_frame_info(frame_info, inp_sh) - prd_f = cfg.predictor_func ( prev_imgs, img, next_imgs ) + prd_f = predictor_func ( prev_imgs, img, next_imgs ) if cfg.super_resolution_mode != 0: prd_f = cfg.superres_func(cfg.super_resolution_mode, prd_f) diff --git a/converters/ConvertMasked.py b/converters/ConvertMasked.py index 1354875..0e9d8cb 100644 --- a/converters/ConvertMasked.py +++ b/converters/ConvertMasked.py @@ -9,7 +9,7 @@ from interact import interact as io from utils.cv2_utils import * -def ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmarks): +def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmarks): #if debug: # debugs = [img_bgr.copy()] @@ -26,7 +26,7 @@ def ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmar out_img = img_bgr.copy() out_merging_mask = None - output_size = cfg.predictor_input_shape[0] + output_size = predictor_input_shape[0] if cfg.super_resolution_mode != 0: output_size *= 2 @@ -36,17 +36,19 @@ def ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmar dst_face_bgr = cv2.warpAffine( img_bgr , face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC ) dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC ) - predictor_input_bgr = cv2.resize (dst_face_bgr, cfg.predictor_input_shape[0:2] ) + predictor_input_bgr = cv2.resize (dst_face_bgr, predictor_input_shape[0:2] ) - if cfg.predictor_masked: - prd_face_bgr, prd_face_mask_a_0 = cfg.predictor_func (predictor_input_bgr) - - prd_face_bgr = np.clip (prd_face_bgr, 0, 1.0 ) - prd_face_mask_a_0 = np.clip (prd_face_mask_a_0, 0.0, 1.0) + predicted = predictor_func (predictor_input_bgr) + if isinstance(predicted, tuple): + #converter return bgr,mask + prd_face_bgr = np.clip (predicted[0], 0, 1.0) + prd_face_mask_a_0 = np.clip (predicted[1], 0, 1.0) + predictor_masked = True else: - predicted = cfg.predictor_func (predictor_input_bgr) + #converter return bgr only, using dst mask prd_face_bgr = np.clip (predicted, 0, 1.0 ) - prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, cfg.predictor_input_shape[0:2] ) + prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, predictor_input_shape[0:2] ) + predictor_masked = False if cfg.super_resolution_mode: #if debug: @@ -57,7 +59,7 @@ def ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmar #if debug: # debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - if cfg.predictor_masked: + if predictor_masked: prd_face_mask_a_0 = cv2.resize (prd_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) else: prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) @@ -198,7 +200,7 @@ def ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmar # debugs += [img_face_mask_aaa.copy()] if 'seamless' not in cfg.mode and cfg.color_transfer_mode != 0: - if cfg.color_transfer_mode == 1: + if cfg.color_transfer_mode == 1: #rct #if debug: # debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] @@ -211,8 +213,8 @@ def ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmar # debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - elif cfg.color_transfer_mode == 2: - #if debug: + elif cfg.color_transfer_mode == 2: #lct + #if debug: # debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] prd_face_bgr = imagelib.linear_color_transfer (prd_face_bgr, dst_face_bgr) @@ -220,7 +222,14 @@ def ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmar #if debug: # debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + elif cfg.color_transfer_mode == 3: #ebs + #if debug: + # debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + prd_face_bgr = cfg.ebs_ct_func ( np.clip( (dst_face_bgr*255), 0, 255).astype(np.uint8), + np.clip( (prd_face_bgr*255), 0, 255).astype(np.uint8), )#prd_face_mask_a + prd_face_bgr = np.clip( prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) + if cfg.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 ) @@ -249,34 +258,39 @@ def ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmar if cfg.mode == 'hist-match-bw': prd_face_bgr = prd_face_bgr.astype(dtype=np.float32) - - out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) - out_img = np.clip(out_img, 0.0, 1.0) - - #if debug: - # debugs += [out_img.copy()] - - if cfg.mode == 'overlay': - pass - + if 'seamless' in cfg.mode: - #mask used for cv2.seamlessClone - img_face_seamless_mask_a = None + #mask used for cv2.seamlessClone img_face_mask_a = img_face_mask_aaa[...,0:1] + + if cfg.mode == 'seamless2': + img_face_mask_a = cv2.warpAffine( img_face_mask_a, face_output_mat, (output_size, output_size), flags=cv2.INTER_CUBIC ) + + img_face_seamless_mask_a = None for i in range(1,10): a = img_face_mask_a > i / 10.0 if len(np.argwhere(a)) == 0: continue - img_face_seamless_mask_a = img_face_mask_aaa[...,0:1].copy() + img_face_seamless_mask_a = img_face_mask_a.copy() img_face_seamless_mask_a[a] = 1.0 img_face_seamless_mask_a[img_face_seamless_mask_a <= i / 10.0] = 0.0 break + + if cfg.mode == 'seamless2': + + face_seamless = imagelib.seamless_clone ( prd_face_bgr, dst_face_bgr, img_face_seamless_mask_a ) + + out_img = cv2.warpAffine( face_seamless, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) + else: + out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) + + out_img = np.clip(out_img, 0.0, 1.0) + if 'seamless' in cfg.mode and cfg.mode != 'seamless2': try: #calc same bounding rect and center point as in cv2.seamlessClone to prevent jittering (not flickering) l,t,w,h = cv2.boundingRect( (img_face_seamless_mask_a*255).astype(np.uint8) ) s_maskx, s_masky = int(l+w/2), int(t+h/2) - out_img = cv2.seamlessClone( (out_img*255).astype(np.uint8), img_bgr_uint8, (img_face_seamless_mask_a*255).astype(np.uint8), (s_maskx,s_masky) , cv2.NORMAL_CLONE ) out_img = out_img.astype(dtype=np.float32) / 255.0 except Exception as e: @@ -301,8 +315,8 @@ def ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmar # debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] face_mask_aaa = cv2.warpAffine( img_face_mask_aaa, face_mat, (output_size, output_size) ) - out_face_bgr = imagelib.reinhard_color_transfer ( np.clip( (out_face_bgr*255).astype(np.uint8), 0, 255), - np.clip( (dst_face_bgr*255).astype(np.uint8), 0, 255), + out_face_bgr = imagelib.reinhard_color_transfer ( np.clip( (out_face_bgr*255), 0, 255).astype(np.uint8), + np.clip( (dst_face_bgr*255), 0, 255).astype(np.uint8), source_mask=face_mask_aaa, target_mask=face_mask_aaa) out_face_bgr = np.clip( out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) @@ -318,7 +332,15 @@ def ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmar #if debug: # debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - + + elif cfg.color_transfer_mode == 3: #ebs + #if debug: + # debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + out_face_bgr = cfg.ebs_ct_func ( np.clip( (dst_face_bgr*255), 0, 255).astype(np.uint8), + np.clip( (out_face_bgr*255), 0, 255).astype(np.uint8), ) + out_face_bgr = np.clip( out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) + if cfg.mode == 'seamless-hist-match': out_face_bgr = imagelib.color_hist_match(out_face_bgr, dst_face_bgr, cfg.hist_match_threshold) @@ -359,14 +381,14 @@ def ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmar return out_img, out_merging_mask -def ConvertMasked (cfg, frame_info): +def ConvertMasked (predictor_func, predictor_input_shape, cfg, frame_info): img_bgr_uint8 = cv2_imread(frame_info.filename) img_bgr_uint8 = imagelib.normalize_channels (img_bgr_uint8, 3) img_bgr = img_bgr_uint8.astype(np.float32) / 255.0 outs = [] for face_num, img_landmarks in enumerate( frame_info.landmarks_list ): - out_img, out_img_merging_mask = ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_landmarks) + out_img, out_img_merging_mask = ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, img_bgr_uint8, img_bgr, img_landmarks) outs += [ (out_img, out_img_merging_mask) ] #Combining multiple face outputs diff --git a/converters/ConverterConfig.py b/converters/ConverterConfig.py index e309238..5832228 100644 --- a/converters/ConverterConfig.py +++ b/converters/ConverterConfig.py @@ -14,16 +14,14 @@ class ConverterConfig(object): TYPE_IMAGE = 3 TYPE_IMAGE_WITH_LANDMARKS = 4 - def __init__(self, type=0, predictor_func=None, - predictor_input_shape=None): + def __init__(self, type=0): self.type = type - self.predictor_func = predictor_func - self.predictor_input_shape = predictor_input_shape self.superres_func = None self.sharpen_func = None self.fanseg_input_size = None self.fanseg_extract_func = None + self.ebs_ct_func = None self.super_res_dict = {0:"None", 1:'RankSRGAN'} self.sharpen_dict = {0:"None", 1:'box', 2:'gaussian'} @@ -84,40 +82,46 @@ class ConverterConfig(object): r += f"super_resolution_mode : {self.super_res_dict[self.super_resolution_mode]}\n" return r +mode_dict = {0:'original', + 1:'overlay', + 2:'hist-match', + 3:'seamless2', + 4:'seamless', + 5:'seamless-hist-match', + 6:'raw-rgb', + 7:'raw-rgb-mask', + 8:'raw-mask-only', + 9:'raw-predicted-only'} + +full_face_mask_mode_dict = {1:'learned', + 2:'dst', + 3:'FAN-prd', + 4:'FAN-dst', + 5:'FAN-prd*FAN-dst', + 6:'learned*FAN-prd*FAN-dst'} + +half_face_mask_mode_dict = {1:'learned', + 2:'dst', + 4:'FAN-dst', + 7:'learned*FAN-dst'} + +ctm_dict = { 0: "None", 1:"rct", 2:"lct", 3:"ebs" } +ctm_str_dict = {None:0, "rct":1, "lct": 2, "ebs":3 } + class ConverterConfigMasked(ConverterConfig): - def __init__(self, predictor_func=None, - predictor_input_shape=None, - predictor_masked=True, - face_type=FaceType.FULL, + def __init__(self, face_type=FaceType.FULL, default_mode = 4, - base_erode_mask_modifier = 0, - base_blur_mask_modifier = 0, - default_erode_mask_modifier = 0, - default_blur_mask_modifier = 0, clip_hborder_mask_per = 0, ): - super().__init__(type=ConverterConfig.TYPE_MASKED, - predictor_func=predictor_func, - predictor_input_shape=predictor_input_shape, - ) - if len(predictor_input_shape) != 3: - raise ValueError("ConverterConfigMasked: predictor_input_shape must be rank 3.") - - if predictor_input_shape[0] != predictor_input_shape[1]: - raise ValueError("ConverterConfigMasked: predictor_input_shape must be a square.") - - self.predictor_masked = predictor_masked + super().__init__(type=ConverterConfig.TYPE_MASKED) + self.face_type = face_type if self.face_type not in [FaceType.FULL, FaceType.HALF]: raise ValueError("ConverterConfigMasked supports only full or half face masks.") self.default_mode = default_mode - self.base_erode_mask_modifier = base_erode_mask_modifier - self.base_blur_mask_modifier = base_blur_mask_modifier - self.default_erode_mask_modifier = default_erode_mask_modifier - self.default_blur_mask_modifier = default_blur_mask_modifier self.clip_hborder_mask_per = clip_hborder_mask_per #default changeable params @@ -133,37 +137,11 @@ class ConverterConfigMasked(ConverterConfig): self.color_degrade_power = 0 self.export_mask_alpha = False - self.mode_dict = {0:'original', - 1:'overlay', - 2:'hist-match', - 3:'hist-match-bw', - 4:'seamless', - 5:'seamless-hist-match', - 6:'raw-rgb', - 7:'raw-rgb-mask', - 8:'raw-mask-only', - 9:'raw-predicted-only'} - - self.full_face_mask_mode_dict = {1:'learned', - 2:'dst', - 3:'FAN-prd', - 4:'FAN-dst', - 5:'FAN-prd*FAN-dst', - 6:'learned*FAN-prd*FAN-dst'} - - self.half_face_mask_mode_dict = {1:'learned', - 2:'dst', - 4:'FAN-dst', - 7:'learned*FAN-dst'} - - self.ctm_dict = { 0: "None", 1:"rct", 2:"lct" } - self.ctm_str_dict = {None:0, "rct":1, "lct": 2 } - def copy(self): return copy.copy(self) def set_mode (self, mode): - self.mode = self.mode_dict.get (mode, self.mode_dict[self.default_mode] ) + self.mode = mode_dict.get (mode, mode_dict[self.default_mode] ) def toggle_masked_hist_match(self): if self.mode == 'hist-match' or self.mode == 'hist-match-bw': @@ -175,16 +153,16 @@ class ConverterConfigMasked(ConverterConfig): def toggle_mask_mode(self): if self.face_type == FaceType.FULL: - a = list( self.full_face_mask_mode_dict.keys() ) + a = list( full_face_mask_mode_dict.keys() ) else: - a = list( self.half_face_mask_mode_dict.keys() ) + a = list( half_face_mask_mode_dict.keys() ) self.mask_mode = a[ (a.index(self.mask_mode)+1) % len(a) ] def add_erode_mask_modifier(self, diff): - self.erode_mask_modifier = np.clip ( self.erode_mask_modifier+diff , -200, 200) + self.erode_mask_modifier = np.clip ( self.erode_mask_modifier+diff , -400, 400) def add_blur_mask_modifier(self, diff): - self.blur_mask_modifier = np.clip ( self.blur_mask_modifier+diff , -200, 200) + self.blur_mask_modifier = np.clip ( self.blur_mask_modifier+diff , -400, 400) def add_motion_blur_power(self, diff): self.motion_blur_power = np.clip ( self.motion_blur_power+diff, 0, 100) @@ -193,7 +171,7 @@ class ConverterConfigMasked(ConverterConfig): self.output_face_scale = np.clip ( self.output_face_scale+diff , -50, 50) def toggle_color_transfer_mode(self): - self.color_transfer_mode = (self.color_transfer_mode+1) % 3 + self.color_transfer_mode = (self.color_transfer_mode+1) % ( max(ctm_dict.keys())+1 ) def add_color_degrade_power(self, diff): self.color_degrade_power = np.clip ( self.color_degrade_power+diff , 0, 100) @@ -204,13 +182,13 @@ class ConverterConfigMasked(ConverterConfig): def ask_settings(self): s = """Choose mode: \n""" - for key in self.mode_dict.keys(): - s += f"""({key}) {self.mode_dict[key]}\n""" + for key in mode_dict.keys(): + s += f"""({key}) {mode_dict[key]}\n""" s += f"""Default: {self.default_mode} : """ mode = io.input_int (s, self.default_mode) - self.mode = self.mode_dict.get (mode, self.mode_dict[self.default_mode] ) + self.mode = mode_dict.get (mode, mode_dict[self.default_mode] ) if 'raw' not in self.mode: if self.mode == 'hist-match' or self.mode == 'hist-match-bw': @@ -221,28 +199,28 @@ class ConverterConfigMasked(ConverterConfig): if self.face_type == FaceType.FULL: s = """Choose mask mode: \n""" - for key in self.full_face_mask_mode_dict.keys(): - s += f"""({key}) {self.full_face_mask_mode_dict[key]}\n""" + for key in full_face_mask_mode_dict.keys(): + s += f"""({key}) {full_face_mask_mode_dict[key]}\n""" s += f"""?:help Default: 1 : """ - self.mask_mode = io.input_int (s, 1, valid_list=self.full_face_mask_mode_dict.keys(), help_message="If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks.") + self.mask_mode = io.input_int (s, 1, valid_list=full_face_mask_mode_dict.keys(), help_message="If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks.") else: s = """Choose mask mode: \n""" - for key in self.half_face_mask_mode_dict.keys(): - s += f"""({key}) {self.half_face_mask_mode_dict[key]}\n""" + for key in half_face_mask_mode_dict.keys(): + s += f"""({key}) {half_face_mask_mode_dict[key]}\n""" s += f"""?:help , Default: 1 : """ - self.mask_mode = io.input_int (s, 1, valid_list=self.half_face_mask_mode_dict.keys(), help_message="If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images.") + self.mask_mode = io.input_int (s, 1, valid_list=half_face_mask_mode_dict.keys(), help_message="If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images.") if 'raw' not in self.mode: - self.erode_mask_modifier = self.base_erode_mask_modifier + np.clip ( io.input_int ("Choose erode mask modifier [-200..200] (skip:%d) : " % (self.default_erode_mask_modifier), self.default_erode_mask_modifier), -200, 200) - self.blur_mask_modifier = self.base_blur_mask_modifier + np.clip ( io.input_int ("Choose blur mask modifier [-200..200] (skip:%d) : " % (self.default_blur_mask_modifier), self.default_blur_mask_modifier), -200, 200) + self.erode_mask_modifier = np.clip ( io.input_int ("Choose erode mask modifier [-400..400] (skip:%d) : " % 0, 0), -400, 400) + self.blur_mask_modifier = np.clip ( io.input_int ("Choose blur mask modifier [-400..400] (skip:%d) : " % 0, 0), -400, 400) self.motion_blur_power = np.clip ( io.input_int ("Choose motion blur power [0..100] (skip:%d) : " % (0), 0), 0, 100) self.output_face_scale = np.clip (io.input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0), -50, 50) if 'raw' not in self.mode: - self.color_transfer_mode = io.input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ", None, ['rct','lct']) - self.color_transfer_mode = self.ctm_str_dict[self.color_transfer_mode] + self.color_transfer_mode = io.input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct/ebs skip:None ) : ", None, ctm_str_dict.keys() ) + self.color_transfer_mode = ctm_str_dict[self.color_transfer_mode] super().ask_settings() @@ -284,9 +262,9 @@ class ConverterConfigMasked(ConverterConfig): r += f"""hist_match_threshold: {self.hist_match_threshold}\n""" if self.face_type == FaceType.FULL: - r += f"""mask_mode: { self.full_face_mask_mode_dict[self.mask_mode] }\n""" + r += f"""mask_mode: { full_face_mask_mode_dict[self.mask_mode] }\n""" else: - r += f"""mask_mode: { self.half_face_mask_mode_dict[self.mask_mode] }\n""" + r += f"""mask_mode: { half_face_mask_mode_dict[self.mask_mode] }\n""" if 'raw' not in self.mode: r += (f"""erode_mask_modifier: {self.erode_mask_modifier}\n""" @@ -296,7 +274,7 @@ class ConverterConfigMasked(ConverterConfig): r += f"""output_face_scale: {self.output_face_scale}\n""" if 'raw' not in self.mode: - r += f"""color_transfer_mode: { self.ctm_dict[self.color_transfer_mode]}\n""" + r += f"""color_transfer_mode: { ctm_dict[self.color_transfer_mode]}\n""" r += super().__str__() @@ -311,14 +289,8 @@ class ConverterConfigMasked(ConverterConfig): class ConverterConfigFaceAvatar(ConverterConfig): - def __init__(self, predictor_func=None, - predictor_input_shape=None, - temporal_face_count=0 - ): - super().__init__(type=ConverterConfig.TYPE_FACE_AVATAR, - predictor_func=predictor_func, - predictor_input_shape=predictor_input_shape - ) + def __init__(self, temporal_face_count=0): + super().__init__(type=ConverterConfig.TYPE_FACE_AVATAR) self.temporal_face_count = temporal_face_count #changeable params diff --git a/ebsynth/__init__.py b/ebsynth/__init__.py new file mode 100644 index 0000000..ce31c32 --- /dev/null +++ b/ebsynth/__init__.py @@ -0,0 +1 @@ +from .ebsynth import color_transfer \ No newline at end of file diff --git a/ebsynth/ebsynth.dll b/ebsynth/ebsynth.dll new file mode 100644 index 0000000..f6c085e Binary files /dev/null and b/ebsynth/ebsynth.dll differ diff --git a/ebsynth/ebsynth.py b/ebsynth/ebsynth.py new file mode 100644 index 0000000..ec2ec08 --- /dev/null +++ b/ebsynth/ebsynth.py @@ -0,0 +1,201 @@ +import os +import sys +from ctypes import * +from pathlib import Path + +import cv2 +import numpy as np + +libebsynth = None +cached_buffer = {} + +EBSYNTH_BACKEND_CPU = 0x0001 +EBSYNTH_BACKEND_CUDA = 0x0002 +EBSYNTH_BACKEND_AUTO = 0x0000 +EBSYNTH_MAX_STYLE_CHANNELS = 8 +EBSYNTH_MAX_GUIDE_CHANNELS = 24 +EBSYNTH_VOTEMODE_PLAIN = 0x0001 # weight = 1 +EBSYNTH_VOTEMODE_WEIGHTED = 0x0002 # weight = 1/(1+error) + + +def _normalize_img_shape (img): + img_len = len(img.shape) + if img_len == 2: + sh, sw = img.shape + sc = 0 + elif img_len == 3: + sh, sw, sc = img.shape + + if sc == 0: + sc = 1 + img = img [...,np.newaxis] + return img + +def run (img_style, guides, + patch_size=5, + num_pyramid_levels=-1, + num_search_vote_iters = 6, + num_patch_match_iters = 4, + stop_threshold = 5, + uniformity_weight = 3500.0, + extraPass3x3 = False, + ): + if patch_size < 3: + raise ValueError ("patch_size is too small") + if patch_size % 2 == 0: + raise ValueError ("patch_size must be an odd number") + if len(guides) == 0: + raise ValueError ("at least one guide must be specified") + + global libebsynth + if libebsynth is None: + if sys.platform[0:3] == 'win': + libebsynth_path = str ( Path(__file__).parent / 'ebsynth.dll' ) + libebsynth = CDLL(libebsynth_path) + else: + #todo: implement for linux + pass + + if libebsynth is not None: + libebsynth.ebsynthRun.argtypes = ( \ + c_int, + c_int, + c_int, + c_int, + c_int, + c_void_p, + c_void_p, + c_int, + c_int, + c_void_p, + c_void_p, + POINTER(c_float), + POINTER(c_float), + c_float, + c_int, + c_int, + c_int, + POINTER(c_int), + POINTER(c_int), + POINTER(c_int), + c_int, + c_void_p, + c_void_p + ) + + if libebsynth is None: + return img_style + + img_style = _normalize_img_shape (img_style) + sh, sw, sc = img_style.shape + t_h, t_w, t_c = 0,0,0 + + if sc > EBSYNTH_MAX_STYLE_CHANNELS: + raise ValueError (f"error: too many style channels {sc}, maximum number is {EBSYNTH_MAX_STYLE_CHANNELS}") + + guides_source = [] + guides_target = [] + guides_weights = [] + + for i in range(len(guides)): + source_guide, target_guide, guide_weight = guides[i] + source_guide = _normalize_img_shape(source_guide) + target_guide = _normalize_img_shape(target_guide) + s_h, s_w, s_c = source_guide.shape + nt_h, nt_w, nt_c = target_guide.shape + + if s_h != sh or s_w != sw: + raise ValueError ("guide source and style resolution must match style resolution.") + + if t_c == 0: + t_h, t_w, t_c = nt_h, nt_w, nt_c + elif nt_h != t_h or nt_w != t_w: + raise ValueError ("guides target resolutions must be equal") + + if s_c != nt_c: + raise ValueError ("guide source and target channels must match exactly.") + + guides_source.append (source_guide) + guides_target.append (target_guide) + + guides_weights += [ guide_weight / s_c ] * s_c + + guides_source = np.concatenate ( guides_source, axis=-1) + guides_target = np.concatenate ( guides_target, axis=-1) + guides_weights = (c_float*len(guides_weights) ) ( *guides_weights ) + + styleWeight = 1.0 + style_weights = [ styleWeight / sc for i in range(sc) ] + style_weights = (c_float*sc) ( *style_weights ) + + + maxPyramidLevels = 0 + for level in range(32,-1,-1): + if min( min(sh, t_h)*pow(2.0, -level), \ + min(sw, t_w)*pow(2.0, -level) ) >= (2*patch_size+1): + maxPyramidLevels = level+1 + break + + if num_pyramid_levels == -1: + num_pyramid_levels = maxPyramidLevels + num_pyramid_levels = min(num_pyramid_levels, maxPyramidLevels) + + num_search_vote_iters_per_level = (c_int*num_pyramid_levels) ( *[num_search_vote_iters]*num_pyramid_levels ) + num_patch_match_iters_per_level = (c_int*num_pyramid_levels) ( *[num_patch_match_iters]*num_pyramid_levels ) + stop_threshold_per_level = (c_int*num_pyramid_levels) ( *[stop_threshold]*num_pyramid_levels ) + + buffer = cached_buffer.get ( (t_h,t_w,sc), None ) + if buffer is None: + buffer = create_string_buffer (t_h*t_w*sc) + cached_buffer[(t_h,t_w,sc)] = buffer + + libebsynth.ebsynthRun (EBSYNTH_BACKEND_CPU, #backend + sc, #numStyleChannels + guides_source.shape[-1], #numGuideChannels + sw, #sourceWidth + sh, #sourceHeight + img_style.tobytes(), #sourceStyleData (width * height * numStyleChannels) bytes, scan-line order + guides_source.tobytes(), #sourceGuideData (width * height * numGuideChannels) bytes, scan-line order + t_w, #targetWidth + t_h, #targetHeight + guides_target.tobytes(), #targetGuideData (width * height * numGuideChannels) bytes, scan-line order + None, #targetModulationData (width * height * numGuideChannels) bytes, scan-line order; pass NULL to switch off the modulation + style_weights, #styleWeights (numStyleChannels) floats + guides_weights, #guideWeights (numGuideChannels) floats + uniformity_weight, #uniformityWeight reasonable values are between 500-15000, 3500 is a good default + patch_size, #patchSize odd sizes only, use 5 for 5x5 patch, 7 for 7x7, etc. + EBSYNTH_VOTEMODE_PLAIN, #voteMode use VOTEMODE_WEIGHTED for sharper result + num_pyramid_levels, #numPyramidLevels + + num_search_vote_iters_per_level, #numSearchVoteItersPerLevel how many search/vote iters to perform at each level (array of ints, coarse first, fine last) + num_patch_match_iters_per_level, #numPatchMatchItersPerLevel how many Patch-Match iters to perform at each level (array of ints, coarse first, fine last) + stop_threshold_per_level, #stopThresholdPerLevel stop improving pixel when its change since last iteration falls under this threshold + 1 if extraPass3x3 else 0, #extraPass3x3 perform additional polishing pass with 3x3 patches at the finest level, use 0 to disable + None, #outputNnfData (width * height * 2) ints, scan-line order; pass NULL to ignore + buffer #outputImageData (width * height * numStyleChannels) bytes, scan-line order + ) + + return np.frombuffer(buffer, dtype=np.uint8).reshape ( (t_h,t_w,sc) ).copy() + +#transfer color from source to target +def color_transfer(img_source, img_target): + guides = [( cv2.cvtColor(img_source, cv2.COLOR_BGR2GRAY), + cv2.cvtColor(img_target, cv2.COLOR_BGR2GRAY), + 1 ) ] + + h,w,c = img_source.shape + result = [] + for i in range(c): + result += [ + run( img_source[...,i:i+1] , guides=guides, + patch_size=11, + num_pyramid_levels=40, + num_search_vote_iters = 6, + num_patch_match_iters = 4, + stop_threshold = 5, + uniformity_weight=500.0, + extraPass3x3=True, + ) + ] + + return np.concatenate( result, axis=-1 ) diff --git a/facelib/FANSeg_256_full_face.h5 b/facelib/FANSeg_256_full_face.h5 index a9b2e59..6886504 100644 Binary files a/facelib/FANSeg_256_full_face.h5 and b/facelib/FANSeg_256_full_face.h5 differ diff --git a/facelib/FANSegmentator.py b/facelib/FANSegmentator.py index e2ae705..1bdc2f3 100644 --- a/facelib/FANSegmentator.py +++ b/facelib/FANSegmentator.py @@ -47,10 +47,30 @@ class FANSegmentator(object): self.model.get_layer (s).set_weights ( d[s] ) except: io.log_err("Unable to load VGG11 pretrained weights from vgg11_enc_weights.npy") + + conv_weights_list = [] + for layer in self.model.layers: + if 'CA.' in layer.name: + conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights + CAInitializerMP ( conv_weights_list ) if training: - #self.model.compile(loss='mse', optimizer=Adam(tf_cpu_mode=2)) - self.model.compile(loss='binary_crossentropy', optimizer=Adam(tf_cpu_mode=2) ) + inp_t = Input ( (resolution, resolution, 3) ) + real_t = Input ( (resolution, resolution, 1) ) + out_t = self.model(inp_t) + + #loss = K.mean(10*K.square(out_t-real_t)) + loss = K.mean(10*K.binary_crossentropy(real_t,out_t) ) + + out_t_diff1 = out_t[:, 1:, :, :] - out_t[:, :-1, :, :] + out_t_diff2 = out_t[:, :, 1:, :] - out_t[:, :, :-1, :] + + total_var_loss = K.mean( 0.1*K.abs(out_t_diff1), axis=[1, 2, 3] ) + K.mean( 0.1*K.abs(out_t_diff2), axis=[1, 2, 3] ) + + opt = Adam(lr=0.0001, beta_1=0.5, beta_2=0.999, tf_cpu_mode=2) + + self.train_func = K.function ( [inp_t, real_t], [K.mean(loss)], opt.get_updates( [loss,total_var_loss], self.model.trainable_weights) ) + def __enter__(self): return self @@ -61,8 +81,9 @@ class FANSegmentator(object): def save_weights(self): self.model.save_weights (str(self.weights_path)) - def train_on_batch(self, inp, outp): - return self.model.train_on_batch(inp, outp) + def train(self, inp, real): + loss, = self.train_func ([inp, real]) + return loss def extract (self, input_image, is_input_tanh=False): input_shape_len = len(input_image.shape) @@ -78,62 +99,62 @@ class FANSegmentator(object): return result @staticmethod - def BuildModel ( resolution, ngf=64, norm='', act='lrelu'): + def BuildModel ( resolution, ngf=64): exec( nnlib.import_all(), locals(), globals() ) inp = Input ( (resolution,resolution,3) ) x = inp - x = FANSegmentator.Flow(ngf=ngf, norm=norm, act=act)(x) + x = FANSegmentator.Flow(ngf=ngf)(x) model = Model(inp,x) return model @staticmethod - def Flow(ngf=64, num_downs=4, norm='', act='lrelu'): + def Flow(ngf=64): exec( nnlib.import_all(), locals(), globals() ) def func(input): x = input x0 = x = Conv2D(ngf, kernel_size=3, strides=1, padding='same', activation='relu', name='features.0')(x) - x = MaxPooling2D()(x) + x = BlurPool(filt_size=3)(x) #x = MaxPooling2D()(x) x1 = x = Conv2D(ngf*2, kernel_size=3, strides=1, padding='same', activation='relu', name='features.3')(x) - x = MaxPooling2D()(x) + x = BlurPool(filt_size=3)(x) x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.6')(x) x2 = x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.8')(x) - x = MaxPooling2D()(x) + x = BlurPool(filt_size=3)(x) x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.11')(x) x3 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.13')(x) - x = MaxPooling2D()(x) + x = BlurPool(filt_size=3)(x) x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.16')(x) x4 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.18')(x) - x = MaxPooling2D()(x) + x = BlurPool(filt_size=3)(x) - x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same')(x) + x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', name='CA.1')(x) - x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x) + x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu', name='CA.2') (x) x = Concatenate(axis=3)([ x, x4]) - x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x) + x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu', name='CA.3') (x) - x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x) + x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu', name='CA.4') (x) x = Concatenate(axis=3)([ x, x3]) - x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x) + x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu', name='CA.5') (x) - x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu') (x) + x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu', name='CA.6') (x) x = Concatenate(axis=3)([ x, x2]) - x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu') (x) + x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu', name='CA.7') (x) - x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu') (x) + x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu', name='CA.8') (x) x = Concatenate(axis=3)([ x, x1]) - x = Conv2D (ngf*2, 3, strides=1, padding='same', activation='relu') (x) + x = Conv2D (ngf*2, 3, strides=1, padding='same', activation='relu', name='CA.9') (x) - x = Conv2DTranspose (ngf // 2, 3, strides=2, padding='same', activation='relu') (x) + x = Conv2DTranspose (ngf // 2, 3, strides=2, padding='same', activation='relu', name='CA.10') (x) x = Concatenate(axis=3)([ x, x0]) - x = Conv2D (ngf, 3, strides=1, padding='same', activation='relu') (x) + x = Conv2D (ngf, 3, strides=1, padding='same', activation='relu', name='CA.11') (x) - return Conv2D(1, 3, strides=1, padding='same', activation='sigmoid')(x) + return Conv2D(1, 3, strides=1, padding='same', activation='sigmoid', name='CA.12')(x) return func diff --git a/facelib/LandmarksProcessor.py b/facelib/LandmarksProcessor.py index e0c3840..c865bbf 100644 --- a/facelib/LandmarksProcessor.py +++ b/facelib/LandmarksProcessor.py @@ -249,52 +249,224 @@ def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0): return mat -def get_image_hull_mask (image_shape, image_landmarks, eyebrows_expand_mod=1.0, ie_polys=None): - if len(image_landmarks) != 68: - raise Exception('get_image_hull_mask works only with 68 landmarks') - int_lmrks = np.array(image_landmarks.copy(), dtype=np.int) - - hull_mask = np.zeros(image_shape[0:2]+(1,),dtype=np.float32) +def expand_eyebrows(lmrks, eyebrows_expand_mod=1.0): + if len(lmrks) != 68: + raise Exception('works only with 68 landmarks') + lmrks = np.array( lmrks.copy(), dtype=np.int ) # #nose - ml_pnt = (int_lmrks[36] + int_lmrks[0]) // 2 - mr_pnt = (int_lmrks[16] + int_lmrks[45]) // 2 + ml_pnt = (lmrks[36] + lmrks[0]) // 2 + mr_pnt = (lmrks[16] + lmrks[45]) // 2 # mid points between the mid points and eye - ql_pnt = (int_lmrks[36] + ml_pnt) // 2 - qr_pnt = (int_lmrks[45] + mr_pnt) // 2 + ql_pnt = (lmrks[36] + ml_pnt) // 2 + qr_pnt = (lmrks[45] + mr_pnt) // 2 # Top of the eye arrays - bot_l = np.array((ql_pnt, int_lmrks[36], int_lmrks[37], int_lmrks[38], int_lmrks[39])) - bot_r = np.array((int_lmrks[42], int_lmrks[43], int_lmrks[44], int_lmrks[45], qr_pnt)) + bot_l = np.array((ql_pnt, lmrks[36], lmrks[37], lmrks[38], lmrks[39])) + bot_r = np.array((lmrks[42], lmrks[43], lmrks[44], lmrks[45], qr_pnt)) # Eyebrow arrays - top_l = int_lmrks[17:22] - top_r = int_lmrks[22:27] + top_l = lmrks[17:22] + top_r = lmrks[22:27] # Adjust eyebrow arrays - int_lmrks[17:22] = top_l + eyebrows_expand_mod * 0.5 * (top_l - bot_l) - int_lmrks[22:27] = top_r + eyebrows_expand_mod * 0.5 * (top_r - bot_r) + lmrks[17:22] = top_l + eyebrows_expand_mod * 0.5 * (top_l - bot_l) + lmrks[22:27] = top_r + eyebrows_expand_mod * 0.5 * (top_r - bot_r) + return lmrks - r_jaw = (int_lmrks[0:9], int_lmrks[17:18]) - l_jaw = (int_lmrks[8:17], int_lmrks[26:27]) - r_cheek = (int_lmrks[17:20], int_lmrks[8:9]) - l_cheek = (int_lmrks[24:27], int_lmrks[8:9]) - nose_ridge = (int_lmrks[19:25], int_lmrks[8:9],) - r_eye = (int_lmrks[17:22], int_lmrks[27:28], int_lmrks[31:36], int_lmrks[8:9]) - l_eye = (int_lmrks[22:27], int_lmrks[27:28], int_lmrks[31:36], int_lmrks[8:9]) - nose = (int_lmrks[27:31], int_lmrks[31:36]) + + + +def get_image_hull_mask (image_shape, image_landmarks, eyebrows_expand_mod=1.0, ie_polys=None, color=(1,) ): + hull_mask = np.zeros(image_shape[0:2]+( len(color),),dtype=np.float32) + + lmrks = expand_eyebrows(image_landmarks, eyebrows_expand_mod) + + r_jaw = (lmrks[0:9], lmrks[17:18]) + l_jaw = (lmrks[8:17], lmrks[26:27]) + r_cheek = (lmrks[17:20], lmrks[8:9]) + l_cheek = (lmrks[24:27], lmrks[8:9]) + nose_ridge = (lmrks[19:25], lmrks[8:9],) + r_eye = (lmrks[17:22], lmrks[27:28], lmrks[31:36], lmrks[8:9]) + l_eye = (lmrks[22:27], lmrks[27:28], lmrks[31:36], lmrks[8:9]) + nose = (lmrks[27:31], lmrks[31:36]) parts = [r_jaw, l_jaw, r_cheek, l_cheek, nose_ridge, r_eye, l_eye, nose] for item in parts: merged = np.concatenate(item) - cv2.fillConvexPoly(hull_mask, cv2.convexHull(merged), 1) + cv2.fillConvexPoly(hull_mask, cv2.convexHull(merged), color ) if ie_polys is not None: ie_polys.overlay_mask(hull_mask) return hull_mask +def alpha_to_color (img_alpha, color): + if len(img_alpha.shape) == 2: + img_alpha = img_alpha[...,None] + h,w,c = img_alpha.shape + result = np.zeros( (h,w, len(color) ), dtype=np.float32 ) + result[:,:] = color + + return result * img_alpha + + + +def get_cmask (image_shape, lmrks, eyebrows_expand_mod=1.0): + h,w,c = image_shape + + hull = get_image_hull_mask (image_shape, lmrks, eyebrows_expand_mod, color=(1,) ) + + result = np.zeros( (h,w,3), dtype=np.float32 ) + + + + def process(w,h, data ): + d = {} + cur_lc = 0 + all_lines = [] + for s, pts_loop_ar in data: + lines = [] + for pts, loop in pts_loop_ar: + pts_len = len(pts) + lines.append ( [ [ pts[i], pts[(i+1) % pts_len ] ] for i in range(pts_len - (0 if loop else 1) ) ] ) + lines = np.concatenate (lines) + + lc = lines.shape[0] + all_lines.append(lines) + d[s] = cur_lc, cur_lc+lc + cur_lc += lc + all_lines = np.concatenate (all_lines, 0) + + #calculate signed distance for all points and lines + line_count = all_lines.shape[0] + pts_count = w*h + + all_lines = np.repeat ( all_lines[None,...], pts_count, axis=0 ).reshape ( (pts_count*line_count,2,2) ) + + pts = np.empty( (h,w,line_count,2), dtype=np.float32 ) + pts[...,1] = np.arange(h)[:,None,None] + pts[...,0] = np.arange(w)[:,None] + pts = pts.reshape ( (h*w*line_count, -1) ) + + a = all_lines[:,0,:] + b = all_lines[:,1,:] + pa = pts-a + ba = b-a + ph = np.clip ( np.einsum('ij,ij->i', pa, ba) / np.einsum('ij,ij->i', ba, ba), 0, 1 ) + dists = npla.norm ( pa - ba*ph[...,None], axis=1).reshape ( (h,w,line_count) ) + + def get_dists(name, thickness=0): + s,e = d[name] + result = dists[...,s:e] + if thickness != 0: + result = np.abs(result)-thickness + return np.min (result, axis=-1) + + return get_dists + + l_eye = lmrks[42:48] + r_eye = lmrks[36:42] + l_brow = lmrks[22:27] + r_brow = lmrks[17:22] + mouth = lmrks[48:60] + + up_nose = np.concatenate( (lmrks[27:31], lmrks[33:34]) ) + down_nose = lmrks[31:36] + nose = np.concatenate ( (up_nose, down_nose) ) + + gdf = process ( w,h, + ( + ('eyes', ((l_eye, True), (r_eye, True)) ), + ('brows', ((l_brow, False), (r_brow,False)) ), + ('up_nose', ((up_nose, False),) ), + ('down_nose', ((down_nose, False),) ), + ('mouth', ((mouth, True),) ), + ) + ) + + + #import code + #code.interact(local=dict(globals(), **locals())) + eyes_fall_dist = w // 32 + eyes_thickness = max( w // 64, 1 ) + + brows_fall_dist = w // 32 + brows_thickness = max( w // 256, 1 ) + + nose_fall_dist = w / 12 + nose_thickness = max( w // 96, 1 ) + + mouth_fall_dist = w // 32 + mouth_thickness = max( w // 64, 1 ) + + eyes_mask = gdf('eyes',eyes_thickness) + eyes_mask = 1-np.clip( eyes_mask/ eyes_fall_dist, 0, 1) + #eyes_mask = np.clip ( 1- ( np.sqrt( np.maximum(eyes_mask,0) ) / eyes_fall_dist ), 0, 1) + #eyes_mask = np.clip ( 1- ( np.cbrt( np.maximum(eyes_mask,0) ) / eyes_fall_dist ), 0, 1) + + brows_mask = gdf('brows', brows_thickness) + brows_mask = 1-np.clip( brows_mask / brows_fall_dist, 0, 1) + #brows_mask = np.clip ( 1- ( np.sqrt( np.maximum(brows_mask,0) ) / brows_fall_dist ), 0, 1) + + mouth_mask = gdf('mouth', mouth_thickness) + mouth_mask = 1-np.clip( mouth_mask / mouth_fall_dist, 0, 1) + #mouth_mask = np.clip ( 1- ( np.sqrt( np.maximum(mouth_mask,0) ) / mouth_fall_dist ), 0, 1) + + def blend(a,b,k): + x = np.clip ( 0.5+0.5*(b-a)/k, 0.0, 1.0 ) + return (a-b)*x+b - k*x*(1.0-x) + + + #nose_mask = (a-b)*x+b - k*x*(1.0-x) + + #nose_mask = np.minimum (up_nose_mask , down_nose_mask ) + #nose_mask = 1-np.clip( nose_mask / nose_fall_dist, 0, 1) + + nose_mask = blend ( gdf('up_nose', nose_thickness), gdf('down_nose', nose_thickness), nose_thickness*3 ) + nose_mask = 1-np.clip( nose_mask / nose_fall_dist, 0, 1) + + up_nose_mask = gdf('up_nose', nose_thickness) + up_nose_mask = 1-np.clip( up_nose_mask / nose_fall_dist, 0, 1) + #up_nose_mask = np.clip ( 1- ( np.cbrt( np.maximum(up_nose_mask,0) ) / nose_fall_dist ), 0, 1) + + down_nose_mask = gdf('down_nose', nose_thickness) + down_nose_mask = 1-np.clip( down_nose_mask / nose_fall_dist, 0, 1) + #down_nose_mask = np.clip ( 1- ( np.cbrt( np.maximum(down_nose_mask,0) ) / nose_fall_dist ), 0, 1) + + #nose_mask = np.clip( up_nose_mask + down_nose_mask, 0, 1 ) + #nose_mask /= np.max(nose_mask) + #nose_mask = np.maximum (up_nose_mask , down_nose_mask ) + #nose_mask = down_nose_mask + + #nose_mask = np.zeros_like(nose_mask) + + eyes_mask = eyes_mask * (1-mouth_mask) + nose_mask = nose_mask * (1-eyes_mask) + + hull_mask = hull[...,0].copy() + hull_mask = hull_mask * (1-eyes_mask) * (1-brows_mask) * (1-nose_mask) * (1-mouth_mask) + + #eyes_mask = eyes_mask * (1-nose_mask) + + mouth_mask= mouth_mask * (1-nose_mask) + + brows_mask = brows_mask * (1-nose_mask)* (1-eyes_mask ) + + hull_mask = alpha_to_color(hull_mask, (0,1,0) ) + eyes_mask = alpha_to_color(eyes_mask, (1,0,0) ) + brows_mask = alpha_to_color(brows_mask, (0,0,1) ) + nose_mask = alpha_to_color(nose_mask, (0,1,1) ) + mouth_mask = alpha_to_color(mouth_mask, (0,0,1) ) + + #nose_mask = np.maximum( up_nose_mask, down_nose_mask ) + + result = hull_mask + mouth_mask+ nose_mask + brows_mask + eyes_mask + result *= hull + #result = np.clip (result, 0, 1) + return result + def get_image_eye_mask (image_shape, image_landmarks): if len(image_landmarks) != 68: raise Exception('get_image_eye_mask works only with 68 landmarks') diff --git a/imagelib/__init__.py b/imagelib/__init__.py index 4d6b2a7..438db88 100644 --- a/imagelib/__init__.py +++ b/imagelib/__init__.py @@ -11,7 +11,7 @@ from .warp import gen_warp_params, warp_by_params from .reduce_colors import reduce_colors -from .color_transfer import color_hist_match, reinhard_color_transfer, linear_color_transfer +from .color_transfer import color_hist_match, reinhard_color_transfer, linear_color_transfer, seamless_clone from .RankSRGAN import RankSRGAN diff --git a/imagelib/color_transfer.py b/imagelib/color_transfer.py index eb66074..ef806f0 100644 --- a/imagelib/color_transfer.py +++ b/imagelib/color_transfer.py @@ -1,6 +1,55 @@ -import numpy as np import cv2 +import numpy as np +import scipy.sparse +from scipy.sparse.linalg import spsolve + +def laplacian_matrix(n, m): + mat_D = scipy.sparse.lil_matrix((m, m)) + mat_D.setdiag(-1, -1) + mat_D.setdiag(4) + mat_D.setdiag(-1, 1) + mat_A = scipy.sparse.block_diag([mat_D] * n).tolil() + mat_A.setdiag(-1, 1*m) + mat_A.setdiag(-1, -1*m) + return mat_A + +def seamless_clone(source, target, mask): + h, w,c = target.shape + result = [] + + mat_A = laplacian_matrix(h, w) + laplacian = mat_A.tocsc() + + mask[0,:] = 1 + mask[-1,:] = 1 + mask[:,0] = 1 + mask[:,-1] = 1 + q = np.argwhere(mask==0) + + k = q[:,1]+q[:,0]*w + mat_A[k, k] = 1 + mat_A[k, k + 1] = 0 + mat_A[k, k - 1] = 0 + mat_A[k, k + w] = 0 + mat_A[k, k - w] = 0 + + mat_A = mat_A.tocsc() + mask_flat = mask.flatten() + for channel in range(c): + + source_flat = source[:, :, channel].flatten() + target_flat = target[:, :, channel].flatten() + + mat_b = laplacian.dot(source_flat)*0.75 + mat_b[mask_flat==0] = target_flat[mask_flat==0] + + x = spsolve(mat_A, mat_b).reshape((h, w)) + result.append (x) + + + return np.clip( np.dstack(result), 0, 1 ) + def reinhard_color_transfer(target, source, clip=False, preserve_paper=False, source_mask=None, target_mask=None): """ Transfers the color distribution from the source to the target diff --git a/main.py b/main.py index bd2b537..f44f498 100644 --- a/main.py +++ b/main.py @@ -106,13 +106,17 @@ if __name__ == "__main__": #if arguments.remove_fanseg: # Util.remove_fanseg_folder (input_path=arguments.input_dir) - + + if arguments.remove_ie_polys: + Util.remove_ie_polys_folder (input_path=arguments.input_dir) + p = subparsers.add_parser( "util", help="Utilities.") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") p.add_argument('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.") p.add_argument('--add-landmarks-debug-images', action="store_true", dest="add_landmarks_debug_images", default=False, help="Add landmarks debug image for aligned faces.") p.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.") #p.add_argument('--remove-fanseg', action="store_true", dest="remove_fanseg", default=False, help="Remove fanseg mask from aligned faces.") + p.add_argument('--remove-ie-polys', action="store_true", dest="remove_ie_polys", default=False, help="Remove ie_polys from aligned faces.") p.set_defaults (func=process_util) @@ -232,13 +236,15 @@ if __name__ == "__main__": def process_labelingtool_edit_mask(arguments): from mainscripts import MaskEditorTool - MaskEditorTool.mask_editor_main (arguments.input_dir, arguments.confirmed_dir, arguments.skipped_dir) + MaskEditorTool.mask_editor_main (arguments.input_dir, arguments.confirmed_dir, arguments.skipped_dir, no_default_mask=arguments.no_default_mask) labeling_parser = subparsers.add_parser( "labelingtool", help="Labeling tool.").add_subparsers() p = labeling_parser.add_parser ( "edit_mask", help="") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.") p.add_argument('--confirmed-dir', required=True, action=fixPathAction, dest="confirmed_dir", help="This is where the labeled faces will be stored.") p.add_argument('--skipped-dir', required=True, action=fixPathAction, dest="skipped_dir", help="This is where the labeled faces will be stored.") + p.add_argument('--no-default-mask', action="store_true", dest="no_default_mask", default=False, help="Don't use default mask.") + p.set_defaults(func=process_labelingtool_edit_mask) def bad_args(arguments): diff --git a/mainscripts/Converter.py b/mainscripts/Converter.py index d437a6f..a666b35 100644 --- a/mainscripts/Converter.py +++ b/mainscripts/Converter.py @@ -73,6 +73,7 @@ class ConvertSubprocessor(Subprocessor): self.device_idx = client_dict['device_idx'] self.device_name = client_dict['device_name'] self.predictor_func = client_dict['predictor_func'] + self.predictor_input_shape = client_dict['predictor_input_shape'] self.superres_func = client_dict['superres_func'] #transfer and set stdin in order to work code.interact in debug subprocess @@ -115,15 +116,21 @@ class ConvertSubprocessor(Subprocessor): return fanseg.extract(*args, **kwargs) self.fanseg_extract_func = fanseg_extract - + + import ebsynth + def ebs_ct(*args, **kwargs): + return ebsynth.color_transfer(*args, **kwargs) + + self.ebs_ct_func = ebs_ct + return None #override def process_data(self, pf): #pf=ProcessingFrame cfg = pf.cfg.copy() - cfg.predictor_func = self.predictor_func cfg.sharpen_func = self.sharpen_func cfg.superres_func = self.superres_func + cfg.ebs_ct_func = self.ebs_ct_func frame_info = pf.frame_info @@ -152,7 +159,7 @@ class ConvertSubprocessor(Subprocessor): cfg.fanseg_extract_func = self.fanseg_extract_func try: - final_img = ConvertMasked (cfg, frame_info) + final_img = ConvertMasked (self.predictor_func, self.predictor_input_shape, cfg, frame_info) except Exception as e: e_str = traceback.format_exc() if 'MemoryError' in e_str: @@ -161,7 +168,8 @@ class ConvertSubprocessor(Subprocessor): raise Exception( 'Error while converting file [%s]: %s' % (filename, e_str) ) elif cfg.type == ConverterConfig.TYPE_FACE_AVATAR: - final_img = ConvertFaceAvatar (cfg, pf.prev_temporal_frame_infos, + final_img = ConvertFaceAvatar (self.predictor_func, self.predictor_input_shape, + cfg, pf.prev_temporal_frame_infos, pf.frame_info, pf.next_temporal_frame_infos ) @@ -179,21 +187,22 @@ class ConvertSubprocessor(Subprocessor): return pf.frame_info.filename #override - def __init__(self, is_interactive, converter_config, frames, output_path, model_iter): + def __init__(self, is_interactive, converter_session_filepath, predictor_func, predictor_input_shape, converter_config, frames, output_path, model_iter): if len (frames) == 0: raise ValueError ("len (frames) == 0") - super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if CONVERTER_DEBUG else 60, io_loop_sleep_time=0.001, initialize_subprocesses_in_serial=False)# if debug == True else 60) + super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if CONVERTER_DEBUG else 60, io_loop_sleep_time=0.001, initialize_subprocesses_in_serial=False) self.is_interactive = is_interactive + self.converter_session_filepath = Path(converter_session_filepath) self.converter_config = converter_config #dummy predict and sleep, tensorflow caching kernels. If remove it, sometime conversion speed can be x2 slower - self.converter_config.predictor_func (dummy_predict=True) + predictor_func (dummy_predict=True) time.sleep(2) - self.predictor_func_host, self.predictor_func = SubprocessFunctionCaller.make_pair(self.converter_config.predictor_func) - self.converter_config.predictor_func = None + self.predictor_func_host, self.predictor_func = SubprocessFunctionCaller.make_pair(predictor_func) + self.predictor_input_shape = predictor_input_shape self.dcscn = None self.ranksrgan = None @@ -211,11 +220,9 @@ class ConvertSubprocessor(Subprocessor): self.prefetch_frame_count = self.process_count = min(6,multiprocessing.cpu_count()) session_data = None - session_dat_path = self.output_path / 'session.dat' - if session_dat_path.exists(): - + if self.is_interactive and self.converter_session_filepath.exists(): try: - with open( str(session_dat_path), "rb") as f: + with open( str(self.converter_session_filepath), "rb") as f: session_data = pickle.loads(f.read()) except Exception as e: pass @@ -246,7 +253,7 @@ class ConvertSubprocessor(Subprocessor): break if frames_equal: - io.log_info ("Using saved session.") + io.log_info ('Using saved session from ' + '/'.join (self.converter_session_filepath.parts[-2:]) ) self.frames = s_frames self.frames_idxs = s_frames_idxs self.frames_done_idxs = s_frames_done_idxs @@ -292,6 +299,7 @@ class ConvertSubprocessor(Subprocessor): yield 'CPU%d' % (i), {}, {'device_idx': i, 'device_name': 'CPU%d' % (i), 'predictor_func': self.predictor_func, + 'predictor_input_shape' : self.predictor_input_shape, 'superres_func': self.superres_func, 'stdin_fd': sys.stdin.fileno() if CONVERTER_DEBUG else None } @@ -332,10 +340,9 @@ class ConvertSubprocessor(Subprocessor): 'frames_done_idxs': self.frames_done_idxs, 'model_iter' : self.model_iter, } - save_path = self.output_path / 'session.dat' - save_path.write_bytes( pickle.dumps(session_data) ) + self.converter_session_filepath.write_bytes( pickle.dumps(session_data) ) - io.log_info ("Session is saved to " + '/'.join (save_path.parts[-2:]) ) + io.log_info ("Session is saved to " + '/'.join (self.converter_session_filepath.parts[-2:]) ) cfg_change_keys = ['`','1', '2', '3', '4', '5', '6', '7', '8', '9', 'q', 'a', 'w', 's', 'e', 'd', 'r', 'f', 't', 'g','y','h','u','j', @@ -550,7 +557,7 @@ class ConvertSubprocessor(Subprocessor): for i in range ( min(len(self.frames_idxs), self.prefetch_frame_count) ): frame = self.frames[ self.frames_idxs[i] ] - if not frame.is_done and not frame.is_processing and frame.cfg is not None: + if not frame.is_done and not frame.is_processing and frame.cfg is not None: frame.is_processing = True return ConvertSubprocessor.ProcessingFrame(idx=frame.idx, cfg=frame.cfg.copy(), @@ -592,8 +599,8 @@ def main (args, device_args): import models model = models.import_model( args['model_name'] )(model_path, device_args=device_args) - - cfg = model.get_ConverterConfig() + converter_session_filepath = model.get_strpath_storage_for_file('converter_session.dat') + predictor_func, predictor_input_shape, cfg = model.get_ConverterConfig() if not is_interactive: cfg.ask_settings() @@ -717,7 +724,10 @@ def main (args, device_args): else: ConvertSubprocessor ( is_interactive = is_interactive, - converter_config = cfg, + converter_session_filepath = converter_session_filepath, + predictor_func = predictor_func, + predictor_input_shape = predictor_input_shape, + converter_config = cfg, frames = frames, output_path = output_path, model_iter = model.get_iter() diff --git a/mainscripts/MaskEditorTool.py b/mainscripts/MaskEditorTool.py index 1d9750f..51018cf 100644 --- a/mainscripts/MaskEditorTool.py +++ b/mainscripts/MaskEditorTool.py @@ -320,7 +320,7 @@ class MaskEditor: def get_ie_polys(self): return self.ie_polys -def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None): +def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None, no_default_mask=False): input_path = Path(input_dir) confirmed_path = Path(confirmed_dir) @@ -334,6 +334,11 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None): if not skipped_path.exists(): skipped_path.mkdir(parents=True) + + if not no_default_mask: + eyebrows_expand_mod = np.clip ( io.input_int ("Default eyebrows expand modifier? (0..400, skip:100) : ", 100), 0, 400 ) / 100.0 + else: + eyebrows_expand_mod = None wnd_name = "MaskEditor tool" io.named_window (wnd_name) @@ -407,7 +412,10 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None): if fanseg_mask is not None: mask = fanseg_mask else: - mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks) + if no_default_mask: + mask = np.zeros ( (target_wh,target_wh,3) ) + else: + mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks, eyebrows_expand_mod=eyebrows_expand_mod) else: img = np.zeros ( (target_wh,target_wh,3) ) mask = np.ones ( (target_wh,target_wh,3) ) @@ -506,7 +514,7 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None): do_save_move_count -= 1 ed.mask_finish() - dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys() ) + dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys(), eyebrows_expand_mod=eyebrows_expand_mod ) done_paths += [ confirmed_path / filepath.name ] done_images_types[filepath.name] = 2 @@ -517,7 +525,7 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None): do_save_count -= 1 ed.mask_finish() - dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys() ) + dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys(), eyebrows_expand_mod=eyebrows_expand_mod ) done_paths += [ filepath ] done_images_types[filepath.name] = 2 diff --git a/mainscripts/Util.py b/mainscripts/Util.py index 0d3e7c4..fa6f470 100644 --- a/mainscripts/Util.py +++ b/mainscripts/Util.py @@ -7,6 +7,33 @@ from utils.cv2_utils import * from facelib import LandmarksProcessor from interact import interact as io +def remove_ie_polys_file (filepath): + filepath = Path(filepath) + + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath) ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath) ) + else: + return + + if dflimg is None: + io.log_err ("%s is not a dfl image file" % (filepath.name) ) + return + + dflimg.remove_ie_polys() + dflimg.embed_and_set( str(filepath) ) + + +def remove_ie_polys_folder(input_path): + input_path = Path(input_path) + + io.log_info ("Removing ie_polys...\r\n") + + for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Removing"): + filepath = Path(filepath) + remove_ie_polys_file(filepath) + def remove_fanseg_file (filepath): filepath = Path(filepath) diff --git a/models/ModelBase.py b/models/ModelBase.py index 7acb2de..991a86b 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -117,9 +117,9 @@ class ModelBase(object): if ask_batch_size and (self.iter == 0 or ask_override): default_batch_size = 0 if self.iter == 0 else self.options.get('batch_size',0) - self.options['batch_size'] = max(0, io.input_int("Batch_size (?:help skip:%d) : " % (default_batch_size), default_batch_size, help_message="Larger batch size is better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually.")) + self.batch_size = max(0, io.input_int("Batch_size (?:help skip:%d) : " % (default_batch_size), default_batch_size, help_message="Larger batch size is better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually.")) else: - self.options['batch_size'] = self.options.get('batch_size', 0) + self.batch_size = self.options.get('batch_size', 0) if ask_sort_by_yaw: if (self.iter == 0 or ask_override): @@ -152,7 +152,7 @@ class ModelBase(object): if self.target_iter == 0 and 'target_iter' in self.options: self.options.pop('target_iter') - self.batch_size = self.options.get('batch_size',0) + #self.batch_size = self.options.get('batch_size',0) self.sort_by_yaw = self.options.get('sort_by_yaw',False) self.random_flip = self.options.get('random_flip',True) @@ -325,14 +325,9 @@ class ModelBase(object): #overridable def get_ConverterConfig(self): - #return ConverterConfig() for the model + #return predictor_func, predictor_input_shape, ConverterConfig() for the model raise NotImplementedError - #overridable - def get_converter(self): - raise NotImplementedError - #return existing or your own converter which derived from base - def get_target_iter(self): return self.target_iter diff --git a/models/Model_AVATAR/Model.py b/models/Model_AVATAR/Model.py index ba4a74e..2c1c3c7 100644 --- a/models/Model_AVATAR/Model.py +++ b/models/Model_AVATAR/Model.py @@ -353,10 +353,7 @@ class AVATARModel(ModelBase): #override def get_ConverterConfig(self): import converters - return converters.ConverterConfigFaceAvatar(predictor_func=self.predictor_func, - predictor_input_shape=(self.df_res, self.df_res, 3), - temporal_face_count=1 - ) + return self.predictor_func, (self.df_res, self.df_res, 3), converters.ConverterConfigFaceAvatar(temporal_face_count=1) @staticmethod def NLayerDiscriminator(ndf=64, n_layers=3): diff --git a/models/Model_DEV_FANSEG/Model.py b/models/Model_DEV_FANSEG/Model.py index 5ccb320..ea24c9a 100644 --- a/models/Model_DEV_FANSEG/Model.py +++ b/models/Model_DEV_FANSEG/Model.py @@ -48,7 +48,7 @@ class Model(ModelBase): self.set_training_data_generators ([ SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, sample_process_options=SampleProcessor.Options(random_flip=True), - output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution' : self.resolution, 'motion_blur':(25, 1) }, + output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution' : self.resolution, 'motion_blur':(25, 5), 'border_replicate':False }, { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_M), 'resolution': self.resolution }, ]), @@ -66,7 +66,7 @@ class Model(ModelBase): def onTrainOneIter(self, generators_samples, generators_list): target_src, target_src_mask = generators_samples[0] - loss = self.fan_seg.train_on_batch( [target_src], [target_src_mask] ) + loss = self.fan_seg.train( target_src, target_src_mask ) return ( ('loss', loss), ) diff --git a/models/Model_DF/Model.py b/models/Model_DF/Model.py index 041aa02..c0b55ac 100644 --- a/models/Model_DF/Model.py +++ b/models/Model_DF/Model.py @@ -121,16 +121,7 @@ class Model(ModelBase): #override def get_ConverterConfig(self): import converters - return converters.ConverterConfigMasked(predictor_func=self.predictor_func, - predictor_input_shape=(128,128,3), - predictor_masked=True, - face_type=FaceType.FULL, - default_mode=4, - base_erode_mask_modifier=30, - base_blur_mask_modifier=0, - default_erode_mask_modifier=0, - default_blur_mask_modifier=0, - ) + return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.FULL, default_mode=4) def Build(self, input_layer): exec(nnlib.code_import_all, locals(), globals()) diff --git a/models/Model_H128/Model.py b/models/Model_H128/Model.py index 49a5ee1..280389a 100644 --- a/models/Model_H128/Model.py +++ b/models/Model_H128/Model.py @@ -129,16 +129,7 @@ class Model(ModelBase): #override def get_ConverterConfig(self): import converters - return converters.ConverterConfigMasked(predictor_func=self.predictor_func, - predictor_input_shape=(128,128,3), - predictor_masked=True, - face_type=FaceType.HALF, - default_mode=4, - base_erode_mask_modifier=100, - base_blur_mask_modifier=100, - default_erode_mask_modifier=0, - default_blur_mask_modifier=0, - ) + return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.HALF, default_mode=4) def Build(self, lighter_ae): exec(nnlib.code_import_all, locals(), globals()) diff --git a/models/Model_H64/Model.py b/models/Model_H64/Model.py index 9d45c32..d288206 100644 --- a/models/Model_H64/Model.py +++ b/models/Model_H64/Model.py @@ -130,16 +130,7 @@ class Model(ModelBase): #override def get_ConverterConfig(self): import converters - return converters.ConverterConfigMasked(predictor_func=self.predictor_func, - predictor_input_shape=(64,64,3), - predictor_masked=True, - face_type=FaceType.HALF, - default_mode=4, - base_erode_mask_modifier=100, - base_blur_mask_modifier=100, - default_erode_mask_modifier=0, - default_blur_mask_modifier=0, - ) + return self.predictor_func, (64,64,3), converters.ConverterConfigMasked(face_type=FaceType.HALF, default_mode=4) def Build(self, lighter_ae): exec(nnlib.code_import_all, locals(), globals()) diff --git a/models/Model_LIAEF128/Model.py b/models/Model_LIAEF128/Model.py index 96a0964..a908677 100644 --- a/models/Model_LIAEF128/Model.py +++ b/models/Model_LIAEF128/Model.py @@ -127,16 +127,7 @@ class Model(ModelBase): #override def get_ConverterConfig(self): import converters - return converters.ConverterConfigMasked(predictor_func=self.predictor_func, - predictor_input_shape=(128,128,3), - predictor_masked=True, - face_type=FaceType.FULL, - default_mode=4, - base_erode_mask_modifier=30, - base_blur_mask_modifier=0, - default_erode_mask_modifier=0, - default_blur_mask_modifier=0, - ) + return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.FULL, default_mode=4) def Build(self, input_layer): exec(nnlib.code_import_all, locals(), globals()) diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index 555880d..3e3eb74 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -483,25 +483,11 @@ class SAEModel(ModelBase): #override def get_ConverterConfig(self): - base_erode_mask_modifier = 30 if self.options['face_type'] == 'f' else 100 - base_blur_mask_modifier = 0 if self.options['face_type'] == 'f' else 100 - - default_erode_mask_modifier = 0 - default_blur_mask_modifier = 100 if (self.options['face_style_power'] or self.options['bg_style_power']) and \ - self.options['face_type'] == 'f' else 0 - face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF import converters - return converters.ConverterConfigMasked(predictor_func=self.predictor_func, - predictor_input_shape=(self.options['resolution'], self.options['resolution'], 3), - predictor_masked=self.options['learn_mask'], - face_type=face_type, + return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), converters.ConverterConfigMasked(face_type=face_type, default_mode = 1 if self.options['apply_random_ct'] or self.options['face_style_power'] or self.options['bg_style_power'] else 4, - base_erode_mask_modifier=base_erode_mask_modifier, - base_blur_mask_modifier=base_blur_mask_modifier, - default_erode_mask_modifier=default_erode_mask_modifier, - default_blur_mask_modifier=default_blur_mask_modifier, clip_hborder_mask_per=0.0625 if (self.options['face_type'] == 'f') else 0, ) diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index 257c0e7..3efe437 100644 --- a/nnlib/nnlib.py +++ b/nnlib/nnlib.py @@ -685,7 +685,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator def CAInitializerMP( conv_weights_list ): #Convolution Aware Initialization https://arxiv.org/abs/1702.06295 data = [ (i, K.int_shape(conv_weights)) for i, conv_weights in enumerate(conv_weights_list) ] - data = sorted(data, key=lambda data: sum(data[1]) ) + data = sorted(data, key=lambda data: np.prod(data[1]) ) result = CAInitializerMPSubprocessor (data, K.floatx(), K.image_data_format() ).run() for idx, weights in result: K.set_value ( conv_weights_list[idx], weights ) diff --git a/samplelib/Sample.py b/samplelib/Sample.py index d22a869..34f3a58 100644 --- a/samplelib/Sample.py +++ b/samplelib/Sample.py @@ -22,7 +22,7 @@ class SampleType(IntEnum): QTY = 5 class Sample(object): - def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask_exist=False): + def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, eyebrows_expand_mod=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask_exist=False): self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE self.filename = filename self.face_type = face_type @@ -30,12 +30,13 @@ class Sample(object): self.landmarks = np.array(landmarks) if landmarks is not None else None self.ie_polys = ie_polys self.pitch_yaw_roll = pitch_yaw_roll + self.eyebrows_expand_mod = eyebrows_expand_mod self.source_filename = source_filename self.mirror = mirror self.close_target_list = close_target_list self.fanseg_mask_exist = fanseg_mask_exist - def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask=None, fanseg_mask_exist=None): + def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, eyebrows_expand_mod=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask=None, fanseg_mask_exist=None): return Sample( sample_type=sample_type if sample_type is not None else self.sample_type, filename=filename if filename is not None else self.filename, @@ -44,6 +45,7 @@ class Sample(object): landmarks=landmarks if landmarks is not None else self.landmarks.copy(), ie_polys=ie_polys if ie_polys is not None else self.ie_polys, pitch_yaw_roll=pitch_yaw_roll if pitch_yaw_roll is not None else self.pitch_yaw_roll, + eyebrows_expand_mod=eyebrows_expand_mod if eyebrows_expand_mod is not None else self.eyebrows_expand_mod, source_filename=source_filename if source_filename is not None else self.source_filename, mirror=mirror if mirror is not None else self.mirror, close_target_list=close_target_list if close_target_list is not None else self.close_target_list, diff --git a/samplelib/SampleLoader.py b/samplelib/SampleLoader.py index 762171f..2d7d467 100644 --- a/samplelib/SampleLoader.py +++ b/samplelib/SampleLoader.py @@ -71,9 +71,10 @@ class SampleLoader: landmarks = dflimg.get_landmarks() pitch_yaw_roll = dflimg.get_pitch_yaw_roll() - if pitch_yaw_roll is None: - pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(landmarks) + eyebrows_expand_mod = dflimg.get_eyebrows_expand_mod() + if pitch_yaw_roll is None: + pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(landmarks) sample_list.append( s.copy_and_set(sample_type=SampleType.FACE, face_type=FaceType.fromString (dflimg.get_face_type()), @@ -81,6 +82,7 @@ class SampleLoader: landmarks=landmarks, ie_polys=dflimg.get_ie_polys(), pitch_yaw_roll=pitch_yaw_roll, + eyebrows_expand_mod=eyebrows_expand_mod, source_filename=dflimg.get_source_filename(), fanseg_mask_exist=dflimg.get_fanseg_mask() is not None, ) ) except: diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index ac78ead..25a6644 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -36,7 +36,7 @@ opts: '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 + 'motion_blur' : (chance_int, range) - chance 0..100 to apply to face (not mask), and max_size of motion blur 'apply_ct' : bool 'normalize_tanh' : bool @@ -116,6 +116,7 @@ class SampleProcessor(object): resolution = opts.get('resolution', 0) types = opts.get('types', [] ) + border_replicate = opts.get('border_replicate', True) 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) @@ -167,7 +168,7 @@ class SampleProcessor(object): transform = (img_type==SPTF.IMG_WARPED_TRANSFORMED or img_type==SPTF.IMG_TRANSFORMED) flip = img_type != SPTF.IMG_WARPED - img = imagelib.warp_by_params (params, img, warp, transform, flip, True) + img = imagelib.warp_by_params (params, img, warp, transform, flip, border_replicate) if mask is not None: mask = imagelib.warp_by_params (params, mask, warp, transform, flip, False) if len(mask.shape) == 2: @@ -176,38 +177,30 @@ class SampleProcessor(object): img = np.concatenate( (img, mask ), -1 ) return img - img = cached_images.get(img_type, None) - if img is None: + img = sample_bgr + + ### Prepare a mask + mask = None + if is_face_sample: + mask = sample.load_fanseg_mask() #using fanseg_mask if exist - img = sample_bgr - mask = None - cur_sample = sample + if mask is None: + if sample.eyebrows_expand_mod is not None: + mask = LandmarksProcessor.get_image_hull_mask (img.shape, sample.landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod ) + else: + mask = LandmarksProcessor.get_image_hull_mask (img.shape, sample.landmarks) - if is_face_sample: - if motion_blur is not None: - chance, mb_range = motion_blur - chance = np.clip(chance, 0, 100) + if sample.ie_polys is not None: + sample.ie_polys.overlay_mask(mask) + ################## + + + if motion_blur is not None: + chance, mb_max_size = motion_blur + 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) ) - - 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 cur_sample.ie_polys is not None: - cur_sample.ie_polys.overlay_mask(mask) - - if sample.face_type == FaceType.MARK_ONLY: - if mask is not None: - img = np.concatenate( (img, mask), -1 ) - else: - img = do_transform (img, mask) - - cached_images[img_type] = img + if np.random.randint(100) < chance: + img = imagelib.LinearMotionBlur (img, np.random.randint( mb_max_size )+1, np.random.randint(360) ) if is_face_sample and target_face_type != SPTF.NONE: target_ft = SampleProcessor.SPTF_FACETYPE_TO_FACETYPE[target_face_type] @@ -215,16 +208,18 @@ class SampleProcessor(object): raise Exception ('sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, target_ft) ) if sample.face_type == FaceType.MARK_ONLY: - img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], target_ft), (sample.shape[0],sample.shape[0]), flags=cv2.INTER_CUBIC ) - - mask = img[...,3:4] if img.shape[2] > 3 else None - img = img[...,0:3] + #first warp to target facetype + img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], target_ft), (sample.shape[0],sample.shape[0]), flags=cv2.INTER_CUBIC ) + mask = cv2.warpAffine( mask, LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], target_ft), (sample.shape[0],sample.shape[0]), flags=cv2.INTER_CUBIC ) + #then apply transforms img = do_transform (img, mask) img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) else: - img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, target_ft), (resolution,resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC ) + img = do_transform (img, mask) + img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, target_ft), (resolution,resolution), borderMode=(cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT), flags=cv2.INTER_CUBIC ) else: + img = do_transform (img, mask) img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) if random_sub_res != 0: diff --git a/utils/DFLJPG.py b/utils/DFLJPG.py index c06e96a..888ff98 100644 --- a/utils/DFLJPG.py +++ b/utils/DFLJPG.py @@ -168,7 +168,8 @@ class DFLJPG(object): source_landmarks=None, image_to_face_mat=None, fanseg_mask=None, - pitch_yaw_roll=None, + pitch_yaw_roll=None, + eyebrows_expand_mod=None, **kwargs ): @@ -193,7 +194,8 @@ class DFLJPG(object): 'source_landmarks': source_landmarks, 'image_to_face_mat': image_to_face_mat, 'fanseg_mask' : fanseg_mask, - 'pitch_yaw_roll' : pitch_yaw_roll + 'pitch_yaw_roll' : pitch_yaw_roll, + 'eyebrows_expand_mod' : eyebrows_expand_mod }) try: @@ -211,6 +213,7 @@ class DFLJPG(object): image_to_face_mat=None, fanseg_mask=None, pitch_yaw_roll=None, + eyebrows_expand_mod=None, **kwargs ): if face_type is None: face_type = self.get_face_type() @@ -222,6 +225,8 @@ class DFLJPG(object): if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat() if fanseg_mask is None: fanseg_mask = self.get_fanseg_mask() if pitch_yaw_roll is None: pitch_yaw_roll = self.get_pitch_yaw_roll() + if eyebrows_expand_mod is None: eyebrows_expand_mod = self.get_eyebrows_expand_mod() + DFLJPG.embed_data (filename, face_type=face_type, landmarks=landmarks, ie_polys=ie_polys, @@ -230,7 +235,12 @@ class DFLJPG(object): source_landmarks=source_landmarks, image_to_face_mat=image_to_face_mat, fanseg_mask=fanseg_mask, - pitch_yaw_roll=pitch_yaw_roll) + pitch_yaw_roll=pitch_yaw_roll, + eyebrows_expand_mod=eyebrows_expand_mod) + + def remove_ie_polys(self): + self.dfl_dict['ie_polys'] = None + def remove_fanseg_mask(self): self.dfl_dict['fanseg_mask'] = None @@ -300,4 +310,6 @@ class DFLJPG(object): return None def get_pitch_yaw_roll(self): return self.dfl_dict.get ('pitch_yaw_roll', None) - + def get_eyebrows_expand_mod(self): + return self.dfl_dict.get ('eyebrows_expand_mod', None) + diff --git a/utils/DFLPNG.py b/utils/DFLPNG.py index d1fb76e..7650f91 100644 --- a/utils/DFLPNG.py +++ b/utils/DFLPNG.py @@ -285,6 +285,7 @@ class DFLPNG(object): image_to_face_mat=None, fanseg_mask=None, pitch_yaw_roll=None, + eyebrows_expand_mod=None, **kwargs ): @@ -309,7 +310,8 @@ class DFLPNG(object): 'source_landmarks': source_landmarks, 'image_to_face_mat':image_to_face_mat, 'fanseg_mask' : fanseg_mask, - 'pitch_yaw_roll' : pitch_yaw_roll + 'pitch_yaw_roll' : pitch_yaw_roll, + 'eyebrows_expand_mod' : eyebrows_expand_mod, }) try: @@ -327,6 +329,7 @@ class DFLPNG(object): image_to_face_mat=None, fanseg_mask=None, pitch_yaw_roll=None, + eyebrows_expand_mod=None, **kwargs ): if face_type is None: face_type = self.get_face_type() @@ -338,6 +341,8 @@ class DFLPNG(object): if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat() if fanseg_mask is None: fanseg_mask = self.get_fanseg_mask() if pitch_yaw_roll is None: pitch_yaw_roll = self.get_pitch_yaw_roll() + if eyebrows_expand_mod is None: eyebrows_expand_mod = self.get_eyebrows_expand_mod() + DFLPNG.embed_data (filename, face_type=face_type, landmarks=landmarks, ie_polys=ie_polys, @@ -346,7 +351,11 @@ class DFLPNG(object): source_landmarks=source_landmarks, image_to_face_mat=image_to_face_mat, fanseg_mask=fanseg_mask, - pitch_yaw_roll=pitch_yaw_roll) + pitch_yaw_roll=pitch_yaw_roll, + eyebrows_expand_mod=eyebrows_expand_mod) + + def remove_ie_polys(self): + self.dfl_dict['ie_polys'] = None def remove_fanseg_mask(self): self.dfl_dict['fanseg_mask'] = None @@ -406,5 +415,8 @@ class DFLPNG(object): return None def get_pitch_yaw_roll(self): return self.dfl_dict.get ('pitch_yaw_roll', None) + def get_eyebrows_expand_mod(self): + return self.dfl_dict.get ('eyebrows_expand_mod', None) + def __str__(self): return "".format(len(self.chunks), **self.__dict__)