added mouth prio

facelib/LandmarksProcessor.py

@@ -433,6 +433,27 @@ def get_image_eye_mask (image_shape, image_landmarks):
 
     return hull_mask
 
+def get_image_mouth_mask (image_shape, image_landmarks):
+    if len(image_landmarks) != 68:
+        raise Exception('get_image_eye_mask works only with 68 landmarks')
+
+    h,w,c = image_shape
+
+    hull_mask = np.zeros( (h,w,1),dtype=np.float32)
+
+    image_landmarks = image_landmarks.astype(np.int)
+
+    cv2.fillConvexPoly( hull_mask, cv2.convexHull( image_landmarks[48, 60]), (1,) )
+
+    dilate = h // 32
+    hull_mask = cv2.dilate(hull_mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(dilate,dilate)), iterations = 1 )
+
+    blur = h // 16
+    blur = blur + (1-blur % 2)
+    hull_mask = cv2.GaussianBlur(hull_mask, (blur, blur) , 0)
+    hull_mask = hull_mask[...,None]
+
+    return hull_mask
 
 def alpha_to_color (img_alpha, color):
     if len(img_alpha.shape) == 2:

models/Model_SAEHD/Model.py

@@ -44,6 +44,7 @@ class SAEHDModel(ModelBase):
         default_d_mask_dims        = self.options['d_mask_dims']        = self.options.get('d_mask_dims', None)
         default_masked_training    = self.options['masked_training']    = self.load_or_def_option('masked_training', True)
         default_eyes_prio          = self.options['eyes_prio']          = self.load_or_def_option('eyes_prio', False)
+        default_mouth_prio         = self.options['mouth_prio']         = self.load_or_def_option('mouth_prio', False)
         default_uniform_yaw        = self.options['uniform_yaw']        = self.load_or_def_option('uniform_yaw', False)
 
         lr_dropout = self.load_or_def_option('lr_dropout', 'n')
@@ -130,6 +131,8 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                 self.options['masked_training']  = io.input_bool ("Masked training", default_masked_training, help_message="This option is available only for 'whole_face' or 'head' type. Masked training clips training area to full_face mask or XSeg mask, thus network will train the faces properly.")
 
             self.options['eyes_prio'] = io.input_bool ("Eyes priority", default_eyes_prio, help_message='Helps to fix eye problems during training like "alien eyes" and wrong eyes direction ( especially on HD architectures ) by forcing the neural network to train eyes with higher priority. before/after https://i.imgur.com/YQHOuSR.jpg ')
+            self.options['mouth_prio'] = io.input_bool ("Mouth priority", default_mouth_prio, help_message='Helps to fix mouth problems during training by forcing the neural network to train mouth with higher priority similar to eyes ')
+
             self.options['uniform_yaw'] = io.input_bool ("Uniform yaw distribution of samples", default_uniform_yaw, help_message='Helps to fix blurry side faces due to small amount of them in the faceset.')
 
         if self.is_first_run() or ask_override:
@@ -175,6 +178,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                           'head' : FaceType.HEAD}[ self.options['face_type'] ]
 
         eyes_prio = self.options['eyes_prio']
+        mouth_prio = self.options['mouth_prio']
 
         archi_split = self.options['archi'].split('-')
 
@@ -363,8 +367,12 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                     # unpack masks from one combined mask
                     gpu_target_srcm      = tf.clip_by_value (gpu_target_srcm_all, 0, 1)
                     gpu_target_dstm      = tf.clip_by_value (gpu_target_dstm_all, 0, 1)
-                    gpu_target_srcm_eyes = tf.clip_by_value (gpu_target_srcm_all-1, 0, 1)
+                    gpu_target_srcm_eye_mouth = tf.clip_by_value (gpu_target_srcm_all-1, 0, 1)
-                    gpu_target_dstm_eyes = tf.clip_by_value (gpu_target_dstm_all-1, 0, 1)
+                    gpu_target_dstm_eye_mouth = tf.clip_by_value (gpu_target_dstm_all-1, 0, 1)
+                    gpu_target_srcm_mouth = tf.clip_by_value (gpu_target_srcm_all-2, 0, 1)
+                    gpu_target_dstm_mouth = tf.clip_by_value (gpu_target_dstm_all-2, 0, 1)
+                    gpu_target_srcm_eyes = tf.clip_by_value (gpu_target_srcm_eye_mouth-gpu_target_srcm_mouth, 0, 1)
+                    gpu_target_dstm_eyes = tf.clip_by_value (gpu_target_dstm_eye_mouth-gpu_target_dstm_mouth, 0, 1)
 
                     gpu_target_srcm_blur = nn.gaussian_blur(gpu_target_srcm,  max(1, resolution // 32) )
                     gpu_target_srcm_blur = tf.clip_by_value(gpu_target_srcm_blur, 0, 0.5) * 2
@@ -393,8 +401,15 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                         gpu_src_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1])
                     gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
 
-                    if eyes_prio:
+                    if eyes_prio or mouth_prio:
-                        gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_srcm_eyes - gpu_pred_src_src*gpu_target_srcm_eyes ), axis=[1,2,3])
+                        if eyes_prio and mouth_prio:
+                            gpu_target_part_mask = gpu_target_srcm_eye_mouth
+                        elif eyes_prio:
+                            gpu_target_part_mask = gpu_target_srcm_eyes
+                        elif mouth_prio:
+                            gpu_target_part_mask = gpu_target_srcm_mouth
+
+                        gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_part_mask - gpu_pred_src_src*gpu_target_part_mask ), axis=[1,2,3])
 
                     gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )
 
@@ -415,8 +430,15 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                     gpu_dst_loss += tf.reduce_mean ( 10*tf.square(  gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
 
 
-                    if eyes_prio:
+                    if eyes_prio or mouth_prio:
-                        gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_dstm_eyes - gpu_pred_dst_dst*gpu_target_dstm_eyes ), axis=[1,2,3])
+                        if eyes_prio and mouth_prio:
+                            gpu_target_part_mask = gpu_target_dstm_eye_mouth
+                        elif eyes_prio:
+                            gpu_target_part_mask = gpu_target_dstm_eyes
+                        elif mouth_prio:
+                            gpu_target_part_mask = gpu_target_dstm_mouth
+
+                        gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_part_mask - gpu_pred_dst_dst*gpu_target_part_mask ), axis=[1,2,3])
 
                     gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )
 

samplelib/SampleProcessor.py

@@ -70,7 +70,15 @@ class SampleProcessor(object):
                 
             def get_eyes_mask():
                 eyes_mask = LandmarksProcessor.get_image_eye_mask (sample_bgr.shape, sample_landmarks)
-                return np.clip(eyes_mask, 0, 1)
+                # set eye masks to 1-2
+                clip = np.clip(eyes_mask, 0, 1)
+                return a[a > 0.1] += 1
+
+            def get_mouth_mask():
+                mouth_mask = LandmarksProcessor.get_image_mouth_mask (sample_bgr.shape, sample_landmarks)
+                # set eye masks to 2-3
+                clip = np.clip(mouth_mask, 0, 1)
+                return a[a > 0.1] += 2
 
             is_face_sample = sample_landmarks is not None
 
@@ -136,8 +144,12 @@ class SampleProcessor(object):
                         elif face_mask_type == SPFMT.EYES:
                             img = get_eyes_mask()
                         elif face_mask_type == SPFMT.FULL_FACE_EYES:
+                            # sets both eyes and mouth mask parts
                             img = get_full_face_mask()                            
-                            img += get_eyes_mask()*img
+                            eye_mask = get_eyes_mask()
+                            img = np.where(eye_mask > 1, eye_mask, img)
+                            mouth_mask = get_mouth_mask()
+                            img = np.where(mouth_mask > 2, mouth_mask, img)
                         else:
                             img = np.zeros ( sample_bgr.shape[0:2]+(1,), dtype=np.float32)