codec/pixel_format/
mod.rs

1//! Pixel-format detection from codec sequence headers.
2//!
3//! Given raw bitstream samples (the same Vec<Vec<u8>> our decoders
4//! consume), parse just enough of the first sequence header to
5//! extract chroma subsampling + luma bit depth, then map to our
6//! PixelFormat enum.
7//!
8//! Why not use the full decoder: our CPU decoders (H.264 openh264,
9//! HEVC Rust, VP9 Rust, rav1d AV1) each have their own parser
10//! entry points, but none of them expose a "just probe the format"
11//! API. NVDEC's sequence_callback tells us, but only after decode
12//! starts. This module gives the pipeline a fast, codec-agnostic
13//! probe path that runs before decoder construction.
14
15use crate::frame::PixelFormat;
16
17mod bitreader;
18mod h264;
19mod hevc;
20mod av1;
21mod mpeg2;
22
23#[cfg(test)]
24mod tests;
25
26pub use h264::*;
27pub use hevc::*;
28pub use av1::*;
29pub use mpeg2::*;
30
31/// Detect pixel format from the first sequence header in `samples`.
32/// Falls back to Yuv420p on any parse failure — that matches the
33/// previous hard-coded behavior so a bad probe doesn't block the
34/// transcode, just the probe payload accuracy.
35pub fn detect(codec: &str, samples: &[Vec<u8>]) -> PixelFormat {
36    if samples.is_empty() {
37        return PixelFormat::Yuv420p;
38    }
39
40    let result = match codec.to_lowercase().as_str() {
41        "h264" | "avc1" | "avc" => h264::detect_h264(&samples[0]),
42        "h265" | "hevc" | "hvc1" | "hev1" => hevc::detect_hevc(&samples[0]),
43        "vp9" | "vp09" => detect_vp9(&samples[0]),
44        "av1" | "av01" => av1::detect_av1(&samples[0]),
45        _ => None,
46    };
47
48    result.unwrap_or(PixelFormat::Yuv420p)
49}
50
51/// Public entry point — dispatch by codec and return `Some((width,
52/// height))` if the sequence header in `samples[0]` is parseable,
53/// `None` otherwise.
54///
55/// Callers should treat `None` as "keep the existing width/height" —
56/// it's load-bearing for MPEG-TS where `StreamInfo` would otherwise
57/// carry `0×0`, but a parse failure on MP4/MKV (which already have
58/// width/height in the sample-entry / track-header) is a no-op.
59pub fn detect_dims(codec: &str, samples: &[Vec<u8>]) -> Option<(u32, u32)> {
60    if samples.is_empty() {
61        return None;
62    }
63    let sample = &samples[0];
64    match codec.to_lowercase().as_str() {
65        "h264" | "avc1" | "avc" | "avc3" => {
66            let info = parse_h264_sps(sample)?;
67            Some((info.width?, info.height?))
68        }
69        "h265" | "hevc" | "hvc1" | "hev1" | "hvc2" | "hev2" => {
70            let info = parse_hevc_sps(sample)?;
71            Some((info.width?, info.height?))
72        }
73        "mpeg2" | "mpeg2video" | "mp2v" => {
74            let info = parse_mpeg2_sequence_header(sample)?;
75            Some((info.width, info.height))
76        }
77        _ => None,
78    }
79}
80
81// ─── VP9 uncompressed-header pixel-format detection ────────────────
82// Private — called only by detect() above. Parses just enough of the
83// VP9 frame header to derive chroma subsampling + bit depth.
84
85fn detect_vp9(sample: &[u8]) -> Option<PixelFormat> {
86    if sample.len() < 2 {
87        return None;
88    }
89    let mut br = bitreader::BitReader::new(sample);
90    let frame_marker = br.read_bits(2)?;
91    if frame_marker != 2 {
92        return None;
93    }
94    let profile_low = br.read_bits(1)?;
95    let profile_high = br.read_bits(1)?;
96    let profile = (profile_high << 1) | profile_low;
97    if profile == 3 {
98        let _reserved_zero = br.read_bits(1)?;
99    }
100    let show_existing_frame = br.read_bits(1)?;
101    if show_existing_frame == 1 {
102        return None;
103    }
104    let frame_type = br.read_bits(1)?;
105    let _show_frame = br.read_bits(1)?;
106    let _error_resilient = br.read_bits(1)?;
107
108    // color_config only appears on keyframes.
109    if frame_type != 0 {
110        return None;
111    }
112
113    // Keyframe sync code: 3 bytes {0x49, 0x83, 0x42}. 24 bits.
114    let sync = br.read_bits(24)?;
115    if sync != 0x498342 {
116        return None;
117    }
118
119    let bit_depth = if profile >= 2 {
120        if br.read_bits(1)? == 0 { 10 } else { 12 }
121    } else {
122        8
123    };
124    let _color_space = br.read_bits(3)?;
125    // color_range + subsampling — layout depends on color_space
126    // For simplicity: for Profile 0/2 the subsampling is 4:2:0. Profile
127    // 1/3 read subsampling_x/y fields to distinguish 4:2:2 vs 4:4:4.
128    let (sx, sy) = if profile == 1 || profile == 3 {
129        let _color_range = br.read_bits(1)?;
130        let sx = br.read_bits(1)?;
131        let sy = br.read_bits(1)?;
132        (sx, sy)
133    } else {
134        (1, 1) // 4:2:0
135    };
136
137    let chroma_idc = match (sx, sy) {
138        (1, 1) => 1, // 4:2:0
139        (1, 0) => 2, // 4:2:2
140        (0, 0) => 3, // 4:4:4
141        _ => 1,
142    };
143
144    Some(PixelFormat::from_chroma_and_depth(chroma_idc, bit_depth))
145}
codec/pixel_format/mod.rs

codec/pixel_format/
mod.rs
