diff --git a/main.py b/main.py index b821910..d2ff3c0 100644 --- a/main.py +++ b/main.py @@ -1,9 +1,11 @@ if __name__ == "__main__": # Fix for linux import multiprocessing + multiprocessing.set_start_method("spawn") from core.leras import nn + nn.initialize_main_env() import os import sys @@ -18,341 +20,463 @@ if __name__ == "__main__": if sys.version_info[0] < 3 or (sys.version_info[0] == 3 and sys.version_info[1] < 6): raise Exception("This program requires at least Python 3.6") + class fixPathAction(argparse.Action): def __call__(self, parser, namespace, values, option_string=None): setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values))) + exit_code = 0 - + parser = argparse.ArgumentParser() subparsers = parser.add_subparsers() + def process_extract(arguments): osex.set_process_lowest_prio() from mainscripts import Extractor - Extractor.main( detector = arguments.detector, - input_path = Path(arguments.input_dir), - output_path = Path(arguments.output_dir), - output_debug = arguments.output_debug, - manual_fix = arguments.manual_fix, - 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, - force_gpu_idxs = [ int(x) for x in arguments.force_gpu_idxs.split(',') ] if arguments.force_gpu_idxs is not None else None, - ) + Extractor.main(detector=arguments.detector, + input_path=Path(arguments.input_dir), + output_path=Path(arguments.output_dir), + output_debug=arguments.output_debug, + manual_fix=arguments.manual_fix, + 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, + force_gpu_idxs=[int(x) for x in arguments.force_gpu_idxs.split( + ',')] if arguments.force_gpu_idxs is not None else None, + ) - p = subparsers.add_parser( "extract", help="Extract the faces from a pictures.") - p.add_argument('--detector', dest="detector", choices=['s3fd','manual'], default=None, help="Type of detector.") - p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") - 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', '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 = subparsers.add_parser("extract", help="Extract the faces from a pictures.") + p.add_argument('--detector', dest="detector", choices=['s3fd', 'manual'], default=None, help="Type of detector.") + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", + help="Input directory. A directory containing the files you wish to process.") + 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', '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.") - 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.") + 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.") + 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.") p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU..") - p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.") + p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, + help="Force to choose GPU indexes separated by comma.") + + p.set_defaults(func=process_extract) - p.set_defaults (func=process_extract) def process_sort(arguments): osex.set_process_lowest_prio() from mainscripts import Sorter - Sorter.main (input_path=Path(arguments.input_dir), sort_by_method=arguments.sort_by_method) + Sorter.main(input_path=Path(arguments.input_dir), sort_by_method=arguments.sort_by_method) + + + p = subparsers.add_parser("sort", help="Sort faces in a directory.") + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", + help="Input directory. A directory containing the files you wish to process.") + p.add_argument('--by', dest="sort_by_method", default=None, choices=( + "blur", "motion-blur", "face-yaw", "face-pitch", "face-source-rect-size", "hist", "hist-dissim", "brightness", + "hue", "black", "origname", "oneface", "final-by-blur", "final-by-size", "absdiff"), + help="Method of sorting. 'origname' sort by original filename to recover original sequence.") + p.set_defaults(func=process_sort) - p = subparsers.add_parser( "sort", help="Sort faces in a directory.") - p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") - p.add_argument('--by', dest="sort_by_method", default=None, choices=("blur", "motion-blur", "face-yaw", "face-pitch", "face-source-rect-size", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final-by-blur", "final-by-size", "absdiff"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." ) - p.set_defaults (func=process_sort) def process_util(arguments): osex.set_process_lowest_prio() from mainscripts import Util if arguments.add_landmarks_debug_images: - Util.add_landmarks_debug_images (input_path=arguments.input_dir) + Util.add_landmarks_debug_images(input_path=arguments.input_dir) if arguments.recover_original_aligned_filename: - Util.recover_original_aligned_filename (input_path=arguments.input_dir) + Util.recover_original_aligned_filename(input_path=arguments.input_dir) if arguments.save_faceset_metadata: - Util.save_faceset_metadata_folder (input_path=arguments.input_dir) + Util.save_faceset_metadata_folder(input_path=arguments.input_dir) if arguments.restore_faceset_metadata: - Util.restore_faceset_metadata_folder (input_path=arguments.input_dir) + Util.restore_faceset_metadata_folder(input_path=arguments.input_dir) if arguments.pack_faceset: - io.log_info ("Performing faceset packing...\r\n") + io.log_info("Performing faceset packing...\r\n") from samplelib import PackedFaceset - PackedFaceset.pack( Path(arguments.input_dir) ) + PackedFaceset.pack(Path(arguments.input_dir)) if arguments.unpack_faceset: - io.log_info ("Performing faceset unpacking...\r\n") + io.log_info("Performing faceset unpacking...\r\n") from samplelib import PackedFaceset - PackedFaceset.unpack( Path(arguments.input_dir) ) - - if arguments.export_faceset_mask: - io.log_info ("Exporting faceset mask..\r\n") - Util.export_faceset_mask( Path(arguments.input_dir) ) + PackedFaceset.unpack(Path(arguments.input_dir)) - p = subparsers.add_parser( "util", help="Utilities.") - p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") - p.add_argument('--add-landmarks-debug-images', action="store_true", dest="add_landmarks_debug_images", default=False, help="Add landmarks debug image for aligned faces.") - p.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.") - p.add_argument('--save-faceset-metadata', action="store_true", dest="save_faceset_metadata", default=False, help="Save faceset metadata to file.") - p.add_argument('--restore-faceset-metadata', action="store_true", dest="restore_faceset_metadata", default=False, help="Restore faceset metadata to file. Image filenames must be the same as used with save.") + if arguments.export_faceset_mask: + io.log_info("Exporting faceset mask..\r\n") + Util.export_faceset_mask(Path(arguments.input_dir)) + + + p = subparsers.add_parser("util", help="Utilities.") + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", + help="Input directory. A directory containing the files you wish to process.") + p.add_argument('--add-landmarks-debug-images', action="store_true", dest="add_landmarks_debug_images", + default=False, help="Add landmarks debug image for aligned faces.") + p.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", + default=False, help="Recover original aligned filename.") + p.add_argument('--save-faceset-metadata', action="store_true", dest="save_faceset_metadata", default=False, + help="Save faceset metadata to file.") + p.add_argument('--restore-faceset-metadata', action="store_true", dest="restore_faceset_metadata", default=False, + help="Restore faceset metadata to file. Image filenames must be the same as used with save.") p.add_argument('--pack-faceset', action="store_true", dest="pack_faceset", default=False, help="") p.add_argument('--unpack-faceset', action="store_true", dest="unpack_faceset", default=False, help="") p.add_argument('--export-faceset-mask', action="store_true", dest="export_faceset_mask", default=False, help="") - p.set_defaults (func=process_util) + p.set_defaults(func=process_util) + def process_train(arguments): osex.set_process_lowest_prio() - - kwargs = {'model_class_name' : arguments.model_name, - 'saved_models_path' : Path(arguments.model_dir), - 'training_data_src_path' : Path(arguments.training_data_src_dir), - 'training_data_dst_path' : Path(arguments.training_data_dst_dir), - 'pretraining_data_path' : Path(arguments.pretraining_data_dir) if arguments.pretraining_data_dir is not None else None, - 'pretrained_model_path' : Path(arguments.pretrained_model_dir) if arguments.pretrained_model_dir is not None else None, - 'no_preview' : arguments.no_preview, - 'force_model_name' : arguments.force_model_name, - 'force_gpu_idxs' : [ int(x) for x in arguments.force_gpu_idxs.split(',') ] if arguments.force_gpu_idxs is not None else None, - 'cpu_only' : arguments.cpu_only, - 'silent_start' : arguments.silent_start, - 'execute_programs' : [ [int(x[0]), x[1] ] for x in arguments.execute_program ], - 'debug' : arguments.debug, + kwargs = {'model_class_name': arguments.model_name, + 'saved_models_path': Path(arguments.model_dir), + 'training_data_src_path': Path(arguments.training_data_src_dir), + 'training_data_dst_path': Path(arguments.training_data_dst_dir), + 'pretraining_data_path': Path( + arguments.pretraining_data_dir) if arguments.pretraining_data_dir is not None else None, + 'pretrained_model_path': Path( + arguments.pretrained_model_dir) if arguments.pretrained_model_dir is not None else None, + 'no_preview': arguments.no_preview, + 'force_model_name': arguments.force_model_name, + 'force_gpu_idxs': [int(x) for x in arguments.force_gpu_idxs.split( + ',')] if arguments.force_gpu_idxs is not None else None, + 'cpu_only': arguments.cpu_only, + 'silent_start': arguments.silent_start, + 'execute_programs': [[int(x[0]), x[1]] for x in arguments.execute_program], + 'debug': arguments.debug, } from mainscripts import Trainer Trainer.main(**kwargs) - p = subparsers.add_parser( "train", help="Trainer") - p.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", help="Dir of extracted SRC faceset.") - p.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", help="Dir of extracted DST faceset.") - p.add_argument('--pretraining-data-dir', action=fixPathAction, dest="pretraining_data_dir", default=None, help="Optional dir of extracted faceset that will be used in pretraining mode.") - p.add_argument('--pretrained-model-dir', action=fixPathAction, dest="pretrained_model_dir", default=None, help="Optional dir of pretrain model files. (Currently only for Quick96).") + + p = subparsers.add_parser("train", help="Trainer") + p.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", + help="Dir of extracted SRC faceset.") + p.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", + help="Dir of extracted DST faceset.") + p.add_argument('--pretraining-data-dir', action=fixPathAction, dest="pretraining_data_dir", default=None, + help="Optional dir of extracted faceset that will be used in pretraining mode.") + p.add_argument('--pretrained-model-dir', action=fixPathAction, dest="pretrained_model_dir", default=None, + help="Optional dir of pretrain model files. (Currently only for Quick96).") p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Saved models dir.") - p.add_argument('--model', required=True, dest="model_name", choices=pathex.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Model class name.") + p.add_argument('--model', required=True, dest="model_name", + choices=pathex.get_all_dir_names_startswith(Path(__file__).parent / 'models', 'Model_'), + help="Model class name.") p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.") - p.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, help="Disable preview window.") - p.add_argument('--force-model-name', dest="force_model_name", default=None, help="Forcing to choose model name from model/ folder.") + p.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, + help="Disable preview window.") + p.add_argument('--force-model-name', dest="force_model_name", default=None, + help="Forcing to choose model name from model/ folder.") p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.") - p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.") - p.add_argument('--silent-start', action="store_true", dest="silent_start", default=False, help="Silent start. Automatically chooses Best GPU and last used model.") - + p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, + help="Force to choose GPU indexes separated by comma.") + p.add_argument('--silent-start', action="store_true", dest="silent_start", default=False, + help="Silent start. Automatically chooses Best GPU and last used model.") + p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+') - p.set_defaults (func=process_train) - + p.set_defaults(func=process_train) + + def process_exportdfm(arguments): osex.set_process_lowest_prio() from mainscripts import ExportDFM - ExportDFM.main(model_class_name = arguments.model_name, saved_models_path = Path(arguments.model_dir)) + ExportDFM.main(model_class_name=arguments.model_name, saved_models_path=Path(arguments.model_dir)) - p = subparsers.add_parser( "exportdfm", help="Export model to use in DeepFaceLive.") + + p = subparsers.add_parser("exportdfm", help="Export model to use in DeepFaceLive.") p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Saved models dir.") - p.add_argument('--model', required=True, dest="model_name", choices=pathex.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Model class name.") - p.set_defaults (func=process_exportdfm) + p.add_argument('--model', required=True, dest="model_name", + choices=pathex.get_all_dir_names_startswith(Path(__file__).parent / 'models', 'Model_'), + help="Model class name.") + p.set_defaults(func=process_exportdfm) + def process_merge(arguments): osex.set_process_lowest_prio() from mainscripts import Merger - Merger.main ( model_class_name = arguments.model_name, - saved_models_path = Path(arguments.model_dir), - force_model_name = arguments.force_model_name, - input_path = Path(arguments.input_dir), - output_path = Path(arguments.output_dir), - output_mask_path = Path(arguments.output_mask_dir), - aligned_path = Path(arguments.aligned_dir) if arguments.aligned_dir is not None else None, - force_gpu_idxs = arguments.force_gpu_idxs, - cpu_only = arguments.cpu_only) + Merger.main(model_class_name=arguments.model_name, + saved_models_path=Path(arguments.model_dir), + force_model_name=arguments.force_model_name, + input_path=Path(arguments.input_dir), + output_path=Path(arguments.output_dir), + output_mask_path=Path(arguments.output_mask_dir), + aligned_path=Path(arguments.aligned_dir) if arguments.aligned_dir is not None else None, + force_gpu_idxs=arguments.force_gpu_idxs, + cpu_only=arguments.cpu_only, + is_interactive=arguments.is_interactive, + config=arguments.config, + subprocess_count=arguments.subprocess_count) - p = subparsers.add_parser( "merge", help="Merger") - p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") - p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the merged files will be stored.") - p.add_argument('--output-mask-dir', required=True, action=fixPathAction, dest="output_mask_dir", help="Output mask directory. This is where the mask files will be stored.") - p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", default=None, help="Aligned directory. This is where the extracted of dst faces stored.") + + p = subparsers.add_parser("merge", help="Merger") + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", + help="Input directory. A directory containing the files you wish to process.") + p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", + help="Output directory. This is where the merged files will be stored.") + p.add_argument('--output-mask-dir', required=True, action=fixPathAction, dest="output_mask_dir", + help="Output mask directory. This is where the mask files will be stored.") + p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", default=None, + help="Aligned directory. This is where the extracted of dst faces stored.") p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.") - p.add_argument('--model', required=True, dest="model_name", choices=pathex.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Model class name.") - p.add_argument('--force-model-name', dest="force_model_name", default=None, help="Forcing to choose model name from model/ folder.") + p.add_argument('--model', required=True, dest="model_name", + choices=pathex.get_all_dir_names_startswith(Path(__file__).parent / 'models', 'Model_'), + help="Model class name.") + p.add_argument('--force-model-name', dest="force_model_name", default=None, + help="Forcing to choose model name from model/ folder.") p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Merge on CPU.") - p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.") + p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, + help="Force to choose GPU indexes separated by comma.") + p.add_argument('--is_interactive', dest="is_interactive", default=None, + help="Forcing to interactivate mode or not.") + p.add_argument('--config', dest="config", default=None, + help="MergerConfig object can be set outside by script") + p.add_argument('--subprocess_count', type=int, dest="subprocess_count", default=-1, + help="Specify the number of threads to process. A low value may affect performance. A high value may result in memory error. The value may not be greater than CPU cores.") p.set_defaults(func=process_merge) - videoed_parser = subparsers.add_parser( "videoed", help="Video processing.").add_subparsers() + videoed_parser = subparsers.add_parser("videoed", help="Video processing.").add_subparsers() + def process_videoed_extract_video(arguments): osex.set_process_lowest_prio() from mainscripts import VideoEd - VideoEd.extract_video (arguments.input_file, arguments.output_dir, arguments.output_ext, arguments.fps) - p = videoed_parser.add_parser( "extract-video", help="Extract images from video file.") - p.add_argument('--input-file', required=True, action=fixPathAction, dest="input_file", help="Input file to be processed. Specify .*-extension to find first file.") - p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted images will be stored.") + VideoEd.extract_video(arguments.input_file, arguments.output_dir, arguments.output_ext, arguments.fps) + + + p = videoed_parser.add_parser("extract-video", help="Extract images from video file.") + p.add_argument('--input-file', required=True, action=fixPathAction, dest="input_file", + help="Input file to be processed. Specify .*-extension to find first file.") + p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", + help="Output directory. This is where the extracted images will be stored.") p.add_argument('--output-ext', dest="output_ext", default=None, help="Image format (extension) of output files.") - p.add_argument('--fps', type=int, dest="fps", default=None, help="How many frames of every second of the video will be extracted. 0 - full fps.") + p.add_argument('--fps', type=int, dest="fps", default=None, + help="How many frames of every second of the video will be extracted. 0 - full fps.") p.set_defaults(func=process_videoed_extract_video) + def process_videoed_cut_video(arguments): osex.set_process_lowest_prio() from mainscripts import VideoEd - VideoEd.cut_video (arguments.input_file, - arguments.from_time, - arguments.to_time, - arguments.audio_track_id, - arguments.bitrate) - p = videoed_parser.add_parser( "cut-video", help="Cut video file.") - p.add_argument('--input-file', required=True, action=fixPathAction, dest="input_file", help="Input file to be processed. Specify .*-extension to find first file.") + VideoEd.cut_video(arguments.input_file, + arguments.from_time, + arguments.to_time, + arguments.audio_track_id, + arguments.bitrate) + + + p = videoed_parser.add_parser("cut-video", help="Cut video file.") + p.add_argument('--input-file', required=True, action=fixPathAction, dest="input_file", + help="Input file to be processed. Specify .*-extension to find first file.") p.add_argument('--from-time', dest="from_time", default=None, help="From time, for example 00:00:00.000") p.add_argument('--to-time', dest="to_time", default=None, help="To time, for example 00:00:00.000") p.add_argument('--audio-track-id', type=int, dest="audio_track_id", default=None, help="Specify audio track id.") p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.") p.set_defaults(func=process_videoed_cut_video) + def process_videoed_denoise_image_sequence(arguments): osex.set_process_lowest_prio() from mainscripts import VideoEd - VideoEd.denoise_image_sequence (arguments.input_dir, arguments.factor) - p = videoed_parser.add_parser( "denoise-image-sequence", help="Denoise sequence of images, keeping sharp edges. Helps to remove pixel shake from the predicted face.") - p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory to be processed.") + VideoEd.denoise_image_sequence(arguments.input_dir, None, arguments.factor) + + + p = videoed_parser.add_parser("denoise-image-sequence", + help="Denoise sequence of images, keeping sharp edges. Helps to remove pixel shake from the predicted face.") + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", + help="Input directory to be processed.") p.add_argument('--factor', type=int, dest="factor", default=None, help="Denoise factor (1-20).") p.set_defaults(func=process_videoed_denoise_image_sequence) + def process_videoed_video_from_sequence(arguments): osex.set_process_lowest_prio() from mainscripts import VideoEd - VideoEd.video_from_sequence (input_dir = arguments.input_dir, - output_file = arguments.output_file, - reference_file = arguments.reference_file, - ext = arguments.ext, - fps = arguments.fps, - bitrate = arguments.bitrate, - include_audio = arguments.include_audio, - lossless = arguments.lossless) + VideoEd.video_from_sequence(input_dir=arguments.input_dir, + output_file=arguments.output_file, + reference_file=arguments.reference_file, + ext=arguments.ext, + fps=arguments.fps, + bitrate=arguments.bitrate, + include_audio=arguments.include_audio, + lossless=arguments.lossless) - p = videoed_parser.add_parser( "video-from-sequence", help="Make video from image sequence.") - p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input file to be processed. Specify .*-extension to find first file.") - p.add_argument('--output-file', required=True, action=fixPathAction, dest="output_file", help="Input file to be processed. Specify .*-extension to find first file.") - p.add_argument('--reference-file', action=fixPathAction, dest="reference_file", help="Reference file used to determine proper FPS and transfer audio from it. Specify .*-extension to find first file.") + + p = videoed_parser.add_parser("video-from-sequence", help="Make video from image sequence.") + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", + help="Input file to be processed. Specify .*-extension to find first file.") + p.add_argument('--output-file', required=True, action=fixPathAction, dest="output_file", + help="Input file to be processed. Specify .*-extension to find first file.") + p.add_argument('--reference-file', action=fixPathAction, dest="reference_file", + help="Reference file used to determine proper FPS and transfer audio from it. Specify .*-extension to find first file.") p.add_argument('--ext', dest="ext", default='png', help="Image format (extension) of input files.") - p.add_argument('--fps', type=int, dest="fps", default=None, help="FPS of output file. Overwritten by reference-file.") + p.add_argument('--fps', type=int, dest="fps", default=None, + help="FPS of output file. Overwritten by reference-file.") p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.") - p.add_argument('--include-audio', action="store_true", dest="include_audio", default=False, help="Include audio from reference file.") + p.add_argument('--include-audio', action="store_true", dest="include_audio", default=False, + help="Include audio from reference file.") p.add_argument('--lossless', action="store_true", dest="lossless", default=False, help="PNG codec.") p.set_defaults(func=process_videoed_video_from_sequence) - facesettool_parser = subparsers.add_parser( "facesettool", help="Faceset tools.").add_subparsers() + facesettool_parser = subparsers.add_parser("facesettool", help="Faceset tools.").add_subparsers() + def process_faceset_enhancer(arguments): osex.set_process_lowest_prio() from mainscripts import FacesetEnhancer - FacesetEnhancer.process_folder ( Path(arguments.input_dir), - cpu_only=arguments.cpu_only, - force_gpu_idxs=arguments.force_gpu_idxs + FacesetEnhancer.process_folder(Path(arguments.input_dir), + cpu_only=arguments.cpu_only, + force_gpu_idxs=arguments.force_gpu_idxs, + is_merge=arguments.is_merge ) - p = facesettool_parser.add_parser ("enhance", help="Enhance details in DFL faceset.") - p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.") + + p = facesettool_parser.add_parser("enhance", help="Enhance details in DFL faceset.") + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", + help="Input directory of aligned faces.") p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Process on CPU.") - p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.") + p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, + help="Force to choose GPU indexes separated by comma.") + p.add_argument('--is_merge', dest="is_merge", default=None, + help="Force to merge enhanced into aligned") p.set_defaults(func=process_faceset_enhancer) - - - p = facesettool_parser.add_parser ("resize", help="Resize DFL faceset.") - p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.") + + p = facesettool_parser.add_parser("resize", help="Resize DFL faceset.") + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", + help="Input directory of aligned faces.") + p.add_argument('--image-size', type=int, dest="image_size", default=512, help="Output image size.") + p.add_argument('--face-type', dest="face_type", + choices=['h','mf','f','wf','head','same'], default='same') + p.add_argument('--is_merge', dest="is_merge", default=None, + help="Force to merge enhanced into aligned") def process_faceset_resizer(arguments): osex.set_process_lowest_prio() from mainscripts import FacesetResizer - FacesetResizer.process_folder ( Path(arguments.input_dir) ) + FacesetResizer.process_folder(Path(arguments.input_dir), + arguments.image_size, arguments.face_type, arguments.is_merge) + + p.set_defaults(func=process_faceset_resizer) + def process_dev_test(arguments): osex.set_process_lowest_prio() from mainscripts import dev_misc - dev_misc.dev_gen_mask_files( arguments.input_dir ) + dev_misc.dev_gen_mask_files(arguments.input_dir) - p = subparsers.add_parser( "dev_test", help="") + + p = subparsers.add_parser("dev_test", help="") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") - p.set_defaults (func=process_dev_test) - + p.set_defaults(func=process_dev_test) + # ========== XSeg - xseg_parser = subparsers.add_parser( "xseg", help="XSeg tools.").add_subparsers() - - p = xseg_parser.add_parser( "editor", help="XSeg editor.") + xseg_parser = subparsers.add_parser("xseg", help="XSeg tools.").add_subparsers() + + p = xseg_parser.add_parser("editor", help="XSeg editor.") + def process_xsegeditor(arguments): osex.set_process_lowest_prio() from XSegEditor import XSegEditor global exit_code - exit_code = XSegEditor.start (Path(arguments.input_dir)) - + exit_code = XSegEditor.start(Path(arguments.input_dir)) + + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") - p.set_defaults (func=process_xsegeditor) - - p = xseg_parser.add_parser( "apply", help="Apply trained XSeg model to the extracted faces.") + p.set_defaults(func=process_xsegeditor) + + p = xseg_parser.add_parser("apply", help="Apply trained XSeg model to the extracted faces.") + def process_xsegapply(arguments): osex.set_process_lowest_prio() from mainscripts import XSegUtil - XSegUtil.apply_xseg (Path(arguments.input_dir), Path(arguments.model_dir)) + XSegUtil.apply_xseg(Path(arguments.input_dir), Path(arguments.model_dir), arguments.force_gpu_idxs) + + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir") - p.set_defaults (func=process_xsegapply) - - - p = xseg_parser.add_parser( "remove", help="Remove applied XSeg masks from the extracted faces.") + p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, + help="Force to choose GPU indexes separated by comma.") + p.set_defaults(func=process_xsegapply) + + p = xseg_parser.add_parser("remove", help="Remove applied XSeg masks from the extracted faces.") + + def process_xsegremove(arguments): osex.set_process_lowest_prio() from mainscripts import XSegUtil - XSegUtil.remove_xseg (Path(arguments.input_dir) ) + XSegUtil.remove_xseg(Path(arguments.input_dir)) + + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") - p.set_defaults (func=process_xsegremove) - - - p = xseg_parser.add_parser( "remove_labels", help="Remove XSeg labels from the extracted faces.") + p.set_defaults(func=process_xsegremove) + + p = xseg_parser.add_parser("remove_labels", help="Remove XSeg labels from the extracted faces.") + + def process_xsegremovelabels(arguments): osex.set_process_lowest_prio() from mainscripts import XSegUtil - XSegUtil.remove_xseg_labels (Path(arguments.input_dir) ) + XSegUtil.remove_xseg_labels(Path(arguments.input_dir)) + + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") - p.set_defaults (func=process_xsegremovelabels) - - - p = xseg_parser.add_parser( "fetch", help="Copies faces containing XSeg polygons in _xseg dir.") + p.set_defaults(func=process_xsegremovelabels) + + p = xseg_parser.add_parser("fetch", help="Copies faces containing XSeg polygons in _xseg dir.") + def process_xsegfetch(arguments): osex.set_process_lowest_prio() from mainscripts import XSegUtil - XSegUtil.fetch_xseg (Path(arguments.input_dir) ) + XSegUtil.fetch_xseg(Path(arguments.input_dir)) + + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") - p.set_defaults (func=process_xsegfetch) - + p.set_defaults(func=process_xsegfetch) + + def bad_args(arguments): parser.print_help() exit(0) + + parser.set_defaults(func=bad_args) arguments = parser.parse_args() arguments.func(arguments) if exit_code == 0: - print ("Done.") - + print("Done.") + exit(exit_code) - + ''' import code code.interact(local=dict(globals(), **locals())) -''' +''' \ No newline at end of file diff --git a/mainscripts/FacesetEnhancer.py b/mainscripts/FacesetEnhancer.py index 3de9cea..0d4bc97 100644 --- a/mainscripts/FacesetEnhancer.py +++ b/mainscripts/FacesetEnhancer.py @@ -122,7 +122,7 @@ class FacesetEnhancerSubprocessor(Subprocessor): return (0, filepath, None) -def process_folder ( dirpath, cpu_only=False, force_gpu_idxs=None ): +def process_folder ( dirpath, cpu_only=False, force_gpu_idxs=None, is_merge=None ): device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(suggest_all_gpu=True) ) \ if not cpu_only else nn.DeviceConfig.CPU() @@ -142,7 +142,9 @@ def process_folder ( dirpath, cpu_only=False, force_gpu_idxs=None ): image_paths = [Path(x) for x in pathex.get_image_paths( dirpath )] result = FacesetEnhancerSubprocessor ( image_paths, output_dirpath, device_config=device_config).run() - is_merge = io.input_bool (f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ?", True) + if is_merge is None: + is_merge = io.input_bool (f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ?", True) + if is_merge: io.log_info (f"Copying processed files to {dirpath_parts}") diff --git a/mainscripts/FacesetResizer.py b/mainscripts/FacesetResizer.py index 4bcd1b8..d824eea 100644 --- a/mainscripts/FacesetResizer.py +++ b/mainscripts/FacesetResizer.py @@ -164,11 +164,14 @@ class FacesetResizerSubprocessor(Subprocessor): return (0, filepath, None) -def process_folder ( dirpath): - - image_size = io.input_int(f"New image size", 512, valid_range=[128,2048]) - - face_type = io.input_str ("Change face type", 'same', ['h','mf','f','wf','head','same']).lower() +def process_folder ( dirpath, image_size=None, face_type=None, is_merge=None): + + if not image_size: + image_size = io.input_int(f"New image size", 512, valid_range=[128,2048]) + + if not face_type: + face_type = io.input_str ("Change face type", 'same', ['h','mf','f','wf','head','same']).lower() + if face_type == 'same': face_type = None else: @@ -195,7 +198,9 @@ def process_folder ( dirpath): image_paths = [Path(x) for x in pathex.get_image_paths( dirpath )] result = FacesetResizerSubprocessor ( image_paths, output_dirpath, image_size, face_type).run() - is_merge = io.input_bool (f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ?", True) + if not is_merge: + is_merge = io.input_bool (f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ?", True) + if is_merge: io.log_info (f"Copying processed files to {dirpath_parts}") diff --git a/mainscripts/Merger.py b/mainscripts/Merger.py index 0703dc1..9ba5483 100644 --- a/mainscripts/Merger.py +++ b/mainscripts/Merger.py @@ -26,7 +26,10 @@ def main (model_class_name=None, output_mask_path=None, aligned_path=None, force_gpu_idxs=None, - cpu_only=None): + cpu_only=None, + is_interactive=None, + config=None, + subprocess_count=0): io.log_info ("Running merger.\r\n") try: @@ -69,12 +72,16 @@ def main (model_class_name=None, place_model_on_cpu=True, run_on_cpu=run_on_cpu) - is_interactive = io.input_bool ("Use interactive merger?", True) if not io.is_colab() else False + if is_interactive is None: + is_interactive = io.input_bool ("Use interactive merger?", True) if not io.is_colab() else False - if not is_interactive: + if not is_interactive and not config: cfg.ask_settings() - - subprocess_count = io.input_int("Number of workers?", max(8, multiprocessing.cpu_count()), + else: + cfg = config + + if subprocess_count <= 0: + subprocess_count = io.input_int("Number of workers?", max(8, multiprocessing.cpu_count()), valid_range=[1, multiprocessing.cpu_count()], help_message="Specify the number of threads to process. A low value may affect performance. A high value may result in memory error. The value may not be greater than CPU cores." ) input_path_image_paths = pathex.get_image_paths(input_path) diff --git a/mainscripts/XSegUtil.py b/mainscripts/XSegUtil.py index c75a14a..4e84c8f 100644 --- a/mainscripts/XSegUtil.py +++ b/mainscripts/XSegUtil.py @@ -13,7 +13,7 @@ from DFLIMG import * from facelib import XSegNet, LandmarksProcessor, FaceType import pickle -def apply_xseg(input_path, model_path): +def apply_xseg(input_path, model_path, force_gpu_idxs=None): if not input_path.exists(): raise ValueError(f'{input_path} not found. Please ensure it exists.') @@ -45,11 +45,11 @@ def apply_xseg(input_path, model_path): io.log_info(f'Applying trained XSeg model to {input_path.name}/ folder.') - device_config = nn.DeviceConfig.ask_choose_device(choose_only_one=True) + device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(choose_only_one=True) ) nn.initialize(device_config) - - - + + + xseg = XSegNet(name='XSeg', load_weights=True, weights_file_root=model_path,