h264: prepare data references before decoding
[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();
e6aaad5c 225 let start = br.tell() / 8;
47ecb2b7
KS
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 }
5f223cdb 280 fn pred_mv(sstate: &mut SliceState, frame_refs: &SimplifiedSliceRefs, 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)]
5f223cdb 340 fn handle_macroblock(&mut self, slice_hdr: &SliceHeader, mb_info: &mut CurrentMBInfo, slice_refs: &SimplifiedSliceRefs) {
47ecb2b7
KS
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 }
754ab49a
KS
365 if !tx_bypass {
366 if !mb_info.transform_size_8x8 {
367 let quant_dc = !mb_info.mb_type.is_intra16x16();
368 for (coded, coeffs) in mb_info.coded[..16].iter_mut().zip(mb_info.coeffs[..16].iter_mut()) {
369 if *coded {
370 idct(coeffs, qp_y, quant_dc);
371 } else if has_dc {
372 idct_dc(coeffs, qp_y, quant_dc);
373 *coded = true;
47ecb2b7 374 }
754ab49a
KS
375 }
376 } else {
377 for i in 0..4 {
378 if mb_info.coded[(i & 1) * 2 + (i & 2) * 4] {
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);
47ecb2b7 381 }
47ecb2b7
KS
382 }
383 }
754ab49a
KS
384 } else if !mb_info.transform_size_8x8 {
385 for i in 0..16 {
386 if !mb_info.coded[i] && has_dc {
387 mb_info.coded[i] = true;
47ecb2b7
KS
388 }
389 }
390 }
391 for chroma in 0..2 {
392 let qp_c = if chroma == 0 { qp_u } else { qp_v };
393 if mb_info.cbpc != 0 {
394 chroma_dc_transform(&mut mb_info.chroma_dc[chroma], qp_c);
395 }
396 for i in 0..4 {
397 let blk_no = 16 + chroma * 4 + i;
398 mb_info.coeffs[blk_no][0] = mb_info.chroma_dc[chroma][i];
399 if mb_info.coded[blk_no] {
400 idct(&mut mb_info.coeffs[blk_no], qp_c, false);
401 } else if mb_info.coeffs[blk_no][0] != 0 {
402 idct_dc(&mut mb_info.coeffs[blk_no], qp_c, false);
403 mb_info.coded[blk_no] = true;
404 }
405 }
406 }
407 if !pps.entropy_coding_mode || mb_info.mb_type.is_skip() || mb_info.mb_type.is_intra() {
408 self.sstate.reset_mb_mv();
409 }
410 if !mb_info.mb_type.is_intra() {
5f223cdb 411 Self::pred_mv(&mut self.sstate, slice_refs, mb_info, self.cur_id, self.temporal_mv, self.sps[self.cur_sps].direct_8x8_inference);
47ecb2b7
KS
412 }
413 if !pps.constrained_intra_pred && mb_info.mb_type != MBType::Intra4x4 && mb_info.mb_type != MBType::Intra8x8 {
414 self.sstate.fill_ipred(IntraPredMode::DC);
415 }
416
417 let xpos = self.sstate.mb_x * 16;
418 let ypos = self.sstate.mb_y * 16;
419 if let Some(ref mut pic) = self.cur_pic {
420 let mut frm = NASimpleVideoFrame::from_video_buf(&mut pic.buf).unwrap();
421 if mb_info.mb_type != MBType::PCM {
422 let weight_mode = if self.pps[self.cur_pps].weighted_pred && slice_hdr.slice_type.is_p() {
423 1
424 } else if slice_hdr.slice_type.is_b() {
425 self.pps[self.cur_pps].weighted_bipred_idc
426 } else {
427 0
428 };
5f223cdb 429 recon_mb(&mut frm, slice_hdr, mb_info, &mut self.sstate, slice_refs, &mut self.mc_dsp, weight_mode);
47ecb2b7
KS
430 } else {
431 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)) {
432 dline[..16].copy_from_slice(src);
433 }
434 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)) {
435 dline[..8].copy_from_slice(src);
436 }
437 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)) {
438 dline[..8].copy_from_slice(src);
439 }
440 }
441/*match mb_info.mb_type {
442MBType::BSkip | MBType::Direct | MBType::B16x16(_) | MBType::B16x8(_, _) | MBType::B8x16(_, _) | MBType::B8x8 => {
443 let dstride = frm.stride[0];
444 let dst = &mut frm.data[frm.offset[0] + self.sstate.mb_x * 16 + self.sstate.mb_y * 16 * dstride..];
445 for el in dst[..16].iter_mut() { *el = 255; }
446 for row in dst.chunks_mut(dstride).skip(1).take(15) {
447 row[0] = 255;
448 }
449},
450_ => {},
451};*/
452 self.sstate.save_ipred_context(&frm);
453 }
454 if let Some(ref mut pic) = self.cur_pic {
455 let mv_info = &mut pic.mv_info;
456 let mb_pos = self.sstate.mb_x + self.sstate.mb_y * mv_info.mb_stride;
457 let mut mb = FrameMBInfo::new();
458 mb.mb_type = mb_info.mb_type.into();
459 for blk4 in 0..16 {
460 mb.mv[blk4] = self.sstate.get_cur_blk4(blk4).mv;
461 }
462 for blk8 in 0..4 {
5f223cdb 463 mb.ref_poc[blk8] = slice_refs.map_refs(self.sstate.get_cur_blk8(blk8).ref_idx);
47ecb2b7
KS
464 mb.ref_idx[blk8] = self.sstate.get_cur_blk8(blk8).ref_idx;
465 }
466 mv_info.mbs[mb_pos] = mb;
467 }
468 if !self.deblock_skip && self.deblock_mode != 1 {
5f223cdb 469 self.sstate.fill_deblock(slice_refs, self.deblock_mode, self.is_s);
47ecb2b7
KS
470 if let Some(ref mut pic) = self.cur_pic {
471 let mut frm = NASimpleVideoFrame::from_video_buf(&mut pic.buf).unwrap();
472 loop_filter_mb(&mut frm, &self.sstate, self.lf_alpha, self.lf_beta);
473 }
474 }
475 self.sstate.next_mb();
476 }
477 fn decode_slice_cavlc(&mut self, br: &mut BitReader, slice_hdr: &SliceHeader, full_size: usize) -> DecoderResult<bool> {
478 const INTRA_CBP: [u8; 48] = [
479 47, 31, 15, 0, 23, 27, 29, 30, 7, 11, 13, 14, 39, 43, 45, 46,
480 16, 3, 5, 10, 12, 19, 21, 26, 28, 35, 37, 42, 44, 1, 2, 4,
481 8, 17, 18, 20, 24, 6, 9, 22, 25, 32, 33, 34, 36, 40, 38, 41
482 ];
483 const INTER_CBP: [u8; 48] = [
484 0, 16, 1, 2, 4, 8, 32, 3, 5, 10, 12, 15, 47, 7, 11, 13,
485 14, 6, 9, 31, 35, 37, 42, 44, 33, 34, 36, 40, 39, 43, 45, 46,
486 17, 18, 20, 24, 19, 21, 26, 28, 23, 27, 29, 30, 22, 25, 38, 41
487 ];
488
e6aaad5c 489 let mut mb_idx = slice_hdr.first_mb_in_slice;
47ecb2b7
KS
490 let mut mb_info = CurrentMBInfo { qp_y: slice_hdr.slice_qp, ..Default::default() };
491 let skip_type = if slice_hdr.slice_type.is_p() { MBType::PSkip } else { MBType::BSkip };
5f223cdb
KS
492
493 let slice_refs = self.frame_refs.cur_refs.clone();
494 let sslice_refs = SimplifiedSliceRefs::new(&slice_refs);
495
47ecb2b7
KS
496 while br.tell() < full_size && mb_idx < self.num_mbs {
497 mb_info.coded = [false; 25];
498 mb_info.ref_l0 = [ZERO_REF; 4];
499 mb_info.ref_l1 = [ZERO_REF; 4];
500 mb_info.mv_l0 = [ZERO_MV; 16];
501 mb_info.mv_l1 = [ZERO_MV; 16];
502 mb_info.chroma_dc = [[0; 4]; 2];
503 mb_info.cbpy = 0;
504 mb_info.cbpc = 0;
505
506 if !slice_hdr.slice_type.is_intra() {
507 let mb_skip_run = br.read_ue()? as usize;
508 validate!(mb_idx + mb_skip_run <= self.num_mbs);
509 mb_info.mb_type = skip_type;
510 for _ in 0..mb_skip_run {
5f223cdb 511 self.handle_macroblock(slice_hdr, &mut mb_info, &sslice_refs);
47ecb2b7
KS
512 mb_idx += 1;
513 }
514 if mb_idx == self.num_mbs || br.tell() >= full_size {
515 break;
516 }
517 }
518 if br.tell() < full_size {
519 if self.is_mbaff && ((mb_idx & 1) == 0) {
520 let _mb_field_decoding = br.read_bool()?;
521 }
522 let mut mb_type = decode_mb_type_cavlc(br, slice_hdr)?;
523 mb_info.mb_type = mb_type;
524 mb_info.transform_size_8x8 = false;
525 if mb_type == MBType::PCM {
526 br.align();
527 for pix in self.ipcm_buf[..256 + 64 + 64].iter_mut() {
528 *pix = br.read(8)? as u8;
529 }
530 self.sstate.fill_ncoded(16);
531 } else {
532 if self.transform_8x8_mode && mb_type == MBType::Intra4x4 {
533 mb_info.transform_size_8x8 = br.read_bool()?;
534 if mb_info.transform_size_8x8 {
535 mb_type = MBType::Intra8x8;
536 mb_info.mb_type = MBType::Intra8x8;
537 }
538 }
539 decode_mb_pred_cavlc(br, slice_hdr, mb_type, &mut self.sstate, &mut mb_info)?;
540 let (cbpy, cbpc) = if let MBType::Intra16x16(_, cbpy, cbpc) = mb_type {
541 (cbpy, cbpc)
542 } else {
543 let cbp_id = br.read_ue()? as usize;
544 validate!(cbp_id < INTRA_CBP.len());
545 let cbp = if mb_type == MBType::Intra4x4 || mb_type == MBType::Intra8x8 {
546 INTRA_CBP[cbp_id]
547 } else {
548 INTER_CBP[cbp_id]
549 };
550 if self.transform_8x8_mode && (cbp & 0xF) != 0 && mb_info.can_have_8x8_tx(self.sps[self.cur_sps].direct_8x8_inference) {
551 mb_info.transform_size_8x8 = br.read_bool()?;
552 }
553 ((cbp & 0xF), (cbp >> 4))
554 };
555 mb_info.cbpy = cbpy;
556 mb_info.cbpc = cbpc;
557 self.sstate.get_cur_mb().cbp = (cbpc << 4) | cbpy;
558 if cbpy != 0 || cbpc != 0 || mb_type.is_intra16x16() {
559 let mb_qp_delta = br.read_se()?;
560 validate!(mb_qp_delta >= -26 && mb_qp_delta <= 25);
561 let new_qp = mb_qp_delta + i32::from(mb_info.qp_y);
562 mb_info.qp_y = if new_qp < 0 {
563 (new_qp + 52) as u8
564 } else if new_qp >= 52 {
565 (new_qp - 52) as u8
566 } else {
567 new_qp as u8
568 };
569 mb_info.coeffs = [[0; 16]; 25];
570 if self.transform_8x8_mode {
571 mb_info.clear_coeffs8x8();
572 }
573 mb_info.chroma_dc = [[0; 4]; 2];
574 decode_residual_cavlc(br, &mut self.sstate, &mut mb_info, &self.cavlc_cb)?;
575 }
576 }
5f223cdb 577 self.handle_macroblock(slice_hdr, &mut mb_info, &sslice_refs);
47ecb2b7
KS
578 }
579 mb_idx += 1;
580 }
581 if let Some(ref mut pic) = self.cur_pic {
582 pic.cur_mb = mb_idx;
583 }
584 Ok(mb_idx == self.num_mbs)
585 }
586 fn decode_slice_cabac(&mut self, cabac: &mut CABAC, slice_hdr: &SliceHeader) -> DecoderResult<bool> {
e6aaad5c 587 let mut mb_idx = slice_hdr.first_mb_in_slice;
47ecb2b7
KS
588 let mut prev_mb_skipped = false;
589 let skip_type = if slice_hdr.slice_type.is_p() { MBType::PSkip } else { MBType::BSkip };
590 let mut last_qp_diff = false;
591
592 let mut mb_info = CurrentMBInfo { qp_y: slice_hdr.slice_qp, ..Default::default() };
593
5f223cdb
KS
594 let slice_refs = self.frame_refs.cur_refs.clone();
595 let sslice_refs = SimplifiedSliceRefs::new(&slice_refs);
596
47ecb2b7
KS
597 while mb_idx < self.num_mbs {
598 mb_info.coded = [false; 25];
599 mb_info.ref_l0 = [ZERO_REF; 4];
600 mb_info.ref_l1 = [ZERO_REF; 4];
601 mb_info.mv_l0 = [ZERO_MV; 16];
602 mb_info.mv_l1 = [ZERO_MV; 16];
603 mb_info.chroma_dc = [[0; 4]; 2];
604 mb_info.cbpy = 0;
605 mb_info.cbpc = 0;
606 let mb_skip = cabac_decode_mbskip(cabac, &self.sstate, slice_hdr);
607 if !mb_skip {
608 if self.is_mbaff && (((mb_idx & 1) == 0) || (prev_mb_skipped && ((mb_idx & 1) == 1))) {
609 let _mb_field_decoding = cabac.decode_bit(70);
610 }
611 let mut mb_type = cabac_decode_mb_type(cabac, slice_hdr, &self.sstate);
612 mb_info.mb_type = mb_type;
613 mb_info.transform_size_8x8 = false;
614 if mb_type == MBType::PCM {
615 let ipcm_size = 256 + 64 + 64;
616 validate!(cabac.pos + ipcm_size <= cabac.src.len());
617 self.ipcm_buf[..ipcm_size].copy_from_slice(&cabac.src[cabac.pos..][..ipcm_size]);
618 cabac.pos += ipcm_size;
619 cabac.reinit()?;
620 last_qp_diff = false;
621 } else {
622 if self.transform_8x8_mode && mb_type == MBType::Intra4x4 {
623 let mut ctx = 0;
624 if self.sstate.get_top_mb().transform_8x8 {
625 ctx += 1;
626 }
627 if self.sstate.get_left_mb().transform_8x8 {
628 ctx += 1;
629 }
630 mb_info.transform_size_8x8 = cabac.decode_bit(399 + ctx);
631 if mb_info.transform_size_8x8 {
632 mb_type = MBType::Intra8x8;
633 mb_info.mb_type = MBType::Intra8x8;
634 }
635 }
636 decode_mb_pred_cabac(cabac, slice_hdr, mb_type, &mut self.sstate, &mut mb_info);
637 let (cbpy, cbpc) = if let MBType::Intra16x16(_, cbpy, cbpc) = mb_type {
638 (cbpy, cbpc)
639 } else {
640 decode_cbp_cabac(cabac, &self.sstate)
641 };
642 if self.transform_8x8_mode && cbpy != 0 && mb_info.can_have_8x8_tx(self.sps[self.cur_sps].direct_8x8_inference) {
643 let mut ctx = 0;
644 if self.sstate.get_top_mb().transform_8x8 {
645 ctx += 1;
646 }
647 if self.sstate.get_left_mb().transform_8x8 {
648 ctx += 1;
649 }
650 mb_info.transform_size_8x8 = cabac.decode_bit(399 + ctx);
651 }
652 if mb_type.is_intra() {
653 self.sstate.get_cur_mb().cmode = mb_info.chroma_ipred;
654 }
655 mb_info.cbpy = cbpy;
656 mb_info.cbpc = cbpc;
657 self.sstate.get_cur_mb().cbp = (cbpc << 4) | cbpy;
658 if cbpy != 0 || cbpc != 0 || mb_type.is_intra16x16() {
659 let mb_qp_delta = decode_mb_qp_delta_cabac(cabac, last_qp_diff as usize);
660 validate!(mb_qp_delta >= -26 && mb_qp_delta <= 25);
661 last_qp_diff = mb_qp_delta != 0;
662 let new_qp = mb_qp_delta + i32::from(mb_info.qp_y);
663 mb_info.qp_y = if new_qp < 0 {
664 (new_qp + 52) as u8
665 } else if new_qp >= 52 {
666 (new_qp - 52) as u8
667 } else {
668 new_qp as u8
669 };
670 mb_info.coeffs = [[0; 16]; 25];
671 if self.transform_8x8_mode {
672 mb_info.clear_coeffs8x8();
673 }
674 mb_info.chroma_dc = [[0; 4]; 2];
675 decode_residual_cabac(cabac, &mut self.sstate, &mut mb_info);
676 } else {
677 last_qp_diff = false;
678 }
679 }
680 } else {
681 mb_info.mb_type = skip_type;
682 mb_info.transform_size_8x8 = false;
683 last_qp_diff = false;
684 }
5f223cdb 685 self.handle_macroblock(slice_hdr, &mut mb_info, &sslice_refs);
47ecb2b7
KS
686 prev_mb_skipped = mb_skip;
687 if !(self.is_mbaff && ((mb_idx & 1) == 0)) && cabac.decode_terminate() {
688 if let Some(ref mut pic) = self.cur_pic {
689 pic.cur_mb = mb_idx + 1;
690 }
691 return Ok(mb_idx + 1 == self.num_mbs);
692 }
693 mb_idx += 1;
694 }
695 Err(DecoderError::InvalidData)
696 }
697}
698
699impl NADecoder for H264Decoder {
700 fn init(&mut self, supp: &mut NADecoderSupport, info: NACodecInfoRef) -> DecoderResult<()> {
701 if let NACodecTypeInfo::Video(vinfo) = info.get_properties() {
702 let fmt = YUV420_FORMAT;
703 let myinfo = NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, false, fmt));
704 self.info = NACodecInfo::new_ref(info.get_name(), myinfo, info.get_extradata()).into_ref();
705
706 let edata = info.get_extradata().unwrap();
707//print!("edata:"); for &el in edata.iter() { print!(" {:02X}", el); } println!();
708 if edata.len() > 11 && &edata[0..4] == b"avcC" {
709 let mut mr = MemoryReader::new_read(edata.as_slice());
710 let mut br = ByteReader::new(&mut mr);
711 let mut nal_buf = Vec::new();
712
713 br.read_skip(4)?;
714 let version = br.read_byte()?;
715 validate!(version == 1);
716 let profile = br.read_byte()?;
717 let _compatibility = br.read_byte()?;
718 let _level = br.read_byte()?;
719 let b = br.read_byte()?;
720 validate!((b & 0xFC) == 0xFC);
721 self.nal_len = (b & 3) + 1;
722 let b = br.read_byte()?;
723 validate!((b & 0xE0) == 0xE0);
724 let num_sps = (b & 0x1F) as usize;
725 for _ in 0..num_sps {
726 let len = br.read_u16be()? as usize;
727 let offset = br.tell() as usize;
728 validate!((br.peek_byte()? & 0x1F) == 7);
729 let _size = unescape_nal(&edata[offset..][..len], &mut nal_buf);
730 self.handle_nal(&nal_buf, supp, true)?;
731 br.read_skip(len)?;
732 }
733 let num_pps = br.read_byte()? as usize;
734 for _ in 0..num_pps {
735 let len = br.read_u16be()? as usize;
736 let offset = br.tell() as usize;
737 validate!((br.peek_byte()? & 0x1F) == 8);
738 let _size = unescape_nal(&edata[offset..][..len], &mut nal_buf);
739 self.handle_nal(&nal_buf, supp, true)?;
740 br.read_skip(len)?;
741 }
742 if br.left() > 0 {
743 match profile {
744 100 | 110 | 122 | 144 => {
745 let b = br.read_byte()?;
746 validate!((b & 0xFC) == 0xFC);
747 // b & 3 -> chroma format
748 let b = br.read_byte()?;
749 validate!((b & 0xF8) == 0xF8);
750 // b & 7 -> luma depth minus 8
751 let b = br.read_byte()?;
752 validate!((b & 0xF8) == 0xF8);
753 // b & 7 -> chroma depth minus 8
754 let num_spsext = br.read_byte()? as usize;
755 for _ in 0..num_spsext {
756 let len = br.read_u16be()? as usize;
757 // parse spsext
758 br.read_skip(len)?;
759 }
760 },
761 _ => {},
762 };
763 }
764 } else {
765 return Err(DecoderError::NotImplemented);
766 }
767
768 self.width = vinfo.get_width();
769 self.height = vinfo.get_height();
770
771 if (self.width == 0 || self.height == 0) && !self.sps.is_empty() {
772 self.width = self.sps[0].pic_width_in_mbs * 16;
773 self.height = self.sps[0].pic_height_in_mbs * 16;
774 }
775
776 let num_bufs = if !self.sps.is_empty() {
e6aaad5c 777 self.sps[0].num_ref_frames + 1
47ecb2b7
KS
778 } else {
779 3
780 }.max(16 + 1);
781 supp.pool_u8.set_dec_bufs(num_bufs);
782 supp.pool_u8.prealloc_video(NAVideoInfo::new(self.width, self.height, false, fmt), 4)?;
783
784 Ok(())
785 } else {
786 Err(DecoderError::InvalidData)
787 }
788 }
789 fn decode(&mut self, supp: &mut NADecoderSupport, pkt: &NAPacket) -> DecoderResult<NAFrameRef> {
790 let src = pkt.get_buffer();
791
792 let mut mr = MemoryReader::new_read(&src);
793 let mut br = ByteReader::new(&mut mr);
794 let mut nal_buf = Vec::with_capacity(src.len());
795 if self.nal_len > 0 {
796 let mut skip_decoding = false;
797 if self.skip_mode != FrameSkipMode::None {
798 let mut pic_type = FrameType::I;
799 let mut is_ref = false;
800 while br.left() > 0 {
801 let size = match self.nal_len {
802 1 => br.read_byte()? as usize,
803 2 => br.read_u16be()? as usize,
804 3 => br.read_u24be()? as usize,
805 4 => br.read_u32be()? as usize,
806 _ => unreachable!(),
807 };
808 validate!(br.left() >= (size as i64));
809 let offset = br.tell() as usize;
810 let size = unescape_nal(&src[offset..][..size], &mut nal_buf);
811 validate!(size > 0);
812 let nal_ref_idc = nal_buf[0] >> 5;
813 let nal_unit_type = nal_buf[0] & 0x1F;
814 if nal_unit_type == 1 || nal_unit_type == 5 {
815 let mut bitr = BitReader::new(&nal_buf[1..], BitReaderMode::BE);
816 let (first_mb, slice_type) = parse_slice_header_minimal(&mut bitr)?;
817 if first_mb == 0 && nal_ref_idc != 0 {
818 is_ref = true;
819 }
820 let new_type = slice_type.to_frame_type();
821 pic_type = match (pic_type, new_type) {
822 (FrameType::I, _) => new_type,
823 (_, FrameType::B) => FrameType::B,
824 _ => pic_type,
825 };
826 }
827 br.read_skip(size)?;
828 }
829 match self.skip_mode {
830 FrameSkipMode::IntraOnly => {
831 skip_decoding = pic_type != FrameType::I;
832 },
833 FrameSkipMode::KeyframesOnly => {
834 if !is_ref {
835 skip_decoding = true;
836 }
837 },
838 _ => {},
839 };
840 br.seek(SeekFrom::Start(0))?;
841 }
842 while br.left() > 0 {
843 let size = match self.nal_len {
844 1 => br.read_byte()? as usize,
845 2 => br.read_u16be()? as usize,
846 3 => br.read_u24be()? as usize,
847 4 => br.read_u32be()? as usize,
848 _ => unreachable!(),
849 };
850 validate!(br.left() >= (size as i64));
851 let offset = br.tell() as usize;
852 let _size = unescape_nal(&src[offset..][..size], &mut nal_buf);
853 self.handle_nal(nal_buf.as_slice(), supp, skip_decoding)?;
854 br.read_skip(size)?;
855 }
856 } else {
857//todo NAL detection
858 unimplemented!();
859 }
860
861 let (bufinfo, ftype, dts) = if self.has_pic && self.cur_pic.is_some() {
862 let mut npic = None;
863 std::mem::swap(&mut self.cur_pic, &mut npic);
864 let cpic = npic.unwrap();
865 let ret = (NABufferType::Video(cpic.buf.clone()), cpic.pic_type, Some(u64::from(cpic.full_id)));
866 if cpic.is_ref {
867 self.frame_refs.add_short_term(cpic.clone(), self.sps[self.cur_sps].num_ref_frames);
868 }
869 if let Some(lt_idx) = cpic.long_term {
870 self.frame_refs.add_long_term(lt_idx, cpic);
871 }
872 ret
873 } else {
874 (NABufferType::None, FrameType::Skip, None)
875 };
876
877 let mut frm = NAFrame::new_from_pkt(pkt, self.info.clone(), bufinfo);
878 frm.set_keyframe(ftype == FrameType::I);
879 if let (Some(mydts), None) = (dts, frm.get_dts()) {
880 frm.set_dts(Some(mydts));
881 }
882 if let Some(dts) = dts {
883 frm.set_id(dts as i64);
884 }
885 frm.set_frame_type(ftype);
886 Ok(frm.into_ref())
887 }
888 fn flush(&mut self) {
889 }
890}
891
892impl NAOptionHandler for H264Decoder {
893 fn get_supported_options(&self) -> &[NAOptionDefinition] { DECODER_OPTIONS }
894 fn set_options(&mut self, options: &[NAOption]) {
895 for option in options.iter() {
896 for opt_def in DECODER_OPTIONS.iter() {
897 if opt_def.check(option).is_ok() {
898 match (option.name, &option.value) {
899 (FRAME_SKIP_OPTION, NAValue::String(ref strval)) => {
900 if let Ok(smode) = FrameSkipMode::from_str(strval) {
901 self.skip_mode = smode;
902 }
903 },
904 (DEBLOCK_SKIP_OPTION, NAValue::Bool(val)) => {
905 self.deblock_skip = *val;
906 },
907 _ => {},
908 }
909 }
910 }
911 }
912 }
913 fn query_option_value(&self, name: &str) -> Option<NAValue> {
914 match name {
915 FRAME_SKIP_OPTION => Some(NAValue::String(self.skip_mode.to_string())),
916 DEBLOCK_SKIP_OPTION => Some(NAValue::Bool(self.deblock_skip)),
917 _ => None,
918 }
919 }
920}
921
922pub fn get_decoder() -> Box<dyn NADecoder + Send> {
923 Box::new(H264Decoder::new())
924}