nihav-itu/src/codecs/h264/pic_ref.rs

   1 use nihav_core::codecs::DecoderResult;
   2 use nihav_core::frame::{FrameType, NAVideoBufferRef};
   3 use nihav_codec_support::codecs::MV;
   4 use super::sets::SeqParameterSet;
   5 use super::slice::*;
   6 use super::types::*;
   7
   8 #[derive(Clone)]
   9 pub struct PictureInfo {
  10     pub id:         u16,
  11     pub full_id:    u32,
  12     pub pic_type:   FrameType,
  13     pub buf:        NAVideoBufferRef<u8>,
  14     pub cur_mb:     usize,
  15     pub is_ref:     bool,
  16     pub long_term:  Option<usize>,
  17
  18     pub mv_info:    FrameMV, //todo replace with refcounted index to a pool
  19 }
  20
  21 #[derive(Clone,Copy,Default, Debug)]
  22 pub struct FrameMBInfo {
  23     pub mb_type:    CompactMBType,
  24     pub ref_poc:    [[u16; 2]; 4],
  25     pub ref_idx:    [[PicRef; 2]; 4],
  26     pub mv:         [[MV; 2]; 16],
  27 }
  28
  29 impl FrameMBInfo {
  30     pub fn new() -> Self { Self::default() }
  31 }
  32
  33 #[derive(Clone)]
  34 pub struct FrameMV {
  35     pub mbs:        Vec<FrameMBInfo>,
  36     pub mb_stride:  usize,
  37 }
  38
  39 impl FrameMV {
  40     pub fn new(mb_w: usize, mb_h: usize) -> Self {
  41         Self {
  42             mbs:        vec![FrameMBInfo::default(); mb_w * mb_h],
  43             mb_stride:  mb_w,
  44         }
  45     }
  46 }
  47
  48 pub struct FrameRefs {
  49     pub ref_pics:   Vec<PictureInfo>,
  50     pub ref_list0:  Vec<Option<PictureInfo>>,
  51     pub ref_list1:  Vec<Option<PictureInfo>>,
  52     pub long_term:  Vec<Option<PictureInfo>>,
  53     pub cur_id:         u32,
  54
  55     prev_poc_msb:       u32,
  56     prev_poc_lsb:       u16,
  57     prev_ref_poc_lsb:   u16,
  58     prev_frame_num:     u16,
  59     frame_num_offset:   u32,
  60     max_frame_num:      i32,
  61 }
  62
  63 impl FrameRefs {
  64     pub fn new() -> Self {
  65         Self {
  66             ref_pics:   Vec::with_capacity(16),
  67             ref_list0:  Vec::with_capacity(3),
  68             ref_list1:  Vec::with_capacity(3),
  69             long_term:  Vec::new(),
  70             cur_id:     0,
  71
  72             prev_poc_msb:       0,
  73             prev_poc_lsb:       0,
  74             prev_ref_poc_lsb:   0,
  75             prev_frame_num:     0,
  76             frame_num_offset:   0,
  77             max_frame_num:      0,
  78         }
  79     }
  80     pub fn calc_picture_num(&mut self, slice_hdr: &SliceHeader, is_idr: bool, ref_id: u8, sps: &SeqParameterSet) -> u32 {
  81         self.max_frame_num = 1 << sps.log2_max_frame_num;
  82         match sps.pic_order_cnt_type {
  83             0 => {
  84                 if is_idr {
  85                     //self.prev_poc_msb = 0;
  86                     self.prev_poc_lsb = 0;
  87                 } else {
  88                     self.prev_poc_lsb = self.prev_ref_poc_lsb;
  89                 }
  90                 let max_poc_lsb = 1 << sps.log2_max_pic_order_cnt_lsb;
  91                 let half_max_poc_lsb = 1 << (sps.log2_max_pic_order_cnt_lsb - 1);
  92                 let cur_lsb = slice_hdr.pic_order_cnt_lsb;
  93                 let poc_msb = if cur_lsb < self.prev_poc_lsb && (self.prev_poc_lsb - cur_lsb >= half_max_poc_lsb) {
  94                         self.prev_poc_msb + max_poc_lsb
  95                     } else if cur_lsb > self.prev_poc_lsb && (cur_lsb - self.prev_poc_lsb > half_max_poc_lsb) {
  96                         self.prev_poc_msb.wrapping_sub(max_poc_lsb)
  97                     } else {
  98                         self.prev_poc_msb
  99                     };
 100                 let poc = poc_msb + u32::from(cur_lsb);
 101                 if ref_id != 0 {
 102                     self.prev_ref_poc_lsb = slice_hdr.pic_order_cnt_lsb;
 103                     self.prev_poc_msb = poc_msb;
 104                 }
 105                 poc
 106             },
 107             1 => {
 108                 let off = if self.prev_frame_num > slice_hdr.frame_num {
 109                         self.frame_num_offset + (1 << sps.log2_max_frame_num)
 110                     } else {
 111                         self.frame_num_offset
 112                     };
 113                 let mut anum = if sps.num_ref_frames_in_pic_order_cnt_cycle != 0 {
 114                         (off as i32) + i32::from(slice_hdr.frame_num)
 115                     } else {
 116                         0
 117                     };
 118                 if ref_id == 0 && anum > 0 {
 119                     anum -= 1;
 120                 }
 121                 let (poc_cycle_cnt, fno_in_poc_cycle) = if anum > 0 {
 122                         let nrf = sps.num_ref_frames_in_pic_order_cnt_cycle as i32;
 123                         ((anum - 1) / nrf, (anum - 1) % nrf)
 124                     } else {
 125                         (0, 0)
 126                     };
 127                 let mut expected_delta = 0;
 128                 for &offset in sps.offset_for_ref_frame[..sps.num_ref_frames_in_pic_order_cnt_cycle].iter() {
 129                     expected_delta += offset;
 130                 }
 131                 let mut expected_poc = if anum > 0 {
 132                         let mut sum = poc_cycle_cnt * expected_delta;
 133                         for &offset in sps.offset_for_ref_frame[..=fno_in_poc_cycle as usize].iter() {
 134                             sum += offset;
 135                         }
 136                         sum
 137                     } else {
 138                         0
 139                     };
 140                 if ref_id == 0 {
 141                     expected_poc += sps.offset_for_non_ref_pic;
 142                 }
 143                 let (top_id, _bottom_id) = if !slice_hdr.field_pic {
 144                         let top_id = expected_poc + slice_hdr.delta_pic_order_cnt[0];
 145                         let bot_id = top_id + sps.offset_for_top_to_bottom_field + slice_hdr.delta_pic_order_cnt[1];
 146                         (top_id, bot_id)
 147                     } else if !slice_hdr.bottom_field {
 148                         (expected_poc + slice_hdr.delta_pic_order_cnt[0], 0)
 149                     } else {
 150                         (0, sps.offset_for_top_to_bottom_field + slice_hdr.delta_pic_order_cnt[1])
 151                     };
 152                 self.prev_frame_num = slice_hdr.frame_num;
 153                 self.frame_num_offset = off;
 154                 top_id as u32
 155             },
 156             _ => {
 157                 if slice_hdr.frame_num < self.prev_frame_num {
 158                     self.frame_num_offset   += 1 << sps.log2_max_frame_num;
 159                 }
 160                 self.prev_frame_num = slice_hdr.frame_num;
 161                 self.frame_num_offset + u32::from(slice_hdr.frame_num)
 162             },
 163         }
 164     }
 165     pub fn apply_adaptive_marking(&mut self, marking: &AdaptiveMarking, cur_id: u16, max_id: u16) -> DecoderResult<()> {
 166         let all_ref_pics = self.ref_pics.clone();
 167
 168         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) {
 169             match op {
 170                 1 => {
 171                     let src_id = cur_id.wrapping_sub(arg1) & (max_id - 1);
 172                     let mut found = false;
 173                     let mut idx = 0;
 174                     for (i, pic) in self.ref_pics.iter().enumerate() {
 175                         if pic.id == src_id {
 176                             found = true;
 177                             idx = i;
 178                             break;
 179                         }
 180                     }
 181                     if found {
 182                         self.ref_pics.remove(idx);
 183                     }
 184                 },
 185                 2 => { // mark long term picture as unused
 186                     let idx = arg1 as usize;
 187                     if idx < self.long_term.len() {
 188                         self.long_term[idx] = None;
 189                     }
 190                 },
 191                 3 => {
 192                     let src_id = cur_id.wrapping_sub(arg1) & (max_id - 1);
 193
 194                     let didx = arg2 as usize;
 195                     for pic in all_ref_pics.iter() {
 196                         if pic.id == src_id {
 197                             if didx < self.long_term.len() {
 198                                 self.long_term[didx] = Some(pic.clone());
 199                             }
 200                             break;
 201                         }
 202                     }
 203                 },
 204                 4 => {
 205                     self.long_term.resize(arg1 as usize, None);
 206                 },
 207                 5 => {
 208                     self.ref_pics.clear();
 209                     self.long_term.clear();
 210                 },
 211                 6 => {
 212                     // assign an long term index to current pic - done elsewhere
 213                 },
 214                 _ => {},
 215             };
 216         }
 217         Ok(())
 218     }
 219     pub fn clear_refs(&mut self) {
 220         self.ref_pics.clear();
 221         self.long_term.clear();
 222     }
 223     #[allow(clippy::cognitive_complexity)]
 224     pub fn select_refs(&mut self, sps: &SeqParameterSet, slice_hdr: &SliceHeader, cur_id: u32) {
 225         self.cur_id = cur_id;
 226         self.ref_list0.clear();
 227         self.ref_list1.clear();
 228         let pic_num_mask = if sps.log2_max_frame_num == 16 {
 229                 0xFFFF
 230             } else {
 231                 (1 << sps.log2_max_frame_num) - 1
 232             };
 233         if !slice_hdr.slice_type.is_intra() {
 234             let has_reordering = slice_hdr.ref_pic_list_reordering_l0;
 235             if !has_reordering {
 236                 let num_ref = slice_hdr.num_ref_idx_l0_active;
 237                 if slice_hdr.slice_type.is_p() {
 238                     if !self.ref_pics.is_empty() {
 239                         for pic in self.ref_pics.iter().rev().take(num_ref) {
 240                             self.ref_list0.push(Some(pic.clone()));
 241                         }
 242                     }
 243                 } else {
 244                     let mut pivot = 0;
 245                     for (i, pic) in self.ref_pics.iter().enumerate() {
 246                         pivot = i;
 247                         if pic.full_id > cur_id {
 248                             break;
 249                         }
 250                     }
 251                     for pic in self.ref_pics[..pivot].iter().rev() {
 252                         if self.ref_list0.len() >= num_ref {
 253                             break;
 254                         }
 255                         self.ref_list0.push(Some(pic.clone()));
 256                     }
 257                     for pic in self.ref_pics.iter().skip(pivot) {
 258                         if self.ref_list0.len() >= num_ref {
 259                             break;
 260                         }
 261                         self.ref_list0.push(Some(pic.clone()));
 262                     }
 263                 }
 264                 if !self.long_term.is_empty() && self.ref_list0.len() < num_ref {
 265                     let copy_size = num_ref - self.ref_list0.len();
 266                     for ltpic in self.long_term.iter().take(copy_size) {
 267                         self.ref_list0.push(ltpic.clone());
 268                     }
 269                 }
 270             } else {
 271                 form_ref_list(&mut self.ref_list0,
 272                               &self.ref_pics, &self.long_term,
 273                               &slice_hdr.reordering_list_l0,
 274                               slice_hdr.frame_num, pic_num_mask);
 275             }
 276             if slice_hdr.slice_type.is_b() {
 277                 let has_reordering = slice_hdr.ref_pic_list_reordering_l1;
 278                 if !has_reordering {
 279                     let num_ref = slice_hdr.num_ref_idx_l1_active;
 280                     let mut pivot = 0;
 281                     for (i, pic) in self.ref_pics.iter().enumerate() {
 282                         pivot = i;
 283                         if pic.full_id > cur_id {
 284                             break;
 285                         }
 286                     }
 287                     for pic in self.ref_pics.iter().skip(pivot) {
 288                         if self.ref_list1.len() >= num_ref {
 289                             break;
 290                         }
 291                         self.ref_list1.push(Some(pic.clone()));
 292                     }
 293                     for pic in self.ref_pics[..pivot].iter().rev() {
 294                         if self.ref_list1.len() >= num_ref {
 295                             break;
 296                         }
 297                         self.ref_list1.push(Some(pic.clone()));
 298                     }
 299                     if !self.long_term.is_empty() && self.ref_list1.len() < num_ref {
 300                         let copy_size = num_ref - self.ref_list1.len();
 301                         for ltpic in self.long_term.iter().take(copy_size) {
 302                             self.ref_list1.push(ltpic.clone());
 303                         }
 304                     }
 305                     if self.ref_list1.len() > 1 && self.ref_list0.len() == self.ref_list1.len() {
 306                         let mut equal = true;
 307                         for (pic1, pic2) in self.ref_list0.iter().zip(self.ref_list1.iter()) {
 308                             match (pic1, pic2) {
 309                                 (Some(p1), Some(p2)) => {
 310                                     if p1.full_id != p2.full_id {
 311                                         equal = false;
 312                                         break;
 313                                     }
 314                                 },
 315                                 (None, None) => {},
 316                                 _ => {
 317                                     equal = false;
 318                                     break;
 319                                 },
 320                             };
 321                         }
 322                         if equal {
 323                             self.ref_list1.swap(0, 1);
 324                         }
 325                     }
 326                 } else {
 327                     form_ref_list(&mut self.ref_list1,
 328                                   &self.ref_pics, &self.long_term,
 329                                   &slice_hdr.reordering_list_l1,
 330                                   slice_hdr.frame_num, pic_num_mask);
 331                 }
 332             }
 333         }
 334     }
 335     pub fn add_short_term(&mut self, cpic: PictureInfo, num_ref_frames: usize) {
 336         if !self.ref_pics.is_empty() && self.ref_pics.len() >= num_ref_frames {
 337             let base_id = i32::from(cpic.id);
 338             let mut min_id  = base_id;
 339             let mut min_idx = 0;
 340             for (i, pic) in self.ref_pics.iter().enumerate() {
 341                 let mut pic_id = i32::from(pic.id);
 342                 if pic_id > base_id {
 343                     pic_id -= self.max_frame_num;
 344                 }
 345                 if pic_id < min_id {
 346                     min_id = pic_id;
 347                     min_idx = i;
 348                 }
 349             }
 350             self.ref_pics.remove(min_idx);
 351         }
 352         if self.ref_pics.is_empty() || self.ref_pics.last().unwrap().full_id < cpic.full_id {
 353             self.ref_pics.push(cpic);
 354         } else {
 355             let mut idx = 0;
 356             for (i, pic) in self.ref_pics.iter().enumerate() {
 357                 if pic.full_id < cpic.full_id {
 358                     idx = i;
 359                 } else {
 360                     break;
 361                 }
 362             }
 363             self.ref_pics.insert(idx + 1, cpic);
 364         }
 365     }
 366     pub fn add_long_term(&mut self, lt_idx: usize, cpic: PictureInfo) {
 367         if lt_idx < self.long_term.len() {
 368             self.long_term[lt_idx] = Some(cpic);
 369         }
 370     }
 371     pub fn select_ref_pic(&self, list_id: u8, ref_id: usize) -> Option<NAVideoBufferRef<u8>> {
 372         let ref_list = if list_id == 0 { &self.ref_list0 } else { &self.ref_list1 };
 373         if ref_list.len() > ref_id {
 374             if let Some(ref pic) = ref_list[ref_id] {
 375                 Some(pic.buf.clone())
 376             } else {
 377                 None
 378             }
 379         } else {
 380             None
 381         }
 382     }
 383     pub fn get_colocated_info(&self, mb_x: usize, mb_y: usize) -> (FrameMBInfo, u16, bool) {
 384         if let Some(ref ref_pic) = &self.ref_list1[0] {
 385             let mv_info = &ref_pic.mv_info;
 386             let mb = mv_info.mbs[mb_x + mb_y * mv_info.mb_stride];
 387             (mb, ref_pic.full_id as u16, ref_pic.long_term.is_some())
 388         } else {
 389             (FrameMBInfo::default(), 0, false)
 390         }
 391     }
 392     pub fn map_ref0(&self, ref0_id: u16) -> (PicRef, bool) {
 393         let mut r0_idx = 0;
 394         let mut long = false;
 395         for (i, rpic0) in self.ref_list0.iter().enumerate() {
 396             if let Some(ref pic) = rpic0 {
 397                 if (pic.full_id as u16) == ref0_id {
 398                     r0_idx = i as u8;
 399                     long = pic.long_term.is_some();
 400                     break;
 401                 }
 402             }
 403         }
 404         (PicRef::new(r0_idx), long)
 405     }
 406     pub fn map_refs(&self, ref_idx: [PicRef; 2]) -> [u16; 2] {
 407         let r0 = ref_idx[0].index();
 408         let r1 = ref_idx[1].index();
 409         let ref0 = if r0 < self.ref_list0.len() {
 410                 if let Some(ref pic) = self.ref_list0[r0] {
 411                     pic.full_id as u16
 412                 } else {
 413                     MISSING_POC
 414                 }
 415             } else {
 416                 MISSING_POC
 417             };
 418         let ref1 = if r1 < self.ref_list1.len() {
 419                 if let Some(ref pic) = self.ref_list1[r1] {
 420                     pic.full_id as u16
 421                 } else {
 422                     MISSING_POC
 423                 }
 424             } else {
 425                 MISSING_POC
 426             };
 427         [ref0, ref1]
 428     }
 429     pub fn cmp_refs(&self, ref1: [PicRef; 2], ref2: [PicRef; 2]) -> bool {
 430         if ref1 != ref2 {
 431             self.cmp_ref(ref1[0], ref2[0], 0) && self.cmp_ref(ref1[1], ref2[1], 1)
 432         } else {
 433             true
 434         }
 435     }
 436     fn cmp_ref(&self, ref1: PicRef, ref2: PicRef, list: u8) -> bool {
 437         if ref1 == ref2 {
 438             true
 439         } else {
 440             let idx0 = ref1.index();
 441             let idx1 = ref2.index();
 442             if idx0 == idx1 {
 443                 return true;
 444             }
 445             let src = if list == 0 { &self.ref_list0 } else { &self.ref_list1 };
 446             if idx0 >= src.len() || idx1 >= src.len() {
 447 //panic!("wrong refs");
 448                 return false;
 449             }
 450             if let (Some(ref pic0), Some(ref pic1)) = (&src[idx0], &src[idx1]) {
 451                 pic0.full_id == pic1.full_id
 452             } else {
 453 //panic!("missing pics");
 454                 false
 455             }
 456         }
 457     }
 458 }
 459
 460 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) {
 461     let mut ref_pic_id = cur_id;
 462     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) {
 463         if op < 2 {
 464             if op == 0 {
 465                 ref_pic_id = ref_pic_id.wrapping_sub(num) & pic_num_mask;
 466             } else {
 467                 ref_pic_id = ref_pic_id.wrapping_add(num) & pic_num_mask;
 468             }
 469             let mut found = false;
 470             for pic in ref_pics.iter() {
 471                 if pic.id == ref_pic_id {
 472                     ref_list.push(Some(pic.clone()));
 473                     found = true;
 474                     break;
 475                 }
 476             }
 477             if !found {
 478                 ref_list.push(None);
 479             }
 480         } else {
 481             let idx = num as usize;
 482             if idx < long_term.len() {
 483                 ref_list.push(long_term[idx].clone());
 484             } else {
 485                 ref_list.push(None);
 486             }
 487         }
 488     }
 489 }