diff --git a/facelib/LandmarksProcessor.py b/facelib/LandmarksProcessor.py index 8d3a327..2181fe1 100644 --- a/facelib/LandmarksProcessor.py +++ b/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: diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 3803504..aa2bc2f 100644 --- a/models/Model_SAEHD/Model.py +++ b/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_dstm_eyes = tf.clip_by_value (gpu_target_dstm_all-1, 0, 1) + gpu_target_srcm_eye_mouth = tf.clip_by_value (gpu_target_srcm_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: - 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 or mouth_prio: + 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: - 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 or mouth_prio: + 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] ) diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index 9d734d0..a951df4 100644 --- a/samplelib/SampleProcessor.py +++ b/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)