Converter:

Session is now saved to the model folder.

blur and erode ranges are increased to -400+400

hist-match-bw is now replaced with seamless2 mode.

Added 'ebs' color transfer mode (works only on Windows).

FANSEG model (used in FAN-x mask modes) is retrained with new model configuration
and now produces better precision and less jitter

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)

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

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

ebsynth/__init__.py

@@ -0,0 +1 @@
+from .ebsynth import color_transfer

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 )

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

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')

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
 

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

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):

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()

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

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)
 

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
 

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):

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), )
         

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())

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())

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())

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())

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,
                                     )
 

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 )

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,

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:

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:

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)
+

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 "<PNG length={length} chunks={}>".format(len(self.chunks), **self.__dict__)