Score flat blocks and always return some.

Adding score of flat blocks so that we can return some flat
blocks in the extremely noisy cases. This is a no-op for easier cases.

Change-Id: I99387a642bc84ceb468e5360d3e023735f693843
diff --git a/aom_dsp/noise_model.c b/aom_dsp/noise_model.c
index 56f39ef..84ac997 100644
--- a/aom_dsp/noise_model.c
+++ b/aom_dsp/noise_model.c
@@ -483,14 +483,34 @@
   }
 }
 
+typedef struct {
+  int index;
+  float score;
+} index_and_score_t;
+
+static int compare_scores(const void *a, const void *b) {
+  const float diff =
+      ((index_and_score_t *)a)->score - ((index_and_score_t *)b)->score;
+  if (diff < 0)
+    return -1;
+  else if (diff > 0)
+    return 1;
+  return 0;
+}
+
 int aom_flat_block_finder_run(const aom_flat_block_finder_t *block_finder,
                               const uint8_t *const data, int w, int h,
                               int stride, uint8_t *flat_blocks) {
+  // The gradient-based features used in this code are based on:
+  //  A. Kokaram, D. Kelly, H. Denman and A. Crawford, "Measuring noise
+  //  correlation for improved video denoising," 2012 19th, ICIP.
+  // The thresholds are more lenient to allow for correct grain modeling
+  // if extreme cases.
   const int block_size = block_finder->block_size;
   const int n = block_size * block_size;
-  const double kTraceThreshold = 0.1 / (32 * 32);
-  const double kRatioThreshold = 1.2;
-  const double kNormThreshold = 0.05 / (32 * 32);
+  const double kTraceThreshold = 0.15 / (32 * 32);
+  const double kRatioThreshold = 1.25;
+  const double kNormThreshold = 0.08 / (32 * 32);
   const double kVarThreshold = 0.005 / (double)n;
   const int num_blocks_w = (w + block_size - 1) / block_size;
   const int num_blocks_h = (h + block_size - 1) / block_size;
@@ -498,13 +518,19 @@
   int bx = 0, by = 0;
   double *plane = (double *)aom_malloc(n * sizeof(*plane));
   double *block = (double *)aom_malloc(n * sizeof(*block));
-  if (plane == NULL || block == NULL) {
+  index_and_score_t *scores = (index_and_score_t *)aom_malloc(
+      num_blocks_w * num_blocks_h * sizeof(*scores));
+  if (plane == NULL || block == NULL || scores == NULL) {
     fprintf(stderr, "Failed to allocate memory for block of size %d\n", n);
     aom_free(plane);
     aom_free(block);
+    aom_free(scores);
     return -1;
   }
 
+#ifdef NOISE_MODEL_LOG_SCORE
+  fprintf(stderr, "score = [");
+#endif
   for (by = 0; by < num_blocks_h; ++by) {
     for (bx = 0; bx < num_blocks_w; ++bx) {
       // Compute gradient covariance matrix.
@@ -535,28 +561,65 @@
       mean /= (block_size - 2) * (block_size - 2);
 
       // Normalize gradients by block_size.
-      Gxx /= (block_size - 2) * (block_size - 2);
-      Gxy /= (block_size - 2) * (block_size - 2);
-      Gyy /= (block_size - 2) * (block_size - 2);
-      var = var / (block_size - 2) * (block_size - 2) - mean * mean;
+      Gxx /= ((block_size - 2) * (block_size - 2));
+      Gxy /= ((block_size - 2) * (block_size - 2));
+      Gyy /= ((block_size - 2) * (block_size - 2));
+      var = var / ((block_size - 2) * (block_size - 2)) - mean * mean;
 
       {
         const double trace = Gxx + Gyy;
         const double det = Gxx * Gyy - Gxy * Gxy;
         const double e1 = (trace + sqrt(trace * trace - 4 * det)) / 2.;
         const double e2 = (trace - sqrt(trace * trace - 4 * det)) / 2.;
-        const double norm = sqrt(Gxx * Gxx + Gxy * Gxy * 2 + Gyy * Gyy);
+        const double norm = e1;  // Spectral norm
+        const double ratio = (e1 / AOMMAX(e2, 1e-6));
         const int is_flat = (trace < kTraceThreshold) &&
-                            (e1 / AOMMAX(e2, 1e-8) < kRatioThreshold) &&
-                            norm < kNormThreshold && var > kVarThreshold;
+                            (ratio < kRatioThreshold) &&
+                            (norm < kNormThreshold) && (var > kVarThreshold);
+        // The following weights are used to combine the above features to give
+        // a sigmoid score for flatness. If the input was normalized to [0,100]
+        // the magnitude of these values would be close to 1 (e.g., weights
+        // corresponding to variance would be a factor of 10000x smaller).
+        // The weights are given in the following order:
+        //    [{var}, {ratio}, {trace}, {norm}, offset]
+        // with one of the most discriminative being simply the variance.
+        const double weights[5] = { -6682, -0.2056, 13087, -12434, 2.5694 };
+        const float score =
+            (float)(1.0 / (1 + exp(-(weights[0] * var + weights[1] * ratio +
+                                     weights[2] * trace + weights[3] * norm +
+                                     weights[4]))));
         flat_blocks[by * num_blocks_w + bx] = is_flat ? 255 : 0;
+        scores[by * num_blocks_w + bx].score = var > kVarThreshold ? score : 0;
+        scores[by * num_blocks_w + bx].index = by * num_blocks_w + bx;
+#ifdef NOISE_MODEL_LOG_SCORE
+        fprintf(stderr, "%g %g %g %g %g %d ", score, var, ratio, trace, norm,
+                is_flat);
+#endif
         num_flat += is_flat;
       }
     }
+#ifdef NOISE_MODEL_LOG_SCORE
+    fprintf(stderr, "\n");
+#endif
   }
-
+#ifdef NOISE_MODEL_LOG_SCORE
+  fprintf(stderr, "];\n");
+#endif
+  // Find the top-scored blocks (most likely to be flat) and set the flat blocks
+  // be the union of the thresholded results and the top 10th percentile of the
+  // scored results.
+  qsort(scores, num_blocks_w * num_blocks_h, sizeof(*scores), &compare_scores);
+  const int top_nth_percentile = num_blocks_w * num_blocks_h * 90 / 100;
+  const float score_threshold = scores[top_nth_percentile].score;
+  for (int i = 0; i < num_blocks_w * num_blocks_h; ++i) {
+    if (scores[i].score >= score_threshold) {
+      num_flat += flat_blocks[scores[i].index] == 0;
+      flat_blocks[scores[i].index] |= 1;
+    }
+  }
   aom_free(block);
   aom_free(plane);
+  aom_free(scores);
   return num_flat;
 }
 
diff --git a/aom_dsp/noise_model.h b/aom_dsp/noise_model.h
index 8864e8a..bc1b20e 100644
--- a/aom_dsp/noise_model.h
+++ b/aom_dsp/noise_model.h
@@ -145,6 +145,13 @@
     int w, int h, int stride, int offsx, int offsy, double *plane,
     double *block);
 
+/*!\brief Runs the flat block finder on the input data.
+ *
+ * Find flat blocks in the input image data. Returns a map of
+ * flat_blocks, where the value of flat_blocks map will be non-zero
+ * when a block is determined to be flat. A higher value indicates a bigger
+ * confidence in the decision.
+ */
 int aom_flat_block_finder_run(const aom_flat_block_finder_t *block_finder,
                               const uint8_t *const data, int w, int h,
                               int stride, uint8_t *flat_blocks);
diff --git a/examples/noise_model.c b/examples/noise_model.c
index 442016a..3fca385 100644
--- a/examples/noise_model.c
+++ b/examples/noise_model.c
@@ -58,8 +58,10 @@
     ARG_DEF("n", "output-grain-table", 1, "Output noise file");
 static const arg_def_t input_denoised_arg =
     ARG_DEF("d", "input-denoised", 1, "Input denoised filename (YUV) only");
+static const arg_def_t flat_block_finder_arg =
+    ARG_DEF("b", "flat-block-finder", 1, "Run the flat block finder");
 static const arg_def_t block_size_arg =
-    ARG_DEF("b", "block_size", 1, "Block size");
+    ARG_DEF("b", "block-size", 1, "Block size");
 static const arg_def_t bit_depth_arg =
     ARG_DEF(NULL, "bit-depth", 1, "Bit depth of input");
 static const arg_def_t use_i420 =
@@ -117,6 +119,8 @@
       noise_args->block_size = atoi(arg.val);
     } else if (arg_match(&arg, &bit_depth_arg, argv)) {
       noise_args->bit_depth = atoi(arg.val);
+    } else if (arg_match(&arg, &flat_block_finder_arg, argv)) {
+      noise_args->run_flat_block_finder = atoi(arg.val);
     } else if (arg_match(&arg, &fps_arg, argv)) {
       noise_args->fps = arg_parse_rational(&arg);
     } else if (arg_match(&arg, &use_i420, argv)) {
@@ -140,7 +144,7 @@
 
 int main(int argc, char *argv[]) {
   noise_model_args_t args = { 0,  0, { 1, 25 }, 0, 0, 0, AOM_IMG_FMT_I420,
-                              32, 8, 0,         0, 1 };
+                              32, 8, 1,         0, 1 };
   aom_image_t raw, denoised;
   FILE *infile = NULL;
   AvxVideoInfo info;
diff --git a/test/noise_model_test.cc b/test/noise_model_test.cc
index 72aa840..d9869ac 100644
--- a/test/noise_model_test.cc
+++ b/test/noise_model_test.cc
@@ -286,7 +286,7 @@
 template <typename T>
 class FlatBlockEstimatorTest : public ::testing::Test, public T {
  public:
-  virtual void SetUp() { random_.Reset(1071); }
+  virtual void SetUp() { random_.Reset(171); }
   typedef std::vector<typename T::data_type_t> VecType;
   VecType data_;
   libaom_test::ACMRandom random_;
@@ -409,12 +409,34 @@
   EXPECT_EQ(0, flat_blocks[1]);
 
   // Next 4 blocks are flat.
-  EXPECT_NE(0, flat_blocks[2]);
-  EXPECT_NE(0, flat_blocks[3]);
-  EXPECT_NE(0, flat_blocks[4]);
-  EXPECT_NE(0, flat_blocks[5]);
+  EXPECT_EQ(255, flat_blocks[2]);
+  EXPECT_EQ(255, flat_blocks[3]);
+  EXPECT_EQ(255, flat_blocks[4]);
+  EXPECT_EQ(255, flat_blocks[5]);
 
-  // Last 2 are not.
+  // Last 2 are not flat by threshold
+  EXPECT_EQ(0, flat_blocks[6]);
+  EXPECT_EQ(0, flat_blocks[7]);
+
+  // Add the noise from non-flat block 1 to every block.
+  for (int y = 0; y < kBlockSize; ++y) {
+    for (int x = 0; x < kBlockSize * num_blocks_w; ++x) {
+      this->data_[y * stride + x] +=
+          (this->data_[y * stride + x % kBlockSize + kBlockSize] -
+           (128 << shift));
+    }
+  }
+  // Now the scored selection will pick the one that is most likely flat (block
+  // 0)
+  EXPECT_EQ(1, aom_flat_block_finder_run(&flat_block_finder,
+                                         (uint8_t *)&this->data_[0], w, h,
+                                         stride, &flat_blocks[0]));
+  EXPECT_EQ(1, flat_blocks[0]);
+  EXPECT_EQ(0, flat_blocks[1]);
+  EXPECT_EQ(0, flat_blocks[2]);
+  EXPECT_EQ(0, flat_blocks[3]);
+  EXPECT_EQ(0, flat_blocks[4]);
+  EXPECT_EQ(0, flat_blocks[5]);
   EXPECT_EQ(0, flat_blocks[6]);
   EXPECT_EQ(0, flat_blocks[7]);