diff --git a/.github/preview.png b/.github/preview.png index 960017c0..b0c29e5f 100644 Binary files a/.github/preview.png and b/.github/preview.png differ diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 39bf14ee..e31ce521 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -30,6 +30,6 @@ jobs: uses: actions/setup-python@v2 with: python-version: '3.10' - - run: python install.py --torch cpu --onnxruntime default + - run: python install.py --torch cpu --onnxruntime default --skip-venv - run: pip install pytest - run: pytest diff --git a/README.md b/README.md index d878859c..45bb9533 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,7 @@ python run.py [options] options: -h, --help show this help message and exit - -s SOURCE_PATH, --source SOURCE_PATH select a source image + -s SOURCE_PATHS, --source SOURCE_PATHS select a source image -t TARGET_PATH, --target TARGET_PATH select a target image or video -o OUTPUT_PATH, --output OUTPUT_PATH specify the output file or directory -v, --version show program's version number and exit @@ -39,9 +39,10 @@ options: misc: --skip-download omit automate downloads and lookups --headless run the program in headless mode + --log-level {error,warn,info,debug} choose from the available log levels execution: - --execution-providers {cpu} [{cpu} ...] choose from the available execution providers + --execution-providers EXECUTION_PROVIDERS [EXECUTION_PROVIDERS ...] choose from the available execution providers (choices: cpu, ...) --execution-thread-count [1-128] specify the number of execution threads --execution-queue-count [1-32] specify the number of execution queries --max-memory [0-128] specify the maximum amount of ram to be used (in gb) @@ -61,8 +62,10 @@ face selector: --reference-frame-number REFERENCE_FRAME_NUMBER specify the number of the reference frame face mask: + --face-mask-types FACE_MASK_TYPES [FACE_MASK_TYPES ...] choose from the available face mask types (choices: box, occlusion, region) --face-mask-blur [0.0-1.0] specify the blur amount for face mask --face-mask-padding FACE_MASK_PADDING [FACE_MASK_PADDING ...] specify the face mask padding (top, right, bottom, left) in percent + --face-mask-regions FACE_MASK_REGIONS [FACE_MASK_REGIONS ...] choose from the available face mask regions (choices: skin, left-eyebrow, right-eyebrow, left-eye, right-eye, eye-glasses, nose, mouth, upper-lip, lower-lip) frame extraction: --trim-frame-start TRIM_FRAME_START specify the start frame for extraction @@ -80,12 +83,12 @@ output creation: frame processors: --frame-processors FRAME_PROCESSORS [FRAME_PROCESSORS ...] choose from the available frame processors (choices: face_debugger, face_enhancer, face_swapper, frame_enhancer, ...) - --face-debugger-items {bbox,kps,face-mask,score} [{bbox,kps,face-mask,score} ...] specify the face debugger items + --face-debugger-items FACE_DEBUGGER_ITEMS [FACE_DEBUGGER_ITEMS ...] specify the face debugger items (choices: bbox, kps, face-mask, score) --face-enhancer-model {codeformer,gfpgan_1.2,gfpgan_1.3,gfpgan_1.4,gpen_bfr_256,gpen_bfr_512,restoreformer} choose the model for the frame processor - --face-enhancer-blend [0-100] specify the blend factor for the frame processor + --face-enhancer-blend [0-100] specify the blend amount for the frame processor --face-swapper-model {blendswap_256,inswapper_128,inswapper_128_fp16,simswap_256,simswap_512_unofficial} choose the model for the frame processor --frame-enhancer-model {real_esrgan_x2plus,real_esrgan_x4plus,real_esrnet_x4plus} choose the model for the frame processor - --frame-enhancer-blend [0-100] specify the blend factor for the frame processor + --frame-enhancer-blend [0-100] specify the blend amount for the frame processor uis: --ui-layouts UI_LAYOUTS [UI_LAYOUTS ...] choose from the available ui layouts (choices: benchmark, webcam, default, ...) diff --git a/facefusion/choices.py b/facefusion/choices.py old mode 100644 new mode 100755 index cadeda67..135586cc --- a/facefusion/choices.py +++ b/facefusion/choices.py @@ -2,8 +2,7 @@ from typing import List import numpy -from facefusion.typing import FaceSelectorMode, FaceAnalyserOrder, FaceAnalyserAge, FaceAnalyserGender, TempFrameFormat, OutputVideoEncoder - +from facefusion.typing import FaceSelectorMode, FaceAnalyserOrder, FaceAnalyserAge, FaceAnalyserGender, FaceMaskType, FaceMaskRegion, TempFrameFormat, OutputVideoEncoder face_analyser_orders : List[FaceAnalyserOrder] = [ 'left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best' ] face_analyser_ages : List[FaceAnalyserAge] = [ 'child', 'teen', 'adult', 'senior' ] @@ -11,6 +10,8 @@ face_analyser_genders : List[FaceAnalyserGender] = [ 'male', 'female' ] face_detector_models : List[str] = [ 'retinaface', 'yunet' ] face_detector_sizes : List[str] = [ '160x160', '320x320', '480x480', '512x512', '640x640', '768x768', '960x960', '1024x1024' ] face_selector_modes : List[FaceSelectorMode] = [ 'reference', 'one', 'many' ] +face_mask_types : List[FaceMaskType] = [ 'box', 'occlusion', 'region' ] +face_mask_regions : List[FaceMaskRegion] = [ 'skin', 'left-eyebrow', 'right-eyebrow', 'left-eye', 'right-eye', 'eye-glasses', 'nose', 'mouth', 'upper-lip', 'lower-lip' ] temp_frame_formats : List[TempFrameFormat] = [ 'jpg', 'png' ] output_video_encoders : List[OutputVideoEncoder] = [ 'libx264', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc' ] diff --git a/facefusion/cli_helper.py b/facefusion/cli_helper.py new file mode 100644 index 00000000..84aa847d --- /dev/null +++ b/facefusion/cli_helper.py @@ -0,0 +1,5 @@ +from typing import List, Any + + +def create_metavar(ranges : List[Any]) -> str: + return '[' + str(ranges[0]) + '-' + str(ranges[-1]) + ']' diff --git a/facefusion/content_analyser.py b/facefusion/content_analyser.py index 2111effa..daa276e9 100644 --- a/facefusion/content_analyser.py +++ b/facefusion/content_analyser.py @@ -10,7 +10,8 @@ import facefusion.globals from facefusion import wording from facefusion.typing import Frame, ModelValue from facefusion.vision import get_video_frame, count_video_frame_total, read_image, detect_fps -from facefusion.utilities import resolve_relative_path, conditional_download +from facefusion.filesystem import resolve_relative_path +from facefusion.download import conditional_download CONTENT_ANALYSER = None THREAD_LOCK : threading.Lock = threading.Lock() @@ -90,7 +91,7 @@ def analyse_video(video_path : str, start_frame : int, end_frame : int) -> bool: frame_range = range(start_frame or 0, end_frame or video_frame_total) rate = 0.0 counter = 0 - with tqdm(total = len(frame_range), desc = wording.get('analysing'), unit = 'frame', ascii = ' =') as progress: + with tqdm(total = len(frame_range), desc = wording.get('analysing'), unit = 'frame', ascii = ' =', disable = facefusion.globals.log_level in [ 'warn', 'error' ]) as progress: for frame_number in frame_range: if frame_number % int(fps) == 0: frame = get_video_frame(video_path, frame_number) diff --git a/facefusion/core.py b/facefusion/core.py index c6c205cb..ad9941d2 100755 --- a/facefusion/core.py +++ b/facefusion/core.py @@ -3,6 +3,7 @@ import os os.environ['OMP_NUM_THREADS'] = '1' import signal +import ssl import sys import warnings import platform @@ -12,92 +13,104 @@ from argparse import ArgumentParser, HelpFormatter import facefusion.choices import facefusion.globals -from facefusion.face_analyser import get_one_face -from facefusion.face_reference import get_face_reference, set_face_reference -from facefusion.vision import get_video_frame, read_image -from facefusion import face_analyser, content_analyser, metadata, wording +from facefusion.face_analyser import get_one_face, get_average_face +from facefusion.face_store import get_reference_faces, append_reference_face +from facefusion.vision import get_video_frame, detect_fps, read_image, read_static_images +from facefusion import face_analyser, face_masker, content_analyser, metadata, logger, wording from facefusion.content_analyser import analyse_image, analyse_video from facefusion.processors.frame.core import get_frame_processors_modules, load_frame_processor_module -from facefusion.utilities import is_image, is_video, detect_fps, compress_image, merge_video, extract_frames, get_temp_frame_paths, restore_audio, create_temp, move_temp, clear_temp, list_module_names, encode_execution_providers, decode_execution_providers, normalize_output_path, normalize_padding, create_metavar, update_status +from facefusion.cli_helper import create_metavar +from facefusion.execution_helper import encode_execution_providers, decode_execution_providers +from facefusion.normalizer import normalize_output_path, normalize_padding +from facefusion.filesystem import is_image, is_video, list_module_names, get_temp_frame_paths, create_temp, move_temp, clear_temp +from facefusion.ffmpeg import extract_frames, compress_image, merge_video, restore_audio onnxruntime.set_default_logger_severity(3) warnings.filterwarnings('ignore', category = UserWarning, module = 'gradio') warnings.filterwarnings('ignore', category = UserWarning, module = 'torchvision') +if platform.system().lower() == 'darwin': + ssl._create_default_https_context = ssl._create_unverified_context + def cli() -> None: signal.signal(signal.SIGINT, lambda signal_number, frame: destroy()) program = ArgumentParser(formatter_class = lambda prog: HelpFormatter(prog, max_help_position = 120), add_help = False) # general - program.add_argument('-s', '--source', help = wording.get('source_help'), dest = 'source_path') + program.add_argument('-s', '--source', action = 'append', help = wording.get('source_help'), dest = 'source_paths') program.add_argument('-t', '--target', help = wording.get('target_help'), dest = 'target_path') program.add_argument('-o', '--output', help = wording.get('output_help'), dest = 'output_path') program.add_argument('-v', '--version', version = metadata.get('name') + ' ' + metadata.get('version'), action = 'version') # misc group_misc = program.add_argument_group('misc') - group_misc.add_argument('--skip-download', help = wording.get('skip_download_help'), dest = 'skip_download', action = 'store_true') - group_misc.add_argument('--headless', help = wording.get('headless_help'), dest = 'headless', action = 'store_true') + group_misc.add_argument('--skip-download', help = wording.get('skip_download_help'), action = 'store_true') + group_misc.add_argument('--headless', help = wording.get('headless_help'), action = 'store_true') + group_misc.add_argument('--log-level', help = wording.get('log_level_help'), default = 'info', choices = logger.get_log_levels()) # execution + execution_providers = encode_execution_providers(onnxruntime.get_available_providers()) group_execution = program.add_argument_group('execution') - group_execution.add_argument('--execution-providers', help = wording.get('execution_providers_help'), dest = 'execution_providers', default = [ 'cpu' ], choices = encode_execution_providers(onnxruntime.get_available_providers()), nargs = '+') - group_execution.add_argument('--execution-thread-count', help = wording.get('execution_thread_count_help'), dest = 'execution_thread_count', type = int, default = 4, choices = facefusion.choices.execution_thread_count_range, metavar = create_metavar(facefusion.choices.execution_thread_count_range)) - group_execution.add_argument('--execution-queue-count', help = wording.get('execution_queue_count_help'), dest = 'execution_queue_count', type = int, default = 1, choices = facefusion.choices.execution_queue_count_range, metavar = create_metavar(facefusion.choices.execution_queue_count_range)) - group_execution.add_argument('--max-memory', help = wording.get('max_memory_help'), dest = 'max_memory', type = int, choices = facefusion.choices.max_memory_range, metavar = create_metavar(facefusion.choices.max_memory_range)) + group_execution.add_argument('--execution-providers', help = wording.get('execution_providers_help').format(choices = ', '.join(execution_providers)), default = [ 'cpu' ], choices = execution_providers, nargs = '+', metavar = 'EXECUTION_PROVIDERS') + group_execution.add_argument('--execution-thread-count', help = wording.get('execution_thread_count_help'), type = int, default = 4, choices = facefusion.choices.execution_thread_count_range, metavar = create_metavar(facefusion.choices.execution_thread_count_range)) + group_execution.add_argument('--execution-queue-count', help = wording.get('execution_queue_count_help'), type = int, default = 1, choices = facefusion.choices.execution_queue_count_range, metavar = create_metavar(facefusion.choices.execution_queue_count_range)) + group_execution.add_argument('--max-memory', help = wording.get('max_memory_help'), type = int, choices = facefusion.choices.max_memory_range, metavar = create_metavar(facefusion.choices.max_memory_range)) # face analyser group_face_analyser = program.add_argument_group('face analyser') - group_face_analyser.add_argument('--face-analyser-order', help = wording.get('face_analyser_order_help'), dest = 'face_analyser_order', default = 'left-right', choices = facefusion.choices.face_analyser_orders) - group_face_analyser.add_argument('--face-analyser-age', help = wording.get('face_analyser_age_help'), dest = 'face_analyser_age', choices = facefusion.choices.face_analyser_ages) - group_face_analyser.add_argument('--face-analyser-gender', help = wording.get('face_analyser_gender_help'), dest = 'face_analyser_gender', choices = facefusion.choices.face_analyser_genders) - group_face_analyser.add_argument('--face-detector-model', help = wording.get('face_detector_model_help'), dest = 'face_detector_model', default = 'retinaface', choices = facefusion.choices.face_detector_models) - group_face_analyser.add_argument('--face-detector-size', help = wording.get('face_detector_size_help'), dest = 'face_detector_size', default = '640x640', choices = facefusion.choices.face_detector_sizes) - group_face_analyser.add_argument('--face-detector-score', help = wording.get('face_detector_score_help'), dest = 'face_detector_score', type = float, default = 0.5, choices = facefusion.choices.face_detector_score_range, metavar = create_metavar(facefusion.choices.face_detector_score_range)) + group_face_analyser.add_argument('--face-analyser-order', help = wording.get('face_analyser_order_help'), default = 'left-right', choices = facefusion.choices.face_analyser_orders) + group_face_analyser.add_argument('--face-analyser-age', help = wording.get('face_analyser_age_help'), choices = facefusion.choices.face_analyser_ages) + group_face_analyser.add_argument('--face-analyser-gender', help = wording.get('face_analyser_gender_help'), choices = facefusion.choices.face_analyser_genders) + group_face_analyser.add_argument('--face-detector-model', help = wording.get('face_detector_model_help'), default = 'retinaface', choices = facefusion.choices.face_detector_models) + group_face_analyser.add_argument('--face-detector-size', help = wording.get('face_detector_size_help'), default = '640x640', choices = facefusion.choices.face_detector_sizes) + group_face_analyser.add_argument('--face-detector-score', help = wording.get('face_detector_score_help'), type = float, default = 0.5, choices = facefusion.choices.face_detector_score_range, metavar = create_metavar(facefusion.choices.face_detector_score_range)) # face selector group_face_selector = program.add_argument_group('face selector') - group_face_selector.add_argument('--face-selector-mode', help = wording.get('face_selector_mode_help'), dest = 'face_selector_mode', default = 'reference', choices = facefusion.choices.face_selector_modes) - group_face_selector.add_argument('--reference-face-position', help = wording.get('reference_face_position_help'), dest = 'reference_face_position', type = int, default = 0) - group_face_selector.add_argument('--reference-face-distance', help = wording.get('reference_face_distance_help'), dest = 'reference_face_distance', type = float, default = 0.6, choices = facefusion.choices.reference_face_distance_range, metavar = create_metavar(facefusion.choices.reference_face_distance_range)) - group_face_selector.add_argument('--reference-frame-number', help = wording.get('reference_frame_number_help'), dest = 'reference_frame_number', type = int, default = 0) + group_face_selector.add_argument('--face-selector-mode', help = wording.get('face_selector_mode_help'), default = 'reference', choices = facefusion.choices.face_selector_modes) + group_face_selector.add_argument('--reference-face-position', help = wording.get('reference_face_position_help'), type = int, default = 0) + group_face_selector.add_argument('--reference-face-distance', help = wording.get('reference_face_distance_help'), type = float, default = 0.6, choices = facefusion.choices.reference_face_distance_range, metavar = create_metavar(facefusion.choices.reference_face_distance_range)) + group_face_selector.add_argument('--reference-frame-number', help = wording.get('reference_frame_number_help'), type = int, default = 0) # face mask group_face_mask = program.add_argument_group('face mask') - group_face_mask.add_argument('--face-mask-blur', help = wording.get('face_mask_blur_help'), dest = 'face_mask_blur', type = float, default = 0.3, choices = facefusion.choices.face_mask_blur_range, metavar = create_metavar(facefusion.choices.face_mask_blur_range)) - group_face_mask.add_argument('--face-mask-padding', help = wording.get('face_mask_padding_help'), dest = 'face_mask_padding', type = int, default = [ 0, 0, 0, 0 ], nargs = '+') + group_face_mask.add_argument('--face-mask-types', help = wording.get('face_mask_types_help').format(choices = ', '.join(facefusion.choices.face_mask_types)), default = [ 'box' ], choices = facefusion.choices.face_mask_types, nargs = '+', metavar = 'FACE_MASK_TYPES') + group_face_mask.add_argument('--face-mask-blur', help = wording.get('face_mask_blur_help'), type = float, default = 0.3, choices = facefusion.choices.face_mask_blur_range, metavar = create_metavar(facefusion.choices.face_mask_blur_range)) + group_face_mask.add_argument('--face-mask-padding', help = wording.get('face_mask_padding_help'), type = int, default = [ 0, 0, 0, 0 ], nargs = '+') + group_face_mask.add_argument('--face-mask-regions', help = wording.get('face_mask_regions_help').format(choices = ', '.join(facefusion.choices.face_mask_regions)), default = facefusion.choices.face_mask_regions, choices = facefusion.choices.face_mask_regions, nargs = '+', metavar = 'FACE_MASK_REGIONS') # frame extraction group_frame_extraction = program.add_argument_group('frame extraction') - group_frame_extraction.add_argument('--trim-frame-start', help = wording.get('trim_frame_start_help'), dest = 'trim_frame_start', type = int) - group_frame_extraction.add_argument('--trim-frame-end', help = wording.get('trim_frame_end_help'), dest = 'trim_frame_end', type = int) - group_frame_extraction.add_argument('--temp-frame-format', help = wording.get('temp_frame_format_help'), dest = 'temp_frame_format', default = 'jpg', choices = facefusion.choices.temp_frame_formats) - group_frame_extraction.add_argument('--temp-frame-quality', help = wording.get('temp_frame_quality_help'), dest = 'temp_frame_quality', type = int, default = 100, choices = facefusion.choices.temp_frame_quality_range, metavar = create_metavar(facefusion.choices.temp_frame_quality_range)) - group_frame_extraction.add_argument('--keep-temp', help = wording.get('keep_temp_help'), dest = 'keep_temp', action = 'store_true') + group_frame_extraction.add_argument('--trim-frame-start', help = wording.get('trim_frame_start_help'), type = int) + group_frame_extraction.add_argument('--trim-frame-end', help = wording.get('trim_frame_end_help'), type = int) + group_frame_extraction.add_argument('--temp-frame-format', help = wording.get('temp_frame_format_help'), default = 'jpg', choices = facefusion.choices.temp_frame_formats) + group_frame_extraction.add_argument('--temp-frame-quality', help = wording.get('temp_frame_quality_help'), type = int, default = 100, choices = facefusion.choices.temp_frame_quality_range, metavar = create_metavar(facefusion.choices.temp_frame_quality_range)) + group_frame_extraction.add_argument('--keep-temp', help = wording.get('keep_temp_help'), action = 'store_true') # output creation group_output_creation = program.add_argument_group('output creation') - group_output_creation.add_argument('--output-image-quality', help = wording.get('output_image_quality_help'), dest = 'output_image_quality', type = int, default = 80, choices = facefusion.choices.output_image_quality_range, metavar = create_metavar(facefusion.choices.output_image_quality_range)) - group_output_creation.add_argument('--output-video-encoder', help = wording.get('output_video_encoder_help'), dest = 'output_video_encoder', default = 'libx264', choices = facefusion.choices.output_video_encoders) - group_output_creation.add_argument('--output-video-quality', help = wording.get('output_video_quality_help'), dest = 'output_video_quality', type = int, default = 80, choices = facefusion.choices.output_video_quality_range, metavar = create_metavar(facefusion.choices.output_video_quality_range)) - group_output_creation.add_argument('--keep-fps', help = wording.get('keep_fps_help'), dest = 'keep_fps', action = 'store_true') - group_output_creation.add_argument('--skip-audio', help = wording.get('skip_audio_help'), dest = 'skip_audio', action = 'store_true') + group_output_creation.add_argument('--output-image-quality', help = wording.get('output_image_quality_help'), type = int, default = 80, choices = facefusion.choices.output_image_quality_range, metavar = create_metavar(facefusion.choices.output_image_quality_range)) + group_output_creation.add_argument('--output-video-encoder', help = wording.get('output_video_encoder_help'), default = 'libx264', choices = facefusion.choices.output_video_encoders) + group_output_creation.add_argument('--output-video-quality', help = wording.get('output_video_quality_help'), type = int, default = 80, choices = facefusion.choices.output_video_quality_range, metavar = create_metavar(facefusion.choices.output_video_quality_range)) + group_output_creation.add_argument('--keep-fps', help = wording.get('keep_fps_help'), action = 'store_true') + group_output_creation.add_argument('--skip-audio', help = wording.get('skip_audio_help'), action = 'store_true') # frame processors available_frame_processors = list_module_names('facefusion/processors/frame/modules') program = ArgumentParser(parents = [ program ], formatter_class = program.formatter_class, add_help = True) group_frame_processors = program.add_argument_group('frame processors') - group_frame_processors.add_argument('--frame-processors', help = wording.get('frame_processors_help').format(choices = ', '.join(available_frame_processors)), dest = 'frame_processors', default = [ 'face_swapper' ], nargs = '+') + group_frame_processors.add_argument('--frame-processors', help = wording.get('frame_processors_help').format(choices = ', '.join(available_frame_processors)), default = [ 'face_swapper' ], nargs = '+') for frame_processor in available_frame_processors: frame_processor_module = load_frame_processor_module(frame_processor) frame_processor_module.register_args(group_frame_processors) # uis group_uis = program.add_argument_group('uis') - group_uis.add_argument('--ui-layouts', help = wording.get('ui_layouts_help').format(choices = ', '.join(list_module_names('facefusion/uis/layouts'))), dest = 'ui_layouts', default = [ 'default' ], nargs = '+') + group_uis.add_argument('--ui-layouts', help = wording.get('ui_layouts_help').format(choices = ', '.join(list_module_names('facefusion/uis/layouts'))), default = [ 'default' ], nargs = '+') run(program) def apply_args(program : ArgumentParser) -> None: args = program.parse_args() # general - facefusion.globals.source_path = args.source_path + facefusion.globals.source_paths = args.source_paths facefusion.globals.target_path = args.target_path - facefusion.globals.output_path = normalize_output_path(facefusion.globals.source_path, facefusion.globals.target_path, args.output_path) + facefusion.globals.output_path = normalize_output_path(facefusion.globals.source_paths, facefusion.globals.target_path, args.output_path) # misc facefusion.globals.skip_download = args.skip_download facefusion.globals.headless = args.headless + facefusion.globals.log_level = args.log_level # execution facefusion.globals.execution_providers = decode_execution_providers(args.execution_providers) facefusion.globals.execution_thread_count = args.execution_thread_count @@ -116,8 +129,10 @@ def apply_args(program : ArgumentParser) -> None: facefusion.globals.reference_face_distance = args.reference_face_distance facefusion.globals.reference_frame_number = args.reference_frame_number # face mask + facefusion.globals.face_mask_types = args.face_mask_types facefusion.globals.face_mask_blur = args.face_mask_blur facefusion.globals.face_mask_padding = normalize_padding(args.face_mask_padding) + facefusion.globals.face_mask_regions = args.face_mask_regions # frame extraction facefusion.globals.trim_frame_start = args.trim_frame_start facefusion.globals.trim_frame_end = args.trim_frame_end @@ -142,8 +157,9 @@ def apply_args(program : ArgumentParser) -> None: def run(program : ArgumentParser) -> None: apply_args(program) + logger.init(facefusion.globals.log_level) limit_resources() - if not pre_check() or not content_analyser.pre_check() or not face_analyser.pre_check(): + if not pre_check() or not content_analyser.pre_check() or not face_analyser.pre_check() or not face_masker.pre_check(): return for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): if not frame_processor_module.pre_check(): @@ -172,25 +188,27 @@ def limit_resources() -> None: memory = facefusion.globals.max_memory * 1024 ** 6 if platform.system().lower() == 'windows': import ctypes + kernel32 = ctypes.windll.kernel32 # type: ignore[attr-defined] kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory)) else: import resource + resource.setrlimit(resource.RLIMIT_DATA, (memory, memory)) def pre_check() -> bool: if sys.version_info < (3, 9): - update_status(wording.get('python_not_supported').format(version = '3.9')) + logger.error(wording.get('python_not_supported').format(version = '3.9'), __name__.upper()) return False if not shutil.which('ffmpeg'): - update_status(wording.get('ffmpeg_not_installed')) + logger.error(wording.get('ffmpeg_not_installed'), __name__.upper()) return False return True def conditional_process() -> None: - conditional_set_face_reference() + conditional_append_reference_faces() for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): if not frame_processor_module.pre_process('output'): return @@ -200,14 +218,21 @@ def conditional_process() -> None: process_video() -def conditional_set_face_reference() -> None: - if 'reference' in facefusion.globals.face_selector_mode and not get_face_reference(): +def conditional_append_reference_faces() -> None: + if 'reference' in facefusion.globals.face_selector_mode and not get_reference_faces(): + source_frames = read_static_images(facefusion.globals.source_paths) + source_face = get_average_face(source_frames) if is_video(facefusion.globals.target_path): reference_frame = get_video_frame(facefusion.globals.target_path, facefusion.globals.reference_frame_number) else: reference_frame = read_image(facefusion.globals.target_path) reference_face = get_one_face(reference_frame, facefusion.globals.reference_face_position) - set_face_reference(reference_face) + append_reference_face('origin', reference_face) + if source_face and reference_face: + for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): + reference_frame = frame_processor_module.get_reference_frame(source_face, reference_face, reference_frame) + reference_face = get_one_face(reference_frame, facefusion.globals.reference_face_position) + append_reference_face(frame_processor_module.__name__, reference_face) def process_image() -> None: @@ -216,18 +241,18 @@ def process_image() -> None: shutil.copy2(facefusion.globals.target_path, facefusion.globals.output_path) # process frame for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): - update_status(wording.get('processing'), frame_processor_module.NAME) - frame_processor_module.process_image(facefusion.globals.source_path, facefusion.globals.output_path, facefusion.globals.output_path) + logger.info(wording.get('processing'), frame_processor_module.NAME) + frame_processor_module.process_image(facefusion.globals.source_paths, facefusion.globals.output_path, facefusion.globals.output_path) frame_processor_module.post_process() # compress image - update_status(wording.get('compressing_image')) + logger.info(wording.get('compressing_image'), __name__.upper()) if not compress_image(facefusion.globals.output_path): - update_status(wording.get('compressing_image_failed')) + logger.error(wording.get('compressing_image_failed'), __name__.upper()) # validate image if is_image(facefusion.globals.output_path): - update_status(wording.get('processing_image_succeed')) + logger.info(wording.get('processing_image_succeed'), __name__.upper()) else: - update_status(wording.get('processing_image_failed')) + logger.error(wording.get('processing_image_failed'), __name__.upper()) def process_video() -> None: @@ -235,40 +260,40 @@ def process_video() -> None: return fps = detect_fps(facefusion.globals.target_path) if facefusion.globals.keep_fps else 25.0 # create temp - update_status(wording.get('creating_temp')) + logger.info(wording.get('creating_temp'), __name__.upper()) create_temp(facefusion.globals.target_path) # extract frames - update_status(wording.get('extracting_frames_fps').format(fps = fps)) + logger.info(wording.get('extracting_frames_fps').format(fps = fps), __name__.upper()) extract_frames(facefusion.globals.target_path, fps) # process frame temp_frame_paths = get_temp_frame_paths(facefusion.globals.target_path) if temp_frame_paths: for frame_processor_module in get_frame_processors_modules(facefusion.globals.frame_processors): - update_status(wording.get('processing'), frame_processor_module.NAME) - frame_processor_module.process_video(facefusion.globals.source_path, temp_frame_paths) + logger.info(wording.get('processing'), frame_processor_module.NAME) + frame_processor_module.process_video(facefusion.globals.source_paths, temp_frame_paths) frame_processor_module.post_process() else: - update_status(wording.get('temp_frames_not_found')) + logger.error(wording.get('temp_frames_not_found'), __name__.upper()) return # merge video - update_status(wording.get('merging_video_fps').format(fps = fps)) + logger.info(wording.get('merging_video_fps').format(fps = fps), __name__.upper()) if not merge_video(facefusion.globals.target_path, fps): - update_status(wording.get('merging_video_failed')) + logger.error(wording.get('merging_video_failed'), __name__.upper()) return # handle audio if facefusion.globals.skip_audio: - update_status(wording.get('skipping_audio')) + logger.info(wording.get('skipping_audio'), __name__.upper()) move_temp(facefusion.globals.target_path, facefusion.globals.output_path) else: - update_status(wording.get('restoring_audio')) + logger.info(wording.get('restoring_audio'), __name__.upper()) if not restore_audio(facefusion.globals.target_path, facefusion.globals.output_path): - update_status(wording.get('restoring_audio_failed')) + logger.warn(wording.get('restoring_audio_skipped'), __name__.upper()) move_temp(facefusion.globals.target_path, facefusion.globals.output_path) # clear temp - update_status(wording.get('clearing_temp')) + logger.info(wording.get('clearing_temp'), __name__.upper()) clear_temp(facefusion.globals.target_path) # validate video if is_video(facefusion.globals.output_path): - update_status(wording.get('processing_video_succeed')) + logger.info(wording.get('processing_video_succeed'), __name__.upper()) else: - update_status(wording.get('processing_video_failed')) + logger.error(wording.get('processing_video_failed'), __name__.upper()) diff --git a/facefusion/download.py b/facefusion/download.py new file mode 100644 index 00000000..d50935f2 --- /dev/null +++ b/facefusion/download.py @@ -0,0 +1,44 @@ +import os +import subprocess +import urllib.request +from typing import List +from concurrent.futures import ThreadPoolExecutor +from functools import lru_cache +from tqdm import tqdm + +import facefusion.globals +from facefusion import wording +from facefusion.filesystem import is_file + + +def conditional_download(download_directory_path : str, urls : List[str]) -> None: + with ThreadPoolExecutor() as executor: + for url in urls: + executor.submit(get_download_size, url) + for url in urls: + download_file_path = os.path.join(download_directory_path, os.path.basename(url)) + initial = os.path.getsize(download_file_path) if is_file(download_file_path) else 0 + total = get_download_size(url) + if initial < total: + with tqdm(total = total, initial = initial, desc = wording.get('downloading'), unit = 'B', unit_scale = True, unit_divisor = 1024, ascii = ' =', disable = facefusion.globals.log_level in [ 'warn', 'error' ]) as progress: + subprocess.Popen([ 'curl', '--create-dirs', '--silent', '--insecure', '--location', '--continue-at', '-', '--output', download_file_path, url ]) + current = initial + while current < total: + if is_file(download_file_path): + current = os.path.getsize(download_file_path) + progress.update(current - progress.n) + + +@lru_cache(maxsize = None) +def get_download_size(url : str) -> int: + try: + response = urllib.request.urlopen(url, timeout = 10) + return int(response.getheader('Content-Length')) + except (OSError, ValueError): + return 0 + + +def is_download_done(url : str, file_path : str) -> bool: + if is_file(file_path): + return get_download_size(url) == os.path.getsize(file_path) + return False diff --git a/facefusion/execution_helper.py b/facefusion/execution_helper.py new file mode 100644 index 00000000..9c66865a --- /dev/null +++ b/facefusion/execution_helper.py @@ -0,0 +1,22 @@ +from typing import List +import onnxruntime + + +def encode_execution_providers(execution_providers : List[str]) -> List[str]: + return [ execution_provider.replace('ExecutionProvider', '').lower() for execution_provider in execution_providers ] + + +def decode_execution_providers(execution_providers: List[str]) -> List[str]: + available_execution_providers = onnxruntime.get_available_providers() + encoded_execution_providers = encode_execution_providers(available_execution_providers) + return [ execution_provider for execution_provider, encoded_execution_provider in zip(available_execution_providers, encoded_execution_providers) if any(execution_provider in encoded_execution_provider for execution_provider in execution_providers) ] + + +def map_device(execution_providers : List[str]) -> str: + if 'CoreMLExecutionProvider' in execution_providers: + return 'mps' + if 'CUDAExecutionProvider' in execution_providers or 'ROCMExecutionProvider' in execution_providers : + return 'cuda' + if 'OpenVINOExecutionProvider' in execution_providers: + return 'mkl' + return 'cpu' diff --git a/facefusion/face_analyser.py b/facefusion/face_analyser.py index d45aadf5..06960e4e 100644 --- a/facefusion/face_analyser.py +++ b/facefusion/face_analyser.py @@ -1,20 +1,21 @@ -from typing import Any, Optional, List, Dict, Tuple +from typing import Any, Optional, List, Tuple import threading import cv2 import numpy import onnxruntime import facefusion.globals -from facefusion.face_cache import get_faces_cache, set_faces_cache +from facefusion.download import conditional_download +from facefusion.face_store import get_static_faces, set_static_faces from facefusion.face_helper import warp_face, create_static_anchors, distance_to_kps, distance_to_bbox, apply_nms -from facefusion.typing import Frame, Face, FaceAnalyserOrder, FaceAnalyserAge, FaceAnalyserGender, ModelValue, Bbox, Kps, Score, Embedding -from facefusion.utilities import resolve_relative_path, conditional_download +from facefusion.filesystem import resolve_relative_path +from facefusion.typing import Frame, Face, FaceSet, FaceAnalyserOrder, FaceAnalyserAge, FaceAnalyserGender, ModelSet, Bbox, Kps, Score, Embedding from facefusion.vision import resize_frame_dimension FACE_ANALYSER = None THREAD_SEMAPHORE : threading.Semaphore = threading.Semaphore() THREAD_LOCK : threading.Lock = threading.Lock() -MODELS : Dict[str, ModelValue] =\ +MODELS : ModelSet =\ { 'face_detector_retinaface': { @@ -174,9 +175,13 @@ def detect_with_yunet(temp_frame : Frame, temp_frame_height : int, temp_frame_wi return bbox_list, kps_list, score_list -def create_faces(frame : Frame, bbox_list : List[Bbox], kps_list : List[Kps], score_list : List[Score]) -> List[Face] : - faces : List[Face] = [] +def create_faces(frame : Frame, bbox_list : List[Bbox], kps_list : List[Kps], score_list : List[Score]) -> List[Face]: + faces = [] if facefusion.globals.face_detector_score > 0: + sort_indices = numpy.argsort(-numpy.array(score_list)) + bbox_list = [ bbox_list[index] for index in sort_indices ] + kps_list = [ kps_list[index] for index in sort_indices ] + score_list = [ score_list[index] for index in sort_indices ] keep_indices = apply_nms(bbox_list, 0.4) for index in keep_indices: bbox = bbox_list[index] @@ -198,7 +203,7 @@ def create_faces(frame : Frame, bbox_list : List[Bbox], kps_list : List[Kps], sc def calc_embedding(temp_frame : Frame, kps : Kps) -> Tuple[Embedding, Embedding]: face_recognizer = get_face_analyser().get('face_recognizer') - crop_frame, matrix = warp_face(temp_frame, kps, 'arcface_v2', (112, 112)) + crop_frame, matrix = warp_face(temp_frame, kps, 'arcface_112_v2', (112, 112)) crop_frame = crop_frame.astype(numpy.float32) / 127.5 - 1 crop_frame = crop_frame[:, :, ::-1].transpose(2, 0, 1) crop_frame = numpy.expand_dims(crop_frame, axis = 0) @@ -213,7 +218,7 @@ def calc_embedding(temp_frame : Frame, kps : Kps) -> Tuple[Embedding, Embedding] def detect_gender_age(frame : Frame, kps : Kps) -> Tuple[int, int]: gender_age = get_face_analyser().get('gender_age') - crop_frame, affine_matrix = warp_face(frame, kps, 'arcface_v2', (96, 96)) + crop_frame, affine_matrix = warp_face(frame, kps, 'arcface_112_v2', (96, 96)) crop_frame = numpy.expand_dims(crop_frame, axis = 0).transpose(0, 3, 1, 2).astype(numpy.float32) prediction = gender_age.run(None, { @@ -234,14 +239,38 @@ def get_one_face(frame : Frame, position : int = 0) -> Optional[Face]: return None +def get_average_face(frames : List[Frame], position : int = 0) -> Optional[Face]: + average_face = None + faces = [] + embedding_list = [] + normed_embedding_list = [] + for frame in frames: + face = get_one_face(frame, position) + if face: + faces.append(face) + embedding_list.append(face.embedding) + normed_embedding_list.append(face.normed_embedding) + if faces: + average_face = Face( + bbox = faces[0].bbox, + kps = faces[0].kps, + score = faces[0].score, + embedding = numpy.mean(embedding_list, axis = 0), + normed_embedding = numpy.mean(normed_embedding_list, axis = 0), + gender = faces[0].gender, + age = faces[0].age + ) + return average_face + + def get_many_faces(frame : Frame) -> List[Face]: try: - faces_cache = get_faces_cache(frame) + faces_cache = get_static_faces(frame) if faces_cache: faces = faces_cache else: faces = extract_faces(frame) - set_faces_cache(frame, faces) + set_static_faces(frame, faces) if facefusion.globals.face_analyser_order: faces = sort_by_order(faces, facefusion.globals.face_analyser_order) if facefusion.globals.face_analyser_age: @@ -253,18 +282,27 @@ def get_many_faces(frame : Frame) -> List[Face]: return [] -def find_similar_faces(frame : Frame, reference_face : Face, face_distance : float) -> List[Face]: +def find_similar_faces(frame : Frame, reference_faces : FaceSet, face_distance : float) -> List[Face]: + similar_faces : List[Face] = [] many_faces = get_many_faces(frame) - similar_faces = [] - if many_faces: - for face in many_faces: - if hasattr(face, 'normed_embedding') and hasattr(reference_face, 'normed_embedding'): - current_face_distance = 1 - numpy.dot(face.normed_embedding, reference_face.normed_embedding) - if current_face_distance < face_distance: - similar_faces.append(face) + + if reference_faces: + for reference_set in reference_faces: + if not similar_faces: + for reference_face in reference_faces[reference_set]: + for face in many_faces: + if compare_faces(face, reference_face, face_distance): + similar_faces.append(face) return similar_faces +def compare_faces(face : Face, reference_face : Face, face_distance : float) -> bool: + if hasattr(face, 'normed_embedding') and hasattr(reference_face, 'normed_embedding'): + current_face_distance = 1 - numpy.dot(face.normed_embedding, reference_face.normed_embedding) + return current_face_distance < face_distance + return False + + def sort_by_order(faces : List[Face], order : FaceAnalyserOrder) -> List[Face]: if order == 'left-right': return sorted(faces, key = lambda face: face.bbox[0]) diff --git a/facefusion/face_cache.py b/facefusion/face_cache.py deleted file mode 100644 index 8730509a..00000000 --- a/facefusion/face_cache.py +++ /dev/null @@ -1,29 +0,0 @@ -from typing import Optional, List, Dict -import hashlib - -from facefusion.typing import Frame, Face - -FACES_CACHE : Dict[str, List[Face]] = {} - - -def get_faces_cache(frame : Frame) -> Optional[List[Face]]: - frame_hash = create_frame_hash(frame) - if frame_hash in FACES_CACHE: - return FACES_CACHE[frame_hash] - return None - - -def set_faces_cache(frame : Frame, faces : List[Face]) -> None: - frame_hash = create_frame_hash(frame) - if frame_hash: - FACES_CACHE[frame_hash] = faces - - -def clear_faces_cache() -> None: - global FACES_CACHE - - FACES_CACHE = {} - - -def create_frame_hash(frame : Frame) -> Optional[str]: - return hashlib.sha1(frame.tobytes()).hexdigest() if frame.any() else None diff --git a/facefusion/face_helper.py b/facefusion/face_helper.py index b635f944..ce7940fd 100644 --- a/facefusion/face_helper.py +++ b/facefusion/face_helper.py @@ -1,14 +1,14 @@ from typing import Any, Dict, Tuple, List -from functools import lru_cache from cv2.typing import Size +from functools import lru_cache import cv2 import numpy -from facefusion.typing import Bbox, Kps, Frame, Matrix, Template, Padding +from facefusion.typing import Bbox, Kps, Frame, Mask, Matrix, Template TEMPLATES : Dict[Template, numpy.ndarray[Any, Any]] =\ { - 'arcface_v1': numpy.array( + 'arcface_112_v1': numpy.array( [ [ 39.7300, 51.1380 ], [ 72.2700, 51.1380 ], @@ -16,7 +16,7 @@ TEMPLATES : Dict[Template, numpy.ndarray[Any, Any]] =\ [ 42.4630, 87.0100 ], [ 69.5370, 87.0100 ] ]), - 'arcface_v2': numpy.array( + 'arcface_112_v2': numpy.array( [ [ 38.2946, 51.6963 ], [ 73.5318, 51.5014 ], @@ -24,7 +24,15 @@ TEMPLATES : Dict[Template, numpy.ndarray[Any, Any]] =\ [ 41.5493, 92.3655 ], [ 70.7299, 92.2041 ] ]), - 'ffhq': numpy.array( + 'arcface_128_v2': numpy.array( + [ + [ 46.2946, 51.6963 ], + [ 81.5318, 51.5014 ], + [ 64.0252, 71.7366 ], + [ 49.5493, 92.3655 ], + [ 78.7299, 92.2041 ] + ]), + 'ffhq_512': numpy.array( [ [ 192.98138, 239.94708 ], [ 318.90277, 240.1936 ], @@ -37,39 +45,23 @@ TEMPLATES : Dict[Template, numpy.ndarray[Any, Any]] =\ def warp_face(temp_frame : Frame, kps : Kps, template : Template, size : Size) -> Tuple[Frame, Matrix]: normed_template = TEMPLATES.get(template) * size[1] / size[0] - affine_matrix = cv2.estimateAffinePartial2D(kps, normed_template, method = cv2.LMEDS)[0] + affine_matrix = cv2.estimateAffinePartial2D(kps, normed_template, method = cv2.RANSAC, ransacReprojThreshold = 100)[0] crop_frame = cv2.warpAffine(temp_frame, affine_matrix, (size[1], size[1]), borderMode = cv2.BORDER_REPLICATE) return crop_frame, affine_matrix -def paste_back(temp_frame : Frame, crop_frame: Frame, affine_matrix : Matrix, face_mask_blur : float, face_mask_padding : Padding) -> Frame: +def paste_back(temp_frame : Frame, crop_frame: Frame, crop_mask : Mask, affine_matrix : Matrix) -> Frame: inverse_matrix = cv2.invertAffineTransform(affine_matrix) temp_frame_size = temp_frame.shape[:2][::-1] - mask_size = tuple(crop_frame.shape[:2]) - mask_frame = create_static_mask_frame(mask_size, face_mask_blur, face_mask_padding) - inverse_mask_frame = cv2.warpAffine(mask_frame, inverse_matrix, temp_frame_size).clip(0, 1) + inverse_crop_mask = cv2.warpAffine(crop_mask, inverse_matrix, temp_frame_size).clip(0, 1) inverse_crop_frame = cv2.warpAffine(crop_frame, inverse_matrix, temp_frame_size, borderMode = cv2.BORDER_REPLICATE) paste_frame = temp_frame.copy() - paste_frame[:, :, 0] = inverse_mask_frame * inverse_crop_frame[:, :, 0] + (1 - inverse_mask_frame) * temp_frame[:, :, 0] - paste_frame[:, :, 1] = inverse_mask_frame * inverse_crop_frame[:, :, 1] + (1 - inverse_mask_frame) * temp_frame[:, :, 1] - paste_frame[:, :, 2] = inverse_mask_frame * inverse_crop_frame[:, :, 2] + (1 - inverse_mask_frame) * temp_frame[:, :, 2] + paste_frame[:, :, 0] = inverse_crop_mask * inverse_crop_frame[:, :, 0] + (1 - inverse_crop_mask) * temp_frame[:, :, 0] + paste_frame[:, :, 1] = inverse_crop_mask * inverse_crop_frame[:, :, 1] + (1 - inverse_crop_mask) * temp_frame[:, :, 1] + paste_frame[:, :, 2] = inverse_crop_mask * inverse_crop_frame[:, :, 2] + (1 - inverse_crop_mask) * temp_frame[:, :, 2] return paste_frame -@lru_cache(maxsize = None) -def create_static_mask_frame(mask_size : Size, face_mask_blur : float, face_mask_padding : Padding) -> Frame: - mask_frame = numpy.ones(mask_size, numpy.float32) - blur_amount = int(mask_size[0] * 0.5 * face_mask_blur) - blur_area = max(blur_amount // 2, 1) - mask_frame[:max(blur_area, int(mask_size[1] * face_mask_padding[0] / 100)), :] = 0 - mask_frame[-max(blur_area, int(mask_size[1] * face_mask_padding[2] / 100)):, :] = 0 - mask_frame[:, :max(blur_area, int(mask_size[0] * face_mask_padding[3] / 100))] = 0 - mask_frame[:, -max(blur_area, int(mask_size[0] * face_mask_padding[1] / 100)):] = 0 - if blur_amount > 0: - mask_frame = cv2.GaussianBlur(mask_frame, (0, 0), blur_amount * 0.25) - return mask_frame - - @lru_cache(maxsize = None) def create_static_anchors(feature_stride : int, anchor_total : int, stride_height : int, stride_width : int) -> numpy.ndarray[Any, Any]: y, x = numpy.mgrid[:stride_height, :stride_width][::-1] diff --git a/facefusion/face_masker.py b/facefusion/face_masker.py new file mode 100755 index 00000000..96d877b7 --- /dev/null +++ b/facefusion/face_masker.py @@ -0,0 +1,128 @@ +from typing import Any, Dict, List +from cv2.typing import Size +from functools import lru_cache +import threading +import cv2 +import numpy +import onnxruntime + +import facefusion.globals +from facefusion.typing import Frame, Mask, Padding, FaceMaskRegion, ModelSet +from facefusion.filesystem import resolve_relative_path +from facefusion.download import conditional_download + +FACE_OCCLUDER = None +FACE_PARSER = None +THREAD_LOCK : threading.Lock = threading.Lock() +MODELS : ModelSet =\ +{ + 'face_occluder': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/face_occluder.onnx', + 'path': resolve_relative_path('../.assets/models/face_occluder.onnx') + }, + 'face_parser': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/face_parser.onnx', + 'path': resolve_relative_path('../.assets/models/face_parser.onnx') + } +} +FACE_MASK_REGIONS : Dict[FaceMaskRegion, int] =\ +{ + 'skin': 1, + 'left-eyebrow': 2, + 'right-eyebrow': 3, + 'left-eye': 4, + 'right-eye': 5, + 'eye-glasses': 6, + 'nose': 10, + 'mouth': 11, + 'upper-lip': 12, + 'lower-lip': 13 +} + + +def get_face_occluder() -> Any: + global FACE_OCCLUDER + + with THREAD_LOCK: + if FACE_OCCLUDER is None: + model_path = MODELS.get('face_occluder').get('path') + FACE_OCCLUDER = onnxruntime.InferenceSession(model_path, providers = facefusion.globals.execution_providers) + return FACE_OCCLUDER + + +def get_face_parser() -> Any: + global FACE_PARSER + + with THREAD_LOCK: + if FACE_PARSER is None: + model_path = MODELS.get('face_parser').get('path') + FACE_PARSER = onnxruntime.InferenceSession(model_path, providers = facefusion.globals.execution_providers) + return FACE_PARSER + + +def clear_face_occluder() -> None: + global FACE_OCCLUDER + + FACE_OCCLUDER = None + + +def clear_face_parser() -> None: + global FACE_PARSER + + FACE_PARSER = None + + +def pre_check() -> bool: + if not facefusion.globals.skip_download: + download_directory_path = resolve_relative_path('../.assets/models') + model_urls =\ + [ + MODELS.get('face_occluder').get('url'), + MODELS.get('face_parser').get('url'), + ] + conditional_download(download_directory_path, model_urls) + return True + + +@lru_cache(maxsize = None) +def create_static_box_mask(crop_size : Size, face_mask_blur : float, face_mask_padding : Padding) -> Mask: + blur_amount = int(crop_size[0] * 0.5 * face_mask_blur) + blur_area = max(blur_amount // 2, 1) + box_mask = numpy.ones(crop_size, numpy.float32) + box_mask[:max(blur_area, int(crop_size[1] * face_mask_padding[0] / 100)), :] = 0 + box_mask[-max(blur_area, int(crop_size[1] * face_mask_padding[2] / 100)):, :] = 0 + box_mask[:, :max(blur_area, int(crop_size[0] * face_mask_padding[3] / 100))] = 0 + box_mask[:, -max(blur_area, int(crop_size[0] * face_mask_padding[1] / 100)):] = 0 + if blur_amount > 0: + box_mask = cv2.GaussianBlur(box_mask, (0, 0), blur_amount * 0.25) + return box_mask + + +def create_occlusion_mask(crop_frame : Frame) -> Mask: + face_occluder = get_face_occluder() + prepare_frame = cv2.resize(crop_frame, face_occluder.get_inputs()[0].shape[1:3][::-1]) + prepare_frame = numpy.expand_dims(prepare_frame, axis = 0).astype(numpy.float32) / 255 + prepare_frame = prepare_frame.transpose(0, 1, 2, 3) + occlusion_mask = face_occluder.run(None, + { + face_occluder.get_inputs()[0].name: prepare_frame + })[0][0] + occlusion_mask = occlusion_mask.transpose(0, 1, 2).clip(0, 1).astype(numpy.float32) + occlusion_mask = cv2.resize(occlusion_mask, crop_frame.shape[:2][::-1]) + return occlusion_mask + + +def create_region_mask(crop_frame : Frame, face_mask_regions : List[FaceMaskRegion]) -> Mask: + face_parser = get_face_parser() + prepare_frame = cv2.flip(cv2.resize(crop_frame, (512, 512)), 1) + prepare_frame = numpy.expand_dims(prepare_frame, axis = 0).astype(numpy.float32)[:, :, ::-1] / 127.5 - 1 + prepare_frame = prepare_frame.transpose(0, 3, 1, 2) + region_mask = face_parser.run(None, + { + face_parser.get_inputs()[0].name: prepare_frame + })[0][0] + region_mask = numpy.isin(region_mask.argmax(0), [ FACE_MASK_REGIONS[region] for region in face_mask_regions ]) + region_mask = cv2.resize(region_mask.astype(numpy.float32), crop_frame.shape[:2][::-1]) + return region_mask diff --git a/facefusion/face_reference.py b/facefusion/face_reference.py deleted file mode 100644 index 72281fe6..00000000 --- a/facefusion/face_reference.py +++ /dev/null @@ -1,21 +0,0 @@ -from typing import Optional - -from facefusion.typing import Face - -FACE_REFERENCE = None - - -def get_face_reference() -> Optional[Face]: - return FACE_REFERENCE - - -def set_face_reference(face : Face) -> None: - global FACE_REFERENCE - - FACE_REFERENCE = face - - -def clear_face_reference() -> None: - global FACE_REFERENCE - - FACE_REFERENCE = None diff --git a/facefusion/face_store.py b/facefusion/face_store.py new file mode 100644 index 00000000..1f0dfa4d --- /dev/null +++ b/facefusion/face_store.py @@ -0,0 +1,47 @@ +from typing import Optional, List +import hashlib + +from facefusion.typing import Frame, Face, FaceStore, FaceSet + +FACE_STORE: FaceStore =\ +{ + 'static_faces': {}, + 'reference_faces': {} +} + + +def get_static_faces(frame : Frame) -> Optional[List[Face]]: + frame_hash = create_frame_hash(frame) + if frame_hash in FACE_STORE['static_faces']: + return FACE_STORE['static_faces'][frame_hash] + return None + + +def set_static_faces(frame : Frame, faces : List[Face]) -> None: + frame_hash = create_frame_hash(frame) + if frame_hash: + FACE_STORE['static_faces'][frame_hash] = faces + + +def clear_static_faces() -> None: + FACE_STORE['static_faces'] = {} + + +def create_frame_hash(frame: Frame) -> Optional[str]: + return hashlib.sha1(frame.tobytes()).hexdigest() if frame.any() else None + + +def get_reference_faces() -> Optional[FaceSet]: + if FACE_STORE['reference_faces']: + return FACE_STORE['reference_faces'] + return None + + +def append_reference_face(name : str, face : Face) -> None: + if name not in FACE_STORE['reference_faces']: + FACE_STORE['reference_faces'][name] = [] + FACE_STORE['reference_faces'][name].append(face) + + +def clear_reference_faces() -> None: + FACE_STORE['reference_faces'] = {} diff --git a/facefusion/ffmpeg.py b/facefusion/ffmpeg.py new file mode 100644 index 00000000..4cbb38e8 --- /dev/null +++ b/facefusion/ffmpeg.py @@ -0,0 +1,81 @@ +from typing import List +import subprocess + +import facefusion.globals +from facefusion import logger +from facefusion.filesystem import get_temp_frames_pattern, get_temp_output_video_path +from facefusion.vision import detect_fps + + +def run_ffmpeg(args : List[str]) -> bool: + commands = [ 'ffmpeg', '-hide_banner', '-loglevel', 'error' ] + commands.extend(args) + try: + subprocess.run(commands, stderr = subprocess.PIPE, check = True) + return True + except subprocess.CalledProcessError as exception: + logger.debug(exception.stderr.decode().strip(), __name__.upper()) + return False + + +def open_ffmpeg(args : List[str]) -> subprocess.Popen[bytes]: + commands = [ 'ffmpeg', '-hide_banner', '-loglevel', 'error' ] + commands.extend(args) + return subprocess.Popen(commands, stdin = subprocess.PIPE) + + +def extract_frames(target_path : str, fps : float) -> bool: + temp_frame_compression = round(31 - (facefusion.globals.temp_frame_quality * 0.31)) + trim_frame_start = facefusion.globals.trim_frame_start + trim_frame_end = facefusion.globals.trim_frame_end + temp_frames_pattern = get_temp_frames_pattern(target_path, '%04d') + commands = [ '-hwaccel', 'auto', '-i', target_path, '-q:v', str(temp_frame_compression), '-pix_fmt', 'rgb24' ] + if trim_frame_start is not None and trim_frame_end is not None: + commands.extend([ '-vf', 'trim=start_frame=' + str(trim_frame_start) + ':end_frame=' + str(trim_frame_end) + ',fps=' + str(fps) ]) + elif trim_frame_start is not None: + commands.extend([ '-vf', 'trim=start_frame=' + str(trim_frame_start) + ',fps=' + str(fps) ]) + elif trim_frame_end is not None: + commands.extend([ '-vf', 'trim=end_frame=' + str(trim_frame_end) + ',fps=' + str(fps) ]) + else: + commands.extend([ '-vf', 'fps=' + str(fps) ]) + commands.extend([ '-vsync', '0', temp_frames_pattern ]) + return run_ffmpeg(commands) + + +def compress_image(output_path : str) -> bool: + output_image_compression = round(31 - (facefusion.globals.output_image_quality * 0.31)) + commands = [ '-hwaccel', 'auto', '-i', output_path, '-q:v', str(output_image_compression), '-y', output_path ] + return run_ffmpeg(commands) + + +def merge_video(target_path : str, fps : float) -> bool: + temp_output_video_path = get_temp_output_video_path(target_path) + temp_frames_pattern = get_temp_frames_pattern(target_path, '%04d') + commands = [ '-hwaccel', 'auto', '-r', str(fps), '-i', temp_frames_pattern, '-c:v', facefusion.globals.output_video_encoder ] + if facefusion.globals.output_video_encoder in [ 'libx264', 'libx265' ]: + output_video_compression = round(51 - (facefusion.globals.output_video_quality * 0.51)) + commands.extend([ '-crf', str(output_video_compression) ]) + if facefusion.globals.output_video_encoder in [ 'libvpx-vp9' ]: + output_video_compression = round(63 - (facefusion.globals.output_video_quality * 0.63)) + commands.extend([ '-crf', str(output_video_compression) ]) + if facefusion.globals.output_video_encoder in [ 'h264_nvenc', 'hevc_nvenc' ]: + output_video_compression = round(51 - (facefusion.globals.output_video_quality * 0.51)) + commands.extend([ '-cq', str(output_video_compression) ]) + commands.extend([ '-pix_fmt', 'yuv420p', '-colorspace', 'bt709', '-y', temp_output_video_path ]) + return run_ffmpeg(commands) + + +def restore_audio(target_path : str, output_path : str) -> bool: + fps = detect_fps(target_path) + trim_frame_start = facefusion.globals.trim_frame_start + trim_frame_end = facefusion.globals.trim_frame_end + temp_output_video_path = get_temp_output_video_path(target_path) + commands = [ '-hwaccel', 'auto', '-i', temp_output_video_path ] + if trim_frame_start is not None: + start_time = trim_frame_start / fps + commands.extend([ '-ss', str(start_time) ]) + if trim_frame_end is not None: + end_time = trim_frame_end / fps + commands.extend([ '-to', str(end_time) ]) + commands.extend([ '-i', target_path, '-c', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-shortest', '-y', output_path ]) + return run_ffmpeg(commands) diff --git a/facefusion/filesystem.py b/facefusion/filesystem.py new file mode 100644 index 00000000..ce972819 --- /dev/null +++ b/facefusion/filesystem.py @@ -0,0 +1,91 @@ +from typing import List, Optional +import glob +import os +import shutil +import tempfile +import filetype +from pathlib import Path + +import facefusion.globals + +TEMP_DIRECTORY_PATH = os.path.join(tempfile.gettempdir(), 'facefusion') +TEMP_OUTPUT_VIDEO_NAME = 'temp.mp4' + + +def get_temp_frame_paths(target_path : str) -> List[str]: + temp_frames_pattern = get_temp_frames_pattern(target_path, '*') + return sorted(glob.glob(temp_frames_pattern)) + + +def get_temp_frames_pattern(target_path : str, temp_frame_prefix : str) -> str: + temp_directory_path = get_temp_directory_path(target_path) + return os.path.join(temp_directory_path, temp_frame_prefix + '.' + facefusion.globals.temp_frame_format) + + +def get_temp_directory_path(target_path : str) -> str: + target_name, _ = os.path.splitext(os.path.basename(target_path)) + return os.path.join(TEMP_DIRECTORY_PATH, target_name) + + +def get_temp_output_video_path(target_path : str) -> str: + temp_directory_path = get_temp_directory_path(target_path) + return os.path.join(temp_directory_path, TEMP_OUTPUT_VIDEO_NAME) + + +def create_temp(target_path : str) -> None: + temp_directory_path = get_temp_directory_path(target_path) + Path(temp_directory_path).mkdir(parents = True, exist_ok = True) + + +def move_temp(target_path : str, output_path : str) -> None: + temp_output_video_path = get_temp_output_video_path(target_path) + if is_file(temp_output_video_path): + if is_file(output_path): + os.remove(output_path) + shutil.move(temp_output_video_path, output_path) + + +def clear_temp(target_path : str) -> None: + temp_directory_path = get_temp_directory_path(target_path) + parent_directory_path = os.path.dirname(temp_directory_path) + if not facefusion.globals.keep_temp and is_directory(temp_directory_path): + shutil.rmtree(temp_directory_path) + if os.path.exists(parent_directory_path) and not os.listdir(parent_directory_path): + os.rmdir(parent_directory_path) + + +def is_file(file_path : str) -> bool: + return bool(file_path and os.path.isfile(file_path)) + + +def is_directory(directory_path : str) -> bool: + return bool(directory_path and os.path.isdir(directory_path)) + + +def is_image(image_path : str) -> bool: + if is_file(image_path): + return filetype.helpers.is_image(image_path) + return False + + +def are_images(image_paths : List[str]) -> bool: + if image_paths: + return all(is_image(image_path) for image_path in image_paths) + return False + + +def is_video(video_path : str) -> bool: + if is_file(video_path): + return filetype.helpers.is_video(video_path) + return False + + +def resolve_relative_path(path : str) -> str: + return os.path.abspath(os.path.join(os.path.dirname(__file__), path)) + + +def list_module_names(path : str) -> Optional[List[str]]: + if os.path.exists(path): + files = os.listdir(path) + return [ Path(file).stem for file in files if not Path(file).stem.startswith(('.', '__')) ] + return None diff --git a/facefusion/globals.py b/facefusion/globals.py index 7d3d7add..fe7aed3f 100755 --- a/facefusion/globals.py +++ b/facefusion/globals.py @@ -1,14 +1,15 @@ from typing import List, Optional -from facefusion.typing import FaceSelectorMode, FaceAnalyserOrder, FaceAnalyserAge, FaceAnalyserGender, OutputVideoEncoder, FaceDetectorModel, FaceRecognizerModel, TempFrameFormat, Padding +from facefusion.typing import LogLevel, FaceSelectorMode, FaceAnalyserOrder, FaceAnalyserAge, FaceAnalyserGender, FaceMaskType, FaceMaskRegion, OutputVideoEncoder, FaceDetectorModel, FaceRecognizerModel, TempFrameFormat, Padding # general -source_path : Optional[str] = None +source_paths : Optional[List[str]] = None target_path : Optional[str] = None output_path : Optional[str] = None # misc skip_download : Optional[bool] = None headless : Optional[bool] = None +log_level : Optional[LogLevel] = None # execution execution_providers : List[str] = [] execution_thread_count : Optional[int] = None @@ -28,8 +29,10 @@ reference_face_position : Optional[int] = None reference_face_distance : Optional[float] = None reference_frame_number : Optional[int] = None # face mask +face_mask_types : Optional[List[FaceMaskType]] = None face_mask_blur : Optional[float] = None face_mask_padding : Optional[Padding] = None +face_mask_regions : Optional[List[FaceMaskRegion]] = None # frame extraction trim_frame_start : Optional[int] = None trim_frame_end : Optional[int] = None diff --git a/facefusion/installer.py b/facefusion/installer.py index dfd17dab..290ef4ee 100644 --- a/facefusion/installer.py +++ b/facefusion/installer.py @@ -1,4 +1,8 @@ from typing import Dict, Tuple +import sys +import os +import platform +import tempfile import subprocess from argparse import ArgumentParser, HelpFormatter @@ -11,32 +15,40 @@ from facefusion import metadata, wording TORCH : Dict[str, str] =\ { 'default': 'default', - 'cpu': 'cpu', - 'cuda': 'cu118', - 'rocm': 'rocm5.6' + 'cpu': 'cpu' } ONNXRUNTIMES : Dict[str, Tuple[str, str]] =\ { - 'default': ('onnxruntime', '1.16.3'), - 'cuda': ('onnxruntime-gpu', '1.16.3'), - 'coreml-legacy': ('onnxruntime-coreml', '1.13.1'), - 'coreml-silicon': ('onnxruntime-silicon', '1.16.0'), - 'directml': ('onnxruntime-directml', '1.16.3'), - 'openvino': ('onnxruntime-openvino', '1.16.0') + 'default': ('onnxruntime', '1.16.3') } +if platform.system().lower() == 'linux' or platform.system().lower() == 'windows': + TORCH['cuda'] = 'cu118' + ONNXRUNTIMES['cuda'] = ('onnxruntime-gpu', '1.16.3') + ONNXRUNTIMES['openvino'] = ('onnxruntime-openvino', '1.16.0') +if platform.system().lower() == 'linux': + TORCH['rocm'] = 'rocm5.6' + ONNXRUNTIMES['directml'] = ('onnxruntime-directml', '1.16.3') + ONNXRUNTIMES['rocm'] = ('onnxruntime-rocm', '1.16.3') +if platform.system().lower() == 'darwin': + ONNXRUNTIMES['coreml-legacy'] = ('onnxruntime-coreml', '1.13.1') + ONNXRUNTIMES['coreml-silicon'] = ('onnxruntime-silicon', '1.16.0') def cli() -> None: program = ArgumentParser(formatter_class = lambda prog: HelpFormatter(prog, max_help_position = 120)) - program.add_argument('--torch', help = wording.get('install_dependency_help').format(dependency = 'torch'), dest = 'torch', choices = TORCH.keys()) - program.add_argument('--onnxruntime', help = wording.get('install_dependency_help').format(dependency = 'onnxruntime'), dest = 'onnxruntime', choices = ONNXRUNTIMES.keys()) + program.add_argument('--torch', help = wording.get('install_dependency_help').format(dependency = 'torch'), choices = TORCH.keys()) + program.add_argument('--onnxruntime', help = wording.get('install_dependency_help').format(dependency = 'onnxruntime'), choices = ONNXRUNTIMES.keys()) + program.add_argument('--skip-venv', help = wording.get('skip_venv_help'), action = 'store_true') program.add_argument('-v', '--version', version = metadata.get('name') + ' ' + metadata.get('version'), action = 'version') run(program) def run(program : ArgumentParser) -> None: args = program.parse_args() + python_id = 'cp' + str(sys.version_info.major) + str(sys.version_info.minor) + if not args.skip_venv: + os.environ['PIP_REQUIRE_VIRTUALENV'] = '1' if args.torch and args.onnxruntime: answers =\ { @@ -54,10 +66,19 @@ def run(program : ArgumentParser) -> None: torch_wheel = TORCH[torch] onnxruntime = answers['onnxruntime'] onnxruntime_name, onnxruntime_version = ONNXRUNTIMES[onnxruntime] - subprocess.call([ 'pip', 'uninstall', 'torch', '-y' ]) + subprocess.call([ 'pip', 'uninstall', 'torch', '-y', '-q' ]) if torch_wheel == 'default': subprocess.call([ 'pip', 'install', '-r', 'requirements.txt' ]) else: subprocess.call([ 'pip', 'install', '-r', 'requirements.txt', '--extra-index-url', 'https://download.pytorch.org/whl/' + torch_wheel ]) - subprocess.call([ 'pip', 'uninstall', 'onnxruntime', onnxruntime_name, '-y' ]) - subprocess.call([ 'pip', 'install', onnxruntime_name + '==' + onnxruntime_version ]) + if onnxruntime != 'rocm': + subprocess.call([ 'pip', 'uninstall', 'onnxruntime', onnxruntime_name, '-y', '-q' ]) + subprocess.call([ 'pip', 'install', onnxruntime_name + '==' + onnxruntime_version ]) + elif python_id in [ 'cp39', 'cp310', 'cp311' ]: + wheel_name = 'onnxruntime_training-' + onnxruntime_version + '+rocm56-' + python_id + '-' + python_id + '-manylinux_2_17_x86_64.manylinux2014_x86_64.whl' + wheel_path = os.path.join(tempfile.gettempdir(), wheel_name) + wheel_url = 'https://download.onnxruntime.ai/' + wheel_name + subprocess.call([ 'curl', '--silent', '--location', '--continue-at', '-', '--output', wheel_path, wheel_url ]) + subprocess.call([ 'pip', 'uninstall', wheel_path, '-y', '-q' ]) + subprocess.call([ 'pip', 'install', wheel_path ]) + os.remove(wheel_path) diff --git a/facefusion/logger.py b/facefusion/logger.py new file mode 100644 index 00000000..56fe3679 --- /dev/null +++ b/facefusion/logger.py @@ -0,0 +1,39 @@ +from typing import Dict +from logging import basicConfig, getLogger, Logger, DEBUG, INFO, WARNING, ERROR + +from facefusion.typing import LogLevel + + +def init(log_level : LogLevel) -> None: + basicConfig(format = None) + get_package_logger().setLevel(get_log_levels()[log_level]) + + +def get_package_logger() -> Logger: + return getLogger('facefusion') + + +def debug(message : str, scope : str) -> None: + get_package_logger().debug('[' + scope + '] ' + message) + + +def info(message : str, scope : str) -> None: + get_package_logger().info('[' + scope + '] ' + message) + + +def warn(message : str, scope : str) -> None: + get_package_logger().warning('[' + scope + '] ' + message) + + +def error(message : str, scope : str) -> None: + get_package_logger().error('[' + scope + '] ' + message) + + +def get_log_levels() -> Dict[LogLevel, int]: + return\ + { + 'error': ERROR, + 'warn': WARNING, + 'info': INFO, + 'debug': DEBUG + } diff --git a/facefusion/metadata.py b/facefusion/metadata.py index 6e167342..a6e7da14 100644 --- a/facefusion/metadata.py +++ b/facefusion/metadata.py @@ -2,7 +2,7 @@ METADATA =\ { 'name': 'FaceFusion', 'description': 'Next generation face swapper and enhancer', - 'version': '2.0.0', + 'version': '2.1.0', 'license': 'MIT', 'author': 'Henry Ruhs', 'url': 'https://facefusion.io' diff --git a/facefusion/normalizer.py b/facefusion/normalizer.py new file mode 100644 index 00000000..eee93e57 --- /dev/null +++ b/facefusion/normalizer.py @@ -0,0 +1,34 @@ +from typing import List, Optional +import os + +from facefusion.filesystem import is_file, is_directory +from facefusion.typing import Padding + + +def normalize_output_path(source_paths : List[str], target_path : str, output_path : str) -> Optional[str]: + if is_file(target_path) and is_directory(output_path): + target_name, target_extension = os.path.splitext(os.path.basename(target_path)) + if source_paths and is_file(source_paths[0]): + source_name, _ = os.path.splitext(os.path.basename(source_paths[0])) + return os.path.join(output_path, source_name + '-' + target_name + target_extension) + return os.path.join(output_path, target_name + target_extension) + if is_file(target_path) and output_path: + _, target_extension = os.path.splitext(os.path.basename(target_path)) + output_name, output_extension = os.path.splitext(os.path.basename(output_path)) + output_directory_path = os.path.dirname(output_path) + if is_directory(output_directory_path) and output_extension: + return os.path.join(output_directory_path, output_name + target_extension) + return None + return output_path + + +def normalize_padding(padding : Optional[List[int]]) -> Optional[Padding]: + if padding and len(padding) == 1: + return tuple([ padding[0], padding[0], padding[0], padding[0] ]) # type: ignore[return-value] + if padding and len(padding) == 2: + return tuple([ padding[0], padding[1], padding[0], padding[1] ]) # type: ignore[return-value] + if padding and len(padding) == 3: + return tuple([ padding[0], padding[1], padding[2], padding[1] ]) # type: ignore[return-value] + if padding and len(padding) == 4: + return tuple(padding) # type: ignore[return-value] + return None diff --git a/facefusion/processors/frame/core.py b/facefusion/processors/frame/core.py index 294ffe7b..fdbbe883 100644 --- a/facefusion/processors/frame/core.py +++ b/facefusion/processors/frame/core.py @@ -8,8 +8,8 @@ from tqdm import tqdm import facefusion.globals from facefusion.typing import Process_Frames -from facefusion import wording -from facefusion.utilities import encode_execution_providers +from facefusion.execution_helper import encode_execution_providers +from facefusion import logger, wording FRAME_PROCESSORS_MODULES : List[ModuleType] = [] FRAME_PROCESSORS_METHODS =\ @@ -22,6 +22,7 @@ FRAME_PROCESSORS_METHODS =\ 'apply_args', 'pre_check', 'pre_process', + 'get_reference_frame', 'process_frame', 'process_frames', 'process_image', @@ -36,7 +37,8 @@ def load_frame_processor_module(frame_processor : str) -> Any: for method_name in FRAME_PROCESSORS_METHODS: if not hasattr(frame_processor_module, method_name): raise NotImplementedError - except ModuleNotFoundError: + except ModuleNotFoundError as exception: + logger.debug(exception.msg, __name__.upper()) sys.exit(wording.get('frame_processor_not_loaded').format(frame_processor = frame_processor)) except NotImplementedError: sys.exit(wording.get('frame_processor_not_implemented').format(frame_processor = frame_processor)) @@ -61,8 +63,8 @@ def clear_frame_processors_modules() -> None: FRAME_PROCESSORS_MODULES = [] -def multi_process_frames(source_path : str, temp_frame_paths : List[str], process_frames : Process_Frames) -> None: - with tqdm(total = len(temp_frame_paths), desc = wording.get('processing'), unit = 'frame', ascii = ' =') as progress: +def multi_process_frames(source_paths : List[str], temp_frame_paths : List[str], process_frames : Process_Frames) -> None: + with tqdm(total = len(temp_frame_paths), desc = wording.get('processing'), unit = 'frame', ascii = ' =', disable = facefusion.globals.log_level in [ 'warn', 'error' ]) as progress: progress.set_postfix( { 'execution_providers': encode_execution_providers(facefusion.globals.execution_providers), @@ -75,7 +77,7 @@ def multi_process_frames(source_path : str, temp_frame_paths : List[str], proces queue_per_future = max(len(temp_frame_paths) // facefusion.globals.execution_thread_count * facefusion.globals.execution_queue_count, 1) while not queue_temp_frame_paths.empty(): payload_temp_frame_paths = pick_queue(queue_temp_frame_paths, queue_per_future) - future = executor.submit(process_frames, source_path, payload_temp_frame_paths, progress.update) + future = executor.submit(process_frames, source_paths, payload_temp_frame_paths, progress.update) futures.append(future) for future_done in as_completed(futures): future_done.result() diff --git a/facefusion/processors/frame/modules/face_debugger.py b/facefusion/processors/frame/modules/face_debugger.py index 75477e5d..9ec4e9e4 100755 --- a/facefusion/processors/frame/modules/face_debugger.py +++ b/facefusion/processors/frame/modules/face_debugger.py @@ -6,15 +6,16 @@ import numpy import facefusion.globals import facefusion.processors.frame.core as frame_processors from facefusion import wording -from facefusion.face_analyser import get_one_face, get_many_faces, find_similar_faces, clear_face_analyser -from facefusion.face_reference import get_face_reference +from facefusion.face_analyser import get_one_face, get_average_face, get_many_faces, find_similar_faces, clear_face_analyser +from facefusion.face_store import get_reference_faces from facefusion.content_analyser import clear_content_analyser -from facefusion.typing import Face, Frame, Update_Process, ProcessMode -from facefusion.vision import read_image, read_static_image, write_image -from facefusion.face_helper import warp_face, create_static_mask_frame +from facefusion.typing import Face, FaceSet, Frame, Update_Process, ProcessMode +from facefusion.vision import read_image, read_static_image, read_static_images, write_image +from facefusion.face_helper import warp_face +from facefusion.face_masker import create_static_box_mask, create_occlusion_mask, create_region_mask, clear_face_occluder, clear_face_parser from facefusion.processors.frame import globals as frame_processors_globals, choices as frame_processors_choices -NAME = 'FACEFUSION.FRAME_PROCESSOR.FACE_DEBUGGER' +NAME = __name__.upper() def get_frame_processor() -> None: @@ -34,7 +35,7 @@ def set_options(key : Literal['model'], value : Any) -> None: def register_args(program : ArgumentParser) -> None: - program.add_argument('--face-debugger-items', help = wording.get('face_debugger_items_help'), dest = 'face_debugger_items', default = [ 'kps', 'face-mask' ], choices = frame_processors_choices.face_debugger_items, nargs = '+') + program.add_argument('--face-debugger-items', help = wording.get('face_debugger_items_help').format(choices = ', '.join(frame_processors_choices.face_debugger_items)), default = [ 'kps', 'face-mask' ], choices = frame_processors_choices.face_debugger_items, nargs = '+', metavar = 'FACE_DEBUGGER_ITEMS') def apply_args(program : ArgumentParser) -> None: @@ -54,6 +55,9 @@ def post_process() -> None: clear_frame_processor() clear_face_analyser() clear_content_analyser() + clear_face_occluder() + clear_face_parser() + read_static_image.cache_clear() def debug_face(source_face : Face, target_face : Face, temp_frame : Frame) -> Frame: @@ -63,14 +67,23 @@ def debug_face(source_face : Face, target_face : Face, temp_frame : Frame) -> Fr if 'bbox' in frame_processors_globals.face_debugger_items: cv2.rectangle(temp_frame, (bounding_box[0], bounding_box[1]), (bounding_box[2], bounding_box[3]), secondary_color, 2) if 'face-mask' in frame_processors_globals.face_debugger_items: - crop_frame, affine_matrix = warp_face(temp_frame, target_face.kps, 'arcface_v2', (128, 128)) + crop_frame, affine_matrix = warp_face(temp_frame, target_face.kps, 'arcface_128_v2', (128, 512)) inverse_matrix = cv2.invertAffineTransform(affine_matrix) temp_frame_size = temp_frame.shape[:2][::-1] - mask_frame = create_static_mask_frame(crop_frame.shape[:2], 0, facefusion.globals.face_mask_padding) - mask_frame[mask_frame > 0] = 255 - inverse_mask_frame = cv2.warpAffine(mask_frame.astype(numpy.uint8), inverse_matrix, temp_frame_size) - inverse_mask_contours = cv2.findContours(inverse_mask_frame, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0] - cv2.drawContours(temp_frame, inverse_mask_contours, 0, primary_color, 2) + crop_mask_list = [] + if 'box' in facefusion.globals.face_mask_types: + crop_mask_list.append(create_static_box_mask(crop_frame.shape[:2][::-1], 0, facefusion.globals.face_mask_padding)) + if 'occlusion' in facefusion.globals.face_mask_types: + crop_mask_list.append(create_occlusion_mask(crop_frame)) + if 'region' in facefusion.globals.face_mask_types: + crop_mask_list.append(create_region_mask(crop_frame, facefusion.globals.face_mask_regions)) + crop_mask = numpy.minimum.reduce(crop_mask_list).clip(0, 1) + crop_mask = (crop_mask * 255).astype(numpy.uint8) + inverse_mask_frame = cv2.warpAffine(crop_mask, inverse_matrix, temp_frame_size) + inverse_mask_frame_edges = cv2.threshold(inverse_mask_frame, 100, 255, cv2.THRESH_BINARY)[1] + inverse_mask_frame_edges[inverse_mask_frame_edges > 0] = 255 + inverse_mask_contours = cv2.findContours(inverse_mask_frame_edges, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)[0] + cv2.drawContours(temp_frame, inverse_mask_contours, -1, primary_color, 2) if bounding_box[3] - bounding_box[1] > 60 and bounding_box[2] - bounding_box[0] > 60: if 'kps' in frame_processors_globals.face_debugger_items: kps = target_face.kps.astype(numpy.int32) @@ -83,9 +96,13 @@ def debug_face(source_face : Face, target_face : Face, temp_frame : Frame) -> Fr return temp_frame -def process_frame(source_face : Face, reference_face : Face, temp_frame : Frame) -> Frame: +def get_reference_frame(source_face : Face, target_face : Face, temp_frame : Frame) -> Frame: + pass + + +def process_frame(source_face : Face, reference_faces : FaceSet, temp_frame : Frame) -> Frame: if 'reference' in facefusion.globals.face_selector_mode: - similar_faces = find_similar_faces(temp_frame, reference_face, facefusion.globals.reference_face_distance) + similar_faces = find_similar_faces(temp_frame, reference_faces, facefusion.globals.reference_face_distance) if similar_faces: for similar_face in similar_faces: temp_frame = debug_face(source_face, similar_face, temp_frame) @@ -101,23 +118,25 @@ def process_frame(source_face : Face, reference_face : Face, temp_frame : Frame) return temp_frame -def process_frames(source_path : str, temp_frame_paths : List[str], update_progress : Update_Process) -> None: - source_face = get_one_face(read_static_image(source_path)) - reference_face = get_face_reference() if 'reference' in facefusion.globals.face_selector_mode else None +def process_frames(source_paths : List[str], temp_frame_paths : List[str], update_progress : Update_Process) -> None: + source_frames = read_static_images(source_paths) + source_face = get_average_face(source_frames) + reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None for temp_frame_path in temp_frame_paths: temp_frame = read_image(temp_frame_path) - result_frame = process_frame(source_face, reference_face, temp_frame) + result_frame = process_frame(source_face, reference_faces, temp_frame) write_image(temp_frame_path, result_frame) update_progress() -def process_image(source_path : str, target_path : str, output_path : str) -> None: - source_face = get_one_face(read_static_image(source_path)) +def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: + source_frames = read_static_images(source_paths) + source_face = get_average_face(source_frames) target_frame = read_static_image(target_path) - reference_face = get_one_face(target_frame, facefusion.globals.reference_face_position) if 'reference' in facefusion.globals.face_selector_mode else None - result_frame = process_frame(source_face, reference_face, target_frame) + reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None + result_frame = process_frame(source_face, reference_faces, target_frame) write_image(output_path, result_frame) -def process_video(source_path : str, temp_frame_paths : List[str]) -> None: - frame_processors.multi_process_frames(source_path, temp_frame_paths, process_frames) +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: + frame_processors.multi_process_frames(source_paths, temp_frame_paths, process_frames) diff --git a/facefusion/processors/frame/modules/face_enhancer.py b/facefusion/processors/frame/modules/face_enhancer.py index 416e4991..7c4b57ae 100755 --- a/facefusion/processors/frame/modules/face_enhancer.py +++ b/facefusion/processors/frame/modules/face_enhancer.py @@ -1,4 +1,4 @@ -from typing import Any, List, Dict, Literal, Optional +from typing import Any, List, Literal, Optional from argparse import ArgumentParser import cv2 import threading @@ -7,69 +7,73 @@ import onnxruntime import facefusion.globals import facefusion.processors.frame.core as frame_processors -from facefusion import wording -from facefusion.face_analyser import get_many_faces, clear_face_analyser +from facefusion import logger, wording +from facefusion.face_analyser import get_many_faces, clear_face_analyser, find_similar_faces, get_one_face from facefusion.face_helper import warp_face, paste_back from facefusion.content_analyser import clear_content_analyser -from facefusion.typing import Face, Frame, Update_Process, ProcessMode, ModelValue, OptionsWithModel -from facefusion.utilities import conditional_download, resolve_relative_path, is_image, is_video, is_file, is_download_done, create_metavar, update_status +from facefusion.face_store import get_reference_faces +from facefusion.typing import Face, FaceSet, Frame, Update_Process, ProcessMode, ModelSet, OptionsWithModel +from facefusion.cli_helper import create_metavar +from facefusion.filesystem import is_file, is_image, is_video, resolve_relative_path +from facefusion.download import conditional_download, is_download_done from facefusion.vision import read_image, read_static_image, write_image from facefusion.processors.frame import globals as frame_processors_globals from facefusion.processors.frame import choices as frame_processors_choices +from facefusion.face_masker import create_static_box_mask, create_occlusion_mask, clear_face_occluder FRAME_PROCESSOR = None THREAD_SEMAPHORE : threading.Semaphore = threading.Semaphore() THREAD_LOCK : threading.Lock = threading.Lock() -NAME = 'FACEFUSION.FRAME_PROCESSOR.FACE_ENHANCER' -MODELS : Dict[str, ModelValue] =\ +NAME = __name__.upper() +MODELS : ModelSet =\ { 'codeformer': { 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/codeformer.onnx', 'path': resolve_relative_path('../.assets/models/codeformer.onnx'), - 'template': 'ffhq', + 'template': 'ffhq_512', 'size': (512, 512) }, 'gfpgan_1.2': { 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.2.onnx', 'path': resolve_relative_path('../.assets/models/gfpgan_1.2.onnx'), - 'template': 'ffhq', + 'template': 'ffhq_512', 'size': (512, 512) }, 'gfpgan_1.3': { 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.3.onnx', 'path': resolve_relative_path('../.assets/models/gfpgan_1.3.onnx'), - 'template': 'ffhq', + 'template': 'ffhq_512', 'size': (512, 512) }, 'gfpgan_1.4': { 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.4.onnx', 'path': resolve_relative_path('../.assets/models/gfpgan_1.4.onnx'), - 'template': 'ffhq', + 'template': 'ffhq_512', 'size': (512, 512) }, 'gpen_bfr_256': { 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_256.onnx', 'path': resolve_relative_path('../.assets/models/gpen_bfr_256.onnx'), - 'template': 'arcface_v2', + 'template': 'arcface_128_v2', 'size': (128, 256) }, 'gpen_bfr_512': { 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_512.onnx', 'path': resolve_relative_path('../.assets/models/gpen_bfr_512.onnx'), - 'template': 'ffhq', + 'template': 'ffhq_512', 'size': (512, 512) }, 'restoreformer': { 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/restoreformer.onnx', 'path': resolve_relative_path('../.assets/models/restoreformer.onnx'), - 'template': 'ffhq', + 'template': 'ffhq_512', 'size': (512, 512) } } @@ -110,8 +114,8 @@ def set_options(key : Literal['model'], value : Any) -> None: def register_args(program : ArgumentParser) -> None: - program.add_argument('--face-enhancer-model', help = wording.get('frame_processor_model_help'), dest = 'face_enhancer_model', default = 'gfpgan_1.4', choices = frame_processors_choices.face_enhancer_models) - program.add_argument('--face-enhancer-blend', help = wording.get('frame_processor_blend_help'), dest = 'face_enhancer_blend', type = int, default = 80, choices = frame_processors_choices.face_enhancer_blend_range, metavar = create_metavar(frame_processors_choices.face_enhancer_blend_range)) + program.add_argument('--face-enhancer-model', help = wording.get('frame_processor_model_help'), default = 'gfpgan_1.4', choices = frame_processors_choices.face_enhancer_models) + program.add_argument('--face-enhancer-blend', help = wording.get('frame_processor_blend_help'), type = int, default = 80, choices = frame_processors_choices.face_enhancer_blend_range, metavar = create_metavar(frame_processors_choices.face_enhancer_blend_range)) def apply_args(program : ArgumentParser) -> None: @@ -132,16 +136,16 @@ def pre_process(mode : ProcessMode) -> bool: model_url = get_options('model').get('url') model_path = get_options('model').get('path') if not facefusion.globals.skip_download and not is_download_done(model_url, model_path): - update_status(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME) + logger.error(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME) return False elif not is_file(model_path): - update_status(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME) + logger.error(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME) return False if mode in [ 'output', 'preview' ] and not is_image(facefusion.globals.target_path) and not is_video(facefusion.globals.target_path): - update_status(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME) + logger.error(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME) return False if mode == 'output' and not facefusion.globals.output_path: - update_status(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME) + logger.error(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME) return False return True @@ -150,6 +154,7 @@ def post_process() -> None: clear_frame_processor() clear_face_analyser() clear_content_analyser() + clear_face_occluder() read_static_image.cache_clear() @@ -158,6 +163,12 @@ def enhance_face(target_face: Face, temp_frame: Frame) -> Frame: model_template = get_options('model').get('template') model_size = get_options('model').get('size') crop_frame, affine_matrix = warp_face(temp_frame, target_face.kps, model_template, model_size) + crop_mask_list =\ + [ + create_static_box_mask(crop_frame.shape[:2][::-1], facefusion.globals.face_mask_blur, (0, 0, 0, 0)) + ] + if 'occlusion' in facefusion.globals.face_mask_types: + crop_mask_list.append(create_occlusion_mask(crop_frame)) crop_frame = prepare_crop_frame(crop_frame) frame_processor_inputs = {} for frame_processor_input in frame_processor.get_inputs(): @@ -168,7 +179,8 @@ def enhance_face(target_face: Face, temp_frame: Frame) -> Frame: with THREAD_SEMAPHORE: crop_frame = frame_processor.run(None, frame_processor_inputs)[0][0] crop_frame = normalize_crop_frame(crop_frame) - paste_frame = paste_back(temp_frame, crop_frame, affine_matrix, facefusion.globals.face_mask_blur, (0, 0, 0, 0)) + crop_mask = numpy.minimum.reduce(crop_mask_list).clip(0, 1) + paste_frame = paste_back(temp_frame, crop_frame, crop_mask, affine_matrix) temp_frame = blend_frame(temp_frame, paste_frame) return temp_frame @@ -195,27 +207,43 @@ def blend_frame(temp_frame : Frame, paste_frame : Frame) -> Frame: return temp_frame -def process_frame(source_face : Face, reference_face : Face, temp_frame : Frame) -> Frame: - many_faces = get_many_faces(temp_frame) - if many_faces: - for target_face in many_faces: +def get_reference_frame(source_face : Face, target_face : Face, temp_frame : Frame) -> Optional[Frame]: + return enhance_face(target_face, temp_frame) + + +def process_frame(source_face : Face, reference_faces : FaceSet, temp_frame : Frame) -> Frame: + if 'reference' in facefusion.globals.face_selector_mode: + similar_faces = find_similar_faces(temp_frame, reference_faces, facefusion.globals.reference_face_distance) + if similar_faces: + for similar_face in similar_faces: + temp_frame = enhance_face(similar_face, temp_frame) + if 'one' in facefusion.globals.face_selector_mode: + target_face = get_one_face(temp_frame) + if target_face: temp_frame = enhance_face(target_face, temp_frame) + if 'many' in facefusion.globals.face_selector_mode: + many_faces = get_many_faces(temp_frame) + if many_faces: + for target_face in many_faces: + temp_frame = enhance_face(target_face, temp_frame) return temp_frame -def process_frames(source_path : str, temp_frame_paths : List[str], update_progress : Update_Process) -> None: +def process_frames(source_path : List[str], temp_frame_paths : List[str], update_progress : Update_Process) -> None: + reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None for temp_frame_path in temp_frame_paths: temp_frame = read_image(temp_frame_path) - result_frame = process_frame(None, None, temp_frame) + result_frame = process_frame(None, reference_faces, temp_frame) write_image(temp_frame_path, result_frame) update_progress() def process_image(source_path : str, target_path : str, output_path : str) -> None: + reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None target_frame = read_static_image(target_path) - result_frame = process_frame(None, None, target_frame) + result_frame = process_frame(None, reference_faces, target_frame) write_image(output_path, result_frame) -def process_video(source_path : str, temp_frame_paths : List[str]) -> None: +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: frame_processors.multi_process_frames(None, temp_frame_paths, process_frames) diff --git a/facefusion/processors/frame/modules/face_swapper.py b/facefusion/processors/frame/modules/face_swapper.py index ed12ca7e..5df02348 100755 --- a/facefusion/processors/frame/modules/face_swapper.py +++ b/facefusion/processors/frame/modules/face_swapper.py @@ -1,4 +1,4 @@ -from typing import Any, List, Dict, Literal, Optional +from typing import Any, List, Literal, Optional from argparse import ArgumentParser import threading import numpy @@ -8,29 +8,31 @@ from onnx import numpy_helper import facefusion.globals import facefusion.processors.frame.core as frame_processors -from facefusion import wording -from facefusion.face_analyser import get_one_face, get_many_faces, find_similar_faces, clear_face_analyser +from facefusion import logger, wording +from facefusion.face_analyser import get_one_face, get_average_face, get_many_faces, find_similar_faces, clear_face_analyser from facefusion.face_helper import warp_face, paste_back -from facefusion.face_reference import get_face_reference +from facefusion.face_store import get_reference_faces from facefusion.content_analyser import clear_content_analyser -from facefusion.typing import Face, Frame, Update_Process, ProcessMode, ModelValue, OptionsWithModel, Embedding -from facefusion.utilities import conditional_download, resolve_relative_path, is_image, is_video, is_file, is_download_done, update_status -from facefusion.vision import read_image, read_static_image, write_image +from facefusion.typing import Face, FaceSet, Frame, Update_Process, ProcessMode, ModelSet, OptionsWithModel, Embedding +from facefusion.filesystem import is_file, is_image, are_images, is_video, resolve_relative_path +from facefusion.download import conditional_download, is_download_done +from facefusion.vision import read_image, read_static_image, read_static_images, write_image from facefusion.processors.frame import globals as frame_processors_globals from facefusion.processors.frame import choices as frame_processors_choices +from facefusion.face_masker import create_static_box_mask, create_occlusion_mask, create_region_mask, clear_face_occluder, clear_face_parser FRAME_PROCESSOR = None MODEL_MATRIX = None THREAD_LOCK : threading.Lock = threading.Lock() -NAME = 'FACEFUSION.FRAME_PROCESSOR.FACE_SWAPPER' -MODELS : Dict[str, ModelValue] =\ +NAME = __name__.upper() +MODELS : ModelSet =\ { 'blendswap_256': { 'type': 'blendswap', 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/blendswap_256.onnx', 'path': resolve_relative_path('../.assets/models/blendswap_256.onnx'), - 'template': 'ffhq', + 'template': 'ffhq_512', 'size': (512, 256), 'mean': [ 0.0, 0.0, 0.0 ], 'standard_deviation': [ 1.0, 1.0, 1.0 ] @@ -40,7 +42,7 @@ MODELS : Dict[str, ModelValue] =\ 'type': 'inswapper', 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128.onnx', 'path': resolve_relative_path('../.assets/models/inswapper_128.onnx'), - 'template': 'arcface_v2', + 'template': 'arcface_128_v2', 'size': (128, 128), 'mean': [ 0.0, 0.0, 0.0 ], 'standard_deviation': [ 1.0, 1.0, 1.0 ] @@ -50,7 +52,7 @@ MODELS : Dict[str, ModelValue] =\ 'type': 'inswapper', 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128_fp16.onnx', 'path': resolve_relative_path('../.assets/models/inswapper_128_fp16.onnx'), - 'template': 'arcface_v2', + 'template': 'arcface_128_v2', 'size': (128, 128), 'mean': [ 0.0, 0.0, 0.0 ], 'standard_deviation': [ 1.0, 1.0, 1.0 ] @@ -60,7 +62,7 @@ MODELS : Dict[str, ModelValue] =\ 'type': 'simswap', 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/simswap_256.onnx', 'path': resolve_relative_path('../.assets/models/simswap_256.onnx'), - 'template': 'arcface_v1', + 'template': 'arcface_112_v1', 'size': (112, 256), 'mean': [ 0.485, 0.456, 0.406 ], 'standard_deviation': [ 0.229, 0.224, 0.225 ] @@ -70,7 +72,7 @@ MODELS : Dict[str, ModelValue] =\ 'type': 'simswap', 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/simswap_512_unofficial.onnx', 'path': resolve_relative_path('../.assets/models/simswap_512_unofficial.onnx'), - 'template': 'arcface_v1', + 'template': 'arcface_112_v1', 'size': (112, 512), 'mean': [ 0.0, 0.0, 0.0 ], 'standard_deviation': [ 1.0, 1.0, 1.0 ] @@ -130,7 +132,7 @@ def set_options(key : Literal['model'], value : Any) -> None: def register_args(program : ArgumentParser) -> None: - program.add_argument('--face-swapper-model', help = wording.get('frame_processor_model_help'), dest = 'face_swapper_model', default = 'inswapper_128', choices = frame_processors_choices.face_swapper_models) + program.add_argument('--face-swapper-model', help = wording.get('frame_processor_model_help'), default = 'inswapper_128', choices = frame_processors_choices.face_swapper_models) def apply_args(program : ArgumentParser) -> None: @@ -156,22 +158,23 @@ def pre_process(mode : ProcessMode) -> bool: model_url = get_options('model').get('url') model_path = get_options('model').get('path') if not facefusion.globals.skip_download and not is_download_done(model_url, model_path): - update_status(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME) + logger.error(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME) return False elif not is_file(model_path): - update_status(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME) + logger.error(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME) return False - if not is_image(facefusion.globals.source_path): - update_status(wording.get('select_image_source') + wording.get('exclamation_mark'), NAME) - return False - elif not get_one_face(read_static_image(facefusion.globals.source_path)): - update_status(wording.get('no_source_face_detected') + wording.get('exclamation_mark'), NAME) + if not are_images(facefusion.globals.source_paths): + logger.error(wording.get('select_image_source') + wording.get('exclamation_mark'), NAME) return False + for source_frame in read_static_images(facefusion.globals.source_paths): + if not get_one_face(source_frame): + logger.error(wording.get('no_source_face_detected') + wording.get('exclamation_mark'), NAME) + return False if mode in [ 'output', 'preview' ] and not is_image(facefusion.globals.target_path) and not is_video(facefusion.globals.target_path): - update_status(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME) + logger.error(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME) return False if mode == 'output' and not facefusion.globals.output_path: - update_status(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME) + logger.error(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME) return False return True @@ -181,6 +184,8 @@ def post_process() -> None: clear_model_matrix() clear_face_analyser() clear_content_analyser() + clear_face_occluder() + clear_face_parser() read_static_image.cache_clear() @@ -190,6 +195,11 @@ def swap_face(source_face : Face, target_face : Face, temp_frame : Frame) -> Fra model_size = get_options('model').get('size') model_type = get_options('model').get('type') crop_frame, affine_matrix = warp_face(temp_frame, target_face.kps, model_template, model_size) + crop_mask_list = [] + if 'box' in facefusion.globals.face_mask_types: + crop_mask_list.append(create_static_box_mask(crop_frame.shape[:2][::-1], facefusion.globals.face_mask_blur, facefusion.globals.face_mask_padding)) + if 'occlusion' in facefusion.globals.face_mask_types: + crop_mask_list.append(create_occlusion_mask(crop_frame)) crop_frame = prepare_crop_frame(crop_frame) frame_processor_inputs = {} for frame_processor_input in frame_processor.get_inputs(): @@ -202,13 +212,16 @@ def swap_face(source_face : Face, target_face : Face, temp_frame : Frame) -> Fra frame_processor_inputs[frame_processor_input.name] = crop_frame crop_frame = frame_processor.run(None, frame_processor_inputs)[0][0] crop_frame = normalize_crop_frame(crop_frame) - temp_frame = paste_back(temp_frame, crop_frame, affine_matrix, facefusion.globals.face_mask_blur, facefusion.globals.face_mask_padding) + if 'region' in facefusion.globals.face_mask_types: + crop_mask_list.append(create_region_mask(crop_frame, facefusion.globals.face_mask_regions)) + crop_mask = numpy.minimum.reduce(crop_mask_list).clip(0, 1) + temp_frame = paste_back(temp_frame, crop_frame, crop_mask, affine_matrix) return temp_frame -def prepare_source_frame(source_face : Face) -> numpy.ndarray[Any, Any]: - source_frame = read_static_image(facefusion.globals.source_path) - source_frame, _ = warp_face(source_frame, source_face.kps, 'arcface_v2', (112, 112)) +def prepare_source_frame(source_face : Face) -> Frame: + source_frame = read_static_image(facefusion.globals.source_paths[0]) + source_frame, _ = warp_face(source_frame, source_face.kps, 'arcface_112_v2', (112, 112)) source_frame = source_frame[:, :, ::-1] / 255.0 source_frame = source_frame.transpose(2, 0, 1) source_frame = numpy.expand_dims(source_frame, axis = 0).astype(numpy.float32) @@ -243,9 +256,13 @@ def normalize_crop_frame(crop_frame : Frame) -> Frame: return crop_frame -def process_frame(source_face : Face, reference_face : Face, temp_frame : Frame) -> Frame: +def get_reference_frame(source_face : Face, target_face : Face, temp_frame : Frame) -> Frame: + return swap_face(source_face, target_face, temp_frame) + + +def process_frame(source_face : Face, reference_faces : FaceSet, temp_frame : Frame) -> Frame: if 'reference' in facefusion.globals.face_selector_mode: - similar_faces = find_similar_faces(temp_frame, reference_face, facefusion.globals.reference_face_distance) + similar_faces = find_similar_faces(temp_frame, reference_faces, facefusion.globals.reference_face_distance) if similar_faces: for similar_face in similar_faces: temp_frame = swap_face(source_face, similar_face, temp_frame) @@ -261,23 +278,25 @@ def process_frame(source_face : Face, reference_face : Face, temp_frame : Frame) return temp_frame -def process_frames(source_path : str, temp_frame_paths : List[str], update_progress : Update_Process) -> None: - source_face = get_one_face(read_static_image(source_path)) - reference_face = get_face_reference() if 'reference' in facefusion.globals.face_selector_mode else None +def process_frames(source_paths : List[str], temp_frame_paths : List[str], update_progress : Update_Process) -> None: + source_frames = read_static_images(source_paths) + source_face = get_average_face(source_frames) + reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None for temp_frame_path in temp_frame_paths: temp_frame = read_image(temp_frame_path) - result_frame = process_frame(source_face, reference_face, temp_frame) + result_frame = process_frame(source_face, reference_faces, temp_frame) write_image(temp_frame_path, result_frame) update_progress() -def process_image(source_path : str, target_path : str, output_path : str) -> None: - source_face = get_one_face(read_static_image(source_path)) +def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: + source_frames = read_static_images(source_paths) + source_face = get_average_face(source_frames) + reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None target_frame = read_static_image(target_path) - reference_face = get_one_face(target_frame, facefusion.globals.reference_face_position) if 'reference' in facefusion.globals.face_selector_mode else None - result_frame = process_frame(source_face, reference_face, target_frame) + result_frame = process_frame(source_face, reference_faces, target_frame) write_image(output_path, result_frame) -def process_video(source_path : str, temp_frame_paths : List[str]) -> None: - frame_processors.multi_process_frames(source_path, temp_frame_paths, process_frames) +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: + frame_processors.multi_process_frames(source_paths, temp_frame_paths, process_frames) diff --git a/facefusion/processors/frame/modules/frame_enhancer.py b/facefusion/processors/frame/modules/frame_enhancer.py index c2194e6d..1b82d703 100644 --- a/facefusion/processors/frame/modules/frame_enhancer.py +++ b/facefusion/processors/frame/modules/frame_enhancer.py @@ -1,4 +1,4 @@ -from typing import Any, List, Dict, Literal, Optional +from typing import Any, List, Literal, Optional from argparse import ArgumentParser import threading import cv2 @@ -7,11 +7,14 @@ from realesrgan import RealESRGANer import facefusion.globals import facefusion.processors.frame.core as frame_processors -from facefusion import wording +from facefusion import logger, wording from facefusion.face_analyser import clear_face_analyser from facefusion.content_analyser import clear_content_analyser -from facefusion.typing import Frame, Face, Update_Process, ProcessMode, ModelValue, OptionsWithModel -from facefusion.utilities import conditional_download, resolve_relative_path, is_file, is_download_done, map_device, create_metavar, update_status +from facefusion.typing import Face, FaceSet, Frame, Update_Process, ProcessMode, ModelSet, OptionsWithModel +from facefusion.cli_helper import create_metavar +from facefusion.execution_helper import map_device +from facefusion.filesystem import is_file, resolve_relative_path +from facefusion.download import conditional_download, is_download_done from facefusion.vision import read_image, read_static_image, write_image from facefusion.processors.frame import globals as frame_processors_globals from facefusion.processors.frame import choices as frame_processors_choices @@ -19,8 +22,8 @@ from facefusion.processors.frame import choices as frame_processors_choices FRAME_PROCESSOR = None THREAD_SEMAPHORE : threading.Semaphore = threading.Semaphore() THREAD_LOCK : threading.Lock = threading.Lock() -NAME = 'FACEFUSION.FRAME_PROCESSOR.FRAME_ENHANCER' -MODELS: Dict[str, ModelValue] =\ +NAME = __name__.upper() +MODELS : ModelSet =\ { 'real_esrgan_x2plus': { @@ -88,8 +91,8 @@ def set_options(key : Literal['model'], value : Any) -> None: def register_args(program : ArgumentParser) -> None: - program.add_argument('--frame-enhancer-model', help = wording.get('frame_processor_model_help'), dest = 'frame_enhancer_model', default = 'real_esrgan_x2plus', choices = frame_processors_choices.frame_enhancer_models) - program.add_argument('--frame-enhancer-blend', help = wording.get('frame_processor_blend_help'), dest = 'frame_enhancer_blend', type = int, default = 80, choices = frame_processors_choices.frame_enhancer_blend_range, metavar = create_metavar(frame_processors_choices.frame_enhancer_blend_range)) + program.add_argument('--frame-enhancer-model', help = wording.get('frame_processor_model_help'), default = 'real_esrgan_x2plus', choices = frame_processors_choices.frame_enhancer_models) + program.add_argument('--frame-enhancer-blend', help = wording.get('frame_processor_blend_help'), type = int, default = 80, choices = frame_processors_choices.frame_enhancer_blend_range, metavar = create_metavar(frame_processors_choices.frame_enhancer_blend_range)) def apply_args(program : ArgumentParser) -> None: @@ -110,13 +113,13 @@ def pre_process(mode : ProcessMode) -> bool: model_url = get_options('model').get('url') model_path = get_options('model').get('path') if not facefusion.globals.skip_download and not is_download_done(model_url, model_path): - update_status(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME) + logger.error(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME) return False elif not is_file(model_path): - update_status(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME) + logger.error(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME) return False if mode == 'output' and not facefusion.globals.output_path: - update_status(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME) + logger.error(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME) return False return True @@ -143,11 +146,15 @@ def blend_frame(temp_frame : Frame, paste_frame : Frame) -> Frame: return temp_frame -def process_frame(source_face : Face, reference_face : Face, temp_frame : Frame) -> Frame: +def get_reference_frame(source_face : Face, target_face : Face, temp_frame : Frame) -> Frame: + pass + + +def process_frame(source_face : Face, reference_faces : FaceSet, temp_frame : Frame) -> Frame: return enhance_frame(temp_frame) -def process_frames(source_path : str, temp_frame_paths : List[str], update_progress : Update_Process) -> None: +def process_frames(source_paths : List[str], temp_frame_paths : List[str], update_progress : Update_Process) -> None: for temp_frame_path in temp_frame_paths: temp_frame = read_image(temp_frame_path) result_frame = process_frame(None, None, temp_frame) @@ -155,11 +162,11 @@ def process_frames(source_path : str, temp_frame_paths : List[str], update_progr update_progress() -def process_image(source_path : str, target_path : str, output_path : str) -> None: +def process_image(source_paths : List[str], target_path : str, output_path : str) -> None: target_frame = read_static_image(target_path) result = process_frame(None, None, target_frame) write_image(output_path, result) -def process_video(source_path : str, temp_frame_paths : List[str]) -> None: +def process_video(source_paths : List[str], temp_frame_paths : List[str]) -> None: frame_processors.multi_process_frames(None, temp_frame_paths, process_frames) diff --git a/facefusion/typing.py b/facefusion/typing.py index 9db019de..2964040f 100755 --- a/facefusion/typing.py +++ b/facefusion/typing.py @@ -1,5 +1,5 @@ -from collections import namedtuple from typing import Any, Literal, Callable, List, Tuple, Dict, TypedDict +from collections import namedtuple import numpy Bbox = numpy.ndarray[Any, Any] @@ -16,14 +16,21 @@ Face = namedtuple('Face', 'gender', 'age' ]) +FaceSet = Dict[str, List[Face]] +FaceStore = TypedDict('FaceStore', +{ + 'static_faces' : FaceSet, + 'reference_faces': FaceSet +}) Frame = numpy.ndarray[Any, Any] +Mask = numpy.ndarray[Any, Any] Matrix = numpy.ndarray[Any, Any] Padding = Tuple[int, int, int, int] Update_Process = Callable[[], None] -Process_Frames = Callable[[str, List[str], Update_Process], None] - -Template = Literal['arcface_v1', 'arcface_v2', 'ffhq'] +Process_Frames = Callable[[List[str], List[str], Update_Process], None] +LogLevel = Literal['error', 'warn', 'info', 'debug'] +Template = Literal['arcface_112_v1', 'arcface_112_v2', 'arcface_128_v2', 'ffhq_512'] ProcessMode = Literal['output', 'preview', 'stream'] FaceSelectorMode = Literal['reference', 'one', 'many'] FaceAnalyserOrder = Literal['left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best'] @@ -31,10 +38,13 @@ FaceAnalyserAge = Literal['child', 'teen', 'adult', 'senior'] FaceAnalyserGender = Literal['male', 'female'] FaceDetectorModel = Literal['retinaface', 'yunet'] FaceRecognizerModel = Literal['arcface_blendswap', 'arcface_inswapper', 'arcface_simswap'] +FaceMaskType = Literal['box', 'occlusion', 'region'] +FaceMaskRegion = Literal['skin', 'left-eyebrow', 'right-eyebrow', 'left-eye', 'right-eye', 'eye-glasses', 'nose', 'mouth', 'upper-lip', 'lower-lip'] TempFrameFormat = Literal['jpg', 'png'] OutputVideoEncoder = Literal['libx264', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc'] ModelValue = Dict[str, Any] +ModelSet = Dict[str, ModelValue] OptionsWithModel = TypedDict('OptionsWithModel', { 'model' : ModelValue diff --git a/facefusion/uis/components/benchmark.py b/facefusion/uis/components/benchmark.py index 09afe152..bc5c08be 100644 --- a/facefusion/uis/components/benchmark.py +++ b/facefusion/uis/components/benchmark.py @@ -7,11 +7,12 @@ import gradio import facefusion.globals from facefusion import wording from facefusion.face_analyser import get_face_analyser -from facefusion.face_cache import clear_faces_cache +from facefusion.face_store import clear_static_faces from facefusion.processors.frame.core import get_frame_processors_modules from facefusion.vision import count_video_frame_total from facefusion.core import limit_resources, conditional_process -from facefusion.utilities import normalize_output_path, clear_temp +from facefusion.normalizer import normalize_output_path +from facefusion.filesystem import clear_temp from facefusion.uis.core import get_ui_component BENCHMARK_RESULTS_DATAFRAME : Optional[gradio.Dataframe] = None @@ -75,7 +76,7 @@ def listen() -> None: def start(benchmark_runs : List[str], benchmark_cycles : int) -> Generator[List[Any], None, None]: - facefusion.globals.source_path = '.assets/examples/source.jpg' + facefusion.globals.source_paths = [ '.assets/examples/source.jpg' ] target_paths = [ BENCHMARKS[benchmark_run] for benchmark_run in benchmark_runs if benchmark_run in BENCHMARKS ] benchmark_results = [] if target_paths: @@ -94,7 +95,7 @@ def pre_process() -> None: def post_process() -> None: - clear_faces_cache() + clear_static_faces() def benchmark(target_path : str, benchmark_cycles : int) -> List[Any]: @@ -102,7 +103,7 @@ def benchmark(target_path : str, benchmark_cycles : int) -> List[Any]: total_fps = 0.0 for i in range(benchmark_cycles): facefusion.globals.target_path = target_path - facefusion.globals.output_path = normalize_output_path(facefusion.globals.source_path, facefusion.globals.target_path, tempfile.gettempdir()) + facefusion.globals.output_path = normalize_output_path(facefusion.globals.source_paths, facefusion.globals.target_path, tempfile.gettempdir()) video_frame_total = count_video_frame_total(facefusion.globals.target_path) start_time = time.perf_counter() conditional_process() diff --git a/facefusion/uis/components/execution.py b/facefusion/uis/components/execution.py index 632d38cf..e8df28fd 100644 --- a/facefusion/uis/components/execution.py +++ b/facefusion/uis/components/execution.py @@ -6,7 +6,7 @@ import facefusion.globals from facefusion import wording from facefusion.face_analyser import clear_face_analyser from facefusion.processors.frame.core import clear_frame_processors_modules -from facefusion.utilities import encode_execution_providers, decode_execution_providers +from facefusion.execution_helper import encode_execution_providers, decode_execution_providers EXECUTION_PROVIDERS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None diff --git a/facefusion/uis/components/face_analyser.py b/facefusion/uis/components/face_analyser.py index 3c701182..cf0b23b2 100644 --- a/facefusion/uis/components/face_analyser.py +++ b/facefusion/uis/components/face_analyser.py @@ -53,7 +53,7 @@ def render() -> None: FACE_DETECTOR_SCORE_SLIDER = gradio.Slider( label = wording.get('face_detector_score_slider_label'), value = facefusion.globals.face_detector_score, - step =facefusion.choices.face_detector_score_range[1] - facefusion.choices.face_detector_score_range[0], + step = facefusion.choices.face_detector_score_range[1] - facefusion.choices.face_detector_score_range[0], minimum = facefusion.choices.face_detector_score_range[0], maximum = facefusion.choices.face_detector_score_range[-1] ) diff --git a/facefusion/uis/components/face_mask.py b/facefusion/uis/components/face_masker.py similarity index 55% rename from facefusion/uis/components/face_mask.py rename to facefusion/uis/components/face_masker.py index 08780fba..978a0199 100755 --- a/facefusion/uis/components/face_mask.py +++ b/facefusion/uis/components/face_masker.py @@ -1,33 +1,49 @@ -from typing import Optional +from typing import Optional, Tuple, List import gradio import facefusion.globals import facefusion.choices from facefusion import wording +from facefusion.typing import FaceMaskType, FaceMaskRegion from facefusion.uis.core import register_ui_component +FACE_MASK_TYPES_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None FACE_MASK_BLUR_SLIDER : Optional[gradio.Slider] = None +FACE_MASK_BOX_GROUP : Optional[gradio.Group] = None +FACE_MASK_REGION_GROUP : Optional[gradio.Group] = None FACE_MASK_PADDING_TOP_SLIDER : Optional[gradio.Slider] = None FACE_MASK_PADDING_RIGHT_SLIDER : Optional[gradio.Slider] = None FACE_MASK_PADDING_BOTTOM_SLIDER : Optional[gradio.Slider] = None FACE_MASK_PADDING_LEFT_SLIDER : Optional[gradio.Slider] = None +FACE_MASK_REGION_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None def render() -> None: + global FACE_MASK_TYPES_CHECKBOX_GROUP global FACE_MASK_BLUR_SLIDER + global FACE_MASK_BOX_GROUP + global FACE_MASK_REGION_GROUP global FACE_MASK_PADDING_TOP_SLIDER global FACE_MASK_PADDING_RIGHT_SLIDER global FACE_MASK_PADDING_BOTTOM_SLIDER global FACE_MASK_PADDING_LEFT_SLIDER + global FACE_MASK_REGION_CHECKBOX_GROUP - FACE_MASK_BLUR_SLIDER = gradio.Slider( - label = wording.get('face_mask_blur_slider_label'), - step = facefusion.choices.face_mask_blur_range[1] - facefusion.choices.face_mask_blur_range[0], - minimum = facefusion.choices.face_mask_blur_range[0], - maximum = facefusion.choices.face_mask_blur_range[-1], - value = facefusion.globals.face_mask_blur + has_box_mask = 'box' in facefusion.globals.face_mask_types + has_region_mask = 'region' in facefusion.globals.face_mask_types + FACE_MASK_TYPES_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = wording.get('face_mask_types_checkbox_group_label'), + choices = facefusion.choices.face_mask_types, + value = facefusion.globals.face_mask_types ) - with gradio.Group(): + with gradio.Group(visible = has_box_mask) as FACE_MASK_BOX_GROUP: + FACE_MASK_BLUR_SLIDER = gradio.Slider( + label = wording.get('face_mask_blur_slider_label'), + step = facefusion.choices.face_mask_blur_range[1] - facefusion.choices.face_mask_blur_range[0], + minimum = facefusion.choices.face_mask_blur_range[0], + maximum = facefusion.choices.face_mask_blur_range[-1], + value = facefusion.globals.face_mask_blur + ) with gradio.Row(): FACE_MASK_PADDING_TOP_SLIDER = gradio.Slider( label = wording.get('face_mask_padding_top_slider_label'), @@ -58,23 +74,50 @@ def render() -> None: maximum = facefusion.choices.face_mask_padding_range[-1], value = facefusion.globals.face_mask_padding[3] ) + with gradio.Row(): + FACE_MASK_REGION_CHECKBOX_GROUP = gradio.CheckboxGroup( + label = wording.get('face_mask_region_checkbox_group_label'), + choices = facefusion.choices.face_mask_regions, + value = facefusion.globals.face_mask_regions, + visible = has_region_mask + ) + register_ui_component('face_mask_types_checkbox_group', FACE_MASK_TYPES_CHECKBOX_GROUP) register_ui_component('face_mask_blur_slider', FACE_MASK_BLUR_SLIDER) register_ui_component('face_mask_padding_top_slider', FACE_MASK_PADDING_TOP_SLIDER) register_ui_component('face_mask_padding_right_slider', FACE_MASK_PADDING_RIGHT_SLIDER) register_ui_component('face_mask_padding_bottom_slider', FACE_MASK_PADDING_BOTTOM_SLIDER) register_ui_component('face_mask_padding_left_slider', FACE_MASK_PADDING_LEFT_SLIDER) + register_ui_component('face_mask_region_checkbox_group', FACE_MASK_REGION_CHECKBOX_GROUP) def listen() -> None: + FACE_MASK_TYPES_CHECKBOX_GROUP.change(update_face_mask_type, inputs = FACE_MASK_TYPES_CHECKBOX_GROUP, outputs = [ FACE_MASK_TYPES_CHECKBOX_GROUP, FACE_MASK_BOX_GROUP, FACE_MASK_REGION_CHECKBOX_GROUP ]) FACE_MASK_BLUR_SLIDER.change(update_face_mask_blur, inputs = FACE_MASK_BLUR_SLIDER) + FACE_MASK_REGION_CHECKBOX_GROUP.change(update_face_mask_regions, inputs = FACE_MASK_REGION_CHECKBOX_GROUP, outputs = FACE_MASK_REGION_CHECKBOX_GROUP) face_mask_padding_sliders = [ FACE_MASK_PADDING_TOP_SLIDER, FACE_MASK_PADDING_RIGHT_SLIDER, FACE_MASK_PADDING_BOTTOM_SLIDER, FACE_MASK_PADDING_LEFT_SLIDER ] for face_mask_padding_slider in face_mask_padding_sliders: face_mask_padding_slider.change(update_face_mask_padding, inputs = face_mask_padding_sliders) +def update_face_mask_type(face_mask_types : List[FaceMaskType]) -> Tuple[gradio.CheckboxGroup, gradio.Group, gradio.CheckboxGroup]: + if not face_mask_types: + face_mask_types = facefusion.choices.face_mask_types + facefusion.globals.face_mask_types = face_mask_types + has_box_mask = 'box' in face_mask_types + has_region_mask = 'region' in face_mask_types + return gradio.CheckboxGroup(value = face_mask_types), gradio.Group(visible = has_box_mask), gradio.CheckboxGroup(visible = has_region_mask) + + def update_face_mask_blur(face_mask_blur : float) -> None: facefusion.globals.face_mask_blur = face_mask_blur def update_face_mask_padding(face_mask_padding_top : int, face_mask_padding_right : int, face_mask_padding_bottom : int, face_mask_padding_left : int) -> None: facefusion.globals.face_mask_padding = (face_mask_padding_top, face_mask_padding_right, face_mask_padding_bottom, face_mask_padding_left) + + +def update_face_mask_regions(face_mask_regions : List[FaceMaskRegion]) -> gradio.CheckboxGroup: + if not face_mask_regions: + face_mask_regions = facefusion.choices.face_mask_regions + facefusion.globals.face_mask_regions = face_mask_regions + return gradio.CheckboxGroup(value = face_mask_regions) diff --git a/facefusion/uis/components/face_selector.py b/facefusion/uis/components/face_selector.py index 5ac5f5ed..90ebf3de 100644 --- a/facefusion/uis/components/face_selector.py +++ b/facefusion/uis/components/face_selector.py @@ -5,12 +5,11 @@ import gradio import facefusion.globals import facefusion.choices from facefusion import wording -from facefusion.face_cache import clear_faces_cache +from facefusion.face_store import clear_static_faces, clear_reference_faces from facefusion.vision import get_video_frame, read_static_image, normalize_frame_color from facefusion.face_analyser import get_many_faces -from facefusion.face_reference import clear_face_reference from facefusion.typing import Frame, FaceSelectorMode -from facefusion.utilities import is_image, is_video +from facefusion.filesystem import is_image, is_video from facefusion.uis.core import get_ui_component, register_ui_component from facefusion.uis.typing import ComponentName @@ -111,8 +110,8 @@ def update_face_selector_mode(face_selector_mode : FaceSelectorMode) -> Tuple[gr def clear_and_update_reference_face_position(event : gradio.SelectData) -> gradio.Gallery: - clear_face_reference() - clear_faces_cache() + clear_reference_faces() + clear_static_faces() update_reference_face_position(event.index) return update_reference_position_gallery() @@ -130,8 +129,8 @@ def update_reference_frame_number(reference_frame_number : int) -> None: def clear_and_update_reference_position_gallery() -> gradio.Gallery: - clear_face_reference() - clear_faces_cache() + clear_reference_faces() + clear_static_faces() return update_reference_position_gallery() diff --git a/facefusion/uis/components/frame_processors.py b/facefusion/uis/components/frame_processors.py index 861e5771..ac687349 100644 --- a/facefusion/uis/components/frame_processors.py +++ b/facefusion/uis/components/frame_processors.py @@ -4,7 +4,7 @@ import gradio import facefusion.globals from facefusion import wording from facefusion.processors.frame.core import load_frame_processor_module, clear_frame_processors_modules -from facefusion.utilities import list_module_names +from facefusion.filesystem import list_module_names from facefusion.uis.core import register_ui_component FRAME_PROCESSORS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None diff --git a/facefusion/uis/components/output.py b/facefusion/uis/components/output.py index 81156e69..5d738857 100644 --- a/facefusion/uis/components/output.py +++ b/facefusion/uis/components/output.py @@ -5,7 +5,8 @@ import facefusion.globals from facefusion import wording from facefusion.core import limit_resources, conditional_process from facefusion.uis.core import get_ui_component -from facefusion.utilities import is_image, is_video, normalize_output_path, clear_temp +from facefusion.normalizer import normalize_output_path +from facefusion.filesystem import is_image, is_video, clear_temp OUTPUT_IMAGE : Optional[gradio.Image] = None OUTPUT_VIDEO : Optional[gradio.Video] = None @@ -45,7 +46,7 @@ def listen() -> None: def start(output_path : str) -> Tuple[gradio.Image, gradio.Video]: - facefusion.globals.output_path = normalize_output_path(facefusion.globals.source_path, facefusion.globals.target_path, output_path) + facefusion.globals.output_path = normalize_output_path(facefusion.globals.source_paths, facefusion.globals.target_path, output_path) limit_resources() conditional_process() if is_image(facefusion.globals.output_path): diff --git a/facefusion/uis/components/output_options.py b/facefusion/uis/components/output_options.py index 900a92cc..6b32a110 100644 --- a/facefusion/uis/components/output_options.py +++ b/facefusion/uis/components/output_options.py @@ -6,7 +6,7 @@ import facefusion.globals import facefusion.choices from facefusion import wording from facefusion.typing import OutputVideoEncoder -from facefusion.utilities import is_image, is_video +from facefusion.filesystem import is_image, is_video from facefusion.uis.typing import ComponentName from facefusion.uis.core import get_ui_component, register_ui_component diff --git a/facefusion/uis/components/preview.py b/facefusion/uis/components/preview.py index 7d0fddfe..9087b0d6 100755 --- a/facefusion/uis/components/preview.py +++ b/facefusion/uis/components/preview.py @@ -4,15 +4,14 @@ import gradio import facefusion.globals from facefusion import wording -from facefusion.core import conditional_set_face_reference -from facefusion.face_cache import clear_faces_cache -from facefusion.typing import Frame, Face -from facefusion.vision import get_video_frame, count_video_frame_total, normalize_frame_color, resize_frame_dimension, read_static_image -from facefusion.face_analyser import get_one_face, clear_face_analyser -from facefusion.face_reference import get_face_reference, clear_face_reference +from facefusion.core import conditional_append_reference_faces +from facefusion.face_store import clear_static_faces, get_reference_faces, clear_reference_faces +from facefusion.typing import Frame, Face, FaceSet +from facefusion.vision import get_video_frame, count_video_frame_total, normalize_frame_color, resize_frame_dimension, read_static_image, read_static_images +from facefusion.face_analyser import get_average_face, clear_face_analyser from facefusion.content_analyser import analyse_frame from facefusion.processors.frame.core import load_frame_processor_module -from facefusion.utilities import is_video, is_image +from facefusion.filesystem import is_image, is_video from facefusion.uis.typing import ComponentName from facefusion.uis.core import get_ui_component, register_ui_component @@ -37,16 +36,17 @@ def render() -> None: 'maximum': 100, 'visible': False } - conditional_set_face_reference() - source_face = get_one_face(read_static_image(facefusion.globals.source_path)) - reference_face = get_face_reference() if 'reference' in facefusion.globals.face_selector_mode else None + conditional_append_reference_faces() + source_frames = read_static_images(facefusion.globals.source_paths) + source_face = get_average_face(source_frames) + reference_faces = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None if is_image(facefusion.globals.target_path): target_frame = read_static_image(facefusion.globals.target_path) - preview_frame = process_preview_frame(source_face, reference_face, target_frame) + preview_frame = process_preview_frame(source_face, reference_faces, target_frame) preview_image_args['value'] = normalize_frame_color(preview_frame) if is_video(facefusion.globals.target_path): temp_frame = get_video_frame(facefusion.globals.target_path, facefusion.globals.reference_frame_number) - preview_frame = process_preview_frame(source_face, reference_face, temp_frame) + preview_frame = process_preview_frame(source_face, reference_faces, temp_frame) preview_image_args['value'] = normalize_frame_color(preview_frame) preview_image_args['visible'] = True preview_frame_slider_args['value'] = facefusion.globals.reference_frame_number @@ -58,7 +58,7 @@ def render() -> None: def listen() -> None: - PREVIEW_FRAME_SLIDER.change(update_preview_image, inputs = PREVIEW_FRAME_SLIDER, outputs = PREVIEW_IMAGE) + PREVIEW_FRAME_SLIDER.release(update_preview_image, inputs = PREVIEW_FRAME_SLIDER, outputs = PREVIEW_IMAGE) multi_one_component_names : List[ComponentName] =\ [ 'source_image', @@ -101,11 +101,13 @@ def listen() -> None: 'frame_enhancer_blend_slider', 'face_selector_mode_dropdown', 'reference_face_distance_slider', + 'face_mask_types_checkbox_group', 'face_mask_blur_slider', 'face_mask_padding_top_slider', 'face_mask_padding_bottom_slider', 'face_mask_padding_left_slider', - 'face_mask_padding_right_slider' + 'face_mask_padding_right_slider', + 'face_mask_region_checkbox_group' ] for component_name in change_one_component_names: component = get_ui_component(component_name) @@ -126,15 +128,17 @@ def listen() -> None: def clear_and_update_preview_image(frame_number : int = 0) -> gradio.Image: clear_face_analyser() - clear_face_reference() - clear_faces_cache() + clear_reference_faces() + clear_static_faces() return update_preview_image(frame_number) def update_preview_image(frame_number : int = 0) -> gradio.Image: - conditional_set_face_reference() - source_face = get_one_face(read_static_image(facefusion.globals.source_path)) - reference_face = get_face_reference() if 'reference' in facefusion.globals.face_selector_mode else None + clear_reference_faces() + conditional_append_reference_faces() + source_frames = read_static_images(facefusion.globals.source_paths) + source_face = get_average_face(source_frames) + reference_face = get_reference_faces() if 'reference' in facefusion.globals.face_selector_mode else None if is_image(facefusion.globals.target_path): target_frame = read_static_image(facefusion.globals.target_path) preview_frame = process_preview_frame(source_face, reference_face, target_frame) @@ -155,7 +159,7 @@ def update_preview_frame_slider() -> gradio.Slider: return gradio.Slider(value = None, maximum = None, visible = False) -def process_preview_frame(source_face : Face, reference_face : Face, temp_frame : Frame) -> Frame: +def process_preview_frame(source_face : Face, reference_faces : FaceSet, temp_frame : Frame) -> Frame: temp_frame = resize_frame_dimension(temp_frame, 640, 640) if analyse_frame(temp_frame): return cv2.GaussianBlur(temp_frame, (99, 99), 0) @@ -164,7 +168,7 @@ def process_preview_frame(source_face : Face, reference_face : Face, temp_frame if frame_processor_module.pre_process('preview'): temp_frame = frame_processor_module.process_frame( source_face, - reference_face, + reference_faces, temp_frame ) return temp_frame diff --git a/facefusion/uis/components/source.py b/facefusion/uis/components/source.py index 37777ea4..5fd7a6e7 100644 --- a/facefusion/uis/components/source.py +++ b/facefusion/uis/components/source.py @@ -1,9 +1,10 @@ -from typing import Any, IO, Optional +from typing import Optional, List import gradio import facefusion.globals from facefusion import wording -from facefusion.utilities import is_image +from facefusion.uis.typing import File +from facefusion.filesystem import are_images from facefusion.uis.core import register_ui_component SOURCE_FILE : Optional[gradio.File] = None @@ -14,9 +15,9 @@ def render() -> None: global SOURCE_FILE global SOURCE_IMAGE - is_source_image = is_image(facefusion.globals.source_path) + are_source_images = are_images(facefusion.globals.source_paths) SOURCE_FILE = gradio.File( - file_count = 'single', + file_count = 'multiple', file_types = [ '.png', @@ -24,11 +25,12 @@ def render() -> None: '.webp' ], label = wording.get('source_file_label'), - value = facefusion.globals.source_path if is_source_image else None + value = facefusion.globals.source_paths if are_source_images else None ) + source_file_names = [ source_file_value['name'] for source_file_value in SOURCE_FILE.value ] if SOURCE_FILE.value else None SOURCE_IMAGE = gradio.Image( - value = SOURCE_FILE.value['name'] if is_source_image else None, - visible = is_source_image, + value = source_file_names[0] if are_source_images else None, + visible = are_source_images, show_label = False ) register_ui_component('source_image', SOURCE_IMAGE) @@ -38,9 +40,10 @@ def listen() -> None: SOURCE_FILE.change(update, inputs = SOURCE_FILE, outputs = SOURCE_IMAGE) -def update(file: IO[Any]) -> gradio.Image: - if file and is_image(file.name): - facefusion.globals.source_path = file.name - return gradio.Image(value = file.name, visible = True) - facefusion.globals.source_path = None +def update(files : List[File]) -> gradio.Image: + file_names = [ file.name for file in files ] if files else None + if are_images(file_names): + facefusion.globals.source_paths = file_names + return gradio.Image(value = file_names[0], visible = True) + facefusion.globals.source_paths = None return gradio.Image(value = None, visible = False) diff --git a/facefusion/uis/components/target.py b/facefusion/uis/components/target.py index b89ac187..307b670c 100644 --- a/facefusion/uis/components/target.py +++ b/facefusion/uis/components/target.py @@ -1,11 +1,11 @@ -from typing import Any, IO, Tuple, Optional +from typing import Tuple, Optional import gradio import facefusion.globals from facefusion import wording -from facefusion.face_cache import clear_faces_cache -from facefusion.face_reference import clear_face_reference -from facefusion.utilities import is_image, is_video +from facefusion.face_store import clear_static_faces, clear_reference_faces +from facefusion.uis.typing import File +from facefusion.filesystem import is_image, is_video from facefusion.uis.core import register_ui_component TARGET_FILE : Optional[gradio.File] = None @@ -50,9 +50,9 @@ def listen() -> None: TARGET_FILE.change(update, inputs = TARGET_FILE, outputs = [ TARGET_IMAGE, TARGET_VIDEO ]) -def update(file : IO[Any]) -> Tuple[gradio.Image, gradio.Video]: - clear_face_reference() - clear_faces_cache() +def update(file : File) -> Tuple[gradio.Image, gradio.Video]: + clear_reference_faces() + clear_static_faces() if file and is_image(file.name): facefusion.globals.target_path = file.name return gradio.Image(value = file.name, visible = True), gradio.Video(value = None, visible = False) diff --git a/facefusion/uis/components/temp_frame.py b/facefusion/uis/components/temp_frame.py index dfab64fe..d07f8365 100644 --- a/facefusion/uis/components/temp_frame.py +++ b/facefusion/uis/components/temp_frame.py @@ -5,7 +5,7 @@ import facefusion.globals import facefusion.choices from facefusion import wording from facefusion.typing import TempFrameFormat -from facefusion.utilities import is_video +from facefusion.filesystem import is_video from facefusion.uis.core import get_ui_component TEMP_FRAME_FORMAT_DROPDOWN : Optional[gradio.Dropdown] = None diff --git a/facefusion/uis/components/trim_frame.py b/facefusion/uis/components/trim_frame.py index 1e6048c3..10d6089a 100644 --- a/facefusion/uis/components/trim_frame.py +++ b/facefusion/uis/components/trim_frame.py @@ -4,7 +4,7 @@ import gradio import facefusion.globals from facefusion import wording from facefusion.vision import count_video_frame_total -from facefusion.utilities import is_video +from facefusion.filesystem import is_video from facefusion.uis.core import get_ui_component TRIM_FRAME_START_SLIDER : Optional[gradio.Slider] = None diff --git a/facefusion/uis/components/webcam.py b/facefusion/uis/components/webcam.py index 0b7ba8d0..d1217a6d 100644 --- a/facefusion/uis/components/webcam.py +++ b/facefusion/uis/components/webcam.py @@ -9,13 +9,13 @@ import gradio from tqdm import tqdm import facefusion.globals -from facefusion import wording +from facefusion import logger, wording from facefusion.content_analyser import analyse_stream from facefusion.typing import Frame, Face -from facefusion.face_analyser import get_one_face +from facefusion.face_analyser import get_average_face from facefusion.processors.frame.core import get_frame_processors_modules -from facefusion.utilities import open_ffmpeg -from facefusion.vision import normalize_frame_color, read_static_image +from facefusion.ffmpeg import open_ffmpeg +from facefusion.vision import normalize_frame_color, read_static_images from facefusion.uis.typing import StreamMode, WebcamMode from facefusion.uis.core import get_ui_component @@ -79,30 +79,34 @@ def listen() -> None: getattr(source_image, method)(stop, cancels = start_event) -def start(mode : WebcamMode, resolution : str, fps : float) -> Generator[Frame, None, None]: +def start(webcam_mode : WebcamMode, resolution : str, fps : float) -> Generator[Frame, None, None]: facefusion.globals.face_selector_mode = 'one' facefusion.globals.face_analyser_order = 'large-small' - source_face = get_one_face(read_static_image(facefusion.globals.source_path)) + source_frames = read_static_images(facefusion.globals.source_paths) + source_face = get_average_face(source_frames) stream = None - if mode in [ 'udp', 'v4l2' ]: - stream = open_stream(mode, resolution, fps) # type: ignore[arg-type] + if webcam_mode in [ 'udp', 'v4l2' ]: + stream = open_stream(webcam_mode, resolution, fps) # type: ignore[arg-type] webcam_width, webcam_height = map(int, resolution.split('x')) webcam_capture = get_webcam_capture() if webcam_capture and webcam_capture.isOpened(): - webcam_capture.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG')) # type: ignore[attr-defined] + webcam_capture.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG')) # type: ignore[attr-defined] webcam_capture.set(cv2.CAP_PROP_FRAME_WIDTH, webcam_width) webcam_capture.set(cv2.CAP_PROP_FRAME_HEIGHT, webcam_height) webcam_capture.set(cv2.CAP_PROP_FPS, fps) for capture_frame in multi_process_capture(source_face, webcam_capture, fps): - if mode == 'inline': + if webcam_mode == 'inline': yield normalize_frame_color(capture_frame) else: - stream.stdin.write(capture_frame.tobytes()) + try: + stream.stdin.write(capture_frame.tobytes()) + except Exception: + clear_webcam_capture() yield None def multi_process_capture(source_face : Face, webcam_capture : cv2.VideoCapture, fps : float) -> Generator[Frame, None, None]: - with tqdm(desc = wording.get('processing'), unit = 'frame', ascii = ' =') as progress: + with tqdm(desc = wording.get('processing'), unit = 'frame', ascii = ' =', disable = facefusion.globals.log_level in [ 'warn', 'error' ]) as progress: with ThreadPoolExecutor(max_workers = facefusion.globals.execution_thread_count) as executor: futures = [] deque_capture_frames : Deque[Frame] = deque() @@ -137,11 +141,15 @@ def process_stream_frame(source_face : Face, temp_frame : Frame) -> Frame: return temp_frame -def open_stream(mode : StreamMode, resolution : str, fps : float) -> subprocess.Popen[bytes]: +def open_stream(stream_mode : StreamMode, resolution : str, fps : float) -> subprocess.Popen[bytes]: commands = [ '-f', 'rawvideo', '-pix_fmt', 'bgr24', '-s', resolution, '-r', str(fps), '-i', '-' ] - if mode == 'udp': + if stream_mode == 'udp': commands.extend([ '-b:v', '2000k', '-f', 'mpegts', 'udp://localhost:27000?pkt_size=1316' ]) - if mode == 'v4l2': - device_name = os.listdir('/sys/devices/virtual/video4linux')[0] - commands.extend([ '-f', 'v4l2', '/dev/' + device_name ]) + if stream_mode == 'v4l2': + try: + device_name = os.listdir('/sys/devices/virtual/video4linux')[0] + if device_name: + commands.extend([ '-f', 'v4l2', '/dev/' + device_name ]) + except FileNotFoundError: + logger.error(wording.get('stream_not_loaded').format(stream_mode = stream_mode), __name__.upper()) return open_ffmpeg(commands) diff --git a/facefusion/uis/core.py b/facefusion/uis/core.py index d8b565e0..9f7b6cd0 100644 --- a/facefusion/uis/core.py +++ b/facefusion/uis/core.py @@ -5,9 +5,9 @@ import sys import gradio import facefusion.globals -from facefusion import metadata, wording +from facefusion import metadata, logger, wording from facefusion.uis.typing import Component, ComponentName -from facefusion.utilities import resolve_relative_path +from facefusion.filesystem import resolve_relative_path UI_COMPONENTS: Dict[ComponentName, Component] = {} UI_LAYOUT_MODULES : List[ModuleType] = [] @@ -27,7 +27,8 @@ def load_ui_layout_module(ui_layout : str) -> Any: for method_name in UI_LAYOUT_METHODS: if not hasattr(ui_layout_module, method_name): raise NotImplementedError - except ModuleNotFoundError: + except ModuleNotFoundError as exception: + logger.debug(exception.msg, __name__.upper()) sys.exit(wording.get('ui_layout_not_loaded').format(ui_layout = ui_layout)) except NotImplementedError: sys.exit(wording.get('ui_layout_not_implemented').format(ui_layout = ui_layout)) diff --git a/facefusion/uis/layouts/benchmark.py b/facefusion/uis/layouts/benchmark.py index 829db2fd..ae9c3202 100644 --- a/facefusion/uis/layouts/benchmark.py +++ b/facefusion/uis/layouts/benchmark.py @@ -1,7 +1,7 @@ import gradio import facefusion.globals -from facefusion.utilities import conditional_download +from facefusion.download import conditional_download from facefusion.uis.components import about, frame_processors, frame_processors_options, execution, execution_thread_count, execution_queue_count, limit_resources, benchmark_options, benchmark diff --git a/facefusion/uis/layouts/default.py b/facefusion/uis/layouts/default.py index 3e9da2e7..4537297e 100755 --- a/facefusion/uis/layouts/default.py +++ b/facefusion/uis/layouts/default.py @@ -1,6 +1,6 @@ import gradio -from facefusion.uis.components import about, frame_processors, frame_processors_options, execution, execution_thread_count, execution_queue_count, limit_resources, temp_frame, output_options, common_options, source, target, output, preview, trim_frame, face_analyser, face_selector, face_mask +from facefusion.uis.components import about, frame_processors, frame_processors_options, execution, execution_thread_count, execution_queue_count, limit_resources, temp_frame, output_options, common_options, source, target, output, preview, trim_frame, face_analyser, face_selector, face_masker def pre_check() -> bool: @@ -47,7 +47,7 @@ def render() -> gradio.Blocks: with gradio.Blocks(): face_selector.render() with gradio.Blocks(): - face_mask.render() + face_masker.render() with gradio.Blocks(): face_analyser.render() return layout @@ -69,7 +69,7 @@ def listen() -> None: preview.listen() trim_frame.listen() face_selector.listen() - face_mask.listen() + face_masker.listen() face_analyser.listen() diff --git a/facefusion/uis/typing.py b/facefusion/uis/typing.py index e104d0fc..b2c57d32 100644 --- a/facefusion/uis/typing.py +++ b/facefusion/uis/typing.py @@ -1,6 +1,7 @@ -from typing import Literal +from typing import Literal, Any, IO import gradio +File = IO[Any] Component = gradio.File or gradio.Image or gradio.Video or gradio.Slider ComponentName = Literal\ [ @@ -17,11 +18,13 @@ ComponentName = Literal\ 'face_detector_model_dropdown', 'face_detector_size_dropdown', 'face_detector_score_slider', + 'face_mask_types_checkbox_group', 'face_mask_blur_slider', 'face_mask_padding_top_slider', 'face_mask_padding_bottom_slider', 'face_mask_padding_left_slider', 'face_mask_padding_right_slider', + 'face_mask_region_checkbox_group', 'frame_processors_checkbox_group', 'face_swapper_model_dropdown', 'face_enhancer_model_dropdown', diff --git a/facefusion/utilities.py b/facefusion/utilities.py deleted file mode 100644 index dcdaae13..00000000 --- a/facefusion/utilities.py +++ /dev/null @@ -1,268 +0,0 @@ -from typing import Any, List, Optional -from concurrent.futures import ThreadPoolExecutor -from functools import lru_cache -from pathlib import Path -from tqdm import tqdm -import glob -import filetype -import os -import platform -import shutil -import ssl -import subprocess -import tempfile -import urllib.request -import onnxruntime - -import facefusion.globals -from facefusion import wording -from facefusion.typing import Padding -from facefusion.vision import detect_fps - -TEMP_DIRECTORY_PATH = os.path.join(tempfile.gettempdir(), 'facefusion') -TEMP_OUTPUT_VIDEO_NAME = 'temp.mp4' - -# monkey patch ssl -if platform.system().lower() == 'darwin': - ssl._create_default_https_context = ssl._create_unverified_context - - -def run_ffmpeg(args : List[str]) -> bool: - commands = [ 'ffmpeg', '-hide_banner', '-loglevel', 'error' ] - commands.extend(args) - try: - subprocess.run(commands, stderr = subprocess.PIPE, check = True) - return True - except subprocess.CalledProcessError: - return False - - -def open_ffmpeg(args : List[str]) -> subprocess.Popen[bytes]: - commands = [ 'ffmpeg', '-hide_banner', '-loglevel', 'error' ] - commands.extend(args) - return subprocess.Popen(commands, stdin = subprocess.PIPE) - - -def extract_frames(target_path : str, fps : float) -> bool: - temp_frame_compression = round(31 - (facefusion.globals.temp_frame_quality * 0.31)) - trim_frame_start = facefusion.globals.trim_frame_start - trim_frame_end = facefusion.globals.trim_frame_end - temp_frames_pattern = get_temp_frames_pattern(target_path, '%04d') - commands = [ '-hwaccel', 'auto', '-i', target_path, '-q:v', str(temp_frame_compression), '-pix_fmt', 'rgb24' ] - if trim_frame_start is not None and trim_frame_end is not None: - commands.extend([ '-vf', 'trim=start_frame=' + str(trim_frame_start) + ':end_frame=' + str(trim_frame_end) + ',fps=' + str(fps) ]) - elif trim_frame_start is not None: - commands.extend([ '-vf', 'trim=start_frame=' + str(trim_frame_start) + ',fps=' + str(fps) ]) - elif trim_frame_end is not None: - commands.extend([ '-vf', 'trim=end_frame=' + str(trim_frame_end) + ',fps=' + str(fps) ]) - else: - commands.extend([ '-vf', 'fps=' + str(fps) ]) - commands.extend([ '-vsync', '0', temp_frames_pattern ]) - return run_ffmpeg(commands) - - -def compress_image(output_path : str) -> bool: - output_image_compression = round(31 - (facefusion.globals.output_image_quality * 0.31)) - commands = [ '-hwaccel', 'auto', '-i', output_path, '-q:v', str(output_image_compression), '-y', output_path ] - return run_ffmpeg(commands) - - -def merge_video(target_path : str, fps : float) -> bool: - temp_output_video_path = get_temp_output_video_path(target_path) - temp_frames_pattern = get_temp_frames_pattern(target_path, '%04d') - commands = [ '-hwaccel', 'auto', '-r', str(fps), '-i', temp_frames_pattern, '-c:v', facefusion.globals.output_video_encoder ] - if facefusion.globals.output_video_encoder in [ 'libx264', 'libx265' ]: - output_video_compression = round(51 - (facefusion.globals.output_video_quality * 0.51)) - commands.extend([ '-crf', str(output_video_compression) ]) - if facefusion.globals.output_video_encoder in [ 'libvpx-vp9' ]: - output_video_compression = round(63 - (facefusion.globals.output_video_quality * 0.63)) - commands.extend([ '-crf', str(output_video_compression) ]) - if facefusion.globals.output_video_encoder in [ 'h264_nvenc', 'hevc_nvenc' ]: - output_video_compression = round(51 - (facefusion.globals.output_video_quality * 0.51)) - commands.extend([ '-cq', str(output_video_compression) ]) - commands.extend([ '-pix_fmt', 'yuv420p', '-colorspace', 'bt709', '-y', temp_output_video_path ]) - return run_ffmpeg(commands) - - -def restore_audio(target_path : str, output_path : str) -> bool: - fps = detect_fps(target_path) - trim_frame_start = facefusion.globals.trim_frame_start - trim_frame_end = facefusion.globals.trim_frame_end - temp_output_video_path = get_temp_output_video_path(target_path) - commands = [ '-hwaccel', 'auto', '-i', temp_output_video_path ] - if trim_frame_start is not None: - start_time = trim_frame_start / fps - commands.extend([ '-ss', str(start_time) ]) - if trim_frame_end is not None: - end_time = trim_frame_end / fps - commands.extend([ '-to', str(end_time) ]) - commands.extend([ '-i', target_path, '-c', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-shortest', '-y', output_path ]) - return run_ffmpeg(commands) - - -def get_temp_frame_paths(target_path : str) -> List[str]: - temp_frames_pattern = get_temp_frames_pattern(target_path, '*') - return sorted(glob.glob(temp_frames_pattern)) - - -def get_temp_frames_pattern(target_path : str, temp_frame_prefix : str) -> str: - temp_directory_path = get_temp_directory_path(target_path) - return os.path.join(temp_directory_path, temp_frame_prefix + '.' + facefusion.globals.temp_frame_format) - - -def get_temp_directory_path(target_path : str) -> str: - target_name, _ = os.path.splitext(os.path.basename(target_path)) - return os.path.join(TEMP_DIRECTORY_PATH, target_name) - - -def get_temp_output_video_path(target_path : str) -> str: - temp_directory_path = get_temp_directory_path(target_path) - return os.path.join(temp_directory_path, TEMP_OUTPUT_VIDEO_NAME) - - -def create_temp(target_path : str) -> None: - temp_directory_path = get_temp_directory_path(target_path) - Path(temp_directory_path).mkdir(parents = True, exist_ok = True) - - -def move_temp(target_path : str, output_path : str) -> None: - temp_output_video_path = get_temp_output_video_path(target_path) - if is_file(temp_output_video_path): - if is_file(output_path): - os.remove(output_path) - shutil.move(temp_output_video_path, output_path) - - -def clear_temp(target_path : str) -> None: - temp_directory_path = get_temp_directory_path(target_path) - parent_directory_path = os.path.dirname(temp_directory_path) - if not facefusion.globals.keep_temp and is_directory(temp_directory_path): - shutil.rmtree(temp_directory_path) - if os.path.exists(parent_directory_path) and not os.listdir(parent_directory_path): - os.rmdir(parent_directory_path) - - -def normalize_output_path(source_path : Optional[str], target_path : Optional[str], output_path : Optional[str]) -> Optional[str]: - if is_file(target_path) and is_directory(output_path): - target_name, target_extension = os.path.splitext(os.path.basename(target_path)) - if is_file(source_path): - source_name, _ = os.path.splitext(os.path.basename(source_path)) - return os.path.join(output_path, source_name + '-' + target_name + target_extension) - return os.path.join(output_path, target_name + target_extension) - if is_file(target_path) and output_path: - _, target_extension = os.path.splitext(os.path.basename(target_path)) - output_name, output_extension = os.path.splitext(os.path.basename(output_path)) - output_directory_path = os.path.dirname(output_path) - if is_directory(output_directory_path) and output_extension: - return os.path.join(output_directory_path, output_name + target_extension) - return None - return output_path - - -def normalize_padding(padding : Optional[List[int]]) -> Optional[Padding]: - if padding and len(padding) == 1: - return tuple([ padding[0], padding[0], padding[0], padding[0] ]) # type: ignore[return-value] - if padding and len(padding) == 2: - return tuple([ padding[0], padding[1], padding[0], padding[1] ]) # type: ignore[return-value] - if padding and len(padding) == 3: - return tuple([ padding[0], padding[1], padding[2], padding[1] ]) # type: ignore[return-value] - if padding and len(padding) == 4: - return tuple(padding) # type: ignore[return-value] - return None - - -def is_file(file_path : str) -> bool: - return bool(file_path and os.path.isfile(file_path)) - - -def is_directory(directory_path : str) -> bool: - return bool(directory_path and os.path.isdir(directory_path)) - - -def is_image(image_path : str) -> bool: - if is_file(image_path): - mimetype = filetype.guess(image_path).mime - return bool(mimetype and mimetype.startswith('image/')) - return False - - -def is_video(video_path : str) -> bool: - if is_file(video_path): - mimetype = filetype.guess(video_path).mime - return bool(mimetype and mimetype.startswith('video/')) - return False - - -def conditional_download(download_directory_path : str, urls : List[str]) -> None: - with ThreadPoolExecutor() as executor: - for url in urls: - executor.submit(get_download_size, url) - for url in urls: - download_file_path = os.path.join(download_directory_path, os.path.basename(url)) - total = get_download_size(url) - if is_file(download_file_path): - initial = os.path.getsize(download_file_path) - else: - initial = 0 - if initial < total: - with tqdm(total = total, initial = initial, desc = wording.get('downloading'), unit = 'B', unit_scale = True, unit_divisor = 1024, ascii = ' =') as progress: - subprocess.Popen([ 'curl', '--create-dirs', '--silent', '--insecure', '--location', '--continue-at', '-', '--output', download_file_path, url ]) - current = initial - while current < total: - if is_file(download_file_path): - current = os.path.getsize(download_file_path) - progress.update(current - progress.n) - - -@lru_cache(maxsize = None) -def get_download_size(url : str) -> int: - try: - response = urllib.request.urlopen(url, timeout = 10) - return int(response.getheader('Content-Length')) - except (OSError, ValueError): - return 0 - - -def is_download_done(url : str, file_path : str) -> bool: - if is_file(file_path): - return get_download_size(url) == os.path.getsize(file_path) - return False - - -def resolve_relative_path(path : str) -> str: - return os.path.abspath(os.path.join(os.path.dirname(__file__), path)) - - -def list_module_names(path : str) -> Optional[List[str]]: - if os.path.exists(path): - files = os.listdir(path) - return [ Path(file).stem for file in files if not Path(file).stem.startswith(('.', '__')) ] - return None - - -def encode_execution_providers(execution_providers : List[str]) -> List[str]: - return [ execution_provider.replace('ExecutionProvider', '').lower() for execution_provider in execution_providers ] - - -def decode_execution_providers(execution_providers: List[str]) -> List[str]: - available_execution_providers = onnxruntime.get_available_providers() - encoded_execution_providers = encode_execution_providers(available_execution_providers) - return [ execution_provider for execution_provider, encoded_execution_provider in zip(available_execution_providers, encoded_execution_providers) if any(execution_provider in encoded_execution_provider for execution_provider in execution_providers) ] - - -def map_device(execution_providers : List[str]) -> str: - if 'CoreMLExecutionProvider' in execution_providers: - return 'mps' - if 'CUDAExecutionProvider' in execution_providers or 'ROCMExecutionProvider' in execution_providers : - return 'cuda' - if 'OpenVINOExecutionProvider' in execution_providers: - return 'mkl' - return 'cpu' - - -def create_metavar(ranges : List[Any]) -> str: - return '[' + str(ranges[0]) + '-' + str(ranges[-1]) + ']' - - -def update_status(message : str, scope : str = 'FACEFUSION.CORE') -> None: - print('[' + scope + '] ' + message) diff --git a/facefusion/vision.py b/facefusion/vision.py index f5ee547d..4706bf7c 100644 --- a/facefusion/vision.py +++ b/facefusion/vision.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Optional, List from functools import lru_cache import cv2 @@ -55,6 +55,14 @@ def read_static_image(image_path : str) -> Optional[Frame]: return read_image(image_path) +def read_static_images(image_paths : List[str]) -> Optional[List[Frame]]: + frames = [] + if image_paths: + for image_path in image_paths: + frames.append(read_static_image(image_path)) + return frames + + def read_image(image_path : str) -> Optional[Frame]: if image_path: return cv2.imread(image_path) diff --git a/facefusion/wording.py b/facefusion/wording.py index c7c739bd..78f3cd3d 100755 --- a/facefusion/wording.py +++ b/facefusion/wording.py @@ -3,13 +3,14 @@ WORDING =\ 'python_not_supported': 'Python version is not supported, upgrade to {version} or higher', 'ffmpeg_not_installed': 'FFMpeg is not installed', 'install_dependency_help': 'select the variant of {dependency} to install', + 'skip_venv_help': 'skip the virtual environment check', 'source_help': 'select a source image', 'target_help': 'select a target image or video', 'output_help': 'specify the output file or directory', 'frame_processors_help': 'choose from the available frame processors (choices: {choices}, ...)', 'frame_processor_model_help': 'choose the model for the frame processor', - 'frame_processor_blend_help': 'specify the blend factor for the frame processor', - 'face_debugger_items_help': 'specify the face debugger items', + 'frame_processor_blend_help': 'specify the blend amount for the frame processor', + 'face_debugger_items_help': 'specify the face debugger items (choices: {choices})', 'ui_layouts_help': 'choose from the available ui layouts (choices: {choices}, ...)', 'keep_fps_help': 'preserve the frames per second (fps) of the target', 'keep_temp_help': 'retain temporary frames after processing', @@ -24,8 +25,10 @@ WORDING =\ 'reference_face_position_help': 'specify the position of the reference face', 'reference_face_distance_help': 'specify the distance between the reference face and the target face', 'reference_frame_number_help': 'specify the number of the reference frame', + 'face_mask_types_help': 'choose from the available face mask types (choices: {choices})', 'face_mask_blur_help': 'specify the blur amount for face mask', 'face_mask_padding_help': 'specify the face mask padding (top, right, bottom, left) in percent', + 'face_mask_regions_help': 'choose from the available face mask regions (choices: {choices})', 'trim_frame_start_help': 'specify the start frame for extraction', 'trim_frame_end_help': 'specify the end frame for extraction', 'temp_frame_format_help': 'specify the image format used for frame extraction', @@ -34,11 +37,12 @@ WORDING =\ 'output_video_encoder_help': 'specify the encoder used for the output video', 'output_video_quality_help': 'specify the quality used for the output video', 'max_memory_help': 'specify the maximum amount of ram to be used (in gb)', - 'execution_providers_help': 'choose from the available execution providers', + 'execution_providers_help': 'choose from the available execution providers (choices: {choices}, ...)', 'execution_thread_count_help': 'specify the number of execution threads', 'execution_queue_count_help': 'specify the number of execution queries', 'skip_download_help': 'omit automate downloads and lookups', 'headless_help': 'run the program in headless mode', + 'log_level_help': 'choose from the available log levels', 'creating_temp': 'Creating temporary resources', 'extracting_frames_fps': 'Extracting frames with {fps} FPS', 'analysing': 'Analysing', @@ -51,7 +55,7 @@ WORDING =\ 'merging_video_failed': 'Merging video failed', 'skipping_audio': 'Skipping audio', 'restoring_audio': 'Restoring audio', - 'restoring_audio_failed': 'Restoring audio failed', + 'restoring_audio_skipped': 'Restoring audio skipped', 'clearing_temp': 'Clearing temporary resources', 'processing_image_succeed': 'Processing to image succeed', 'processing_image_failed': 'Processing to image failed', @@ -67,6 +71,7 @@ WORDING =\ 'frame_processor_not_implemented': 'Frame processor {frame_processor} not implemented correctly', 'ui_layout_not_loaded': 'UI layout {ui_layout} could not be loaded', 'ui_layout_not_implemented': 'UI layout {ui_layout} not implemented correctly', + 'stream_not_loaded': 'Stream {stream_mode} could not be loaded', 'donate_button_label': 'DONATE', 'start_button_label': 'START', 'stop_button_label': 'STOP', @@ -86,11 +91,13 @@ WORDING =\ 'face_selector_mode_dropdown_label': 'FACE SELECTOR MODE', 'reference_face_gallery_label': 'REFERENCE FACE', 'reference_face_distance_slider_label': 'REFERENCE FACE DISTANCE', + 'face_mask_types_checkbox_group_label': 'FACE MASK TYPES', 'face_mask_blur_slider_label': 'FACE MASK BLUR', 'face_mask_padding_top_slider_label': 'FACE MASK PADDING TOP', 'face_mask_padding_bottom_slider_label': 'FACE MASK PADDING BOTTOM', 'face_mask_padding_left_slider_label': 'FACE MASK PADDING LEFT', 'face_mask_padding_right_slider_label': 'FACE MASK PADDING RIGHT', + 'face_mask_region_checkbox_group_label': 'FACE MASK REGIONS', 'max_memory_slider_label': 'MAX MEMORY', 'output_image_or_video_label': 'OUTPUT', 'output_path_textbox_label': 'OUTPUT PATH', diff --git a/requirements.txt b/requirements.txt index 66a16b95..d6929853 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,11 +1,11 @@ basicsr==1.4.2 filetype==1.2.0 gradio==3.50.2 -numpy==1.26.1 +numpy==1.26.2 onnx==1.15.0 -onnxruntime==1.16.0 +onnxruntime==1.16.3 opencv-python==4.8.1.78 psutil==5.9.6 realesrgan==0.3.0 -torch==2.1.0 +torch==2.1.1 tqdm==4.66.1 diff --git a/tests/test_cli.py b/tests/test_cli.py index 65104eab..0935222d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3,7 +3,7 @@ import sys import pytest from facefusion import wording -from facefusion.utilities import conditional_download +from facefusion.download import conditional_download @pytest.fixture(scope = 'module', autouse = True) @@ -18,7 +18,7 @@ def before_all() -> None: def test_image_to_image() -> None: commands = [ sys.executable, 'run.py', '-s', '.assets/examples/source.jpg', '-t', '.assets/examples/target-1080p.jpg', '-o', '.assets/examples', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE) + run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) assert run.returncode == 0 assert wording.get('processing_image_succeed') in run.stdout.decode() @@ -26,7 +26,7 @@ def test_image_to_image() -> None: def test_image_to_video() -> None: commands = [ sys.executable, 'run.py', '-s', '.assets/examples/source.jpg', '-t', '.assets/examples/target-1080p.mp4', '-o', '.assets/examples', '--trim-frame-end', '10', '--headless' ] - run = subprocess.run(commands, stdout = subprocess.PIPE) + run = subprocess.run(commands, stdout = subprocess.PIPE, stderr = subprocess.STDOUT) assert run.returncode == 0 assert wording.get('processing_video_succeed') in run.stdout.decode() diff --git a/tests/test_download.py b/tests/test_download.py new file mode 100644 index 00000000..f80c44b8 --- /dev/null +++ b/tests/test_download.py @@ -0,0 +1,23 @@ +import pytest + +from facefusion.download import conditional_download, get_download_size, is_download_done + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download('.assets/examples', + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4' + ]) + + +def test_get_download_size() -> None: + assert get_download_size('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4') == 191675 + assert get_download_size('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-360p.mp4') == 370732 + assert get_download_size('invalid') == 0 + + +def test_is_download_done() -> None: + assert is_download_done('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4', '.assets/examples/target-240p.mp4') is True + assert is_download_done('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4','invalid') is False + assert is_download_done('invalid', 'invalid') is False diff --git a/tests/test_execution_helper.py b/tests/test_execution_helper.py new file mode 100644 index 00000000..5d199123 --- /dev/null +++ b/tests/test_execution_helper.py @@ -0,0 +1,9 @@ +from facefusion.execution_helper import encode_execution_providers, decode_execution_providers + + +def test_encode_execution_providers() -> None: + assert encode_execution_providers([ 'CPUExecutionProvider' ]) == [ 'cpu' ] + + +def test_decode_execution_providers() -> None: + assert decode_execution_providers([ 'cpu' ]) == [ 'CPUExecutionProvider' ] diff --git a/tests/test_ffmpeg.py b/tests/test_ffmpeg.py new file mode 100644 index 00000000..b67ed759 --- /dev/null +++ b/tests/test_ffmpeg.py @@ -0,0 +1,100 @@ +import glob +import subprocess +import pytest + +import facefusion.globals +from facefusion.filesystem import get_temp_directory_path, create_temp, clear_temp +from facefusion.download import conditional_download +from facefusion.ffmpeg import extract_frames + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download('.assets/examples', + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4' + ]) + subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=25', '.assets/examples/target-240p-25fps.mp4' ]) + subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=30', '.assets/examples/target-240p-30fps.mp4' ]) + subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=60', '.assets/examples/target-240p-60fps.mp4' ]) + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + facefusion.globals.trim_frame_start = None + facefusion.globals.trim_frame_end = None + facefusion.globals.temp_frame_quality = 80 + facefusion.globals.temp_frame_format = 'jpg' + + +def test_extract_frames() -> None: + target_paths =\ + [ + '.assets/examples/target-240p-25fps.mp4', + '.assets/examples/target-240p-30fps.mp4', + '.assets/examples/target-240p-60fps.mp4' + ] + for target_path in target_paths: + temp_directory_path = get_temp_directory_path(target_path) + create_temp(target_path) + + assert extract_frames(target_path, 30.0) is True + assert len(glob.glob1(temp_directory_path, '*.jpg')) == 324 + + clear_temp(target_path) + + +def test_extract_frames_with_trim_start() -> None: + facefusion.globals.trim_frame_start = 224 + data_provider =\ + [ + ('.assets/examples/target-240p-25fps.mp4', 55), + ('.assets/examples/target-240p-30fps.mp4', 100), + ('.assets/examples/target-240p-60fps.mp4', 212) + ] + for target_path, frame_total in data_provider: + temp_directory_path = get_temp_directory_path(target_path) + create_temp(target_path) + + assert extract_frames(target_path, 30.0) is True + assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total + + clear_temp(target_path) + + +def test_extract_frames_with_trim_start_and_trim_end() -> None: + facefusion.globals.trim_frame_start = 124 + facefusion.globals.trim_frame_end = 224 + data_provider =\ + [ + ('.assets/examples/target-240p-25fps.mp4', 120), + ('.assets/examples/target-240p-30fps.mp4', 100), + ('.assets/examples/target-240p-60fps.mp4', 50) + ] + for target_path, frame_total in data_provider: + temp_directory_path = get_temp_directory_path(target_path) + create_temp(target_path) + + assert extract_frames(target_path, 30.0) is True + assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total + + clear_temp(target_path) + + +def test_extract_frames_with_trim_end() -> None: + facefusion.globals.trim_frame_end = 100 + data_provider =\ + [ + ('.assets/examples/target-240p-25fps.mp4', 120), + ('.assets/examples/target-240p-30fps.mp4', 100), + ('.assets/examples/target-240p-60fps.mp4', 50) + ] + for target_path, frame_total in data_provider: + temp_directory_path = get_temp_directory_path(target_path) + create_temp(target_path) + + assert extract_frames(target_path, 30.0) is True + assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total + + clear_temp(target_path) diff --git a/tests/test_filesystem.py b/tests/test_filesystem.py new file mode 100644 index 00000000..c5047421 --- /dev/null +++ b/tests/test_filesystem.py @@ -0,0 +1,31 @@ +from facefusion.filesystem import is_file, is_directory, is_image, are_images, is_video + + +def test_is_file() -> None: + assert is_file('.assets/examples/source.jpg') is True + assert is_file('.assets/examples') is False + assert is_file('invalid') is False + + +def test_is_directory() -> None: + assert is_directory('.assets/examples') is True + assert is_directory('.assets/examples/source.jpg') is False + assert is_directory('invalid') is False + + +def test_is_image() -> None: + assert is_image('.assets/examples/source.jpg') is True + assert is_image('.assets/examples/target-240p.mp4') is False + assert is_image('invalid') is False + + +def test_are_images() -> None: + assert are_images([ '.assets/examples/source.jpg' ]) is True + assert are_images([ '.assets/examples/source.jpg', '.assets/examples/target-240p.mp4' ]) is False + assert are_images([ 'invalid' ]) is False + + +def test_is_video() -> None: + assert is_video('.assets/examples/target-240p.mp4') is True + assert is_video('.assets/examples/source.jpg') is False + assert is_video('invalid') is False diff --git a/tests/test_normalizer.py b/tests/test_normalizer.py new file mode 100644 index 00000000..0547e12d --- /dev/null +++ b/tests/test_normalizer.py @@ -0,0 +1,25 @@ +import platform + +from facefusion.normalizer import normalize_output_path, normalize_padding + + +def test_normalize_output_path() -> None: + if platform.system().lower() != 'windows': + assert normalize_output_path([ '.assets/examples/source.jpg' ], None, '.assets/examples/target-240p.mp4') == '.assets/examples/target-240p.mp4' + assert normalize_output_path(None, '.assets/examples/target-240p.mp4', '.assets/examples/target-240p.mp4') == '.assets/examples/target-240p.mp4' + assert normalize_output_path(None, '.assets/examples/target-240p.mp4', '.assets/examples') == '.assets/examples/target-240p.mp4' + assert normalize_output_path([ '.assets/examples/source.jpg' ], '.assets/examples/target-240p.mp4', '.assets/examples') == '.assets/examples/source-target-240p.mp4' + assert normalize_output_path(None, '.assets/examples/target-240p.mp4', '.assets/examples/output.mp4') == '.assets/examples/output.mp4' + assert normalize_output_path(None, '.assets/examples/target-240p.mp4', '.assets/output.mov') == '.assets/output.mp4' + assert normalize_output_path(None, '.assets/examples/target-240p.mp4', '.assets/examples/invalid') is None + assert normalize_output_path(None, '.assets/examples/target-240p.mp4', '.assets/invalid/output.mp4') is None + assert normalize_output_path(None, '.assets/examples/target-240p.mp4', 'invalid') is None + assert normalize_output_path([ '.assets/examples/source.jpg' ], '.assets/examples/target-240p.mp4', None) is None + + +def test_normalize_padding() -> None: + assert normalize_padding([ 0, 0, 0, 0 ]) == (0, 0, 0, 0) + assert normalize_padding([ 1 ]) == (1, 1, 1, 1) + assert normalize_padding([ 1, 2 ]) == (1, 2, 1, 2) + assert normalize_padding([ 1, 2, 3 ]) == (1, 2, 3, 2) + assert normalize_padding(None) is None diff --git a/tests/test_utilities.py b/tests/test_utilities.py deleted file mode 100644 index 934b3400..00000000 --- a/tests/test_utilities.py +++ /dev/null @@ -1,169 +0,0 @@ -import glob -import platform -import subprocess -import pytest - -import facefusion.globals -from facefusion.utilities import conditional_download, extract_frames, create_temp, get_temp_directory_path, clear_temp, normalize_output_path, normalize_padding, is_file, is_directory, is_image, is_video, get_download_size, is_download_done, encode_execution_providers, decode_execution_providers - - -@pytest.fixture(scope = 'module', autouse = True) -def before_all() -> None: - facefusion.globals.temp_frame_quality = 100 - facefusion.globals.trim_frame_start = None - facefusion.globals.trim_frame_end = None - facefusion.globals.temp_frame_format = 'png' - conditional_download('.assets/examples', - [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', - 'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4' - ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=25', '.assets/examples/target-240p-25fps.mp4' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=30', '.assets/examples/target-240p-30fps.mp4' ]) - subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=60', '.assets/examples/target-240p-60fps.mp4' ]) - - -@pytest.fixture(scope = 'function', autouse = True) -def before_each() -> None: - facefusion.globals.trim_frame_start = None - facefusion.globals.trim_frame_end = None - facefusion.globals.temp_frame_quality = 90 - facefusion.globals.temp_frame_format = 'jpg' - - -def test_extract_frames() -> None: - target_paths =\ - [ - '.assets/examples/target-240p-25fps.mp4', - '.assets/examples/target-240p-30fps.mp4', - '.assets/examples/target-240p-60fps.mp4' - ] - for target_path in target_paths: - temp_directory_path = get_temp_directory_path(target_path) - create_temp(target_path) - - assert extract_frames(target_path, 30.0) is True - assert len(glob.glob1(temp_directory_path, '*.jpg')) == 324 - - clear_temp(target_path) - - -def test_extract_frames_with_trim_start() -> None: - facefusion.globals.trim_frame_start = 224 - data_provider =\ - [ - ('.assets/examples/target-240p-25fps.mp4', 55), - ('.assets/examples/target-240p-30fps.mp4', 100), - ('.assets/examples/target-240p-60fps.mp4', 212) - ] - for target_path, frame_total in data_provider: - temp_directory_path = get_temp_directory_path(target_path) - create_temp(target_path) - - assert extract_frames(target_path, 30.0) is True - assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total - - clear_temp(target_path) - - -def test_extract_frames_with_trim_start_and_trim_end() -> None: - facefusion.globals.trim_frame_start = 124 - facefusion.globals.trim_frame_end = 224 - data_provider =\ - [ - ('.assets/examples/target-240p-25fps.mp4', 120), - ('.assets/examples/target-240p-30fps.mp4', 100), - ('.assets/examples/target-240p-60fps.mp4', 50) - ] - for target_path, frame_total in data_provider: - temp_directory_path = get_temp_directory_path(target_path) - create_temp(target_path) - - assert extract_frames(target_path, 30.0) is True - assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total - - clear_temp(target_path) - - -def test_extract_frames_with_trim_end() -> None: - facefusion.globals.trim_frame_end = 100 - data_provider =\ - [ - ('.assets/examples/target-240p-25fps.mp4', 120), - ('.assets/examples/target-240p-30fps.mp4', 100), - ('.assets/examples/target-240p-60fps.mp4', 50) - ] - for target_path, frame_total in data_provider: - temp_directory_path = get_temp_directory_path(target_path) - create_temp(target_path) - - assert extract_frames(target_path, 30.0) is True - assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total - - clear_temp(target_path) - - -def test_normalize_output_path() -> None: - if platform.system().lower() != 'windows': - assert normalize_output_path('.assets/examples/source.jpg', None, '.assets/examples/target-240p.mp4') == '.assets/examples/target-240p.mp4' - assert normalize_output_path(None, '.assets/examples/target-240p.mp4', '.assets/examples/target-240p.mp4') == '.assets/examples/target-240p.mp4' - assert normalize_output_path(None, '.assets/examples/target-240p.mp4', '.assets/examples') == '.assets/examples/target-240p.mp4' - assert normalize_output_path('.assets/examples/source.jpg', '.assets/examples/target-240p.mp4', '.assets/examples') == '.assets/examples/source-target-240p.mp4' - assert normalize_output_path(None, '.assets/examples/target-240p.mp4', '.assets/examples/output.mp4') == '.assets/examples/output.mp4' - assert normalize_output_path(None, '.assets/examples/target-240p.mp4', '.assets/output.mov') == '.assets/output.mp4' - assert normalize_output_path(None, '.assets/examples/target-240p.mp4', '.assets/examples/invalid') is None - assert normalize_output_path(None, '.assets/examples/target-240p.mp4', '.assets/invalid/output.mp4') is None - assert normalize_output_path(None, '.assets/examples/target-240p.mp4', 'invalid') is None - assert normalize_output_path('.assets/examples/source.jpg', '.assets/examples/target-240p.mp4', None) is None - - -def test_normalize_padding() -> None: - assert normalize_padding([ 0, 0, 0, 0 ]) == (0, 0, 0, 0) - assert normalize_padding([ 1 ]) == (1, 1, 1, 1) - assert normalize_padding([ 1, 2 ]) == (1, 2, 1, 2) - assert normalize_padding([ 1, 2, 3 ]) == (1, 2, 3, 2) - assert normalize_padding(None) is None - - -def test_is_file() -> None: - assert is_file('.assets/examples/source.jpg') is True - assert is_file('.assets/examples') is False - assert is_file('invalid') is False - - -def test_is_directory() -> None: - assert is_directory('.assets/examples') is True - assert is_directory('.assets/examples/source.jpg') is False - assert is_directory('invalid') is False - - -def test_is_image() -> None: - assert is_image('.assets/examples/source.jpg') is True - assert is_image('.assets/examples/target-240p.mp4') is False - assert is_image('invalid') is False - - -def test_is_video() -> None: - assert is_video('.assets/examples/target-240p.mp4') is True - assert is_video('.assets/examples/source.jpg') is False - assert is_video('invalid') is False - - -def test_get_download_size() -> None: - assert get_download_size('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4') == 191675 - assert get_download_size('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-360p.mp4') == 370732 - assert get_download_size('invalid') == 0 - - -def test_is_download_done() -> None: - assert is_download_done('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4', '.assets/examples/target-240p.mp4') is True - assert is_download_done('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4','invalid') is False - assert is_download_done('invalid', 'invalid') is False - - -def test_encode_execution_providers() -> None: - assert encode_execution_providers([ 'CPUExecutionProvider' ]) == [ 'cpu' ] - - -def test_decode_execution_providers() -> None: - assert decode_execution_providers([ 'cpu' ]) == [ 'CPUExecutionProvider' ] diff --git a/tests/test_vision.py b/tests/test_vision.py index f77af049..5b51e629 100644 --- a/tests/test_vision.py +++ b/tests/test_vision.py @@ -1,17 +1,12 @@ import subprocess import pytest -import facefusion.globals -from facefusion.utilities import conditional_download +from facefusion.download import conditional_download from facefusion.vision import get_video_frame, detect_fps, count_video_frame_total @pytest.fixture(scope = 'module', autouse = True) def before_all() -> None: - facefusion.globals.temp_frame_quality = 100 - facefusion.globals.trim_frame_start = None - facefusion.globals.trim_frame_end = None - facefusion.globals.temp_frame_format = 'png' conditional_download('.assets/examples', [ 'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg', @@ -22,14 +17,6 @@ def before_all() -> None: subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=60', '.assets/examples/target-240p-60fps.mp4' ]) -@pytest.fixture(scope = 'function', autouse = True) -def before_each() -> None: - facefusion.globals.trim_frame_start = None - facefusion.globals.trim_frame_end = None - facefusion.globals.temp_frame_quality = 90 - facefusion.globals.temp_frame_format = 'jpg' - - def test_get_video_frame() -> None: assert get_video_frame('.assets/examples/target-240p-25fps.mp4') is not None assert get_video_frame('invalid') is None