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