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..cca2a196 100644 --- a/facefusion/video_manager.py +++ b/facefusion/video_manager.py @@ -24,71 +24,77 @@ def get_reader(video_path : str) -> VideoReader: '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) #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 seek_video_reader(video_reader : VideoReader, frame_number : int) -> 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) + read_video_frame(video_reader) - if not video_reader.get('position') == frame_position: - refresh_video_reader(video_reader, frame_position) + if not video_reader.get('frame_number') == frame_number: + refresh_video_reader(video_reader, frame_number) -def refresh_video_reader(video_reader : VideoReader, frame_position : int) -> None: +def refresh_video_reader(video_reader : VideoReader, frame_number : int) -> 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 read_video_frame(video_reader : VideoReader) -> Optional[VisionFrame]: width, height = video_reader.get('metadata').get('resolution') - frame_size = width * height * 3 + channel_total = 3 + frame_size = width * height * channel_total frame_buffer = video_reader.get('process').stdout.read(frame_size) 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) + video_reader['frame_number'] = video_reader.get('frame_number') + 1 + return numpy.frombuffer(frame_buffer, numpy.uint8).reshape(height, width, channel_total) + 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') + decode_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 decode_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: + refresh_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 +113,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..d0561375 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) - 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.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 @@ -100,7 +94,7 @@ 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) + 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..ad6b0392 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_writer, get_reader, get_writer, read_video_frame, read_video_frames, refresh_video_reader, seek_video_reader, write_video_frame from .helper import get_test_example_file, get_test_examples_directory @@ -52,26 +52,26 @@ def test_get_reader() -> None: 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 video_reader.get('frame_number') == 0 assert get_reader(get_test_example_file('target-240p-25fps.mp4')) is video_reader #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: +def test_seek_video_reader() -> None: video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4')) - conditional_set_video_reader_position(video_reader, 50) + seek_video_reader(video_reader, 50) - assert video_reader.get('position') == 50 + assert video_reader.get('frame_number') == 50 - conditional_set_video_reader_position(video_reader, 10) + seek_video_reader(video_reader, 10) - assert video_reader.get('position') == 10 + assert video_reader.get('frame_number') == 10 - conditional_set_video_reader_position(video_reader, 200) + seek_video_reader(video_reader, 200) - assert video_reader.get('position') == 200 + assert video_reader.get('frame_number') == 200 #todo: needs review - [testing] question if the assertions are good @@ -81,51 +81,51 @@ def test_refresh_video_reader() -> None: 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) + vision_frame = read_video_frame(video_reader) assert numpy.array_equal(vision_frame, 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: +def test_read_video_frame() -> None: video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4')) - vision_frame = read_video_reader_frame(video_reader) + vision_frame = read_video_frame(video_reader) assert vision_frame.shape == (226, 426, 3) - assert video_reader.get('position') == 1 + assert video_reader.get('frame_number') == 1 - conditional_set_video_reader_position(video_reader, 269) - vision_frame = read_video_reader_frame(video_reader) + seek_video_reader(video_reader, 269) + vision_frame = read_video_frame(video_reader) assert vision_frame.shape == (226, 426, 3) - assert read_video_reader_frame(video_reader) is None + assert read_video_frame(video_reader) is 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_read_video_reader_window() -> None: +def test_read_video_frames() -> None: video_reader = get_reader(get_test_example_file('target-240p-25fps.mp4')) - frame_set = read_video_reader_window(video_reader, 0, 4) + frame_set = read_video_frames(video_reader, 0, 4) assert sorted(frame_set) == [ 0, 1, 2, 3, 4 ] - position = video_reader.get('position') - frame_set = read_video_reader_window(video_reader, 1, 3) + frame_number = video_reader.get('frame_number') + frame_set = read_video_frames(video_reader, 1, 3) - assert video_reader.get('position') == position + assert video_reader.get('frame_number') == frame_number assert sorted(frame_set) == [ 1, 2, 3 ] - read_video_reader_window(video_reader, 21, 25) + 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) + frame_set = read_video_frames(video_reader, 268, 275) assert sorted(frame_set) == [ 268, 269 ] @@ -142,15 +142,15 @@ def test_get_writer() -> 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_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_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) + vision_frame = read_video_frame(video_reader) + write_video_frame(video_writer, vision_frame) assert close_video_writer(video_writer) is True @@ -170,8 +170,8 @@ def test_close_video_writer() -> None: create_temp_directory(target_path) video_reader = get_reader(target_path) 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) + vision_frame = read_video_frame(video_reader) + write_video_frame(video_writer, vision_frame) assert close_video_writer(video_writer) is True @@ -183,8 +183,8 @@ def test_clear_video_pool() -> None: create_temp_directory(target_path) video_reader = get_reader(target_path) 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) + vision_frame = read_video_frame(video_reader) + write_video_frame(video_writer, vision_frame) clear_video_pool() if is_windows():