scuffle_h264/
sps.rs

1use std::io;
2
3use byteorder::{BigEndian, ReadBytesExt};
4use bytes::Bytes;
5use scuffle_bytes_util::BitReader;
6use scuffle_expgolomb::BitReaderExpGolombExt;
7
8#[derive(Debug, Clone, PartialEq)]
9/// Sequence parameter set
10/// ISO/IEC-14496-10-2022 - 7.3.2
11pub struct Sps {
12    pub profile_idc: u8,
13    pub level_idc: u8,
14    pub ext: Option<SpsExtended>,
15    pub width: u64,
16    pub height: u64,
17    pub frame_rate: f64,
18    pub color_config: Option<ColorConfig>,
19}
20
21#[derive(Debug, Clone, PartialEq)]
22/// Color config for SPS
23pub struct ColorConfig {
24    pub full_range: bool,
25    pub color_primaries: u8,
26    pub transfer_characteristics: u8,
27    pub matrix_coefficients: u8,
28}
29
30impl Sps {
31    pub fn parse(data: Bytes) -> io::Result<Self> {
32        let mut vec = Vec::with_capacity(data.len());
33
34        // We need to remove the emulation prevention byte
35        // This is BARELY documented in the spec, but it's there.
36        // ISO/IEC-14496-10-2022 - 3.1.48
37        let mut i = 0;
38        while i < data.len() - 3 {
39            if data[i] == 0x00 && data[i + 1] == 0x00 && data[i + 2] == 0x03 {
40                vec.push(0x00);
41                vec.push(0x00);
42                i += 3;
43            } else {
44                vec.push(data[i]);
45                i += 1;
46            }
47        }
48
49        let mut bit_reader = BitReader::new_from_slice(vec);
50
51        let forbidden_zero_bit = bit_reader.read_bit()?;
52        if forbidden_zero_bit {
53            return Err(io::Error::new(io::ErrorKind::InvalidData, "Forbidden zero bit is set"));
54        }
55
56        bit_reader.seek_bits(2)?; // nal_ref_idc
57
58        let nal_unit_type = bit_reader.read_bits(5)?;
59        if nal_unit_type != 7 {
60            return Err(io::Error::new(io::ErrorKind::InvalidData, "NAL unit type is not SPS"));
61        }
62
63        let profile_idc = bit_reader.read_u8()?;
64        bit_reader.seek_bits(
65            1 // constraint_set0_flag
66            + 1 // constraint_set1_flag
67            + 1 // constraint_set2_flag
68            + 1 // constraint_set3_flag
69            + 4, // reserved_zero_4bits
70        )?;
71
72        let level_idc = bit_reader.read_u8()?;
73        bit_reader.read_exp_golomb()?; // seq_parameter_set_id
74
75        let sps_ext = match profile_idc {
76            100 | 110 | 122 | 244 | 44 | 83 | 86 | 118 | 128 | 138 | 139 | 134 | 135 => {
77                Some(SpsExtended::parse(&mut bit_reader)?)
78            }
79            _ => None,
80        };
81
82        bit_reader.read_exp_golomb()?; // log2_max_frame_num_minus4
83        let pic_order_cnt_type = bit_reader.read_exp_golomb()?;
84        if pic_order_cnt_type == 0 {
85            bit_reader.read_exp_golomb()?; // log2_max_pic_order_cnt_lsb_minus4
86        } else if pic_order_cnt_type == 1 {
87            bit_reader.seek_bits(1)?; // delta_pic_order_always_zero_flag
88            bit_reader.read_signed_exp_golomb()?; // offset_for_non_ref_pic
89            bit_reader.read_signed_exp_golomb()?; // offset_for_top_to_bottom_field
90            let num_ref_frames_in_pic_order_cnt_cycle = bit_reader.read_exp_golomb()?;
91            for _ in 0..num_ref_frames_in_pic_order_cnt_cycle {
92                bit_reader.read_signed_exp_golomb()?; // offset_for_ref_frame
93            }
94        }
95
96        bit_reader.read_exp_golomb()?; // max_num_ref_frames
97        bit_reader.read_bit()?; // gaps_in_frame_num_value_allowed_flag
98        let pic_width_in_mbs_minus1 = bit_reader.read_exp_golomb()?; // pic_width_in_mbs_minus1
99        let pic_height_in_map_units_minus1 = bit_reader.read_exp_golomb()?; // pic_height_in_map_units_minus1
100        let frame_mbs_only_flag = bit_reader.read_bit()?;
101        if !frame_mbs_only_flag {
102            bit_reader.seek_bits(1)?; // mb_adaptive_frame_field_flag
103        }
104
105        bit_reader.seek_bits(1)?; // direct_8x8_inference_flag
106
107        let mut frame_crop_left_offset = 0;
108        let mut frame_crop_right_offset = 0;
109        let mut frame_crop_top_offset = 0;
110        let mut frame_crop_bottom_offset = 0;
111
112        if bit_reader.read_bit()? {
113            // frame_cropping_flag
114            frame_crop_left_offset = bit_reader.read_exp_golomb()?; // frame_crop_left_offset
115            frame_crop_right_offset = bit_reader.read_exp_golomb()?; // frame_crop_right_offset
116            frame_crop_top_offset = bit_reader.read_exp_golomb()?; // frame_crop_top_offset
117            frame_crop_bottom_offset = bit_reader.read_exp_golomb()?; // frame_crop_bottom_offset
118        }
119
120        let width = ((pic_width_in_mbs_minus1 + 1) * 16) - frame_crop_right_offset * 2 - frame_crop_left_offset * 2;
121        let height = ((2 - frame_mbs_only_flag as u64) * (pic_height_in_map_units_minus1 + 1) * 16)
122            - frame_crop_bottom_offset * 2
123            - frame_crop_top_offset * 2;
124
125        let mut frame_rate = 0.0;
126
127        let vui_parameters_present_flag = bit_reader.read_bit()?;
128
129        let mut color_config = None;
130
131        if vui_parameters_present_flag {
132            // We do want to read the VUI parameters to get the frame rate.
133
134            // aspect_ratio_info_present_flag
135            if bit_reader.read_bit()? {
136                let aspect_ratio_idc = bit_reader.read_u8()?;
137                if aspect_ratio_idc == 255 {
138                    bit_reader.seek_bits(16)?; // sar_width
139                    bit_reader.seek_bits(16)?; // sar_height
140                }
141            }
142
143            // overscan_info_present_flag
144            if bit_reader.read_bit()? {
145                bit_reader.seek_bits(1)?; // overscan_appropriate_flag
146            }
147
148            // video_signal_type_present_flag
149            if bit_reader.read_bit()? {
150                bit_reader.seek_bits(3)?; // video_format
151                let full_range = bit_reader.read_bit()?; // video_full_range_flag
152
153                let color_primaries;
154                let transfer_characteristics;
155                let matrix_coefficients;
156
157                if bit_reader.read_bit()? {
158                    // colour_description_present_flag
159                    color_primaries = bit_reader.read_u8()?; // colour_primaries
160                    transfer_characteristics = bit_reader.read_u8()?; // transfer_characteristics
161                    matrix_coefficients = bit_reader.read_u8()?; // matrix_coefficients
162                } else {
163                    color_primaries = 2; // UNSPECIFIED
164                    transfer_characteristics = 2; // UNSPECIFIED
165                    matrix_coefficients = 2; // UNSPECIFIED
166                }
167
168                color_config = Some(ColorConfig {
169                    full_range,
170                    color_primaries,
171                    transfer_characteristics,
172                    matrix_coefficients,
173                });
174            }
175
176            // chroma_loc_info_present_flag
177            if bit_reader.read_bit()? {
178                bit_reader.read_exp_golomb()?; // chroma_sample_loc_type_top_field
179                bit_reader.read_exp_golomb()?; // chroma_sample_loc_type_bottom_field
180            }
181
182            // timing_info_present_flag
183            if bit_reader.read_bit()? {
184                let num_units_in_tick = bit_reader.read_u32::<BigEndian>()?;
185                let time_scale = bit_reader.read_u32::<BigEndian>()?;
186
187                if num_units_in_tick == 0 {
188                    return Err(io::Error::new(io::ErrorKind::InvalidData, "num_units_in_tick cannot be zero"));
189                }
190
191                frame_rate = time_scale as f64 / (2.0 * num_units_in_tick as f64);
192            }
193        }
194
195        Ok(Sps {
196            profile_idc,
197            level_idc,
198            ext: sps_ext,
199            width,
200            height,
201            frame_rate,
202            color_config,
203        })
204    }
205}
206
207#[derive(Debug, Clone, PartialEq)]
208/// Sequence parameter set extension.
209/// ISO/IEC-14496-10-2022 - 7.3.2
210pub struct SpsExtended {
211    pub chroma_format_idc: u64,       // ue(v)
212    pub bit_depth_luma_minus8: u64,   // ue(v)
213    pub bit_depth_chroma_minus8: u64, // ue(v)
214}
215
216impl SpsExtended {
217    pub fn parse<T: io::Read>(reader: &mut BitReader<T>) -> io::Result<Self> {
218        let chroma_format_idc = reader.read_exp_golomb()?;
219        if chroma_format_idc == 3 {
220            reader.read_bit()?;
221        }
222
223        let bit_depth_luma_minus8 = reader.read_exp_golomb()?;
224        let bit_depth_chroma_minus8 = reader.read_exp_golomb()?;
225        reader.read_bit()?; // qpprime_y_zero_transform_bypass_flag
226
227        if reader.read_bit()? {
228            // seq_scaling_matrix_present_flag
229            // We need to read the scaling matrices here, but we don't need them
230            // for decoding, so we just skip them.
231            let count = if chroma_format_idc != 3 { 8 } else { 12 };
232            for i in 0..count {
233                if reader.read_bit()? {
234                    let size = if i < 6 { 16 } else { 64 };
235                    let mut next_scale = 8;
236                    for _ in 0..size {
237                        let delta_scale = reader.read_signed_exp_golomb()?;
238                        next_scale = (next_scale + delta_scale + 256) % 256;
239                        if next_scale == 0 {
240                            break;
241                        }
242                    }
243                }
244            }
245        }
246
247        Ok(SpsExtended {
248            chroma_format_idc,
249            bit_depth_luma_minus8,
250            bit_depth_chroma_minus8,
251        })
252    }
253}