diff --git a/facelib/LandmarksProcessor.py b/facelib/LandmarksProcessor.py index a82d1a1..5f1cac9 100644 --- a/facelib/LandmarksProcessor.py +++ b/facelib/LandmarksProcessor.py @@ -247,7 +247,7 @@ def mirror_landmarks (landmarks, val): result[:,0] = val - result[:,0] - 1 return result -def draw_landmarks (image, image_landmarks, color): +def draw_landmarks (image, image_landmarks, color=(0,255,0)): if len(image_landmarks) != 68: raise Exception('get_image_eye_mask works only with 68 landmarks') @@ -274,7 +274,7 @@ def draw_landmarks (image, image_landmarks, color): def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type): image_utils.draw_rect (image, rect, (255,0,0), 2 ) - draw_landmarks(image, image_landmarks, (0,255,0) ) + draw_landmarks(image, image_landmarks) image_to_face_mat = get_transform_mat (image_landmarks, face_size, face_type) points = transform_points ( [ (0,0), (0,face_size-1), (face_size-1, face_size-1), (face_size-1,0) ], image_to_face_mat, True) diff --git a/main.py b/main.py index 4247da1..eee0e7f 100644 --- a/main.py +++ b/main.py @@ -70,9 +70,13 @@ if __name__ == "__main__": if arguments.convert_png_to_jpg: Util.convert_png_to_jpg_folder (input_path=arguments.input_dir) + if arguments.add_landmarks_debug_images: + Util.add_landmarks_debug_images (input_path=arguments.input_dir) + util_parser = subparsers.add_parser( "util", help="Utilities.") util_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") util_parser.add_argument('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.") + util_parser.add_argument('--add-landmarks-debug-images', action="store_true", dest="add_landmarks_debug_images", default=False, help="Add landmarks debug image for aligned faces.") util_parser.set_defaults (func=process_util) def process_train(arguments): diff --git a/mainscripts/Util.py b/mainscripts/Util.py index 39a9e69..22e8455 100644 --- a/mainscripts/Util.py +++ b/mainscripts/Util.py @@ -48,3 +48,26 @@ def convert_png_to_jpg_folder (input_path): for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Converting", ascii=True): filepath = Path(filepath) convert_png_to_jpg_file(filepath) + +def add_landmarks_debug_images(input_path): + print ("Adding landmarks debug images...") + + for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Processing", ascii=True): + filepath = Path(filepath) + + img = cv2_imread(str(filepath)) + + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True ) + else: + print ("%s is not a dfl image file" % (filepath.name) ) + continue + + if not (dflimg is None or img is None): + face_landmarks = dflimg.get_landmarks() + LandmarksProcessor.draw_landmarks(img, face_landmarks) + + output_file = '{}{}'.format( str(Path(str(input_path)) / filepath.stem), '_debug.jpg') + cv2_imwrite(output_file, img, [int(cv2.IMWRITE_JPEG_QUALITY), 50] )