Fix subsampled AR-filter coefficient estimation

Also, if chroma estimation fails, default to returning a null filter instead
of nonesense.

Change-Id: Ica3df3264c13929ff87ed0f0ec211c95d9cfc0e4
diff --git a/aom_dsp/noise_model.c b/aom_dsp/noise_model.c
index beed8b3..702ff4d 100644
--- a/aom_dsp/noise_model.c
+++ b/aom_dsp/noise_model.c
@@ -131,6 +131,17 @@
   return aom_noise_strength_solver_init(&state->strength_solver, kNumBins);
 }
 
+static void set_chroma_coefficient_fallback_soln(aom_equation_system_t *eqns) {
+  const double kTolerance = 1e-6;
+  const int last = eqns->n - 1;
+  // Set all of the AR coefficients to zero, but try to solve for correlation
+  // with the luma channel
+  memset(eqns->x, 0, sizeof(*eqns->x) * eqns->n);
+  if (fabs(eqns->A[last * eqns->n + last]) > kTolerance) {
+    eqns->x[last] = eqns->b[last] / eqns->A[last * eqns->n + last];
+  }
+}
+
 int aom_noise_strength_lut_init(aom_noise_strength_lut_t *lut, int num_points) {
   if (!lut) return 0;
   lut->points = (double(*)[2])aom_malloc(num_points * sizeof(*lut->points));
@@ -228,6 +239,8 @@
   solver->num_bins = num_bins;
   solver->min_intensity = 0;
   solver->max_intensity = 255;
+  solver->total = 0;
+  solver->num_equations = 0;
   return equation_system_init(&solver->eqns, num_bins);
 }
 
@@ -574,16 +587,15 @@
   double *buffer = (double *)aom_malloc(sizeof(*buffer) * (num_coords + 1));
   const int n = noise_model->latest_state[c].eqns.n;
   int bx, by;
-  (void)sub_log2;
 
   if (!buffer) {
     fprintf(stderr, "Unable to allocate buffer of size %d\n", num_coords + 1);
     return 0;
   }
   for (by = 0; by < num_blocks_h; ++by) {
-    const int y_o = by * block_size;
+    const int y_o = by * (block_size >> sub_log2[1]);
     for (bx = 0; bx < num_blocks_w; ++bx) {
-      const int x_o = bx * block_size;
+      const int x_o = bx * (block_size >> sub_log2[0]);
       int x_start = 0, y_start = 0, x_end = 0, y_end = 0;
       int x, y, i, j;
       if (!flat_blocks[by * num_blocks_w + bx]) {
@@ -591,16 +603,13 @@
       }
       y_start = (by > 0 && flat_blocks[(by - 1) * num_blocks_w + bx]) ? 0 : lag;
       x_start = (bx > 0 && flat_blocks[by * num_blocks_w + bx - 1]) ? 0 : lag;
-      y_end = AOMMIN(
-          h - by * block_size,
-          (by + 1 < num_blocks_h && flat_blocks[(by + 1) * num_blocks_w + bx])
-              ? block_size
-              : block_size - lag);
+      y_end = AOMMIN((h >> sub_log2[1]) - by * (block_size >> sub_log2[1]),
+                     block_size >> sub_log2[1]);
       x_end = AOMMIN(
-          w - bx * block_size - lag,
+          (w >> sub_log2[0]) - bx * (block_size >> sub_log2[0]) - lag,
           (bx + 1 < num_blocks_w && flat_blocks[by * num_blocks_w + bx + 1])
-              ? block_size
-              : block_size - lag);
+              ? (block_size >> sub_log2[0])
+              : ((block_size >> sub_log2[0]) - lag));
       for (y = y_start; y < y_end; ++y) {
         for (x = x_start; x < x_end; ++x) {
           double val = 0;
@@ -609,7 +618,7 @@
             const int dy_i = noise_model->coords[i][1];
             const int x_i = x_o + x + dx_i;
             const int y_i = y_o + y + dy_i;
-            assert(x_i < w && y_i < h);
+            assert(x_i < (w >> sub_log2[0]) && y_i < (h >> sub_log2[1]));
             buffer[i] = ((double)(data[y_i * stride + x_i]) -
                          (double)(denoised[y_i * stride + x_i]));
           }
@@ -617,11 +626,20 @@
                 ((double)denoised[(y_o + y) * stride + (x_o + x)]);
 
           // For the color channels we must also consider the correlation with
-          // the luma channel.
+          // the luma channel
           if (alt_data && alt_denoised) {
-            buffer[num_coords] =
-                ((double)alt_data[(y_o + y) * alt_stride + (x_o + x)]) -
-                ((double)alt_denoised[(y_o + y) * alt_stride + (x_o + x)]);
+            double avg_data = 0, avg_denoised = 0;
+            int num_samples = 0;
+            for (int dy_i = 0; dy_i < (1 << sub_log2[1]); dy_i++) {
+              const int y_up = ((y_o + y) << sub_log2[1]) + dy_i;
+              for (int dx_i = 0; dx_i < (1 << sub_log2[0]); dx_i++) {
+                const int x_up = ((x_o + x) << sub_log2[0]) + dx_i;
+                avg_data += alt_data[y_up * alt_stride + x_up];
+                avg_denoised += alt_denoised[y_up * alt_stride + x_up];
+                num_samples++;
+              }
+            }
+            buffer[num_coords] = (avg_data - avg_denoised) / num_samples;
           }
 
           for (i = 0; i < n; ++i) {
@@ -640,71 +658,79 @@
   return 1;
 }
 
-void add_noise_std_observations(aom_noise_model_t *noise_model, int c,
-                                const double *coeffs, const uint8_t *const data,
-                                const uint8_t *const denoised, int w, int h,
-                                int stride, const uint8_t *const alt_data,
-                                const uint8_t *const alt_denoised,
-                                int alt_stride,
-                                const uint8_t *const flat_blocks,
-                                int block_size, int num_blocks_w,
-                                int num_blocks_h) {
+static void add_noise_std_observations(
+    aom_noise_model_t *noise_model, int c, const double *coeffs,
+    const uint8_t *const data, const uint8_t *const denoised, int w, int h,
+    int stride, int sub_log2[2], const uint8_t *const alt_data,
+    const uint8_t *const alt_denoised, int alt_stride,
+    const uint8_t *const flat_blocks, int block_size, int num_blocks_w,
+    int num_blocks_h) {
   const int lag = noise_model->params.lag;
   const int num_coords = noise_model->n;
   aom_noise_strength_solver_t *noise_strength_solver =
       &noise_model->latest_state[c].strength_solver;
-  int bx = 0, by = 0;
 
-  for (by = 0; by < num_blocks_h; ++by) {
-    const int y_o = by * block_size;
-    for (bx = 0; bx < num_blocks_w; ++bx) {
-      const int x_o = bx * block_size;
+  for (int by = 0; by < num_blocks_h; ++by) {
+    const int y_o = by * (block_size >> sub_log2[1]);
+    for (int bx = 0; bx < num_blocks_w; ++bx) {
+      const int x_o = bx * (block_size >> sub_log2[0]);
       if (!flat_blocks[by * num_blocks_w + bx]) {
         continue;
       }
-      const double block_mean =
-          get_block_mean(alt_data ? alt_data : data, w, h,
-                         alt_data ? alt_stride : stride, x_o, y_o, block_size);
+      const double block_mean = get_block_mean(
+          alt_data ? alt_data : data, w, h, alt_data ? alt_stride : stride,
+          x_o << sub_log2[0], y_o << sub_log2[1], block_size);
       double noise_var = 0;
       int num_samples_in_block = 0;
-      int y_start =
+      const int y_start =
           (by > 0 && flat_blocks[(by - 1) * num_blocks_w + bx]) ? 0 : lag;
-      int x_start =
+      const int x_start =
           (bx > 0 && flat_blocks[by * num_blocks_w + bx - 1]) ? 0 : lag;
-      int y_end =
-          (by + 1 < num_blocks_h && flat_blocks[(by + 1) * num_blocks_w + bx])
-              ? block_size
-              : block_size - lag;
-      int x_end =
+      const int y_end =
+          AOMMIN((h >> sub_log2[1]) - by * (block_size >> sub_log2[1]),
+                 block_size >> sub_log2[1]);
+      const int x_end = AOMMIN(
+          (w >> sub_log2[0]) - bx * (block_size >> sub_log2[0]) - lag,
           (bx + 1 < num_blocks_w && flat_blocks[by * num_blocks_w + bx + 1])
-              ? block_size
-              : block_size - lag;
-      for (int y = y_start; y < y_end; ++y) {
-        for (int x = x_start; x < x_end; ++x) {
+              ? (block_size >> sub_log2[0])
+              : ((block_size >> sub_log2[0]) - lag));
+      for (int y = y_start; y < y_end; y++) {
+        for (int x = x_start; x < x_end; x++) {
           const double actual =
-              ((double)(data[(y_o + y) * stride + (x_o + x)]) -
-               (double)(denoised[(y_o + y) * stride + (x_o + x)]));
+              (double)data[(y_o + y) * stride + (x_o + x)] -
+              (double)denoised[(y_o + y) * stride + (x_o + x)];
           double sum = 0;
           for (int i = 0; i < num_coords; ++i) {
             const int dx_i = noise_model->coords[i][0];
             const int dy_i = noise_model->coords[i][1];
-            const int x_i = x_o + x + dx_i;
-            const int y_i = y_o + y + dy_i;
+            const int x_i = (x_o + x + dx_i);
+            const int y_i = (y_o + y + dy_i);
             sum += coeffs[i] * ((double)(data[y_i * stride + x_i]) -
                                 (double)(denoised[y_i * stride + x_i]));
           }
           if (alt_data && alt_denoised) {
-            sum += coeffs[num_coords] *
-                   ((double)(alt_data[(y_o + y) * stride + (x_o + x)]) -
-                    (double)(alt_denoised[(y_o + y) * stride + (x_o + x)]));
+            double avg_data = 0, avg_denoised = 0;
+            int n = 0;
+            for (int dy_i = 0; dy_i < (1 << sub_log2[1]); dy_i++) {
+              const int y_up = ((y_o + y) << sub_log2[1]) + dy_i;
+              for (int dx_i = 0; dx_i < (1 << sub_log2[0]); dx_i++) {
+                const int x_up = ((x_o + x) << sub_log2[0]) + dx_i;
+                avg_data += alt_data[y_up * alt_stride + x_up];
+                avg_denoised += alt_denoised[y_up * alt_stride + x_up];
+                n++;
+              }
+            }
+            sum += coeffs[num_coords] * (avg_data - avg_denoised) / n;
           }
           noise_var += (sum - actual) * (sum - actual);
           num_samples_in_block++;
         }
       }
-      const double noise_std = sqrt(noise_var / num_samples_in_block);
-      aom_noise_strength_solver_add_measurement(noise_strength_solver,
-                                                block_mean, noise_std);
+      if (num_samples_in_block > block_size) {
+        const double noise_std = sqrt(noise_var / num_samples_in_block);
+        aom_noise_strength_solver_add_measurement(noise_strength_solver,
+                                                  block_mean, noise_std);
+      }
     }
   }
 }
@@ -712,7 +738,7 @@
 aom_noise_status_t aom_noise_model_update(
     aom_noise_model_t *const noise_model, const uint8_t *const data[3],
     const uint8_t *const denoised[3], int w, int h, int stride[3],
-    int chroma_sub[2], const uint8_t *const flat_blocks, int block_size) {
+    int chroma_sub_log2[2], const uint8_t *const flat_blocks, int block_size) {
   const int num_blocks_w = (w + block_size - 1) / block_size;
   const int num_blocks_h = (h + block_size - 1) / block_size;
   int y_model_different = 0;
@@ -748,11 +774,11 @@
   }
 
   for (channel = 0; channel < 3; ++channel) {
+    int no_subsampling[2] = { 0, 0 };
     const uint8_t *alt_data = channel > 0 ? data[0] : 0;
     const uint8_t *alt_denoised = channel > 0 ? denoised[0] : 0;
-    int *sub = channel > 0 ? chroma_sub : 0;
+    int *sub = channel > 0 ? chroma_sub_log2 : no_subsampling;
     if (!data[channel] || !denoised[channel]) break;
-
     if (!add_block_observations(noise_model, channel, data[channel],
                                 denoised[channel], w, h, stride[channel], sub,
                                 alt_data, alt_denoised, stride[0], flat_blocks,
@@ -762,13 +788,19 @@
     }
 
     if (!equation_system_solve(&noise_model->latest_state[channel].eqns)) {
-      fprintf(stderr, "Solving latest noise equation system failed!\n");
-      return AOM_NOISE_STATUS_INTERNAL_ERROR;
+      if (channel > 0) {
+        set_chroma_coefficient_fallback_soln(
+            &noise_model->latest_state[channel].eqns);
+      } else {
+        fprintf(stderr, "Solving latest noise equation system failed %d!\n",
+                channel);
+        return AOM_NOISE_STATUS_INTERNAL_ERROR;
+      }
     }
 
     add_noise_std_observations(
         noise_model, channel, noise_model->latest_state[channel].eqns.x,
-        data[channel], denoised[channel], w, h, stride[channel], alt_data,
+        data[channel], denoised[channel], w, h, stride[channel], sub, alt_data,
         alt_denoised, stride[0], flat_blocks, block_size, num_blocks_w,
         num_blocks_h);
 
@@ -791,8 +823,14 @@
     equation_system_add(&noise_model->combined_state[channel].eqns,
                         &noise_model->latest_state[channel].eqns);
     if (!equation_system_solve(&noise_model->combined_state[channel].eqns)) {
-      fprintf(stderr, "Solving combined noise equation failed!\n");
-      return AOM_NOISE_STATUS_INTERNAL_ERROR;
+      if (channel > 0) {
+        set_chroma_coefficient_fallback_soln(
+            &noise_model->combined_state[channel].eqns);
+      } else {
+        fprintf(stderr, "Solving combined noise equation system failed %d!\n",
+                channel);
+        return AOM_NOISE_STATUS_INTERNAL_ERROR;
+      }
     }
 
     noise_strength_solver_add(
diff --git a/test/noise_model_test.cc b/test/noise_model_test.cc
index 2312a2f..ebbbd3e 100644
--- a/test/noise_model_test.cc
+++ b/test/noise_model_test.cc
@@ -301,21 +301,16 @@
 
     data_.resize(kWidth * kHeight * 3);
     denoised_.resize(kWidth * kHeight * 3);
-    noise_.resize(kWidth * kHeight);
+    noise_.resize(kWidth * kHeight * 3);
     renoise_.resize(kWidth * kHeight);
     flat_blocks_.resize(kNumBlocksX * kNumBlocksY);
 
-    data_ptr_[0] = &data_[0];
-    data_ptr_[1] = &data_[kWidth * kHeight];
-    data_ptr_[2] = &data_[kWidth * kHeight * 2];
-
-    denoised_ptr_[0] = &denoised_[0];
-    denoised_ptr_[1] = &denoised_[kWidth * kHeight];
-    denoised_ptr_[2] = &denoised_[kWidth * kHeight * 2];
-
-    strides_[0] = kWidth;
-    strides_[1] = kWidth;
-    strides_[2] = kWidth;
+    for (int c = 0, offset = 0; c < 3; ++c, offset += kWidth * kHeight) {
+      data_ptr_[c] = &data_[offset];
+      noise_ptr_[c] = &noise_[offset];
+      denoised_ptr_[c] = &denoised_[offset];
+      strides_[c] = kWidth;
+    }
     chroma_sub_[0] = 0;
     chroma_sub_[1] = 0;
   }
@@ -333,6 +328,7 @@
 
   uint8_t *data_ptr_[3];
   uint8_t *denoised_ptr_[3];
+  double *noise_ptr_[3];
   int strides_[3];
   int chroma_sub_[2];
 };
@@ -421,7 +417,7 @@
 }
 
 TEST_F(NoiseModelUpdateTest, UpdateSuccessForScaledWhiteNoise) {
-  const double kCoeffEps = 0.05;
+  const double kCoeffEps = 0.055;
   const double kLowStd = 1;
   const double kHighStd = 4;
   for (int y = 0; y < kHeight; ++y) {
@@ -493,26 +489,39 @@
   const double kStd = 4;
   const double kStdEps = 0.3;
   const int kBlockSize = 16;
-  const double kCoeffEps = 0.05;
-  const double kCoeffs[24] = {
-    0.02884, -0.03356, 0.00633,  0.01757,  0.02849,  -0.04620,
-    0.02833, -0.07178, 0.07076,  -0.11603, -0.10413, -0.16571,
-    0.05158, -0.07969, 0.02640,  -0.07191, 0.02530,  0.41968,
-    0.21450, -0.00702, -0.01401, -0.03676, -0.08713, 0.44196,
+  const double kCoeffEps = 0.06;
+  // Use different coefficients for each channel
+  const double kCoeffs[3][24] = {
+    { 0.02884, -0.03356, 0.00633,  0.01757,  0.02849,  -0.04620,
+      0.02833, -0.07178, 0.07076,  -0.11603, -0.10413, -0.16571,
+      0.05158, -0.07969, 0.02640,  -0.07191, 0.02530,  0.41968,
+      0.21450, -0.00702, -0.01401, -0.03676, -0.08713, 0.44196 },
+    { 0.00269, -0.01291, -0.01513, 0.07234,  0.03208,   0.00477,
+      0.00226, -0.00254, 0.03533,  0.12841,  -0.25970,  -0.06336,
+      0.05238, -0.00845, -0.03118, 0.09043,  -0.36558,  0.48903,
+      0.00595, -0.11938, 0.02106,  0.095956, -0.350139, 0.59305 },
+    { -0.00643, -0.01080, -0.01466, 0.06951, 0.03707,  -0.00482,
+      0.00817,  -0.00909, 0.02949,  0.12181, -0.25210, -0.07886,
+      0.06083,  -0.01210, -0.03108, 0.08944, -0.35875, 0.49150,
+      0.00415,  -0.12905, 0.02870,  0.09740, -0.34610, 0.58824 },
   };
   ASSERT_EQ(model_.n, kNumCoeffs);
-  aom_noise_synth(model_.params.lag, model_.n, model_.coords, kCoeffs,
-                  &noise_[0], kWidth, kHeight);
+  chroma_sub_[0] = chroma_sub_[1] = 1;
+
   flat_blocks_.assign(flat_blocks_.size(), 1);
 
-  // Add noise onto a planar image
-  for (int y = 0; y < kHeight; ++y) {
-    for (int x = 0; x < kWidth; ++x) {
-      for (int c = 0; c < 3; ++c) {
+  // Add different noise onto each plane
+  for (int c = 0; c < 3; ++c) {
+    aom_noise_synth(model_.params.lag, model_.n, model_.coords, kCoeffs[c],
+                    noise_ptr_[c], kWidth, kHeight);
+    const int x_shift = c > 0 ? chroma_sub_[0] : 0;
+    const int y_shift = c > 0 ? chroma_sub_[1] : 0;
+    for (int y = 0; y < (kHeight >> y_shift); ++y) {
+      for (int x = 0; x < (kWidth >> x_shift); ++x) {
         const uint8_t value = 64 + x / 2 + y / 4;
-        denoised_ptr_[c][y * kWidth + x] = value;
         data_ptr_[c][y * kWidth + x] =
-            uint8_t(value + noise_[y * kWidth + x] * kStd);
+            uint8_t(value + noise_ptr_[c][y * strides_[c] + x] * kStd);
+        denoised_ptr_[c][y * strides_[c] + x] = value;
       }
     }
   }
@@ -523,36 +532,33 @@
 
   // For the Y plane, the solved coefficients should be close to the original
   const int n = model_.n;
-  for (int i = 0; i < n; ++i) {
-    EXPECT_NEAR(kCoeffs[i], model_.latest_state[0].eqns.x[i], kCoeffEps);
-    EXPECT_NEAR(kCoeffs[i], model_.combined_state[0].eqns.x[i], kCoeffEps);
-  }
-
-  // Check chroma planes are completely correlated with the Y data
-  for (int c = 1; c < 3; ++c) {
-    // The AR coefficients should be close to zero
-    for (int i = 0; i < model_.n; ++i) {
-      EXPECT_NEAR(0, model_.latest_state[c].eqns.x[i], kCoeffEps);
-      EXPECT_NEAR(0, model_.combined_state[c].eqns.x[i], kCoeffEps);
+  for (int c = 0; c < 3; ++c) {
+    for (int i = 0; i < n; ++i) {
+      EXPECT_NEAR(kCoeffs[c][i], model_.latest_state[c].eqns.x[i], kCoeffEps);
+      EXPECT_NEAR(kCoeffs[c][i], model_.combined_state[c].eqns.x[i], kCoeffEps);
     }
-    // We should have high correlation between the Y plane
-    EXPECT_NEAR(1, model_.latest_state[c].eqns.x[n], kCoeffEps);
-    EXPECT_NEAR(1, model_.combined_state[c].eqns.x[n], kCoeffEps);
+    // The chroma planes should be uncorrelated with the luma plane
+    if (c > 0) {
+      EXPECT_NEAR(0, model_.latest_state[c].eqns.x[n], kCoeffEps);
+      EXPECT_NEAR(0, model_.combined_state[c].eqns.x[n], kCoeffEps);
+    }
+    // Correlation between the coefficient vector and the fitted coefficients
+    // should be close to 1.
+    EXPECT_LT(0.98, aom_normalized_cross_correlation(
+                        model_.latest_state[c].eqns.x, kCoeffs[c], kNumCoeffs));
+
+    aom_noise_synth(model_.params.lag, model_.n, model_.coords,
+                    model_.latest_state[c].eqns.x, &renoise_[0], kWidth,
+                    kHeight);
+
+    EXPECT_TRUE(aom_noise_data_validate(&renoise_[0], kWidth, kHeight));
   }
 
-  // Correlation between the coefficient vector and the fitted coefficients
-  // should be close to 1.
-  EXPECT_LT(0.99, aom_normalized_cross_correlation(
-                      model_.latest_state[0].eqns.x, kCoeffs, kNumCoeffs));
-
-  aom_noise_synth(model_.params.lag, model_.n, model_.coords,
-                  model_.latest_state[0].eqns.x, &renoise_[0], kWidth, kHeight);
-
-  EXPECT_TRUE(aom_noise_data_validate(&renoise_[0], kWidth, kHeight));
-
-  // Check noise variance
-  for (int i = 0; i < model_.latest_state[0].strength_solver.eqns.n; ++i) {
-    EXPECT_NEAR(kStd, model_.latest_state[0].strength_solver.eqns.x[i],
-                kStdEps);
+  // Check fitted noise strength
+  for (int c = 0; c < 3; ++c) {
+    for (int i = 0; i < model_.latest_state[c].strength_solver.eqns.n; ++i) {
+      EXPECT_NEAR(kStd, model_.latest_state[c].strength_solver.eqns.x[i],
+                  kStdEps);
+    }
   }
 }