mirrors/facefusion
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added progress bar to ffmpeg extract and merge

Browse Source
This commit is contained in:
chuckkay 2024-09-07 08:01:56 -04:00 committed by GitHub
parent 96282f192f
commit ea992e4f92
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 244 additions and 175 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

79
facefusion/ffmpeg.py
View File

@ -2,14 +2,14 @@ import os
import subprocess
import tempfile
from typing import List, Optional
from tqdm import tqdm # Import tqdm for progress bar
import filetype
import re
from facefusion import logger, process_manager, state_manager
from facefusion.filesystem import remove_file
from facefusion.temp_helper import get_temp_file_path, get_temp_frames_pattern
from facefusion.typing import AudioBuffer, Fps, OutputVideoPreset
from facefusion.vision import restrict_video_fps
from facefusion.vision import count_video_frame_total, restrict_video_fps
def run_ffmpeg(args: List[str]) -> subprocess.Popen[bytes]:
@ -54,14 +54,45 @@ def extract_frames(target_path : str, temp_video_resolution : str, temp_video_fp
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)])
frame_count = (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)])
target_frame_count = count_video_frame_total(target_path)
frame_count = (target_frame_count - trim_frame_start)
elif isinstance(trim_frame_end, int):
commands.extend(['-vf', 'trim=end_frame=' + str(trim_frame_end) + ',fps=' + str(temp_video_fps)])
frame_count = trim_frame_end
else:
commands.extend(['-vf', 'fps=' + str(temp_video_fps)])
frame_count = count_video_frame_total(target_path)
commands.extend(['-vsync', '0', temp_frames_pattern])
return run_ffmpeg(commands).returncode == 0
# Run ffmpeg and monitor progress
process = subprocess.Popen(['ffmpeg'] + commands, stderr=subprocess.PIPE, text=True)
pbar = tqdm(total=frame_count, desc="Extracting frames", unit = 'frame', ascii = ' =')
frame_re = re.compile(r'frame=\s*(\d+)')
previous_frame = 0
while True:
output = process.stderr.readline()
if output == '' and process.poll() is not None:
break
if not process_manager.is_processing():
process.terminate()
pbar.close()
return False # Indicate the process was canceled
if output:
match = frame_re.search(output)
if match:
frame = int(match.group(1))
pbar.update(frame - previous_frame)
previous_frame = frame
pbar.update(frame_count - previous_frame) # Ensure the progress bar reaches 100%
pbar.close()
process.wait()
return process.returncode == 0
def merge_video(target_path: str, output_video_resolution: str, output_video_fps: Fps) -> bool:
@ -85,8 +116,46 @@ def merge_video(target_path : str, output_video_resolution : str, output_video_f
if state_manager.get_item('output_video_encoder') in ['h264_videotoolbox', 'hevc_videotoolbox']:
commands.extend(['-q:v', str(state_manager.get_item('output_video_quality'))])
commands.extend(['-vf', 'framerate=fps=' + str(output_video_fps), '-pix_fmt', 'yuv420p', '-colorspace', 'bt709', '-y', temp_file_path])
return run_ffmpeg(commands).returncode == 0
# Calculate frame count
trim_frame_start = state_manager.get_item('trim_frame_start')
trim_frame_end = state_manager.get_item('trim_frame_end')
if isinstance(trim_frame_start, int) and isinstance(trim_frame_end, int):
frame_count = (trim_frame_end - trim_frame_start)
elif isinstance(trim_frame_start, int):
target_frame_count = count_video_frame_total(target_path)
frame_count = (target_frame_count - trim_frame_start)
elif isinstance(trim_frame_end, int):
frame_count = trim_frame_end
else:
frame_count = count_video_frame_total(target_path)
# Run ffmpeg and monitor progress
process = subprocess.Popen(['ffmpeg'] + commands, stderr=subprocess.PIPE, text=True)
pbar = tqdm(total=frame_count, desc="Merging video", unit = 'frame', ascii = ' =')
frame_re = re.compile(r'frame=\s*(\d+)')
previous_frame = 0
while True:
output = process.stderr.readline()
if output == '' and process.poll() is not None:
break
if not process_manager.is_processing():
process.terminate()
pbar.close()
return False # Indicate the process was canceled
if output:
match = frame_re.search(output)
if match:
frame = int(match.group(1))
pbar.update(frame - previous_frame)
previous_frame = frame
pbar.update(frame_count - previous_frame) # Ensure the progress bar reaches 100%
pbar.close()
process.wait()
return process.returncode == 0
def concat_video(output_path : str, temp_output_paths : List[str]) -> bool:
concat_video_path = tempfile.mktemp()