DeepFaceLab/main.py
Colombo d46fb5cfd3 fixed mask editor
added FacesetEnhancer
4.2.other) data_src util faceset enhance best GPU.bat
4.2.other) data_src util faceset enhance multi GPU.bat

FacesetEnhancer greatly increases details in your source face set,
same as Gigapixel enhancer, but in fully automatic mode.
In OpenCL build it works on CPU only.

Please consider a donation.
2019-12-26 21:27:10 +04:00

354 lines
23 KiB
Python

if __name__ == "__main__":
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 pathlib import Path
from 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")
class fixPathAction(argparse.Action):
def __call__(self, parser, namespace, values, option_string=None):
setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values)))
parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers()
def process_extract(arguments):
os_utils.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,
}
)
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-dir', action=fixPathAction, dest="debug_dir", help="Writes debug images to this 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.set_defaults (func=process_extract)
def process_dev_extract_vggface2_dataset(arguments):
os_utils.set_process_lowest_prio()
from mainscripts import dev_misc
dev_misc.extract_vggface2_dataset( arguments.input_dir,
device_args={'cpu_only' : arguments.cpu_only,
'multi_gpu' : arguments.multi_gpu,
}
)
p = subparsers.add_parser( "dev_extract_vggface2_dataset", help="")
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('--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="Extract on CPU.")
p.set_defaults (func=process_dev_extract_vggface2_dataset)
def process_dev_extract_umd_csv(arguments):
os_utils.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,
'multi_gpu' : arguments.multi_gpu,
}
)
p = subparsers.add_parser( "dev_extract_umd_csv", help="")
p.add_argument('--input-csv-file', required=True, action=fixPathAction, dest="input_csv_file", help="input_csv_file")
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="Extract on CPU.")
p.set_defaults (func=process_dev_extract_umd_csv)
def process_dev_apply_celebamaskhq(arguments):
os_utils.set_process_lowest_prio()
from mainscripts import dev_misc
dev_misc.apply_celebamaskhq( arguments.input_dir )
p = subparsers.add_parser( "dev_apply_celebamaskhq", help="")
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
p.set_defaults (func=process_dev_apply_celebamaskhq)
def process_dev_test(arguments):
os_utils.set_process_lowest_prio()
from mainscripts import dev_misc
dev_misc.dev_test( arguments.input_dir )
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)
def process_sort(arguments):
os_utils.set_process_lowest_prio()
from mainscripts import Sorter
Sorter.main (input_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.set_defaults (func=process_sort)
def process_util(arguments):
os_utils.set_process_lowest_prio()
from mainscripts import Util
if arguments.convert_png_to_jpg:
Util.convert_png_to_jpg_folder (input_path=arguments.input_dir)
if arguments.add_landmarks_debug_images:
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)
#if arguments.remove_fanseg:
# Util.remove_fanseg_folder (input_path=arguments.input_dir)
if arguments.remove_ie_polys:
Util.remove_ie_polys_folder (input_path=arguments.input_dir)
if arguments.save_faceset_metadata:
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)
if arguments.pack_faceset:
io.log_info ("Performing faceset packing...\r\n")
from samplelib import PackedFaceset
PackedFaceset.pack( Path(arguments.input_dir) )
if arguments.unpack_faceset:
io.log_info ("Performing faceset unpacking...\r\n")
from samplelib import PackedFaceset
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('--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.add_argument('--remove-fanseg', action="store_true", dest="remove_fanseg", default=False, help="Remove fanseg mask from aligned faces.")
p.add_argument('--remove-ie-polys', action="store_true", dest="remove_ie_polys", default=False, help="Remove ie_polys from aligned faces.")
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.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,
}
from mainscripts import Trainer
Trainer.main(args, device_args)
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('--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.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,
}
device_args = {'cpu_only' : arguments.cpu_only,
'force_gpu_idx' : arguments.force_gpu_idx,
}
from mainscripts import Converter
Converter.main (args, device_args)
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.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('--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)
videoed_parser = subparsers.add_parser( "videoed", help="Video processing.").add_subparsers()
def process_videoed_extract_video(arguments):
os_utils.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.")
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.set_defaults(func=process_videoed_extract_video)
def process_videoed_cut_video(arguments):
os_utils.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.")
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):
os_utils.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.")
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('--ext', dest="ext", default=None, help="Image format (extension) of input files.")
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):
os_utils.set_process_lowest_prio()
from mainscripts import VideoEd
VideoEd.video_from_sequence (arguments.input_dir,
arguments.output_file,
arguments.reference_file,
arguments.ext,
arguments.fps,
arguments.bitrate,
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.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('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.")
p.add_argument('--lossless', action="store_true", dest="lossless", default=False, help="PNG codec.")
p.set_defaults(func=process_videoed_video_from_sequence)
def process_labelingtool_edit_mask(arguments):
from mainscripts import MaskEditorTool
MaskEditorTool.mask_editor_main (arguments.input_dir, arguments.confirmed_dir, arguments.skipped_dir, no_default_mask=arguments.no_default_mask)
labeling_parser = subparsers.add_parser( "labelingtool", help="Labeling tool.").add_subparsers()
p = labeling_parser.add_parser ( "edit_mask", help="")
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.")
p.add_argument('--confirmed-dir', required=True, action=fixPathAction, dest="confirmed_dir", help="This is where the labeled faces will be stored.")
p.add_argument('--skipped-dir', required=True, action=fixPathAction, dest="skipped_dir", help="This is where the labeled faces will be stored.")
p.add_argument('--no-default-mask', action="store_true", dest="no_default_mask", default=False, help="Don't use default mask.")
p.set_defaults(func=process_labelingtool_edit_mask)
facesettool_parser = subparsers.add_parser( "facesettool", help="Faceset tools.").add_subparsers()
def process_faceset_enhancer(arguments):
os_utils.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 )
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.set_defaults(func=process_faceset_enhancer)
"""
def process_relight_faceset(arguments):
os_utils.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()
from mainscripts import FacesetRelighter
FacesetRelighter.delete_relighted (arguments.input_dir)
p = facesettool_parser.add_parser ("relight", help="Synthesize new faces from existing ones by relighting them. With the relighted faces neural network will better reproduce face shadows.")
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.")
p.add_argument('--lighten', action="store_true", dest="lighten", default=None, help="Lighten the faces.")
p.add_argument('--random-one', action="store_true", dest="random_one", default=None, help="Relight the faces only with one random direction, otherwise relight with all directions.")
p.set_defaults(func=process_relight_faceset)
p = facesettool_parser.add_parser ("delete_relighted", help="Delete relighted faces.")
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.")
p.set_defaults(func=process_delete_relighted)
"""
def bad_args(arguments):
parser.print_help()
exit(0)
parser.set_defaults(func=bad_args)
arguments = parser.parse_args()
arguments.func(arguments)
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()))
'''