fix colors

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)

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 = []
 

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

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']

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)
 	)