|
|
- #!/usr/bin/env python3
- # -*- coding: utf-8 -*-
- """
- DeepFaceLab 主入口文件
- 这是 DeepFaceLab 的命令行界面入口,提供了多个子命令用于人脸处理、模型训练、视频编辑等功能。
- 主要功能包括:
- - 人脸提取 (extract)
- - 人脸排序 (sort)
- - 工具集 (util)
- - 模型训练 (train)
- - 模型导出 (exportdfm)
- - 结果合并 (merge)
- - 视频编辑 (videoed)
- - 人脸集工具 (facesettool)
- - XSeg 工具 (xseg)
- """
- if __name__ == "__main__":
- # 修复 Linux 系统上的多进程问题
- import multiprocessing
- multiprocessing.set_start_method("spawn")
- # 初始化神经网络环境
- from core.leras import nn
- nn.initialize_main_env()
-
- # 导入必要的系统模块
- import os
- import sys
- import time
- import argparse
- # 导入项目核心模块
- from core import pathex
- from core import osex
- from pathlib import Path
- from core.interact import interact as io
- # 检查 Python 版本,要求至少 3.6
- 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")
- # 自定义的路径处理 Action 类,用于将路径转换为绝对路径
- class fixPathAction(argparse.Action):
- def __call__(self, parser, namespace, values, option_string=None):
- setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values)))
- exit_code = 0
-
- # 创建主命令行解析器
- parser = argparse.ArgumentParser()
- subparsers = parser.add_subparsers()
- # ========== extract 子命令:从图片中提取人脸 ==========
- def process_extract(arguments):
- """处理人脸提取命令"""
- # 设置进程优先级为最低,避免影响系统性能
- osex.set_process_lowest_prio()
- # 导入提取器模块
- from mainscripts import Extractor
- # 调用提取器的主函数
- 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,
- max_faces_from_image = arguments.max_faces_from_image,
- image_size = arguments.image_size,
- jpeg_quality = arguments.jpeg_quality,
- 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,
- )
- # 创建 extract 子命令解析器
- 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('--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', 'whole_face', 'head', 'mark_only'], default=None)
- p.add_argument('--max-faces-from-image', type=int, dest="max_faces_from_image", default=None, help="Max faces from image.")
- p.add_argument('--image-size', type=int, dest="image_size", default=None, help="Output image size.")
- p.add_argument('--jpeg-quality', type=int, dest="jpeg_quality", default=None, help="Jpeg quality.")
- 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..")
- 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)
- # ========== sort 子命令:对人脸进行排序 ==========
- def process_sort(arguments):
- """处理人脸排序命令"""
- osex.set_process_lowest_prio()
- from mainscripts import Sorter
- 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', dest="sort_by_method", default=None, choices=("blur", "motion-blur", "face-yaw", "face-pitch", "face-source-rect-size", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final-by-blur", "final-by-size", "absdiff"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
- p.set_defaults (func=process_sort)
- # ========== util 子命令:实用工具集合 ==========
- def process_util(arguments):
- """处理工具命令"""
- osex.set_process_lowest_prio()
- from mainscripts import Util
- 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.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) )
-
- if arguments.export_faceset_mask:
- io.log_info ("Exporting faceset mask..\r\n")
- Util.export_faceset_mask( 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('--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('--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="Pack faceset into a binary format.")
- p.add_argument('--unpack-faceset', action="store_true", dest="unpack_faceset", default=False, help="Unpack faceset from binary format.")
- p.add_argument('--export-faceset-mask', action="store_true", dest="export_faceset_mask", default=False, help="Export faceset masks.")
- p.set_defaults (func=process_util)
- # ========== train 子命令:训练模型 ==========
- def process_train(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),
- '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' : [ int(x) for x in arguments.force_gpu_idxs.split(',') ] if arguments.force_gpu_idxs is not None else None,
- 'cpu_only' : arguments.cpu_only,
- 'silent_start' : arguments.silent_start,
- 'execute_programs' : [ [int(x[0]), x[1] ] for x in arguments.execute_program ],
- 'debug' : arguments.debug,
- }
- from mainscripts import Trainer
- 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('--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-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.")
- p.add_argument('--silent-start', action="store_true", dest="silent_start", default=False, help="Silent start. Automatically chooses Best GPU and last used model.")
-
- p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+')
- p.set_defaults (func=process_train)
-
- # ========== exportdfm 子命令:导出模型用于 DeepFaceLive ==========
- def process_exportdfm(arguments):
- """处理模型导出命令"""
- osex.set_process_lowest_prio()
- from mainscripts import ExportDFM
- ExportDFM.main(model_class_name = arguments.model_name, saved_models_path = Path(arguments.model_dir))
- p = subparsers.add_parser( "exportdfm", help="Export model to use in DeepFaceLive.")
- 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.set_defaults (func=process_exportdfm)
- # ========== merge 子命令:合并结果 ==========
- def process_merge(arguments):
- """处理结果合并命令"""
- osex.set_process_lowest_prio()
- from mainscripts import Merger
- Merger.main ( model_class_name = arguments.model_name,
- saved_models_path = Path(arguments.model_dir),
- force_model_name = arguments.force_model_name,
- input_path = Path(arguments.input_dir),
- output_path = Path(arguments.output_dir),
- output_mask_path = Path(arguments.output_mask_dir),
- aligned_path = Path(arguments.aligned_dir) if arguments.aligned_dir is not None else None,
- force_gpu_idxs = arguments.force_gpu_idxs,
- cpu_only = arguments.cpu_only)
- p = subparsers.add_parser( "merge", help="Merger")
- 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 merged files will be stored.")
- p.add_argument('--output-mask-dir', required=True, action=fixPathAction, dest="output_mask_dir", help="Output mask directory. This is where the mask 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=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 子命令:视频编辑工具 ==========
- videoed_parser = subparsers.add_parser( "videoed", help="Video processing.").add_subparsers()
- # videoed extract-video 子命令:从视频中提取帧
- def process_videoed_extract_video(arguments):
- """处理视频帧提取命令"""
- 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.")
- 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)
- # videoed cut-video 子命令:裁剪视频
- def process_videoed_cut_video(arguments):
- """处理视频裁剪命令"""
- osex.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)
- # videoed denoise-image-sequence 子命令:降噪图像序列
- def process_videoed_denoise_image_sequence(arguments):
- """处理图像序列降噪命令"""
- osex.set_process_lowest_prio()
- from mainscripts import VideoEd
- VideoEd.denoise_image_sequence (arguments.input_dir, arguments.factor)
- p = videoed_parser.add_parser( "denoise-image-sequence", help="Denoise sequence of images, keeping sharp edges. Helps to remove pixel shake from the predicted face.")
- p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory to be processed.")
- p.add_argument('--factor', type=int, dest="factor", default=None, help="Denoise factor (1-20).")
- p.set_defaults(func=process_videoed_denoise_image_sequence)
- # videoed video-from-sequence 子命令:从图像序列创建视频
- def process_videoed_video_from_sequence(arguments):
- """处理从图像序列创建视频命令"""
- osex.set_process_lowest_prio()
- from mainscripts import VideoEd
- VideoEd.video_from_sequence (input_dir = arguments.input_dir,
- output_file = arguments.output_file,
- reference_file = arguments.reference_file,
- ext = arguments.ext,
- fps = arguments.fps,
- bitrate = arguments.bitrate,
- include_audio = arguments.include_audio,
- lossless = 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('--include-audio', action="store_true", dest="include_audio", default=False, help="Include audio from reference file.")
- p.add_argument('--lossless', action="store_true", dest="lossless", default=False, help="PNG codec.")
- p.set_defaults(func=process_videoed_video_from_sequence)
- # ========== facesettool 子命令:人脸集工具 ==========
- facesettool_parser = subparsers.add_parser( "facesettool", help="Faceset tools.").add_subparsers()
- # facesettool enhance 子命令:增强人脸集细节
- def process_faceset_enhancer(arguments):
- """处理人脸集增强命令"""
- osex.set_process_lowest_prio()
- from mainscripts import FacesetEnhancer
- 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('--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)
-
-
- p = facesettool_parser.add_parser ("resize", help="Resize DFL faceset.")
- p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.")
- def process_faceset_resizer(arguments):
- """处理人脸集调整大小命令"""
- osex.set_process_lowest_prio()
- from mainscripts import FacesetResizer
- FacesetResizer.process_folder ( Path(arguments.input_dir) )
- p.set_defaults(func=process_faceset_resizer)
- # ========== dev_test 子命令:开发测试用 ==========
- def process_dev_test(arguments):
- """处理开发测试命令"""
- osex.set_process_lowest_prio()
- from mainscripts import dev_misc
- dev_misc.dev_gen_mask_files( arguments.input_dir )
- p = subparsers.add_parser( "dev_test", help="Development test command.")
- p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
- p.set_defaults (func=process_dev_test)
-
- # ========== XSeg 子命令:XSeg 工具 ==========
- xseg_parser = subparsers.add_parser( "xseg", help="XSeg tools.").add_subparsers()
-
- # xseg editor 子命令:XSeg 编辑器
- p = xseg_parser.add_parser( "editor", help="XSeg editor.")
- def process_xsegeditor(arguments):
- """处理 XSeg 编辑器命令"""
- osex.set_process_lowest_prio()
- from XSegEditor import XSegEditor
- global exit_code
- exit_code = XSegEditor.start (Path(arguments.input_dir))
-
- p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
- p.set_defaults (func=process_xsegeditor)
-
- # xseg apply 子命令:应用训练好的 XSeg 模型到提取的人脸
- p = xseg_parser.add_parser( "apply", help="Apply trained XSeg model to the extracted faces.")
- def process_xsegapply(arguments):
- """处理应用 XSeg 模型命令"""
- osex.set_process_lowest_prio()
- from mainscripts import XSegUtil
- XSegUtil.apply_xseg (Path(arguments.input_dir), Path(arguments.model_dir))
- p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
- p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir")
- p.set_defaults (func=process_xsegapply)
-
- # xseg remove 子命令:从提取的人脸中移除应用的 XSeg 掩码
- p = xseg_parser.add_parser( "remove", help="Remove applied XSeg masks from the extracted faces.")
- def process_xsegremove(arguments):
- """处理移除 XSeg 掩码命令"""
- osex.set_process_lowest_prio()
- from mainscripts import XSegUtil
- XSegUtil.remove_xseg (Path(arguments.input_dir) )
- p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
- p.set_defaults (func=process_xsegremove)
-
- # xseg remove_labels 子命令:从提取的人脸中移除 XSeg 标签
- p = xseg_parser.add_parser( "remove_labels", help="Remove XSeg labels from the extracted faces.")
- def process_xsegremovelabels(arguments):
- """处理移除 XSeg 标签命令"""
- osex.set_process_lowest_prio()
- from mainscripts import XSegUtil
- XSegUtil.remove_xseg_labels (Path(arguments.input_dir) )
- p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
- p.set_defaults (func=process_xsegremovelabels)
-
- # xseg fetch 子命令:复制包含 XSeg 多边形的人脸到 <input_dir>_xseg 目录
- p = xseg_parser.add_parser( "fetch", help="Copies faces containing XSeg polygons in <input_dir>_xseg dir.")
- def process_xsegfetch(arguments):
- """处理获取 XSeg 标记人脸命令"""
- osex.set_process_lowest_prio()
- from mainscripts import XSegUtil
- XSegUtil.fetch_xseg (Path(arguments.input_dir) )
- p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
- p.set_defaults (func=process_xsegfetch)
-
- # 默认命令:显示帮助信息
- def bad_args(arguments):
- parser.print_help()
- exit(0)
- parser.set_defaults(func=bad_args)
- # 解析命令行参数并执行对应的功能
- arguments = parser.parse_args()
- arguments.func(arguments)
- # 输出完成信息并退出
- if exit_code == 0:
- print ("Done.")
-
- exit(exit_code)
-
- # 调试代码(被注释掉)
- '''import code
- code.interact(local=dict(globals(), **locals()))'''
复制代码parser规则理解:
举例说明 :
- 当用户运行 python main.py extract --input-dir images --output-dir faces
- parser.parse_args() 会解析出子命令是 "extract",并获取 input-dir 和 output-dir 的值
- arguments 对象将包含 arguments.func = process_extract 、 arguments.input_dir = "images" 、 arguments.output_dir = "faces" 等属性
- arguments.func(arguments) 相当于 process_extract(arguments)
然后执行 process_extract(arguments) 函数,开始人脸提取过程
|
|