hevc_parser/
utils.rs

1use anyhow::{Result, anyhow};
2use bitvec_helpers::bitstream_io_writer::BitstreamIoWriter;
3
4use super::{Frame, NAL_AUD, NALUStartCode};
5
6pub fn clear_start_code_emulation_prevention_3_byte(data: &[u8]) -> Vec<u8> {
7    let len = data.len();
8
9    if len > 2 {
10        let mut unescaped_bytes: Vec<u8> = Vec::with_capacity(len);
11        unescaped_bytes.push(data[0]);
12        unescaped_bytes.push(data[1]);
13
14        for i in 2..len {
15            if !(data[i - 2] == 0 && data[i - 1] == 0 && data[i] == 3) {
16                unescaped_bytes.push(data[i]);
17            }
18        }
19
20        unescaped_bytes
21    } else {
22        data.to_owned()
23    }
24}
25
26/// Within the NAL unit, the following three-byte sequences shall not occur at any byte-aligned position:
27///   - 0x000000
28///   - 0x000001
29///   - 0x000002
30pub fn add_start_code_emulation_prevention_3_byte(data: &mut Vec<u8>) {
31    let mut count = data.len();
32    let mut i = 0;
33
34    while i < count {
35        if i > 2 && data[i - 2] == 0 && data[i - 1] == 0 && data[i] <= 3 {
36            data.insert(i, 3);
37            count += 1;
38        }
39
40        i += 1;
41    }
42}
43
44pub fn aud_for_frame(frame: &Frame, start_code: Option<NALUStartCode>) -> Result<Vec<u8>> {
45    let pic_type: u8 = match &frame.frame_type {
46        2 => 0, // I
47        1 => 1, // P, I
48        0 => 2, // B, P, I
49        _ => 7,
50    };
51
52    let mut data = if let Some(sc) = start_code {
53        sc.slice().to_vec()
54    } else {
55        Vec::new()
56    };
57
58    let mut writer = BitstreamIoWriter::with_capacity(24);
59
60    writer.write_bit(false)?; // forbidden_zero_bit
61
62    writer.write::<6, u8>(NAL_AUD)?; // nal_unit_type
63    writer.write_const::<6, 0>()?; // nuh_layer_id
64    writer.write_const::<3, 1>()?; // nuh_temporal_id_plus1
65
66    writer.write::<3, u8>(pic_type)?; // pic_type
67
68    // rbsp_trailing_bits()
69    writer.write_bit(true)?; // rbsp_stop_one_bit
70
71    // rbsp_alignment_zero_bit
72    writer.byte_align()?;
73
74    data.extend_from_slice(
75        writer
76            .as_slice()
77            .ok_or_else(|| anyhow!("Unaligned bytes"))?,
78    );
79
80    Ok(data)
81}