diff --git a/facefusion/content_analyser.py b/facefusion/content_analyser.py index cf21f06f..23de238c 100644 --- a/facefusion/content_analyser.py +++ b/facefusion/content_analyser.py @@ -9,7 +9,7 @@ from facefusion.download import conditional_download_hashes, conditional_downloa from facefusion.filesystem import resolve_relative_path from facefusion.thread_helper import conditional_thread_semaphore from facefusion.types import Detection, DownloadScope, DownloadSet, Fps, InferencePool, ModelSet, VisionFrame -from facefusion.vision import detect_video_fps, fit_contain_frame, read_image, read_video_frame +from facefusion.vision import detect_video_fps, fit_contain_frame, read_image, sample_video_frames STREAM_COUNTER = 0 @@ -158,15 +158,18 @@ def analyse_image(image_path : str) -> bool: def analyse_video(video_path : str, trim_frame_start : int, trim_frame_end : int) -> bool: video_fps = detect_video_fps(video_path) frame_range = range(trim_frame_start, trim_frame_end) + vision_frames = sample_video_frames(video_path, trim_frame_start, trim_frame_end, int(video_fps)) rate = 0.0 total = 0 counter = 0 + sample_index = 0 with tqdm(total = len(frame_range), desc = translator.get('analysing'), unit = 'frame', ascii = ' =', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress: for frame_number in frame_range: - if frame_number % int(video_fps) == 0: - vision_frame = read_video_frame(video_path, frame_number) + if frame_number % int(video_fps) == 0 and sample_index < len(vision_frames): + vision_frame = vision_frames[sample_index] + sample_index += 1 if numpy.any(vision_frame): total += 1 diff --git a/facefusion/core.py b/facefusion/core.py index c8168c39..adeb5dcc 100755 --- a/facefusion/core.py +++ b/facefusion/core.py @@ -103,7 +103,7 @@ def pre_check() -> bool: def common_pre_check() -> bool: content_analyser_content = inspect.getsource(content_analyser).encode() - return hash_helper.create_hash(content_analyser_content) == '975d67d6' + return hash_helper.create_hash(content_analyser_content) == 'dfca29ec' def processors_pre_check() -> bool: diff --git a/facefusion/ffmpeg_builder.py b/facefusion/ffmpeg_builder.py index 53b1bd40..6769161c 100644 --- a/facefusion/ffmpeg_builder.py +++ b/facefusion/ffmpeg_builder.py @@ -107,6 +107,14 @@ def set_frame_quality(frame_quality : int) -> List[Command]: return [ '-q:v', str(frame_quality) ] +def select_frame_samples(frame_start : int, frame_end : int, frame_stride : int) -> List[Command]: + return [ '-vf', 'select=between(n\\,' + str(frame_start) + '\\,' + str(frame_end - 1) + ')*not(mod(n\\,' + str(frame_stride) + '))' ] + + +def set_frame_total(frame_total : int) -> List[Command]: + return [ '-frames:v', str(frame_total) ] + + def select_frame_range(frame_start : int, frame_end : int, video_fps : Fps) -> List[Command]: if isinstance(frame_start, int) and isinstance(frame_end, int): return [ '-vf', 'trim=start_frame=' + str(frame_start) + ':end_frame=' + str(frame_end) + ',fps=' + str(video_fps) ] diff --git a/facefusion/video_manager.py b/facefusion/video_manager.py index fc04e854..7fd3a2f8 100644 --- a/facefusion/video_manager.py +++ b/facefusion/video_manager.py @@ -1,9 +1,11 @@ +import os import subprocess from typing import Optional, Tuple, cast import numpy from facefusion import ffmpeg_builder, state_manager +from facefusion.common_helper import is_windows from facefusion.ffprobe import extract_video_metadata from facefusion.filesystem import get_file_format from facefusion.temp_helper import get_temp_file_path @@ -29,6 +31,30 @@ def create_video_reader_process(video_path : str, frame_position : int, video_fp return subprocess.Popen(commands, stdout = subprocess.PIPE, stderr = subprocess.DEVNULL) +def create_video_sampler_process(video_path : str, frame_start : int, frame_end : int, frame_stride : int) -> subprocess.Popen[bytes]: + sample_total = len([ frame_number for frame_number in range(frame_start, frame_end) if frame_number % frame_stride == 0 ]) + sampler_thread_count = 4 + commands = ffmpeg_builder.chain( + ffmpeg_builder.set_thread_count(sampler_thread_count), + ffmpeg_builder.set_input(video_path), + ffmpeg_builder.select_frame_samples(frame_start, frame_end, frame_stride), + ffmpeg_builder.prevent_frame_drop(), + ffmpeg_builder.set_frame_total(sample_total), + ffmpeg_builder.enforce_pixel_format('bgr24'), + ffmpeg_builder.set_output_format('rawvideo'), + ffmpeg_builder.cast_stream() + ) + commands = ffmpeg_builder.run(commands) + + if is_windows(): + return subprocess.Popen(commands, stdout = subprocess.PIPE, stderr = subprocess.DEVNULL, creationflags = subprocess.IDLE_PRIORITY_CLASS) + return subprocess.Popen(commands, stdout = subprocess.PIPE, stderr = subprocess.DEVNULL, preexec_fn = demote_process_priority) + + +def demote_process_priority() -> None: + os.nice(10) + + def get_video_reader(video_path : str) -> VideoReader: if video_path not in VIDEO_POOL_SET.get('reader'): video_metadata = extract_video_metadata(video_path) diff --git a/facefusion/vision.py b/facefusion/vision.py index b007f903..aca98960 100644 --- a/facefusion/vision.py +++ b/facefusion/vision.py @@ -11,7 +11,7 @@ from facefusion.ffprobe import extract_video_metadata from facefusion.filesystem import get_file_extension, is_image, is_video from facefusion.thread_helper import thread_lock, thread_semaphore from facefusion.types import ColorMode, Duration, Fps, Mask, Orientation, Resolution, Scale, VisionFrame -from facefusion.video_manager import conditional_set_video_reader_position, get_video_reader, read_video_reader_frame, read_video_reader_window +from facefusion.video_manager import conditional_set_video_reader_position, create_video_sampler_process, get_video_reader, read_video_reader_frame, read_video_reader_window def read_static_images(image_paths : List[str], color_mode : ColorMode = 'rgb') -> List[VisionFrame]: @@ -93,6 +93,27 @@ def read_video_frame(video_path : str, frame_number : int = 0) -> Optional[Visio return None +def sample_video_frames(video_path : str, frame_start : int, frame_end : int, frame_stride : int) -> List[VisionFrame]: + vision_frames = [] + + if is_video(video_path): + width, height = extract_video_metadata(video_path).get('resolution') + frame_size = width * height * 3 + sample_total = len([ frame_number for frame_number in range(frame_start, frame_end) if frame_number % frame_stride == 0 ]) + video_sampler = create_video_sampler_process(video_path, frame_start, frame_end, frame_stride) + + for _ in range(sample_total): + frame_buffer = video_sampler.stdout.read(frame_size) + + if len(frame_buffer) == frame_size: + vision_frames.append(numpy.frombuffer(frame_buffer, dtype = numpy.uint8).reshape(height, width, 3)) + + video_sampler.stdout.close() + video_sampler.wait() + + return vision_frames + + @lru_cache(maxsize = 2) def read_static_video_chunk(video_path : str, chunk_number : int, chunk_size : int) -> Dict[int, VisionFrame]: return read_video_chunk(video_path, chunk_number, chunk_size) diff --git a/facefusion/workflows/image_to_video.py b/facefusion/workflows/image_to_video.py index 1bde61d3..4d532b20 100644 --- a/facefusion/workflows/image_to_video.py +++ b/facefusion/workflows/image_to_video.py @@ -1,4 +1,4 @@ -from concurrent.futures import ThreadPoolExecutor, as_completed +from concurrent.futures import Future, ThreadPoolExecutor, as_completed from functools import partial from typing import List, Tuple @@ -21,43 +21,51 @@ from facefusion.workflows.core import is_process_stopping def process(start_time : float) -> ErrorCode: - tasks =\ - [ - setup, - extract_frames, - process_disk_frames, - merge_frames, - restore_audio, - partial(finalize_video, start_time) - ] + trim_frame_start, trim_frame_end = restrict_trim_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) - if state_manager.get_item('workflow_mode') == 'stream': + with ThreadPoolExecutor(max_workers = 1) as analyse_executor: + analyse_future = analyse_executor.submit(analyse_video, state_manager.get_item('target_path'), trim_frame_start, trim_frame_end) tasks =\ [ setup, - process_stream_frames, + extract_frames, + process_disk_frames, + merge_frames, + partial(enforce_analysis, analyse_future), restore_audio, partial(finalize_video, start_time) ] - process_manager.start() - for task in tasks: - error_code = task() #type:ignore[operator] + if state_manager.get_item('workflow_mode') == 'stream': + tasks =\ + [ + setup, + process_stream_frames, + partial(enforce_analysis, analyse_future), + restore_audio, + partial(finalize_video, start_time) + ] + process_manager.start() - if error_code > 0: - process_manager.end() - return error_code + for task in tasks: + error_code = task() #type:ignore[operator] - process_manager.end() + if error_code > 0: + process_manager.end() + return error_code + + process_manager.end() + return 0 + + +def enforce_analysis(analyse_future : Future[bool]) -> ErrorCode: + if analyse_future.result(): + clear_temp_directory(state_manager.get_item('target_path')) + return 3 return 0 def setup() -> ErrorCode: - trim_frame_start, trim_frame_end = restrict_trim_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) - - if analyse_video(state_manager.get_item('target_path'), trim_frame_start, trim_frame_end): - return 3 - if clear_temp_directory(state_manager.get_item('target_path')): logger.debug(translator.get('clearing_temp'), __name__)