]> git.nihav.org Git - nihav.git/blame - nihav-itu/src/codecs/h264/decoder_mt.rs
avimux: do not record palette change chunks in OpenDML index
[nihav.git] / nihav-itu / src / codecs / h264 / decoder_mt.rs
CommitLineData
11d7aef2
KS
1use nihav_core::codecs::*;
2use nihav_core::io::bitreader::*;
3
4use super::*;
5use super::dispatch::*;
6
7pub struct FrameDecoder {
8 pub slices: Vec<(SliceHeader, usize, SliceRefs, Vec<u8>)>,
9 pub cur_pic: PictureInfo,
10 sps: Arc<SeqParameterSet>,
11 pps: Arc<PicParameterSet>,
12 pub num_mbs: usize,
13 mc_dsp: H264MC,
14 dispatch: Shareable<ThreadDispatcher>,
15 sstate: SliceState,
16 cavlc_cb: Arc<CAVLCTables>,
17 ipcm_buf: [u8; 256 + 64 + 64],
18 is_mbaff: bool,
19 deblock_skip: bool,
20}
21
22impl FrameDecoder {
23 pub fn decode_slice(&mut self, hdr: &SliceHeader, hdr_size: usize, refs: &SliceRefs, nal: &[u8]) -> DecoderResult<usize> {
24 self.sstate.reset(self.sps.pic_width_in_mbs, self.sps.pic_height_in_mbs, hdr.first_mb_in_slice);
25
26 let mut full_size = nal.len() * 8;
27 for &byte in nal.iter().rev() {
28 if byte == 0 {
29 full_size -= 8;
30 } else {
31 full_size -= (byte.trailing_zeros() + 1) as usize;
32 break;
33 }
34 }
35 validate!(full_size > 0);
36
5f223cdb
KS
37 let sslice_refs = SimplifiedSliceRefs::new(refs);
38
11d7aef2 39 let mut br = BitReader::new(&nal[hdr_size / 8..], BitReaderMode::BE);
fe64781d
KS
40 let mut dst_pic = self.cur_pic.clone();
41 let mut dst_frm = NASimpleVideoFrame::from_video_buf(&mut dst_pic.buf).unwrap();
11d7aef2
KS
42 if !self.pps.entropy_coding_mode {
43 br.skip((hdr_size & 7) as u32)?;
fe64781d 44 self.decode_slice_cavlc(&mut br, full_size - (hdr_size & !7), hdr, &sslice_refs, &mut dst_frm)
11d7aef2
KS
45 } else {
46 let csrc = &nal[(hdr_size + 7) / 8..];
47 validate!(csrc.len() >= 2);
48 let mut cabac = CABAC::new(csrc, hdr.slice_type, hdr.slice_qp, hdr.cabac_init_idc as usize)?;
fe64781d 49 self.decode_slice_cabac(&mut cabac, hdr, &sslice_refs, &mut dst_frm)
11d7aef2
KS
50 }
51 }
fe64781d 52 fn decode_slice_cavlc(&mut self, br: &mut BitReader, full_size: usize, slice_hdr: &SliceHeader, refs: &SimplifiedSliceRefs, frm: &mut NASimpleVideoFrame<u8>) -> DecoderResult<usize> {
11d7aef2
KS
53 const INTRA_CBP: [u8; 48] = [
54 47, 31, 15, 0, 23, 27, 29, 30, 7, 11, 13, 14, 39, 43, 45, 46,
55 16, 3, 5, 10, 12, 19, 21, 26, 28, 35, 37, 42, 44, 1, 2, 4,
56 8, 17, 18, 20, 24, 6, 9, 22, 25, 32, 33, 34, 36, 40, 38, 41
57 ];
58 const INTER_CBP: [u8; 48] = [
59 0, 16, 1, 2, 4, 8, 32, 3, 5, 10, 12, 15, 47, 7, 11, 13,
60 14, 6, 9, 31, 35, 37, 42, 44, 33, 34, 36, 40, 39, 43, 45, 46,
61 17, 18, 20, 24, 19, 21, 26, 28, 23, 27, 29, 30, 22, 25, 38, 41
62 ];
63
e6aaad5c 64 let mut mb_idx = slice_hdr.first_mb_in_slice;
11d7aef2
KS
65 let mut mb_info = CurrentMBInfo { qp_y: slice_hdr.slice_qp, ..Default::default() };
66 let skip_type = if slice_hdr.slice_type.is_p() { MBType::PSkip } else { MBType::BSkip };
67 while br.tell() < full_size && mb_idx < self.num_mbs {
68 mb_info.coded = [false; 25];
69 mb_info.ref_l0 = [ZERO_REF; 4];
70 mb_info.ref_l1 = [ZERO_REF; 4];
71 mb_info.mv_l0 = [ZERO_MV; 16];
72 mb_info.mv_l1 = [ZERO_MV; 16];
73 mb_info.chroma_dc = [[0; 4]; 2];
74 mb_info.cbpy = 0;
75 mb_info.cbpc = 0;
76
77 if !slice_hdr.slice_type.is_intra() {
78 let mb_skip_run = br.read_ue()? as usize;
79 validate!(mb_idx + mb_skip_run <= self.num_mbs);
80 mb_info.mb_type = skip_type;
81 for _ in 0..mb_skip_run {
fe64781d 82 self.handle_macroblock(slice_hdr, &mut mb_info, refs, frm)?;
11d7aef2
KS
83 mb_idx += 1;
84 }
85 if mb_idx == self.num_mbs || br.tell() >= full_size {
86 break;
87 }
88 }
89 if br.tell() < full_size {
90 if self.is_mbaff && ((mb_idx & 1) == 0) {
91 let _mb_field_decoding = br.read_bool()?;
92 }
93 let mut mb_type = decode_mb_type_cavlc(br, slice_hdr)?;
94 mb_info.mb_type = mb_type;
95 mb_info.transform_size_8x8 = false;
96 if mb_type == MBType::PCM {
97 br.align();
98 for pix in self.ipcm_buf[..256 + 64 + 64].iter_mut() {
99 *pix = br.read(8)? as u8;
100 }
101 self.sstate.fill_ncoded(16);
102 } else {
103 if self.pps.transform_8x8_mode && mb_type == MBType::Intra4x4 {
104 mb_info.transform_size_8x8 = br.read_bool()?;
105 if mb_info.transform_size_8x8 {
106 mb_type = MBType::Intra8x8;
107 mb_info.mb_type = MBType::Intra8x8;
108 }
109 }
110 decode_mb_pred_cavlc(br, slice_hdr, mb_type, &mut self.sstate, &mut mb_info)?;
111 let (cbpy, cbpc) = if let MBType::Intra16x16(_, cbpy, cbpc) = mb_type {
112 (cbpy, cbpc)
113 } else {
114 let cbp_id = br.read_ue()? as usize;
115 validate!(cbp_id < INTRA_CBP.len());
116 let cbp = if mb_type == MBType::Intra4x4 || mb_type == MBType::Intra8x8 {
117 INTRA_CBP[cbp_id]
118 } else {
119 INTER_CBP[cbp_id]
120 };
121 if self.pps.transform_8x8_mode && (cbp & 0xF) != 0 && mb_info.can_have_8x8_tx(self.sps.direct_8x8_inference) {
122 mb_info.transform_size_8x8 = br.read_bool()?;
123 }
124 ((cbp & 0xF), (cbp >> 4))
125 };
126 mb_info.cbpy = cbpy;
127 mb_info.cbpc = cbpc;
128 self.sstate.get_cur_mb().cbp = (cbpc << 4) | cbpy;
129 if cbpy != 0 || cbpc != 0 || mb_type.is_intra16x16() {
130 let mb_qp_delta = br.read_se()?;
131 validate!(mb_qp_delta >= -26 && mb_qp_delta <= 25);
132 let new_qp = mb_qp_delta + i32::from(mb_info.qp_y);
133 mb_info.qp_y = if new_qp < 0 {
134 (new_qp + 52) as u8
135 } else if new_qp >= 52 {
136 (new_qp - 52) as u8
137 } else {
138 new_qp as u8
139 };
140 mb_info.coeffs = [[0; 16]; 25];
141 if self.pps.transform_8x8_mode {
142 mb_info.clear_coeffs8x8();
143 }
144 mb_info.chroma_dc = [[0; 4]; 2];
145 decode_residual_cavlc(br, &mut self.sstate, &mut mb_info, &self.cavlc_cb)?;
146 }
147 }
fe64781d 148 self.handle_macroblock(slice_hdr, &mut mb_info, refs, frm)?;
11d7aef2
KS
149 }
150 mb_idx += 1;
151 if let Ok(disp) = self.dispatch.read() {
152 disp.update_pos(self.cur_pic.full_id, mb_idx);
153 }
154 }
155 Ok(mb_idx)
156 }
fe64781d 157 fn decode_slice_cabac(&mut self, cabac: &mut CABAC, slice_hdr: &SliceHeader, refs: &SimplifiedSliceRefs, frm: &mut NASimpleVideoFrame<u8>) -> DecoderResult<usize> {
e6aaad5c 158 let mut mb_idx = slice_hdr.first_mb_in_slice;
11d7aef2
KS
159 let mut prev_mb_skipped = false;
160 let skip_type = if slice_hdr.slice_type.is_p() { MBType::PSkip } else { MBType::BSkip };
161 let mut last_qp_diff = false;
162
163 let mut mb_info = CurrentMBInfo { qp_y: slice_hdr.slice_qp, ..Default::default() };
164
165 while mb_idx < self.num_mbs {
166 mb_info.coded = [false; 25];
167 mb_info.ref_l0 = [ZERO_REF; 4];
168 mb_info.ref_l1 = [ZERO_REF; 4];
169 mb_info.mv_l0 = [ZERO_MV; 16];
170 mb_info.mv_l1 = [ZERO_MV; 16];
171 mb_info.chroma_dc = [[0; 4]; 2];
172 mb_info.cbpy = 0;
173 mb_info.cbpc = 0;
174 let mb_skip = cabac_decode_mbskip(cabac, &self.sstate, slice_hdr);
175 if !mb_skip {
176 if self.is_mbaff && (((mb_idx & 1) == 0) || (prev_mb_skipped && ((mb_idx & 1) == 1))) {
177 let _mb_field_decoding = cabac.decode_bit(70);
178 }
179 let mut mb_type = cabac_decode_mb_type(cabac, slice_hdr, &self.sstate);
180 mb_info.mb_type = mb_type;
181 mb_info.transform_size_8x8 = false;
182 if mb_type == MBType::PCM {
183 let ipcm_size = 256 + 64 + 64;
184 validate!(cabac.pos + ipcm_size <= cabac.src.len());
185 self.ipcm_buf[..ipcm_size].copy_from_slice(&cabac.src[cabac.pos..][..ipcm_size]);
186 cabac.pos += ipcm_size;
187 cabac.reinit()?;
188 last_qp_diff = false;
189 } else {
190 if self.pps.transform_8x8_mode && mb_type == MBType::Intra4x4 {
191 let mut ctx = 0;
192 if self.sstate.get_top_mb().transform_8x8 {
193 ctx += 1;
194 }
195 if self.sstate.get_left_mb().transform_8x8 {
196 ctx += 1;
197 }
198 mb_info.transform_size_8x8 = cabac.decode_bit(399 + ctx);
199 if mb_info.transform_size_8x8 {
200 mb_type = MBType::Intra8x8;
201 mb_info.mb_type = MBType::Intra8x8;
202 }
203 }
204 decode_mb_pred_cabac(cabac, slice_hdr, mb_type, &mut self.sstate, &mut mb_info);
205 let (cbpy, cbpc) = if let MBType::Intra16x16(_, cbpy, cbpc) = mb_type {
206 (cbpy, cbpc)
207 } else {
208 decode_cbp_cabac(cabac, &self.sstate)
209 };
210 if self.pps.transform_8x8_mode && cbpy != 0 && mb_info.can_have_8x8_tx(self.sps.direct_8x8_inference) {
211 let mut ctx = 0;
212 if self.sstate.get_top_mb().transform_8x8 {
213 ctx += 1;
214 }
215 if self.sstate.get_left_mb().transform_8x8 {
216 ctx += 1;
217 }
218 mb_info.transform_size_8x8 = cabac.decode_bit(399 + ctx);
219 }
220 if mb_type.is_intra() {
221 self.sstate.get_cur_mb().cmode = mb_info.chroma_ipred;
222 }
223 mb_info.cbpy = cbpy;
224 mb_info.cbpc = cbpc;
225 self.sstate.get_cur_mb().cbp = (cbpc << 4) | cbpy;
226 if cbpy != 0 || cbpc != 0 || mb_type.is_intra16x16() {
227 let mb_qp_delta = decode_mb_qp_delta_cabac(cabac, last_qp_diff as usize);
228 validate!(mb_qp_delta >= -26 && mb_qp_delta <= 25);
229 last_qp_diff = mb_qp_delta != 0;
230 let new_qp = mb_qp_delta + i32::from(mb_info.qp_y);
231 mb_info.qp_y = if new_qp < 0 {
232 (new_qp + 52) as u8
233 } else if new_qp >= 52 {
234 (new_qp - 52) as u8
235 } else {
236 new_qp as u8
237 };
238 mb_info.coeffs = [[0; 16]; 25];
239 if self.pps.transform_8x8_mode {
240 mb_info.clear_coeffs8x8();
241 }
242 mb_info.chroma_dc = [[0; 4]; 2];
243 decode_residual_cabac(cabac, &mut self.sstate, &mut mb_info);
244 } else {
245 last_qp_diff = false;
246 }
247 }
248 } else {
249 mb_info.mb_type = skip_type;
250 mb_info.transform_size_8x8 = false;
251 last_qp_diff = false;
252 }
fe64781d 253 self.handle_macroblock(slice_hdr, &mut mb_info, refs, frm)?;
11d7aef2
KS
254 prev_mb_skipped = mb_skip;
255 if !(self.is_mbaff && ((mb_idx & 1) == 0)) && cabac.decode_terminate() {
256 if let Ok(disp) = self.dispatch.read() {
257 disp.update_pos(self.cur_pic.full_id, mb_idx + 1);
258 }
259 return Ok(mb_idx + 1);
260 }
261 mb_idx += 1;
262 if let Ok(disp) = self.dispatch.read() {
263 disp.update_pos(self.cur_pic.full_id, mb_idx);
264 }
265 }
266 Err(DecoderError::InvalidData)
267 }
268 #[allow(clippy::cognitive_complexity)]
fe64781d 269 fn handle_macroblock(&mut self, slice_hdr: &SliceHeader, mb_info: &mut CurrentMBInfo, refs: &SimplifiedSliceRefs, frm: &mut NASimpleVideoFrame<u8>) -> DecoderResult<()> {
11d7aef2
KS
270 let qp_y = mb_info.qp_y;
271 let qpr = ((qp_y as i8) + self.pps.chroma_qp_index_offset).max(0).min(51) as usize;
272 let qp_u = CHROMA_QUANTS[qpr];
273 let qpb = ((qp_y as i8) + self.pps.second_chroma_qp_index_offset).max(0).min(51) as usize;
274 let qp_v = CHROMA_QUANTS[qpb];
275
276 let tx_bypass = qp_y == 0 && self.sps.qpprime_y_zero_transform_bypass;
277
278 self.sstate.get_cur_mb().mb_type = mb_info.mb_type.into();
279 if mb_info.mb_type != MBType::PCM {
280 self.sstate.get_cur_mb().qp_y = qp_y;
281 self.sstate.get_cur_mb().qp_u = qp_u;
282 self.sstate.get_cur_mb().qp_v = qp_v;
283 self.sstate.get_cur_mb().transform_8x8 = mb_info.transform_size_8x8;
284 }
285 let has_dc = mb_info.mb_type.is_intra16x16() && mb_info.coded[24];
286 if has_dc {
287 idct_luma_dc(&mut mb_info.coeffs[24], qp_y);
288 for i in 0..16 {
289 mb_info.coeffs[i][0] = mb_info.coeffs[24][i];
290 }
291 }
292 if !mb_info.transform_size_8x8 {
293 let quant_dc = !mb_info.mb_type.is_intra16x16();
fe64781d
KS
294 if quant_dc {
295 for i in 0..16 {
296 if mb_info.coded[i] {
297 if !tx_bypass {
298 idct(&mut mb_info.coeffs[i], qp_y);
299 }
300 } else if has_dc {
301 if !tx_bypass {
302 idct_dc(&mut mb_info.coeffs[i], qp_y, quant_dc);
303 }
304 mb_info.coded[i] = true;
11d7aef2 305 }
fe64781d
KS
306 }
307 } else {
308 for i in 0..16 {
309 if mb_info.coded[i] {
310 if !tx_bypass {
311 idct_skip_dc(&mut mb_info.coeffs[i], qp_y);
312 }
313 } else if has_dc {
314 if !tx_bypass {
315 idct_dc(&mut mb_info.coeffs[i], qp_y, quant_dc);
316 }
317 mb_info.coded[i] = true;
11d7aef2 318 }
11d7aef2
KS
319 }
320 }
321 } else {
322 for i in 0..4 {
323 if mb_info.coded[(i & 1) * 2 + (i & 2) * 4] && !tx_bypass {
324 dequant8x8(&mut mb_info.coeffs8x8[i].coeffs, &self.pps.scaling_list_8x8[!mb_info.mb_type.is_intra() as usize]);
325 idct8x8(&mut mb_info.coeffs8x8[i].coeffs, qp_y);
326 }
327 }
328 }
329 for chroma in 0..2 {
330 let qp_c = if chroma == 0 { qp_u } else { qp_v };
331 if mb_info.cbpc != 0 {
332 chroma_dc_transform(&mut mb_info.chroma_dc[chroma], qp_c);
333 }
334 for i in 0..4 {
335 let blk_no = 16 + chroma * 4 + i;
336 mb_info.coeffs[blk_no][0] = mb_info.chroma_dc[chroma][i];
337 if mb_info.coded[blk_no] {
fe64781d 338 idct_skip_dc(&mut mb_info.coeffs[blk_no], qp_c);
11d7aef2
KS
339 } else if mb_info.coeffs[blk_no][0] != 0 {
340 idct_dc(&mut mb_info.coeffs[blk_no], qp_c, false);
341 mb_info.coded[blk_no] = true;
342 }
343 }
344 }
345 if !self.pps.entropy_coding_mode || mb_info.mb_type.is_skip() || mb_info.mb_type.is_intra() {
346 self.sstate.reset_mb_mv();
347 }
348 if !mb_info.mb_type.is_intra() {
349 let temporal_mv = !slice_hdr.direct_spatial_mv_pred;
350 let cur_id = self.cur_pic.full_id as u16;
351 // wait for the reference macroblock MV to be available
352 if matches!(mb_info.mb_type, MBType::Direct | MBType::BSkip | MBType::B8x8) {
353 if let Some(ref_id) = refs.get_ref_id(0, mb_info.ref_l1[0].index()) {
354 wait_for_mb(&self.dispatch, &self.sstate, self.sstate.mb_x * 16, self.sstate.mb_y * 16, ZERO_MV, ref_id)?;
355 }
356 }
357 Self::pred_mv(&mut self.sstate, refs, mb_info, cur_id, temporal_mv, self.sps.direct_8x8_inference);
358 }
359 if !self.pps.constrained_intra_pred && mb_info.mb_type != MBType::Intra4x4 && mb_info.mb_type != MBType::Intra8x8 {
360 self.sstate.fill_ipred(IntraPredMode::DC);
361 }
362
363 let xpos = self.sstate.mb_x * 16;
364 let ypos = self.sstate.mb_y * 16;
11d7aef2
KS
365 if mb_info.mb_type != MBType::PCM {
366 let weight_mode = if self.pps.weighted_pred && slice_hdr.slice_type.is_p() {
367 1
368 } else if slice_hdr.slice_type.is_b() {
369 self.pps.weighted_bipred_idc
370 } else {
371 0
372 };
fe64781d 373 recon_mb_mt(frm, slice_hdr, mb_info, &mut self.sstate, refs, &mut self.mc_dsp, weight_mode, &self.dispatch)?;
11d7aef2
KS
374 } else {
375 for (dline, src) in frm.data[frm.offset[0] + xpos + ypos * frm.stride[0]..].chunks_mut(frm.stride[0]).take(16).zip(self.ipcm_buf.chunks(16)) {
376 dline[..16].copy_from_slice(src);
377 }
378 for (dline, src) in frm.data[frm.offset[1] + xpos/2 + ypos/2 * frm.stride[1]..].chunks_mut(frm.stride[1]).take(8).zip(self.ipcm_buf[256..].chunks(8)) {
379 dline[..8].copy_from_slice(src);
380 }
381 for (dline, src) in frm.data[frm.offset[2] + xpos/2 + ypos/2 * frm.stride[2]..].chunks_mut(frm.stride[2]).take(8).zip(self.ipcm_buf[256 + 64..].chunks(8)) {
382 dline[..8].copy_from_slice(src);
383 }
384 }
fe64781d 385 self.sstate.save_ipred_context(frm);
11d7aef2
KS
386
387 let mv_info = &mut self.cur_pic.mv_info;
388 let mb_pos = self.sstate.mb_x + self.sstate.mb_y * mv_info.mb_stride;
389 let mut mb = FrameMBInfo::new();
390 mb.mb_type = mb_info.mb_type.into();
391 for blk4 in 0..16 {
392 mb.mv[blk4] = self.sstate.get_cur_blk4(blk4).mv;
393 }
394 for blk8 in 0..4 {
395 mb.ref_poc[blk8] = refs.map_refs(self.sstate.get_cur_blk8(blk8).ref_idx);
396 mb.ref_idx[blk8] = self.sstate.get_cur_blk8(blk8).ref_idx;
397 }
398 mv_info.mbs[mb_pos] = mb;
399
400 let deblock_mode = slice_hdr.disable_deblocking_filter_idc;
401 if !self.deblock_skip && deblock_mode != 1 {
402 let is_s = slice_hdr.slice_type == SliceType::SI || slice_hdr.slice_type == SliceType::SP;
403 self.sstate.fill_deblock(refs, deblock_mode, is_s);
404 let mut frm = NASimpleVideoFrame::from_video_buf(&mut self.cur_pic.buf).unwrap();
405 let lf_alpha = slice_hdr.slice_alpha_c0_offset;
406 let lf_beta = slice_hdr.slice_beta_offset;
407 loop_filter_mb(&mut frm, &self.sstate, lf_alpha, lf_beta);
408 }
409 self.sstate.next_mb();
410 Ok(())
411 }
412
5f223cdb 413 fn pred_mv(sstate: &mut SliceState, frame_refs: &SimplifiedSliceRefs, mb_info: &mut CurrentMBInfo, cur_id: u16, temporal_mv: bool, direct_8x8: bool) {
11d7aef2
KS
414 let mb_type = mb_info.mb_type;
415 if !mb_type.is_4x4() {
416 let (pw, ph) = mb_type.size();
417 let mut xoff = 0;
418 let mut yoff = 0;
419 if mb_type == MBType::Direct || mb_type == MBType::BSkip {
420 sstate.predict_direct_mb(frame_refs, temporal_mv, direct_8x8, cur_id);
421 }
422 for part in 0..mb_type.num_parts() {
423 if !mb_type.is_l1(part) {
424 match mb_type {
425 MBType::PSkip => sstate.predict_pskip(),
426 MBType::BSkip | MBType::Direct => {
427 },
428 _ => {
429 sstate.predict(xoff, yoff, pw, ph, 0,
430 mb_info.mv_l0[part], mb_info.ref_l0[part]);
431 },
432 };
433 }
434 if !mb_type.is_l0(part) && mb_type != MBType::BSkip && mb_type != MBType::Direct {
435 sstate.predict(xoff, yoff, pw, ph, 1, mb_info.mv_l1[part], mb_info.ref_l1[part]);
436 }
437 if pw != 16 {
438 xoff += pw;
439 } else {
440 yoff += ph;
441 }
442 }
443 } else {
444 for part in 0..4 {
445 let sub_type = mb_info.sub_mb_type[part];
446 let mut xoff = (part & 1) * 8;
447 let mut yoff = (part & 2) * 4;
448 let orig_x = xoff;
449 let (pw, ph) = sub_type.size();
450 for subpart in 0..sub_type.num_parts() {
451 if sub_type != SubMBType::Direct8x8 {
452 if !sub_type.is_l1() {
453 sstate.predict(xoff, yoff, pw, ph, 0, mb_info.mv_l0[part * 4 + subpart], mb_info.ref_l0[part]);
454 }
455 if !sub_type.is_l0() {
456 sstate.predict(xoff, yoff, pw, ph, 1, mb_info.mv_l1[part * 4 + subpart], mb_info.ref_l1[part]);
457 }
458 } else {
459 for sblk in 0..4 {
460 sstate.predict_direct_sub(frame_refs, temporal_mv, direct_8x8, cur_id, (xoff / 4) + (sblk & 1) + (yoff / 4) * 4 + (sblk & 2) * 2);
461 }
462 }
463 xoff += pw;
464 if xoff == orig_x + 8 {
465 xoff -= 8;
466 yoff += ph;
467 }
468 }
469 }
470 }
471 }
472}
473
474struct H264MTDecoder {
475 info: NACodecInfoRef,
476 nal_len: u8,
477 dispatch: Shareable<ThreadDispatcher>,
478 frame_refs: FrameRefs,
479 skip_mode: FrameSkipMode,
480 sps: Vec<Arc<SeqParameterSet>>,
481 cur_sps: usize,
482 pps: Vec<Arc<PicParameterSet>>,
483 cur_pps: usize,
484 cur_fdec: Option<FrameDecoder>,
485 cavlc_cb: Arc<CAVLCTables>,
486 deblock_skip: bool,
487 max_last_poc: u32,
488 poc_base: u32,
25f6045b
KS
489 disp_w: usize,
490 disp_h: usize,
11d7aef2
KS
491}
492
493impl H264MTDecoder {
494 fn new() -> Self {
495 Self {
496 info: NACodecInfoRef::default(),
497 nal_len: 0,
498 dispatch: Arc::new(RwLock::new(ThreadDispatcher::new())),
499 frame_refs: FrameRefs::new(),
500 skip_mode: FrameSkipMode::default(),
501 sps: Vec::new(),
502 cur_sps: 0,
503 pps: Vec::new(),
504 cur_pps: 0,
505 cur_fdec: None,
506 cavlc_cb: Arc::new(CAVLCTables::new()),
507 deblock_skip: false,
508 max_last_poc: 0,
509 poc_base: 0,
25f6045b
KS
510 disp_w: 0,
511 disp_h: 0,
11d7aef2
KS
512 }
513 }
514 fn handle_nal(&mut self, src: Vec<u8>, supp: &mut NADecoderSupport, skip_decoding: bool, user_id: u32, time: NATimeInfo) -> DecoderResult<()> {
515 validate!(!src.is_empty());
516 validate!((src[0] & 0x80) == 0);
517 let nal_ref_idc = src[0] >> 5;
518 let nal_unit_type = src[0] & 0x1F;
519
520 let mut full_size = src.len() * 8;
521 for &byte in src.iter().rev() {
522 if byte == 0 {
523 full_size -= 8;
524 } else {
525 full_size -= (byte.trailing_zeros() + 1) as usize;
526 break;
527 }
528 }
529 validate!(full_size > 0);
530 match nal_unit_type {
531 1 | 5 if !skip_decoding => {
532 let is_idr = nal_unit_type == 5;
533 let mut br = BitReader::new(&src[..(full_size + 7)/8], BitReaderMode::BE);
534 br.skip(8)?;
535
536 let slice_hdr = parse_slice_header(&mut br, self.sps.as_slice(), self.pps.as_slice(), is_idr, nal_ref_idc)?;
537 let hdr_size = br.tell();
538 validate!(br.tell() < full_size);
539 let full_id;
540 if slice_hdr.first_mb_in_slice == 0 {
541 validate!(self.cur_fdec.is_none());
542 for (i, pps) in self.pps.iter().enumerate() {
543 if pps.pic_parameter_set_id == slice_hdr.pic_parameter_set_id {
544 self.cur_pps = i;
545 break;
546 }
547 }
548 for (i, sps) in self.sps.iter().enumerate() {
549 if sps.seq_parameter_set_id == self.pps[self.cur_pps].seq_parameter_set_id {
550 self.cur_sps = i;
551 break;
552 }
553 }
554
555 let mut cur_full_id = self.frame_refs.calc_picture_num(&slice_hdr, is_idr, nal_ref_idc, &self.sps[self.cur_sps]) + self.poc_base;
556 if is_idr {
557 if cur_full_id <= self.max_last_poc {
558 self.poc_base = self.max_last_poc + 2 - (cur_full_id - self.poc_base);
559 cur_full_id = self.max_last_poc + 2;
560 }
561 }
562 self.max_last_poc = self.max_last_poc.max(cur_full_id);
563 full_id = cur_full_id;
564
565 let sps = &self.sps[self.cur_sps];
566 if sps.chroma_format_idc != 1 || sps.bit_depth_luma != 8 || sps.bit_depth_chroma != 8 {
567 println!(" chroma fmt {} bits {}/{}", sps.chroma_format_idc, sps.bit_depth_luma, sps.bit_depth_chroma);
568 return Err(DecoderError::NotImplemented);
569 }
570
571 if is_idr {
572 self.frame_refs.clear_refs();
573 }
574
575 let width = sps.pic_width_in_mbs << 4;
576 let height = sps.pic_height_in_mbs << 4;
577 let num_mbs = sps.pic_width_in_mbs * sps.pic_height_in_mbs;
578
71efc2dc 579 let mut mc_dsp = H264MC::new();
11d7aef2
KS
580 mc_dsp.set_dimensions(width, height);
581
582 let is_mbaff = sps.mb_adaptive_frame_field && !slice_hdr.field_pic;
583 if is_mbaff {
584 println!("MBAFF");
585 return Err(DecoderError::NotImplemented);
586 }
587 if !sps.frame_mbs_only {
588 println!("PAFF?");
589 return Err(DecoderError::NotImplemented);
590 }
591
592 let cur_vinfo = supp.pool_u8.get_info();
25f6045b
KS
593 let (w, h) = if ((self.disp_w + 15) & !15) == width && ((self.disp_h + 15) & !15) == height {
594 (self.disp_w, self.disp_h)
595 } else {
596 (width, height)
597 };
598 let tmp_vinfo = NAVideoInfo::new(w, h, false, YUV420_FORMAT);
11d7aef2
KS
599 if cur_vinfo != Some(tmp_vinfo) {
600 supp.pool_u8.reset();
601 supp.pool_u8.prealloc_video(tmp_vinfo, 4)?;
602 }
603
604 let buf = if let Some(pic) = supp.pool_u8.get_free() {
605 pic
606 } else {
607 if supp.pool_u8.get_num_used() > 256 {
608 return Err(DecoderError::AllocError);
609 }
610 if let Ok(nbuf) = alloc_video_buffer(tmp_vinfo, 4) {
611 let vbuf = nbuf.get_vbuf().unwrap();
612 supp.pool_u8.add_frame(vbuf.clone());
613 vbuf
614 } else {
615 return Err(DecoderError::AllocError);
616 }
617 };
618
619 let cur_pic = PictureInfo {
620 id: slice_hdr.frame_num,
621 full_id, user_id, time,
622 pic_type: slice_hdr.slice_type.to_frame_type(),
623 buf,
624 cur_mb: 0,
625 is_ref: nal_ref_idc != 0,
626 is_idr,
627 long_term: get_long_term_id(is_idr, &slice_hdr),
628 mv_info: NABufferRef::new(FrameMV::new(sps.pic_width_in_mbs, sps.pic_height_in_mbs)),
629 };
630
631 self.cur_fdec = Some(FrameDecoder{
632 slices: Vec::new(),
633 sstate: SliceState::new(),
634 ipcm_buf: [0; 256 + 64 + 64],
635 //width, height,
636 num_mbs,
637 sps: Arc::clone(sps),
638 pps: Arc::clone(&self.pps[self.cur_pps]),
639 dispatch: Arc::clone(&self.dispatch),
640 cavlc_cb: Arc::clone(&self.cavlc_cb),
641 mc_dsp,
642 cur_pic,
643 is_mbaff,
644 deblock_skip: self.deblock_skip,
645 });
646 } else {
647 if let Some(ref mut fdec) = self.cur_fdec {
648 let new_type = slice_hdr.slice_type.to_frame_type();
649 let pic = &mut fdec.cur_pic;
650 pic.pic_type = match (pic.pic_type, new_type) {
651 (FrameType::I, _) => new_type,
652 (_, FrameType::B) => FrameType::B,
653 _ => pic.pic_type,
654 };
655 full_id = pic.full_id;
656 } else {
657 return Ok(());
658 }
659 }
660
661 let sps = &self.sps[self.cur_sps];
662
663 self.frame_refs.select_refs(sps, &slice_hdr, full_id);
664
665 if slice_hdr.adaptive_ref_pic_marking_mode {
666 self.frame_refs.apply_adaptive_marking(&slice_hdr.adaptive_ref_pic_marking, slice_hdr.frame_num, 1 << self.sps[self.cur_sps].log2_max_frame_num)?;
667 }
668 if let Some(ref mut fdec) = self.cur_fdec {
669 fdec.slices.push((slice_hdr, hdr_size, self.frame_refs.cur_refs.clone(), src));
670 }
671 },
672 2 => { // slice data partition A
673 //slice header
674 //slice id = read_ue()
675 //cat 2 slice data (all but MB layer residual)
676 return Err(DecoderError::NotImplemented);
677 },
678 3 => { // slice data partition B
679 //slice id = read_ue()
680 //if pps.redundant_pic_cnt_present { redundant_pic_cnt = read_ue() }
681 //cat 3 slice data (MB layer residual)
682 return Err(DecoderError::NotImplemented);
683 },
684 4 => { // slice data partition C
685 //slice id = read_ue()
686 //if pps.redundant_pic_cnt_present { redundant_pic_cnt = read_ue() }
687 //cat 4 slice data (MB layer residual)
688 return Err(DecoderError::NotImplemented);
689 },
690 6 => {}, //SEI
691 7 => {
692 let sps = parse_sps(&src[1..])?;
693 self.sps.push(Arc::new(sps));
694 },
695 8 => {
696 validate!(full_size >= 8 + 16);
697 let pps = parse_pps(&src[1..], self.sps.as_slice(), full_size - 8)?;
698 let mut found = false;
699 for stored_pps in self.pps.iter_mut() {
700 if stored_pps.pic_parameter_set_id == pps.pic_parameter_set_id {
701 *stored_pps = Arc::clone(&pps);
702 found = true;
703 break;
704 }
705 }
706 if !found {
707 self.pps.push(pps);
708 }
709 },
710 9 => { // access unit delimiter
711 },
712 10 => {}, //end of sequence
713 11 => {}, //end of stream
714 12 => {}, //filler
715 _ => {},
716 };
717
718 Ok(())
719 }
720}
721
722impl NADecoderMT for H264MTDecoder {
723 fn init(&mut self, supp: &mut NADecoderSupport, info: NACodecInfoRef, nthreads: usize) -> DecoderResult<()> {
724 if let NACodecTypeInfo::Video(vinfo) = info.get_properties() {
725 let fmt = YUV420_FORMAT;
726 let myinfo = NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, false, fmt));
727 self.info = NACodecInfo::new_ref(info.get_name(), myinfo, info.get_extradata()).into_ref();
728
729 let edata = info.get_extradata().unwrap();
730//print!("edata:"); for &el in edata.iter() { print!(" {:02X}", el); } println!();
731 if edata.len() > 11 && &edata[0..4] == b"avcC" {
732 let mut mr = MemoryReader::new_read(edata.as_slice());
733 let mut br = ByteReader::new(&mut mr);
734
735 br.read_skip(4)?;
736 let version = br.read_byte()?;
737 validate!(version == 1);
738 let profile = br.read_byte()?;
739 let _compatibility = br.read_byte()?;
740 let _level = br.read_byte()?;
741 let b = br.read_byte()?;
344e803e 742 //validate!((b & 0xFC) == 0xFC);
11d7aef2
KS
743 self.nal_len = (b & 3) + 1;
744 let b = br.read_byte()?;
344e803e 745 //validate!((b & 0xE0) == 0xE0);
11d7aef2
KS
746 let num_sps = (b & 0x1F) as usize;
747 for _ in 0..num_sps {
748 let len = br.read_u16be()? as usize;
749 let offset = br.tell() as usize;
750 validate!((br.peek_byte()? & 0x1F) == 7);
751 let mut nal_buf = Vec::new();
752 let _size = unescape_nal(&edata[offset..][..len], &mut nal_buf);
753 self.handle_nal(nal_buf, supp, true, 0, NATimeInfo::new(None, None, None, 0, 0))?;
754 br.read_skip(len)?;
755 }
756 let num_pps = br.read_byte()? as usize;
757 for _ in 0..num_pps {
758 let len = br.read_u16be()? as usize;
759 let offset = br.tell() as usize;
760 validate!((br.peek_byte()? & 0x1F) == 8);
761 let mut nal_buf = Vec::new();
762 let _size = unescape_nal(&edata[offset..][..len], &mut nal_buf);
763 self.handle_nal(nal_buf, supp, true, 0, NATimeInfo::new(None, None, None, 0, 0))?;
764 br.read_skip(len)?;
765 }
766 if br.left() > 0 {
767 match profile {
768 100 | 110 | 122 | 144 => {
769 let b = br.read_byte()?;
e12dabcd
KS
770 // some encoders put something different here
771 if (b & 0xFC) != 0xFC {
772 return Ok(());
773 }
11d7aef2
KS
774 // b & 3 -> chroma format
775 let b = br.read_byte()?;
776 validate!((b & 0xF8) == 0xF8);
777 // b & 7 -> luma depth minus 8
778 let b = br.read_byte()?;
779 validate!((b & 0xF8) == 0xF8);
780 // b & 7 -> chroma depth minus 8
781 let num_spsext = br.read_byte()? as usize;
782 for _ in 0..num_spsext {
783 let len = br.read_u16be()? as usize;
784 // parse spsext
785 br.read_skip(len)?;
786 }
787 },
788 _ => {},
789 };
790 }
791 } else {
792 return Err(DecoderError::NotImplemented);
793 }
794
795 let mut width = vinfo.get_width();
796 let mut height = vinfo.get_height();
25f6045b
KS
797 self.disp_w = width;
798 self.disp_h = height;
11d7aef2
KS
799
800 if (width == 0 || height == 0) && !self.sps.is_empty() {
801 width = self.sps[0].pic_width_in_mbs * 16;
802 height = self.sps[0].pic_height_in_mbs * 16;
803 }
804
805 let num_bufs = if !self.sps.is_empty() {
e6aaad5c 806 self.sps[0].num_ref_frames + 1
11d7aef2
KS
807 } else {
808 3
809 }.max(16 + 1);
810 if let Ok(ref mut sd) = self.dispatch.write() {
811 sd.max_threads = nthreads;
812 } else {
813 return Err(DecoderError::Bug);
814 }
815 supp.pool_u8.set_dec_bufs(num_bufs + nthreads);
816 supp.pool_u8.prealloc_video(NAVideoInfo::new(width, height, false, fmt), 4)?;
817
818 Ok(())
819 } else {
820 Err(DecoderError::InvalidData)
821 }
822 }
823 fn can_take_input(&mut self) -> bool {
824 if let Ok(ref sd) = self.dispatch.read() {
825 sd.can_decode_more()
826 } else {
827 false
828 }
829 }
830 fn queue_pkt(&mut self, supp: &mut NADecoderSupport, pkt: &NAPacket, user_id: u32) -> DecoderResult<bool> {
831 if !self.can_take_input() {
832 return Ok(false);
833 }
834
835 let src = pkt.get_buffer();
836
837 let mut mr = MemoryReader::new_read(&src);
838 let mut br = ByteReader::new(&mut mr);
839 let mut nal_buf = Vec::with_capacity(src.len());
840
841 if self.nal_len > 0 {
842 let mut skip_decoding = false;
843 if self.skip_mode != FrameSkipMode::None {
844 let mut pic_type = FrameType::I;
845 let mut is_ref = false;
846 while br.left() > 0 {
847 let size = match self.nal_len {
848 1 => br.read_byte()? as usize,
849 2 => br.read_u16be()? as usize,
850 3 => br.read_u24be()? as usize,
851 4 => br.read_u32be()? as usize,
852 _ => unreachable!(),
853 };
854 validate!(br.left() >= (size as i64));
855 let offset = br.tell() as usize;
856 let size = unescape_nal(&src[offset..][..size], &mut nal_buf);
857 validate!(size > 0);
858 let nal_ref_idc = nal_buf[0] >> 5;
859 let nal_unit_type = nal_buf[0] & 0x1F;
860 if nal_unit_type == 1 || nal_unit_type == 5 {
861 let mut bitr = BitReader::new(&nal_buf[1..], BitReaderMode::BE);
862 let (first_mb, slice_type) = parse_slice_header_minimal(&mut bitr)?;
863 if first_mb == 0 && nal_ref_idc != 0 {
864 is_ref = true;
865 }
866 let new_type = slice_type.to_frame_type();
867 pic_type = match (pic_type, new_type) {
868 (FrameType::I, _) => new_type,
869 (_, FrameType::B) => FrameType::B,
870 _ => pic_type,
871 };
872 }
873 br.read_skip(size)?;
874 }
875 match self.skip_mode {
876 FrameSkipMode::IntraOnly => {
877 skip_decoding = pic_type != FrameType::I;
878 },
879 FrameSkipMode::KeyframesOnly => {
880 if !is_ref {
881 skip_decoding = true;
882 }
883 },
884 _ => {},
885 };
886 br.seek(SeekFrom::Start(0))?;
887 }
888
889 let mut initial_ref_frames = Vec::new();
890 self.frame_refs.fill_ref_nums(&mut initial_ref_frames);
891
892 while br.left() > 0 {
893 let size = match self.nal_len {
894 1 => br.read_byte()? as usize,
895 2 => br.read_u16be()? as usize,
896 3 => br.read_u24be()? as usize,
897 4 => br.read_u32be()? as usize,
898 _ => unreachable!(),
899 };
900 validate!(br.left() >= (size as i64));
901 let offset = br.tell() as usize;
902 let mut cur_nal_buf = Vec::with_capacity(size);
903 let _size = unescape_nal(&src[offset..][..size], &mut cur_nal_buf);
904 self.handle_nal(cur_nal_buf, supp, skip_decoding, user_id, pkt.ts)?;
905 br.read_skip(size)?;
906 }
907 let mut fdec = None;
908 std::mem::swap(&mut fdec, &mut self.cur_fdec);
909 if let Some(fdc) = fdec {
910 let cpic = &fdc.cur_pic;
911 if cpic.is_ref {
912 self.frame_refs.add_short_term(cpic.clone(), self.sps[self.cur_sps].num_ref_frames);
913 }
914 if let Some(lt_idx) = cpic.long_term {
915 self.frame_refs.add_long_term(lt_idx, cpic.clone());
916 }
917 let mut ref_frames = Vec::new();
918 self.frame_refs.fill_ref_nums(&mut ref_frames);
919 queue_decoding(&mut self.dispatch, fdc, &initial_ref_frames, &ref_frames);
920 }
921 } else {
922//todo NAL detection
923 unimplemented!();
924 }
925 Ok(true)
926 }
927 fn has_output(&mut self) -> bool {
928 if let Ok(ref ds) = self.dispatch.read() {
929 ds.has_output()
930 } else {
931 panic!("can't peek into status");
932 }
933 }
934 fn get_frame(&mut self) -> (DecoderResult<NAFrameRef>, u32) {
935 match wait_for_one(&mut self.dispatch) {
936 Ok(cpic) => {
937 let bufinfo = NABufferType::Video(cpic.buf.clone());
938 let ftype = cpic.pic_type;
939 let dts = Some(u64::from(cpic.full_id));
940 let mut frm = NAFrame::new(cpic.time, ftype, cpic.is_idr, self.info.clone(), bufinfo);
941 if let (Some(mydts), None) = (dts, frm.get_dts()) {
942 frm.set_dts(Some(mydts));
943 }
944 frm.set_id(cpic.user_id as i64);
945 (Ok(frm.into_ref()), cpic.user_id)
946 },
947 Err((err, id)) => (Err(err), id),
948 }
949 }
950 fn flush(&mut self) {
951 clear_threads(&mut self.dispatch);
f2cb96ff 952 self.frame_refs.clear_refs();
11d7aef2
KS
953 }
954}
955
956impl NAOptionHandler for H264MTDecoder {
957 fn get_supported_options(&self) -> &[NAOptionDefinition] { DECODER_OPTIONS }
958 fn set_options(&mut self, options: &[NAOption]) {
959 for option in options.iter() {
960 for opt_def in DECODER_OPTIONS.iter() {
961 if opt_def.check(option).is_ok() {
962 match (option.name, &option.value) {
963 (FRAME_SKIP_OPTION, NAValue::String(ref strval)) => {
964 if let Ok(smode) = FrameSkipMode::from_str(strval) {
965 self.skip_mode = smode;
966 }
967 },
968 (DEBLOCK_SKIP_OPTION, NAValue::Bool(val)) => {
969 self.deblock_skip = *val;
970 },
971 _ => {},
972 }
973 }
974 }
975 }
976 }
977 fn query_option_value(&self, name: &str) -> Option<NAValue> {
978 match name {
979 FRAME_SKIP_OPTION => Some(NAValue::String(self.skip_mode.to_string())),
980 DEBLOCK_SKIP_OPTION => Some(NAValue::Bool(self.deblock_skip)),
981 _ => None,
982 }
983 }
984}
985
986pub fn get_decoder_mt() -> Box<dyn NADecoderMT + Send> {
987 Box::new(H264MTDecoder::new())
988}