h264_stream/
lib.rs

Help
1// Copyright 2019 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use anyhow::ensure;
6use byteorder::{BigEndian, WriteBytesExt};
7use fidl_fuchsia_media::FormatDetails;
8use std::path::Path;
9use std::{fs, mem};
10use stream_processor_test::*;
11
12const SEPARATE_SPS_PPS: bool = true;
13// written to vector as big-endian
14const NAL_START_CODE: u32 = 1;
15
16/// Represents an H264 elementary stream.
17pub struct H264Stream {
18    data: Vec<u8>,
19}
20
21impl H264Stream {
22    /// Adds data to raw H264 elementary stream
23    pub fn append(&mut self, data: &mut Vec<u8>) {
24        self.data.append(data);
25    }
26
27    /// Constructs an H264 elementary stream from a file with raw elementary stream data.
28    pub fn from_file(filename: impl AsRef<Path>) -> Result<Self> {
29        Ok(H264Stream::from(fs::read(filename)?))
30    }
31
32    /// Returns an iterator over H264 NALs that does not copy.
33    pub fn nal_iter(&self) -> impl Iterator<Item = H264Nal<'_>> {
34        H264NalIter { data: &self.data, pos: 0 }
35    }
36
37    /// Returns an iterator over H264 chunks (pictures and anything preceding each picture) that
38    /// does not copy.
39    //
40    // TODO(https://fxbug.dev/42084549): Make H264ChunkIter capable of iterating NALs or iterating pictures with any
41    // non-picture NALs in front of each picture prepended (the current behavior).
42    fn chunk_iter(&self) -> impl Iterator<Item = H264Chunk<'_>> {
43        H264ChunkIter::create(&self.data)
44    }
45}
46
47impl From<Vec<u8>> for H264Stream {
48    fn from(data: Vec<u8>) -> Self {
49        Self { data }
50    }
51}
52
53impl ElementaryStream for H264Stream {
54    fn format_details(&self, version_ordinal: u64) -> FormatDetails {
55        FormatDetails {
56            format_details_version_ordinal: Some(version_ordinal),
57            mime_type: Some(String::from("video/h264")),
58            ..Default::default()
59        }
60    }
61
62    fn is_access_units(&self) -> bool {
63        true
64    }
65
66    fn stream<'a>(&'a self) -> Box<dyn Iterator<Item = ElementaryStreamChunk> + 'a> {
67        Box::new(self.chunk_iter().map(|chunk| ElementaryStreamChunk {
68            start_access_unit: true,
69            known_end_access_unit: true,
70            data: chunk.data.to_vec(),
71            significance: match chunk.nal_kind {
72                Some(H264NalKind::IDR) | Some(H264NalKind::NonIDR) => {
73                    Significance::Video(VideoSignificance::Picture)
74                }
75                _ => Significance::Video(VideoSignificance::NotPicture),
76            },
77            timestamp: None,
78        }))
79    }
80}
81
82pub struct H264SeiItuT35 {
83    pub country_code: u8,
84    pub country_code_extension: u8,
85    pub payload: Vec<u8>,
86}
87
88impl H264SeiItuT35 {
89    pub const COUNTRY_CODE_UNITED_STATES: u8 = 0xb5;
90
91    pub fn as_bytes(&self) -> Result<Vec<u8>> {
92        const ITU_T35_PAYLOAD_TYPE: u8 = 4;
93
94        let mut bytes = vec![];
95        bytes.write_u32::<BigEndian>(NAL_START_CODE)?;
96        bytes.write_u8(SEI_CODE)?;
97        bytes.write_u8(ITU_T35_PAYLOAD_TYPE)?;
98        bytes.write_u8(u8::try_from(self.payload_size())?)?;
99        bytes.write_u8(self.country_code)?;
100        bytes.write_u8(self.country_code_extension)?;
101        bytes.append(&mut self.payload.clone());
102        Ok(bytes)
103    }
104
105    fn payload_size(&self) -> usize {
106        // At the time this is called, the country_code and country_code_extension bytes haven't
107        // yet been written.
108        mem::size_of::<u8>() + mem::size_of::<u8>() + self.payload.len()
109    }
110}
111
112/// Represents an AVCC stream.
113pub struct H264AVCCStream {
114    // NAL units with 4-byte length prefixes but no start codes.
115    nal_units: Vec<(H264NalKind, Vec<u8>)>,
116    oob_data: Vec<u8>,
117}
118
119impl H264AVCCStream {
120    pub fn from_annexb_stream(annexb: H264Stream) -> Result<Self> {
121        let mut oob_data = Vec::new();
122
123        let mut found_pps = false;
124        for nal in annexb.nal_iter() {
125            if nal.kind == H264NalKind::SPS {
126                let without_start = remove_nal_start(nal.data);
127                ensure!(oob_data.is_empty(), "Too many SPS NAL units");
128                let profile = without_start[1];
129                let profile_compatibility = without_start[2];
130                let level = without_start[3];
131                let nalu_size_minus1 = 3;
132                let num_sps_nals = 1;
133                oob_data.extend(&[
134                    1,
135                    profile,
136                    profile_compatibility,
137                    level,
138                    0xfc | nalu_size_minus1, // reserved | nalu length size
139                    0xe0 | num_sps_nals,     // reserved | Num  SPS nals
140                ]);
141
142                oob_data.extend((without_start.len() as u16).to_be_bytes());
143                oob_data.extend(without_start);
144            } else if nal.kind == H264NalKind::PPS {
145                ensure!(!found_pps, "Too many PPS NAL units");
146                found_pps = true;
147                let without_start = remove_nal_start(nal.data);
148                let num_pps_nals = 1;
149                oob_data.push(num_pps_nals);
150                oob_data.extend((without_start.len() as u16).to_be_bytes());
151                oob_data.extend(without_start);
152            }
153        }
154
155        let nal_units = annexb
156            .nal_iter()
157            .filter_map(|nal_unit| {
158                match nal_unit.kind {
159                    H264NalKind::SPS | H264NalKind::PPS => return None,
160                    _ => (),
161                }
162                let without_start = remove_nal_start(nal_unit.data);
163                let mut chunk_data = Vec::new();
164                chunk_data.extend((without_start.len() as u32).to_be_bytes());
165                chunk_data.extend(without_start);
166                Some((nal_unit.kind, chunk_data))
167            })
168            .collect();
169        Ok(Self { nal_units, oob_data })
170    }
171}
172
173impl ElementaryStream for H264AVCCStream {
174    fn format_details(&self, version_ordinal: u64) -> FormatDetails {
175        FormatDetails {
176            format_details_version_ordinal: Some(version_ordinal),
177            mime_type: Some(String::from("video/h264")),
178            oob_bytes: Some(self.oob_data.clone()),
179            ..Default::default()
180        }
181    }
182
183    fn is_access_units(&self) -> bool {
184        true
185    }
186
187    fn stream<'a>(&'a self) -> Box<dyn Iterator<Item = ElementaryStreamChunk> + 'a> {
188        Box::new(self.nal_units.iter().map(|chunk| ElementaryStreamChunk {
189            start_access_unit: true,
190            known_end_access_unit: true,
191            data: chunk.1.clone(),
192            significance: match chunk.0 {
193                H264NalKind::IDR | H264NalKind::NonIDR => {
194                    Significance::Video(VideoSignificance::Picture)
195                }
196                _ => Significance::Video(VideoSignificance::NotPicture),
197            },
198            timestamp: None,
199        }))
200    }
201}
202
203#[derive(Debug)]
204pub struct H264Nal<'a> {
205    pub kind: H264NalKind,
206    pub data: &'a [u8],
207}
208
209// Returns the NAL unit with a 3-byte or 4-byte start code removed.
210fn remove_nal_start<'a>(input: &'a [u8]) -> &'a [u8] {
211    let start_len = if input[2] == 1 { 3 } else { 4 };
212    assert!(input[start_len - 1] == 1);
213    &input[start_len..]
214}
215
216#[derive(Debug, Clone, Copy, PartialEq)]
217pub enum H264NalKind {
218    NonIDR,
219    IDR,
220    SEI,
221    SPS,
222    PPS,
223    Unknown,
224}
225
226const NON_IDR_PICTURE_CODE: u8 = 1;
227const IDR_PICTURE_CODE: u8 = 5;
228const SEI_CODE: u8 = 6;
229const SPS_CODE: u8 = 7;
230const PPS_CODE: u8 = 8;
231
232impl H264NalKind {
233    fn from_header(header: u8) -> Self {
234        let kind = header & 0x1f;
235        match kind {
236            kind if kind == IDR_PICTURE_CODE => H264NalKind::IDR,
237            kind if kind == NON_IDR_PICTURE_CODE => H264NalKind::NonIDR,
238            kind if kind == SPS_CODE => H264NalKind::SPS,
239            kind if kind == PPS_CODE => H264NalKind::PPS,
240            kind if kind == SEI_CODE => H264NalKind::SEI,
241            _ => H264NalKind::Unknown,
242        }
243    }
244}
245
246struct H264NalStart<'a> {
247    /// Position in the h264 stream of the start.
248    pos: usize,
249    /// All the data from the start of the NAL onward.
250    data: &'a [u8],
251    kind: H264NalKind,
252}
253
254/// An iterator over NALs in an H264 stream.
255struct H264NalIter<'a> {
256    data: &'a [u8],
257    pos: usize,
258}
259
260impl<'a> H264NalIter<'a> {
261    fn next_nal(&self, pos: usize) -> Option<H264Nal<'a>> {
262        // This won't need to search if pos already at a start code.
263        let nal_start = self.next_nal_start(pos)?;
264        // We search 3 bytes after the found nal's start, because that will
265        // ensure we don't just find the same start code again.
266        match self.next_nal_start(nal_start.pos + 3) {
267            Some(next_start) => Some(H264Nal {
268                kind: nal_start.kind,
269                data: &nal_start.data[0..(next_start.pos - nal_start.pos)],
270            }),
271            None => Some(H264Nal { kind: nal_start.kind, data: nal_start.data }),
272        }
273    }
274
275    fn next_nal_start(&self, pos: usize) -> Option<H264NalStart<'a>> {
276        // This search size will find 3 and 4 byte start codes, and the
277        // header value.
278        const NAL_SEARCH_SIZE: usize = 5;
279
280        let data = self.data.get(pos..)?;
281        data.windows(NAL_SEARCH_SIZE).enumerate().find_map(|(i, candidate)| match candidate {
282            [0, 0, 0, 1, h] | [0, 0, 1, h, _] => Some(H264NalStart {
283                pos: pos + i,
284                data: data.get(i..).expect("Getting slice starting where we just matched"),
285                kind: H264NalKind::from_header(*h),
286            }),
287            _ => None,
288        })
289    }
290}
291
292impl<'a> Iterator for H264NalIter<'a> {
293    type Item = H264Nal<'a>;
294    fn next(&mut self) -> Option<Self::Item> {
295        let nal = self.next_nal(self.pos);
296        self.pos += nal.as_ref().map(|n| n.data.len()).unwrap_or(0);
297        nal
298    }
299}
300
301#[derive(Debug)]
302pub struct H264Chunk<'a> {
303    /// Only NonIDR, IDR are used, or not set if last data in file needs to be emitted but IDR or
304    /// NonIDR wasn't last in the file.
305    pub nal_kind: Option<H264NalKind>,
306    pub data: &'a [u8],
307}
308
309/// An iterator over groupings of NALs in an H264 stream.
310struct H264ChunkIter<'a> {
311    nal_iter: H264NalIter<'a>,
312    chunk_start_pos: Option<usize>,
313}
314
315impl<'a> H264ChunkIter<'a> {
316    fn create(data: &'a [u8]) -> H264ChunkIter<'a> {
317        H264ChunkIter { nal_iter: H264NalIter { data: data, pos: 0 }, chunk_start_pos: None }
318    }
319}
320
321impl<'a> Iterator for H264ChunkIter<'a> {
322    type Item = H264Chunk<'a>;
323    fn next(&mut self) -> Option<Self::Item> {
324        let chunk = loop {
325            let maybe_nal = self.nal_iter.next();
326            match maybe_nal {
327                Some(nal) => {
328                    if let None = self.chunk_start_pos {
329                        self.chunk_start_pos = Some(self.nal_iter.pos - nal.data.len());
330                    }
331                    let queue_chunk = match nal.kind {
332                        // picture
333                        H264NalKind::IDR | H264NalKind::NonIDR => true,
334                        // not picture
335                        _ => SEPARATE_SPS_PPS,
336                    };
337                    if queue_chunk {
338                        break Some(H264Chunk {
339                            nal_kind: Some(nal.kind),
340                            data: &self.nal_iter.data
341                                [self.chunk_start_pos.unwrap()..self.nal_iter.pos],
342                        });
343                    } else {
344                        continue;
345                    }
346                }
347                None => {
348                    if self.chunk_start_pos.expect("Zero NALs?") == self.nal_iter.data.len() {
349                        // done
350                        break None;
351                    } else {
352                        break Some(H264Chunk {
353                            nal_kind: None,
354                            data: &self.nal_iter.data[self.chunk_start_pos.unwrap()..],
355                        });
356                    }
357                }
358            }
359        };
360        *self.chunk_start_pos.as_mut().unwrap() +=
361            chunk.as_ref().map(|c| c.data.len()).unwrap_or(0);
362        chunk
363    }
364}
365
366#[cfg(test)]
367mod test {
368    use super::*;
369
370    #[fuchsia::test]
371    fn test_nal_iter() {
372        let sps_pps_idr = [
373            0x00, 0x00, 0x00, 0x01, 0x67, 0x4d, 0x00, 0x28, 0xf4, 0x03, 0xc0, 0x11, 0x3f, 0x2a,
374            0x00, 0x00, 0x00, 0x01, 0x68, 0xee, 0x38, 0x80, 0x00, 0x00, 0x00, 0x01, 0x65, 0x88,
375            0x84, 0x01,
376        ];
377
378        let iter = H264NalIter { data: &sps_pps_idr, pos: 0 };
379
380        let result: Vec<H264Nal<'_>> = iter.collect();
381
382        assert_eq!(result[0].kind, H264NalKind::SPS);
383        assert_eq!(result[1].kind, H264NalKind::PPS);
384        assert_eq!(result[2].kind, H264NalKind::IDR);
385    }
386
387    #[fuchsia::test]
388    fn test_chunk_iter() {
389        let sps_pps_idr = [
390            0x00, 0x00, 0x00, 0x01, 0x67, 0x4d, 0x00, 0x28, 0xf4, 0x03, 0xc0, 0x11, 0x3f, 0x2a,
391            0x00, 0x00, 0x00, 0x01, 0x68, 0xee, 0x38, 0x80, 0x00, 0x00, 0x00, 0x01, 0x65, 0x88,
392            0x84, 0x01,
393        ];
394
395        let iter = H264ChunkIter::create(&sps_pps_idr);
396
397        let result: Vec<H264Chunk<'_>> = iter.collect();
398
399        if SEPARATE_SPS_PPS {
400            assert_eq!(result[0].nal_kind.expect("nal"), H264NalKind::SPS);
401            assert_eq!(result[1].nal_kind.expect("nal"), H264NalKind::PPS);
402            assert_eq!(result[2].nal_kind.expect("nal"), H264NalKind::IDR);
403        } else {
404            assert_eq!(result[0].nal_kind.expect("nal"), H264NalKind::IDR);
405        }
406    }
407}
h264_stream/lib.rs

h264_stream/
lib.rs
