1use crate::bitreader::BitReader;
2use crate::eg::{read_se, read_ue};
3use crate::nal::NalUnitType;
4use crate::pps::Pps;
5use crate::sps::Sps;
6use crate::{Error, Result};
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum SliceType {
10 P = 0,
11 B = 1,
12 I = 2,
13 Sp = 3,
14 Si = 4,
15}
16
17impl SliceType {
18 pub fn from_value(value: u32) -> Option<Self> {
19 match value % 5 {
20 0 => Some(SliceType::P),
21 1 => Some(SliceType::B),
22 2 => Some(SliceType::I),
23 3 => Some(SliceType::Sp),
24 4 => Some(SliceType::Si),
25 _ => None,
26 }
27 }
28}
29
30#[derive(Debug, Clone)]
31pub struct SliceHeader {
32 pub first_mb_in_slice: u32,
33 pub slice_type: SliceType,
34 pub pic_parameter_set_id: u8,
35 pub colour_plane_id: u8,
36 pub frame_num: u32,
37 pub field_pic_flag: bool,
38 pub bottom_field_flag: bool,
39 pub idr_pic_id: u32,
40 pub pic_order_cnt_lsb: u32,
41 pub delta_pic_order_cnt_bottom: i32,
42 pub delta_pic_order_cnt: [i32; 2],
43 pub redundant_pic_cnt: u32,
44 pub direct_spatial_mv_pred_flag: bool,
45 pub num_ref_idx_active_override_flag: bool,
46 pub num_ref_idx_l0_active_minus1: u32,
47 pub num_ref_idx_l1_active_minus1: u32,
48}
49
50impl SliceHeader {
51 pub fn parse(rbsp: &[u8], nal_type: NalUnitType, sps: &Sps, pps: &Pps) -> Result<Self> {
52 let mut reader = BitReader::new(rbsp);
53
54 let first_mb_in_slice = read_ue(&mut reader)?;
55
56 let slice_type_value = read_ue(&mut reader)?;
57 let slice_type = SliceType::from_value(slice_type_value)
58 .ok_or_else(|| Error::SliceParseError("Invalid slice type".into()))?;
59
60 let pic_parameter_set_id = read_ue(&mut reader)?;
61 if pic_parameter_set_id > 255 {
62 return Err(Error::SliceParseError("Invalid PPS ID".into()));
63 }
64
65 let mut colour_plane_id = 0;
66 if sps.separate_colour_plane_flag {
67 colour_plane_id = reader.read_bits(2)? as u8;
68 }
69
70 let frame_num_bits = sps.log2_max_frame_num_minus4 + 4;
71 let frame_num = reader.read_bits(frame_num_bits as u32)?;
72
73 let mut field_pic_flag = false;
74 let mut bottom_field_flag = false;
75
76 if !sps.frame_mbs_only_flag {
77 field_pic_flag = reader.read_flag()?;
78 if field_pic_flag {
79 bottom_field_flag = reader.read_flag()?;
80 }
81 }
82
83 let mut idr_pic_id = 0;
84 if nal_type == NalUnitType::IdrSlice {
85 idr_pic_id = read_ue(&mut reader)?;
86 }
87
88 let mut pic_order_cnt_lsb = 0;
89 let mut delta_pic_order_cnt_bottom = 0;
90 let mut delta_pic_order_cnt = [0, 0];
91
92 if sps.pic_order_cnt_type == 0 {
93 let pic_order_cnt_lsb_bits = sps.log2_max_pic_order_cnt_lsb_minus4 + 4;
94 pic_order_cnt_lsb = reader.read_bits(pic_order_cnt_lsb_bits as u32)?;
95
96 if pps.bottom_field_pic_order_in_frame_present_flag && !field_pic_flag {
97 delta_pic_order_cnt_bottom = read_se(&mut reader)?;
98 }
99 } else if sps.pic_order_cnt_type == 1 && !sps.delta_pic_order_always_zero_flag {
100 delta_pic_order_cnt[0] = read_se(&mut reader)?;
101
102 if pps.bottom_field_pic_order_in_frame_present_flag && !field_pic_flag {
103 delta_pic_order_cnt[1] = read_se(&mut reader)?;
104 }
105 }
106
107 let mut redundant_pic_cnt = 0;
108 if pps.redundant_pic_cnt_present_flag {
109 redundant_pic_cnt = read_ue(&mut reader)?;
110 }
111
112 let mut direct_spatial_mv_pred_flag = false;
113 if slice_type == SliceType::B {
114 direct_spatial_mv_pred_flag = reader.read_flag()?;
115 }
116
117 let mut num_ref_idx_active_override_flag = false;
118 let mut num_ref_idx_l0_active_minus1 = pps.num_ref_idx_l0_default_active_minus1 as u32;
119 let mut num_ref_idx_l1_active_minus1 = pps.num_ref_idx_l1_default_active_minus1 as u32;
120
121 if slice_type == SliceType::P || slice_type == SliceType::Sp || slice_type == SliceType::B {
122 num_ref_idx_active_override_flag = reader.read_flag()?;
123
124 if num_ref_idx_active_override_flag {
125 num_ref_idx_l0_active_minus1 = read_ue(&mut reader)?;
126
127 if slice_type == SliceType::B {
128 num_ref_idx_l1_active_minus1 = read_ue(&mut reader)?;
129 }
130 }
131 }
132
133 Ok(SliceHeader {
134 first_mb_in_slice,
135 slice_type,
136 pic_parameter_set_id: pic_parameter_set_id as u8,
137 colour_plane_id,
138 frame_num,
139 field_pic_flag,
140 bottom_field_flag,
141 idr_pic_id,
142 pic_order_cnt_lsb,
143 delta_pic_order_cnt_bottom,
144 delta_pic_order_cnt,
145 redundant_pic_cnt,
146 direct_spatial_mv_pred_flag,
147 num_ref_idx_active_override_flag,
148 num_ref_idx_l0_active_minus1,
149 num_ref_idx_l1_active_minus1,
150 })
151 }
152}
153
154#[derive(Debug, Clone, PartialEq, Eq, Hash)]
155pub struct PictureId {
156 pub frame_num: u32,
157 pub pic_parameter_set_id: u8,
158 pub idr_pic_id: Option<u32>,
159 pub pic_order_cnt_lsb: Option<u32>,
160 pub delta_pic_order_cnt: Option<[i32; 2]>,
161 pub field_pic_flag: bool,
162 pub bottom_field_flag: bool,
163}
164
165impl PictureId {
166 pub fn from_slice_header(header: &SliceHeader, nal_type: NalUnitType, sps: &Sps) -> Self {
167 let idr_pic_id = if nal_type == NalUnitType::IdrSlice {
168 Some(header.idr_pic_id)
169 } else {
170 None
171 };
172
173 let pic_order_cnt_lsb = if sps.pic_order_cnt_type == 0 {
174 Some(header.pic_order_cnt_lsb)
175 } else {
176 None
177 };
178
179 let delta_pic_order_cnt = if sps.pic_order_cnt_type == 1 {
180 Some(header.delta_pic_order_cnt)
181 } else {
182 None
183 };
184
185 PictureId {
186 frame_num: header.frame_num,
187 pic_parameter_set_id: header.pic_parameter_set_id,
188 idr_pic_id,
189 pic_order_cnt_lsb,
190 delta_pic_order_cnt,
191 field_pic_flag: header.field_pic_flag,
192 bottom_field_flag: header.bottom_field_flag,
193 }
194 }
195}
196
197#[cfg(test)]
198mod tests {
199 use super::*;
200
201 #[test]
202 fn test_slice_type_conversion() {
203 assert_eq!(SliceType::from_value(0), Some(SliceType::P));
204 assert_eq!(SliceType::from_value(1), Some(SliceType::B));
205 assert_eq!(SliceType::from_value(2), Some(SliceType::I));
206 assert_eq!(SliceType::from_value(5), Some(SliceType::P));
207 assert_eq!(SliceType::from_value(7), Some(SliceType::I));
208 }
209}