Skip to main content

vk_video/
parser.rs

1mod au_splitter;
2mod nalu_parser;
3mod nalu_splitter;
4
5#[cfg(vulkan)]
6pub(crate) mod decoder_instructions;
7#[cfg(vulkan)]
8pub(crate) mod reference_manager;
9
10pub mod h264 {
11    use super::au_splitter::AUSplitter;
12    use super::nalu_parser::NalReceiver;
13    use super::nalu_splitter::NALUSplitter;
14    use h264_reader::annexb::AnnexBReader;
15    use h264_reader::push::NalAccumulator;
16    use std::sync::mpsc;
17
18    pub use super::au_splitter::AccessUnit;
19    pub use super::nalu_parser::{Nalu, ParsedNalu};
20    #[cfg(feature = "expose-parsers")]
21    pub use h264_reader::nal as nal_types;
22
23    #[derive(Debug, thiserror::Error)]
24    pub enum H264ParserError {
25        #[error("Bitstreams that allow gaps in frame_num are not supported")]
26        GapsInFrameNumNotSupported,
27
28        #[error("Streams containing fields instead of frames are not supported")]
29        FieldsNotSupported,
30
31        #[error("Error while parsing a NAL header: {0:?}")]
32        NalHeaderParseError(h264_reader::nal::NalHeaderError),
33
34        #[error("Error while parsing SPS: {0:?}")]
35        SpsParseError(h264_reader::nal::sps::SpsError),
36
37        #[error("Error while parsing PPS: {0:?}")]
38        PpsParseError(h264_reader::nal::pps::PpsError),
39
40        #[error("Error while parsing a slice: {0:?}")]
41        SliceParseError(h264_reader::nal::slice::SliceHeaderError),
42    }
43
44    /// H264 parser for Annex B format
45    pub struct H264Parser {
46        reader: AnnexBReader<NalAccumulator<NalReceiver>>,
47        receiver: mpsc::Receiver<Result<ParsedNalu, H264ParserError>>,
48        nalu_splitter: NALUSplitter,
49        au_splitter: AUSplitter,
50    }
51
52    impl Default for H264Parser {
53        fn default() -> Self {
54            let (tx, rx) = mpsc::channel();
55
56            H264Parser {
57                reader: AnnexBReader::accumulate(NalReceiver::new(tx)),
58                receiver: rx,
59                nalu_splitter: NALUSplitter::default(),
60                au_splitter: AUSplitter::default(),
61            }
62        }
63    }
64
65    impl H264Parser {
66        /// Parses nalus in Annex B format.
67        /// Returns [`AccessUnit`]s representing whole frame
68        pub fn parse(
69            &mut self,
70            bytes: &[u8],
71            pts: Option<u64>,
72        ) -> Result<Vec<AccessUnit>, H264ParserError> {
73            let nalus = self.nalu_splitter.push(bytes, pts);
74            let nalus = nalus.into_iter().map(|(nalu_bytes, pts)| {
75                self.reader.push(&nalu_bytes);
76
77                let parsed_nalu = self.receiver.try_recv().unwrap();
78                parsed_nalu.map(|parsed_nalu| Nalu {
79                    parsed: parsed_nalu,
80                    raw_bytes: nalu_bytes.into_boxed_slice(),
81                    pts,
82                })
83            });
84
85            let mut access_units = Vec::new();
86            for nalu in nalus {
87                let nalu = nalu?;
88
89                let Some(au) = self.au_splitter.put_nalu(nalu) else {
90                    continue;
91                };
92
93                access_units.push(au);
94            }
95
96            Ok(access_units)
97        }
98
99        pub fn flush(&mut self) -> Result<Vec<AccessUnit>, H264ParserError> {
100            let nalus = self.nalu_splitter.flush();
101            let nalus = nalus.into_iter().map(|(nalu_bytes, pts)| {
102                self.reader.push(&nalu_bytes);
103
104                let parsed_nalu = self.receiver.try_recv().unwrap();
105                parsed_nalu.map(|parsed_nalu| Nalu {
106                    parsed: parsed_nalu,
107                    raw_bytes: nalu_bytes.into_boxed_slice(),
108                    pts,
109                })
110            });
111
112            let mut access_units = Vec::new();
113            for nalu in nalus {
114                let nalu = nalu?;
115
116                let Some(au) = self.au_splitter.put_nalu(nalu) else {
117                    continue;
118                };
119
120                access_units.push(au);
121            }
122
123            if let Some(au) = self.au_splitter.flush() {
124                access_units.push(au);
125            }
126
127            Ok(access_units)
128        }
129    }
130}