diff --git a/facefusion/ffprobe.py b/facefusion/ffprobe.py index 3d01365c..99bf3ad5 100644 --- a/facefusion/ffprobe.py +++ b/facefusion/ffprobe.py @@ -1,4 +1,5 @@ import subprocess +from functools import lru_cache from typing import Dict, List from facefusion import ffprobe_builder @@ -10,16 +11,9 @@ def run_ffprobe(commands : List[Command]) -> subprocess.Popen[Buffer]: return subprocess.Popen(commands, stderr = subprocess.PIPE, stdout = subprocess.PIPE) -def probe_entries(media_path : str, entries : List[str]) -> Dict[str, str]: +def parse_entries(output : Buffer) -> 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() @@ -31,14 +25,50 @@ def probe_entries(media_path : str, entries : List[str]) -> Dict[str, str]: return media_entries -def extract_audio_metadata(audio_path : str) -> AudioMetadata: - audio_entries = probe_entries(audio_path, [ 'duration', 'sample_rate', 'channels', 'bit_rate' ]) +def probe_audio_entries(audio_path : str, entries : List[str]) -> Dict[str, str]: + commands = ffprobe_builder.chain( + ffprobe_builder.select_stream('a:0'), + ffprobe_builder.show_stream_entries(entries), + ffprobe_builder.format_to_key_value(), + ffprobe_builder.set_input(audio_path) + ) + output, _ = run_ffprobe(commands).communicate() - duration = float(audio_entries.get('duration')) + return parse_entries(output) + + +def probe_video_entries(video_path : str, entries : List[str]) -> Dict[str, str]: + commands = ffprobe_builder.chain( + ffprobe_builder.select_stream('v:0'), + ffprobe_builder.show_stream_entries(entries), + ffprobe_builder.format_to_key_value(), + ffprobe_builder.set_input(video_path) + ) + output, _ = run_ffprobe(commands).communicate() + + return parse_entries(output) + + +def probe_format_entries(media_path : str, entries : List[str]) -> Dict[str, str]: + commands = ffprobe_builder.chain( + ffprobe_builder.show_format_entries(entries), + ffprobe_builder.format_to_key_value(), + ffprobe_builder.set_input(media_path) + ) + output, _ = run_ffprobe(commands).communicate() + + return parse_entries(output) + + +def extract_audio_metadata(audio_path : str) -> AudioMetadata: + audio_entries = probe_audio_entries(audio_path, [ 'sample_rate', 'channels' ]) + format_entries = probe_format_entries(audio_path, [ 'duration', 'bit_rate' ]) + + duration = float(format_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')) + bit_rate = int(format_entries.get('bit_rate')) audio_metadata : AudioMetadata =\ { @@ -52,15 +82,18 @@ def extract_audio_metadata(audio_path : str) -> AudioMetadata: return audio_metadata +@lru_cache(maxsize = 128) def extract_video_metadata(video_path : str) -> VideoMetadata: - video_entries = probe_entries(video_path, [ 'duration', 'width', 'height', 'r_frame_rate', 'bit_rate' ]) + video_entries = probe_video_entries(video_path, [ 'width', 'height', 'r_frame_rate', 'color_transfer' ]) + format_entries = probe_format_entries(video_path, [ 'duration', 'bit_rate' ]) - duration = float(video_entries.get('duration')) + duration = float(format_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')) + bit_rate = int(format_entries.get('bit_rate')) + color_transfer = video_entries.get('color_transfer', 'unknown') video_metadata : VideoMetadata =\ { @@ -68,7 +101,8 @@ def extract_video_metadata(video_path : str) -> VideoMetadata: 'frame_total' : frame_total, 'fps' : fps, 'resolution' : (width, height), - 'bit_rate' : bit_rate + 'bit_rate' : bit_rate, + 'color_transfer' : color_transfer } return video_metadata diff --git a/facefusion/ffprobe_builder.py b/facefusion/ffprobe_builder.py index f1cbccc2..22e7d146 100644 --- a/facefusion/ffprobe_builder.py +++ b/facefusion/ffprobe_builder.py @@ -13,10 +13,18 @@ def chain(*commands : List[Command]) -> List[Command]: return list(itertools.chain(*commands)) -def show_entries(entries : List[str]) -> List[Command]: +def select_stream(stream : str) -> List[Command]: + return [ '-select_streams', stream ] + + +def show_stream_entries(entries : List[str]) -> List[Command]: return [ '-show_entries', 'stream=' + ','.join(entries) ] +def show_format_entries(entries : List[str]) -> List[Command]: + return [ '-show_entries', 'format=' + ','.join(entries) ] + + def format_to_key_value() -> List[Command]: return [ '-of', 'default=noprint_wrappers=1' ] diff --git a/facefusion/types.py b/facefusion/types.py index 9272ce08..5c0b8f26 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -77,6 +77,7 @@ CameraPoolSet = TypedDict('CameraPoolSet', }) ColorMode = Literal['rgb', 'rgba'] +ColorTransfer : TypeAlias = str VisionFrame : TypeAlias = NDArray[Any] Mask : TypeAlias = NDArray[Any] Points : TypeAlias = NDArray[Any] @@ -120,7 +121,8 @@ VideoMetadata = TypedDict('VideoMetadata', 'frame_total' : int, 'fps' : Fps, 'resolution' : Resolution, - 'bit_rate' : BitRate + 'bit_rate' : BitRate, + 'color_transfer' : ColorTransfer }) ProcessState = Literal['checking', 'processing', 'stopping', 'pending'] diff --git a/facefusion/vision.py b/facefusion/vision.py index 79ddb536..6d37db27 100644 --- a/facefusion/vision.py +++ b/facefusion/vision.py @@ -6,6 +6,7 @@ import cv2 import numpy from cv2.typing import Size +from facefusion import ffprobe from facefusion.common_helper import is_windows from facefusion.filesystem import get_file_extension, is_image, is_video from facefusion.thread_helper import thread_lock, thread_semaphore @@ -140,12 +141,7 @@ def select_video_frames(video_path : str, frame_number : int = 0, frame_offset : def count_video_frame_total(video_path : str) -> int: if is_video(video_path): - video_capture = get_video_capture(video_path) - - if video_capture and video_capture.isOpened(): - with thread_semaphore(): - video_frame_total = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT)) - return video_frame_total + return ffprobe.extract_video_metadata(video_path).get('frame_total') return 0 @@ -160,12 +156,7 @@ def predict_video_frame_total(video_path : str, fps : Fps, trim_frame_start : in def detect_video_fps(video_path : str) -> Optional[float]: if is_video(video_path): - video_capture = get_video_capture(video_path) - - if video_capture and video_capture.isOpened(): - with thread_semaphore(): - video_fps = video_capture.get(cv2.CAP_PROP_FPS) - return video_fps + return ffprobe.extract_video_metadata(video_path).get('fps') return None @@ -213,13 +204,7 @@ def restrict_trim_frame(video_path : str, trim_frame_start : Optional[int], trim def detect_video_resolution(video_path : str) -> Optional[Resolution]: if is_video(video_path): - video_capture = get_video_capture(video_path) - - if video_capture and video_capture.isOpened(): - with thread_semaphore(): - width = video_capture.get(cv2.CAP_PROP_FRAME_WIDTH) - height = video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT) - return int(width), int(height) + return ffprobe.extract_video_metadata(video_path).get('resolution') return None diff --git a/tests/test_ffprobe.py b/tests/test_ffprobe.py index 86a47831..4718dcf1 100644 --- a/tests/test_ffprobe.py +++ b/tests/test_ffprobe.py @@ -19,6 +19,7 @@ def before_all() -> None: ]) 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.mkv') ]) subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-t', '1', get_test_example_file('target-240p-1s.mov') ]) @@ -35,7 +36,7 @@ def test_extract_audio_metadata() -> None: 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 + assert audio_metadata.get('bit_rate') == 1536328 def test_extract_video_metadata() -> None: @@ -44,7 +45,15 @@ def test_extract_video_metadata() -> None: 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 + assert video_metadata.get('bit_rate') == 141981 + assert video_metadata.get('color_transfer') == 'smpte170m' + + video_metadata = extract_video_metadata(get_test_example_file('target-240p-1s.mkv')) + + assert video_metadata.get('fps') == 25.0 + assert video_metadata.get('duration') == 1.0 + assert video_metadata.get('frame_total') == 25 + assert video_metadata.get('resolution') == (426, 226) video_metadata = extract_video_metadata(get_test_example_file('target-240p-1s.mov')) diff --git a/tests/test_ffprobe_builder.py b/tests/test_ffprobe_builder.py index 88aaf641..050b9097 100644 --- a/tests/test_ffprobe_builder.py +++ b/tests/test_ffprobe_builder.py @@ -1,7 +1,7 @@ from shutil import which from facefusion import ffprobe_builder -from facefusion.ffprobe_builder import chain, format_to_key_value, run, set_input, show_entries +from facefusion.ffprobe_builder import chain, format_to_key_value, run, select_stream, set_input, show_stream_entries def test_run() -> None: @@ -10,15 +10,21 @@ def test_run() -> None: def test_chain() -> None: assert chain( - ffprobe_builder.show_entries([ 'sample_rate' ]), + ffprobe_builder.select_stream('a:0'), + ffprobe_builder.show_stream_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' ] + ) == [ '-select_streams', 'a:0', '-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_select_stream() -> None: + assert select_stream('a:0') == [ '-select_streams', 'a:0' ] + assert select_stream('v:0') == [ '-select_streams', 'v:0' ] + + +def test_show_stream_entries() -> None: + assert show_stream_entries([ 'duration' ]) == [ '-show_entries', 'stream=duration' ] + assert show_stream_entries([ 'duration', 'sample_rate' ]) == [ '-show_entries', 'stream=duration,sample_rate' ] def test_format_to_key_value() -> None: