Merge pull request #11 from facefusion/develop

Develop

facefusion/utilities.py

@@ -34,14 +34,14 @@ 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:
@@ -49,7 +49,7 @@ def extract_frames(target_path : str, fps : float = 30) -> bool:
 	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) ])
 	elif trim_frame_start is not None:
@@ -76,17 +76,24 @@ def create_video(target_path : str, fps : float = 30) -> bool:
 
 
 def restore_audio(target_path : str, output_path : str) -> None:
+	fps = detect_fps(target_path)
 	trim_frame_start = facefusion.globals.trim_frame_start
 	trim_frame_end = facefusion.globals.trim_frame_end
 	temp_output_path = get_temp_output_path(target_path)
 	commands = [ '-hwaccel', 'auto', '-i', temp_output_path, '-i', target_path ]
-	if trim_frame_start is not None and trim_frame_end is not None:
+	if trim_frame_start is None and trim_frame_end is None:
-		commands.extend([ '-filter:v', 'select=between(n,' + str(trim_frame_start) + ',' + str(trim_frame_end) + ')' ])
+		commands.extend([ '-c:a', 'copy' ])
-	elif trim_frame_start is not None:
+	else:
-		commands.extend([ '-filter:v', 'select=gt(n,' + str(trim_frame_start) + ')' ])
+		if trim_frame_start is not None:
-	elif trim_frame_end is not None:
+			start_time = trim_frame_start / fps
-		commands.extend([ '-filter:v', 'select=lt(n,' + str(trim_frame_end) + ')' ])
+			commands.extend([ '-ss', str(start_time) ])
-	commands.extend([ '-c:a', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-y', output_path ])
+		else:
+			commands.extend([ '-ss', '0' ])
+		if trim_frame_end is not None:
+			end_time = trim_frame_end / fps
+			commands.extend([ '-to', str(end_time) ])
+		commands.extend([ '-c:a', 'aac' ])
+	commands.extend([ '-map', '0:v:0', '-map', '1:a:0', '-y', output_path ])
 	done = run_ffmpeg(commands)
 	if not done:
 		move_temp(target_path, output_path)

tests/test_cli.py

@@ -6,7 +6,7 @@ from facefusion.utilities import conditional_download
 
 
 @pytest.fixture(scope = 'module', autouse = True)
-def setup() -> None:
+def before_all() -> None:
 	conditional_download('.assets/examples',
 	[
 		'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg',
@@ -18,6 +18,7 @@ def setup() -> None:
 def test_image_to_image() -> None:
 	commands = [ 'python', 'run.py', '-s', '.assets/examples/source.jpg', '-t', '.assets/examples/target-1080p.jpg', '-o', '.assets/examples' ]
 	run = subprocess.run(commands, stdout = subprocess.PIPE)
+
 	assert run.returncode == 0
 	assert wording.get('processing_image_succeed') in run.stdout.decode()
 
@@ -25,5 +26,6 @@ def test_image_to_image() -> None:
 def test_image_to_video() -> None:
 	commands = [ 'python', 'run.py', '-s', '.assets/examples/source.jpg', '-t', '.assets/examples/target-1080p.mp4', '-o', '.assets/examples', '--trim-frame-end', '10' ]
 	run = subprocess.run(commands, stdout = subprocess.PIPE)
+
 	assert run.returncode == 0
 	assert wording.get('processing_video_succeed') in run.stdout.decode()

tests/test_utilities.py

@@ -1,21 +1,107 @@
+import glob
 import subprocess
 import pytest
 
-from facefusion.utilities import conditional_download, detect_fps
+import facefusion.globals
+from facefusion.utilities import conditional_download, detect_fps, extract_frames, create_temp, get_temp_directory_path, clear_temp
 
 
 @pytest.fixture(scope = 'module', autouse = True)
-def setup() -> None:
+def before_all() -> None:
+	facefusion.globals.temp_frame_quality = 100
+	facefusion.globals.trim_frame_start = None
+	facefusion.globals.trim_frame_end = None
+	facefusion.globals.temp_frame_format = 'png'
 	conditional_download('.assets/examples',
 	[
-		'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-1080p.mp4'
+		'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4'
 	])
-	subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-1080p.mp4', '-vf', 'fps=25', '.assets/examples/target-1080p-25fps.mp4' ])
+	subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=25', '.assets/examples/target-240p-25fps.mp4' ])
-	subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-1080p.mp4', '-vf', 'fps=30', '.assets/examples/target-1080p-30fps.mp4' ])
+	subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=30', '.assets/examples/target-240p-30fps.mp4' ])
-	subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-1080p.mp4', '-vf', 'fps=60', '.assets/examples/target-1080p-60fps.mp4' ])
+	subprocess.run([ 'ffmpeg', '-i', '.assets/examples/target-240p.mp4', '-vf', 'fps=60', '.assets/examples/target-240p-60fps.mp4' ])
+
+
+@pytest.fixture(scope = 'function', autouse = True)
+def before_each() -> None:
+	facefusion.globals.trim_frame_start = None
+	facefusion.globals.trim_frame_end = None
+	facefusion.globals.temp_frame_quality = 90
+	facefusion.globals.temp_frame_format = 'jpg'
 
 
 def test_detect_fps() -> None:
-	assert detect_fps('.assets/examples/target-1080p-25fps.mp4') == 25.0
+	assert detect_fps('.assets/examples/target-240p-25fps.mp4') == 25.0
-	assert detect_fps('.assets/examples/target-1080p-30fps.mp4') == 30.0
+	assert detect_fps('.assets/examples/target-240p-30fps.mp4') == 30.0
-	assert detect_fps('.assets/examples/target-1080p-60fps.mp4') == 60.0
+	assert detect_fps('.assets/examples/target-240p-60fps.mp4') == 60.0
+
+
+def test_extract_frames() -> None:
+	target_paths =\
+	[
+		'.assets/examples/target-240p-25fps.mp4',
+		'.assets/examples/target-240p-30fps.mp4',
+		'.assets/examples/target-240p-60fps.mp4'
+	]
+	for target_path in target_paths:
+		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')) == 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)