mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-06 04:52:13 -07:00
fix
This commit is contained in:
iperov
parent
6e59f1a150
commit
5bd3f875fa
2 changed files with 68 additions and 63 deletions
85
main.py
85
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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
@ -530,7 +531,8 @@ class FinalLoaderSubprocessor(Subprocessor):
|
|||
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)
|
||||
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) > 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]
|
||||
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
|
||||
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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue