1 use {
2 super::{define::h264_nal_type, errors::Mpeg4AvcHevcError},
3 byteorder::BigEndian,
4 bytes::BytesMut,
5 bytesio::{bytes_reader::BytesReader, bytes_writer::BytesWriter},
6 std::vec::Vec,
7 };
8
9 use super::errors::MpegErrorValue;
10 use h264_decoder::sps::SpsParser;
11
12 const H264_START_CODE: [u8; 4] = [0x00, 0x00, 0x00, 0x01];
13
14 #[derive(Clone, Default)]
15 pub struct Sps {
16 // pub size: u16,
17 pub data: BytesMut,
18 }
19
20 impl Sps {
new() -> Self21 pub fn new() -> Self {
22 Self {
23 // size: 0,
24 data: BytesMut::new(),
25 }
26 }
len(&self) -> usize27 pub fn len(&self) -> usize {
28 self.data.len()
29 }
30
is_empty(&self) -> bool31 pub fn is_empty(&self) -> bool {
32 self.len() == 0
33 }
34 }
35
36 #[derive(Clone, Default)]
37 pub struct Pps {
38 // pub size: u16,
39 pub data: BytesMut,
40 }
41
42 impl Pps {
new() -> Self43 pub fn new() -> Self {
44 Self {
45 // size: 0,
46 data: BytesMut::new(),
47 }
48 }
len(&self) -> usize49 pub fn len(&self) -> usize {
50 self.data.len()
51 }
is_empty(&self) -> bool52 pub fn is_empty(&self) -> bool {
53 self.len() == 0
54 }
55 }
56
57 #[derive(Default)]
58 pub struct Mpeg4Avc {
59 pub profile: u8,
60 pub compatibility: u8,
61 pub level: u8,
62 pub nalu_length: u8,
63 pub width: u32,
64 pub height: u32,
65
66 pub nb_sps: u8,
67 pub nb_pps: u8,
68
69 pub sps: Vec<Sps>,
70 pub pps: Vec<Pps>,
71
72 pub sps_annexb_data: BytesWriter, // pice together all the sps data
73 pub pps_annexb_data: BytesWriter, // pice together all the pps data
74
75 //extension
76 pub chroma_format_idc: u8,
77 pub bit_depth_luma_minus8: u8,
78 pub bit_depth_chroma_minus8: u8,
79 // data: Vec<u8>, //[u8; 4 * 1024],
80 // off: i32,
81 }
82
print(data: BytesMut)83 pub fn print(data: BytesMut) {
84 println!("==========={}", data.len());
85 let mut idx = 0;
86 for i in data {
87 print!("{i:02X} ");
88 idx += 1;
89 if idx % 16 == 0 {
90 println!()
91 }
92 }
93
94 println!("===========")
95 }
96
97 impl Mpeg4Avc {
new() -> Self98 pub fn new() -> Self {
99 Self {
100 profile: 0,
101 compatibility: 0,
102 level: 0,
103 nalu_length: 0,
104 width: 0,
105 height: 0,
106
107 nb_pps: 0,
108 nb_sps: 0,
109
110 sps: Vec::new(),
111 pps: Vec::new(),
112
113 sps_annexb_data: BytesWriter::new(),
114 pps_annexb_data: BytesWriter::new(),
115
116 chroma_format_idc: 0,
117 bit_depth_chroma_minus8: 0,
118 bit_depth_luma_minus8: 0,
119 }
120 }
121 }
122
123 #[derive(Default)]
124 pub struct Mpeg4AvcProcessor {
125 pub mpeg4_avc: Mpeg4Avc,
126 }
127
128 impl Mpeg4AvcProcessor {
new() -> Self129 pub fn new() -> Self {
130 Self {
131 mpeg4_avc: Mpeg4Avc::new(),
132 }
133 }
134
clear_sps_data(&mut self)135 pub fn clear_sps_data(&mut self) {
136 self.mpeg4_avc.sps.clear();
137 self.mpeg4_avc.sps_annexb_data.clear();
138 }
139
clear_pps_data(&mut self)140 pub fn clear_pps_data(&mut self) {
141 self.mpeg4_avc.pps.clear();
142 self.mpeg4_avc.pps_annexb_data.clear();
143 }
144
decoder_configuration_record_load( &mut self, bytes_reader: &mut BytesReader, ) -> Result<&mut Self, Mpeg4AvcHevcError>145 pub fn decoder_configuration_record_load(
146 &mut self,
147 bytes_reader: &mut BytesReader,
148 ) -> Result<&mut Self, Mpeg4AvcHevcError> {
149 /*version */
150 bytes_reader.read_u8()?;
151 /*avc profile*/
152 self.mpeg4_avc.profile = bytes_reader.read_u8()?;
153 /*avc compatibility*/
154 self.mpeg4_avc.compatibility = bytes_reader.read_u8()?;
155 /*avc level*/
156 self.mpeg4_avc.level = bytes_reader.read_u8()?;
157 /*nalu length*/
158 self.mpeg4_avc.nalu_length = (bytes_reader.read_u8()? & 0x03) + 1;
159
160 /*number of SPS NALUs */
161 self.mpeg4_avc.nb_sps = bytes_reader.read_u8()? & 0x1F;
162
163 if self.mpeg4_avc.nb_sps > 0 {
164 self.clear_sps_data();
165 }
166
167 for i in 0..self.mpeg4_avc.nb_sps as usize {
168 /*SPS size*/
169 let sps_data_size = bytes_reader.read_u16::<BigEndian>()?;
170 let sps_data = Sps {
171 // size: sps_data_size,
172 /*SPS data*/
173 data: bytes_reader.read_bytes(sps_data_size as usize)?,
174 };
175
176 let mut sps_reader = BytesReader::new(sps_data.clone().data);
177 /*parse SPS data to get video resolution(widthxheight) */
178 let nal_type = sps_reader.read_u8()?;
179 if (nal_type & 0x1f) != h264_nal_type::H264_NAL_SPS {
180 return Err(Mpeg4AvcHevcError {
181 value: MpegErrorValue::SPSNalunitTypeNotCorrect,
182 });
183 }
184 let mut sps_parser = SpsParser::new(sps_reader);
185 (self.mpeg4_avc.width, self.mpeg4_avc.height) = sps_parser.parse()?;
186
187 log::info!("mpeg4 avc profile: {}", self.mpeg4_avc.profile);
188 log::info!("mpeg4 avc compatibility: {}", self.mpeg4_avc.compatibility);
189 log::info!("mpeg4 avc level: {}", self.mpeg4_avc.level);
190 log::info!(
191 "mpeg4 avc resolution: {}x{}",
192 self.mpeg4_avc.width,
193 self.mpeg4_avc.height
194 );
195
196 self.mpeg4_avc.sps.push(sps_data);
197 self.mpeg4_avc.sps_annexb_data.write(&H264_START_CODE)?;
198 self.mpeg4_avc
199 .sps_annexb_data
200 .write(&self.mpeg4_avc.sps[i].data[..])?;
201 }
202 /*number of PPS NALUs*/
203 self.mpeg4_avc.nb_pps = bytes_reader.read_u8()?;
204
205 if self.mpeg4_avc.nb_pps > 0 {
206 self.clear_pps_data();
207 }
208
209 for i in 0..self.mpeg4_avc.nb_pps as usize {
210 let pps_data_size = bytes_reader.read_u16::<BigEndian>()?;
211 let pps_data = Pps {
212 // size: pps_data_size,
213 data: bytes_reader.read_bytes(pps_data_size as usize)?,
214 };
215
216 self.mpeg4_avc.pps.push(pps_data);
217 self.mpeg4_avc.pps_annexb_data.write(&H264_START_CODE)?;
218 self.mpeg4_avc
219 .pps_annexb_data
220 .write(&self.mpeg4_avc.pps[i].data[..])?;
221 }
222 /*clear the left bytes*/
223 bytes_reader.extract_remaining_bytes();
224
225 Ok(self)
226 }
227 //https://stackoverflow.com/questions/28678615/efficiently-insert-or-replace-multiple-elements-in-the-middle-or-at-the-beginnin
h264_mp4toannexb( &mut self, bytes_reader: &mut BytesReader, ) -> Result<BytesMut, Mpeg4AvcHevcError>228 pub fn h264_mp4toannexb(
229 &mut self,
230 bytes_reader: &mut BytesReader,
231 ) -> Result<BytesMut, Mpeg4AvcHevcError> {
232 let mut bytes_writer = BytesWriter::new();
233
234 let mut sps_pps_flag = false;
235 while !bytes_reader.is_empty() {
236 let size = self.read_nalu_size(bytes_reader)?;
237 let nalu_type = bytes_reader.advance_u8()? & 0x1f;
238
239 match nalu_type {
240 h264_nal_type::H264_NAL_PPS | h264_nal_type::H264_NAL_SPS => {
241 sps_pps_flag = true;
242 }
243 h264_nal_type::H264_NAL_IDR => {
244 if !sps_pps_flag {
245 sps_pps_flag = true;
246
247 bytes_writer
248 .prepend(&self.mpeg4_avc.pps_annexb_data.get_current_bytes()[..])?;
249 bytes_writer
250 .prepend(&self.mpeg4_avc.sps_annexb_data.get_current_bytes()[..])?;
251 }
252 }
253 _ => {}
254 }
255
256 bytes_writer.write(&H264_START_CODE)?;
257 let data = bytes_reader.read_bytes(size as usize)?;
258 bytes_writer.write(&data[..])?;
259 }
260
261 Ok(bytes_writer.extract_current_bytes())
262 }
263
read_nalu_size(&mut self, bytes_reader: &mut BytesReader) -> Result<u32, Mpeg4AvcHevcError>264 pub fn read_nalu_size(&mut self, bytes_reader: &mut BytesReader) -> Result<u32, Mpeg4AvcHevcError> {
265 let mut size: u32 = 0;
266
267 for _ in 0..self.mpeg4_avc.nalu_length {
268 size = bytes_reader.read_u8()? as u32 + (size << 8);
269 }
270 Ok(size)
271 }
272
write_nalu_size( &mut self, writer: &mut BytesWriter, length: usize, ) -> Result<(), Mpeg4AvcHevcError>273 pub fn write_nalu_size(
274 &mut self,
275 writer: &mut BytesWriter,
276 length: usize,
277 ) -> Result<(), Mpeg4AvcHevcError> {
278 let nalu_length = self.mpeg4_avc.nalu_length;
279 for i in 0..nalu_length {
280 let shift = (nalu_length - i - 1) * 8;
281 let num = ((length >> shift) & 0xFF) as u8;
282 writer.write_u8(num)?;
283 }
284 Ok(())
285 }
286
nalus_to_mpeg4avc(&mut self, nalus: Vec<BytesMut>) -> Result<BytesMut, Mpeg4AvcHevcError>287 pub fn nalus_to_mpeg4avc(&mut self, nalus: Vec<BytesMut>) -> Result<BytesMut, Mpeg4AvcHevcError> {
288 let mut bytes_writer = BytesWriter::new();
289
290 for nalu in nalus {
291 let length = nalu.len();
292 self.write_nalu_size(&mut bytes_writer, length)?;
293 bytes_writer.write(&nalu)?;
294 }
295
296 Ok(bytes_writer.extract_current_bytes())
297 }
298
decoder_configuration_record_save(&mut self) -> Result<BytesMut, Mpeg4AvcHevcError>299 pub fn decoder_configuration_record_save(&mut self) -> Result<BytesMut, Mpeg4AvcHevcError> {
300 let mut bytes_writer = BytesWriter::new();
301
302 bytes_writer.write_u8(1)?;
303 bytes_writer.write_u8(self.mpeg4_avc.profile)?;
304 bytes_writer.write_u8(self.mpeg4_avc.compatibility)?;
305 bytes_writer.write_u8(self.mpeg4_avc.level)?;
306 bytes_writer.write_u8((self.mpeg4_avc.nalu_length - 1) | 0xFC)?;
307
308 //sps
309 bytes_writer.write_u8(self.mpeg4_avc.nb_sps | 0xE0)?;
310 for i in 0..self.mpeg4_avc.nb_sps as usize {
311 bytes_writer.write_u16::<BigEndian>(self.mpeg4_avc.sps[i].len() as u16)?;
312 bytes_writer.write(&self.mpeg4_avc.sps[i].data[..])?;
313 }
314
315 //pps
316 bytes_writer.write_u8(self.mpeg4_avc.nb_pps)?;
317 for i in 0..self.mpeg4_avc.nb_pps as usize {
318 bytes_writer.write_u16::<BigEndian>(self.mpeg4_avc.pps[i].len() as u16)?;
319 bytes_writer.write(&self.mpeg4_avc.pps[i].data[..])?
320 }
321
322 match self.mpeg4_avc.profile {
323 100 | 110 | 122 | 244 | 44 | 83 | 86 | 118 | 128 | 138 | 139 | 134 => {
324 bytes_writer.write_u8(0xFC | self.mpeg4_avc.chroma_format_idc)?;
325 bytes_writer.write_u8(0xF8 | self.mpeg4_avc.bit_depth_luma_minus8)?;
326 bytes_writer.write_u8(0xF8 | self.mpeg4_avc.bit_depth_chroma_minus8)?;
327 bytes_writer.write_u8(0)?;
328 }
329 _ => {}
330 }
331
332 Ok(bytes_writer.extract_current_bytes())
333 }
334 }
335
336 #[cfg(test)]
337 mod tests {
338 use bytes::BytesMut;
339 use bytesio::{bytes_reader::BytesReader, bytes_writer::BytesWriter};
340
341 #[test]
test_bytes_to_bigend()342 fn test_bytes_to_bigend() {
343 let mut size: u32 = 0;
344 let mut b = BytesMut::new();
345 b.extend_from_slice(b"\0\0\x03\xe8");
346 let mut bytes_reader = BytesReader::new(b);
347
348 for _ in 0..4 {
349 size = bytes_reader.read_u8().unwrap() as u32 + (size << 8);
350 }
351 println!("size: {size}");
352 }
353 #[test]
test_bigend_to_bytes()354 fn test_bigend_to_bytes() {
355 let size = 1000;
356 let length = 4;
357 let mut bytes_writer = BytesWriter::new();
358
359 for i in 0..length {
360 let shift = (length - i - 1) * 8;
361 let num = ((size >> shift) & 0xFF) as u8;
362 bytes_writer.write_u8(num).unwrap();
363 }
364 println!("num: {:?}", bytes_writer.extract_current_bytes());
365 }
366 }
367