diff --git a/facefusion/frame_store.py b/facefusion/frame_store.py index 6ce448b6..843ce21f 100644 --- a/facefusion/frame_store.py +++ b/facefusion/frame_store.py @@ -30,5 +30,6 @@ def reduce_frames(id : str, frame_min : int, frame_max : int) -> None: FRAME_STORE_SET[id] = select_frame_set(id, frame_min, frame_max) -def clear_frames() -> None: - FRAME_STORE_SET.clear() +def clear_frames(id : str) -> None: + if id in FRAME_STORE_SET: + del FRAME_STORE_SET[id] diff --git a/facefusion/types.py b/facefusion/types.py index 5e946f58..7407a821 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -127,15 +127,16 @@ VideoMetadata = TypedDict('VideoMetadata', #todo: question if the body of VideoReader and VideoWriter needs all the keys VideoReader = TypedDict('VideoReader', { + 'id' : str, 'process' : subprocess.Popen[bytes], 'file_path' : str, 'metadata' : VideoMetadata, - 'position' : int, - 'frame_set' : VisionFrameSet + 'position' : int }) VideoReaderSet : TypeAlias = Dict[str, VideoReader] VideoWriter = TypedDict('VideoWriter', { + 'id' : str, 'process' : subprocess.Popen[bytes], 'file_path' : str, 'metadata' : VideoMetadata diff --git a/facefusion/video_manager.py b/facefusion/video_manager.py index e202c201..fcfbbae4 100644 --- a/facefusion/video_manager.py +++ b/facefusion/video_manager.py @@ -1,8 +1,10 @@ +import uuid from typing import Optional import numpy -from facefusion import ffmpeg, ffprobe +from facefusion import ffmpeg, ffprobe, frame_store +from facefusion.common_helper import get_first, get_last from facefusion.types import Fps, Resolution, VideoMetadata, VideoPoolSet, VideoReader, VideoWriter, VisionFrame, VisionFrameSet VIDEO_POOL_SET : VideoPoolSet =\ @@ -19,11 +21,11 @@ def get_reader(video_path : str) -> VideoReader: VIDEO_POOL_SET['reader'][video_path] =\ { + 'id': uuid.uuid4().hex, 'process': ffmpeg.create_video_reader(video_path, 0, video_metadata), 'file_path': video_path, 'metadata': video_metadata, - 'position': 0, - 'frame_set': {} + 'position': 0 } return VIDEO_POOL_SET.get('reader').get(video_path) @@ -48,7 +50,6 @@ def refresh_video_reader(video_reader : VideoReader, frame_position : int) -> No video_reader.get('process').wait() 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['frame_set'].clear() #todo: needs review - [decoding] [critical: high] partial pipe read returns none and desyncs position from the actual stream @@ -65,23 +66,30 @@ def read_video_reader_frame(video_reader : VideoReader) -> Optional[VisionFrame] #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: - frame_set = video_reader.get('frame_set') + id = video_reader.get('id') + frame_set = frame_store.get_frame_store(id) buffer_margin = 16 + frame_gaps = [] - if frame_start not in frame_set and (frame_start < video_reader.get('position') or frame_start > video_reader.get('position') + buffer_margin): - refresh_video_reader(video_reader, frame_start) + 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(video_reader.get('position'), frame_end + 1): - vision_frame = read_video_reader_frame(video_reader) + if frame_gaps: + frame_position = get_first(frame_gaps) + skip_total = frame_position - video_reader.get('position') - if numpy.any(vision_frame): - frame_set[frame_number] = vision_frame + if skip_total < 0 or skip_total > buffer_margin: + refresh_video_reader(video_reader, frame_position) - for frame_number in list(frame_set): - if frame_number < frame_start - buffer_margin: - del frame_set[frame_number] + for decode_index in range(video_reader.get('position'), get_last(frame_gaps) + 1): + vision_frame = read_video_reader_frame(video_reader) - return frame_set + if numpy.any(vision_frame): + frame_store.set_frame(id, decode_index, vision_frame) + + frame_store.reduce_frames(id, frame_start - buffer_margin, frame_end + buffer_margin) + return frame_store.select_frame_set(id, frame_start, frame_end) #todo: needs review - [lifecycle] [critical: low] pooled writer keyed by target_path, metadata forwarded by the caller @@ -89,6 +97,7 @@ def get_writer(target_path : str, video_metadata : VideoMetadata, temp_video_fps if target_path not in VIDEO_POOL_SET.get('writer'): VIDEO_POOL_SET['writer'][target_path] =\ { + 'id': uuid.uuid4().hex, 'process': ffmpeg.create_video_writer(target_path, temp_video_fps, temp_video_resolution, output_video_resolution, output_video_fps), 'file_path': target_path, 'metadata': video_metadata @@ -114,6 +123,7 @@ def clear_video_pool() -> None: for video_reader in VIDEO_POOL_SET.get('reader').values(): video_reader.get('process').kill() video_reader.get('process').wait() + frame_store.clear_frames(video_reader.get('id')) for video_writer in VIDEO_POOL_SET.get('writer').values(): video_writer.get('process').kill() diff --git a/tests/test_frame_store.py b/tests/test_frame_store.py index 057a121f..84c87929 100644 --- a/tests/test_frame_store.py +++ b/tests/test_frame_store.py @@ -18,7 +18,8 @@ def before_all() -> None: @pytest.fixture(scope = 'function', autouse = True) def before_each() -> None: - clear_frames() + clear_frames('reader-1') + clear_frames('reader-2') def test_get_frame_store() -> None: @@ -65,6 +66,8 @@ def test_clear_frames() -> None: target_frame = read_video_frame(get_test_example_file('target-240p.mp4'), 0) set_frame('reader-1', 0, target_frame) - clear_frames() + set_frame('reader-2', 0, target_frame) + clear_frames('reader-1') assert get_frame_store('reader-1') == {} + assert get_frame_store('reader-2').get(0) is target_frame diff --git a/tests/test_video_manager.py b/tests/test_video_manager.py index 8347d611..d6e4e55c 100644 --- a/tests/test_video_manager.py +++ b/tests/test_video_manager.py @@ -6,6 +6,7 @@ import pytest from facefusion import ffmpeg, ffmpeg_builder, process_manager, state_manager 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_frame from .helper import get_test_example_file, get_test_examples_directory @@ -112,26 +113,22 @@ def test_read_video_reader_window() -> None: assert sorted(frame_set) == [ 0, 1, 2, 3, 4 ] - frame_set = read_video_reader_window(video_reader, 100, 104) + position = video_reader.get('position') + frame_set = read_video_reader_window(video_reader, 1, 3) - assert sorted(frame_set) == [ 100, 101, 102, 103, 104 ] + assert video_reader.get('position') == position + assert sorted(frame_set) == [ 1, 2, 3 ] + + read_video_reader_window(video_reader, 21, 25) + + assert min(get_frame_store(video_reader.get('id'))) == 5 + 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 ] -#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_evict_video_reader_buffer() -> None: - video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4')) - read_video_reader_window(video_reader, 0, 4) - frame_set = read_video_reader_window(video_reader, 21, 25) - - assert min(frame_set) == 5 - assert max(frame_set) == 25 - - #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: