add meaningful tests

facefusion/ffprobe.py

@@ -32,6 +32,20 @@ def get_audio_entries(audio_path : str) -> Dict[str, str]:
 	return audio_entries
 
 
+def detect_audio_codec(audio_path : str) -> Optional[str]:
+	commands = ffprobe_builder.chain(
+		ffprobe_builder.show_entries([ 'codec_name' ]),
+		ffprobe_builder.format_to_value(),
+		ffprobe_builder.set_input(audio_path)
+	)
+	process = run_ffprobe(commands)
+	output, _ = process.communicate()
+
+	if output:
+		return output.decode().strip().split(os.linesep)[0]
+	return None
+
+
 def detect_audio_sample_rate(audio_path : str) -> Optional[int]:
 	audio_entries = get_audio_entries(audio_path)
 	sample_rate = audio_entries.get('sample_rate')
@@ -58,3 +72,17 @@ def detect_audio_frame_total(audio_path : str) -> Optional[int]:
 	if audio_duration and audio_sample_rate:
 		return int(float(audio_duration) * int(audio_sample_rate))
 	return None
+
+
+def detect_video_codec(video_path : str) -> Optional[str]:
+	commands = ffprobe_builder.chain(
+		ffprobe_builder.show_entries([ 'codec_name' ]),
+		ffprobe_builder.format_to_value(),
+		ffprobe_builder.set_input(video_path)
+	)
+	process = run_ffprobe(commands)
+	output, _ = process.communicate()
+
+	if output:
+		return output.decode().strip().split(os.linesep)[0]
+	return None

tests/test_ffmpeg.py

@@ -8,7 +8,8 @@ import facefusion.ffmpeg
 from facefusion import process_manager, state_manager
 from facefusion.download import conditional_download
 from facefusion.ffmpeg import concat_video, extract_frames, merge_video, read_audio_buffer, replace_audio, restore_audio, sanitize_audio, sanitize_image, sanitize_video, spawn_frames
-from facefusion.filesystem import copy_file, is_audio, is_image, is_video
+from facefusion.ffprobe import detect_audio_codec, detect_video_codec
+from facefusion.filesystem import copy_file, is_image
 from facefusion.temp_helper import clear_temp_directory, create_temp_directory, get_temp_file_path, resolve_temp_frame_paths
 from facefusion.types import EncoderSet
 from .assert_helper import create_media_reader, get_test_example_file, get_test_examples_directory, get_test_output_path, prepare_test_output_directory
@@ -32,6 +33,7 @@ def before_all() -> None:
 		subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.mp3'), '-i', get_test_example_file('target-240p.mp4'), '-ar', '16000', get_test_example_file('target-240p-16khz.' + output_video_format) ])
 
 	subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.mp3'), '-i', get_test_example_file('target-240p.mp4'), '-ar', '48000', get_test_example_file('target-240p-48khz.mp4') ])
+	subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-c:v', 'libx265', '-an', get_test_example_file('target-240p-h265.mp4') ])
 	state_manager.init_item('temp_path', tempfile.gettempdir())
 	state_manager.init_item('temp_frame_format', 'png')
 	state_manager.init_item('output_audio_encoder', 'aac')
@@ -213,32 +215,38 @@ def test_replace_audio() -> None:
 
 
 def test_sanitize_audio() -> None:
-	audio_path = get_test_example_file('source.mp3')
-	output_strict_path = get_test_output_path('output-strict.mp3')
-	output_moderate_path = get_test_output_path('output-moderate.mp3')
+	file_path = get_test_example_file('source.wav')
+	output_paths =\
+	[
+		get_test_output_path('test-sanitize-audio-strict.mp3'),
+		get_test_output_path('test-sanitize-audio-moderate.wav')
+	]
 
-	assert sanitize_audio('mp3', create_media_reader(audio_path), output_strict_path, 'strict') is True
-	assert is_audio(output_strict_path) is True
+	assert sanitize_audio('wav', create_media_reader(file_path), output_paths[0], 'strict') is True
+	assert detect_audio_codec(output_paths[0]) == 'mp3'
 
-	assert sanitize_audio('mp3', create_media_reader(audio_path), output_moderate_path, 'moderate') is True
-	assert is_audio(output_moderate_path) is True
+	assert sanitize_audio('wav', create_media_reader(file_path), output_paths[1], 'moderate') is True
+	assert detect_audio_codec(output_paths[1]) == 'pcm_s16le'
 
 
 def test_sanitize_image() -> None:
-	source_path = get_test_example_file('source.jpg')
-	output_path = get_test_output_path('output.jpg')
+	file_path = get_test_example_file('source.jpg')
+	output_path = get_test_output_path('test-sanitize-image.jpg')
 
-	assert sanitize_image('jpeg', create_media_reader(source_path), output_path) is True
+	assert sanitize_image('jpeg', create_media_reader(file_path), output_path) is True
 	assert is_image(output_path) is True
 
 
 def test_sanitize_video() -> None:
-	video_path = get_test_example_file('target-240p.mp4')
-	output_strict_path = get_test_output_path('output-strict.mp4')
-	output_moderate_path = get_test_output_path('output-moderate.mp4')
+	file_path = get_test_example_file('target-240p-h265.mp4')
+	output_paths =\
+	[
+		get_test_output_path('test-sanitize-video-strict.mp4'),
+		get_test_output_path('test-sanitize-video-moderate.mp4')
+	]
 
-	assert sanitize_video('mp4', create_media_reader(video_path), output_strict_path, 'strict') is True
-	assert is_video(output_strict_path) is True
+	assert sanitize_video('mp4', create_media_reader(file_path), output_paths[0], 'strict') is True
+	assert detect_video_codec(output_paths[0]) == 'h264'
 
-	assert sanitize_video('mp4', create_media_reader(video_path), output_moderate_path, 'moderate') is True
-	assert is_video(output_moderate_path) is True
+	assert sanitize_video('mp4', create_media_reader(file_path), output_paths[1], 'moderate') is True
+	assert detect_video_codec(output_paths[1]) == 'hevc'