route window read through frame_store (#1196)

* route window read through frame_store

* update proper id

* restore todos

* restore todos

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]

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

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()

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

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: