github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 04:52:13 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix

Browse source
This commit is contained in:
iperov 2019-03-14 10:50:53 +04:00
parent 6e59f1a150
commit 5bd3f875fa
2 changed files with 68 additions and 63 deletions
Download patch file Download diff file Expand all files Collapse all files

85
main.py
View file

@ -37,27 +37,27 @@ if __name__ == "__main__":
} }
) )
extract_parser = subparsers.add_parser( "extract", help="Extract the faces from a pictures.") p = 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.") p.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.") p.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.") p.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'") 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'")
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.") 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.")
extract_parser.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.") p.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.") p.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.") 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.")
extract_parser.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.") p.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.") p.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.set_defaults (func=process_extract)
def process_sort(arguments): def process_sort(arguments):
from mainscripts import Sorter from mainscripts import Sorter
Sorter.main (input_path=arguments.input_dir, sort_by_method=arguments.sort_by_method) 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.") p = 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.") p.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." ) 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." )
sort_parser.set_defaults (func=process_sort) p.set_defaults (func=process_sort)
def process_util(arguments): def process_util(arguments):
from mainscripts import Util from mainscripts import Util
@ -71,12 +71,12 @@ if __name__ == "__main__":
if arguments.recover_original_aligned_filename: 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)
util_parser = subparsers.add_parser( "util", help="Utilities.") p = 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.") p.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.") p.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.") 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.")
util_parser.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.") p.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.set_defaults (func=process_util)
def process_train(arguments): def process_train(arguments):
args = {'training_data_src_dir' : arguments.training_data_src_dir, args = {'training_data_src_dir' : arguments.training_data_src_dir,
@ -92,17 +92,16 @@ if __name__ == "__main__":
from mainscripts import Trainer from mainscripts import Trainer
Trainer.main(args, device_args) Trainer.main(args, device_args)
train_parser = subparsers.add_parser( "train", help="Trainer") p = 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.") p.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.") p.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.") p.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") 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")
train_parser.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, help="Disable preview window.") p.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.") p.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.") p.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.") 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)
train_parser.set_defaults (func=process_train)
def process_convert(arguments): def process_convert(arguments):
args = {'input_dir' : arguments.input_dir, args = {'input_dir' : arguments.input_dir,
@ -118,16 +117,16 @@ if __name__ == "__main__":
from mainscripts import Converter from mainscripts import Converter
Converter.main (args, device_args) Converter.main (args, device_args)
convert_parser = subparsers.add_parser( "convert", help="Converter") p = 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.") p.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.") p.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.") 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.")
convert_parser.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.") p.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") 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")
convert_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug converter.") p.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.") p.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.") p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Convert on CPU.")
convert_parser.set_defaults(func=process_convert) p.set_defaults(func=process_convert)
videoed_parser = subparsers.add_parser( "videoed", help="Video processing.").add_subparsers() videoed_parser = subparsers.add_parser( "videoed", help="Video processing.").add_subparsers()

46
mainscripts/Sorter.py
View file

@ -509,6 +509,7 @@ class FinalLoaderSubprocessor(Subprocessor):
#override #override
def on_initialize(self, client_dict): def on_initialize(self, client_dict):
self.log_info ('Running on %s.' % (client_dict['device_name']) ) self.log_info ('Running on %s.' % (client_dict['device_name']) )
self.include_by_blur = client_dict['include_by_blur']
#override #override
def process_data(self, data): def process_data(self, data):
@ -529,8 +530,9 @@ class FinalLoaderSubprocessor(Subprocessor):
bgr = cv2_imread(str(filepath)) bgr = cv2_imread(str(filepath))
if bgr is None: if bgr is None:
raise Exception ("Unable to load %s" % (filepath.name) ) 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() ) pitch, yaw = LandmarksProcessor.estimate_pitch_yaw ( dflimg.get_landmarks() )
hist = cv2.calcHist([gray], [0], None, [256], [0, 256]) hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
@ -546,9 +548,10 @@ class FinalLoaderSubprocessor(Subprocessor):
return data[0] return data[0]
#override #override
def __init__(self, img_list ): def __init__(self, img_list, include_by_blur ):
self.img_list = img_list self.img_list = img_list
self.include_by_blur = include_by_blur
self.result = [] self.result = []
self.result_trash = [] self.result_trash = []
@ -566,7 +569,8 @@ class FinalLoaderSubprocessor(Subprocessor):
def process_info_generator(self): def process_info_generator(self):
for i in range(0, min(multiprocessing.cpu_count(), 8) ): for i in range(0, min(multiprocessing.cpu_count(), 8) ):
yield 'CPU%d' % (i), {}, {'device_idx': i, 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 #override
@ -592,12 +596,12 @@ class FinalLoaderSubprocessor(Subprocessor):
def get_result(self): def get_result(self):
return self.result, self.result_trash 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.") io.log_info ("Performing final sort.")
target_count = io.input_int ("Target number of images? (default:2000) : ", 2000) 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 = [] final_img_list = []
grads = 128 grads = 128
@ -631,22 +635,23 @@ def sort_final(input_path):
imgs_per_grad += total_lack // grads imgs_per_grad += total_lack // grads
sharpned_imgs_per_grad = imgs_per_grad*10 sharpned_imgs_per_grad = imgs_per_grad*10
for g in io.progress_bar_generator ( range (grads), "Sort by blur"): if include_by_blur:
img_list = yaws_sample_list[g] for g in io.progress_bar_generator ( range (grads), "Sort by blur"):
if img_list is None: img_list = yaws_sample_list[g]
continue if img_list is None:
continue
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
if len(img_list) > imgs_per_grad*2: if len(img_list) > imgs_per_grad*2:
trash_img_list += img_list[len(img_list) // 2:] trash_img_list += img_list[len(img_list) // 2:]
img_list = img_list[0: 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]
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"): for g in io.progress_bar_generator ( range (grads), "Sort by hist"):
img_list = yaws_sample_list[g] 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 == '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 == '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': 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) final_process (input_path, img_list, trash_img_list)