VP9 Denoiser denoises after mode/bsize search

In vp8, statistics are collected about the different modes as they are searched.
This process is more complicated due to the variable block size. Fields were
added to the PICM_MODE_CONTEXT struct to hold this information for each point in
the search. The information is then taken from the appropriate part of the tree
during denoising.

Change-Id: I89261ab77ad637821287ae157dfdf694702b8e77

6 files changed, 81 insertions, 67 deletions

diff --git a/vp9/encoder/vp9_context_tree.h b/vp9/encoder/vp9_context_tree.h
index 872223b75..6d76914e9 100644
--- a/vp9/encoder/vp9_context_tree.h
+++ b/vp9/encoder/vp9_context_tree.h
@@ -41,6 +41,15 @@ typedef struct {
   int64_t tx_rd_diff[TX_MODES];
   int64_t best_filter_diff[SWITCHABLE_FILTER_CONTEXTS];
 
+#if CONFIG_DENOISING
+  unsigned int newmv_sse;
+  unsigned int zeromv_sse;
+  PREDICTION_MODE best_sse_inter_mode;
+  int_mv best_sse_mv;
+  MV_REFERENCE_FRAME best_reference_frame;
+  MV_REFERENCE_FRAME best_zeromv_reference_frame;
+#endif
+
   // motion vector cache for adaptive motion search control in partition
   // search loop
   MV pred_mv[MAX_REF_FRAMES];
diff --git a/vp9/encoder/vp9_denoiser.c b/vp9/encoder/vp9_denoiser.c
index dbf8cd740..ca84a1b39 100644
--- a/vp9/encoder/vp9_denoiser.c
+++ b/vp9/encoder/vp9_denoiser.c
@@ -13,6 +13,7 @@
 #include "vpx_scale/yv12config.h"
 #include "vpx/vpx_integer.h"
 #include "vp9/common/vp9_reconinter.h"
+#include "vp9/encoder/vp9_context_tree.h"
 #include "vp9/encoder/vp9_denoiser.h"
 
 /* The VP9 denoiser is a work-in-progress. It currently is only designed to work
@@ -183,9 +184,11 @@ static VP9_DENOISER_DECISION perform_motion_compensation(VP9_DENOISER *denoiser,
                                                          BLOCK_SIZE bs,
                                                          int increase_denoising,
                                                          int mi_row,
-                                                         int mi_col) {
+                                                         int mi_col,
+                                                         PICK_MODE_CONTEXT *ctx
+                                                         ) {
   int mv_col, mv_row;
-  int sse_diff = denoiser->zero_mv_sse - denoiser->best_sse;
+  int sse_diff = ctx->zeromv_sse - ctx->newmv_sse;
   MV_REFERENCE_FRAME frame;
   MACROBLOCKD *filter_mbd = &mb->e_mbd;
   MB_MODE_INFO *mbmi = &filter_mbd->mi[0]->mbmi;
@@ -204,29 +207,29 @@ static VP9_DENOISER_DECISION perform_motion_compensation(VP9_DENOISER *denoiser,
     saved_dst[i] = filter_mbd->plane[i].dst;
   }
 
-  mv_col = denoiser->best_sse_mv.as_mv.col;
-  mv_row = denoiser->best_sse_mv.as_mv.row;
+  mv_col = ctx->best_sse_mv.as_mv.col;
+  mv_row = ctx->best_sse_mv.as_mv.row;
 
-  frame = denoiser->best_reference_frame;
+  frame = ctx->best_reference_frame;
 
   // If the best reference frame uses inter-prediction and there is enough of a
   // difference in sum-squared-error, use it.
   if (frame != INTRA_FRAME &&
       sse_diff > sse_diff_thresh(bs, increase_denoising, mv_row, mv_col)) {
-    mbmi->ref_frame[0] = denoiser->best_reference_frame;
-    mbmi->mode = denoiser->best_sse_inter_mode;
-    mbmi->mv[0] = denoiser->best_sse_mv;
+    mbmi->ref_frame[0] = ctx->best_reference_frame;
+    mbmi->mode = ctx->best_sse_inter_mode;
+    mbmi->mv[0] = ctx->best_sse_mv;
   } else {
     // Otherwise, use the zero reference frame.
-    frame = denoiser->best_zeromv_reference_frame;
+    frame = ctx->best_zeromv_reference_frame;
 
-    mbmi->ref_frame[0] = denoiser->best_zeromv_reference_frame;
+    mbmi->ref_frame[0] = ctx->best_zeromv_reference_frame;
     mbmi->mode = ZEROMV;
     mbmi->mv[0].as_int = 0;
 
-    denoiser->best_sse_inter_mode = ZEROMV;
-    denoiser->best_sse_mv.as_int = 0;
-    denoiser->best_sse = denoiser->zero_mv_sse;
+    ctx->best_sse_inter_mode = ZEROMV;
+    ctx->best_sse_mv.as_int = 0;
+    ctx->newmv_sse = ctx->zeromv_sse;
   }
 
   // Set the pointers in the MACROBLOCKD to point to the buffers in the denoiser
@@ -278,10 +281,10 @@ static VP9_DENOISER_DECISION perform_motion_compensation(VP9_DENOISER *denoiser,
     filter_mbd->plane[i].dst = saved_dst[i];
   }
 
-  mv_row = denoiser->best_sse_mv.as_mv.row;
-  mv_col = denoiser->best_sse_mv.as_mv.col;
+  mv_row = ctx->best_sse_mv.as_mv.row;
+  mv_col = ctx->best_sse_mv.as_mv.col;
 
-  if (denoiser->best_sse > sse_thresh(bs, increase_denoising)) {
+  if (ctx->newmv_sse > sse_thresh(bs, increase_denoising)) {
     return COPY_BLOCK;
   }
   if (mv_row * mv_row + mv_col * mv_col >
@@ -292,7 +295,8 @@ static VP9_DENOISER_DECISION perform_motion_compensation(VP9_DENOISER *denoiser,
 }
 
 void vp9_denoiser_denoise(VP9_DENOISER *denoiser, MACROBLOCK *mb,
-                          int mi_row, int mi_col, BLOCK_SIZE bs) {
+                          int mi_row, int mi_col, BLOCK_SIZE bs,
+                          PICK_MODE_CONTEXT *ctx) {
   VP9_DENOISER_DECISION decision = FILTER_BLOCK;
   YV12_BUFFER_CONFIG avg = denoiser->running_avg_y[INTRA_FRAME];
   YV12_BUFFER_CONFIG mc_avg = denoiser->mc_running_avg_y;
@@ -303,7 +307,7 @@ void vp9_denoiser_denoise(VP9_DENOISER *denoiser, MACROBLOCK *mb,
 
   decision = perform_motion_compensation(denoiser, mb, bs,
                                          denoiser->increase_denoising,
-                                         mi_row, mi_col);
+                                         mi_row, mi_col, ctx);
 
   if (decision == FILTER_BLOCK) {
     decision = denoiser_filter(src.buf, src.stride,
@@ -362,24 +366,25 @@ void vp9_denoiser_update_frame_info(VP9_DENOISER *denoiser,
   }
 }
 
-void vp9_denoiser_reset_frame_stats(VP9_DENOISER *denoiser) {
-  denoiser->zero_mv_sse = UINT_MAX;
-  denoiser->best_sse = UINT_MAX;
+void vp9_denoiser_reset_frame_stats(PICK_MODE_CONTEXT *ctx) {
+  ctx->zeromv_sse = UINT_MAX;
+  ctx->newmv_sse = UINT_MAX;
 }
 
 void vp9_denoiser_update_frame_stats(VP9_DENOISER *denoiser, MB_MODE_INFO *mbmi,
-                                     unsigned int sse, PREDICTION_MODE mode) {
+                                     unsigned int sse, PREDICTION_MODE mode,
+                                     PICK_MODE_CONTEXT *ctx) {
   // TODO(tkopp): Use both MVs if possible
-  if (mbmi->mv[0].as_int == 0 && sse < denoiser->zero_mv_sse) {
-    denoiser->zero_mv_sse = sse;
-    denoiser->best_zeromv_reference_frame = mbmi->ref_frame[0];
+  if (mbmi->mv[0].as_int == 0 && sse < ctx->zeromv_sse) {
+    ctx->zeromv_sse = sse;
+    ctx->best_zeromv_reference_frame = mbmi->ref_frame[0];
   }
 
-  if (mbmi->mv[0].as_int != 0 && sse < denoiser->best_sse) {
-    denoiser->best_sse = sse;
-    denoiser->best_sse_inter_mode = mode;
-    denoiser->best_sse_mv = mbmi->mv[0];
-    denoiser->best_reference_frame = mbmi->ref_frame[0];
+  if (mode == NEWMV) {
+    ctx->newmv_sse = sse;
+    ctx->best_sse_inter_mode = mode;
+    ctx->best_sse_mv = mbmi->mv[0];
+    ctx->best_reference_frame = mbmi->ref_frame[0];
   }
 }
 
diff --git a/vp9/encoder/vp9_denoiser.h b/vp9/encoder/vp9_denoiser.h
index cbb6423c2..d93846ff9 100644
--- a/vp9/encoder/vp9_denoiser.h
+++ b/vp9/encoder/vp9_denoiser.h
@@ -26,14 +26,7 @@ typedef enum vp9_denoiser_decision {
 typedef struct vp9_denoiser {
   YV12_BUFFER_CONFIG running_avg_y[MAX_REF_FRAMES];
   YV12_BUFFER_CONFIG mc_running_avg_y;
-
-  unsigned int zero_mv_sse;
-  unsigned int best_sse;
   int increase_denoising;
-  PREDICTION_MODE best_sse_inter_mode;
-  int_mv best_sse_mv;
-  MV_REFERENCE_FRAME best_reference_frame;
-  MV_REFERENCE_FRAME best_zeromv_reference_frame;
 } VP9_DENOISER;
 
 void vp9_denoiser_update_frame_info(VP9_DENOISER *denoiser,
@@ -44,12 +37,14 @@ void vp9_denoiser_update_frame_info(VP9_DENOISER *denoiser,
                                     int refresh_last_frame);
 
 void vp9_denoiser_denoise(VP9_DENOISER *denoiser, MACROBLOCK *mb,
-                          int mi_row, int mi_col, BLOCK_SIZE bs);
+                          int mi_row, int mi_col, BLOCK_SIZE bs,
+                          PICK_MODE_CONTEXT *ctx);
 
-void vp9_denoiser_reset_frame_stats(VP9_DENOISER *denoiser);
+void vp9_denoiser_reset_frame_stats(PICK_MODE_CONTEXT *ctx);
 
 void vp9_denoiser_update_frame_stats(VP9_DENOISER *denoiser, MB_MODE_INFO *mbmi,
-                                     unsigned int sse, PREDICTION_MODE mode);
+                                     unsigned int sse, PREDICTION_MODE mode,
+                                     PICK_MODE_CONTEXT *ctx);
 
 int vp9_denoiser_alloc(VP9_DENOISER *denoiser, int width, int height,
                        int ssx, int ssy, int border);
diff --git a/vp9/encoder/vp9_encodeframe.c b/vp9/encoder/vp9_encodeframe.c
index 7ba7cae72..defc2361f 100644
--- a/vp9/encoder/vp9_encodeframe.c
+++ b/vp9/encoder/vp9_encodeframe.c
@@ -1333,6 +1333,13 @@ static void encode_b_rt(VP9_COMP *cpi, const TileInfo *const tile,
   set_offsets(cpi, tile, mi_row, mi_col, bsize);
   update_state_rt(cpi, ctx, mi_row, mi_col, bsize);
 
+#if CONFIG_DENOISING
+  if (cpi->oxcf.noise_sensitivity > 0 && output_enabled) {
+    vp9_denoiser_denoise(&cpi->denoiser, &cpi->mb, mi_row, mi_col,
+                         MAX(BLOCK_8X8, bsize), ctx);
+  }
+#endif
+
   encode_superblock(cpi, tp, output_enabled, mi_row, mi_col, bsize, ctx);
   update_stats(cpi);
 
@@ -2395,7 +2402,7 @@ static TX_MODE select_tx_mode(const VP9_COMP *cpi) {
 static void nonrd_pick_sb_modes(VP9_COMP *cpi, const TileInfo *const tile,
                                 int mi_row, int mi_col,
                                 int *rate, int64_t *dist,
-                                BLOCK_SIZE bsize) {
+                                BLOCK_SIZE bsize, PICK_MODE_CONTEXT *ctx) {
   VP9_COMMON *const cm = &cpi->common;
   MACROBLOCK *const x = &cpi->mb;
   MACROBLOCKD *const xd = &x->e_mbd;
@@ -2411,7 +2418,7 @@ static void nonrd_pick_sb_modes(VP9_COMP *cpi, const TileInfo *const tile,
   if (vp9_segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP))
     set_mode_info_seg_skip(x, cm->tx_mode, rate, dist, bsize);
   else
-    vp9_pick_inter_mode(cpi, x, tile, mi_row, mi_col, rate, dist, bsize);
+    vp9_pick_inter_mode(cpi, x, tile, mi_row, mi_col, rate, dist, bsize, ctx);
 
   duplicate_mode_info_in_sb(cm, xd, mi_row, mi_col, bsize);
 }
@@ -2528,7 +2535,7 @@ static void nonrd_pick_partition(VP9_COMP *cpi, const TileInfo *const tile,
   // PARTITION_NONE
   if (partition_none_allowed) {
     nonrd_pick_sb_modes(cpi, tile, mi_row, mi_col,
-                        &this_rate, &this_dist, bsize);
+                        &this_rate, &this_dist, bsize, ctx);
     ctx->mic.mbmi = xd->mi[0]->mbmi;
     ctx->skip_txfm = x->skip_txfm;
     ctx->skip = x->skip;
@@ -2611,7 +2618,8 @@ static void nonrd_pick_partition(VP9_COMP *cpi, const TileInfo *const tile,
       load_pred_mv(x, ctx);
 
     nonrd_pick_sb_modes(cpi, tile, mi_row, mi_col,
-                        &this_rate, &this_dist, subsize);
+                        &this_rate, &this_dist, subsize,
+                        &pc_tree->horizontal[0]);
 
     pc_tree->horizontal[0].mic.mbmi = xd->mi[0]->mbmi;
     pc_tree->horizontal[0].skip_txfm = x->skip_txfm;
@@ -2622,7 +2630,8 @@ static void nonrd_pick_partition(VP9_COMP *cpi, const TileInfo *const tile,
     if (sum_rd < best_rd && mi_row + ms < cm->mi_rows) {
       load_pred_mv(x, ctx);
       nonrd_pick_sb_modes(cpi, tile, mi_row + ms, mi_col,
-                          &this_rate, &this_dist, subsize);
+                          &this_rate, &this_dist, subsize,
+                          &pc_tree->horizontal[1]);
 
       pc_tree->horizontal[1].mic.mbmi = xd->mi[0]->mbmi;
       pc_tree->horizontal[1].skip_txfm = x->skip_txfm;
@@ -2654,7 +2663,8 @@ static void nonrd_pick_partition(VP9_COMP *cpi, const TileInfo *const tile,
       load_pred_mv(x, ctx);
 
     nonrd_pick_sb_modes(cpi, tile, mi_row, mi_col,
-                        &this_rate, &this_dist, subsize);
+                        &this_rate, &this_dist, subsize,
+                        &pc_tree->vertical[0]);
     pc_tree->vertical[0].mic.mbmi = xd->mi[0]->mbmi;
     pc_tree->vertical[0].skip_txfm = x->skip_txfm;
     pc_tree->vertical[0].skip = x->skip;
@@ -2662,7 +2672,8 @@ static void nonrd_pick_partition(VP9_COMP *cpi, const TileInfo *const tile,
     if (sum_rd < best_rd && mi_col + ms < cm->mi_cols) {
       load_pred_mv(x, ctx);
       nonrd_pick_sb_modes(cpi, tile, mi_row, mi_col + ms,
-                          &this_rate, &this_dist, subsize);
+                          &this_rate, &this_dist, subsize,
+                          &pc_tree->vertical[1]);
       pc_tree->vertical[1].mic.mbmi = xd->mi[0]->mbmi;
       pc_tree->vertical[1].skip_txfm = x->skip_txfm;
       pc_tree->vertical[1].skip = x->skip;
@@ -2752,19 +2763,21 @@ static void nonrd_use_partition(VP9_COMP *cpi,
 
   switch (partition) {
     case PARTITION_NONE:
-      nonrd_pick_sb_modes(cpi, tile, mi_row, mi_col, totrate, totdist, subsize);
+      nonrd_pick_sb_modes(cpi, tile, mi_row, mi_col, totrate, totdist,
+                          subsize, &pc_tree->none);
       pc_tree->none.mic.mbmi = xd->mi[0]->mbmi;
       pc_tree->none.skip_txfm = x->skip_txfm;
       pc_tree->none.skip = x->skip;
       break;
     case PARTITION_VERT:
-      nonrd_pick_sb_modes(cpi, tile, mi_row, mi_col, totrate, totdist, subsize);
+      nonrd_pick_sb_modes(cpi, tile, mi_row, mi_col, totrate, totdist,
+                          subsize, &pc_tree->vertical[0]);
       pc_tree->vertical[0].mic.mbmi = xd->mi[0]->mbmi;
       pc_tree->vertical[0].skip_txfm = x->skip_txfm;
       pc_tree->vertical[0].skip = x->skip;
       if (mi_col + hbs < cm->mi_cols) {
         nonrd_pick_sb_modes(cpi, tile, mi_row, mi_col + hbs,
-                            &rate, &dist, subsize);
+                            &rate, &dist, subsize, &pc_tree->vertical[1]);
         pc_tree->vertical[1].mic.mbmi = xd->mi[0]->mbmi;
         pc_tree->vertical[1].skip_txfm = x->skip_txfm;
         pc_tree->vertical[1].skip = x->skip;
@@ -2776,13 +2789,14 @@ static void nonrd_use_partition(VP9_COMP *cpi,
       }
       break;
     case PARTITION_HORZ:
-      nonrd_pick_sb_modes(cpi, tile, mi_row, mi_col, totrate, totdist, subsize);
+      nonrd_pick_sb_modes(cpi, tile, mi_row, mi_col, totrate, totdist,
+                          subsize, &pc_tree->horizontal[0]);
       pc_tree->horizontal[0].mic.mbmi = xd->mi[0]->mbmi;
       pc_tree->horizontal[0].skip_txfm = x->skip_txfm;
       pc_tree->horizontal[0].skip = x->skip;
       if (mi_row + hbs < cm->mi_rows) {
         nonrd_pick_sb_modes(cpi, tile, mi_row + hbs, mi_col,
-                            &rate, &dist, subsize);
+                            &rate, &dist, subsize, &pc_tree->horizontal[0]);
         pc_tree->horizontal[1].mic.mbmi = xd->mi[0]->mbmi;
         pc_tree->horizontal[1].skip_txfm = x->skip_txfm;
         pc_tree->horizontal[1].skip = x->skip;
diff --git a/vp9/encoder/vp9_pickmode.c b/vp9/encoder/vp9_pickmode.c
index 99f4897f4..80da5dab2 100644
--- a/vp9/encoder/vp9_pickmode.c
+++ b/vp9/encoder/vp9_pickmode.c
@@ -350,7 +350,8 @@ int64_t vp9_pick_inter_mode(VP9_COMP *cpi, MACROBLOCK *x,
                             int mi_row, int mi_col,
                             int *returnrate,
                             int64_t *returndistortion,
-                            BLOCK_SIZE bsize) {
+                            BLOCK_SIZE bsize,
+                            PICK_MODE_CONTEXT *ctx) {
   MACROBLOCKD *xd = &x->e_mbd;
   MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
   struct macroblock_plane *const p = &x->plane[0];
@@ -405,12 +406,6 @@ int64_t vp9_pick_inter_mode(VP9_COMP *cpi, MACROBLOCK *x,
   PRED_BUFFER *this_mode_pred = NULL;
   int i;
 
-#if CONFIG_DENOISING
-  if (cpi->oxcf.noise_sensitivity > 0) {
-    vp9_denoiser_reset_frame_stats(&cpi->denoiser);
-  }
-#endif
-
   if (cpi->sf.reuse_inter_pred_sby) {
     for (i = 0; i < 3; i++) {
       tmp[i].data = &pred_buf[pixels_in_block * i];
@@ -616,7 +611,8 @@ int64_t vp9_pick_inter_mode(VP9_COMP *cpi, MACROBLOCK *x,
 
 #if CONFIG_DENOISING
       if (cpi->oxcf.noise_sensitivity > 0) {
-        vp9_denoiser_update_frame_stats(&cpi->denoiser, mbmi, sse_y, this_mode);
+        vp9_denoiser_update_frame_stats(&cpi->denoiser, mbmi, sse_y,
+                                        this_mode, ctx);
       }
 #endif
 
@@ -733,11 +729,5 @@ int64_t vp9_pick_inter_mode(VP9_COMP *cpi, MACROBLOCK *x,
     }
   }
 
-#if CONFIG_DENOISING
-  if (cpi->oxcf.noise_sensitivity > 0) {
-    vp9_denoiser_denoise(&cpi->denoiser, x, mi_row, mi_col, bsize);
-  }
-#endif
-
   return INT64_MAX;
 }
diff --git a/vp9/encoder/vp9_pickmode.h b/vp9/encoder/vp9_pickmode.h
index 3d89974fc..49c6feb88 100644
--- a/vp9/encoder/vp9_pickmode.h
+++ b/vp9/encoder/vp9_pickmode.h
@@ -28,7 +28,8 @@ int64_t vp9_pick_inter_mode(VP9_COMP *cpi, MACROBLOCK *x,
                             int mi_row, int mi_col,
                             int *returnrate,
                             int64_t *returndistortion,
-                            BLOCK_SIZE bsize);
+                            BLOCK_SIZE bsize,
+                            PICK_MODE_CONTEXT *ctx);
 
 #ifdef __cplusplus
 }  // extern "C"