From 2571f23bc4850ad86eb464f7d563a0e932b73c3b Mon Sep 17 00:00:00 2001 From: iperov Date: Tue, 22 Jan 2019 12:46:54 +0400 Subject: [PATCH] Extractor added new feature --manual-output-debug-fix, Performs manual reextract input-dir frames which were deleted from [output_dir]_debug\ dir. --- main.py | 6 ++++-- mainscripts/Extractor.py | 42 +++++++++++++++++++++++++++++++--------- 2 files changed, 37 insertions(+), 11 deletions(-) diff --git a/main.py b/main.py index 1c0e000..dace3c7 100644 --- a/main.py +++ b/main.py @@ -38,6 +38,7 @@ if __name__ == "__main__": multi_gpu=arguments.multi_gpu, cpu_only=arguments.cpu_only, manual_fix=arguments.manual_fix, + manual_output_debug_fix=arguments.manual_output_debug_fix, manual_window_size=arguments.manual_window_size ) @@ -49,7 +50,8 @@ if __name__ == "__main__": extract_parser.add_argument('--detector', dest="detector", choices=['dlib','mt','manual'], default='dlib', help="Type of detector. Default 'dlib'. 'mt' (MTCNNv1) - faster, better, almost no jitter, perfect for gathering thousands faces for src-set. It is also good for dst-set, but can generate false faces in frames where main face not recognized! In this case for dst-set use either 'dlib' with '--manual-fix' or '--detector manual'. Manual detector suitable only for dst-set.") extract_parser.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.") extract_parser.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.") - extract_parser.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=0, help="Manual fix window size. Example: 1368. Default: frame size.") + extract_parser.add_argument('--manual-output-debug-fix', action="store_true", dest="manual_output_debug_fix", default=False, help="Performs manual reextract input-dir frames which were deleted from [output_dir]_debug\ dir.") + extract_parser.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.") extract_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.") @@ -61,7 +63,7 @@ if __name__ == "__main__": sort_parser = subparsers.add_parser( "sort", help="Sort faces in a directory.") sort_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") - sort_parser.add_argument('--by', required=True, dest="sort_by_method", choices=("blur", "face", "face-dissim", "face-yaw", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "final"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." ) + sort_parser.add_argument('--by', required=True, dest="sort_by_method", choices=("blur", "face", "face-dissim", "face-yaw", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "final", "test"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." ) sort_parser.set_defaults (func=process_sort) def process_train(arguments): diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index e42ec83..9a4dd9d 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -311,7 +311,7 @@ class ExtractSubprocessor(SubprocessorBase): faces = data[1] if self.debug: - debug_output_file = '{}_{}'.format( str(Path(str(self.output_path) + '_debug') / filename_path.stem), 'debug.jpg') + debug_output_file = '{}{}'.format( str(Path(str(self.output_path) + '_debug') / filename_path.stem), '.jpg') debug_image = image.copy() for (face_idx, face) in enumerate(faces): @@ -397,7 +397,7 @@ face_type 'full_face' 'avatar' ''' -def main (input_dir, output_dir, debug, detector='mt', multi_gpu=True, cpu_only=False, manual_fix=False, manual_window_size=0, image_size=256, face_type='full_face'): +def main (input_dir, output_dir, debug, detector='mt', multi_gpu=True, cpu_only=False, manual_fix=False, manual_output_debug_fix=False, manual_window_size=1368, image_size=256, face_type='full_face'): print ("Running extractor.\r\n") input_path = Path(input_dir) @@ -409,20 +409,44 @@ def main (input_dir, output_dir, debug, detector='mt', multi_gpu=True, cpu_only= return if output_path.exists(): - for filename in Path_utils.get_image_paths(output_path): - Path(filename).unlink() + if not manual_output_debug_fix: + for filename in Path_utils.get_image_paths(output_path): + Path(filename).unlink() else: output_path.mkdir(parents=True, exist_ok=True) + if manual_output_debug_fix: + debug = True + detector = 'manual' + print('Performing re-extract frames which were deleted from _debug directory.') + + input_path_image_paths = Path_utils.get_image_unique_filestem_paths(input_path, verbose=True) + if debug: debug_output_path = Path(str(output_path) + '_debug') - if debug_output_path.exists(): - for filename in Path_utils.get_image_paths(debug_output_path): - Path(filename).unlink() + + if manual_output_debug_fix: + if not debug_output_path.exists(): + print ("%s not found " % ( str(debug_output_path) )) + return + + debug_imgs = Path_utils.get_image_paths(debug_output_path) + + new_input_path_image_paths = [] + for i in input_path_image_paths: + i_stem = Path(i).stem + + if not any ( [ i_stem == Path(d).stem for d in debug_imgs] ): + new_input_path_image_paths += [i] + + input_path_image_paths = new_input_path_image_paths else: - debug_output_path.mkdir(parents=True, exist_ok=True) + if debug_output_path.exists(): + for filename in Path_utils.get_image_paths(debug_output_path): + Path(filename).unlink() + else: + debug_output_path.mkdir(parents=True, exist_ok=True) - input_path_image_paths = Path_utils.get_image_unique_filestem_paths(input_path, verbose=True) images_found = len(input_path_image_paths) faces_detected = 0 if images_found != 0: