diff --git a/facefusion/content_analyser.py b/facefusion/content_analyser.py index cf21f06f..a53a291f 100644 --- a/facefusion/content_analyser.py +++ b/facefusion/content_analyser.py @@ -4,12 +4,12 @@ from typing import Tuple import numpy from tqdm import tqdm -from facefusion import inference_manager, state_manager, translator +from facefusion import inference_manager, state_manager, translator, video_manager from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url 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 STREAM_COUNTER = 0 @@ -161,13 +161,15 @@ def analyse_video(video_path : str, trim_frame_start : int, trim_frame_end : int rate = 0.0 total = 0 counter = 0 + video_reader = video_manager.get_reader(video_path, 'analyse_video') + video_manager.seek_video_reader(video_reader, trim_frame_start) 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) + vision_frame = video_manager.read_video_frame(video_reader) + if frame_number % int(video_fps) == 0: if numpy.any(vision_frame): total += 1 diff --git a/facefusion/core.py b/facefusion/core.py index 058eb86f..ec55a4fe 100755 --- a/facefusion/core.py +++ b/facefusion/core.py @@ -100,7 +100,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) == '5f04e67c' def processors_pre_check() -> bool: diff --git a/facefusion/ffmpeg.py b/facefusion/ffmpeg.py index 8034cc61..cf2f782d 100644 --- a/facefusion/ffmpeg.py +++ b/facefusion/ffmpeg.py @@ -69,9 +69,9 @@ def open_ffmpeg(commands : List[Command]) -> subprocess.Popen[bytes]: return subprocess.Popen(commands, stdin = subprocess.PIPE, stdout = subprocess.PIPE) -def create_video_reader(video_path : str, frame_position : int, video_metadata : VideoReaderMetadata) -> subprocess.Popen[bytes]: +def create_video_reader(video_path : str, frame_number : int, video_metadata : VideoReaderMetadata) -> subprocess.Popen[bytes]: commands = ffmpeg_builder.chain( - ffmpeg_builder.seek_to(frame_position / video_metadata.get('fps')), + ffmpeg_builder.seek_to(frame_number / video_metadata.get('fps')), ffmpeg_builder.set_input(video_path), ffmpeg_builder.restrict_color_transfer(video_metadata.get('color_transfer')), ffmpeg_builder.prevent_frame_drop(), diff --git a/facefusion/types.py b/facefusion/types.py index 1e59f849..6d592368 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -135,7 +135,7 @@ VideoReader = TypedDict('VideoReader', 'file_path' : str, 'process' : subprocess.Popen[bytes], 'metadata' : VideoReaderMetadata, - 'position' : int + 'frame_number' : int }) VideoReaderSet : TypeAlias = Dict[str, VideoReader] VideoWriter = TypedDict('VideoWriter', diff --git a/facefusion/video_manager.py b/facefusion/video_manager.py index 8da514ce..128c2f9d 100644 --- a/facefusion/video_manager.py +++ b/facefusion/video_manager.py @@ -1,5 +1,7 @@ +import hashlib import uuid -from typing import Optional +from io import BufferedReader +from typing import Optional, cast import numpy @@ -14,81 +16,97 @@ VIDEO_POOL_SET : VideoPoolSet =\ } -def get_reader(video_path : str) -> VideoReader: - if video_path not in VIDEO_POOL_SET.get('reader'): +def get_reader(video_path : str, context : str) -> VideoReader: + reader_id = hashlib.sha1((video_path + '_' + context).encode()).hexdigest() + + if reader_id not in VIDEO_POOL_SET.get('reader'): video_metadata = ffprobe.extract_static_video_metadata(video_path) - VIDEO_POOL_SET['reader'][video_path] =\ + VIDEO_POOL_SET['reader'][reader_id] =\ { - 'id': uuid.uuid4().hex, + 'id': reader_id, 'file_path': video_path, 'process': ffmpeg.create_video_reader(video_path, 0, video_metadata), 'metadata': video_metadata, - 'position': 0 + 'frame_number': 0 } - return VIDEO_POOL_SET.get('reader').get(video_path) + return VIDEO_POOL_SET.get('reader').get(reader_id) -#todo: needs review - [seeking] [critical: high] forward skip up to 128 frames by draining the pipe, everything else refreshes the process -def conditional_set_video_reader_position(video_reader : VideoReader, frame_position : int) -> None: - skip_total = frame_position - video_reader.get('position') +def conditional_seek_video_reader(video_reader : VideoReader, frame_number : int = 0) -> None: + frame_number = min(video_reader.get('metadata').get('frame_total'), frame_number) + skip_total = frame_number - video_reader.get('frame_number') skip_margin = 128 if 0 < skip_total <= skip_margin: - for _ in range(skip_total): - read_video_reader_frame(video_reader) + drain_video_reader(video_reader, skip_total) - if not video_reader.get('position') == frame_position: - refresh_video_reader(video_reader, frame_position) + if not video_reader.get('frame_number') == frame_number: + seek_video_reader(video_reader, frame_number) -def refresh_video_reader(video_reader : VideoReader, frame_position : int) -> None: +def seek_video_reader(video_reader : VideoReader, frame_number : int = 0) -> None: close_video_reader(video_reader) - video_reader['process'] = ffmpeg.create_video_reader(video_reader.get('file_path'), frame_position, video_reader.get('metadata')) - video_reader['position'] = frame_position + video_reader['process'] = ffmpeg.create_video_reader(video_reader.get('file_path'), frame_number, video_reader.get('metadata')) + video_reader['frame_number'] = frame_number -#todo: needs review - [decoding] [critical: high] partial pipe read returns none and desyncs position from the actual stream -def read_video_reader_frame(video_reader : VideoReader) -> Optional[VisionFrame]: +def drain_video_reader(video_reader : VideoReader, skip_total : int) -> None: width, height = video_reader.get('metadata').get('resolution') - frame_size = width * height * 3 - frame_buffer = video_reader.get('process').stdout.read(frame_size) + channel_total = 3 + frame_size = width * height * channel_total + + for _ in range(skip_total): + video_reader.get('process').stdout.read(frame_size) + + video_reader['frame_number'] = video_reader.get('frame_number') + skip_total + + +def read_video_frame(video_reader : VideoReader) -> Optional[VisionFrame]: + width, height = video_reader.get('metadata').get('resolution') + channel_total = 3 + video_stream = cast(BufferedReader, video_reader.get('process').stdout) + vision_frame = numpy.empty(width * height * channel_total, numpy.uint8) + + if video_stream.readinto(vision_frame) == vision_frame.size: + video_reader['frame_number'] = video_reader.get('frame_number') + 1 + return vision_frame.reshape(height, width, channel_total) - if len(frame_buffer) == frame_size: - video_reader['position'] = video_reader.get('position') + 1 - return numpy.frombuffer(frame_buffer, numpy.uint8).reshape(height, width, 3) return None -#todo: needs review - [memory] [critical: high] frame_set keeps decoded frames in ram, eviction only trims below frame_start minus buffer_margin -def read_video_reader_window(video_reader : VideoReader, frame_start : int, frame_end : int) -> VisionFrameSet: - id = video_reader.get('id') - frame_set = frame_store.get_frame_store(id) - buffer_margin = 16 +def read_video_frames(video_reader : VideoReader, frame_start : int, frame_end : int) -> VisionFrameSet: + reader_id = video_reader.get('id') + frame_set = frame_store.get_frame_store(reader_id) keep_margin = 4 frame_gaps = [] - for frame_index in range(frame_start, frame_end + 1): - if frame_index not in frame_set: - frame_gaps.append(frame_index) + for frame_number in range(frame_start, frame_end + 1): + if frame_number not in frame_set: + frame_gaps.append(frame_number) if frame_gaps: - frame_position = get_first(frame_gaps) - skip_total = frame_position - video_reader.get('position') + collect_video_frames(video_reader, get_first(frame_gaps), get_last(frame_gaps)) - if skip_total < 0 or skip_total > buffer_margin: - refresh_video_reader(video_reader, frame_position) + frame_store.reduce_frames(reader_id, frame_start - keep_margin, frame_end + keep_margin) + return frame_store.select_frame_set(reader_id, frame_start, frame_end) - for decode_index in range(video_reader.get('position'), get_last(frame_gaps) + 1): - vision_frame = read_video_reader_frame(video_reader) - if numpy.any(vision_frame): - frame_store.set_frame(id, decode_index, vision_frame) +def collect_video_frames(video_reader : VideoReader, frame_start : int, frame_end : int) -> None: + reader_id = video_reader.get('id') + skip_total = frame_start - video_reader.get('frame_number') + skip_margin = 16 - frame_store.reduce_frames(id, frame_start - keep_margin, frame_end + keep_margin) - return frame_store.select_frame_set(id, frame_start, frame_end) + if skip_total < 0 or skip_total > skip_margin: + seek_video_reader(video_reader, frame_start) + + for frame_number in range(video_reader.get('frame_number'), frame_end + 1): + vision_frame = read_video_frame(video_reader) + + if numpy.any(vision_frame): + frame_store.set_frame(reader_id, frame_number, vision_frame) def close_video_reader(video_reader : VideoReader) -> None: @@ -107,14 +125,14 @@ def get_writer(video_path : str, temp_video_fps : Fps, temp_video_resolution : R 'metadata': { 'fps': output_video_fps, - 'resolution': output_video_resolution, + 'resolution': output_video_resolution } } return VIDEO_POOL_SET.get('writer').get(video_path) -def write_video_writer(video_writer : VideoWriter, vision_frame : VisionFrame) -> None: +def write_video_frame(video_writer : VideoWriter, vision_frame : VisionFrame) -> None: video_writer.get('process').stdin.write(vision_frame.tobytes()) diff --git a/facefusion/vision.py b/facefusion/vision.py index 7b01ccd8..384b967b 100644 --- a/facefusion/vision.py +++ b/facefusion/vision.py @@ -75,23 +75,17 @@ def read_static_video_frame(video_path : str, frame_number : int = 0) -> Optiona return read_video_frame(video_path, frame_number) -#todo: needs review - [decoding] [critical: medium] cv2 capture replaced by the pooled ffmpeg reader, position clamped to the frame_total estimate def read_video_frame(video_path : str, frame_number : int = 0) -> Optional[VisionFrame]: if is_video(video_path): - video_reader = video_manager.get_reader(video_path) + video_reader = video_manager.get_reader(video_path, 'read_video_frame') - if video_reader: - video_frame_position = min(video_reader.get('metadata').get('frame_total'), frame_number) - - with thread_semaphore(): - video_manager.conditional_set_video_reader_position(video_reader, video_frame_position) - return video_manager.read_video_reader_frame(video_reader) + with thread_semaphore(): + video_manager.conditional_seek_video_reader(video_reader, frame_number) + return video_manager.read_video_frame(video_reader) return None -#todo: needs review - [decoding] [critical: high] window read replaces the chunk cache, out of range frames fall back to empty vision frames -#todo: question to restore the chunk_size approach over the window read def select_video_frames(video_path : str, frame_number : int = 0, frame_offset : int = 2) -> List[VisionFrame]: vision_frames = [] frame_start = frame_number - frame_offset @@ -99,8 +93,8 @@ def select_video_frames(video_path : str, frame_number : int = 0, frame_offset : if is_video(video_path): with thread_lock(): - video_reader = video_manager.get_reader(video_path) - frame_set = video_manager.read_video_reader_window(video_reader, max(frame_start, 0), frame_end) + video_reader = video_manager.get_reader(video_path, 'select_video_frames') + frame_set = video_manager.read_video_frames(video_reader, max(frame_start, 0), frame_end) for frame_number in range(frame_start, frame_end + 1): vision_frame = create_empty_vision_frame() diff --git a/facefusion/workflows/image_to_video.py b/facefusion/workflows/image_to_video.py index 827ee097..b0ad862b 100644 --- a/facefusion/workflows/image_to_video.py +++ b/facefusion/workflows/image_to_video.py @@ -186,7 +186,7 @@ def write_stream_frame(video_writer : VideoWriter, futures : List[Future[Tuple[i if not future.cancelled(): _, temp_vision_frame = future.result() - video_manager.write_video_writer(video_writer, temp_vision_frame) + video_manager.write_video_frame(video_writer, temp_vision_frame) progress.update() diff --git a/tests/test_video_manager.py b/tests/test_video_manager.py index a75be714..980fcd84 100644 --- a/tests/test_video_manager.py +++ b/tests/test_video_manager.py @@ -9,7 +9,7 @@ from facefusion.download import conditional_download from facefusion.ffprobe import extract_video_metadata from facefusion.frame_store import get_frame_store from facefusion.temp_helper import create_temp_directory, get_temp_file_path -from facefusion.video_manager import clear_video_pool, close_video_writer, conditional_set_video_reader_position, get_reader, get_writer, read_video_reader_frame, read_video_reader_window, refresh_video_reader, write_video_writer +from facefusion.video_manager import clear_video_pool, close_video_reader, close_video_writer, collect_video_frames, conditional_seek_video_reader, drain_video_reader, get_reader, get_writer, read_video_frame, read_video_frames, seek_video_reader, write_video_frame from .helper import get_test_example_file, get_test_examples_directory @@ -43,95 +43,108 @@ def before_each() -> None: clear_video_pool() -#todo: needs review - [testing] question if the assertions are good -#todo: run mutation testing, strip down to the minimum, test with real data def test_get_reader() -> None: - video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4')) + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') video_metadata = video_reader.get('metadata') assert video_metadata.get('resolution') == (426, 226) assert video_metadata.get('fps') == 25.0 assert video_metadata.get('frame_total') == 270 - assert video_reader.get('position') == 0 - assert get_reader(get_test_example_file('target-240p-25fps.mp4')) is video_reader + assert get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') is video_reader + assert not get_reader(get_test_example_file('target-240p-25fps.mp4'), 'select_video_frames').get('id') == video_reader.get('id') -#todo: needs review - [testing] question if the assertions are good -#todo: run mutation testing, strip down to the minimum, test with real data -def test_conditional_set_video_reader_position() -> None: - video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4')) - - conditional_set_video_reader_position(video_reader, 50) - - assert video_reader.get('position') == 50 - - conditional_set_video_reader_position(video_reader, 10) - - assert video_reader.get('position') == 10 - - conditional_set_video_reader_position(video_reader, 200) - - assert video_reader.get('position') == 200 - - -#todo: needs review - [testing] question if the assertions are good -#todo: run mutation testing, strip down to the minimum, test with real data -def test_refresh_video_reader() -> None: - video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4')) +def test_conditional_seek_video_reader() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') sequential_frames = {} for frame_number in range(30): - sequential_frames[frame_number] = read_video_reader_frame(video_reader) + sequential_frames[frame_number] = read_video_frame(video_reader) for frame_number in [ 5, 17, 29 ]: - refresh_video_reader(video_reader, frame_number) - vision_frame = read_video_reader_frame(video_reader) + conditional_seek_video_reader(video_reader, frame_number) - assert numpy.array_equal(vision_frame, sequential_frames.get(frame_number)) is True + assert numpy.array_equal(read_video_frame(video_reader), sequential_frames.get(frame_number)) is True -#todo: needs review - [testing] question if the assertions are good -#todo: run mutation testing, strip down to the minimum, test with real data -def test_read_video_reader_frame() -> None: - video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4')) - vision_frame = read_video_reader_frame(video_reader) +def test_seek_video_reader() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') + sequential_frames = {} - assert vision_frame.shape == (226, 426, 3) - assert video_reader.get('position') == 1 + for frame_number in range(30): + sequential_frames[frame_number] = read_video_frame(video_reader) - conditional_set_video_reader_position(video_reader, 269) - vision_frame = read_video_reader_frame(video_reader) + for frame_number in [ 5, 17, 29 ]: + seek_video_reader(video_reader, frame_number) - assert vision_frame.shape == (226, 426, 3) - assert read_video_reader_frame(video_reader) is None + assert numpy.array_equal(read_video_frame(video_reader), sequential_frames.get(frame_number)) is True -#todo: needs review - [testing] question if the assertions are good -#todo: run mutation testing, strip down to the minimum, test with real data -def test_read_video_reader_window() -> None: - video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4')) - frame_set = read_video_reader_window(video_reader, 0, 4) +def test_drain_video_reader() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') - assert sorted(frame_set) == [ 0, 1, 2, 3, 4 ] + drain_video_reader(video_reader, 10) - position = video_reader.get('position') - frame_set = read_video_reader_window(video_reader, 1, 3) + assert video_reader.get('frame_number') == 10 - assert video_reader.get('position') == position - assert sorted(frame_set) == [ 1, 2, 3 ] + vision_frame = read_video_frame(video_reader) + seek_video_reader(video_reader, 10) - read_video_reader_window(video_reader, 21, 25) + assert numpy.array_equal(vision_frame, read_video_frame(video_reader)) is True + + +def test_read_video_frame() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') + + assert read_video_frame(video_reader).shape == (226, 426, 3) + assert video_reader.get('frame_number') == 1 + + seek_video_reader(video_reader, 269) + + assert read_video_frame(video_reader).shape == (226, 426, 3) + assert read_video_frame(video_reader) is None + + +def test_read_video_frames() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'read_video_frame') + + assert sorted(read_video_frames(video_reader, 0, 4)) == [ 0, 1, 2, 3, 4 ] + + frame_number = video_reader.get('frame_number') + + assert sorted(read_video_frames(video_reader, 1, 3)) == [ 1, 2, 3 ] + assert video_reader.get('frame_number') == frame_number + + read_video_frames(video_reader, 21, 25) assert min(get_frame_store(video_reader.get('id'))) == 17 assert max(get_frame_store(video_reader.get('id'))) == 25 - - frame_set = read_video_reader_window(video_reader, 268, 275) - - assert sorted(frame_set) == [ 268, 269 ] + assert sorted(read_video_frames(video_reader, 268, 275)) == [ 268, 269 ] + + +def test_collect_video_frames() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'select_video_frames') + + collect_video_frames(video_reader, 20, 24) + + assert sorted(get_frame_store(video_reader.get('id'))) == [ 20, 21, 22, 23, 24 ] + assert video_reader.get('frame_number') == 25 + + +def test_close_video_reader() -> None: + video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4'), 'select_video_frames') + read_video_frames(video_reader, 0, 4) + close_video_reader(video_reader) + + if is_windows(): + assert video_reader.get('process').returncode == 1 + + if is_linux() or is_macos(): + assert video_reader.get('process').returncode == -9 + + assert get_frame_store(video_reader.get('id')) == {} -#todo: needs review - [testing] question if the assertions are good -#todo: run mutation testing, strip down to the minimum, test with real data def test_get_writer() -> None: target_path = get_test_example_file('target-240p-25fps.mp4') create_temp_directory(target_path) @@ -140,17 +153,14 @@ def test_get_writer() -> None: assert get_writer(target_path, 25.0, (426, 226), (426, 226), 25.0) is video_writer -#todo: needs review - [testing] question if the assertions are good -#todo: run mutation testing, strip down to the minimum, test with real data -def test_write_video_writer() -> None: +def test_write_video_frame() -> None: target_path = get_test_example_file('target-240p-25fps.mp4') create_temp_directory(target_path) - video_reader = get_reader(target_path) + video_reader = get_reader(target_path, 'read_video_frame') video_writer = get_writer(target_path, 25.0, (426, 226), (426, 226), 25.0) for frame_number in range(25): - vision_frame = read_video_reader_frame(video_reader) - write_video_writer(video_writer, vision_frame) + write_video_frame(video_writer, read_video_frame(video_reader)) assert close_video_writer(video_writer) is True @@ -163,28 +173,22 @@ def test_write_video_writer() -> None: assert video_metadata.get('color_transfer') == 'bt709' -#todo: needs review - [testing] question if the assertions are good -#todo: run mutation testing, strip down to the minimum, test with real data def test_close_video_writer() -> None: target_path = get_test_example_file('target-240p-30fps.mp4') create_temp_directory(target_path) - video_reader = get_reader(target_path) + video_reader = get_reader(target_path, 'read_video_frame') video_writer = get_writer(target_path, 30.0, (426, 226), (426, 226), 30.0) - vision_frame = read_video_reader_frame(video_reader) - write_video_writer(video_writer, vision_frame) + write_video_frame(video_writer, read_video_frame(video_reader)) assert close_video_writer(video_writer) is True -#todo: needs review - [testing] question if the assertions are good -#todo: run mutation testing, strip down to the minimum, test with real data def test_clear_video_pool() -> None: target_path = get_test_example_file('target-240p-25fps.mp4') create_temp_directory(target_path) - video_reader = get_reader(target_path) + video_reader = get_reader(target_path, 'read_video_frame') video_writer = get_writer(target_path, 25.0, (426, 226), (426, 226), 25.0) - vision_frame = read_video_reader_frame(video_reader) - write_video_writer(video_writer, vision_frame) + write_video_frame(video_writer, read_video_frame(video_reader)) clear_video_pool() if is_windows():