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