Skip to content

Commit 2008559

Browse files
quietvoidquietvoid
committed
Add multiple SEI message support
Fixes #3
1 parent 5a004c2 commit 2008559

File tree

11 files changed

+152
-129
lines changed

11 files changed

+152
-129
lines changed

‎Cargo.toml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "hevc_parser"
3-
version = "0.4.7"
3+
version = "0.5.0"
44
authors = ["quietvoid"]
55
edition = "2021"
66
rust-version = "1.56.1"
@@ -10,8 +10,8 @@ repository = "https://github.com/quietvoid/hevc_parser"
1010

1111
[dependencies]
1212
nom = "7.1.1"
13-
bitvec_helpers = "1.0.2"
14-
anyhow = "1.0.62"
13+
bitvec_helpers = "2.0.0"
14+
anyhow = "1.0.65"
1515

1616
regex = { version = "1.6.0", optional = true }
1717

‎src/hevc/hrd_parameters.rs‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -24,22 +24,22 @@ impl HrdParameters {
2424
subpic_params_present = bs.get()?;
2525

2626
if subpic_params_present {
27-
bs.skip_n(8); // tick_divisor_minus2
28-
bs.skip_n(5); // du_cpb_removal_delay_increment_length_minus1
29-
bs.skip_n(1); // sub_pic_cpb_params_in_pic_timing_sei_flag
30-
bs.skip_n(5); // dpb_output_delay_du_length_minus1
27+
bs.skip_n(8)?; // tick_divisor_minus2
28+
bs.skip_n(5)?; // du_cpb_removal_delay_increment_length_minus1
29+
bs.skip_n(1)?; // sub_pic_cpb_params_in_pic_timing_sei_flag
30+
bs.skip_n(5)?; // dpb_output_delay_du_length_minus1
3131
}
3232

33-
bs.skip_n(4); // bit_rate_scale
34-
bs.skip_n(4); // cpb_size_scale
33+
bs.skip_n(4)?; // bit_rate_scale
34+
bs.skip_n(4)?; // cpb_size_scale
3535

3636
if subpic_params_present {
37-
bs.skip_n(4); // cpb_size_du_scale
37+
bs.skip_n(4)?; // cpb_size_du_scale
3838
}
3939

40-
bs.skip_n(5); // initial_cpb_removal_delay_length_minus1
41-
bs.skip_n(5); // au_cpb_removal_delay_length_minus1
42-
bs.skip_n(5); // dpb_output_delay_length_minus1
40+
bs.skip_n(5)?; // initial_cpb_removal_delay_length_minus1
41+
bs.skip_n(5)?; // au_cpb_removal_delay_length_minus1
42+
bs.skip_n(5)?; // dpb_output_delay_length_minus1
4343
}
4444
}
4545

@@ -86,7 +86,7 @@ impl SubLayerHrdParameter {
8686
bs.get_ue()?; // bit_rate_du_value_minus1
8787
}
8888

89-
bs.skip_n(1); // cbr_flag
89+
bs.skip_n(1)?; // cbr_flag
9090
}
9191

9292
Ok(())

‎src/hevc/mod.rs‎

Lines changed: 3 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
use anyhow::{bail, Result};
2-
31
use self::slice::SliceNAL;
42

53
use super::{BitVecReader, NALUStartCode};
@@ -8,6 +6,7 @@ pub(crate) mod hrd_parameters;
86
pub(crate) mod pps;
97
pub(crate) mod profile_tier_level;
108
pub(crate) mod scaling_list_data;
9+
pub mod sei;
1110
pub(crate) mod short_term_rps;
1211
pub(crate) mod slice;
1312
pub(crate) mod sps;
@@ -46,6 +45,8 @@ pub const NAL_UNSPEC63: u8 = 63;
4645

4746
pub const USER_DATA_REGISTERED_ITU_T_35: u8 = 4;
4847

48+
pub use sei::SeiMessage;
49+
4950
#[derive(Default, Debug, Clone)]
5051
pub struct NALUnit {
5152
pub start: usize,
@@ -73,73 +74,6 @@ pub struct Frame {
7374
pub first_slice: SliceNAL,
7475
}
7576

76-
#[derive(Default, Debug, Clone)]
77-
pub struct SeiMessage {
78-
num_payload_type_ff_bytes: usize,
79-
last_payload_type_byte: u8,
80-
81-
num_payload_size_ff_bytes: usize,
82-
last_payload_size_byte: u8,
83-
84-
pub payload_type: u8,
85-
pub payload_size: usize,
86-
}
87-
88-
impl SeiMessage {
89-
pub fn from_bytes(data: &[u8]) -> Result<SeiMessage> {
90-
let mut reader = BitVecReader::new(data.to_vec());
91-
92-
SeiMessage::parse(&mut reader)
93-
}
94-
95-
pub fn parse(reader: &mut BitVecReader) -> Result<SeiMessage> {
96-
// forbidden_zero_bit
97-
reader.skip_n(1);
98-
99-
let nal_type = reader.get_n::<u8>(6);
100-
101-
if nal_type != NAL_SEI_PREFIX {
102-
bail!("NAL type {} is not SEI_PREFIX", nal_type);
103-
}
104-
105-
if reader.available() < 9 && matches!(nal_type, NAL_EOS_NUT | NAL_EOB_NUT) {
106-
} else {
107-
reader.skip_n(6); // nuh_layer_id
108-
reader.skip_n(3); // temporal_id
109-
}
110-
111-
let mut msg = SeiMessage {
112-
last_payload_type_byte: reader.get_n(8),
113-
..Default::default()
114-
};
115-
116-
while msg.last_payload_type_byte == 0xFF {
117-
msg.num_payload_type_ff_bytes += 1;
118-
msg.last_payload_type_byte = reader.get_n(8);
119-
120-
msg.payload_type += 255;
121-
}
122-
123-
msg.payload_type += msg.last_payload_type_byte;
124-
125-
msg.last_payload_size_byte = reader.get_n(8);
126-
while msg.last_payload_size_byte == 0xFF {
127-
msg.num_payload_size_ff_bytes += 1;
128-
msg.last_payload_size_byte = reader.get_n(8);
129-
130-
msg.payload_size += 255;
131-
}
132-
133-
msg.payload_size += msg.last_payload_size_byte as usize;
134-
135-
if msg.payload_size > reader.available() {
136-
bail!("Payload size is larger than NALU size");
137-
}
138-
139-
Ok(msg)
140-
}
141-
}
142-
14377
impl NALUnit {
14478
pub fn is_type_slice(nal_type: u8) -> bool {
14579
matches!(

‎src/hevc/pps.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ impl PPSNAL {
6161

6262
pps.dependent_slice_segments_enabled_flag = bs.get()?;
6363
pps.output_flag_present_flag = bs.get()?;
64-
pps.num_extra_slice_header_bits = bs.get_n(3);
64+
pps.num_extra_slice_header_bits = bs.get_n(3)?;
6565
pps.sign_data_hiding_flag = bs.get()?;
6666
pps.cabac_init_present_flag = bs.get()?;
6767
pps.num_ref_idx_l0_default_active = bs.get_ue()? + 1;

‎src/hevc/profile_tier_level.rs‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,9 @@ pub struct ProfileTierLevel {
2828

2929
impl ProfileTierLevel {
3030
pub fn parse(&mut self, bs: &mut BitVecReader, max_sub_layers: u8) -> Result<()> {
31-
self.general_profile_space = bs.get_n(2);
31+
self.general_profile_space = bs.get_n(2)?;
3232
self.general_tier_flag = bs.get()?;
33-
self.general_profile_idc = bs.get_n(5);
33+
self.general_profile_idc = bs.get_n(5)?;
3434

3535
for _ in 0..32 {
3636
self.general_profile_compatibility_flag.push(bs.get()?);
@@ -40,9 +40,9 @@ impl ProfileTierLevel {
4040
self.general_interlaced_source_flag = bs.get()?;
4141
self.general_non_packed_constraint_flag = bs.get()?;
4242
self.general_frame_only_constraint_flag = bs.get()?;
43-
bs.skip_n(32);
44-
bs.skip_n(12);
45-
self.general_level_idc = bs.get_n(8);
43+
bs.skip_n(32)?;
44+
bs.skip_n(12)?;
45+
self.general_level_idc = bs.get_n(8)?;
4646

4747
let max_sub_layers_minus1 = max_sub_layers - 1;
4848
for _ in 0..max_sub_layers_minus1 {
@@ -52,15 +52,15 @@ impl ProfileTierLevel {
5252

5353
if max_sub_layers_minus1 > 0 {
5454
for _ in max_sub_layers_minus1..8 {
55-
bs.skip_n(2);
55+
bs.skip_n(2)?;
5656
}
5757
}
5858

5959
for i in 0..max_sub_layers_minus1 as usize {
6060
if self.sub_layer_profile_present_flag[i] {
61-
self.sub_layer_profile_space.push(bs.get_n(2));
61+
self.sub_layer_profile_space.push(bs.get_n(2)?);
6262
self.sub_layer_tier_flag.push(bs.get()?);
63-
self.sub_layer_profile_idc.push(bs.get_n(5));
63+
self.sub_layer_profile_idc.push(bs.get_n(5)?);
6464

6565
for _ in 0..32 {
6666
self.sub_layer_profile_compatibility_flag.push(bs.get()?);
@@ -71,12 +71,12 @@ impl ProfileTierLevel {
7171
self.sub_layer_non_packed_constraint_flag.push(bs.get()?);
7272
self.sub_layer_frame_only_constraint_flag.push(bs.get()?);
7373

74-
bs.skip_n(32);
75-
bs.skip_n(12);
74+
bs.skip_n(32)?;
75+
bs.skip_n(12)?;
7676
}
7777

7878
if self.sub_layer_level_present_flag[i] {
79-
self.sub_layer_level_idc.push(bs.get_n(8));
79+
self.sub_layer_level_idc.push(bs.get_n(8)?);
8080
} else {
8181
self.sub_layer_level_idc.push(1);
8282
}

‎src/hevc/sei.rs‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
use super::{NAL_EOB_NUT, NAL_EOS_NUT, NAL_SEI_PREFIX, NAL_SEI_SUFFIX};
2+
use anyhow::{bail, Result};
3+
use bitvec_helpers::bitslice_reader::BitSliceReader;
4+
5+
#[derive(Default, Debug, Clone)]
6+
pub struct SeiMessage {
7+
num_payload_type_ff_bytes: usize,
8+
last_payload_type_byte: u8,
9+
10+
num_payload_size_ff_bytes: usize,
11+
last_payload_size_byte: u8,
12+
13+
// Offset of the messame in the input slice
14+
pub msg_offset: usize,
15+
16+
pub payload_type: u8,
17+
pub payload_offset: usize,
18+
pub payload_size: usize,
19+
}
20+
21+
impl SeiMessage {
22+
/// Assumes the data does not contain any `emulation_prevention_three_byte`s
23+
pub fn parse_sei_rbsp(data: &[u8]) -> Result<Vec<SeiMessage>> {
24+
let mut reader = BitSliceReader::new(data);
25+
26+
// forbidden_zero_bit
27+
reader.skip_n(1)?;
28+
29+
let nal_type = reader.get_n::<u8>(6)?;
30+
31+
if nal_type != NAL_SEI_PREFIX && nal_type != NAL_SEI_SUFFIX {
32+
bail!("NAL type {} is not SEI", nal_type);
33+
}
34+
35+
if reader.available() < 9 && matches!(nal_type, NAL_EOS_NUT | NAL_EOB_NUT) {
36+
} else {
37+
reader.skip_n(6)?; // nuh_layer_id
38+
reader.skip_n(3)?; // temporal_id
39+
}
40+
41+
let mut messages = Vec::new();
42+
43+
loop {
44+
messages.push(Self::parse_sei_message(&mut reader)?);
45+
46+
if reader.available() <= 8 {
47+
break;
48+
}
49+
}
50+
51+
Ok(messages)
52+
}
53+
54+
fn parse_sei_message(reader: &mut BitSliceReader) -> Result<SeiMessage> {
55+
let mut msg = SeiMessage {
56+
msg_offset: reader.position() / 8,
57+
last_payload_type_byte: reader.get_n(8)?,
58+
..Default::default()
59+
};
60+
61+
while msg.last_payload_type_byte == 0xFF {
62+
msg.num_payload_type_ff_bytes += 1;
63+
msg.last_payload_type_byte = reader.get_n(8)?;
64+
65+
msg.payload_type += 255;
66+
}
67+
68+
msg.payload_type += msg.last_payload_type_byte;
69+
70+
msg.last_payload_size_byte = reader.get_n(8)?;
71+
while msg.last_payload_size_byte == 0xFF {
72+
msg.num_payload_size_ff_bytes += 1;
73+
msg.last_payload_size_byte = reader.get_n(8)?;
74+
75+
msg.payload_size += 255;
76+
}
77+
78+
msg.payload_size += msg.last_payload_size_byte as usize;
79+
msg.payload_offset = reader.position() / 8;
80+
81+
if msg.payload_size > reader.available() {
82+
bail!("Payload size is larger than NALU size");
83+
}
84+
85+
reader.skip_n(msg.payload_size * 8)?;
86+
87+
Ok(msg)
88+
}
89+
}

‎src/hevc/slice.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ impl SliceNAL {
3434

3535
if is_irap_nal(nal) {
3636
slice.key_frame = true;
37-
bs.skip_n(1); // no_output_of_prior_pics_flag
37+
bs.skip_n(1)?; // no_output_of_prior_pics_flag
3838
}
3939

4040
slice.pps_id = bs.get_ue()?;
@@ -55,7 +55,7 @@ impl SliceNAL {
5555
let pic_size = (sps.ctb_width * sps.ctb_height) as f64;
5656
let slice_address_length = pic_size.log2().ceil() as usize;
5757

58-
slice.slice_segment_addr = bs.get_n(slice_address_length);
58+
slice.slice_segment_addr = bs.get_n(slice_address_length)?;
5959
} else {
6060
slice.dependent_slice_segment_flag = false;
6161
}
@@ -65,21 +65,21 @@ impl SliceNAL {
6565
}
6666

6767
for _ in 0..pps.num_extra_slice_header_bits {
68-
bs.skip_n(1); // slice_reserved_undetermined_flag
68+
bs.skip_n(1)?; // slice_reserved_undetermined_flag
6969
}
7070

7171
slice.slice_type = bs.get_ue()?;
7272

7373
if pps.output_flag_present_flag {
74-
bs.skip_n(1);
74+
bs.skip_n(1)?;
7575
}
7676

7777
if sps.separate_colour_plane_flag {
78-
bs.skip_n(2);
78+
bs.skip_n(2)?;
7979
}
8080

8181
if !is_idr_nal(nal) {
82-
slice.pic_order_cnt_lsb = bs.get_n(sps.log2_max_poc_lsb as usize);
82+
slice.pic_order_cnt_lsb = bs.get_n(sps.log2_max_poc_lsb as usize)?;
8383
slice.output_picture_number = compute_poc(sps, *poc_tid0, slice.pic_order_cnt_lsb, nal);
8484
} else {
8585
slice.output_picture_number = 0;

‎src/hevc/sps.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -88,11 +88,11 @@ pub struct SPSNAL {
8888
impl SPSNAL {
8989
pub fn parse(bs: &mut BitVecReader) -> Result<SPSNAL> {
9090
let mut sps = SPSNAL {
91-
vps_id: bs.get_n(4),
91+
vps_id: bs.get_n(4)?,
9292
..Default::default()
9393
};
9494

95-
sps.max_sub_layers = bs.get_n::<u8>(3) + 1;
95+
sps.max_sub_layers = bs.get_n::<u8>(3)? + 1;
9696
sps.temporal_id_nesting_flag = bs.get()?;
9797

9898
sps.ptl.parse(bs, sps.max_sub_layers)?;
@@ -165,8 +165,8 @@ impl SPSNAL {
165165
sps.pcm_enabled_flag = bs.get()?;
166166

167167
if sps.pcm_enabled_flag {
168-
sps.pcm_bit_depth = bs.get_n::<u8>(4) + 1;
169-
sps.pcm_bit_depth_chroma = bs.get_n::<u8>(4) + 1;
168+
sps.pcm_bit_depth = bs.get_n::<u8>(4)? + 1;
169+
sps.pcm_bit_depth_chroma = bs.get_n::<u8>(4)? + 1;
170170
sps.pcm_log2_min_pcm_cb_size = bs.get_ue()? + 3;
171171
sps.pcm_log2_max_pcm_cb_size = bs.get_ue()? + sps.pcm_log2_min_pcm_cb_size;
172172

@@ -189,7 +189,7 @@ impl SPSNAL {
189189

190190
for _ in 0..sps.num_long_term_ref_pics_sps {
191191
sps.lt_ref_pic_poc_lsb_sps
192-
.push(bs.get_n(sps.log2_max_poc_lsb as usize));
192+
.push(bs.get_n(sps.log2_max_poc_lsb as usize)?);
193193
sps.used_by_curr_pic_lt_sps_flag.push(bs.get()?);
194194
}
195195
}

0 commit comments

Comments
 (0)