diff --git a/main.py b/main.py index 386b683..fb64f96 100644 --- a/main.py +++ b/main.py @@ -38,6 +38,7 @@ if __name__ == "__main__": manual_output_debug_fix = arguments.manual_output_debug_fix, manual_window_size = arguments.manual_window_size, face_type = arguments.face_type, + max_faces_from_image = arguments.max_faces_from_image, image_size = arguments.image_size, jpeg_quality = arguments.jpeg_quality, cpu_only = arguments.cpu_only, @@ -51,6 +52,7 @@ if __name__ == "__main__": p.add_argument('--output-debug', action="store_true", dest="output_debug", default=None, help="Writes debug images to _debug\ directory.") p.add_argument('--no-output-debug', action="store_false", dest="output_debug", default=None, help="Don't writes debug images to _debug\ directory.") p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'whole_face', 'head', 'mark_only'], default=None) + p.add_argument('--max-faces-from-image', type=int, dest="max_faces_from_image", default=None, help="Max faces from image.") p.add_argument('--image-size', type=int, dest="image_size", default=None, help="Output image size.") p.add_argument('--jpeg-quality', type=int, dest="jpeg_quality", default=None, help="Jpeg quality.") p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.") diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index 5d7d393..9b5476a 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -151,7 +151,7 @@ class ExtractSubprocessor(Subprocessor): if len(rects) != 0: data.rects_rotation = rot break - if max_faces_from_image != 0 and len(data.rects) > 1: + if max_faces_from_image is not None and len(data.rects) > 1: data.rects = data.rects[0:max_faces_from_image] return data @@ -730,9 +730,9 @@ def main(detector=None, manual_output_debug_fix=False, manual_window_size=1368, face_type='full_face', + max_faces_from_image=None, image_size=None, jpeg_quality=None, - max_faces_from_image=0, cpu_only = False, force_gpu_idxs = None, ): @@ -741,26 +741,34 @@ def main(detector=None, io.log_err ('Input directory not found. Please ensure it exists.') return + if not output_path.exists(): + output_path.mkdir(parents=True, exist_ok=True) + if face_type is not None: face_type = FaceType.fromString(face_type) if face_type is None: - if manual_output_debug_fix and output_path.exists(): + if manual_output_debug_fix: files = pathex.get_image_paths(output_path) if len(files) != 0: dflimg = DFLIMG.load(Path(files[0])) if dflimg is not None and dflimg.has_data(): face_type = FaceType.fromString ( dflimg.get_face_type() ) - if output_path.exists(): - if not manual_output_debug_fix and input_path != output_path: - output_images_paths = pathex.get_image_paths(output_path) - if len(output_images_paths) > 0: + input_image_paths = pathex.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info) + output_images_paths = pathex.get_image_paths(output_path) + + if not manual_output_debug_fix and len(output_images_paths) > 0: + if len(output_images_paths) > 128 and io.input_bool ("Continue extraction?", True, help_message="Extraction can be continued, but you must specify the same options again."): + try: + input_image_paths = input_image_paths[ [ Path(x).stem for x in input_image_paths ].index ( Path(output_images_paths[-128]).stem.split('_')[0] ) : ] + except: + io.log_err("Error in fetching the last index. Extraction cannot be continued.") + return + elif input_path != output_path: io.input(f"\n WARNING !!! \n {output_path} contains files! \n They will be deleted. \n Press enter to continue.\n") for filename in output_images_paths: Path(filename).unlink() - else: - output_path.mkdir(parents=True, exist_ok=True) device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(choose_only_one=detector=='manual', suggest_all_gpu=True) ) \ if not cpu_only else nn.DeviceConfig.CPU() @@ -771,6 +779,9 @@ def main(detector=None, 'wf' : FaceType.WHOLE_FACE, 'head' : FaceType.HEAD}[face_type] + if max_faces_from_image is None: + max_faces_from_image = io.input_int(f"Max number of faces from image", 0, help_message="If you extract a src faceset that has frames with a large number of faces, it is advisable to set max faces to 3 to speed up extraction. 0 - unlimited") + if image_size is None: image_size = io.input_int(f"Image size", 512 if face_type < FaceType.HEAD else 768, valid_range=[256,2048], help_message="Output image size. The higher image size, the worse face-enhancer works. Use higher than 512 value only if the source image is sharp enough and the face does not need to be enhanced.") @@ -788,8 +799,6 @@ def main(detector=None, if output_debug is None: output_debug = io.input_bool (f"Write debug images to {output_debug_path.name}?", False) - input_path_image_paths = pathex.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info) - if manual_output_debug_fix: if not output_debug_path.exists(): io.log_err(f'{output_debug_path} not found. Re-extract faces with "Write debug images" option.') @@ -798,9 +807,9 @@ def main(detector=None, detector = 'manual' io.log_info('Performing re-extract frames which were deleted from _debug directory.') - input_path_image_paths = DeletedFilesSearcherSubprocessor (input_path_image_paths, pathex.get_image_paths(output_debug_path) ).run() - input_path_image_paths = sorted (input_path_image_paths) - io.log_info('Found %d images.' % (len(input_path_image_paths))) + input_image_paths = DeletedFilesSearcherSubprocessor (input_image_paths, pathex.get_image_paths(output_debug_path) ).run() + input_image_paths = sorted (input_image_paths) + io.log_info('Found %d images.' % (len(input_image_paths))) else: if output_debug_path.exists(): for filename in pathex.get_image_paths(output_debug_path): @@ -808,19 +817,19 @@ def main(detector=None, else: output_debug_path.mkdir(parents=True, exist_ok=True) - images_found = len(input_path_image_paths) + images_found = len(input_image_paths) faces_detected = 0 if images_found != 0: if detector == 'manual': io.log_info ('Performing manual extract...') - data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_path_image_paths ], 'landmarks-manual', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None, manual_window_size=manual_window_size, device_config=device_config).run() + data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_image_paths ], 'landmarks-manual', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None, manual_window_size=manual_window_size, device_config=device_config).run() io.log_info ('Performing 3rd pass...') data = ExtractSubprocessor (data, 'final', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None, final_output_path=output_path, device_config=device_config).run() else: io.log_info ('Extracting faces...') - data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_path_image_paths ], + data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_image_paths ], 'all', image_size, jpeg_quality,