From fef34fda77344a793279f38107cba67d5a2e6535 Mon Sep 17 00:00:00 2001
From: Adrien Maglo <adrien@videolabs.io>
Date: Tue, 26 Feb 2019 15:49:27 +0100
Subject: [PATCH] diamond_me: save only selected frame motion vectors

Save them by reference frame types instead of picture slot.
Do not add several times the zero motion vector to the predictor list.
---
 src/encoder.rs   | 36 ++++++++++++++++++++++++++++
 src/me.rs        | 61 ++++++++++++++++++++++++++----------------------
 src/partition.rs |  4 ++++
 src/rdo.rs       |  9 +------
 4 files changed, 74 insertions(+), 36 deletions(-)

diff --git a/src/encoder.rs b/src/encoder.rs
index 2689f997..073773ad 100644
--- a/src/encoder.rs
+++ b/src/encoder.rs
@@ -1110,6 +1110,20 @@ pub fn motion_compensate<T: Pixel>(
   }
 }
 
+pub fn save_block_motion<T: Pixel>(
+   fs: &mut FrameState<T>,
+   w_in_b: usize, _h_in_b: usize,
+   bsize: BlockSize, bo: &BlockOffset,
+   ref_frame: usize, mv: MotionVector,
+) {
+  let frame_mvs = &mut fs.frame_mvs;
+  for mi_y in (bo.y)..(bo.y + bsize.height_mi()) {
+    for mi_x in (bo.x)..(bo.x + bsize.width_mi()) {
+      frame_mvs[ref_frame][mi_y][mi_x] = mv;
+    }
+  }
+}
+
 pub fn encode_block_a<T: Pixel>(
   seq: &Sequence, fs: &FrameState<T>,
   cw: &mut ContextWriter, w: &mut dyn Writer,
@@ -1559,6 +1573,14 @@ fn encode_partition_bottomup<T: Pixel>(
 
     let mode_decision = rdo_mode_decision(fi, fs, cw, bsize, bo, spmvs);
 
+    if !mode_decision.pred_mode_luma.is_intra() {
+      // Fill the saved motion structure
+      save_block_motion(
+        fs, fi.w_in_b, fi.h_in_b, mode_decision.bsize, &mode_decision.bo,
+        mode_decision.ref_frames[0] - LAST_FRAME, mode_decision.mvs[0]
+      );
+    }
+
     rd_cost = mode_decision.rd_cost + cost;
 
     best_partition = PartitionType::PARTITION_NONE;
@@ -1680,6 +1702,14 @@ fn encode_partition_bottomup<T: Pixel>(
         for mode in rdo_output.part_modes.clone() {
           assert!(subsize == mode.bsize);
           let offset = mode.bo.clone();
+
+          if !mode.pred_mode_luma.is_intra() {
+            save_block_motion(
+              fs, fi.w_in_b, fi.h_in_b, mode.bsize, &mode.bo,
+              mode.ref_frames[0] - LAST_FRAME, mode.mvs[0]
+            );
+          }
+
           // FIXME: redundant block re-encode
           encode_block_with_modes(fi, fs, cw, w_pre_cdef, w_post_cdef,
                                   mode.bsize, &offset, &mode, rdo_type);
@@ -1850,6 +1880,12 @@ fn encode_partition_topdown<T: Pixel>(
           }
           mode_chroma = mode_luma;
         }
+
+        save_block_motion(
+          fs, fi.w_in_b, fi.h_in_b,
+          part_decision.bsize, &part_decision.bo,
+          part_decision.ref_frames[0] - LAST_FRAME, part_decision.mvs[0]
+        );
       }
 
       // FIXME: every final block that has gone through the RDO decision process is encoded twice
diff --git a/src/me.rs b/src/me.rs
index 9aa324b3..95044655 100644
--- a/src/me.rs
+++ b/src/me.rs
@@ -265,65 +265,70 @@ fn get_mv_range(
 pub fn get_subset_predictors<T: Pixel>(
   fi: &FrameInvariants<T>, bo: &BlockOffset, cmv: MotionVector,
   frame_mvs: &FrameMotionVectors, frame_ref_opt: &Option<Arc<ReferenceFrame<T>>>,
-  ref_slot: usize
+  ref_frame_id: usize
 ) -> (Vec<MotionVector>) {
   let mut predictors = Vec::new();
 
+  // Zero motion vector
+  predictors.push(MotionVector::default());
+
+  // Coarse motion estimation.
+  predictors.push(cmv.quantize_to_fullpel());
+
   // EPZS subset A and B predictors.
 
+  let mut median_preds = Vec::new();
   if bo.x > 0 {
     let left = frame_mvs[bo.y][bo.x - 1];
-    predictors.push(left);
+    median_preds.push(left);
+    if !left.is_zero() { predictors.push(left); }
   }
   if bo.y > 0 {
     let top = frame_mvs[bo.y - 1][bo.x];
-    predictors.push(top);
+    median_preds.push(top);
+    if !top.is_zero() { predictors.push(top); }
 
     if bo.x < fi.w_in_b - 1 {
       let top_right = frame_mvs[bo.y - 1][bo.x + 1];
-      predictors.push(top_right);
+      median_preds.push(top_right);
+      if !top_right.is_zero() { predictors.push(top_right); }
     }
   }
 
-  if !predictors.is_empty() {
+  if !median_preds.is_empty() {
     let mut median_mv = MotionVector::default();
-    for mv in predictors.iter() {
+    for mv in median_preds.iter() {
       median_mv = median_mv + *mv;
     }
-    median_mv = median_mv / (predictors.len() as i16);
-
-    predictors.push(median_mv.quantize_to_fullpel());
+    median_mv = median_mv / (median_preds.len() as i16);
+    let median_mv_quant = median_mv.quantize_to_fullpel();
+    if !median_mv_quant.is_zero() { predictors.push(median_mv_quant); }
   }
 
-  predictors.push(MotionVector::default());
-
-  // Coarse motion estimation.
-
-  predictors.push(cmv.quantize_to_fullpel());
-
   // EPZS subset C predictors.
 
   if let Some(ref frame_ref) = frame_ref_opt {
-    let prev_frame_mvs = &frame_ref.frame_mvs[ref_slot];
+    let prev_frame_mvs = &frame_ref.frame_mvs[ref_frame_id];
 
     if bo.x > 0 {
       let left = prev_frame_mvs[bo.y][bo.x - 1];
-      predictors.push(left);
+      if !left.is_zero() { predictors.push(left); }
     }
     if bo.y > 0 {
       let top = prev_frame_mvs[bo.y - 1][bo.x];
-      predictors.push(top);
+      if !top.is_zero() { predictors.push(top); }
     }
     if bo.x < fi.w_in_b - 1 {
       let right = prev_frame_mvs[bo.y][bo.x + 1];
-      predictors.push(right);
+      if !right.is_zero() { predictors.push(right); }
     }
     if bo.y < fi.h_in_b - 1 {
       let bottom = prev_frame_mvs[bo.y + 1][bo.x];
-      predictors.push(bottom);
+      if !bottom.is_zero() { predictors.push(bottom); }
     }
 
-    predictors.push(prev_frame_mvs[bo.y][bo.x]);
+    let previous = prev_frame_mvs[bo.y][bo.x];
+    if !previous.is_zero() { predictors.push(previous); }
   }
 
   predictors
@@ -333,7 +338,7 @@ pub trait MotionEstimation {
   fn full_pixel_me<T: Pixel>(
     fi: &FrameInvariants<T>, fs: &FrameState<T>, rec: &Arc<ReferenceFrame<T>>, po: &PlaneOffset,
     bo: &BlockOffset, lambda: u32,
-    ref_slot: usize, cmv: MotionVector, pmv: [MotionVector; 2],
+    cmv: MotionVector, pmv: [MotionVector; 2],
     mvx_min: isize, mvx_max: isize, mvy_min: isize, mvy_max: isize,
     blk_w: usize, blk_h: usize, best_mv: &mut MotionVector,
     lowest_cost: &mut u64, ref_frame: usize
@@ -351,7 +356,7 @@ pub trait MotionEstimation {
   fn motion_estimation<T: Pixel> (
     fi: &FrameInvariants<T>, fs: &FrameState<T>, bsize: BlockSize,
     bo: &BlockOffset, ref_frame: usize, cmv: MotionVector,
-    pmv: [MotionVector; 2], ref_slot: usize
+    pmv: [MotionVector; 2]
   ) -> MotionVector {
     match fi.rec_buffer.frames[fi.ref_frames[ref_frame - LAST_FRAME] as usize]
     {
@@ -373,7 +378,7 @@ pub trait MotionEstimation {
         let mut lowest_cost = std::u64::MAX;
         let mut best_mv = MotionVector::default();
 
-        Self::full_pixel_me(fi, fs, rec, &po, bo, lambda, ref_slot, cmv, pmv,
+        Self::full_pixel_me(fi, fs, rec, &po, bo, lambda, cmv, pmv,
                            mvx_min, mvx_max, mvy_min, mvy_max, blk_w, blk_h,
                            &mut best_mv, &mut lowest_cost, ref_frame);
 
@@ -397,15 +402,15 @@ pub struct FullSearch {}
 impl MotionEstimation for DiamondSearch {
   fn full_pixel_me<T: Pixel>(
     fi: &FrameInvariants<T>, fs: &FrameState<T>, rec: &Arc<ReferenceFrame<T>>,
-    po: &PlaneOffset, bo: &BlockOffset, lambda: u32, ref_slot: usize,
+    po: &PlaneOffset, bo: &BlockOffset, lambda: u32,
     cmv: MotionVector, pmv: [MotionVector; 2], mvx_min: isize, mvx_max: isize,
     mvy_min: isize, mvy_max: isize, blk_w: usize, blk_h: usize,
     best_mv: &mut MotionVector, lowest_cost: &mut u64, ref_frame: usize
   ) {
-    let frame_mvs = &fs.frame_mvs[ref_slot];
+    let frame_mvs = &fs.frame_mvs[ref_frame - LAST_FRAME];
     let frame_ref = &fi.rec_buffer.frames[fi.ref_frames[0] as usize];
     let predictors =
-      get_subset_predictors(fi, bo, cmv, frame_mvs, frame_ref, ref_slot);
+      get_subset_predictors(fi, bo, cmv, frame_mvs, frame_ref, ref_frame - LAST_FRAME);
 
     diamond_me_search(
       fi,
@@ -465,7 +470,7 @@ impl MotionEstimation for DiamondSearch {
 impl MotionEstimation for FullSearch {
   fn full_pixel_me<T: Pixel>(
     fi: &FrameInvariants<T>, fs: &FrameState<T>, rec: &Arc<ReferenceFrame<T>>,
-    po: &PlaneOffset, _bo: &BlockOffset, lambda: u32, _ref_slot: usize,
+    po: &PlaneOffset, _bo: &BlockOffset, lambda: u32,
     cmv: MotionVector, pmv: [MotionVector; 2], mvx_min: isize, mvx_max: isize,
     mvy_min: isize, mvy_max: isize, blk_w: usize, blk_h: usize,
     best_mv: &mut MotionVector, lowest_cost: &mut u64, _ref_frame: usize
diff --git a/src/partition.rs b/src/partition.rs
index fad70a7e..9b0e141a 100644
--- a/src/partition.rs
+++ b/src/partition.rs
@@ -761,6 +761,10 @@ impl MotionVector {
       col: (self.col / 8) * 8
     }
   }
+
+  pub fn is_zero(self) -> bool {
+    self.row == 0 && self.col == 0
+  }
 }
 
 pub const NEWMV_MODE_CONTEXTS: usize = 7;
diff --git a/src/rdo.rs b/src/rdo.rs
index 7ceecf81..69c33ddb 100644
--- a/src/rdo.rs
+++ b/src/rdo.rs
@@ -514,15 +514,8 @@ pub fn rdo_mode_decision<T: Pixel>(
       let ref_slot = ref_slot_set[i] as usize;
       let cmv = pmvs[ref_slot].unwrap();
 
-      let b_me = motion_estimation(fi, fs, bsize, bo, ref_frames[0], cmv, pmv, ref_slot);
 
-      // Fill the saved motion structure.
-      let frame_mvs = &mut fs.frame_mvs[ref_slot as usize];
-      for mi_y in (bo.y)..(bo.y + bsize.height_mi()) {
-        for mi_x in (bo.x)..(bo.x + bsize.width_mi()) {
-          frame_mvs[mi_y][mi_x] = b_me;
-        }
-      }
+      let b_me = motion_estimation(fi, fs, bsize, bo, ref_frames[0], cmv, pmv);
 
       mvs_from_me.push([
         b_me,
-- 
GitLab