From a515d7fc3f7b956a172bd2eebfaf3620b2ba3eff Mon Sep 17 00:00:00 2001 From: Henry Ruhs Date: Thu, 23 Jul 2026 14:38:22 +0200 Subject: [PATCH] restrict hdr color transfer and tag the merge output as bt709 (#1188) * restrict hdr color transfer and tag the merge output as bt709 Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Tbcd6VWCiU4BQP1gywPr2a * restrict hdr color transfer and tag the merge output as bt709 Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Tbcd6VWCiU4BQP1gywPr2a * full video migration * compose the hdr fixture via the builder chain Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Tbcd6VWCiU4BQP1gywPr2a * compose the test fixtures via the builder and run_ffmpeg Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Tbcd6VWCiU4BQP1gywPr2a --------- Co-authored-by: Claude Opus 4.8 --- facefusion/ffmpeg.py | 11 +++- facefusion/ffmpeg_builder.py | 12 +++- facefusion/ffprobe.py | 9 +++ facefusion/types.py | 1 + facefusion/vision.py | 6 +- tests/test_ffmpeg.py | 108 +++++++++++++++++++++++++---------- tests/test_ffmpeg_builder.py | 14 ++++- 7 files changed, 123 insertions(+), 38 deletions(-) diff --git a/facefusion/ffmpeg.py b/facefusion/ffmpeg.py index a8ca32c1..42545158 100644 --- a/facefusion/ffmpeg.py +++ b/facefusion/ffmpeg.py @@ -7,7 +7,7 @@ from typing import List, Optional, cast from tqdm import tqdm import facefusion.choices -from facefusion import ffmpeg_builder, logger, process_manager, state_manager, translator +from facefusion import ffmpeg_builder, ffprobe, logger, process_manager, state_manager, translator from facefusion.filesystem import get_file_format, remove_file from facefusion.temp_helper import get_temp_file_path, get_temp_frame_pattern from facefusion.types import AudioBuffer, AudioEncoder, Command, EncoderSet, Fps, Resolution, UpdateProgress, VideoEncoder, VideoFormat @@ -108,6 +108,7 @@ def get_available_encoder_set() -> EncoderSet: def extract_frames(target_path : str, temp_video_resolution : Resolution, temp_video_fps : Fps, trim_frame_start : int, trim_frame_end : int) -> bool: + color_transfer = ffprobe.extract_static_video_metadata(target_path).get('color_transfer') extract_frame_total = predict_video_frame_total(target_path, temp_video_fps, trim_frame_start, trim_frame_end) temp_frame_pattern = get_temp_frame_pattern(target_path, '%08d') commands = ffmpeg_builder.chain( @@ -115,7 +116,10 @@ def extract_frames(target_path : str, temp_video_resolution : Resolution, temp_v ffmpeg_builder.set_media_resolution(pack_resolution(temp_video_resolution)), ffmpeg_builder.set_frame_quality(0), ffmpeg_builder.enforce_pixel_format('rgb24'), - ffmpeg_builder.select_frame_range(trim_frame_start, trim_frame_end, temp_video_fps), + ffmpeg_builder.concat( + ffmpeg_builder.select_frame_range(trim_frame_start, trim_frame_end, temp_video_fps), + ffmpeg_builder.restrict_color_transfer(color_transfer) + ), ffmpeg_builder.prevent_frame_drop(), ffmpeg_builder.set_start_number(trim_frame_start), ffmpeg_builder.set_output(temp_frame_pattern) @@ -239,7 +243,8 @@ def merge_video(target_path : str, temp_video_fps : Fps, output_video_resolution ffmpeg_builder.set_video_preset(output_video_encoder, output_video_preset), ffmpeg_builder.concat( ffmpeg_builder.set_video_fps(output_video_fps), - ffmpeg_builder.keep_video_alpha(output_video_encoder) + ffmpeg_builder.keep_video_alpha(output_video_encoder), + ffmpeg_builder.convert_color_space('bt709') ), ffmpeg_builder.set_pixel_format(output_video_encoder), ffmpeg_builder.force_output(temp_video_path) diff --git a/facefusion/ffmpeg_builder.py b/facefusion/ffmpeg_builder.py index 60b6bd2c..bbda986a 100644 --- a/facefusion/ffmpeg_builder.py +++ b/facefusion/ffmpeg_builder.py @@ -5,7 +5,7 @@ from typing import List, Optional import numpy from facefusion.filesystem import get_file_format -from facefusion.types import AudioEncoder, Command, CommandSet, Duration, Fps, StreamMode, VideoEncoder, VideoFormat, VideoPreset +from facefusion.types import AudioEncoder, ColorSpace, ColorTransfer, Command, CommandSet, Duration, Fps, StreamMode, VideoEncoder, VideoFormat, VideoPreset def run(commands : List[Command]) -> List[Command]: @@ -113,6 +113,16 @@ def prevent_frame_drop() -> List[Command]: return [ '-vsync', '0' ] +def restrict_color_transfer(color_transfer : ColorTransfer) -> List[Command]: + if color_transfer in [ 'smpte2084', 'arib-std-b67' ]: + return [ '-vf', 'scale=out_primaries=bt709:out_transfer=bt709:intent=perceptual' ] + return [] + + +def convert_color_space(color_space : ColorSpace) -> List[Command]: + return [ '-vf', 'scale=out_color_matrix=' + color_space + ':out_range=tv:out_primaries=' + color_space + ':out_transfer=' + color_space ] + + def select_media_range(frame_start : int, frame_end : int, media_fps : Fps) -> List[Command]: commands = [] diff --git a/facefusion/ffprobe.py b/facefusion/ffprobe.py index 99bf3ad5..24af62d4 100644 --- a/facefusion/ffprobe.py +++ b/facefusion/ffprobe.py @@ -60,6 +60,11 @@ def probe_format_entries(media_path : str, entries : List[str]) -> Dict[str, str return parse_entries(output) +@lru_cache(maxsize = 128) +def extract_static_audio_metadata(audio_path : str) -> AudioMetadata: + return extract_audio_metadata(audio_path) + + 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' ]) @@ -83,6 +88,10 @@ def extract_audio_metadata(audio_path : str) -> AudioMetadata: @lru_cache(maxsize = 128) +def extract_static_video_metadata(video_path : str) -> VideoMetadata: + return extract_video_metadata(video_path) + + def extract_video_metadata(video_path : str) -> VideoMetadata: video_entries = probe_video_entries(video_path, [ 'width', 'height', 'r_frame_rate', 'color_transfer' ]) format_entries = probe_format_entries(video_path, [ 'duration', 'bit_rate' ]) diff --git a/facefusion/types.py b/facefusion/types.py index 25241924..fc0c6790 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -80,6 +80,7 @@ CameraPoolSet = TypedDict('CameraPoolSet', }) ColorMode = Literal['rgb', 'rgba'] +ColorSpace = Literal['bt601', 'bt709', 'bt2020'] ColorTransfer : TypeAlias = str VisionFrame : TypeAlias = NDArray[Any] Mask : TypeAlias = NDArray[Any] diff --git a/facefusion/vision.py b/facefusion/vision.py index 6d37db27..0aa8ad16 100644 --- a/facefusion/vision.py +++ b/facefusion/vision.py @@ -141,7 +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): - return ffprobe.extract_video_metadata(video_path).get('frame_total') + return ffprobe.extract_static_video_metadata(video_path).get('frame_total') return 0 @@ -156,7 +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): - return ffprobe.extract_video_metadata(video_path).get('fps') + return ffprobe.extract_static_video_metadata(video_path).get('fps') return None @@ -204,7 +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): - return ffprobe.extract_video_metadata(video_path).get('resolution') + return ffprobe.extract_static_video_metadata(video_path).get('resolution') return None diff --git a/tests/test_ffmpeg.py b/tests/test_ffmpeg.py index 5fdeca1a..393600ba 100644 --- a/tests/test_ffmpeg.py +++ b/tests/test_ffmpeg.py @@ -1,16 +1,17 @@ import os -import subprocess import tempfile import pytest import facefusion.ffmpeg -from facefusion import process_manager, state_manager +from facefusion import ffmpeg, ffmpeg_builder, process_manager, state_manager from facefusion.download import conditional_download from facefusion.ffmpeg import concat_video, extract_frames, merge_video, read_audio_buffer, replace_audio, restore_audio +from facefusion.ffprobe import extract_video_metadata from facefusion.filesystem import copy_file from facefusion.temp_helper import clear_temp_directory, create_temp_directory, get_temp_file_path, resolve_temp_frame_set from facefusion.types import EncoderSet +from facefusion.vision import read_image from .helper import get_test_example_file, get_test_examples_directory, get_test_output_file, prepare_test_output_directory @@ -23,15 +24,52 @@ def before_all() -> None: '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'), get_test_example_file('source.wav') ]) - subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=25', get_test_example_file('target-240p-25fps.mp4') ]) - subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=30', get_test_example_file('target-240p-30fps.mp4') ]) - subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vf', 'fps=60', get_test_example_file('target-240p-60fps.mp4') ]) + ffmpeg.run_ffmpeg( + ffmpeg_builder.chain( + ffmpeg_builder.set_input(get_test_example_file('source.mp3')), + ffmpeg_builder.set_output(get_test_example_file('source.wav')) + ) + ) + + for video_fps in [ 25, 30, 60 ]: + ffmpeg.run_ffmpeg( + ffmpeg_builder.chain( + ffmpeg_builder.set_input(get_test_example_file('target-240p.mp4')), + ffmpeg_builder.set_video_fps(video_fps), + ffmpeg_builder.set_output(get_test_example_file('target-240p-' + str(video_fps) + 'fps.mp4')) + ) + ) + + ffmpeg.run_ffmpeg( + ffmpeg_builder.chain( + ffmpeg_builder.set_input(get_test_example_file('target-240p.mp4')), + [ + '-vf', + 'scale=out_transfer=smpte2084' + ], + ffmpeg_builder.set_output(get_test_example_file('target-240p-smpte2084.mp4')) + ) + ) for output_video_format in [ 'avi', 'm4v', 'mkv', 'mov', 'mp4', 'webm', 'wmv' ]: - subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.mp3'), '-i', get_test_example_file('target-240p.mp4'), '-ar', '16000', get_test_example_file('target-240p-16khz.' + output_video_format) ]) + ffmpeg.run_ffmpeg( + ffmpeg_builder.chain( + ffmpeg_builder.set_input(get_test_example_file('source.mp3')), + ffmpeg_builder.set_input(get_test_example_file('target-240p.mp4')), + ffmpeg_builder.set_audio_sample_rate(16000), + ffmpeg_builder.set_output(get_test_example_file('target-240p-16khz.' + output_video_format)) + ) + ) + + ffmpeg.run_ffmpeg( + ffmpeg_builder.chain( + ffmpeg_builder.set_input(get_test_example_file('source.mp3')), + ffmpeg_builder.set_input(get_test_example_file('target-240p.mp4')), + ffmpeg_builder.set_audio_sample_rate(48000), + ffmpeg_builder.set_output(get_test_example_file('target-240p-48khz.mp4')) + ) + ) - subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.mp3'), '-i', get_test_example_file('target-240p.mp4'), '-ar', '48000', get_test_example_file('target-240p-48khz.mp4') ]) state_manager.init_item('temp_path', tempfile.gettempdir()) state_manager.init_item('temp_frame_format', 'png') state_manager.init_item('output_audio_encoder', 'aac') @@ -67,43 +105,49 @@ def test_get_available_encoder_set() -> None: def test_extract_frames() -> None: test_set =\ [ - (get_test_example_file('target-240p-25fps.mp4'), 0, 270, 324), - (get_test_example_file('target-240p-25fps.mp4'), 224, 270, 55), - (get_test_example_file('target-240p-25fps.mp4'), 124, 224, 120), - (get_test_example_file('target-240p-25fps.mp4'), 0, 100, 120), - (get_test_example_file('target-240p-30fps.mp4'), 0, 324, 324), - (get_test_example_file('target-240p-30fps.mp4'), 224, 324, 100), - (get_test_example_file('target-240p-30fps.mp4'), 124, 224, 100), - (get_test_example_file('target-240p-30fps.mp4'), 0, 100, 100), - (get_test_example_file('target-240p-60fps.mp4'), 0, 648, 324), - (get_test_example_file('target-240p-60fps.mp4'), 224, 648, 212), - (get_test_example_file('target-240p-60fps.mp4'), 124, 224, 50), - (get_test_example_file('target-240p-60fps.mp4'), 0, 100, 50) + (get_test_example_file('target-240p-25fps.mp4'), 0, 270, 324, 55, 250), + (get_test_example_file('target-240p-25fps.mp4'), 224, 270, 55, 55, 250), + (get_test_example_file('target-240p-25fps.mp4'), 124, 224, 120, 55, 250), + (get_test_example_file('target-240p-25fps.mp4'), 0, 100, 120, 55, 250), + (get_test_example_file('target-240p-30fps.mp4'), 0, 324, 324, 55, 250), + (get_test_example_file('target-240p-30fps.mp4'), 224, 324, 100, 55, 250), + (get_test_example_file('target-240p-30fps.mp4'), 124, 224, 100, 55, 250), + (get_test_example_file('target-240p-30fps.mp4'), 0, 100, 100, 55, 250), + (get_test_example_file('target-240p-60fps.mp4'), 0, 648, 324, 55, 250), + (get_test_example_file('target-240p-60fps.mp4'), 224, 648, 212, 55, 250), + (get_test_example_file('target-240p-60fps.mp4'), 124, 224, 50, 55, 250), + (get_test_example_file('target-240p-60fps.mp4'), 0, 100, 50, 55, 250), + (get_test_example_file('target-240p-smpte2084.mp4'), 0, 1, 1, 32, 190) ] - for target_path, trim_frame_start, trim_frame_end, frame_total in test_set: + for target_path, trim_frame_start, trim_frame_end, frame_total, frame_std, frame_max in test_set: create_temp_directory(target_path) assert extract_frames(target_path, (452, 240), 30.0, trim_frame_start, trim_frame_end) is True assert len(resolve_temp_frame_set(target_path)) == frame_total + temp_vision_frame = read_image(resolve_temp_frame_set(target_path).get(trim_frame_start)) + + assert temp_vision_frame.std() > frame_std + assert temp_vision_frame.max() > frame_max + clear_temp_directory(target_path) def test_merge_video() -> None: - target_paths =\ + test_set =\ [ - get_test_example_file('target-240p-16khz.avi'), - get_test_example_file('target-240p-16khz.m4v'), - get_test_example_file('target-240p-16khz.mkv'), - get_test_example_file('target-240p-16khz.mp4'), - get_test_example_file('target-240p-16khz.mov'), - get_test_example_file('target-240p-16khz.webm'), - get_test_example_file('target-240p-16khz.wmv') + (get_test_example_file('target-240p-16khz.avi'), [ 'bt709', 'unknown' ]), + (get_test_example_file('target-240p-16khz.m4v'), [ 'bt709' ]), + (get_test_example_file('target-240p-16khz.mkv'), [ 'bt709' ]), + (get_test_example_file('target-240p-16khz.mp4'), [ 'bt709' ]), + (get_test_example_file('target-240p-16khz.mov'), [ 'bt709' ]), + (get_test_example_file('target-240p-16khz.webm'), [ 'bt709' ]), + (get_test_example_file('target-240p-16khz.wmv'), [ 'bt709' ]) ] output_video_encoders = get_available_encoder_set().get('video') - for target_path in target_paths: + for target_path, color_transfers in test_set: for output_video_encoder in output_video_encoders: state_manager.init_item('output_video_encoder', output_video_encoder) create_temp_directory(target_path) @@ -111,6 +155,10 @@ def test_merge_video() -> None: assert merge_video(target_path, 25.0, (452, 240), 25.0, 0, 1) is True + video_metadata = extract_video_metadata(get_temp_file_path(target_path)) + + assert video_metadata.get('color_transfer') in color_transfers + clear_temp_directory(target_path) state_manager.init_item('output_video_encoder', 'libx264') diff --git a/tests/test_ffmpeg_builder.py b/tests/test_ffmpeg_builder.py index 2a7b9df1..5be7f1f0 100644 --- a/tests/test_ffmpeg_builder.py +++ b/tests/test_ffmpeg_builder.py @@ -1,7 +1,7 @@ from shutil import which from facefusion import ffmpeg_builder -from facefusion.ffmpeg_builder import chain, concat, keep_video_alpha, run, select_frame_range, set_audio_quality, set_audio_sample_size, set_faststart, set_stream_mode, set_video_encoder, set_video_fps, set_video_quality, set_video_tag +from facefusion.ffmpeg_builder import chain, concat, convert_color_space, keep_video_alpha, restrict_color_transfer, run, select_frame_range, set_audio_quality, set_audio_sample_size, set_faststart, set_stream_mode, set_video_encoder, set_video_fps, set_video_quality, set_video_tag def test_run() -> None: @@ -48,6 +48,18 @@ def test_select_frame_range() -> None: assert select_frame_range(None, None, 30) == [ '-vf', 'fps=30' ] +def test_restrict_color_transfer() -> None: + assert restrict_color_transfer('smpte2084') == [ '-vf', 'scale=out_primaries=bt709:out_transfer=bt709:intent=perceptual' ] + assert restrict_color_transfer('arib-std-b67') == [ '-vf', 'scale=out_primaries=bt709:out_transfer=bt709:intent=perceptual' ] + assert restrict_color_transfer('invalid') == [] + + +def test_convert_color_space() -> None: + assert convert_color_space('bt601') == [ '-vf', 'scale=out_color_matrix=bt601:out_range=tv:out_primaries=bt601:out_transfer=bt601' ] + assert convert_color_space('bt709') == [ '-vf', 'scale=out_color_matrix=bt709:out_range=tv:out_primaries=bt709:out_transfer=bt709' ] + assert convert_color_space('bt2020') == [ '-vf', 'scale=out_color_matrix=bt2020:out_range=tv:out_primaries=bt2020:out_transfer=bt2020' ] + + def test_set_audio_sample_size() -> None: assert set_audio_sample_size(16) == [ '-f', 's16le' ] assert set_audio_sample_size(32) == [ '-f', 's32le' ]