Add multiple SEI message support

Fixes #3

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "hevc_parser"
-version = "0.4.7"
+version = "0.5.0"
 authors = ["quietvoid"]
 edition = "2021"
 rust-version = "1.56.1"
@@ -10,8 +10,8 @@ repository = "https://github.com/quietvoid/hevc_parser"
 
 [dependencies]
 nom = "7.1.1"
-bitvec_helpers = "1.0.2"
-anyhow = "1.0.62"
+bitvec_helpers = "2.0.0"
+anyhow = "1.0.65"
 
 regex = { version = "1.6.0", optional = true }
 

src/hevc/hrd_parameters.rs

@@ -24,22 +24,22 @@ impl HrdParameters {
                 subpic_params_present = bs.get()?;
 
                 if subpic_params_present {
-                    bs.skip_n(8); // tick_divisor_minus2
-                    bs.skip_n(5); // du_cpb_removal_delay_increment_length_minus1
-                    bs.skip_n(1); // sub_pic_cpb_params_in_pic_timing_sei_flag
-                    bs.skip_n(5); // dpb_output_delay_du_length_minus1
+                    bs.skip_n(8)?; // tick_divisor_minus2
+                    bs.skip_n(5)?; // du_cpb_removal_delay_increment_length_minus1
+                    bs.skip_n(1)?; // sub_pic_cpb_params_in_pic_timing_sei_flag
+                    bs.skip_n(5)?; // dpb_output_delay_du_length_minus1
                 }
 
-                bs.skip_n(4); // bit_rate_scale
-                bs.skip_n(4); // cpb_size_scale
+                bs.skip_n(4)?; // bit_rate_scale
+                bs.skip_n(4)?; // cpb_size_scale
 
                 if subpic_params_present {
-                    bs.skip_n(4); // cpb_size_du_scale
+                    bs.skip_n(4)?; // cpb_size_du_scale
                 }
 
-                bs.skip_n(5); // initial_cpb_removal_delay_length_minus1
-                bs.skip_n(5); // au_cpb_removal_delay_length_minus1
-                bs.skip_n(5); // dpb_output_delay_length_minus1
+                bs.skip_n(5)?; // initial_cpb_removal_delay_length_minus1
+                bs.skip_n(5)?; // au_cpb_removal_delay_length_minus1
+                bs.skip_n(5)?; // dpb_output_delay_length_minus1
             }
         }
 
@@ -86,7 +86,7 @@ impl SubLayerHrdParameter {
                 bs.get_ue()?; // bit_rate_du_value_minus1
             }
 
-            bs.skip_n(1); // cbr_flag
+            bs.skip_n(1)?; // cbr_flag
         }
 
         Ok(())

src/hevc/mod.rs

@@ -1,5 +1,3 @@
-use anyhow::{bail, Result};
-
 use self::slice::SliceNAL;
 
 use super::{BitVecReader, NALUStartCode};
@@ -8,6 +6,7 @@ pub(crate) mod hrd_parameters;
 pub(crate) mod pps;
 pub(crate) mod profile_tier_level;
 pub(crate) mod scaling_list_data;
+pub mod sei;
 pub(crate) mod short_term_rps;
 pub(crate) mod slice;
 pub(crate) mod sps;
@@ -46,6 +45,8 @@ pub const NAL_UNSPEC63: u8 = 63;
 
 pub const USER_DATA_REGISTERED_ITU_T_35: u8 = 4;
 
+pub use sei::SeiMessage;
+
 #[derive(Default, Debug, Clone)]
 pub struct NALUnit {
     pub start: usize,
@@ -73,73 +74,6 @@ pub struct Frame {
     pub first_slice: SliceNAL,
 }
 
-#[derive(Default, Debug, Clone)]
-pub struct SeiMessage {
-    num_payload_type_ff_bytes: usize,
-    last_payload_type_byte: u8,
-
-    num_payload_size_ff_bytes: usize,
-    last_payload_size_byte: u8,
-
-    pub payload_type: u8,
-    pub payload_size: usize,
-}
-
-impl SeiMessage {
-    pub fn from_bytes(data: &[u8]) -> Result<SeiMessage> {
-        let mut reader = BitVecReader::new(data.to_vec());
-
-        SeiMessage::parse(&mut reader)
-    }
-
-    pub fn parse(reader: &mut BitVecReader) -> Result<SeiMessage> {
-        // forbidden_zero_bit
-        reader.skip_n(1);
-
-        let nal_type = reader.get_n::<u8>(6);
-
-        if nal_type != NAL_SEI_PREFIX {
-            bail!("NAL type {} is not SEI_PREFIX", nal_type);
-        }
-
-        if reader.available() < 9 && matches!(nal_type, NAL_EOS_NUT | NAL_EOB_NUT) {
-        } else {
-            reader.skip_n(6); // nuh_layer_id
-            reader.skip_n(3); // temporal_id
-        }
-
-        let mut msg = SeiMessage {
-            last_payload_type_byte: reader.get_n(8),
-            ..Default::default()
-        };
-
-        while msg.last_payload_type_byte == 0xFF {
-            msg.num_payload_type_ff_bytes += 1;
-            msg.last_payload_type_byte = reader.get_n(8);
-
-            msg.payload_type += 255;
-        }
-
-        msg.payload_type += msg.last_payload_type_byte;
-
-        msg.last_payload_size_byte = reader.get_n(8);
-        while msg.last_payload_size_byte == 0xFF {
-            msg.num_payload_size_ff_bytes += 1;
-            msg.last_payload_size_byte = reader.get_n(8);
-
-            msg.payload_size += 255;
-        }
-
-        msg.payload_size += msg.last_payload_size_byte as usize;
-
-        if msg.payload_size > reader.available() {
-            bail!("Payload size is larger than NALU size");
-        }
-
-        Ok(msg)
-    }
-}
-
 impl NALUnit {
     pub fn is_type_slice(nal_type: u8) -> bool {
         matches!(

src/hevc/pps.rs

@@ -61,7 +61,7 @@ impl PPSNAL {
 
         pps.dependent_slice_segments_enabled_flag = bs.get()?;
         pps.output_flag_present_flag = bs.get()?;
-        pps.num_extra_slice_header_bits = bs.get_n(3);
+        pps.num_extra_slice_header_bits = bs.get_n(3)?;
         pps.sign_data_hiding_flag = bs.get()?;
         pps.cabac_init_present_flag = bs.get()?;
         pps.num_ref_idx_l0_default_active = bs.get_ue()? + 1;

src/hevc/profile_tier_level.rs

@@ -28,9 +28,9 @@ pub struct ProfileTierLevel {
 
 impl ProfileTierLevel {
     pub fn parse(&mut self, bs: &mut BitVecReader, max_sub_layers: u8) -> Result<()> {
-        self.general_profile_space = bs.get_n(2);
+        self.general_profile_space = bs.get_n(2)?;
         self.general_tier_flag = bs.get()?;
-        self.general_profile_idc = bs.get_n(5);
+        self.general_profile_idc = bs.get_n(5)?;
 
         for _ in 0..32 {
             self.general_profile_compatibility_flag.push(bs.get()?);
@@ -40,9 +40,9 @@ impl ProfileTierLevel {
         self.general_interlaced_source_flag = bs.get()?;
         self.general_non_packed_constraint_flag = bs.get()?;
         self.general_frame_only_constraint_flag = bs.get()?;
-        bs.skip_n(32);
-        bs.skip_n(12);
-        self.general_level_idc = bs.get_n(8);
+        bs.skip_n(32)?;
+        bs.skip_n(12)?;
+        self.general_level_idc = bs.get_n(8)?;
 
         let max_sub_layers_minus1 = max_sub_layers - 1;
         for _ in 0..max_sub_layers_minus1 {
@@ -52,15 +52,15 @@ impl ProfileTierLevel {
 
         if max_sub_layers_minus1 > 0 {
             for _ in max_sub_layers_minus1..8 {
-                bs.skip_n(2);
+                bs.skip_n(2)?;
             }
         }
 
         for i in 0..max_sub_layers_minus1 as usize {
             if self.sub_layer_profile_present_flag[i] {
-                self.sub_layer_profile_space.push(bs.get_n(2));
+                self.sub_layer_profile_space.push(bs.get_n(2)?);
                 self.sub_layer_tier_flag.push(bs.get()?);
-                self.sub_layer_profile_idc.push(bs.get_n(5));
+                self.sub_layer_profile_idc.push(bs.get_n(5)?);
 
                 for _ in 0..32 {
                     self.sub_layer_profile_compatibility_flag.push(bs.get()?);
@@ -71,12 +71,12 @@ impl ProfileTierLevel {
                 self.sub_layer_non_packed_constraint_flag.push(bs.get()?);
                 self.sub_layer_frame_only_constraint_flag.push(bs.get()?);
 
-                bs.skip_n(32);
-                bs.skip_n(12);
+                bs.skip_n(32)?;
+                bs.skip_n(12)?;
             }
 
             if self.sub_layer_level_present_flag[i] {
-                self.sub_layer_level_idc.push(bs.get_n(8));
+                self.sub_layer_level_idc.push(bs.get_n(8)?);
             } else {
                 self.sub_layer_level_idc.push(1);
             }

src/hevc/sei.rs

@@ -0,0 +1,89 @@
+use super::{NAL_EOB_NUT, NAL_EOS_NUT, NAL_SEI_PREFIX, NAL_SEI_SUFFIX};
+use anyhow::{bail, Result};
+use bitvec_helpers::bitslice_reader::BitSliceReader;
+
+#[derive(Default, Debug, Clone)]
+pub struct SeiMessage {
+    num_payload_type_ff_bytes: usize,
+    last_payload_type_byte: u8,
+
+    num_payload_size_ff_bytes: usize,
+    last_payload_size_byte: u8,
+
+    // Offset of the messame in the input slice
+    pub msg_offset: usize,
+
+    pub payload_type: u8,
+    pub payload_offset: usize,
+    pub payload_size: usize,
+}
+
+impl SeiMessage {
+    /// Assumes the data does not contain any `emulation_prevention_three_byte`s
+    pub fn parse_sei_rbsp(data: &[u8]) -> Result<Vec<SeiMessage>> {
+        let mut reader = BitSliceReader::new(data);
+
+        // forbidden_zero_bit
+        reader.skip_n(1)?;
+
+        let nal_type = reader.get_n::<u8>(6)?;
+
+        if nal_type != NAL_SEI_PREFIX && nal_type != NAL_SEI_SUFFIX {
+            bail!("NAL type {} is not SEI", nal_type);
+        }
+
+        if reader.available() < 9 && matches!(nal_type, NAL_EOS_NUT | NAL_EOB_NUT) {
+        } else {
+            reader.skip_n(6)?; // nuh_layer_id
+            reader.skip_n(3)?; // temporal_id
+        }
+
+        let mut messages = Vec::new();
+
+        loop {
+            messages.push(Self::parse_sei_message(&mut reader)?);
+
+            if reader.available() <= 8 {
+                break;
+            }
+        }
+
+        Ok(messages)
+    }
+
+    fn parse_sei_message(reader: &mut BitSliceReader) -> Result<SeiMessage> {
+        let mut msg = SeiMessage {
+            msg_offset: reader.position() / 8,
+            last_payload_type_byte: reader.get_n(8)?,
+            ..Default::default()
+        };
+
+        while msg.last_payload_type_byte == 0xFF {
+            msg.num_payload_type_ff_bytes += 1;
+            msg.last_payload_type_byte = reader.get_n(8)?;
+
+            msg.payload_type += 255;
+        }
+
+        msg.payload_type += msg.last_payload_type_byte;
+
+        msg.last_payload_size_byte = reader.get_n(8)?;
+        while msg.last_payload_size_byte == 0xFF {
+            msg.num_payload_size_ff_bytes += 1;
+            msg.last_payload_size_byte = reader.get_n(8)?;
+
+            msg.payload_size += 255;
+        }
+
+        msg.payload_size += msg.last_payload_size_byte as usize;
+        msg.payload_offset = reader.position() / 8;
+
+        if msg.payload_size > reader.available() {
+            bail!("Payload size is larger than NALU size");
+        }
+
+        reader.skip_n(msg.payload_size * 8)?;
+
+        Ok(msg)
+    }
+}

src/hevc/slice.rs

@@ -34,7 +34,7 @@ impl SliceNAL {
 
         if is_irap_nal(nal) {
             slice.key_frame = true;
-            bs.skip_n(1); // no_output_of_prior_pics_flag
+            bs.skip_n(1)?; // no_output_of_prior_pics_flag
         }
 
         slice.pps_id = bs.get_ue()?;
@@ -55,7 +55,7 @@ impl SliceNAL {
             let pic_size = (sps.ctb_width * sps.ctb_height) as f64;
             let slice_address_length = pic_size.log2().ceil() as usize;
 
-            slice.slice_segment_addr = bs.get_n(slice_address_length);
+            slice.slice_segment_addr = bs.get_n(slice_address_length)?;
         } else {
             slice.dependent_slice_segment_flag = false;
         }
@@ -65,21 +65,21 @@ impl SliceNAL {
         }
 
         for _ in 0..pps.num_extra_slice_header_bits {
-            bs.skip_n(1); // slice_reserved_undetermined_flag
+            bs.skip_n(1)?; // slice_reserved_undetermined_flag
         }
 
         slice.slice_type = bs.get_ue()?;
 
         if pps.output_flag_present_flag {
-            bs.skip_n(1);
+            bs.skip_n(1)?;
         }
 
         if sps.separate_colour_plane_flag {
-            bs.skip_n(2);
+            bs.skip_n(2)?;
         }
 
         if !is_idr_nal(nal) {
-            slice.pic_order_cnt_lsb = bs.get_n(sps.log2_max_poc_lsb as usize);
+            slice.pic_order_cnt_lsb = bs.get_n(sps.log2_max_poc_lsb as usize)?;
             slice.output_picture_number = compute_poc(sps, *poc_tid0, slice.pic_order_cnt_lsb, nal);
         } else {
             slice.output_picture_number = 0;

src/hevc/sps.rs

@@ -88,11 +88,11 @@ pub struct SPSNAL {
 impl SPSNAL {
     pub fn parse(bs: &mut BitVecReader) -> Result<SPSNAL> {
         let mut sps = SPSNAL {
-            vps_id: bs.get_n(4),
+            vps_id: bs.get_n(4)?,
             ..Default::default()
         };
 
-        sps.max_sub_layers = bs.get_n::<u8>(3) + 1;
+        sps.max_sub_layers = bs.get_n::<u8>(3)? + 1;
         sps.temporal_id_nesting_flag = bs.get()?;
 
         sps.ptl.parse(bs, sps.max_sub_layers)?;
@@ -165,8 +165,8 @@ impl SPSNAL {
         sps.pcm_enabled_flag = bs.get()?;
 
         if sps.pcm_enabled_flag {
-            sps.pcm_bit_depth = bs.get_n::<u8>(4) + 1;
-            sps.pcm_bit_depth_chroma = bs.get_n::<u8>(4) + 1;
+            sps.pcm_bit_depth = bs.get_n::<u8>(4)? + 1;
+            sps.pcm_bit_depth_chroma = bs.get_n::<u8>(4)? + 1;
             sps.pcm_log2_min_pcm_cb_size = bs.get_ue()? + 3;
             sps.pcm_log2_max_pcm_cb_size = bs.get_ue()? + sps.pcm_log2_min_pcm_cb_size;
 
@@ -189,7 +189,7 @@ impl SPSNAL {
 
             for _ in 0..sps.num_long_term_ref_pics_sps {
                 sps.lt_ref_pic_poc_lsb_sps
-                    .push(bs.get_n(sps.log2_max_poc_lsb as usize));
+                    .push(bs.get_n(sps.log2_max_poc_lsb as usize)?);
                 sps.used_by_curr_pic_lt_sps_flag.push(bs.get()?);
             }
         }