diff --git a/main.py b/main.py index 0c7aeba..1921451 100644 --- a/main.py +++ b/main.py @@ -37,27 +37,27 @@ if __name__ == "__main__": } ) - extract_parser = subparsers.add_parser( "extract", help="Extract the faces from a pictures.") - extract_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") - extract_parser.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted files will be stored.") - extract_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Writes debug images to [output_dir]_debug\ directory.") - extract_parser.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'avatar', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'") - extract_parser.add_argument('--detector', dest="detector", choices=['dlib','mt','s3fd','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-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.") - extract_parser.set_defaults (func=process_extract) + p = subparsers.add_parser( "extract", help="Extract the faces from a pictures.") + 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('--debug', action="store_true", dest="debug", default=False, help="Writes debug images to [output_dir]_debug\ directory.") + p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'avatar', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'") + p.add_argument('--detector', dest="detector", choices=['dlib','mt','s3fd','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.") + p.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.") + 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. Forces to use MT extractor.") + p.set_defaults (func=process_extract) def process_sort(arguments): from mainscripts import Sorter Sorter.main (input_path=arguments.input_dir, sort_by_method=arguments.sort_by_method) - 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", "face-pitch", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final", "test"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." ) - sort_parser.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', required=True, dest="sort_by_method", choices=("blur", "face", "face-dissim", "face-yaw", "face-pitch", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final", "final-no-blur", "test"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." ) + p.set_defaults (func=process_sort) def process_util(arguments): from mainscripts import Util @@ -71,12 +71,12 @@ if __name__ == "__main__": if arguments.recover_original_aligned_filename: Util.recover_original_aligned_filename (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.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.") - util_parser.set_defaults (func=process_util) + 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('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.") + 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.set_defaults (func=process_util) def process_train(arguments): args = {'training_data_src_dir' : arguments.training_data_src_dir, @@ -92,17 +92,16 @@ if __name__ == "__main__": from mainscripts import Trainer Trainer.main(args, device_args) - train_parser = subparsers.add_parser( "train", help="Trainer") - train_parser.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", help="Dir of src-set.") - train_parser.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", help="Dir of dst-set.") - train_parser.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.") - train_parser.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model") - train_parser.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, help="Disable preview window.") - train_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.") - train_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.") - train_parser.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.") - - train_parser.set_defaults (func=process_train) + 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 src-set.") + p.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", help="Dir of dst-set.") + 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=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model") + p.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, help="Disable preview window.") + p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.") + p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.") + p.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.") + p.set_defaults (func=process_train) def process_convert(arguments): args = {'input_dir' : arguments.input_dir, @@ -118,16 +117,16 @@ if __name__ == "__main__": from mainscripts import Converter Converter.main (args, device_args) - convert_parser = subparsers.add_parser( "convert", help="Converter") - convert_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") - convert_parser.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the converted files will be stored.") - convert_parser.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored. Not used in AVATAR model.") - convert_parser.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.") - convert_parser.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model") - convert_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug converter.") - convert_parser.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.") - convert_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Convert on CPU.") - convert_parser.set_defaults(func=process_convert) + p = subparsers.add_parser( "convert", help="Converter") + 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 converted files will be stored.") + p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored. Not used in AVATAR model.") + 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=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model") + p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug converter.") + p.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.") + p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Convert on CPU.") + p.set_defaults(func=process_convert) videoed_parser = subparsers.add_parser( "videoed", help="Video processing.").add_subparsers() diff --git a/mainscripts/Sorter.py b/mainscripts/Sorter.py index 7c24ff0..2b4e0d8 100644 --- a/mainscripts/Sorter.py +++ b/mainscripts/Sorter.py @@ -509,6 +509,7 @@ class FinalLoaderSubprocessor(Subprocessor): #override def on_initialize(self, client_dict): self.log_info ('Running on %s.' % (client_dict['device_name']) ) + self.include_by_blur = client_dict['include_by_blur'] #override def process_data(self, data): @@ -529,8 +530,9 @@ class FinalLoaderSubprocessor(Subprocessor): bgr = cv2_imread(str(filepath)) if bgr is None: raise Exception ("Unable to load %s" % (filepath.name) ) - - sharpness = estimate_sharpness(bgr) + + gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) + sharpness = estimate_sharpness(gray) if self.include_by_blur else 0 pitch, yaw = LandmarksProcessor.estimate_pitch_yaw ( dflimg.get_landmarks() ) hist = cv2.calcHist([gray], [0], None, [256], [0, 256]) @@ -546,9 +548,10 @@ class FinalLoaderSubprocessor(Subprocessor): return data[0] #override - def __init__(self, img_list ): + def __init__(self, img_list, include_by_blur ): self.img_list = img_list + self.include_by_blur = include_by_blur self.result = [] self.result_trash = [] @@ -566,7 +569,8 @@ class FinalLoaderSubprocessor(Subprocessor): def process_info_generator(self): for i in range(0, min(multiprocessing.cpu_count(), 8) ): yield 'CPU%d' % (i), {}, {'device_idx': i, - 'device_name': 'CPU%d' % (i) + 'device_name': 'CPU%d' % (i), + 'include_by_blur': self.include_by_blur } #override @@ -592,12 +596,12 @@ class FinalLoaderSubprocessor(Subprocessor): def get_result(self): return self.result, self.result_trash -def sort_final(input_path): +def sort_final(input_path, include_by_blur=True): io.log_info ("Performing final sort.") target_count = io.input_int ("Target number of images? (default:2000) : ", 2000) - img_list, trash_img_list = FinalLoaderSubprocessor( Path_utils.get_image_paths(input_path) ).run() + img_list, trash_img_list = FinalLoaderSubprocessor( Path_utils.get_image_paths(input_path), include_by_blur ).run() final_img_list = [] grads = 128 @@ -631,22 +635,23 @@ def sort_final(input_path): imgs_per_grad += total_lack // grads sharpned_imgs_per_grad = imgs_per_grad*10 - for g in io.progress_bar_generator ( range (grads), "Sort by blur"): - img_list = yaws_sample_list[g] - if img_list is None: - continue + if include_by_blur: + for g in io.progress_bar_generator ( range (grads), "Sort by blur"): + img_list = yaws_sample_list[g] + if img_list is None: + continue - img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) - - if len(img_list) > imgs_per_grad*2: - trash_img_list += img_list[len(img_list) // 2:] - img_list = img_list[0: len(img_list) // 2] - - if len(img_list) > sharpned_imgs_per_grad: - trash_img_list += img_list[sharpned_imgs_per_grad:] - img_list = img_list[0:sharpned_imgs_per_grad] + img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) + + if len(img_list) > imgs_per_grad*2: + trash_img_list += img_list[len(img_list) // 2:] + img_list = img_list[0: len(img_list) // 2] - yaws_sample_list[g] = img_list + if len(img_list) > sharpned_imgs_per_grad: + trash_img_list += img_list[sharpned_imgs_per_grad:] + img_list = img_list[0:sharpned_imgs_per_grad] + + yaws_sample_list[g] = img_list for g in io.progress_bar_generator ( range (grads), "Sort by hist"): img_list = yaws_sample_list[g] @@ -737,5 +742,6 @@ def main (input_path, sort_by_method): elif sort_by_method == 'origname': img_list, trash_img_list = sort_by_origname (input_path) elif sort_by_method == 'oneface': img_list, trash_img_list = sort_by_oneface_in_image (input_path) elif sort_by_method == 'final': img_list, trash_img_list = sort_final (input_path) + elif sort_by_method == 'final-no-blur': img_list, trash_img_list = sort_final (input_path, include_by_blur=False) final_process (input_path, img_list, trash_img_list)