1 use nihav_core::codecs::*;
2 use nihav_codec_support::codecs::ZERO_MV;
3 use super::super::vp78::PredMode;
4 use super::super::vp78dsp::*;
5 use super::super::vp7dsp::*;
8 use super::mb_coding::*;
10 use super::motion_est::*;
13 const MBT_Q_OFFSET: usize = 3;
15 pub struct LoopParams {
16 pub loop_sharpness: u8,
17 pub loop_filter_level: u8,
21 pub struct FrameEncoder {
24 pub loop_params: LoopParams,
26 sblocks: Vec<SrcBlock>,
35 me_mode: MVSearchMode,
37 mc_buf1: NAVideoBufferRef<u8>,
38 mc_buf2: NAVideoBufferRef<u8>,
39 mv_search_last: Box<dyn MVSearch + Send>,
40 mv_search_gold: Box<dyn MVSearch + Send>,
44 pub fn new(mc_buf1: NAVideoBufferRef<u8>, mc_buf2: NAVideoBufferRef<u8>) -> Self {
45 let me_mode = MVSearchMode::default();
58 pctx: PredContext::new(),
60 loop_params: LoopParams {
67 mv_search_last: me_mode.create_search(),
68 mv_search_gold: me_mode.create_search(),
72 pub fn resize(&mut self, mb_w: usize, mb_h: usize) {
76 self.pctx.resize(mb_w, mb_h);
79 self.sblocks.reserve(mb_w * mb_h);
81 self.res.reserve(mb_w * mb_h);
83 self.mbtypes.reserve(mb_w * mb_h);
85 self.recon.reserve(mb_w * mb_h);
86 self.features.clear();
87 self.features.reserve(mb_w * mb_h);
89 pub fn set_me_params(&mut self, me_mode: MVSearchMode, me_range: i16, version: u8) {
90 self.me_range = me_range;
91 if self.me_mode != me_mode {
92 self.me_mode = me_mode;
93 self.mv_search_last = me_mode.create_search();
94 self.mv_search_gold = me_mode.create_search();
96 self.pctx.version = version;
98 pub fn load_frame(&mut self, vbuf: &NAVideoBuffer<u8>) {
99 load_blocks(vbuf, &mut self.sblocks);
102 pub fn mb_tree_search(&mut self, ref_frm: NAVideoBufferRef<u8>, mb_map: &[usize], new_mb_map: &mut [usize], mb_weights: &mut [usize]) {
103 let mut mv_est = MVEstimator::new(ref_frm, self.mc_buf1.clone(), self.me_range);
104 self.mv_search_last.preinit(&mv_est);
106 new_mb_map.copy_from_slice(mb_map);
107 for (mb_y, mb_row) in self.sblocks.chunks(self.mb_w).enumerate() {
108 for (mb_x, blk) in mb_row.iter().enumerate() {
109 let (mv, _) = self.mv_search_last.search_mb(&mut mv_est, blk, mb_x, mb_y);
112 let new_x = ((((mb_x as isize) * 64 + (mv.x as isize) + 32) >> 6).max(0) as usize).min(self.mb_w - 1);
113 let new_y = ((((mb_y as isize) * 64 + (mv.y as isize) + 32) >> 6).max(0) as usize).min(self.mb_h - 1);
114 let nidx = new_x + new_y * self.mb_w;
115 new_mb_map[mb_idx] = mb_map[nidx];
117 mb_weights[new_mb_map[mb_idx]] += 1;
123 pub fn intra_blocks(&mut self, base_q: usize, metric: &RateDistMetric, models: &VP7Models, mbt_map: Option<&[usize]>) {
124 self.mbtypes.clear();
126 self.pctx.reset_intra();
129 self.features.clear();
131 self.has_features = false;
132 if base_q > MBT_Q_OFFSET {
133 if let Some(map) = mbt_map {
134 let sum: usize = map.iter().sum();
135 let size = map.len();
136 let avg = (sum + size / 2) / size;
137 for &val in map.iter() {
139 self.features.push(1);
140 self.has_features = true;
142 self.features.push(0);
146 for _ in 0..(self.mb_w * self.mb_h) {
147 self.features.push(0);
152 let mut imctx = IntraModePredCtx {
157 ipred_y: IPredContext::default(),
158 ipred_u: IPredContext::default(),
159 ipred_v: IPredContext::default(),
160 pctx: BlockPCtx::default(),
163 for (mb_y, mb_row) in self.sblocks.chunks_mut(self.mb_w).enumerate() {
164 imctx.ipred_y.has_top = mb_y != 0;
165 imctx.ipred_u.has_top = mb_y != 0;
166 imctx.ipred_v.has_top = mb_y != 0;
168 for (mb_x, sblk) in mb_row.iter().enumerate() {
169 self.pctx.fill_ipred(0, mb_x, &mut imctx.ipred_y);
170 self.pctx.fill_ipred(1, mb_x, &mut imctx.ipred_u);
171 self.pctx.fill_ipred(2, mb_x, &mut imctx.ipred_v);
172 self.pctx.fill_pctx(mb_x, &mut imctx.pctx);
173 if self.has_features {
174 imctx.q = if self.features[mb_x + mb_y * self.mb_w] != 0 {
175 base_q - MBT_Q_OFFSET
181 let mut res = Residue::new();
182 let mut newblk = SrcBlock::default();
184 imctx.tr = self.pctx.get_ipred_tr(mb_x);
185 let mut mb_type = select_intra_mode(sblk, &mut newblk, &mut res, &imctx, MAX_DIST, MBType::InterNoMV(false, [0;4]));
187 let use_i4 = match mb_type {
188 MBType::Intra(best_ymode, best_cmode) => {
189 sblk.apply_ipred_luma(best_ymode, &imctx.ipred_y, &mut res);
190 newblk.fill_ipred_luma(best_ymode, &imctx.ipred_y);
191 sblk.apply_ipred_chroma(best_cmode, &imctx.ipred_u, &imctx.ipred_v, &mut res);
192 newblk.fill_ipred_chroma(best_cmode, &imctx.ipred_u, &imctx.ipred_v);
196 self.pctx.ymodes.set_mode(mb_x, best_ymode);
200 MBType::Intra4x4(ref i4_modes, ref mut i4ctx, best_cmode) => {
201 sblk.apply_ipred_chroma(best_cmode, &imctx.ipred_u, &imctx.ipred_v, &mut res);
202 newblk.fill_ipred_chroma(best_cmode, &imctx.ipred_u, &imctx.ipred_v);
205 self.pctx.ymodes.set_modes4x4(mb_x, i4_modes, i4ctx);
213 self.pctx.set_nz(mb_x, &res);
214 let mut recon = res.clone();
216 self.mbtypes.push(mb_type);
219 recon.add_residue(&mut newblk);
221 recon.add_residue_chroma(&mut newblk);
224 self.pctx.update_mb(&newblk, mb_x);
225 self.recon.push(newblk);
227 self.pctx.update_mb_row();
230 pub fn inter_blocks(&mut self, q: usize, metric: &RateDistMetric, models: &VP7Models, last_frame: &NABufferType, gold_frame: &NABufferType) {
231 self.has_features = false;
233 let mut mv_est_last = MVEstimator::new(last_frame.get_vbuf().unwrap(), self.mc_buf1.clone(), self.me_range);
234 self.mv_search_last.preinit(&mv_est_last);
235 let mut mv_est_gold = if let Some(gbuf) = gold_frame.get_vbuf() {
236 let mv_est = MVEstimator::new(gbuf, self.mc_buf2.clone(), self.me_range);
237 self.mv_search_gold.preinit(&mv_est);
243 self.mbtypes.clear();
245 self.pctx.save_dc_pred();
248 self.features.clear();
250 let mut imctx = IntraModePredCtx {
255 ipred_y: IPredContext::default(),
256 ipred_u: IPredContext::default(),
257 ipred_v: IPredContext::default(),
258 pctx: BlockPCtx::default(),
261 for (mb_y, mb_row) in self.sblocks.chunks_mut(self.mb_w).enumerate() {
262 imctx.ipred_y.has_top = mb_y != 0;
263 imctx.ipred_u.has_top = mb_y != 0;
264 imctx.ipred_v.has_top = mb_y != 0;
266 for (mb_x, sblk) in mb_row.iter().enumerate() {
267 self.pctx.fill_ipred(0, mb_x, &mut imctx.ipred_y);
268 self.pctx.fill_ipred(1, mb_x, &mut imctx.ipred_u);
269 self.pctx.fill_ipred(2, mb_x, &mut imctx.ipred_v);
270 self.pctx.fill_pctx(mb_x, &mut imctx.pctx);
272 let mut res = Residue::new();
273 let mut newblk = SrcBlock::default();
275 let (mvprobs, nearest_mv, near_mv, pred_mv) = self.pctx.find_mv_pred(mb_x, mb_y);
277 let (mv, _dist) = self.mv_search_last.search_mb(&mut mv_est_last, sblk, mb_x, mb_y);
279 mv_est_last.get_mb(&mut newblk, mb_x, mb_y, mv);
280 let mv_nits_dist = metric.calc_metric(0, inter_mv_nits(mv, &mvprobs, nearest_mv, near_mv, pred_mv, models));
281 let last_dist = calc_inter_mb_dist(sblk, &newblk, &mut res, &imctx, self.pctx.get_y2_dc_pred(true)) + mv_nits_dist;
283 let (gmv, gold_dist) = if last_dist > SMALL_DIST {
284 if let Some(ref mut mv_est) = &mut mv_est_gold {
285 let (gmv, _gdist) = self.mv_search_gold.search_mb(mv_est, sblk, mb_x, mb_y);
286 mv_est.get_mb(&mut newblk, mb_x, mb_y, gmv);
287 let mv_nits_dist = metric.calc_metric(0, inter_mv_nits(gmv, &mvprobs, nearest_mv, near_mv, pred_mv, models));
288 let gdist = calc_inter_mb_dist(sblk, &newblk, &mut res, &imctx, self.pctx.get_y2_dc_pred(false)) + mv_nits_dist;
297 let (last, mut inter_dist, mv, mv_est) = if last_dist < gold_dist {
298 (true, last_dist, mv, &mut mv_est_last)
299 } else if let Some (ref mut mv_est) = &mut mv_est_gold {
300 (false, gold_dist, gmv, mv_est)
305 let mut mb_type = if mv == ZERO_MV {
306 MBType::InterNoMV(last, mvprobs)
307 } else if mv == nearest_mv {
308 MBType::InterNearest(last, mvprobs)
309 } else if mv == near_mv {
310 MBType::InterNear(last, mvprobs)
312 MBType::InterMV(last, mvprobs, mv - pred_mv)
314 if inter_dist > SMALL_DIST {
315 if let MBType::InterMV(_, _, _) = mb_type { // xxx: maybe do it for all types?
316 let mv_search = if last { &mut self.mv_search_last } else { &mut self.mv_search_gold };
317 if let Some((mbt, dist)) = try_inter_split(sblk, &mut newblk, &mut res, mvprobs, nearest_mv, near_mv, pred_mv, last, mb_x, mb_y, mv_search, mv_est, &mut self.pctx, &imctx, inter_dist) {
324 if inter_dist > SMALL_DIST {
325 imctx.tr = self.pctx.get_ipred_tr(mb_x);
326 mb_type = select_intra_mode(sblk, &mut newblk, &mut res, &imctx, inter_dist, mb_type);
329 self.mbtypes.push(mb_type);
332 MBType::Intra(ymode, cmode) => {
333 newblk.fill_ipred_luma(ymode, &imctx.ipred_y);
334 newblk.fill_ipred_chroma(cmode, &imctx.ipred_u, &imctx.ipred_v);
335 self.pctx.ymodes.set_mode(mb_x, ymode);
336 self.pctx.fill_mv(mb_x, mb_y, ZERO_MV);
338 MBType::Intra4x4(ref i4_modes, ref mut i4ctx, cmode) => {
339 newblk.fill_ipred_chroma(cmode, &imctx.ipred_u, &imctx.ipred_v);
340 self.pctx.ymodes.set_modes4x4(mb_x, i4_modes, i4ctx);
341 self.pctx.fill_mv(mb_x, mb_y, ZERO_MV);
343 MBType::InterNoMV(_, _) |
344 MBType::InterNearest(_, _) |
345 MBType::InterNear(_, _) |
346 MBType::InterMV(_, _, _) => {
347 mv_est.get_mb(&mut newblk, mb_x, mb_y, mv);
348 self.pctx.fill_mv(mb_x, mb_y, mv);
349 self.pctx.ymodes.set_mode(mb_x, PredMode::Inter);
351 MBType::InterSplitMV(_, _, _, _, _) => {
352 self.pctx.ymodes.set_mode(mb_x, PredMode::Inter);
353 recon_split_mb(&mut newblk, mb_x, mb_y, &self.pctx.mvs, self.pctx.mv_stride, mv_est);
356 if let MBType::Intra4x4(_, _, _) = mb_type {
357 res.set_chroma_from_diff(&sblk.chroma, &newblk.chroma);
360 res.set_luma_from_diff(&sblk.luma, &newblk.luma);
361 res.set_chroma_from_diff(&sblk.chroma, &newblk.chroma);
364 if !mb_type.is_intra() {
365 requant_y2_dc(&mut res.dcs[0], q);
366 self.pctx.predict_y2_dc(&mut res.dcs[0], last);
371 self.pctx.set_nz(mb_x, &res);
372 let mut recon = res.clone();
374 self.features.push(0);
375 if let MBType::Intra4x4(_, _, _) = mb_type {
376 recon.add_residue_chroma(&mut newblk);
378 recon.add_residue(&mut newblk);
380 self.pctx.update_mb(&newblk, mb_x);
381 self.recon.push(newblk);
383 self.pctx.update_mb_row();
386 pub fn encode_features(&self, bc: &mut BoolEncoder, q: usize, models: &VP7Models) -> EncoderResult<()> {
387 if self.has_features {
388 // first feature - quantiser
389 bc.put_bool(true, 128)?;
390 bc.put_byte(models.feature_present[0])?;
391 for &prob in models.feature_tree_probs[0].iter() {
392 bc.put_bool(prob != 255, 128)?;
397 bc.put_bool(true, 128)?;
398 bc.put_bits((q - MBT_Q_OFFSET) as u32, 7)?;
400 bc.put_bool(false, 128)?; // other quants
405 bc.put_bool(false, 128)?;
409 bc.put_bool(false, 128)?;
414 pub fn encode_mb_types(&self, bc: &mut BoolEncoder, is_intra: bool, models: &VP7Models) -> EncoderResult<()> {
415 for (mb_type, &feature) in self.mbtypes.iter().zip(self.features.iter()) {
416 if self.has_features {
417 bc.encode_feature(0, if feature == 0 { None } else { Some(0) }, models)?;
419 bc.encode_mb_type(is_intra, mb_type, models)?;
423 pub fn encode_residues(&mut self, bc: &mut BoolEncoder, models: &VP7Models) -> EncoderResult<()> {
425 //self.pctx.restore_dc_pred();
426 for (_mb_y, mb_row) in self.res.chunks(self.mb_w).enumerate() {
427 for (mb_x, blk) in mb_row.iter().enumerate() {
429 let pctx = (self.pctx.nz_y2_left as u8) + (self.pctx.nz_y2_top[mb_x] as u8);
430 bc.encode_subblock(&blk.dcs, 1, pctx, models)?;
431 let has_nz = blk.dcs.has_nz();
432 self.pctx.nz_y2_left = has_nz;
433 self.pctx.nz_y2_top[mb_x] = has_nz;
435 let ytype = if blk.has_dc { 0 } else { 3 };
436 for (y, blk_row) in blk.luma.chunks(4).enumerate() {
437 for (x, blk) in blk_row.iter().enumerate() {
438 let pctx = (self.pctx.nz_y_left[y] as u8) + (self.pctx.nz_y_top[mb_x * 4 + x] as u8);
439 bc.encode_subblock(blk, ytype, pctx, models)?;
440 let has_nz = blk.has_nz();
441 self.pctx.nz_y_left[y] = has_nz;
442 self.pctx.nz_y_top[mb_x * 4 + x] = has_nz;
446 for (c, chroma) in blk.chroma.iter().enumerate() {
447 for (y, blk_row) in chroma.chunks(2).enumerate() {
448 for (x, blk) in blk_row.iter().enumerate() {
449 let pctx = (self.pctx.nz_c_left[c][y] as u8) + (self.pctx.nz_c_top[c][mb_x * 2 + x] as u8);
450 bc.encode_subblock(blk, 2, pctx, models)?;
451 let has_nz = blk.has_nz();
452 self.pctx.nz_c_left[c][y] = has_nz;
453 self.pctx.nz_c_top[c][mb_x * 2 + x] = has_nz;
458 self.pctx.update_mb_row();
462 pub fn generate_models(&mut self, is_intra: bool, stats: &mut VP7ModelsStat) {
464 let est = Estimator::new();
466 if self.has_features {
467 for &feat in self.features.iter() {
468 est.estimate_feature(0, if feat == 0 { None } else { Some(0) }, stats);
471 for (mbt_row, mb_row) in self.mbtypes.chunks(self.mb_w).zip(self.res.chunks(self.mb_w)) {
472 for (mb_x, (mbtype, blk)) in mbt_row.iter().zip(mb_row.iter()).enumerate() {
473 est.estimate_mb_type(is_intra, mbtype, stats);
475 let pctx = (self.pctx.nz_y2_left as u8) + (self.pctx.nz_y2_top[mb_x] as u8);
476 est.estimate_subblock(&blk.dcs, 1, pctx, stats);
477 let has_nz = blk.dcs.has_nz();
478 self.pctx.nz_y2_left = has_nz;
479 self.pctx.nz_y2_top[mb_x] = has_nz;
481 let ytype = if blk.has_dc { 0 } else { 3 };
482 for (y, blk_row) in blk.luma.chunks(4).enumerate() {
483 for (x, blk) in blk_row.iter().enumerate() {
484 let pctx = (self.pctx.nz_y_left[y] as u8) + (self.pctx.nz_y_top[mb_x * 4 + x] as u8);
485 est.estimate_subblock(blk, ytype, pctx, stats);
486 let has_nz = blk.has_nz();
487 self.pctx.nz_y_left[y] = has_nz;
488 self.pctx.nz_y_top[mb_x * 4 + x] = has_nz;
492 for (c, chroma) in blk.chroma.iter().enumerate() {
493 for (y, blk_row) in chroma.chunks(2).enumerate() {
494 for (x, blk) in blk_row.iter().enumerate() {
495 let pctx = (self.pctx.nz_c_left[c][y] as u8) + (self.pctx.nz_c_top[c][mb_x * 2 + x] as u8);
496 est.estimate_subblock(blk, 2, pctx, stats);
497 let has_nz = blk.has_nz();
498 self.pctx.nz_c_left[c][y] = has_nz;
499 self.pctx.nz_c_top[c][mb_x * 2 + x] = has_nz;
504 self.pctx.update_mb_row();
507 pub fn reconstruct_frame(&mut self, frm: &mut NASimpleVideoFrame<u8>, is_intra: bool) {
508 let mut yidx = frm.offset[0];
509 let mut uidx = frm.offset[1];
510 let mut vidx = frm.offset[2];
511 let ystride = frm.stride[0];
512 let ustride = frm.stride[1];
513 let vstride = frm.stride[2];
515 for (mb_y, (f_row, mb_row)) in self.features.chunks(self.mb_w).zip(self.recon.chunks(self.mb_w)).enumerate() {
516 for (mb_x, (&feature, sblk)) in f_row.iter().zip(mb_row.iter()).enumerate() {
517 let dst = &mut frm.data[yidx + mb_x * 16..];
518 for (dst, src) in dst.chunks_mut(ystride).zip(sblk.luma.chunks(16)) {
519 dst[..16].copy_from_slice(src);
521 let dst = &mut frm.data[uidx + mb_x * 8..];
522 for (dst, src) in dst.chunks_mut(ustride).zip(sblk.chroma[0].chunks(8)) {
523 dst[..8].copy_from_slice(src);
525 let dst = &mut frm.data[vidx + mb_x * 8..];
526 for (dst, src) in dst.chunks_mut(vstride).zip(sblk.chroma[1].chunks(8)) {
527 dst[..8].copy_from_slice(src);
530 let loop_str = if feature != 2 {
531 self.loop_params.loop_filter_level
533 loop_filter_mb(frm, mb_x, mb_y, loop_str, &self.loop_params, is_intra);
535 yidx += ystride * 16;
542 fn loop_filter_mb(dframe: &mut NASimpleVideoFrame<u8>, mb_x: usize, mb_y: usize, loop_str: u8, loop_params: &LoopParams, is_intra: bool) {
543 const HIGH_EDGE_VAR_THR: [[u8; 64]; 2] = [
545 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
546 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
547 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
548 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3
550 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
551 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
552 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2,
553 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2
556 let edge_thr = i16::from(loop_str) + 2;
557 let luma_thr = i16::from(loop_str);
558 let chroma_thr = i16::from(loop_str) * 2;
559 let inner_thr = if loop_params.loop_sharpness == 0 {
562 let bound1 = i16::from(9 - loop_params.loop_sharpness);
563 let shift = (loop_params.loop_sharpness + 3) >> 2;
564 (i16::from(loop_str) >> shift).min(bound1)
566 let hev_thr = i16::from(HIGH_EDGE_VAR_THR[if is_intra { 1 } else { 0 }][loop_str as usize]);
568 let ystride = dframe.stride[0];
569 let ustride = dframe.stride[1];
570 let vstride = dframe.stride[2];
571 let ypos = dframe.offset[0] + mb_x * 16 + mb_y * 16 * ystride;
572 let upos = dframe.offset[1] + mb_x * 8 + mb_y * 8 * ustride;
573 let vpos = dframe.offset[2] + mb_x * 8 + mb_y * 8 * vstride;
575 let (loop_edge, loop_inner) = if loop_params.lf_simple {
576 (simple_loop_filter as LoopFilterFunc, simple_loop_filter as LoopFilterFunc)
578 (normal_loop_filter_edge as LoopFilterFunc, normal_loop_filter_inner as LoopFilterFunc)
582 loop_edge(dframe.data, ypos, 1, ystride, 16, edge_thr, inner_thr, hev_thr);
583 loop_edge(dframe.data, upos, 1, ustride, 8, edge_thr, inner_thr, hev_thr);
584 loop_edge(dframe.data, vpos, 1, vstride, 8, edge_thr, inner_thr, hev_thr);
587 loop_edge(dframe.data, ypos, ystride, 1, 16, edge_thr, inner_thr, hev_thr);
588 loop_edge(dframe.data, upos, ustride, 1, 8, edge_thr, inner_thr, hev_thr);
589 loop_edge(dframe.data, vpos, vstride, 1, 8, edge_thr, inner_thr, hev_thr);
593 loop_inner(dframe.data, ypos + y * 4 * ystride, ystride, 1, 16, luma_thr, inner_thr, hev_thr);
595 loop_inner(dframe.data, upos + 4 * ustride, ustride, 1, 8, chroma_thr, inner_thr, hev_thr);
596 loop_inner(dframe.data, vpos + 4 * vstride, vstride, 1, 8, chroma_thr, inner_thr, hev_thr);
599 loop_inner(dframe.data, ypos + x * 4, 1, ystride, 16, luma_thr, inner_thr, hev_thr);
601 loop_inner(dframe.data, upos + 4, 1, ustride, 8, chroma_thr, inner_thr, hev_thr);
602 loop_inner(dframe.data, vpos + 4, 1, vstride, 8, chroma_thr, inner_thr, hev_thr);