Fix out of range for trim frame, Fix ffmpeg extraction count

facefusion/content_analyser.py

@@ -9,7 +9,7 @@ from facefusion.download import conditional_download_hashes, conditional_downloa
 from facefusion.filesystem import resolve_relative_path
 from facefusion.thread_helper import conditional_thread_semaphore
 from facefusion.typing import DownloadScope, Fps, InferencePool, ModelOptions, ModelSet, VisionFrame
-from facefusion.vision import count_video_frame_total, detect_video_fps, get_video_frame, read_image
+from facefusion.vision import detect_video_fps, get_video_frame, read_image, restrict_trim_frame
 
 PROBABILITY_LIMIT = 0.80
 RATE_LIMIT = 10
@@ -108,10 +108,10 @@ def analyse_image(image_path : str) -> bool:
 
 
 @lru_cache(maxsize = None)
-def analyse_video(video_path : str, start_frame : int, end_frame : int) -> bool:
-	video_frame_total = count_video_frame_total(video_path)
+def analyse_video(video_path : str, trim_frame_start : int, trim_frame_end : int) -> bool:
 	video_fps = detect_video_fps(video_path)
-	frame_range = range(start_frame or 0, end_frame or video_frame_total)
+	trim_frame_start, trim_frame_end = restrict_trim_frame(video_path, trim_frame_start, trim_frame_end)
+	frame_range = range(trim_frame_start, trim_frame_end)
 	rate = 0.0
 	counter = 0
 

facefusion/ffmpeg.py

@@ -11,7 +11,7 @@ from facefusion import logger, process_manager, state_manager, wording
 from facefusion.filesystem import remove_file
 from facefusion.temp_helper import get_temp_file_path, get_temp_frame_paths, get_temp_frames_pattern
 from facefusion.typing import AudioBuffer, Fps, OutputVideoPreset, UpdateProgress
-from facefusion.vision import count_video_frame_total, detect_video_duration, restrict_video_fps
+from facefusion.vision import count_trim_frame_total, detect_video_duration, restrict_video_fps
 
 
 def run_ffmpeg_with_progress(args: List[str], update_progress : UpdateProgress) -> subprocess.Popen[bytes]:
@@ -74,21 +74,18 @@ def log_debug(process : subprocess.Popen[bytes]) -> None:
 
 
 def extract_frames(target_path : str, temp_video_resolution : str, temp_video_fps : Fps) -> bool:
-	extract_frame_total = count_video_frame_total(state_manager.get_item('target_path'))
 	trim_frame_start = state_manager.get_item('trim_frame_start')
 	trim_frame_end = state_manager.get_item('trim_frame_end')
+	extract_frame_total = count_trim_frame_total(state_manager.get_item('target_path'), trim_frame_start, trim_frame_end)
 	temp_frames_pattern = get_temp_frames_pattern(target_path, '%08d')
 	commands = [ '-i', target_path, '-s', str(temp_video_resolution), '-q:v', '0' ]
 
 	if isinstance(trim_frame_start, int) and isinstance(trim_frame_end, int):
 		commands.extend([ '-vf', 'trim=start_frame=' + str(trim_frame_start) + ':end_frame=' + str(trim_frame_end) + ',fps=' + str(temp_video_fps) ])
-		extract_frame_total = trim_frame_end - trim_frame_start
 	elif isinstance(trim_frame_start, int):
 		commands.extend([ '-vf', 'trim=start_frame=' + str(trim_frame_start) + ',fps=' + str(temp_video_fps) ])
-		extract_frame_total -= trim_frame_start
 	elif isinstance(trim_frame_end, int):
 		commands.extend([ '-vf', 'trim=end_frame=' + str(trim_frame_end) + ',fps=' + str(temp_video_fps) ])
-		extract_frame_total -= trim_frame_end
 	else:
 		commands.extend([ '-vf', 'fps=' + str(temp_video_fps) ])
 	commands.extend([ '-vsync', '0', temp_frames_pattern ])
@@ -99,10 +96,10 @@ def extract_frames(target_path : str, temp_video_resolution : str, temp_video_fp
 
 
 def merge_video(target_path : str, output_video_resolution : str, output_video_fps: Fps) -> bool:
-	merge_frame_total = len(get_temp_frame_paths(target_path))
 	output_video_encoder = state_manager.get_item('output_video_encoder')
 	output_video_quality = state_manager.get_item('output_video_quality')
 	output_video_preset = state_manager.get_item('output_video_preset')
+	merge_frame_total = len(get_temp_frame_paths(target_path))
 	temp_video_fps = restrict_video_fps(target_path, output_video_fps)
 	temp_file_path = get_temp_file_path(target_path)
 	temp_frames_pattern = get_temp_frames_pattern(target_path, '%08d')

facefusion/vision.py

@@ -122,11 +122,36 @@ def restrict_video_fps(video_path : str, fps : Fps) -> Fps:
 def detect_video_duration(video_path : str) -> Duration:
 	video_frame_total = count_video_frame_total(video_path)
 	video_fps = detect_video_fps(video_path)
+
 	if video_frame_total and video_fps:
 		return video_frame_total / video_fps
 	return 0
 
 
+def count_trim_frame_total(video_path : str, trim_frame_start : int, trim_frame_end : int) -> int:
+	trim_frame_start, trim_frame_end = restrict_trim_frame(video_path, trim_frame_start, trim_frame_end)
+
+	return trim_frame_end - trim_frame_start
+
+
+def restrict_trim_frame(video_path : str, trim_frame_start : int, trim_frame_end : int) -> Tuple[int, int]:
+	video_frame_total = count_video_frame_total(video_path)
+
+	if isinstance(trim_frame_start, int) and trim_frame_start < 0:
+		trim_frame_start = 0
+	if isinstance(trim_frame_end, int) and trim_frame_end > video_frame_total:
+		trim_frame_end = video_frame_total
+
+	if isinstance(trim_frame_start, int) and isinstance(trim_frame_end, int):
+		return trim_frame_start, trim_frame_end
+	if isinstance(trim_frame_start, int):
+		return trim_frame_start, video_frame_total
+	if isinstance(trim_frame_end, int):
+		return 0, trim_frame_end
+
+	return 0, video_frame_total
+
+
 def detect_video_resolution(video_path : str) -> Optional[Resolution]:
 	if is_video(video_path):
 		if is_windows():

tests/test_audio.py

@@ -17,8 +17,8 @@ def before_all() -> None:
 
 
 def test_get_audio_frame() -> None:
-	assert get_audio_frame(get_test_example_file('source.mp3'), 25) is not None
-	assert get_audio_frame(get_test_example_file('source.wav'), 25) is not None
+	assert hasattr(get_audio_frame(get_test_example_file('source.mp3'), 25), '__array_interface__')
+	assert hasattr(get_audio_frame(get_test_example_file('source.wav'), 25), '__array_interface__')
 	assert get_audio_frame('invalid', 25) is None
 
 

tests/test_vision.py

@@ -3,7 +3,7 @@ import subprocess
 import pytest
 
 from facefusion.download import conditional_download
-from facefusion.vision import calc_histogram_difference, count_video_frame_total, create_image_resolutions, create_video_resolutions, detect_image_resolution, detect_video_duration, detect_video_fps, detect_video_resolution, get_video_frame, match_frame_color, normalize_resolution, pack_resolution, read_image, restrict_image_resolution, restrict_video_fps, restrict_video_resolution, unpack_resolution
+from facefusion.vision import calc_histogram_difference, count_trim_frame_total, count_video_frame_total, create_image_resolutions, create_video_resolutions, detect_image_resolution, detect_video_duration, detect_video_fps, detect_video_resolution, get_video_frame, match_frame_color, normalize_resolution, pack_resolution, read_image, restrict_image_resolution, restrict_trim_frame, restrict_video_fps, restrict_video_resolution, unpack_resolution
 from .helper import get_test_example_file, get_test_examples_directory
 
 
@@ -50,7 +50,7 @@ def test_create_image_resolutions() -> None:
 
 
 def test_get_video_frame() -> None:
-	assert get_video_frame(get_test_example_file('target-240p-25fps.mp4')) is not None
+	assert hasattr(get_video_frame(get_test_example_file('target-240p-25fps.mp4')), '__array_interface__')
 	assert get_video_frame('invalid') is None
 
 
@@ -79,6 +79,22 @@ def test_detect_video_duration() -> None:
 	assert detect_video_duration('invalid') == 0
 
 
+def test_count_trim_frame_total() -> None:
+	assert count_trim_frame_total(get_test_example_file('target-240p.mp4'), 0, 200) == 200
+	assert count_trim_frame_total(get_test_example_file('target-240p.mp4'), 70, 270) == 200
+	assert count_trim_frame_total(get_test_example_file('target-240p.mp4'), -10, None) == 270
+	assert count_trim_frame_total(get_test_example_file('target-240p.mp4'), None, 280) == 270
+	assert count_trim_frame_total(get_test_example_file('target-240p.mp4'), None, None) == 270
+
+
+def test_restrict_trim_frame() -> None:
+	assert restrict_trim_frame(get_test_example_file('target-240p.mp4'), 0, 200) == (0, 200)
+	assert restrict_trim_frame(get_test_example_file('target-240p.mp4'), 70, 270) == (70, 270)
+	assert restrict_trim_frame(get_test_example_file('target-240p.mp4'), -10, None) == (0, 270)
+	assert restrict_trim_frame(get_test_example_file('target-240p.mp4'), None, 280) == (0, 270)
+	assert restrict_trim_frame(get_test_example_file('target-240p.mp4'), None, None) == (0, 270)
+
+
 def test_detect_video_resolution() -> None:
 	assert detect_video_resolution(get_test_example_file('target-240p.mp4')) == (426, 226)
 	assert detect_video_resolution(get_test_example_file('target-240p-90deg.mp4')) == (226, 426)