DFL-2.0 initial branch commit

This commit is contained in:
Colombo 2020-01-21 18:43:39 +04:00
parent 52a67a61b3
commit 38b85108b3
154 changed files with 5251 additions and 9414 deletions

178
main.py
View file

@ -1,14 +1,17 @@
if __name__ == "__main__":
from core.leras import nn
nn.initialize_main_env()
import os
import sys
import time
import argparse
import multiprocessing
multiprocessing.set_start_method("spawn")
from utils import Path_utils
from utils import os_utils
from core import pathex
from core import osex
from pathlib import Path
from interact import interact as io
from core.interact import interact as io
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")
@ -21,36 +24,37 @@ if __name__ == "__main__":
subparsers = parser.add_subparsers()
def process_extract(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import Extractor
Extractor.main( arguments.input_dir,
arguments.output_dir,
arguments.debug_dir,
arguments.detector,
arguments.manual_fix,
arguments.manual_output_debug_fix,
arguments.manual_window_size,
face_type=arguments.face_type,
device_args={'cpu_only' : arguments.cpu_only,
'multi_gpu' : arguments.multi_gpu,
}
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,
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('--debug-dir', action=fixPathAction, dest="debug_dir", help="Writes debug images to this directory.")
p.add_argument('--output-debug', action="store_true", dest="output_debug", default=None, help="Writes debug images to <output-dir>_debug\ directory.")
p.add_argument('--no-output-debug', action="store_false", dest="output_debug", default=None, help="Don't writes debug images to <output-dir>_debug\ directory.")
p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'full_face_no_align', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'")
p.add_argument('--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.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.set_defaults (func=process_extract)
def process_dev_extract_vggface2_dataset(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import dev_misc
dev_misc.extract_vggface2_dataset( arguments.input_dir,
device_args={'cpu_only' : arguments.cpu_only,
@ -65,7 +69,7 @@ if __name__ == "__main__":
p.set_defaults (func=process_dev_extract_vggface2_dataset)
def process_dev_extract_umd_csv(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import dev_misc
dev_misc.extract_umd_csv( arguments.input_csv_file,
device_args={'cpu_only' : arguments.cpu_only,
@ -81,7 +85,7 @@ if __name__ == "__main__":
def process_dev_apply_celebamaskhq(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import dev_misc
dev_misc.apply_celebamaskhq( arguments.input_dir )
@ -90,7 +94,7 @@ if __name__ == "__main__":
p.set_defaults (func=process_dev_apply_celebamaskhq)
def process_dev_test(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import dev_misc
dev_misc.dev_test( arguments.input_dir )
@ -99,17 +103,17 @@ if __name__ == "__main__":
p.set_defaults (func=process_dev_test)
def process_sort(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import Sorter
Sorter.main (input_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', 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", "vggface", "absdiff", "test"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
p.add_argument('--by', dest="sort_by_method", default=None, choices=("blur", "face-yaw", "face-pitch", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final", "absdiff"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
p.set_defaults (func=process_sort)
def process_util(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import Util
if arguments.convert_png_to_jpg:
@ -158,65 +162,71 @@ if __name__ == "__main__":
p.set_defaults (func=process_util)
def process_train(arguments):
os_utils.set_process_lowest_prio()
args = {'training_data_src_dir' : arguments.training_data_src_dir,
'training_data_dst_dir' : arguments.training_data_dst_dir,
'pretraining_data_dir' : arguments.pretraining_data_dir,
'model_path' : arguments.model_dir,
'model_name' : arguments.model_name,
'no_preview' : arguments.no_preview,
'debug' : arguments.debug,
'execute_programs' : [ [int(x[0]), x[1] ] for x in arguments.execute_program ]
}
device_args = {'cpu_only' : arguments.cpu_only,
'force_gpu_idx' : arguments.force_gpu_idx,
}
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' : arguments.force_gpu_idxs,
'cpu_only' : arguments.cpu_only,
'execute_programs' : [ [int(x[0]), x[1] ] for x in arguments.execute_program ],
'debug' : arguments.debug,
}
from mainscripts import Trainer
Trainer.main(args, device_args)
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('--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('--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('--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('--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.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.")
p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+')
p.set_defaults (func=process_train)
def process_convert(arguments):
os_utils.set_process_lowest_prio()
args = {'training_data_src_dir' : arguments.training_data_src_dir,
'input_dir' : arguments.input_dir,
'output_dir' : arguments.output_dir,
'aligned_dir' : arguments.aligned_dir,
'model_dir' : arguments.model_dir,
'model_name' : arguments.model_name,
def process_merge(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) if arguments.training_data_src_dir is not None else None,
'force_model_name' : arguments.force_model_name,
'input_path' : Path(arguments.input_dir),
'output_path' : Path(arguments.output_dir),
'aligned_path' : Path(arguments.aligned_dir) if arguments.aligned_dir is not None else None,
'cpu_only' : arguments.cpu_only,
'force_gpu_idxs' : arguments.force_gpu_idxs,
}
device_args = {'cpu_only' : arguments.cpu_only,
'force_gpu_idx' : arguments.force_gpu_idx,
}
from mainscripts import Converter
Converter.main (args, device_args)
from mainscripts import Merger
Merger.main (**kwargs)
p = subparsers.add_parser( "convert", help="Converter")
p.add_argument('--training-data-src-dir', action=fixPathAction, dest="training_data_src_dir", help="(optional, may be required by some models) Dir of extracted SRC faceset.")
p = subparsers.add_parser( "merge", help="Merger")
p.add_argument('--training-data-src-dir', action=fixPathAction, dest="training_data_src_dir", default=None, help="(optional, may be required by some models) Dir of extracted SRC faceset.")
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.")
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('--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=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model")
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)
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.set_defaults(func=process_merge)
videoed_parser = subparsers.add_parser( "videoed", help="Video processing.").add_subparsers()
def process_videoed_extract_video(arguments):
os_utils.set_process_lowest_prio()
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.")
@ -227,7 +237,7 @@ if __name__ == "__main__":
p.set_defaults(func=process_videoed_extract_video)
def process_videoed_cut_video(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import VideoEd
VideoEd.cut_video (arguments.input_file,
arguments.from_time,
@ -243,7 +253,7 @@ if __name__ == "__main__":
p.set_defaults(func=process_videoed_cut_video)
def process_videoed_denoise_image_sequence(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import VideoEd
VideoEd.denoise_image_sequence (arguments.input_dir, arguments.ext, arguments.factor)
p = videoed_parser.add_parser( "denoise-image-sequence", help="Denoise sequence of images, keeping sharp edges. This allows you to make the final fake more believable, since the neural network is not able to make a detailed skin texture, but it makes the edges quite clear. Therefore, if the whole frame is more `blurred`, then a fake will seem more believable. Especially true for scenes of the film, which are usually very clear.")
@ -253,7 +263,7 @@ if __name__ == "__main__":
p.set_defaults(func=process_videoed_denoise_image_sequence)
def process_videoed_video_from_sequence(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import VideoEd
VideoEd.video_from_sequence (arguments.input_dir,
arguments.output_file,
@ -289,25 +299,28 @@ if __name__ == "__main__":
facesettool_parser = subparsers.add_parser( "facesettool", help="Faceset tools.").add_subparsers()
def process_faceset_enhancer(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import FacesetEnhancer
FacesetEnhancer.process_folder ( Path(arguments.input_dir), multi_gpu=arguments.multi_gpu, cpu_only=arguments.cpu_only )
FacesetEnhancer.process_folder ( Path(arguments.input_dir),
cpu_only=arguments.cpu_only,
force_gpu_idxs=arguments.force_gpu_idxs
)
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('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.")
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.set_defaults(func=process_faceset_enhancer)
"""
def process_relight_faceset(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import FacesetRelighter
FacesetRelighter.relight (arguments.input_dir, arguments.lighten, arguments.random_one)
def process_delete_relighted(arguments):
os_utils.set_process_lowest_prio()
osex.set_process_lowest_prio()
from mainscripts import FacesetRelighter
FacesetRelighter.delete_relighted (arguments.input_dir)
@ -332,21 +345,6 @@ if __name__ == "__main__":
print ("Done.")
"""
Suppressing error with keras 2.2.4+ on python exit:
Exception ignored in: <bound method BaseSession._Callable.__del__ of <tensorflow.python.client.session.BaseSession._Callable object at 0x000000001BDEA9B0>>
Traceback (most recent call last):
File "D:\DeepFaceLab\_internal\bin\lib\site-packages\tensorflow\python\client\session.py", line 1413, in __del__
AttributeError: 'NoneType' object has no attribute 'raise_exception_on_not_ok_status'
reproduce: https://github.com/keras-team/keras/issues/11751 ( still no solution )
"""
outnull_file = open(os.devnull, 'w')
os.dup2 ( outnull_file.fileno(), sys.stderr.fileno() )
sys.stderr = outnull_file
'''
import code
code.interact(local=dict(globals(), **locals()))