h264: prepare data references before decoding
[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
cc818ce8
KS
7const AVG_BUF_VINFO: NAVideoInfo = NAVideoInfo { width: 32, height: 32, flipped: false, format: YUV420_FORMAT, bits: 12 };
8
11d7aef2
KS
9pub struct FrameDecoder {
10 pub slices: Vec<(SliceHeader, usize, SliceRefs, Vec<u8>)>,
11 pub cur_pic: PictureInfo,
12 sps: Arc<SeqParameterSet>,
13 pps: Arc<PicParameterSet>,
14 pub num_mbs: usize,
15 mc_dsp: H264MC,
16 dispatch: Shareable<ThreadDispatcher>,
17 sstate: SliceState,
18 cavlc_cb: Arc<CAVLCTables>,
19 ipcm_buf: [u8; 256 + 64 + 64],
20 is_mbaff: bool,
21 deblock_skip: bool,
22}
23
24impl FrameDecoder {
25 pub fn decode_slice(&mut self, hdr: &SliceHeader, hdr_size: usize, refs: &SliceRefs, nal: &[u8]) -> DecoderResult<usize> {
26 self.sstate.reset(self.sps.pic_width_in_mbs, self.sps.pic_height_in_mbs, hdr.first_mb_in_slice);
27
28 let mut full_size = nal.len() * 8;
29 for &byte in nal.iter().rev() {
30 if byte == 0 {
31 full_size -= 8;
32 } else {
33 full_size -= (byte.trailing_zeros() + 1) as usize;
34 break;
35 }
36 }
37 validate!(full_size > 0);
38
5f223cdb
KS
39 let sslice_refs = SimplifiedSliceRefs::new(refs);
40
11d7aef2
KS
41 let mut br = BitReader::new(&nal[hdr_size / 8..], BitReaderMode::BE);
42 if !self.pps.entropy_coding_mode {
43 br.skip((hdr_size & 7) as u32)?;
5f223cdb 44 self.decode_slice_cavlc(&mut br, full_size - (hdr_size & !7), hdr, &sslice_refs)
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)?;
5f223cdb 49 self.decode_slice_cabac(&mut cabac, hdr, &sslice_refs)
11d7aef2
KS
50 }
51 }
5f223cdb 52 fn decode_slice_cavlc(&mut self, br: &mut BitReader, full_size: usize, slice_hdr: &SliceHeader, refs: &SimplifiedSliceRefs) -> 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 {
82 self.handle_macroblock(slice_hdr, &mut mb_info, refs)?;
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 }
148 self.handle_macroblock(slice_hdr, &mut mb_info, refs)?;
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 }
5f223cdb 157 fn decode_slice_cabac(&mut self, cabac: &mut CABAC, slice_hdr: &SliceHeader, refs: &SimplifiedSliceRefs) -> 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 }
253 self.handle_macroblock(slice_hdr, &mut mb_info, refs)?;
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)]
5f223cdb 269 fn handle_macroblock(&mut self, slice_hdr: &SliceHeader, mb_info: &mut CurrentMBInfo, refs: &SimplifiedSliceRefs) -> 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();
294 for i in 0..16 {
295 if mb_info.coded[i] {
296 if !tx_bypass {
297 idct(&mut mb_info.coeffs[i], qp_y, quant_dc);
298 }
299 } else if has_dc {
300 if !tx_bypass {
301 idct_dc(&mut mb_info.coeffs[i], qp_y, quant_dc);
302 }
303 mb_info.coded[i] = true;
304 }
305 }
306 } else {
307 for i in 0..4 {
308 if mb_info.coded[(i & 1) * 2 + (i & 2) * 4] && !tx_bypass {
309 dequant8x8(&mut mb_info.coeffs8x8[i].coeffs, &self.pps.scaling_list_8x8[!mb_info.mb_type.is_intra() as usize]);
310 idct8x8(&mut mb_info.coeffs8x8[i].coeffs, qp_y);
311 }
312 }
313 }
314 for chroma in 0..2 {
315 let qp_c = if chroma == 0 { qp_u } else { qp_v };
316 if mb_info.cbpc != 0 {
317 chroma_dc_transform(&mut mb_info.chroma_dc[chroma], qp_c);
318 }
319 for i in 0..4 {
320 let blk_no = 16 + chroma * 4 + i;
321 mb_info.coeffs[blk_no][0] = mb_info.chroma_dc[chroma][i];
322 if mb_info.coded[blk_no] {
323 idct(&mut mb_info.coeffs[blk_no], qp_c, false);
324 } else if mb_info.coeffs[blk_no][0] != 0 {
325 idct_dc(&mut mb_info.coeffs[blk_no], qp_c, false);
326 mb_info.coded[blk_no] = true;
327 }
328 }
329 }
330 if !self.pps.entropy_coding_mode || mb_info.mb_type.is_skip() || mb_info.mb_type.is_intra() {
331 self.sstate.reset_mb_mv();
332 }
333 if !mb_info.mb_type.is_intra() {
334 let temporal_mv = !slice_hdr.direct_spatial_mv_pred;
335 let cur_id = self.cur_pic.full_id as u16;
336 // wait for the reference macroblock MV to be available
337 if matches!(mb_info.mb_type, MBType::Direct | MBType::BSkip | MBType::B8x8) {
338 if let Some(ref_id) = refs.get_ref_id(0, mb_info.ref_l1[0].index()) {
339 wait_for_mb(&self.dispatch, &self.sstate, self.sstate.mb_x * 16, self.sstate.mb_y * 16, ZERO_MV, ref_id)?;
340 }
341 }
342 Self::pred_mv(&mut self.sstate, refs, mb_info, cur_id, temporal_mv, self.sps.direct_8x8_inference);
343 }
344 if !self.pps.constrained_intra_pred && mb_info.mb_type != MBType::Intra4x4 && mb_info.mb_type != MBType::Intra8x8 {
345 self.sstate.fill_ipred(IntraPredMode::DC);
346 }
347
348 let xpos = self.sstate.mb_x * 16;
349 let ypos = self.sstate.mb_y * 16;
350 let mut frm = NASimpleVideoFrame::from_video_buf(&mut self.cur_pic.buf).unwrap();
351 if mb_info.mb_type != MBType::PCM {
352 let weight_mode = if self.pps.weighted_pred && slice_hdr.slice_type.is_p() {
353 1
354 } else if slice_hdr.slice_type.is_b() {
355 self.pps.weighted_bipred_idc
356 } else {
357 0
358 };
359 recon_mb_mt(&mut frm, slice_hdr, mb_info, &mut self.sstate, refs, &mut self.mc_dsp, weight_mode, &self.dispatch)?;
360 } else {
361 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)) {
362 dline[..16].copy_from_slice(src);
363 }
364 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)) {
365 dline[..8].copy_from_slice(src);
366 }
367 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)) {
368 dline[..8].copy_from_slice(src);
369 }
370 }
371 self.sstate.save_ipred_context(&frm);
372
373 let mv_info = &mut self.cur_pic.mv_info;
374 let mb_pos = self.sstate.mb_x + self.sstate.mb_y * mv_info.mb_stride;
375 let mut mb = FrameMBInfo::new();
376 mb.mb_type = mb_info.mb_type.into();
377 for blk4 in 0..16 {
378 mb.mv[blk4] = self.sstate.get_cur_blk4(blk4).mv;
379 }
380 for blk8 in 0..4 {
381 mb.ref_poc[blk8] = refs.map_refs(self.sstate.get_cur_blk8(blk8).ref_idx);
382 mb.ref_idx[blk8] = self.sstate.get_cur_blk8(blk8).ref_idx;
383 }
384 mv_info.mbs[mb_pos] = mb;
385
386 let deblock_mode = slice_hdr.disable_deblocking_filter_idc;
387 if !self.deblock_skip && deblock_mode != 1 {
388 let is_s = slice_hdr.slice_type == SliceType::SI || slice_hdr.slice_type == SliceType::SP;
389 self.sstate.fill_deblock(refs, deblock_mode, is_s);
390 let mut frm = NASimpleVideoFrame::from_video_buf(&mut self.cur_pic.buf).unwrap();
391 let lf_alpha = slice_hdr.slice_alpha_c0_offset;
392 let lf_beta = slice_hdr.slice_beta_offset;
393 loop_filter_mb(&mut frm, &self.sstate, lf_alpha, lf_beta);
394 }
395 self.sstate.next_mb();
396 Ok(())
397 }
398
5f223cdb 399 fn pred_mv(sstate: &mut SliceState, frame_refs: &SimplifiedSliceRefs, mb_info: &mut CurrentMBInfo, cur_id: u16, temporal_mv: bool, direct_8x8: bool) {
11d7aef2
KS
400 let mb_type = mb_info.mb_type;
401 if !mb_type.is_4x4() {
402 let (pw, ph) = mb_type.size();
403 let mut xoff = 0;
404 let mut yoff = 0;
405 if mb_type == MBType::Direct || mb_type == MBType::BSkip {
406 sstate.predict_direct_mb(frame_refs, temporal_mv, direct_8x8, cur_id);
407 }
408 for part in 0..mb_type.num_parts() {
409 if !mb_type.is_l1(part) {
410 match mb_type {
411 MBType::PSkip => sstate.predict_pskip(),
412 MBType::BSkip | MBType::Direct => {
413 },
414 _ => {
415 sstate.predict(xoff, yoff, pw, ph, 0,
416 mb_info.mv_l0[part], mb_info.ref_l0[part]);
417 },
418 };
419 }
420 if !mb_type.is_l0(part) && mb_type != MBType::BSkip && mb_type != MBType::Direct {
421 sstate.predict(xoff, yoff, pw, ph, 1, mb_info.mv_l1[part], mb_info.ref_l1[part]);
422 }
423 if pw != 16 {
424 xoff += pw;
425 } else {
426 yoff += ph;
427 }
428 }
429 } else {
430 for part in 0..4 {
431 let sub_type = mb_info.sub_mb_type[part];
432 let mut xoff = (part & 1) * 8;
433 let mut yoff = (part & 2) * 4;
434 let orig_x = xoff;
435 let (pw, ph) = sub_type.size();
436 for subpart in 0..sub_type.num_parts() {
437 if sub_type != SubMBType::Direct8x8 {
438 if !sub_type.is_l1() {
439 sstate.predict(xoff, yoff, pw, ph, 0, mb_info.mv_l0[part * 4 + subpart], mb_info.ref_l0[part]);
440 }
441 if !sub_type.is_l0() {
442 sstate.predict(xoff, yoff, pw, ph, 1, mb_info.mv_l1[part * 4 + subpart], mb_info.ref_l1[part]);
443 }
444 } else {
445 for sblk in 0..4 {
446 sstate.predict_direct_sub(frame_refs, temporal_mv, direct_8x8, cur_id, (xoff / 4) + (sblk & 1) + (yoff / 4) * 4 + (sblk & 2) * 2);
447 }
448 }
449 xoff += pw;
450 if xoff == orig_x + 8 {
451 xoff -= 8;
452 yoff += ph;
453 }
454 }
455 }
456 }
457 }
458}
459
460struct H264MTDecoder {
461 info: NACodecInfoRef,
462 nal_len: u8,
463 dispatch: Shareable<ThreadDispatcher>,
464 frame_refs: FrameRefs,
465 skip_mode: FrameSkipMode,
466 sps: Vec<Arc<SeqParameterSet>>,
467 cur_sps: usize,
468 pps: Vec<Arc<PicParameterSet>>,
469 cur_pps: usize,
470 cur_fdec: Option<FrameDecoder>,
471 cavlc_cb: Arc<CAVLCTables>,
472 deblock_skip: bool,
473 max_last_poc: u32,
474 poc_base: u32,
cc818ce8 475 avg_pool: NAVideoBufferPool<u8>,
11d7aef2
KS
476}
477
478impl H264MTDecoder {
479 fn new() -> Self {
480 Self {
481 info: NACodecInfoRef::default(),
482 nal_len: 0,
483 dispatch: Arc::new(RwLock::new(ThreadDispatcher::new())),
484 frame_refs: FrameRefs::new(),
485 skip_mode: FrameSkipMode::default(),
486 sps: Vec::new(),
487 cur_sps: 0,
488 pps: Vec::new(),
489 cur_pps: 0,
490 cur_fdec: None,
491 cavlc_cb: Arc::new(CAVLCTables::new()),
492 deblock_skip: false,
493 max_last_poc: 0,
494 poc_base: 0,
cc818ce8 495 avg_pool: NAVideoBufferPool::new(8),
11d7aef2
KS
496 }
497 }
498 fn handle_nal(&mut self, src: Vec<u8>, supp: &mut NADecoderSupport, skip_decoding: bool, user_id: u32, time: NATimeInfo) -> DecoderResult<()> {
499 validate!(!src.is_empty());
500 validate!((src[0] & 0x80) == 0);
501 let nal_ref_idc = src[0] >> 5;
502 let nal_unit_type = src[0] & 0x1F;
503
504 let mut full_size = src.len() * 8;
505 for &byte in src.iter().rev() {
506 if byte == 0 {
507 full_size -= 8;
508 } else {
509 full_size -= (byte.trailing_zeros() + 1) as usize;
510 break;
511 }
512 }
513 validate!(full_size > 0);
514 match nal_unit_type {
515 1 | 5 if !skip_decoding => {
516 let is_idr = nal_unit_type == 5;
517 let mut br = BitReader::new(&src[..(full_size + 7)/8], BitReaderMode::BE);
518 br.skip(8)?;
519
520 let slice_hdr = parse_slice_header(&mut br, self.sps.as_slice(), self.pps.as_slice(), is_idr, nal_ref_idc)?;
521 let hdr_size = br.tell();
522 validate!(br.tell() < full_size);
523 let full_id;
524 if slice_hdr.first_mb_in_slice == 0 {
525 validate!(self.cur_fdec.is_none());
526 for (i, pps) in self.pps.iter().enumerate() {
527 if pps.pic_parameter_set_id == slice_hdr.pic_parameter_set_id {
528 self.cur_pps = i;
529 break;
530 }
531 }
532 for (i, sps) in self.sps.iter().enumerate() {
533 if sps.seq_parameter_set_id == self.pps[self.cur_pps].seq_parameter_set_id {
534 self.cur_sps = i;
535 break;
536 }
537 }
538
539 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;
540 if is_idr {
541 if cur_full_id <= self.max_last_poc {
542 self.poc_base = self.max_last_poc + 2 - (cur_full_id - self.poc_base);
543 cur_full_id = self.max_last_poc + 2;
544 }
545 }
546 self.max_last_poc = self.max_last_poc.max(cur_full_id);
547 full_id = cur_full_id;
548
549 let sps = &self.sps[self.cur_sps];
550 if sps.chroma_format_idc != 1 || sps.bit_depth_luma != 8 || sps.bit_depth_chroma != 8 {
551 println!(" chroma fmt {} bits {}/{}", sps.chroma_format_idc, sps.bit_depth_luma, sps.bit_depth_chroma);
552 return Err(DecoderError::NotImplemented);
553 }
554
555 if is_idr {
556 self.frame_refs.clear_refs();
557 }
558
559 let width = sps.pic_width_in_mbs << 4;
560 let height = sps.pic_height_in_mbs << 4;
561 let num_mbs = sps.pic_width_in_mbs * sps.pic_height_in_mbs;
562
cc818ce8
KS
563 let avg_buf = if let Some(buf) = self.avg_pool.get_free() {
564 buf
565 } else {
566 let new_avg_buf = alloc_video_buffer(AVG_BUF_VINFO, 4).unwrap().get_vbuf().unwrap();
567 self.avg_pool.add_frame(new_avg_buf.clone());
568 new_avg_buf
569 };
11d7aef2
KS
570 let mut mc_dsp = H264MC::new(avg_buf);
571 mc_dsp.set_dimensions(width, height);
572
573 let is_mbaff = sps.mb_adaptive_frame_field && !slice_hdr.field_pic;
574 if is_mbaff {
575 println!("MBAFF");
576 return Err(DecoderError::NotImplemented);
577 }
578 if !sps.frame_mbs_only {
579 println!("PAFF?");
580 return Err(DecoderError::NotImplemented);
581 }
582
583 let cur_vinfo = supp.pool_u8.get_info();
584 let tmp_vinfo = NAVideoInfo::new(width, height, false, YUV420_FORMAT);
585 if cur_vinfo != Some(tmp_vinfo) {
586 supp.pool_u8.reset();
587 supp.pool_u8.prealloc_video(tmp_vinfo, 4)?;
588 }
589
590 let buf = if let Some(pic) = supp.pool_u8.get_free() {
591 pic
592 } else {
593 if supp.pool_u8.get_num_used() > 256 {
594 return Err(DecoderError::AllocError);
595 }
596 if let Ok(nbuf) = alloc_video_buffer(tmp_vinfo, 4) {
597 let vbuf = nbuf.get_vbuf().unwrap();
598 supp.pool_u8.add_frame(vbuf.clone());
599 vbuf
600 } else {
601 return Err(DecoderError::AllocError);
602 }
603 };
604
605 let cur_pic = PictureInfo {
606 id: slice_hdr.frame_num,
607 full_id, user_id, time,
608 pic_type: slice_hdr.slice_type.to_frame_type(),
609 buf,
610 cur_mb: 0,
611 is_ref: nal_ref_idc != 0,
612 is_idr,
613 long_term: get_long_term_id(is_idr, &slice_hdr),
614 mv_info: NABufferRef::new(FrameMV::new(sps.pic_width_in_mbs, sps.pic_height_in_mbs)),
615 };
616
617 self.cur_fdec = Some(FrameDecoder{
618 slices: Vec::new(),
619 sstate: SliceState::new(),
620 ipcm_buf: [0; 256 + 64 + 64],
621 //width, height,
622 num_mbs,
623 sps: Arc::clone(sps),
624 pps: Arc::clone(&self.pps[self.cur_pps]),
625 dispatch: Arc::clone(&self.dispatch),
626 cavlc_cb: Arc::clone(&self.cavlc_cb),
627 mc_dsp,
628 cur_pic,
629 is_mbaff,
630 deblock_skip: self.deblock_skip,
631 });
632 } else {
633 if let Some(ref mut fdec) = self.cur_fdec {
634 let new_type = slice_hdr.slice_type.to_frame_type();
635 let pic = &mut fdec.cur_pic;
636 pic.pic_type = match (pic.pic_type, new_type) {
637 (FrameType::I, _) => new_type,
638 (_, FrameType::B) => FrameType::B,
639 _ => pic.pic_type,
640 };
641 full_id = pic.full_id;
642 } else {
643 return Ok(());
644 }
645 }
646
647 let sps = &self.sps[self.cur_sps];
648
649 self.frame_refs.select_refs(sps, &slice_hdr, full_id);
650
651 if slice_hdr.adaptive_ref_pic_marking_mode {
652 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)?;
653 }
654 if let Some(ref mut fdec) = self.cur_fdec {
655 fdec.slices.push((slice_hdr, hdr_size, self.frame_refs.cur_refs.clone(), src));
656 }
657 },
658 2 => { // slice data partition A
659 //slice header
660 //slice id = read_ue()
661 //cat 2 slice data (all but MB layer residual)
662 return Err(DecoderError::NotImplemented);
663 },
664 3 => { // slice data partition B
665 //slice id = read_ue()
666 //if pps.redundant_pic_cnt_present { redundant_pic_cnt = read_ue() }
667 //cat 3 slice data (MB layer residual)
668 return Err(DecoderError::NotImplemented);
669 },
670 4 => { // slice data partition C
671 //slice id = read_ue()
672 //if pps.redundant_pic_cnt_present { redundant_pic_cnt = read_ue() }
673 //cat 4 slice data (MB layer residual)
674 return Err(DecoderError::NotImplemented);
675 },
676 6 => {}, //SEI
677 7 => {
678 let sps = parse_sps(&src[1..])?;
679 self.sps.push(Arc::new(sps));
680 },
681 8 => {
682 validate!(full_size >= 8 + 16);
683 let pps = parse_pps(&src[1..], self.sps.as_slice(), full_size - 8)?;
684 let mut found = false;
685 for stored_pps in self.pps.iter_mut() {
686 if stored_pps.pic_parameter_set_id == pps.pic_parameter_set_id {
687 *stored_pps = Arc::clone(&pps);
688 found = true;
689 break;
690 }
691 }
692 if !found {
693 self.pps.push(pps);
694 }
695 },
696 9 => { // access unit delimiter
697 },
698 10 => {}, //end of sequence
699 11 => {}, //end of stream
700 12 => {}, //filler
701 _ => {},
702 };
703
704 Ok(())
705 }
706}
707
708impl NADecoderMT for H264MTDecoder {
709 fn init(&mut self, supp: &mut NADecoderSupport, info: NACodecInfoRef, nthreads: usize) -> DecoderResult<()> {
710 if let NACodecTypeInfo::Video(vinfo) = info.get_properties() {
711 let fmt = YUV420_FORMAT;
712 let myinfo = NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, false, fmt));
713 self.info = NACodecInfo::new_ref(info.get_name(), myinfo, info.get_extradata()).into_ref();
714
715 let edata = info.get_extradata().unwrap();
716//print!("edata:"); for &el in edata.iter() { print!(" {:02X}", el); } println!();
717 if edata.len() > 11 && &edata[0..4] == b"avcC" {
718 let mut mr = MemoryReader::new_read(edata.as_slice());
719 let mut br = ByteReader::new(&mut mr);
720
721 br.read_skip(4)?;
722 let version = br.read_byte()?;
723 validate!(version == 1);
724 let profile = br.read_byte()?;
725 let _compatibility = br.read_byte()?;
726 let _level = br.read_byte()?;
727 let b = br.read_byte()?;
728 validate!((b & 0xFC) == 0xFC);
729 self.nal_len = (b & 3) + 1;
730 let b = br.read_byte()?;
731 validate!((b & 0xE0) == 0xE0);
732 let num_sps = (b & 0x1F) as usize;
733 for _ in 0..num_sps {
734 let len = br.read_u16be()? as usize;
735 let offset = br.tell() as usize;
736 validate!((br.peek_byte()? & 0x1F) == 7);
737 let mut nal_buf = Vec::new();
738 let _size = unescape_nal(&edata[offset..][..len], &mut nal_buf);
739 self.handle_nal(nal_buf, supp, true, 0, NATimeInfo::new(None, None, None, 0, 0))?;
740 br.read_skip(len)?;
741 }
742 let num_pps = br.read_byte()? as usize;
743 for _ in 0..num_pps {
744 let len = br.read_u16be()? as usize;
745 let offset = br.tell() as usize;
746 validate!((br.peek_byte()? & 0x1F) == 8);
747 let mut nal_buf = Vec::new();
748 let _size = unescape_nal(&edata[offset..][..len], &mut nal_buf);
749 self.handle_nal(nal_buf, supp, true, 0, NATimeInfo::new(None, None, None, 0, 0))?;
750 br.read_skip(len)?;
751 }
752 if br.left() > 0 {
753 match profile {
754 100 | 110 | 122 | 144 => {
755 let b = br.read_byte()?;
756 validate!((b & 0xFC) == 0xFC);
757 // b & 3 -> chroma format
758 let b = br.read_byte()?;
759 validate!((b & 0xF8) == 0xF8);
760 // b & 7 -> luma depth minus 8
761 let b = br.read_byte()?;
762 validate!((b & 0xF8) == 0xF8);
763 // b & 7 -> chroma depth minus 8
764 let num_spsext = br.read_byte()? as usize;
765 for _ in 0..num_spsext {
766 let len = br.read_u16be()? as usize;
767 // parse spsext
768 br.read_skip(len)?;
769 }
770 },
771 _ => {},
772 };
773 }
774 } else {
775 return Err(DecoderError::NotImplemented);
776 }
777
778 let mut width = vinfo.get_width();
779 let mut height = vinfo.get_height();
780
781 if (width == 0 || height == 0) && !self.sps.is_empty() {
782 width = self.sps[0].pic_width_in_mbs * 16;
783 height = self.sps[0].pic_height_in_mbs * 16;
784 }
785
786 let num_bufs = if !self.sps.is_empty() {
e6aaad5c 787 self.sps[0].num_ref_frames + 1
11d7aef2
KS
788 } else {
789 3
790 }.max(16 + 1);
791 if let Ok(ref mut sd) = self.dispatch.write() {
792 sd.max_threads = nthreads;
793 } else {
794 return Err(DecoderError::Bug);
795 }
796 supp.pool_u8.set_dec_bufs(num_bufs + nthreads);
797 supp.pool_u8.prealloc_video(NAVideoInfo::new(width, height, false, fmt), 4)?;
798
cc818ce8
KS
799 self.avg_pool.prealloc_video(AVG_BUF_VINFO, 4)?;
800
11d7aef2
KS
801 Ok(())
802 } else {
803 Err(DecoderError::InvalidData)
804 }
805 }
806 fn can_take_input(&mut self) -> bool {
807 if let Ok(ref sd) = self.dispatch.read() {
808 sd.can_decode_more()
809 } else {
810 false
811 }
812 }
813 fn queue_pkt(&mut self, supp: &mut NADecoderSupport, pkt: &NAPacket, user_id: u32) -> DecoderResult<bool> {
814 if !self.can_take_input() {
815 return Ok(false);
816 }
817
818 let src = pkt.get_buffer();
819
820 let mut mr = MemoryReader::new_read(&src);
821 let mut br = ByteReader::new(&mut mr);
822 let mut nal_buf = Vec::with_capacity(src.len());
823
824 if self.nal_len > 0 {
825 let mut skip_decoding = false;
826 if self.skip_mode != FrameSkipMode::None {
827 let mut pic_type = FrameType::I;
828 let mut is_ref = false;
829 while br.left() > 0 {
830 let size = match self.nal_len {
831 1 => br.read_byte()? as usize,
832 2 => br.read_u16be()? as usize,
833 3 => br.read_u24be()? as usize,
834 4 => br.read_u32be()? as usize,
835 _ => unreachable!(),
836 };
837 validate!(br.left() >= (size as i64));
838 let offset = br.tell() as usize;
839 let size = unescape_nal(&src[offset..][..size], &mut nal_buf);
840 validate!(size > 0);
841 let nal_ref_idc = nal_buf[0] >> 5;
842 let nal_unit_type = nal_buf[0] & 0x1F;
843 if nal_unit_type == 1 || nal_unit_type == 5 {
844 let mut bitr = BitReader::new(&nal_buf[1..], BitReaderMode::BE);
845 let (first_mb, slice_type) = parse_slice_header_minimal(&mut bitr)?;
846 if first_mb == 0 && nal_ref_idc != 0 {
847 is_ref = true;
848 }
849 let new_type = slice_type.to_frame_type();
850 pic_type = match (pic_type, new_type) {
851 (FrameType::I, _) => new_type,
852 (_, FrameType::B) => FrameType::B,
853 _ => pic_type,
854 };
855 }
856 br.read_skip(size)?;
857 }
858 match self.skip_mode {
859 FrameSkipMode::IntraOnly => {
860 skip_decoding = pic_type != FrameType::I;
861 },
862 FrameSkipMode::KeyframesOnly => {
863 if !is_ref {
864 skip_decoding = true;
865 }
866 },
867 _ => {},
868 };
869 br.seek(SeekFrom::Start(0))?;
870 }
871
872 let mut initial_ref_frames = Vec::new();
873 self.frame_refs.fill_ref_nums(&mut initial_ref_frames);
874
875 while br.left() > 0 {
876 let size = match self.nal_len {
877 1 => br.read_byte()? as usize,
878 2 => br.read_u16be()? as usize,
879 3 => br.read_u24be()? as usize,
880 4 => br.read_u32be()? as usize,
881 _ => unreachable!(),
882 };
883 validate!(br.left() >= (size as i64));
884 let offset = br.tell() as usize;
885 let mut cur_nal_buf = Vec::with_capacity(size);
886 let _size = unescape_nal(&src[offset..][..size], &mut cur_nal_buf);
887 self.handle_nal(cur_nal_buf, supp, skip_decoding, user_id, pkt.ts)?;
888 br.read_skip(size)?;
889 }
890 let mut fdec = None;
891 std::mem::swap(&mut fdec, &mut self.cur_fdec);
892 if let Some(fdc) = fdec {
893 let cpic = &fdc.cur_pic;
894 if cpic.is_ref {
895 self.frame_refs.add_short_term(cpic.clone(), self.sps[self.cur_sps].num_ref_frames);
896 }
897 if let Some(lt_idx) = cpic.long_term {
898 self.frame_refs.add_long_term(lt_idx, cpic.clone());
899 }
900 let mut ref_frames = Vec::new();
901 self.frame_refs.fill_ref_nums(&mut ref_frames);
902 queue_decoding(&mut self.dispatch, fdc, &initial_ref_frames, &ref_frames);
903 }
904 } else {
905//todo NAL detection
906 unimplemented!();
907 }
908 Ok(true)
909 }
910 fn has_output(&mut self) -> bool {
911 if let Ok(ref ds) = self.dispatch.read() {
912 ds.has_output()
913 } else {
914 panic!("can't peek into status");
915 }
916 }
917 fn get_frame(&mut self) -> (DecoderResult<NAFrameRef>, u32) {
918 match wait_for_one(&mut self.dispatch) {
919 Ok(cpic) => {
920 let bufinfo = NABufferType::Video(cpic.buf.clone());
921 let ftype = cpic.pic_type;
922 let dts = Some(u64::from(cpic.full_id));
923 let mut frm = NAFrame::new(cpic.time, ftype, cpic.is_idr, self.info.clone(), bufinfo);
924 if let (Some(mydts), None) = (dts, frm.get_dts()) {
925 frm.set_dts(Some(mydts));
926 }
927 frm.set_id(cpic.user_id as i64);
928 (Ok(frm.into_ref()), cpic.user_id)
929 },
930 Err((err, id)) => (Err(err), id),
931 }
932 }
933 fn flush(&mut self) {
934 clear_threads(&mut self.dispatch);
f2cb96ff 935 self.frame_refs.clear_refs();
11d7aef2
KS
936 }
937}
938
939impl NAOptionHandler for H264MTDecoder {
940 fn get_supported_options(&self) -> &[NAOptionDefinition] { DECODER_OPTIONS }
941 fn set_options(&mut self, options: &[NAOption]) {
942 for option in options.iter() {
943 for opt_def in DECODER_OPTIONS.iter() {
944 if opt_def.check(option).is_ok() {
945 match (option.name, &option.value) {
946 (FRAME_SKIP_OPTION, NAValue::String(ref strval)) => {
947 if let Ok(smode) = FrameSkipMode::from_str(strval) {
948 self.skip_mode = smode;
949 }
950 },
951 (DEBLOCK_SKIP_OPTION, NAValue::Bool(val)) => {
952 self.deblock_skip = *val;
953 },
954 _ => {},
955 }
956 }
957 }
958 }
959 }
960 fn query_option_value(&self, name: &str) -> Option<NAValue> {
961 match name {
962 FRAME_SKIP_OPTION => Some(NAValue::String(self.skip_mode.to_string())),
963 DEBLOCK_SKIP_OPTION => Some(NAValue::Bool(self.deblock_skip)),
964 _ => None,
965 }
966 }
967}
968
969pub fn get_decoder_mt() -> Box<dyn NADecoderMT + Send> {
970 Box::new(H264MTDecoder::new())
971}