Fix/extrating merging progress bar total (#868)

* Introduce predict_video_frame_total() to fix progress bar while extracting and merging

* Introduce predict_video_frame_total() to fix progress bar while extracting and merging

* Use floor over ceil

facefusion/core.py

@@ -451,7 +451,7 @@ def process_video(start_time : float) -> ErrorCode:
 		return 1
 
 	logger.info(wording.get('merging_video').format(resolution = state_manager.get_item('output_video_resolution'), fps = state_manager.get_item('output_video_fps')), __name__)
-	if merge_video(state_manager.get_item('target_path'), temp_video_fps, state_manager.get_item('output_video_resolution'), state_manager.get_item('output_video_fps')):
+	if merge_video(state_manager.get_item('target_path'), temp_video_fps, state_manager.get_item('output_video_resolution'), state_manager.get_item('output_video_fps'), trim_frame_start, trim_frame_end):
 		logger.debug(wording.get('merging_video_succeed'), __name__)
 	else:
 		if is_process_stopping():

facefusion/ffmpeg.py

@@ -8,9 +8,9 @@ from tqdm import tqdm
 import facefusion.choices
 from facefusion import ffmpeg_builder, logger, process_manager, state_manager, wording
 from facefusion.filesystem import get_file_format, remove_file
-from facefusion.temp_helper import get_temp_file_path, get_temp_frames_pattern, resolve_temp_frame_paths
+from facefusion.temp_helper import get_temp_file_path, get_temp_frames_pattern
 from facefusion.types import AudioBuffer, Commands, EncoderSet, Fps, UpdateProgress
-from facefusion.vision import count_trim_frame_total, detect_video_duration, detect_video_fps
+from facefusion.vision import detect_video_duration, detect_video_fps, predict_video_frame_total
 
 
 def run_ffmpeg_with_progress(commands : Commands, update_progress : UpdateProgress) -> subprocess.Popen[bytes]:
@@ -100,7 +100,7 @@ def get_available_encoder_set() -> EncoderSet:
 
 
 def extract_frames(target_path : str, temp_video_resolution : str, temp_video_fps : Fps, trim_frame_start : int, trim_frame_end : int) -> bool:
-	extract_frame_total = count_trim_frame_total(target_path, trim_frame_start, trim_frame_end)
+	extract_frame_total = predict_video_frame_total(target_path, temp_video_fps, trim_frame_start, trim_frame_end)
 	temp_frames_pattern = get_temp_frames_pattern(target_path, '%08d')
 	commands = ffmpeg_builder.chain(
 		ffmpeg_builder.set_input(target_path),
@@ -200,11 +200,11 @@ def replace_audio(target_path : str, audio_path : str, output_path : str) -> boo
 	return run_ffmpeg(commands).returncode == 0
 
 
-def merge_video(target_path : str, temp_video_fps : Fps, output_video_resolution : str, output_video_fps : Fps) -> bool:
+def merge_video(target_path : str, temp_video_fps : Fps, output_video_resolution : str, output_video_fps : Fps, trim_frame_start : int, trim_frame_end : int) -> bool:
 	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(resolve_temp_frame_paths(target_path))
+	merge_frame_total = predict_video_frame_total(target_path, output_video_fps, trim_frame_start, trim_frame_end)
 	temp_file_path = get_temp_file_path(target_path)
 	temp_frames_pattern = get_temp_frames_pattern(target_path, '%08d')
 

facefusion/vision.py

@@ -1,3 +1,4 @@
+import math
 from functools import lru_cache
 from typing import List, Optional, Tuple
 
@@ -101,6 +102,14 @@ def count_video_frame_total(video_path : str) -> int:
 	return 0
 
 
+def predict_video_frame_total(video_path : str, fps : Fps, trim_frame_start : int, trim_frame_end : int) -> int:
+	if is_video(video_path):
+		target_video_fps = detect_video_fps(video_path)
+		extract_frame_total = count_trim_frame_total(video_path, trim_frame_start, trim_frame_end) * fps / target_video_fps
+		return math.floor(extract_frame_total)
+	return 0
+
+
 def detect_video_fps(video_path : str) -> Optional[float]:
 	if is_video(video_path):
 		video_capture = cv2.VideoCapture(video_path)

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_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, match_frame_color, normalize_resolution, pack_resolution, read_image, read_video_frame, restrict_image_resolution, restrict_trim_frame, restrict_video_fps, restrict_video_resolution, unpack_resolution, write_image
+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, match_frame_color, normalize_resolution, pack_resolution, predict_video_frame_total, read_image, read_video_frame, restrict_image_resolution, restrict_trim_frame, restrict_video_fps, restrict_video_resolution, unpack_resolution, write_image
 from .helper import get_test_example_file, get_test_examples_directory, get_test_output_file, prepare_test_output_directory
 
 
@@ -80,6 +80,13 @@ def test_count_video_frame_total() -> None:
 	assert count_video_frame_total('invalid') == 0
 
 
+def test_predict_video_frame_total() -> None:
+	assert predict_video_frame_total(get_test_example_file('target-240p-25fps.mp4'), 12.5, 0, 100) == 50
+	assert predict_video_frame_total(get_test_example_file('target-240p-25fps.mp4'), 25, 0, 100) == 100
+	assert predict_video_frame_total(get_test_example_file('target-240p-25fps.mp4'), 25, 0, 200) == 200
+	assert predict_video_frame_total('invalid', 25, 0, 100) == 0
+
+
 def test_detect_video_fps() -> None:
 	assert detect_video_fps(get_test_example_file('target-240p-25fps.mp4')) == 25.0
 	assert detect_video_fps(get_test_example_file('target-240p-30fps.mp4')) == 30.0