viser_ffmpeg/
color.rs

1//! Bit depth, pixel format, and HDR color metadata helpers.
2
3use crate::{Codec, StreamInfo};
4
5/// Snapshot of source video color characteristics for encode preservation.
6#[derive(Debug, Clone, Default, PartialEq, Eq)]
7pub struct SourceFormat {
8    /// Preferred output pixel format (e.g. `yuv420p10le`).
9    pub pix_fmt: String,
10    /// Effective bit depth (8, 10, 12, or 16).
11    pub bit_depth: u8,
12    /// Color primaries from probe (e.g. `bt2020`).
13    pub color_primaries: String,
14    /// Color transfer from probe (e.g. `smpte2084`).
15    pub color_transfer: String,
16    /// Color matrix / space from probe.
17    pub color_space: String,
18    /// Whether the stream carries HDR signaling.
19    pub is_hdr: bool,
20}
21
22impl SourceFormat {
23    /// Builds a format snapshot from a probed video stream.
24    pub fn from_stream(stream: &StreamInfo) -> Self {
25        let bit_depth = bit_depth(stream);
26        let pix_fmt = if stream.pix_fmt.is_empty() {
27            yuv420p_for_depth(bit_depth).to_string()
28        } else {
29            stream.pix_fmt.clone()
30        };
31        Self {
32            pix_fmt,
33            bit_depth,
34            color_primaries: stream.color_primaries.clone(),
35            color_transfer: stream.color_transfer.clone(),
36            color_space: stream.color_space.clone(),
37            is_hdr: stream.is_hdr(),
38        }
39    }
40
41    /// Returns `true` when the source should be encoded at more than 8 bits per sample.
42    pub fn is_high_bit_depth(&self) -> bool {
43        self.bit_depth > 8
44    }
45}
46
47/// Returns the effective bit depth of a video stream.
48pub fn bit_depth(stream: &StreamInfo) -> u8 {
49    if stream.bits_per_raw_sample >= 10 {
50        return stream.bits_per_raw_sample.clamp(8, 16) as u8;
51    }
52    if stream.pix_fmt.contains("16") {
53        return 16;
54    }
55    if stream.pix_fmt.contains("12") {
56        return 12;
57    }
58    if stream.pix_fmt.contains("10") {
59        return 10;
60    }
61    8
62}
63
64/// Default 4:2:0 pixel format for a given bit depth.
65pub fn yuv420p_for_depth(depth: u8) -> &'static str {
66    match depth {
67        10 => "yuv420p10le",
68        12 => "yuv420p12le",
69        16 => "yuv420p16le",
70        _ => "yuv420p",
71    }
72}
73
74/// PSNR peak value for a given bit depth.
75pub fn psnr_peak(depth: u8) -> f64 {
76    match depth {
77        10 => 1023.0,
78        12 => 4095.0,
79        16 => 65535.0,
80        _ => 255.0,
81    }
82}
83
84/// Returns whether a software codec can encode at the requested bit depth.
85pub fn codec_supports_bit_depth(codec: Codec, depth: u8) -> bool {
86    if depth <= 8 {
87        return true;
88    }
89    match codec {
90        Codec::X264 | Codec::X265 | Codec::SvtAv1 => true,
91        _ => false,
92    }
93}
94
95/// FFmpeg output arguments that preserve source bit depth and HDR metadata.
96pub fn encode_color_args(codec: Codec, format: &SourceFormat) -> Vec<String> {
97    let mut args = Vec::new();
98
99    if format.is_high_bit_depth()
100        && codec.is_software()
101        && codec_supports_bit_depth(codec, format.bit_depth)
102    {
103        args.extend(["-pix_fmt".into(), format.pix_fmt.clone()]);
104        match codec {
105            Codec::X264 => args.extend(["-profile:v".into(), "high10".into()]),
106            Codec::X265 => args.extend(["-x265-params".into(), x265_params(format)]),
107            Codec::SvtAv1 => {}
108            _ => {}
109        }
110    }
111
112    if format.is_hdr {
113        append_color_metadata(&mut args, format);
114        if codec == Codec::X265 {
115            merge_x265_color_params(&mut args, format);
116        }
117    }
118
119    args
120}
121
122fn append_color_metadata(args: &mut Vec<String>, format: &SourceFormat) {
123    if !format.color_primaries.is_empty() {
124        args.extend(["-color_primaries".into(), format.color_primaries.clone()]);
125    }
126    if !format.color_transfer.is_empty() {
127        args.extend(["-color_trc".into(), format.color_transfer.clone()]);
128    }
129    if !format.color_space.is_empty() {
130        args.extend(["-colorspace".into(), format.color_space.clone()]);
131    }
132}
133
134fn x265_params(format: &SourceFormat) -> String {
135    let mut parts = Vec::new();
136    if format.bit_depth > 8 {
137        parts.push("profile=main10".into());
138    }
139    if format.is_hdr {
140        if !format.color_primaries.is_empty() {
141            parts.push(format!("colorprim={}", format.color_primaries));
142        }
143        if !format.color_transfer.is_empty() {
144            parts.push(format!("transfer={}", format.color_transfer));
145        }
146        if !format.color_space.is_empty() {
147            parts.push(format!("colormatrix={}", format.color_space));
148        }
149    }
150    parts.join(":")
151}
152
153fn merge_x265_color_params(args: &mut Vec<String>, format: &SourceFormat) {
154    let color = x265_params(format);
155    if color.is_empty() {
156        return;
157    }
158    if let Some(idx) = args.iter().position(|a| a == "-x265-params") {
159        let existing = args.get(idx + 1).cloned().unwrap_or_default();
160        let merged = if existing.is_empty() { color } else { format!("{existing}:{color}") };
161        args[idx + 1] = merged;
162    } else {
163        args.extend(["-x265-params".into(), color]);
164    }
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170    use crate::StreamInfo;
171
172    fn base_stream() -> StreamInfo {
173        StreamInfo {
174            index: 0,
175            codec_name: "h264".into(),
176            codec_long_name: String::new(),
177            codec_type: "video".into(),
178            profile: String::new(),
179            width: 1920,
180            height: 1080,
181            pix_fmt: "yuv420p".into(),
182            level: 0,
183            field_order: String::new(),
184            color_range: String::new(),
185            color_space: "bt709".into(),
186            color_transfer: "bt709".into(),
187            color_primaries: "bt709".into(),
188            duration: 0.0,
189            bit_rate: 0,
190            nb_frames: 0,
191            r_frame_rate: "24/1".into(),
192            avg_frame_rate: "24/1".into(),
193            sample_rate: 0,
194            channels: 0,
195            channel_layout: String::new(),
196            bits_per_raw_sample: 8,
197        }
198    }
199
200    #[test]
201    fn test_bit_depth_from_pix_fmt() {
202        let mut stream = base_stream();
203        stream.pix_fmt = "yuv420p10le".into();
204        assert_eq!(bit_depth(&stream), 10);
205    }
206
207    #[test]
208    fn test_source_format_high_bit_depth() {
209        let mut stream = base_stream();
210        stream.pix_fmt = "yuv420p10le".into();
211        let format = SourceFormat::from_stream(&stream);
212        assert_eq!(format.bit_depth, 10);
213        assert_eq!(format.pix_fmt, "yuv420p10le");
214        assert!(format.is_high_bit_depth());
215    }
216
217    #[test]
218    fn test_encode_color_args_x265_10bit_hdr() {
219        let mut stream = base_stream();
220        stream.pix_fmt = "yuv420p10le".into();
221        stream.color_transfer = "smpte2084".into();
222        stream.color_primaries = "bt2020".into();
223        stream.color_space = "bt2020nc".into();
224        let format = SourceFormat::from_stream(&stream);
225        let args = encode_color_args(Codec::X265, &format);
226        assert!(args.windows(2).any(|w| w[0] == "-pix_fmt" && w[1] == "yuv420p10le"));
227        assert!(args.iter().any(|a| a.contains("profile=main10")));
228        assert!(args.iter().any(|a| a.contains("transfer=smpte2084")));
229    }
230
231    #[test]
232    fn test_psnr_peak_scaling() {
233        assert_eq!(psnr_peak(8), 255.0);
234        assert_eq!(psnr_peak(10), 1023.0);
235    }
236}
viser_ffmpeg/color.rs

viser_ffmpeg/
color.rs
