viser_ffmpeg/
color.rs

1//! Bit depth, pixel format, and HDR color metadata helpers.
2
3use crate::{Codec, Hdr10Metadata, 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    /// HDR10 static metadata (mastering display + MaxCLL/MaxFALL), when probed.
21    pub hdr10: Option<Hdr10Metadata>,
22}
23
24impl SourceFormat {
25    /// Builds a format snapshot from a probed video stream.
26    pub fn from_stream(stream: &StreamInfo) -> Self {
27        let bit_depth = bit_depth(stream);
28        let pix_fmt = if stream.pix_fmt.is_empty() {
29            yuv420p_for_depth(bit_depth).to_string()
30        } else {
31            stream.pix_fmt.clone()
32        };
33        Self {
34            pix_fmt,
35            bit_depth,
36            color_primaries: stream.color_primaries.clone(),
37            color_transfer: stream.color_transfer.clone(),
38            color_space: stream.color_space.clone(),
39            is_hdr: stream.is_hdr(),
40            hdr10: None,
41        }
42    }
43
44    /// Returns `true` when the source should be encoded at more than 8 bits per sample.
45    pub fn is_high_bit_depth(&self) -> bool {
46        self.bit_depth > 8
47    }
48
49    /// Probes and attaches HDR10 static metadata (mastering display + MaxCLL)
50    /// when the source is HDR, so it can be re-signalled on the encode.
51    ///
52    /// A no-op for SDR sources. Probe failures are swallowed (best-effort): a
53    /// missing mastering-display block degrades to colour-primary signalling
54    /// rather than failing the encode.
55    pub async fn enrich_hdr10(mut self, path: &str) -> Self {
56        if self.is_hdr {
57            if let Ok(Some(md)) = crate::probe_hdr10_metadata(path).await {
58                self.hdr10 = Some(md);
59            }
60        }
61        self
62    }
63}
64
65/// Returns the effective bit depth of a video stream.
66pub fn bit_depth(stream: &StreamInfo) -> u8 {
67    if stream.bits_per_raw_sample >= 10 {
68        return stream.bits_per_raw_sample.clamp(8, 16) as u8;
69    }
70    if stream.pix_fmt.contains("16") {
71        return 16;
72    }
73    if stream.pix_fmt.contains("12") {
74        return 12;
75    }
76    if stream.pix_fmt.contains("10") {
77        return 10;
78    }
79    8
80}
81
82/// Default 4:2:0 pixel format for a given bit depth.
83pub fn yuv420p_for_depth(depth: u8) -> &'static str {
84    match depth {
85        10 => "yuv420p10le",
86        12 => "yuv420p12le",
87        16 => "yuv420p16le",
88        _ => "yuv420p",
89    }
90}
91
92/// PSNR peak value for a given bit depth.
93pub fn psnr_peak(depth: u8) -> f64 {
94    match depth {
95        10 => 1023.0,
96        12 => 4095.0,
97        16 => 65535.0,
98        _ => 255.0,
99    }
100}
101
102/// Returns whether a software codec can encode at the requested bit depth.
103pub fn codec_supports_bit_depth(codec: Codec, depth: u8) -> bool {
104    if depth <= 8 {
105        return true;
106    }
107    match codec {
108        Codec::X264 | Codec::X265 | Codec::SvtAv1 | Codec::Vp9 => true,
109        _ => false,
110    }
111}
112
113/// FFmpeg output arguments that preserve source bit depth and HDR metadata.
114pub fn encode_color_args(codec: Codec, format: &SourceFormat) -> Vec<String> {
115    let mut args = Vec::new();
116
117    if format.is_high_bit_depth()
118        && codec.is_software()
119        && codec_supports_bit_depth(codec, format.bit_depth)
120    {
121        args.extend(["-pix_fmt".into(), format.pix_fmt.clone()]);
122        match codec {
123            Codec::X264 => args.extend(["-profile:v".into(), "high10".into()]),
124            Codec::X265 => args.extend(["-x265-params".into(), x265_params(format)]),
125            Codec::SvtAv1 => {}
126            _ => {}
127        }
128    }
129
130    if format.is_hdr {
131        append_color_metadata(&mut args, format);
132        match codec {
133            Codec::X265 => merge_x265_color_params(&mut args, format),
134            Codec::SvtAv1 => {
135                let params = svtav1_hdr_params(format);
136                if !params.is_empty() {
137                    // A second `-svtav1-params` may be added by the rate-control
138                    // builder; `coalesce_svtav1_params` merges them before the
139                    // encode runs (the last flag would otherwise win outright).
140                    args.extend(["-svtav1-params".into(), params]);
141                }
142            }
143            _ => {}
144        }
145    }
146
147    args
148}
149
150fn append_color_metadata(args: &mut Vec<String>, format: &SourceFormat) {
151    if !format.color_primaries.is_empty() {
152        args.extend(["-color_primaries".into(), format.color_primaries.clone()]);
153    }
154    if !format.color_transfer.is_empty() {
155        args.extend(["-color_trc".into(), format.color_transfer.clone()]);
156    }
157    if !format.color_space.is_empty() {
158        args.extend(["-colorspace".into(), format.color_space.clone()]);
159    }
160}
161
162fn x265_params(format: &SourceFormat) -> String {
163    let mut parts = Vec::new();
164    if format.bit_depth > 8 {
165        parts.push("profile=main10".into());
166    }
167    if format.is_hdr {
168        if !format.color_primaries.is_empty() {
169            parts.push(format!("colorprim={}", format.color_primaries));
170        }
171        if !format.color_transfer.is_empty() {
172            parts.push(format!("transfer={}", format.color_transfer));
173        }
174        if !format.color_space.is_empty() {
175            parts.push(format!("colormatrix={}", format.color_space));
176        }
177        if let Some(hdr10) = &format.hdr10 {
178            if let Some(display) = &hdr10.mastering_display {
179                parts.push(format!("master-display={}", display.to_x265_string()));
180            }
181            if let Some(max_cll) = hdr10.max_cll {
182                // x265 expects "MaxCLL,MaxFALL"; MaxFALL defaults to 0 when absent.
183                parts.push(format!("max-cll={},{}", max_cll, hdr10.max_fall.unwrap_or(0)));
184            }
185        }
186    }
187    parts.join(":")
188}
189
190/// SVT-AV1 `-svtav1-params` HDR10 static-metadata fragment (without the flag).
191///
192/// SVT-AV1 shares x265's `mastering-display` grammar but spells content light
193/// `content-light=MaxCLL,MaxFALL`. Colour primaries/transfer/matrix are carried
194/// by the standard `-color_*` options [`append_color_metadata`] emits.
195fn svtav1_hdr_params(format: &SourceFormat) -> String {
196    let mut parts = Vec::new();
197    if let Some(hdr10) = &format.hdr10 {
198        if let Some(display) = &hdr10.mastering_display {
199            parts.push(format!("mastering-display={}", display.to_svtav1_string()));
200        }
201        if let Some(max_cll) = hdr10.max_cll {
202            parts.push(format!("content-light={},{}", max_cll, hdr10.max_fall.unwrap_or(0)));
203        }
204    }
205    parts.join(":")
206}
207
208fn merge_x265_color_params(args: &mut Vec<String>, format: &SourceFormat) {
209    let color = x265_params(format);
210    if color.is_empty() {
211        return;
212    }
213    if let Some(idx) = args.iter().position(|a| a == "-x265-params") {
214        let existing = args.get(idx + 1).cloned().unwrap_or_default();
215        let merged = if existing.is_empty() { color } else { format!("{existing}:{color}") };
216        args[idx + 1] = merged;
217    } else {
218        args.extend(["-x265-params".into(), color]);
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225    use crate::StreamInfo;
226
227    fn base_stream() -> StreamInfo {
228        StreamInfo {
229            index: 0,
230            codec_name: "h264".into(),
231            codec_long_name: String::new(),
232            codec_type: "video".into(),
233            profile: String::new(),
234            width: 1920,
235            height: 1080,
236            pix_fmt: "yuv420p".into(),
237            level: 0,
238            field_order: String::new(),
239            color_range: String::new(),
240            color_space: "bt709".into(),
241            color_transfer: "bt709".into(),
242            color_primaries: "bt709".into(),
243            duration: 0.0,
244            bit_rate: 0,
245            nb_frames: 0,
246            r_frame_rate: "24/1".into(),
247            avg_frame_rate: "24/1".into(),
248            sample_rate: 0,
249            channels: 0,
250            channel_layout: String::new(),
251            bits_per_raw_sample: 8,
252        }
253    }
254
255    #[test]
256    fn test_bit_depth_from_pix_fmt() {
257        let mut stream = base_stream();
258        stream.pix_fmt = "yuv420p10le".into();
259        assert_eq!(bit_depth(&stream), 10);
260    }
261
262    #[test]
263    fn test_source_format_high_bit_depth() {
264        let mut stream = base_stream();
265        stream.pix_fmt = "yuv420p10le".into();
266        let format = SourceFormat::from_stream(&stream);
267        assert_eq!(format.bit_depth, 10);
268        assert_eq!(format.pix_fmt, "yuv420p10le");
269        assert!(format.is_high_bit_depth());
270    }
271
272    #[test]
273    fn test_encode_color_args_x265_10bit_hdr() {
274        let mut stream = base_stream();
275        stream.pix_fmt = "yuv420p10le".into();
276        stream.color_transfer = "smpte2084".into();
277        stream.color_primaries = "bt2020".into();
278        stream.color_space = "bt2020nc".into();
279        let format = SourceFormat::from_stream(&stream);
280        let args = encode_color_args(Codec::X265, &format);
281        assert!(args.windows(2).any(|w| w[0] == "-pix_fmt" && w[1] == "yuv420p10le"));
282        assert!(args.iter().any(|a| a.contains("profile=main10")));
283        assert!(args.iter().any(|a| a.contains("transfer=smpte2084")));
284    }
285
286    #[test]
287    fn test_psnr_peak_scaling() {
288        assert_eq!(psnr_peak(8), 255.0);
289        assert_eq!(psnr_peak(10), 1023.0);
290    }
291
292    #[test]
293    fn test_encode_color_args_emits_hdr10_metadata() {
294        use crate::{Hdr10Metadata, MasteringDisplay};
295        let mut stream = base_stream();
296        stream.pix_fmt = "yuv420p10le".into();
297        stream.color_transfer = "smpte2084".into();
298        stream.color_primaries = "bt2020".into();
299        stream.color_space = "bt2020nc".into();
300        let mut format = SourceFormat::from_stream(&stream);
301        format.hdr10 = Some(Hdr10Metadata {
302            mastering_display: Some(MasteringDisplay {
303                green_x: 13250,
304                green_y: 34500,
305                blue_x: 7500,
306                blue_y: 3000,
307                red_x: 34000,
308                red_y: 16000,
309                white_x: 15635,
310                white_y: 16450,
311                max_luminance: 10_000_000,
312                min_luminance: 50,
313            }),
314            max_cll: Some(1000),
315            max_fall: Some(400),
316        });
317        let args = encode_color_args(Codec::X265, &format);
318        let params = args
319            .windows(2)
320            .find(|w| w[0] == "-x265-params")
321            .map(|w| w[1].clone())
322            .expect("x265-params present");
323        assert!(
324            params.contains("master-display=G(13250,34500)B(7500,3000)R(34000,16000)WP(15635,16450)L(10000000,50)"),
325            "got: {params}"
326        );
327        assert!(params.contains("max-cll=1000,400"), "got: {params}");
328    }
329
330    #[test]
331    fn test_encode_color_args_svtav1_hdr10_metadata() {
332        use crate::{Hdr10Metadata, MasteringDisplay};
333        let mut stream = base_stream();
334        stream.pix_fmt = "yuv420p10le".into();
335        stream.color_transfer = "smpte2084".into();
336        stream.color_primaries = "bt2020".into();
337        stream.color_space = "bt2020nc".into();
338        let mut format = SourceFormat::from_stream(&stream);
339        format.hdr10 = Some(Hdr10Metadata {
340            mastering_display: Some(MasteringDisplay {
341                green_x: 13250,
342                green_y: 34500,
343                blue_x: 7500,
344                blue_y: 3000,
345                red_x: 34000,
346                red_y: 16000,
347                white_x: 15635,
348                white_y: 16450,
349                max_luminance: 10_000_000,
350                min_luminance: 50,
351            }),
352            max_cll: Some(1000),
353            max_fall: Some(400),
354        });
355        let args = encode_color_args(Codec::SvtAv1, &format);
356        let params = args
357            .windows(2)
358            .find(|w| w[0] == "-svtav1-params")
359            .map(|w| w[1].clone())
360            .expect("svtav1-params present");
361        assert!(
362            params.contains("mastering-display=G(0.265,0.69)B(0.15,0.06)R(0.68,0.32)WP(0.3127,0.329)L(1000,0.005)"),
363            "got: {params}"
364        );
365        assert!(params.contains("content-light=1000,400"), "got: {params}");
366        // AV1 carries colour primaries/transfer via the standard -color_* options.
367        assert!(args.windows(2).any(|w| w[0] == "-color_trc" && w[1] == "smpte2084"));
368    }
369
370    #[test]
371    fn test_encode_color_args_no_hdr10_when_sdr() {
372        let mut stream = base_stream();
373        stream.pix_fmt = "yuv420p10le".into();
374        let format = SourceFormat::from_stream(&stream);
375        let args = encode_color_args(Codec::X265, &format);
376        assert!(!args.iter().any(|a| a.contains("master-display")));
377        assert!(!args.iter().any(|a| a.contains("max-cll")));
378    }
379}
viser_ffmpeg/color.rs

viser_ffmpeg/
color.rs
