Skip to main content

pito_codecs/hevc/
hvcc.rs

1use crate::media::{Error, Result};
2use crate::syntax::bits::{annexb_nals, unescape_rbsp};
3use crate::syntax::hevc::params::Sps;
4use crate::syntax::hevc::slice::nal_type;
5
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct Hvcc {
8    pub profile: [u8; 12],
9    pub chroma_format: u8,
10    pub bit_depth_luma: u8,
11    pub bit_depth_chroma: u8,
12    pub temporal_layers: u8,
13    pub temporal_id_nested: bool,
14    pub nal_length_size: usize,
15    pub width: u32,
16    pub height: u32,
17    pub vps: Vec<Vec<u8>>,
18    pub sps: Vec<Vec<u8>>,
19    pub pps: Vec<Vec<u8>>,
20    pub sei: Vec<Vec<u8>>,
21}
22
23fn invalid(msg: &str) -> Error {
24    Error::Invalid(format!("hvcC: {msg}"))
25}
26
27impl Hvcc {
28    pub fn parse(data: &[u8]) -> Result<Hvcc> {
29        if data.len() < 23 {
30            return Err(invalid("too short"));
31        }
32        if data[0] != 1 {
33            return Err(Error::Unsupported(format!("hvcC version {}", data[0])));
34        }
35        let mut profile = [0u8; 12];
36        profile.copy_from_slice(&data[1..13]);
37        let mut config = Hvcc {
38            profile,
39            chroma_format: data[16] & 3,
40            bit_depth_luma: (data[17] & 7) + 8,
41            bit_depth_chroma: (data[18] & 7) + 8,
42            temporal_layers: (data[21] >> 3) & 7,
43            temporal_id_nested: data[21] & 4 != 0,
44            nal_length_size: (data[21] & 3) as usize + 1,
45            width: 0,
46            height: 0,
47            vps: Vec::new(),
48            sps: Vec::new(),
49            pps: Vec::new(),
50            sei: Vec::new(),
51        };
52        let arrays = data[22] as usize;
53        let mut pos = 23;
54        for _ in 0..arrays {
55            let head = data.get(pos..pos + 3).ok_or_else(|| invalid("truncated"))?;
56            let kind = head[0] & 0x3f;
57            let count = u16::from_be_bytes([head[1], head[2]]) as usize;
58            pos += 3;
59            for _ in 0..count {
60                let len = data.get(pos..pos + 2).ok_or_else(|| invalid("truncated"))?;
61                let len = u16::from_be_bytes([len[0], len[1]]) as usize;
62                pos += 2;
63                let nal = data
64                    .get(pos..pos + len)
65                    .ok_or_else(|| invalid("truncated"))?
66                    .to_vec();
67                pos += len;
68                config.push(kind, nal);
69            }
70        }
71        config.read_size();
72        Ok(config)
73    }
74
75    pub fn from_annexb(stream: &[u8]) -> Result<Hvcc> {
76        let mut config = Hvcc {
77            profile: [0; 12],
78            chroma_format: 1,
79            bit_depth_luma: 8,
80            bit_depth_chroma: 8,
81            temporal_layers: 1,
82            temporal_id_nested: false,
83            nal_length_size: 4,
84            width: 0,
85            height: 0,
86            vps: Vec::new(),
87            sps: Vec::new(),
88            pps: Vec::new(),
89            sei: Vec::new(),
90        };
91        for nal in annexb_nals(stream) {
92            let Some(&first) = nal.first() else { continue };
93            let kind = (first >> 1) & 0x3f;
94            if !(nal_type::VPS..=nal_type::PPS).contains(&kind) && kind != nal_type::PREFIX_SEI {
95                continue;
96            }
97            let list = config.list(kind);
98            if list.is_some_and(|l| !l.iter().any(|n| n == nal)) {
99                config.push(kind, nal.to_vec());
100            }
101        }
102        let sps = config
103            .sps
104            .first()
105            .ok_or_else(|| invalid("no SPS in the stream"))?;
106        let rbsp = unescape_rbsp(sps.get(2..).unwrap_or_default());
107        let head = rbsp.get(..13).ok_or_else(|| invalid("SPS too short"))?;
108        config.profile.copy_from_slice(&head[1..13]);
109        config.temporal_layers = ((head[0] >> 1) & 7) + 1;
110        config.temporal_id_nested = head[0] & 1 != 0;
111        let parsed = Sps::parse(&rbsp).map_err(|e| e.media("hevc"))?;
112        config.chroma_format = parsed.chroma_format_idc as u8;
113        config.bit_depth_luma = parsed.bit_depth_luma as u8;
114        config.bit_depth_chroma = parsed.bit_depth_chroma as u8;
115        config.read_size();
116        Ok(config)
117    }
118
119    pub fn to_bytes(&self) -> Vec<u8> {
120        let mut out = vec![1];
121        out.extend_from_slice(&self.profile);
122        out.extend_from_slice(&[0xf0, 0x00, 0xfc]);
123        out.push(0xfc | (self.chroma_format & 3));
124        out.push(0xf8 | (self.bit_depth_luma.saturating_sub(8) & 7));
125        out.push(0xf8 | (self.bit_depth_chroma.saturating_sub(8) & 7));
126        out.extend_from_slice(&[0, 0]);
127        let size = (self.nal_length_size.clamp(1, 4) - 1) as u8;
128        out.push(((self.temporal_layers & 7) << 3) | ((self.temporal_id_nested as u8) << 2) | size);
129        let arrays = [
130            (nal_type::VPS, &self.vps),
131            (nal_type::SPS, &self.sps),
132            (nal_type::PPS, &self.pps),
133            (nal_type::PREFIX_SEI, &self.sei),
134        ];
135        let present: Vec<_> = arrays.iter().filter(|(_, l)| !l.is_empty()).collect();
136        out.push(present.len() as u8);
137        for (kind, list) in present {
138            out.push(0x80 | kind);
139            out.extend_from_slice(&(list.len() as u16).to_be_bytes());
140            for nal in list.iter() {
141                out.extend_from_slice(&(nal.len() as u16).to_be_bytes());
142                out.extend_from_slice(nal);
143            }
144        }
145        out
146    }
147
148    pub fn annexb(&self) -> Vec<u8> {
149        let mut out = Vec::new();
150        for nal in self
151            .vps
152            .iter()
153            .chain(&self.sps)
154            .chain(&self.pps)
155            .chain(&self.sei)
156        {
157            out.extend_from_slice(&[0, 0, 0, 1]);
158            out.extend_from_slice(nal);
159        }
160        out
161    }
162
163    pub fn profile_idc(&self) -> u8 {
164        self.profile[0] & 0x1f
165    }
166
167    pub fn level_idc(&self) -> u8 {
168        self.profile[11]
169    }
170
171    fn list(&self, kind: u8) -> Option<&Vec<Vec<u8>>> {
172        match kind {
173            nal_type::VPS => Some(&self.vps),
174            nal_type::SPS => Some(&self.sps),
175            nal_type::PPS => Some(&self.pps),
176            nal_type::PREFIX_SEI => Some(&self.sei),
177            _ => None,
178        }
179    }
180
181    fn push(&mut self, kind: u8, nal: Vec<u8>) {
182        match kind {
183            nal_type::VPS => self.vps.push(nal),
184            nal_type::SPS => self.sps.push(nal),
185            nal_type::PPS => self.pps.push(nal),
186            nal_type::PREFIX_SEI => self.sei.push(nal),
187            _ => {}
188        }
189    }
190
191    fn read_size(&mut self) {
192        let Some(sps) = self.sps.first() else { return };
193        let Ok(parsed) = Sps::parse(&unescape_rbsp(sps.get(2..).unwrap_or_default())) else {
194            return;
195        };
196        let (_, _, w, h) = parsed.crop_rect();
197        self.width = w;
198        self.height = h;
199    }
200}