diff --git a/main.py b/main.py index 37358b7..ce1540c 100644 --- a/main.py +++ b/main.py @@ -46,7 +46,7 @@ if __name__ == "__main__": p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted files will be stored.") 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', 'full_face_no_align', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'") + p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'whole_face', 'head', 'mark_only'], default=None) p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.") p.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.") p.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.") diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index 2a14453..908e723 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -690,14 +690,30 @@ def main(detector=None, cpu_only = False, force_gpu_idxs = None, ): - face_type = FaceType.fromString(face_type) - - image_size = 512 if face_type < FaceType.HEAD else 768 - + if not input_path.exists(): io.log_err ('Input directory not found. Please ensure it exists.') return + 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(): + 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 face_type is None: + face_type = io.input_str ("Face type", 'wf', ['f','wf','head'], help_message="Full face / whole face / head. 'Whole face' covers full area of face include forehead. 'head' covers full head, but requires XSeg for src and dst faceset.").lower() + face_type = {'f' : FaceType.FULL, + 'wf' : FaceType.WHOLE_FACE, + 'head' : FaceType.HEAD}[face_type] + + image_size = 512 if face_type < FaceType.HEAD else 768 + if detector is None: io.log_info ("Choose detector type.") io.log_info ("[0] S3FD")