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);
151 for _ in 0..(self.mb_w * self.mb_h) {
152 self.features.push(0);
156 let mut imctx = IntraModePredCtx {
161 ipred_y: IPredContext::default(),
162 ipred_u: IPredContext::default(),
163 ipred_v: IPredContext::default(),
164 pctx: BlockPCtx::default(),
167 for (mb_y, mb_row) in self.sblocks.chunks_mut(self.mb_w).enumerate() {
168 imctx.ipred_y.has_top = mb_y != 0;
169 imctx.ipred_u.has_top = mb_y != 0;
170 imctx.ipred_v.has_top = mb_y != 0;
172 for (mb_x, sblk) in mb_row.iter().enumerate() {
173 self.pctx.fill_ipred(0, mb_x, &mut imctx.ipred_y);
174 self.pctx.fill_ipred(1, mb_x, &mut imctx.ipred_u);
175 self.pctx.fill_ipred(2, mb_x, &mut imctx.ipred_v);
176 self.pctx.fill_pctx(mb_x, &mut imctx.pctx);
177 if self.has_features {
178 imctx.q = if self.features[mb_x + mb_y * self.mb_w] != 0 {
179 base_q - MBT_Q_OFFSET
185 let mut res = Residue::new();
186 let mut newblk = SrcBlock::default();
188 imctx.tr = self.pctx.get_ipred_tr(mb_x);
189 let mut mb_type = select_intra_mode(sblk, &mut newblk, &mut res, &imctx, MAX_DIST, MBType::InterNoMV(false, [0;4]));
191 let use_i4 = match mb_type {
192 MBType::Intra(best_ymode, best_cmode) => {
193 sblk.apply_ipred_luma(best_ymode, &imctx.ipred_y, &mut res);
194 newblk.fill_ipred_luma(best_ymode, &imctx.ipred_y);
195 sblk.apply_ipred_chroma(best_cmode, &imctx.ipred_u, &imctx.ipred_v, &mut res);
196 newblk.fill_ipred_chroma(best_cmode, &imctx.ipred_u, &imctx.ipred_v);
200 self.pctx.ymodes.set_mode(mb_x, best_ymode);
204 MBType::Intra4x4(ref i4_modes, ref mut i4ctx, best_cmode) => {
205 sblk.apply_ipred_chroma(best_cmode, &imctx.ipred_u, &imctx.ipred_v, &mut res);
206 newblk.fill_ipred_chroma(best_cmode, &imctx.ipred_u, &imctx.ipred_v);
209 self.pctx.ymodes.set_modes4x4(mb_x, i4_modes, i4ctx);
217 self.pctx.set_nz(mb_x, &res);
218 let mut recon = res.clone();
220 self.mbtypes.push(mb_type);
223 recon.add_residue(&mut newblk);
225 recon.add_residue_chroma(&mut newblk);
228 self.pctx.update_mb(&newblk, mb_x);
229 self.recon.push(newblk);
231 self.pctx.update_mb_row();
234 pub fn inter_blocks(&mut self, q: usize, metric: &RateDistMetric, models: &VP7Models, last_frame: &NABufferType, gold_frame: &NABufferType) {
235 self.has_features = false;
237 let mut mv_est_last = MVEstimator::new(last_frame.get_vbuf().unwrap(), self.mc_buf1.clone(), self.me_range);
238 self.mv_search_last.preinit(&mv_est_last);
239 let mut mv_est_gold = if let Some(gbuf) = gold_frame.get_vbuf() {
240 let mv_est = MVEstimator::new(gbuf, self.mc_buf2.clone(), self.me_range);
241 self.mv_search_gold.preinit(&mv_est);
247 self.mbtypes.clear();
249 self.pctx.save_dc_pred();
252 self.features.clear();
254 let mut imctx = IntraModePredCtx {
259 ipred_y: IPredContext::default(),
260 ipred_u: IPredContext::default(),
261 ipred_v: IPredContext::default(),
262 pctx: BlockPCtx::default(),
265 for (mb_y, mb_row) in self.sblocks.chunks_mut(self.mb_w).enumerate() {
266 imctx.ipred_y.has_top = mb_y != 0;
267 imctx.ipred_u.has_top = mb_y != 0;
268 imctx.ipred_v.has_top = mb_y != 0;
270 for (mb_x, sblk) in mb_row.iter().enumerate() {
271 self.pctx.fill_ipred(0, mb_x, &mut imctx.ipred_y);
272 self.pctx.fill_ipred(1, mb_x, &mut imctx.ipred_u);
273 self.pctx.fill_ipred(2, mb_x, &mut imctx.ipred_v);
274 self.pctx.fill_pctx(mb_x, &mut imctx.pctx);
276 let mut res = Residue::new();
277 let mut newblk = SrcBlock::default();
279 let (mvprobs, nearest_mv, near_mv, pred_mv) = self.pctx.find_mv_pred(mb_x, mb_y);
281 let (mv, _dist) = self.mv_search_last.search_mb(&mut mv_est_last, sblk, mb_x, mb_y);
283 mv_est_last.get_mb(&mut newblk, mb_x, mb_y, mv);
284 let mv_nits_dist = metric.calc_metric(0, inter_mv_nits(mv, &mvprobs, nearest_mv, near_mv, pred_mv, models));
285 let last_dist = calc_inter_mb_dist(sblk, &newblk, &mut res, &imctx, self.pctx.get_y2_dc_pred(true)) + mv_nits_dist;
287 let (gmv, gold_dist) = if last_dist > SMALL_DIST {
288 if let Some(ref mut mv_est) = &mut mv_est_gold {
289 let (gmv, _gdist) = self.mv_search_gold.search_mb(mv_est, sblk, mb_x, mb_y);
290 mv_est.get_mb(&mut newblk, mb_x, mb_y, gmv);
291 let mv_nits_dist = metric.calc_metric(0, inter_mv_nits(gmv, &mvprobs, nearest_mv, near_mv, pred_mv, models));
292 let gdist = calc_inter_mb_dist(sblk, &newblk, &mut res, &imctx, self.pctx.get_y2_dc_pred(false)) + mv_nits_dist;
301 let (last, mut inter_dist, mv, mv_est) = if last_dist < gold_dist {
302 (true, last_dist, mv, &mut mv_est_last)
303 } else if let Some (ref mut mv_est) = &mut mv_est_gold {
304 (false, gold_dist, gmv, mv_est)
309 let mut mb_type = if mv == ZERO_MV {
310 MBType::InterNoMV(last, mvprobs)
311 } else if mv == nearest_mv {
312 MBType::InterNearest(last, mvprobs)
313 } else if mv == near_mv {
314 MBType::InterNear(last, mvprobs)
316 MBType::InterMV(last, mvprobs, mv - pred_mv)
318 if inter_dist > SMALL_DIST {
319 if let MBType::InterMV(_, _, _) = mb_type { // xxx: maybe do it for all types?
320 let mv_search = if last { &mut self.mv_search_last } else { &mut self.mv_search_gold };
321 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) {
328 if inter_dist > SMALL_DIST {
329 imctx.tr = self.pctx.get_ipred_tr(mb_x);
330 mb_type = select_intra_mode(sblk, &mut newblk, &mut res, &imctx, inter_dist, mb_type);
333 self.mbtypes.push(mb_type);
336 MBType::Intra(ymode, cmode) => {
337 newblk.fill_ipred_luma(ymode, &imctx.ipred_y);
338 newblk.fill_ipred_chroma(cmode, &imctx.ipred_u, &imctx.ipred_v);
339 self.pctx.ymodes.set_mode(mb_x, ymode);
340 self.pctx.fill_mv(mb_x, mb_y, ZERO_MV);
342 MBType::Intra4x4(ref i4_modes, ref mut i4ctx, cmode) => {
343 newblk.fill_ipred_chroma(cmode, &imctx.ipred_u, &imctx.ipred_v);
344 self.pctx.ymodes.set_modes4x4(mb_x, i4_modes, i4ctx);
345 self.pctx.fill_mv(mb_x, mb_y, ZERO_MV);
347 MBType::InterNoMV(_, _) |
348 MBType::InterNearest(_, _) |
349 MBType::InterNear(_, _) |
350 MBType::InterMV(_, _, _) => {
351 mv_est.get_mb(&mut newblk, mb_x, mb_y, mv);
352 self.pctx.fill_mv(mb_x, mb_y, mv);
353 self.pctx.ymodes.set_mode(mb_x, PredMode::Inter);
355 MBType::InterSplitMV(_, _, _, _, _) => {
356 self.pctx.ymodes.set_mode(mb_x, PredMode::Inter);
357 recon_split_mb(&mut newblk, mb_x, mb_y, &self.pctx.mvs, self.pctx.mv_stride, mv_est);
360 if let MBType::Intra4x4(_, _, _) = mb_type {
361 res.set_chroma_from_diff(&sblk.chroma, &newblk.chroma);
364 res.set_luma_from_diff(&sblk.luma, &newblk.luma);
365 res.set_chroma_from_diff(&sblk.chroma, &newblk.chroma);
368 if !mb_type.is_intra() {
369 requant_y2_dc(&mut res.dcs[0], q);
370 self.pctx.predict_y2_dc(&mut res.dcs[0], last);
375 self.pctx.set_nz(mb_x, &res);
376 let mut recon = res.clone();
378 self.features.push(0);
379 if let MBType::Intra4x4(_, _, _) = mb_type {
380 recon.add_residue_chroma(&mut newblk);
382 recon.add_residue(&mut newblk);
384 self.pctx.update_mb(&newblk, mb_x);
385 self.recon.push(newblk);
387 self.pctx.update_mb_row();
390 pub fn encode_features(&self, bc: &mut BoolEncoder, q: usize, models: &VP7Models) -> EncoderResult<()> {
391 if self.has_features {
392 // first feature - quantiser
393 bc.put_bool(true, 128)?;
394 bc.put_byte(models.feature_present[0])?;
395 for &prob in models.feature_tree_probs[0].iter() {
396 bc.put_bool(prob != 255, 128)?;
401 bc.put_bool(true, 128)?;
402 bc.put_bits((q - MBT_Q_OFFSET) as u32, 7)?;
404 bc.put_bool(false, 128)?; // other quants
409 bc.put_bool(false, 128)?;
413 bc.put_bool(false, 128)?;
418 pub fn encode_mb_types(&self, bc: &mut BoolEncoder, is_intra: bool, models: &VP7Models) -> EncoderResult<()> {
419 for (mb_type, &feature) in self.mbtypes.iter().zip(self.features.iter()) {
420 if self.has_features {
421 bc.encode_feature(0, if feature == 0 { None } else { Some(0) }, models)?;
423 bc.encode_mb_type(is_intra, mb_type, models)?;
427 pub fn encode_residues(&mut self, bc: &mut BoolEncoder, models: &VP7Models) -> EncoderResult<()> {
429 //self.pctx.restore_dc_pred();
430 for (_mb_y, mb_row) in self.res.chunks(self.mb_w).enumerate() {
431 for (mb_x, blk) in mb_row.iter().enumerate() {
433 let pctx = (self.pctx.nz_y2_left as u8) + (self.pctx.nz_y2_top[mb_x] as u8);
434 bc.encode_subblock(&blk.dcs, 1, pctx, models)?;
435 let has_nz = blk.dcs.has_nz();
436 self.pctx.nz_y2_left = has_nz;
437 self.pctx.nz_y2_top[mb_x] = has_nz;
439 let ytype = if blk.has_dc { 0 } else { 3 };
440 for (y, blk_row) in blk.luma.chunks(4).enumerate() {
441 for (x, blk) in blk_row.iter().enumerate() {
442 let pctx = (self.pctx.nz_y_left[y] as u8) + (self.pctx.nz_y_top[mb_x * 4 + x] as u8);
443 bc.encode_subblock(blk, ytype, pctx, models)?;
444 let has_nz = blk.has_nz();
445 self.pctx.nz_y_left[y] = has_nz;
446 self.pctx.nz_y_top[mb_x * 4 + x] = has_nz;
450 for (c, chroma) in blk.chroma.iter().enumerate() {
451 for (y, blk_row) in chroma.chunks(2).enumerate() {
452 for (x, blk) in blk_row.iter().enumerate() {
453 let pctx = (self.pctx.nz_c_left[c][y] as u8) + (self.pctx.nz_c_top[c][mb_x * 2 + x] as u8);
454 bc.encode_subblock(blk, 2, pctx, models)?;
455 let has_nz = blk.has_nz();
456 self.pctx.nz_c_left[c][y] = has_nz;
457 self.pctx.nz_c_top[c][mb_x * 2 + x] = has_nz;
462 self.pctx.update_mb_row();
466 pub fn generate_models(&mut self, is_intra: bool, stats: &mut VP7ModelsStat) {
468 let est = Estimator::new();
470 if self.has_features {
471 for &feat in self.features.iter() {
472 est.estimate_feature(0, if feat == 0 { None } else { Some(0) }, stats);
475 for (mbt_row, mb_row) in self.mbtypes.chunks(self.mb_w).zip(self.res.chunks(self.mb_w)) {
476 for (mb_x, (mbtype, blk)) in mbt_row.iter().zip(mb_row.iter()).enumerate() {
477 est.estimate_mb_type(is_intra, mbtype, stats);
479 let pctx = (self.pctx.nz_y2_left as u8) + (self.pctx.nz_y2_top[mb_x] as u8);
480 est.estimate_subblock(&blk.dcs, 1, pctx, stats);
481 let has_nz = blk.dcs.has_nz();
482 self.pctx.nz_y2_left = has_nz;
483 self.pctx.nz_y2_top[mb_x] = has_nz;
485 let ytype = if blk.has_dc { 0 } else { 3 };
486 for (y, blk_row) in blk.luma.chunks(4).enumerate() {
487 for (x, blk) in blk_row.iter().enumerate() {
488 let pctx = (self.pctx.nz_y_left[y] as u8) + (self.pctx.nz_y_top[mb_x * 4 + x] as u8);
489 est.estimate_subblock(blk, ytype, pctx, stats);
490 let has_nz = blk.has_nz();
491 self.pctx.nz_y_left[y] = has_nz;
492 self.pctx.nz_y_top[mb_x * 4 + x] = has_nz;
496 for (c, chroma) in blk.chroma.iter().enumerate() {
497 for (y, blk_row) in chroma.chunks(2).enumerate() {
498 for (x, blk) in blk_row.iter().enumerate() {
499 let pctx = (self.pctx.nz_c_left[c][y] as u8) + (self.pctx.nz_c_top[c][mb_x * 2 + x] as u8);
500 est.estimate_subblock(blk, 2, pctx, stats);
501 let has_nz = blk.has_nz();
502 self.pctx.nz_c_left[c][y] = has_nz;
503 self.pctx.nz_c_top[c][mb_x * 2 + x] = has_nz;
508 self.pctx.update_mb_row();
511 pub fn reconstruct_frame(&mut self, frm: &mut NASimpleVideoFrame<u8>, is_intra: bool) {
512 let mut yidx = frm.offset[0];
513 let mut uidx = frm.offset[1];
514 let mut vidx = frm.offset[2];
515 let ystride = frm.stride[0];
516 let ustride = frm.stride[1];
517 let vstride = frm.stride[2];
519 for (mb_y, (f_row, mb_row)) in self.features.chunks(self.mb_w).zip(self.recon.chunks(self.mb_w)).enumerate() {
520 for (mb_x, (&feature, sblk)) in f_row.iter().zip(mb_row.iter()).enumerate() {
521 let dst = &mut frm.data[yidx + mb_x * 16..];
522 for (dst, src) in dst.chunks_mut(ystride).zip(sblk.luma.chunks(16)) {
523 dst[..16].copy_from_slice(src);
525 let dst = &mut frm.data[uidx + mb_x * 8..];
526 for (dst, src) in dst.chunks_mut(ustride).zip(sblk.chroma[0].chunks(8)) {
527 dst[..8].copy_from_slice(src);
529 let dst = &mut frm.data[vidx + mb_x * 8..];
530 for (dst, src) in dst.chunks_mut(vstride).zip(sblk.chroma[1].chunks(8)) {
531 dst[..8].copy_from_slice(src);
534 let loop_str = if feature != 2 {
535 self.loop_params.loop_filter_level
537 loop_filter_mb(frm, mb_x, mb_y, loop_str, &self.loop_params, is_intra);
539 yidx += ystride * 16;
546 fn loop_filter_mb(dframe: &mut NASimpleVideoFrame<u8>, mb_x: usize, mb_y: usize, loop_str: u8, loop_params: &LoopParams, is_intra: bool) {
547 const HIGH_EDGE_VAR_THR: [[u8; 64]; 2] = [
549 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
550 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
551 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
552 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3
554 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
555 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
556 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2,
557 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2
560 let edge_thr = i16::from(loop_str) + 2;
561 let luma_thr = i16::from(loop_str);
562 let chroma_thr = i16::from(loop_str) * 2;
563 let inner_thr = if loop_params.loop_sharpness == 0 {
566 let bound1 = i16::from(9 - loop_params.loop_sharpness);
567 let shift = (loop_params.loop_sharpness + 3) >> 2;
568 (i16::from(loop_str) >> shift).min(bound1)
570 let hev_thr = i16::from(HIGH_EDGE_VAR_THR[if is_intra { 1 } else { 0 }][loop_str as usize]);
572 let ystride = dframe.stride[0];
573 let ustride = dframe.stride[1];
574 let vstride = dframe.stride[2];
575 let ypos = dframe.offset[0] + mb_x * 16 + mb_y * 16 * ystride;
576 let upos = dframe.offset[1] + mb_x * 8 + mb_y * 8 * ustride;
577 let vpos = dframe.offset[2] + mb_x * 8 + mb_y * 8 * vstride;
579 let (loop_edge, loop_inner) = if loop_params.lf_simple {
580 (simple_loop_filter as LoopFilterFunc, simple_loop_filter as LoopFilterFunc)
582 (normal_loop_filter_edge as LoopFilterFunc, normal_loop_filter_inner as LoopFilterFunc)
586 loop_edge(dframe.data, ypos, 1, ystride, 16, edge_thr, inner_thr, hev_thr);
587 loop_edge(dframe.data, upos, 1, ustride, 8, edge_thr, inner_thr, hev_thr);
588 loop_edge(dframe.data, vpos, 1, vstride, 8, edge_thr, inner_thr, hev_thr);
591 loop_edge(dframe.data, ypos, ystride, 1, 16, edge_thr, inner_thr, hev_thr);
592 loop_edge(dframe.data, upos, ustride, 1, 8, edge_thr, inner_thr, hev_thr);
593 loop_edge(dframe.data, vpos, vstride, 1, 8, edge_thr, inner_thr, hev_thr);
597 loop_inner(dframe.data, ypos + y * 4 * ystride, ystride, 1, 16, luma_thr, inner_thr, hev_thr);
599 loop_inner(dframe.data, upos + 4 * ustride, ustride, 1, 8, chroma_thr, inner_thr, hev_thr);
600 loop_inner(dframe.data, vpos + 4 * vstride, vstride, 1, 8, chroma_thr, inner_thr, hev_thr);
603 loop_inner(dframe.data, ypos + x * 4, 1, ystride, 16, luma_thr, inner_thr, hev_thr);
605 loop_inner(dframe.data, upos + 4, 1, ustride, 8, chroma_thr, inner_thr, hev_thr);
606 loop_inner(dframe.data, vpos + 4, 1, vstride, 8, chroma_thr, inner_thr, hev_thr);