From 28a38a5ab2e8442763f724e1e3d4f7062c2969a3 Mon Sep 17 00:00:00 2001
From: fbossen <frank@bossentech.com>
Date: Wed, 12 Sep 2018 15:52:25 -0230
Subject: [PATCH] Add test for NEARMV mode (#523)

* Correctly handle NEAR mode when less than 2 entried in MV stack

If there are fewer than 2 entries in the motion vector stack, then
the NEAR motion vector is considered to be the same vector as the
global motion vector (currently 0,0), as per 7.10.2.12
"Extra search process" in the spec.

* Add support for 2nd and 3rd NEAR motion vector

* Adaptively select best MV mode
---
 src/encoder.rs   | 41 +++++++++++++++++++++++++++++++++++++----
 src/partition.rs |  4 +++-
 src/predict.rs   |  3 +++
 src/rdo.rs       | 10 ++++++++++
 4 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/src/encoder.rs b/src/encoder.rs
index 68986ce7..8bf74a38 100644
--- a/src/encoder.rs
+++ b/src/encoder.rs
@@ -1380,6 +1380,28 @@ pub fn encode_block_b(seq: &Sequence, fi: &FrameInvariants, fs: &mut FrameState,
                 MvSubpelPrecision::MV_SUBPEL_LOW_PRECISION
               };
               cw.write_mv(w, &mv, &ref_mv, mv_precision);
+            } else if luma_mode >= PredictionMode::NEAR0MV && luma_mode <= PredictionMode::NEAR2MV {
+              let ref_mv_idx = luma_mode as usize - PredictionMode::NEAR0MV as usize + 1;
+              let num_mv_found = mv_stack.len();
+              if luma_mode != PredictionMode::NEAR0MV { assert!(num_mv_found > ref_mv_idx); }
+
+              for idx in 1..3 {
+                if num_mv_found > idx + 1 {
+                  let drl_mode = ref_mv_idx > idx;
+                  let ctx: usize = (mv_stack[idx].weight < REF_CAT_LEVEL) as usize
+                    + (mv_stack[idx + 1].weight < REF_CAT_LEVEL) as usize;
+
+                  cw.write_drl_mode(w, drl_mode, ctx);
+                  if !drl_mode { break; }
+                }
+              }
+              if mv_stack.len() > 1 {
+                assert!(mv_stack[ref_mv_idx].this_mv.row == mv.row);
+                assert!(mv_stack[ref_mv_idx].this_mv.col == mv.col);
+              } else {
+                assert!(0 == mv.row);
+                assert!(0 == mv.col);
+              }
             } else if luma_mode == PredictionMode::NEARESTMV {
               if mv_stack.len() > 0 {
                 assert!(mv_stack[0].this_mv.row == mv.row);
@@ -1827,11 +1849,22 @@ fn encode_partition_topdown(seq: &Sequence, fi: &FrameInvariants, fs: &mut Frame
             let mut mv_stack = Vec::new();
             let mode_context = cw.find_mvrefs(bo, ref_frame, &mut mv_stack, bsize, false);
 
-            if mode_luma == PredictionMode::NEARESTMV &&
-                (mv_stack.len() > 0 && (mv_stack[0].this_mv.row != mv.row || mv_stack[0].this_mv.col != mv.col) ||
-                 mv_stack.len() == 0 && (0 != mv.row || 0 != mv.col)) {
+            if !mode_luma.is_intra() && mode_luma != PredictionMode::GLOBALMV {
               mode_luma = PredictionMode::NEWMV;
-              mode_chroma = PredictionMode::NEWMV;
+              for (c, m) in mv_stack.iter().take(4)
+                .zip([PredictionMode::NEARESTMV, PredictionMode::NEAR0MV,
+                      PredictionMode::NEAR1MV, PredictionMode::NEAR2MV].iter()) {
+                if c.this_mv.row == mv.row && c.this_mv.col == mv.col {
+                  mode_luma = *m;
+                }
+              }
+              if mode_luma == PredictionMode::NEWMV && mv.row == 0 && mv.col == 0 {
+                mode_luma =
+                  if mv_stack.len() == 0 { PredictionMode::NEARESTMV }
+                  else if mv_stack.len() == 1 { PredictionMode::NEAR0MV }
+                  else { PredictionMode::GLOBALMV };
+              }
+              mode_chroma = mode_luma;
             }
 
             // FIXME: every final block that has gone through the RDO decision process is encoded twice
diff --git a/src/partition.rs b/src/partition.rs
index 0ebd1995..bb9686d8 100755
--- a/src/partition.rs
+++ b/src/partition.rs
@@ -344,7 +344,9 @@ pub enum PredictionMode {
   PAETH_PRED,
   UV_CFL_PRED,
   NEARESTMV,
-  NEARMV,
+  NEAR0MV,
+  NEAR1MV,
+  NEAR2MV,
   GLOBALMV,
   NEWMV,
   // Compound ref compound modes
diff --git a/src/predict.rs b/src/predict.rs
index 22158208..7b185e2f 100755
--- a/src/predict.rs
+++ b/src/predict.rs
@@ -44,6 +44,9 @@ pub static RAV1E_INTRA_MODES_MINIMAL: &'static [PredictionMode] = &[
 pub static RAV1E_INTER_MODES: &'static [PredictionMode] = &[
   PredictionMode::GLOBALMV,
   PredictionMode::NEARESTMV,
+  PredictionMode::NEAR0MV,
+  PredictionMode::NEAR1MV,
+  PredictionMode::NEAR2MV,
   PredictionMode::NEWMV
 ];
 
diff --git a/src/rdo.rs b/src/rdo.rs
index d890a7d9..27fef29b 100755
--- a/src/rdo.rs
+++ b/src/rdo.rs
@@ -294,6 +294,9 @@ pub fn rdo_mode_decision(
   for &luma_mode in &mode_set {
     assert!(fi.frame_type == FrameType::INTER || luma_mode.is_intra());
 
+    if luma_mode == PredictionMode::NEAR1MV && mv_stack.len() < 3 { continue; }
+    if luma_mode == PredictionMode::NEAR2MV && mv_stack.len() < 4 { continue; }
+
     let mut mode_set_chroma = vec![luma_mode];
 
     if is_chroma_block
@@ -316,6 +319,13 @@ pub fn rdo_mode_decision(
       } else {
         MotionVector { row: 0, col: 0 }
       },
+      PredictionMode::NEAR0MV => if mv_stack.len() > 1 {
+        mv_stack[1].this_mv
+      } else {
+        MotionVector { row: 0, col: 0 }
+      },
+      PredictionMode::NEAR1MV | PredictionMode::NEAR2MV =>
+          mv_stack[luma_mode as usize - PredictionMode::NEAR0MV as usize + 1].this_mv,
       _ => MotionVector { row: 0, col: 0 }
     };
 
-- 
GitLab