From e991d743c350a4daae393301306f2ff7c549123e Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 22 Jul 2026 21:19:12 +0200 Subject: [PATCH] fix colors --- facefusion/ffmpeg.py | 10 ++++++++-- facefusion/ffmpeg_builder.py | 12 +++++++++++- facefusion/ffprobe.py | 6 ++++-- facefusion/types.py | 4 +++- facefusion/video_manager.py | 13 +++++++++++-- 5 files changed, 37 insertions(+), 8 deletions(-) diff --git a/facefusion/ffmpeg.py b/facefusion/ffmpeg.py index f6764987..48f3db53 100644 --- a/facefusion/ffmpeg.py +++ b/facefusion/ffmpeg.py @@ -8,6 +8,7 @@ from tqdm import tqdm import facefusion.choices from facefusion import ffmpeg_builder, logger, process_manager, state_manager, translator +from facefusion.ffprobe import extract_video_metadata 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, VideoFormat @@ -108,6 +109,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 = extract_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 +117,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.tonemap_video(color_transfer) + ), ffmpeg_builder.prevent_frame_drop(), ffmpeg_builder.set_start_number(trim_frame_start), ffmpeg_builder.set_output(temp_frame_pattern) @@ -239,7 +244,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_colorspace('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 6769161c..07dc77f9 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, ColorTransfer, Command, CommandSet, Duration, Fps, StreamMode, VideoEncoder, VideoFormat, VideoPreset def run(commands : List[Command]) -> List[Command]: @@ -129,6 +129,16 @@ def prevent_frame_drop() -> List[Command]: return [ '-vsync', '0' ] +def tonemap_video(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_colorspace(color_space : str) -> 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 a37266f3..3fa0cbcc 100644 --- a/facefusion/ffprobe.py +++ b/facefusion/ffprobe.py @@ -61,7 +61,7 @@ def probe_format_entries(video_path : str, entries : List[str]) -> Dict[str, str @lru_cache(maxsize = 128) def extract_video_metadata(video_path : str) -> VideoMetadata: - video_entries = probe_video_entries(video_path, [ 'duration', 'width', 'height', 'r_frame_rate', 'bit_rate' ]) + video_entries = probe_video_entries(video_path, [ 'duration', 'width', 'height', 'r_frame_rate', 'bit_rate', 'color_transfer' ]) if video_entries.get('duration') == 'N/A': video_entries['duration'] = probe_format_entries(video_path, [ 'duration' ]).get('duration') @@ -75,6 +75,7 @@ def extract_video_metadata(video_path : str) -> VideoMetadata: width = int(video_entries.get('width')) height = int(video_entries.get('height')) bit_rate = int(video_entries.get('bit_rate')) + color_transfer = video_entries.get('color_transfer', 'unknown') video_metadata : VideoMetadata =\ { @@ -82,7 +83,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/types.py b/facefusion/types.py index 29b55d2d..2b060c86 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -118,13 +118,15 @@ Padding : TypeAlias = Tuple[int, int, int, int] Margin : TypeAlias = Tuple[int, int, int, int] Orientation = Literal['landscape', 'portrait'] Resolution : TypeAlias = Tuple[int, int] +ColorTransfer : TypeAlias = str VideoMetadata = TypedDict('VideoMetadata', { 'duration' : Duration, '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/video_manager.py b/facefusion/video_manager.py index cfd66960..7cb42d12 100644 --- a/facefusion/video_manager.py +++ b/facefusion/video_manager.py @@ -19,9 +19,11 @@ VIDEO_POOL_SET : VideoPoolSet =\ def create_video_reader_process(video_path : str, frame_position : int, video_fps : Fps) -> subprocess.Popen[bytes]: + color_transfer = extract_video_metadata(video_path).get('color_transfer') commands = ffmpeg_builder.chain( ffmpeg_builder.set_input_seek(frame_position / video_fps), ffmpeg_builder.set_input(video_path), + ffmpeg_builder.tonemap_video(color_transfer), ffmpeg_builder.prevent_frame_drop(), ffmpeg_builder.enforce_pixel_format('bgr24'), ffmpeg_builder.set_output_format('rawvideo'), @@ -32,12 +34,16 @@ def create_video_reader_process(video_path : str, frame_position : int, video_fp def create_video_sampler_process(video_path : str, frame_start : int, frame_end : int, frame_stride : int) -> subprocess.Popen[bytes]: + color_transfer = extract_video_metadata(video_path).get('color_transfer') sample_total = len([ frame_number for frame_number in range(frame_start, frame_end) if frame_number % frame_stride == 0 ]) sampler_thread_count = 4 commands = ffmpeg_builder.chain( ffmpeg_builder.set_thread_count(sampler_thread_count), ffmpeg_builder.set_input(video_path), - ffmpeg_builder.select_frame_samples(frame_start, frame_end, frame_stride), + ffmpeg_builder.concat( + ffmpeg_builder.select_frame_samples(frame_start, frame_end, frame_stride), + ffmpeg_builder.tonemap_video(color_transfer) + ), ffmpeg_builder.prevent_frame_drop(), ffmpeg_builder.set_frame_total(sample_total), ffmpeg_builder.enforce_pixel_format('bgr24'), @@ -162,7 +168,10 @@ def create_video_writer_process(target_path : str, temp_video_fps : Fps, temp_vi ffmpeg_builder.set_video_tag(output_video_encoder, temp_video_format), ffmpeg_builder.set_video_quality(output_video_encoder, output_video_quality), ffmpeg_builder.set_video_preset(output_video_encoder, output_video_preset), - ffmpeg_builder.set_video_fps(output_video_fps), + ffmpeg_builder.concat( + ffmpeg_builder.set_video_fps(output_video_fps), + ffmpeg_builder.convert_colorspace('bt709') + ), ffmpeg_builder.set_pixel_format(output_video_encoder), ffmpeg_builder.force_output(temp_video_path) )