1 use nihav_core::codecs::DecoderResult;
2 use nihav_core::frame::{FrameType, NAVideoBufferRef, NATimeInfo};
3 use nihav_core::refs::*;
4 use nihav_codec_support::codecs::MV;
5 use super::sets::SeqParameterSet;
10 pub struct PictureInfo {
15 pub pic_type: FrameType,
16 pub buf: NAVideoBufferRef<u8>,
20 pub long_term: Option<usize>,
22 pub mv_info: NABufferRef<FrameMV>,
25 #[derive(Clone,Copy,Default, Debug)]
26 pub struct FrameMBInfo {
27 pub mb_type: CompactMBType,
28 pub ref_poc: [[u16; 2]; 4],
29 pub ref_idx: [[PicRef; 2]; 4],
30 pub mv: [[MV; 2]; 16],
34 pub fn new() -> Self { Self::default() }
39 pub mbs: Vec<FrameMBInfo>,
44 pub fn new(mb_w: usize, mb_h: usize) -> Self {
46 mbs: vec![FrameMBInfo::default(); mb_w * mb_h],
53 pub struct SliceRefs {
54 pub ref_list0: Vec<Option<PictureInfo>>,
55 pub ref_list1: Vec<Option<PictureInfo>>,
61 pub fn get_ref_id(&self, list_id: u8, ref_id: usize) -> Option<u32> {
62 let ref_list = if list_id == 0 { &self.ref_list0 } else { &self.ref_list1 };
63 if ref_list.len() > ref_id {
64 ref_list[ref_id].as_ref().map(|pic| pic.full_id)
69 pub fn select_ref_pic(&self, list_id: u8, ref_id: usize) -> Option<NAVideoBufferRef<u8>> {
70 let ref_list = if list_id == 0 { &self.ref_list0 } else { &self.ref_list1 };
71 if ref_list.len() > ref_id {
72 ref_list[ref_id].as_ref().map(|pic| pic.buf.clone())
77 pub fn get_colocated_info(&self, mb_x: usize, mb_y: usize) -> (FrameMBInfo, u16, bool) {
78 if let Some(ref ref_pic) = &self.ref_list1[0] {
79 let mv_info = &ref_pic.mv_info;
80 let mb = mv_info.mbs[mb_x + mb_y * mv_info.mb_stride];
81 (mb, ref_pic.full_id as u16, ref_pic.long_term.is_some())
83 (FrameMBInfo::default(), 0, false)
86 pub fn map_ref0(&self, ref0_id: u16) -> (PicRef, bool) {
89 for (i, rpic0) in self.ref_list0.iter().enumerate() {
90 if let Some(ref pic) = rpic0 {
91 if (pic.full_id as u16) == ref0_id {
93 long = pic.long_term.is_some();
98 (PicRef::new(r0_idx), long)
100 pub fn map_refs(&self, ref_idx: [PicRef; 2]) -> [u16; 2] {
101 let r0 = ref_idx[0].index();
102 let r1 = ref_idx[1].index();
103 let ref0 = if r0 < self.ref_list0.len() {
104 if let Some(ref pic) = self.ref_list0[r0] {
112 let ref1 = if r1 < self.ref_list1.len() {
113 if let Some(ref pic) = self.ref_list1[r1] {
123 pub fn cmp_refs(&self, ref1: [PicRef; 2], ref2: [PicRef; 2]) -> bool {
125 self.cmp_ref(ref1[0], ref2[0], 0) && self.cmp_ref(ref1[1], ref2[1], 1)
130 fn cmp_ref(&self, ref1: PicRef, ref2: PicRef, list: u8) -> bool {
134 let idx0 = ref1.index();
135 let idx1 = ref2.index();
139 let src = if list == 0 { &self.ref_list0 } else { &self.ref_list1 };
140 if idx0 >= src.len() || idx1 >= src.len() {
141 //panic!("wrong refs");
144 if let (Some(ref pic0), Some(ref pic1)) = (&src[idx0], &src[idx1]) {
145 pic0.full_id == pic1.full_id
147 //panic!("missing pics");
155 pub struct SimplePictureInfo<'a> {
157 pub buf: SimpleFrame<'a>,
159 pub mv_info: &'a FrameMV,
163 pub struct SimplifiedSliceRefs<'a> {
164 pub ref_list0: Vec<Option<SimplePictureInfo<'a>>>,
165 pub ref_list1: Vec<Option<SimplePictureInfo<'a>>>,
169 impl<'a> SimplifiedSliceRefs<'a> {
170 pub fn new(srefs: &'a SliceRefs) -> Self {
171 let mut ref_list0 = Vec::with_capacity(srefs.ref_list0.len());
172 let mut ref_list1 = Vec::with_capacity(srefs.ref_list1.len());
173 for entry in srefs.ref_list0.iter() {
174 ref_list0.push(entry.as_ref().map(|pic| SimplePictureInfo {
175 full_id: pic.full_id,
176 buf: SimpleFrame::new(&pic.buf),
177 long_term: pic.long_term.is_some(),
178 mv_info: &pic.mv_info,
181 for entry in srefs.ref_list1.iter() {
182 ref_list1.push(entry.as_ref().map(|pic| SimplePictureInfo {
183 full_id: pic.full_id,
184 buf: SimpleFrame::new(&pic.buf),
185 long_term: pic.long_term.is_some(),
186 mv_info: &pic.mv_info,
190 cur_id: srefs.cur_id,
194 pub fn get_ref_id(&self, list_id: u8, ref_id: usize) -> Option<u32> {
195 let ref_list = if list_id == 0 { &self.ref_list0 } else { &self.ref_list1 };
196 if ref_list.len() > ref_id {
197 ref_list[ref_id].as_ref().map(|pic| pic.full_id)
202 pub fn select_ref_pic(&self, list_id: u8, ref_id: usize) -> Option<&SimpleFrame> {
203 let ref_list = if list_id == 0 { &self.ref_list0 } else { &self.ref_list1 };
204 if ref_list.len() > ref_id {
205 ref_list[ref_id].as_ref().map(|pic| &pic.buf)
210 pub fn get_colocated_info(&self, mb_x: usize, mb_y: usize) -> (FrameMBInfo, u16, bool) {
211 if let Some(ref ref_pic) = &self.ref_list1[0] {
212 let mv_info = ref_pic.mv_info;
213 let mb = mv_info.mbs[mb_x + mb_y * mv_info.mb_stride];
214 (mb, ref_pic.full_id as u16, ref_pic.long_term)
216 (FrameMBInfo::default(), 0, false)
219 pub fn map_ref0(&self, ref0_id: u16) -> (PicRef, bool) {
221 let mut long = false;
222 for (i, rpic0) in self.ref_list0.iter().enumerate() {
223 if let Some(ref pic) = rpic0 {
224 if (pic.full_id as u16) == ref0_id {
226 long = pic.long_term;
231 (PicRef::new(r0_idx), long)
233 pub fn map_refs(&self, ref_idx: [PicRef; 2]) -> [u16; 2] {
234 let r0 = ref_idx[0].index();
235 let r1 = ref_idx[1].index();
236 let ref0 = if r0 < self.ref_list0.len() {
237 if let Some(ref pic) = self.ref_list0[r0] {
245 let ref1 = if r1 < self.ref_list1.len() {
246 if let Some(ref pic) = self.ref_list1[r1] {
256 pub fn cmp_refs(&self, ref1: [PicRef; 2], ref2: [PicRef; 2]) -> bool {
258 self.cmp_ref(ref1[0], ref2[0], 0) && self.cmp_ref(ref1[1], ref2[1], 1)
263 fn cmp_ref(&self, ref1: PicRef, ref2: PicRef, list: u8) -> bool {
267 let idx0 = ref1.index();
268 let idx1 = ref2.index();
272 let src = if list == 0 { &self.ref_list0 } else { &self.ref_list1 };
273 if idx0 >= src.len() || idx1 >= src.len() {
274 //panic!("wrong refs");
277 if let (Some(ref pic0), Some(ref pic1)) = (&src[idx0], &src[idx1]) {
278 pic0.full_id == pic1.full_id
280 //panic!("missing pics");
287 pub struct FrameRefs {
288 pub ref_pics: Vec<PictureInfo>,
289 pub cur_refs: SliceRefs,
290 pub long_term: Vec<Option<PictureInfo>>,
294 prev_ref_poc_lsb: u16,
296 frame_num_offset: u32,
301 pub fn new() -> Self {
303 ref_pics: Vec::with_capacity(16),
304 cur_refs: SliceRefs {
305 ref_list0: Vec::with_capacity(3),
306 ref_list1: Vec::with_capacity(3),
309 long_term: Vec::new(),
319 pub fn fill_ref_nums(&self, dst: &mut Vec<u32>) {
320 for pic in self.ref_pics.iter() {
321 if !dst.contains(&pic.full_id) {
322 dst.push(pic.full_id);
325 for pic in self.long_term.iter().flatten() {
326 if !dst.contains(&pic.full_id) {
327 dst.push(pic.full_id);
331 pub fn calc_picture_num(&mut self, slice_hdr: &SliceHeader, is_idr: bool, ref_id: u8, sps: &SeqParameterSet) -> u32 {
332 self.max_frame_num = 1 << sps.log2_max_frame_num;
333 match sps.pic_order_cnt_type {
336 //self.prev_poc_msb = 0;
337 self.prev_poc_lsb = 0;
339 self.prev_poc_lsb = self.prev_ref_poc_lsb;
341 let max_poc_lsb = 1 << sps.log2_max_pic_order_cnt_lsb;
342 let half_max_poc_lsb = 1 << (sps.log2_max_pic_order_cnt_lsb - 1);
343 let cur_lsb = slice_hdr.pic_order_cnt_lsb;
344 let poc_msb = if cur_lsb < self.prev_poc_lsb && (self.prev_poc_lsb - cur_lsb >= half_max_poc_lsb) {
345 self.prev_poc_msb + max_poc_lsb
346 } else if cur_lsb > self.prev_poc_lsb && (cur_lsb - self.prev_poc_lsb > half_max_poc_lsb) {
347 self.prev_poc_msb.wrapping_sub(max_poc_lsb)
351 let poc = poc_msb + u32::from(cur_lsb);
353 self.prev_ref_poc_lsb = slice_hdr.pic_order_cnt_lsb;
354 self.prev_poc_msb = poc_msb;
359 let off = if self.prev_frame_num > slice_hdr.frame_num {
360 self.frame_num_offset + (1 << sps.log2_max_frame_num)
362 self.frame_num_offset
364 let mut anum = if sps.num_ref_frames_in_pic_order_cnt_cycle != 0 {
365 (off as i32) + i32::from(slice_hdr.frame_num)
369 if ref_id == 0 && anum > 0 {
372 let (poc_cycle_cnt, fno_in_poc_cycle) = if anum > 0 {
373 let nrf = sps.num_ref_frames_in_pic_order_cnt_cycle as i32;
374 ((anum - 1) / nrf, (anum - 1) % nrf)
378 let mut expected_delta = 0;
379 for &offset in sps.offset_for_ref_frame[..sps.num_ref_frames_in_pic_order_cnt_cycle].iter() {
380 expected_delta += offset;
382 let mut expected_poc = if anum > 0 {
383 let mut sum = poc_cycle_cnt * expected_delta;
384 for &offset in sps.offset_for_ref_frame[..=fno_in_poc_cycle as usize].iter() {
392 expected_poc += sps.offset_for_non_ref_pic;
394 let (top_id, _bottom_id) = if !slice_hdr.field_pic {
395 let top_id = expected_poc + slice_hdr.delta_pic_order_cnt[0];
396 let bot_id = top_id + sps.offset_for_top_to_bottom_field + slice_hdr.delta_pic_order_cnt[1];
398 } else if !slice_hdr.bottom_field {
399 (expected_poc + slice_hdr.delta_pic_order_cnt[0], 0)
401 (0, sps.offset_for_top_to_bottom_field + slice_hdr.delta_pic_order_cnt[1])
403 self.prev_frame_num = slice_hdr.frame_num;
404 self.frame_num_offset = off;
408 if slice_hdr.frame_num < self.prev_frame_num {
409 self.frame_num_offset += 1 << sps.log2_max_frame_num;
411 self.prev_frame_num = slice_hdr.frame_num;
412 self.frame_num_offset + u32::from(slice_hdr.frame_num)
416 pub fn apply_adaptive_marking(&mut self, marking: &AdaptiveMarking, cur_id: u16, max_id: u16) -> DecoderResult<()> {
417 let all_ref_pics = self.ref_pics.clone();
419 for (&op, (&arg1, &arg2)) in marking.memory_management_control_op.iter().zip(marking.operation_arg.iter().zip(marking.operation_arg2.iter())).take(marking.num_ops) {
422 let src_id = cur_id.wrapping_sub(arg1) & (max_id - 1);
423 let mut found = false;
425 for (i, pic) in self.ref_pics.iter().enumerate() {
426 if pic.id == src_id {
433 self.ref_pics.remove(idx);
436 2 => { // mark long term picture as unused
437 let idx = arg1 as usize;
438 if idx < self.long_term.len() {
439 self.long_term[idx] = None;
443 let src_id = cur_id.wrapping_sub(arg1) & (max_id - 1);
445 let didx = arg2 as usize;
446 for pic in all_ref_pics.iter() {
447 if pic.id == src_id {
448 if didx < self.long_term.len() {
449 self.long_term[didx] = Some(pic.clone());
456 self.long_term.resize(arg1 as usize, None);
459 self.ref_pics.clear();
460 self.long_term.clear();
463 // assign an long term index to current pic - done elsewhere
470 pub fn clear_refs(&mut self) {
471 self.ref_pics.clear();
472 self.long_term.clear();
474 #[allow(clippy::cognitive_complexity)]
475 pub fn select_refs(&mut self, sps: &SeqParameterSet, slice_hdr: &SliceHeader, cur_id: u32) {
476 self.cur_refs.cur_id = cur_id;
477 self.cur_refs.ref_list0.clear();
478 self.cur_refs.ref_list1.clear();
479 let pic_num_mask = if sps.log2_max_frame_num == 16 {
482 (1 << sps.log2_max_frame_num) - 1
485 if !slice_hdr.slice_type.is_intra() {
486 let has_reordering = slice_hdr.ref_pic_list_reordering_l0;
488 let num_ref = slice_hdr.num_ref_idx_l0_active;
489 if slice_hdr.slice_type.is_p() {
490 if !self.ref_pics.is_empty() {
491 for pic in self.ref_pics.iter().rev().take(num_ref) {
492 self.cur_refs.ref_list0.push(Some(pic.clone()));
497 for (i, pic) in self.ref_pics.iter().enumerate() {
499 if pic.full_id > cur_id {
503 for pic in self.ref_pics[..pivot].iter().rev() {
504 if self.cur_refs.ref_list0.len() >= num_ref {
507 self.cur_refs.ref_list0.push(Some(pic.clone()));
509 for pic in self.ref_pics.iter().skip(pivot) {
510 if self.cur_refs.ref_list0.len() >= num_ref {
513 self.cur_refs.ref_list0.push(Some(pic.clone()));
516 if !self.long_term.is_empty() && self.cur_refs.ref_list0.len() < num_ref {
517 let copy_size = num_ref - self.cur_refs.ref_list0.len();
518 for ltpic in self.long_term.iter().take(copy_size) {
519 self.cur_refs.ref_list0.push(ltpic.clone());
523 form_ref_list(&mut self.cur_refs.ref_list0,
524 &self.ref_pics, &self.long_term,
525 &slice_hdr.reordering_list_l0,
526 slice_hdr.frame_num, pic_num_mask);
528 if slice_hdr.slice_type.is_b() {
529 let has_reordering = slice_hdr.ref_pic_list_reordering_l1;
531 let num_ref = slice_hdr.num_ref_idx_l1_active;
533 for (i, pic) in self.ref_pics.iter().enumerate() {
535 if pic.full_id > cur_id {
539 for pic in self.ref_pics.iter().skip(pivot) {
540 if self.cur_refs.ref_list1.len() >= num_ref {
543 self.cur_refs.ref_list1.push(Some(pic.clone()));
545 for pic in self.ref_pics[..pivot].iter().rev() {
546 if self.cur_refs.ref_list1.len() >= num_ref {
549 self.cur_refs.ref_list1.push(Some(pic.clone()));
551 if !self.long_term.is_empty() && self.cur_refs.ref_list1.len() < num_ref {
552 let copy_size = num_ref - self.cur_refs.ref_list1.len();
553 for ltpic in self.long_term.iter().take(copy_size) {
554 self.cur_refs.ref_list1.push(ltpic.clone());
557 if self.cur_refs.ref_list1.len() > 1 && self.cur_refs.ref_list0.len() == self.cur_refs.ref_list1.len() {
558 let mut equal = true;
559 for (pic1, pic2) in self.cur_refs.ref_list0.iter().zip(self.cur_refs.ref_list1.iter()) {
561 (Some(p1), Some(p2)) => {
562 if p1.full_id != p2.full_id {
575 self.cur_refs.ref_list1.swap(0, 1);
579 form_ref_list(&mut self.cur_refs.ref_list1,
580 &self.ref_pics, &self.long_term,
581 &slice_hdr.reordering_list_l1,
582 slice_hdr.frame_num, pic_num_mask);
587 pub fn add_short_term(&mut self, cpic: PictureInfo, num_ref_frames: usize) {
588 if !self.ref_pics.is_empty() && self.ref_pics.len() >= num_ref_frames {
589 let base_id = i32::from(cpic.id);
590 let mut min_id = base_id;
592 for (i, pic) in self.ref_pics.iter().enumerate() {
593 let mut pic_id = i32::from(pic.id);
594 if pic_id > base_id {
595 pic_id -= self.max_frame_num;
602 self.ref_pics.remove(min_idx);
604 if self.ref_pics.is_empty() || self.ref_pics.last().unwrap().full_id < cpic.full_id {
605 self.ref_pics.push(cpic);
608 for (i, pic) in self.ref_pics.iter().enumerate() {
609 if pic.full_id < cpic.full_id {
615 self.ref_pics.insert(idx + 1, cpic);
618 pub fn add_long_term(&mut self, lt_idx: usize, cpic: PictureInfo) {
619 if lt_idx < self.long_term.len() {
620 self.long_term[lt_idx] = Some(cpic);
625 fn form_ref_list(ref_list: &mut Vec<Option<PictureInfo>>, ref_pics: &[PictureInfo], long_term: &[Option<PictureInfo>], reord_info: &ReorderingInfo, cur_id: u16, pic_num_mask: u16) {
626 let mut ref_pic_id = cur_id;
627 for (&op, &num) in reord_info.reordering_of_pic_nums_idc.iter().zip(reord_info.abs_diff_or_num.iter()).take(reord_info.num_ops) {
630 ref_pic_id = ref_pic_id.wrapping_sub(num) & pic_num_mask;
632 ref_pic_id = ref_pic_id.wrapping_add(num) & pic_num_mask;
634 let mut found = false;
635 for pic in ref_pics.iter() {
636 if pic.id == ref_pic_id {
637 ref_list.push(Some(pic.clone()));
646 let idx = num as usize;
647 if idx < long_term.len() {
648 ref_list.push(long_term[idx].clone());