diff --git a/facefusion/ffprobe.py b/facefusion/ffprobe.py new file mode 100644 index 00000000..3d01365c --- /dev/null +++ b/facefusion/ffprobe.py @@ -0,0 +1,84 @@ +import subprocess +from typing import Dict, List + +from facefusion import ffprobe_builder +from facefusion.types import AudioMetadata, Buffer, Command, Fps, VideoMetadata + + +def run_ffprobe(commands : List[Command]) -> subprocess.Popen[Buffer]: + commands = ffprobe_builder.run(commands) + return subprocess.Popen(commands, stderr = subprocess.PIPE, stdout = subprocess.PIPE) + + +def probe_entries(media_path : str, entries : List[str]) -> Dict[str, str]: + media_entries = {} + + commands = ffprobe_builder.chain( + ffprobe_builder.show_entries(entries), + ffprobe_builder.format_to_key_value(), + ffprobe_builder.set_input(media_path) + ) + output, _ = run_ffprobe(commands).communicate() + + if output: + lines = output.decode().strip().splitlines() + + for line in lines: + if '=' in line: + key, value = line.split('=', 1) + media_entries[key] = value + + return media_entries + + +def extract_audio_metadata(audio_path : str) -> AudioMetadata: + audio_entries = probe_entries(audio_path, [ 'duration', 'sample_rate', 'channels', 'bit_rate' ]) + + duration = float(audio_entries.get('duration')) + sample_rate = int(audio_entries.get('sample_rate')) + frame_total = int(duration * sample_rate) + channel_total = int(audio_entries.get('channels')) + bit_rate = int(audio_entries.get('bit_rate')) + + audio_metadata : AudioMetadata =\ + { + 'duration' : duration, + 'frame_total' : frame_total, + 'channel_total' : channel_total, + 'sample_rate' : sample_rate, + 'bit_rate' : bit_rate + } + + return audio_metadata + + +def extract_video_metadata(video_path : str) -> VideoMetadata: + video_entries = probe_entries(video_path, [ 'duration', 'width', 'height', 'r_frame_rate', 'bit_rate' ]) + + duration = float(video_entries.get('duration')) + fps = extract_video_fps(video_entries.get('r_frame_rate')) + frame_total = int(duration * fps) + width = int(video_entries.get('width')) + height = int(video_entries.get('height')) + bit_rate = int(video_entries.get('bit_rate')) + + video_metadata : VideoMetadata =\ + { + 'duration' : duration, + 'frame_total' : frame_total, + 'fps' : fps, + 'resolution' : (width, height), + 'bit_rate' : bit_rate + } + + return video_metadata + + +def extract_video_fps(frame_rate : str) -> Fps: + if frame_rate and '/' in frame_rate: + numerator, denominator = frame_rate.split('/') + + if int(numerator) and int(denominator): + return int(numerator) / int(denominator) + + return 0.0 diff --git a/facefusion/ffprobe_builder.py b/facefusion/ffprobe_builder.py new file mode 100644 index 00000000..f1cbccc2 --- /dev/null +++ b/facefusion/ffprobe_builder.py @@ -0,0 +1,25 @@ +import itertools +import shutil +from typing import List + +from facefusion.types import Command + + +def run(commands : List[Command]) -> List[Command]: + return [ shutil.which('ffprobe'), '-loglevel', 'error' ] + commands + + +def chain(*commands : List[Command]) -> List[Command]: + return list(itertools.chain(*commands)) + + +def show_entries(entries : List[str]) -> List[Command]: + return [ '-show_entries', 'stream=' + ','.join(entries) ] + + +def format_to_key_value() -> List[Command]: + return [ '-of', 'default=noprint_wrappers=1' ] + + +def set_input(input_path : str) -> List[Command]: + return [ '-i', input_path ] diff --git a/facefusion/types.py b/facefusion/types.py index 75a58752..9272ce08 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -95,13 +95,33 @@ MelFilterBank : TypeAlias = NDArray[Any] Voice : TypeAlias = NDArray[Any] VoiceChunk : TypeAlias = NDArray[Any] +BitRate : TypeAlias = int +SampleRate : TypeAlias = int Fps : TypeAlias = float Duration : TypeAlias = float + +Buffer : TypeAlias = bytes Color : TypeAlias = Tuple[int, int, int, int] Padding : TypeAlias = Tuple[int, int, int, int] Margin : TypeAlias = Tuple[int, int, int, int] Orientation = Literal['landscape', 'portrait'] Resolution : TypeAlias = Tuple[int, int] +AudioMetadata = TypedDict('AudioMetadata', +{ + 'duration' : Duration, + 'frame_total' : int, + 'channel_total' : int, + 'sample_rate' : SampleRate, + 'bit_rate' : BitRate +}) +VideoMetadata = TypedDict('VideoMetadata', +{ + 'duration' : Duration, + 'frame_total' : int, + 'fps' : Fps, + 'resolution' : Resolution, + 'bit_rate' : BitRate +}) ProcessState = Literal['checking', 'processing', 'stopping', 'pending'] Args : TypeAlias = Dict[str, Any] diff --git a/tests/test_ffprobe.py b/tests/test_ffprobe.py new file mode 100644 index 00000000..86a47831 --- /dev/null +++ b/tests/test_ffprobe.py @@ -0,0 +1,53 @@ +import subprocess + +import pytest + +from facefusion import process_manager +from facefusion.download import conditional_download +from facefusion.ffprobe import extract_audio_metadata, extract_video_metadata +from .helper import get_test_example_file, get_test_examples_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + process_manager.start() + + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.mp3', + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.mp3'), '-t', '1.9', '-ar', '48000', '-ac', '2', get_test_example_file('source-48000khz-2ch.wav') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-t', '1', get_test_example_file('target-240p-1s.mov') ]) + + +def test_extract_audio_metadata() -> None: + audio_metadata = extract_audio_metadata(get_test_example_file('source.mp3')) + + assert audio_metadata.get('sample_rate') == 44100 + assert audio_metadata.get('channel_total') == 1 + assert audio_metadata.get('frame_total') == 167039 + assert audio_metadata.get('bit_rate') == 128000 + + audio_metadata = extract_audio_metadata(get_test_example_file('source-48000khz-2ch.wav')) + + assert audio_metadata.get('sample_rate') == 48000 + assert audio_metadata.get('channel_total') == 2 + assert audio_metadata.get('frame_total') == 91200 + assert audio_metadata.get('bit_rate') == 1536000 + + +def test_extract_video_metadata() -> None: + video_metadata = extract_video_metadata(get_test_example_file('target-240p.mp4')) + + assert video_metadata.get('fps') == 25.0 + assert video_metadata.get('duration') == 10.8 + assert video_metadata.get('resolution') == (426, 226) + assert video_metadata.get('bit_rate') == 138754 + + video_metadata = extract_video_metadata(get_test_example_file('target-240p-1s.mov')) + + assert video_metadata.get('fps') == 25.0 + assert video_metadata.get('duration') == 1.0 + assert video_metadata.get('resolution') == (426, 226) diff --git a/tests/test_ffprobe_builder.py b/tests/test_ffprobe_builder.py new file mode 100644 index 00000000..88aaf641 --- /dev/null +++ b/tests/test_ffprobe_builder.py @@ -0,0 +1,30 @@ +from shutil import which + +from facefusion import ffprobe_builder +from facefusion.ffprobe_builder import chain, format_to_key_value, run, set_input, show_entries + + +def test_run() -> None: + assert run([ '-v', 'error' ]) == [ which('ffprobe'), '-loglevel', 'error', '-v', 'error' ] + + +def test_chain() -> None: + assert chain( + ffprobe_builder.show_entries([ 'sample_rate' ]), + ffprobe_builder.format_to_key_value(), + ffprobe_builder.set_input('audio.mp3') + ) == [ '-show_entries', 'stream=sample_rate', '-of', 'default=noprint_wrappers=1', '-i', 'audio.mp3' ] + + +def test_show_entries() -> None: + assert show_entries([ 'duration' ]) == [ '-show_entries', 'stream=duration' ] + assert show_entries([ 'duration', 'sample_rate']) == [ '-show_entries', 'stream=duration,sample_rate' ] + + +def test_format_to_key_value() -> None: + assert format_to_key_value() == [ '-of', 'default=noprint_wrappers=1' ] + + +def test_set_input() -> None: + assert set_input('input.mp3') == [ '-i', 'input.mp3' ] + assert set_input('input.wav') == [ '-i', 'input.wav' ]