diff --git a/facefusion/video_manager.py b/facefusion/video_manager.py index cca2a196..6a41b05e 100644 --- a/facefusion/video_manager.py +++ b/facefusion/video_manager.py @@ -30,8 +30,7 @@ def get_reader(video_path : str) -> VideoReader: return VIDEO_POOL_SET.get('reader').get(video_path) -#todo: needs review - [seeking] [critical: high] forward skip up to 128 frames by draining the pipe, everything else refreshes the process -def seek_video_reader(video_reader : VideoReader, frame_number : int) -> None: +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 @@ -41,10 +40,10 @@ def seek_video_reader(video_reader : VideoReader, frame_number : int) -> None: read_video_frame(video_reader) if not video_reader.get('frame_number') == frame_number: - refresh_video_reader(video_reader, frame_number) + seek_video_reader(video_reader, frame_number) -def refresh_video_reader(video_reader : VideoReader, frame_number : 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_number, video_reader.get('metadata')) @@ -88,7 +87,7 @@ def decode_video_frames(video_reader : VideoReader, frame_start : int, frame_end skip_margin = 16 if skip_total < 0 or skip_total > skip_margin: - refresh_video_reader(video_reader, frame_start) + 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) diff --git a/facefusion/vision.py b/facefusion/vision.py index d0561375..8cdb134c 100644 --- a/facefusion/vision.py +++ b/facefusion/vision.py @@ -80,7 +80,7 @@ def read_video_frame(video_path : str, frame_number : int = 0) -> Optional[Visio video_reader = video_manager.get_reader(video_path) with thread_semaphore(): - video_manager.seek_video_reader(video_reader, frame_number) + video_manager.conditional_seek_video_reader(video_reader, frame_number) return video_manager.read_video_frame(video_reader) return None diff --git a/tests/test_video_manager.py b/tests/test_video_manager.py index ad6b0392..e133f422 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, get_reader, get_writer, read_video_frame, read_video_frames, refresh_video_reader, seek_video_reader, write_video_frame +from facefusion.video_manager import clear_video_pool, close_video_writer, conditional_seek_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 @@ -58,25 +58,25 @@ def test_get_reader() -> None: #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_seek_video_reader() -> None: +def test_conditional_seek_video_reader() -> None: video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4')) - seek_video_reader(video_reader, 50) + conditional_seek_video_reader(video_reader, 50) assert video_reader.get('frame_number') == 50 - seek_video_reader(video_reader, 10) + conditional_seek_video_reader(video_reader, 10) assert video_reader.get('frame_number') == 10 - seek_video_reader(video_reader, 200) + conditional_seek_video_reader(video_reader, 200) assert video_reader.get('frame_number') == 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: +def test_seek_video_reader() -> None: video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4')) sequential_frames = {} @@ -84,7 +84,7 @@ def test_refresh_video_reader() -> None: sequential_frames[frame_number] = read_video_frame(video_reader) for frame_number in [ 5, 17, 29 ]: - refresh_video_reader(video_reader, frame_number) + seek_video_reader(video_reader, frame_number) vision_frame = read_video_frame(video_reader) assert numpy.array_equal(vision_frame, sequential_frames.get(frame_number)) is True @@ -99,7 +99,7 @@ def test_read_video_frame() -> None: assert vision_frame.shape == (226, 426, 3) assert video_reader.get('frame_number') == 1 - seek_video_reader(video_reader, 269) + conditional_seek_video_reader(video_reader, 269) vision_frame = read_video_frame(video_reader) assert vision_frame.shape == (226, 426, 3)