]> git.nihav.org Git - nihav.git/blob - nihav-itu/src/codecs/h264/pic_ref.rs
h264: fix sliding window frame reference dropping
[nihav.git] / 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 }