Finish tests for extract_frames()

facefusion/utilities.py

@@ -34,30 +34,30 @@ def run_ffmpeg(args : List[str]) -> bool:
 		return False
 
 
-def detect_fps(target_path : str) -> float:
+def detect_fps(target_path : str) -> Optional[float]:
 	commands = [ 'ffprobe', '-v', 'error', '-select_streams', 'v:0', '-show_entries', 'stream=r_frame_rate', '-of', 'default=noprint_wrappers = 1:nokey = 1', target_path ]
 	output = subprocess.check_output(commands).decode().strip().split('/')
 	try:
 		numerator, denominator = map(int, output)
 		return numerator / denominator
 	except (ValueError, ZeroDivisionError):
-		return 30
+		return None
 
 
-def extract_frames(target_path : str, fps : float = 30) -> bool:
+def extract_frames(target_path : str, fps : float = 30.0) -> bool:
 	temp_directory_path = get_temp_directory_path(target_path)
 	temp_frame_quality = round(31 - (facefusion.globals.temp_frame_quality * 0.31))
 	trim_frame_start = facefusion.globals.trim_frame_start
 	trim_frame_end = facefusion.globals.trim_frame_end
-	commands = [ '-hwaccel', 'auto', '-i', target_path, '-q:v', str(temp_frame_quality), '-pix_fmt', 'rgb24' ]
+	commands = [ '-hwaccel', 'auto', '-i', target_path, '-q:v', str(temp_frame_quality), '-pix_fmt', 'rgb24', ]
 	if trim_frame_start is not None and trim_frame_end is not None:
-		commands.extend(['-vf', 'trim=start_frame=' + str(trim_frame_start) + ':end_frame=' + str(trim_frame_end) + ',fps=' + str(fps)])
+		commands.extend([ '-vf', 'trim=start_frame=' + str(trim_frame_start) + ':end_frame=' + str(trim_frame_end) + ',fps=' + str(fps) ])
 	elif trim_frame_start is not None:
-		commands.extend(['-vf', 'trim=start_frame=' + str(trim_frame_start) + ',fps=' + str(fps)])
+		commands.extend([ '-vf', 'trim=start_frame=' + str(trim_frame_start) + ',fps=' + str(fps) ])
 	elif trim_frame_end is not None:
-		commands.extend(['-vf', 'trim=end_frame=' + str(trim_frame_end) + ',fps=' + str(fps)])
+		commands.extend([ '-vf', 'trim=end_frame=' + str(trim_frame_end) + ',fps=' + str(fps) ])
 	else:
-		commands.extend(['-vf', 'fps=' + str(fps)])
+		commands.extend([ '-vf', 'fps=' + str(fps) ])
 	commands.extend([os.path.join(temp_directory_path, '%04d.' + facefusion.globals.temp_frame_format)])
 	return run_ffmpeg(commands)
 
@@ -68,7 +68,7 @@ def create_video(target_path : str, fps : float = 30) -> bool:
 	output_video_quality = round(51 - (facefusion.globals.output_video_quality * 0.5))
 	commands = [ '-hwaccel', 'auto', '-r', str(fps), '-i', os.path.join(temp_directory_path, '%04d.' + facefusion.globals.temp_frame_format), '-c:v', facefusion.globals.output_video_encoder ]
 	if facefusion.globals.output_video_encoder in [ 'libx264', 'libx265', 'libvpx' ]:
-		commands.extend(['-crf', str(output_video_quality)])
+		commands.extend([ '-crf', str(output_video_quality) ])
 	if facefusion.globals.output_video_encoder in [ 'h264_nvenc', 'hevc_nvenc' ]:
 		commands.extend([ '-cq', str(output_video_quality) ])
 	commands.extend([ '-pix_fmt', 'yuv420p', '-vf', 'colorspace=bt709:iall=bt601-6-625', '-y', temp_output_path ])

tests/test_utilities.py

@@ -50,3 +50,58 @@ def test_extract_frames() -> None:
 		assert len(glob.glob1(temp_directory_path, '*.jpg')) == 324
 
 		clear_temp(target_path)
+
+
+def test_extract_frames_with_trim_start() -> None:
+	facefusion.globals.trim_frame_start = 224
+	data_provider =\
+	[
+		('.assets/examples/target-240p-25fps.mp4', 55),
+		('.assets/examples/target-240p-30fps.mp4', 100),
+		('.assets/examples/target-240p-60fps.mp4', 212)
+	]
+	for target_path, frame_total in data_provider:
+		temp_directory_path = get_temp_directory_path(target_path)
+		create_temp(target_path)
+
+		assert extract_frames(target_path, 30) is True
+		assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total
+
+		clear_temp(target_path)
+
+
+def test_extract_frames_with_trim_start_and_trim_end() -> None:
+	facefusion.globals.trim_frame_start = 224
+	facefusion.globals.trim_frame_end = 324
+	data_provider =\
+	[
+		('.assets/examples/target-240p-25fps.mp4', 55),
+		('.assets/examples/target-240p-30fps.mp4', 100),
+		('.assets/examples/target-240p-60fps.mp4', 50)
+	]
+	for target_path, frame_total in data_provider:
+		temp_directory_path = get_temp_directory_path(target_path)
+		create_temp(target_path)
+
+		assert extract_frames(target_path, 30) is True
+		assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total
+
+		clear_temp(target_path)
+
+
+def test_extract_frames_with_trim_end() -> None:
+	facefusion.globals.trim_frame_end = 100
+	data_provider =\
+	[
+		('.assets/examples/target-240p-25fps.mp4', 120),
+		('.assets/examples/target-240p-30fps.mp4', 100),
+		('.assets/examples/target-240p-60fps.mp4', 50)
+	]
+	for target_path, frame_total in data_provider:
+		temp_directory_path = get_temp_directory_path(target_path)
+		create_temp(target_path)
+
+		assert extract_frames(target_path, 30) is True
+		assert len(glob.glob1(temp_directory_path, '*.jpg')) == frame_total
+
+		clear_temp(target_path)