diff options
author | Neil Birkbeck <birkbeck@google.com> | 2017-12-21 13:18:32 -0800 |
---|---|---|
committer | Yaowu Xu <yaowu@google.com> | 2018-01-17 02:21:12 +0000 |
commit | ed25a610ce6535cf3cfd639e5259317c239564ed (patch) | |
tree | 18652efc41ec9b7b74a62952d2f0a9174c6f5c4c | |
parent | 22a51d9e96d99734a3d46fc60d73197c3c30d42d (diff) |
Add utilities to aom_dsp for modeling correlated noise.
The auto-regressive model allows for different window shapes
and different lag sizes.
Although most likely to be used as a reference for modeling
noise in AV1, the model is currently parameterized more generally
than AV1 needs.
I will add an example (hopefully with a denoiser) in future
commits.
Change-Id: I1ba1067543601c2c01db4970d42766bb35da77f0
-rw-r--r-- | aom_dsp/aom_dsp.cmake | 4 | ||||
-rw-r--r-- | aom_dsp/noise_model.c | 811 | ||||
-rw-r--r-- | aom_dsp/noise_model.h | 229 | ||||
-rw-r--r-- | aom_dsp/noise_util.c | 149 | ||||
-rw-r--r-- | aom_dsp/noise_util.h | 35 | ||||
-rw-r--r-- | test/noise_model_test.cc | 563 | ||||
-rw-r--r-- | test/test.cmake | 1 |
7 files changed, 1792 insertions, 0 deletions
diff --git a/aom_dsp/aom_dsp.cmake b/aom_dsp/aom_dsp.cmake index f676c4026..e21a6d95d 100644 --- a/aom_dsp/aom_dsp.cmake +++ b/aom_dsp/aom_dsp.cmake @@ -302,6 +302,10 @@ if (CONFIG_AV1_ENCODER) "${AOM_ROOT}/aom_dsp/bitwriter.h" "${AOM_ROOT}/aom_dsp/bitwriter_buffer.c" "${AOM_ROOT}/aom_dsp/bitwriter_buffer.h" + "${AOM_ROOT}/aom_dsp/noise_util.h" + "${AOM_ROOT}/aom_dsp/noise_util.c" + "${AOM_ROOT}/aom_dsp/noise_model.c" + "${AOM_ROOT}/aom_dsp/noise_model.c" "${AOM_ROOT}/aom_dsp/psnr.c" "${AOM_ROOT}/aom_dsp/psnr.h" "${AOM_ROOT}/aom_dsp/sad.c" diff --git a/aom_dsp/noise_model.c b/aom_dsp/noise_model.c new file mode 100644 index 000000000..beed8b35c --- /dev/null +++ b/aom_dsp/noise_model.c @@ -0,0 +1,811 @@ +/* + * Copyright (c) 2017, Alliance for Open Media. All rights reserved + * + * This source code is subject to the terms of the BSD 2 Clause License and + * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License + * was not distributed with this source code in the LICENSE file, you can + * obtain it at www.aomedia.org/license/software. If the Alliance for Open + * Media Patent License 1.0 was not distributed with this source code in the + * PATENTS file, you can obtain it at www.aomedia.org/license/patent. + */ + +#include <math.h> +#include <stdio.h> +#include <stdlib.h> +#include <string.h> + +#include "aom_dsp/aom_dsp_common.h" +#include "aom_dsp/noise_model.h" +#include "aom_dsp/noise_util.h" +#include "aom_mem/aom_mem.h" +#include "av1/encoder/mathutils.h" + +#define kLowPolyNumParams 3 +static const double kBlockNormalization = 255.0; +static const int kMaxLag = 4; + +static double get_block_mean(const uint8_t *data, int w, int h, int stride, + int x_o, int y_o, int block_size) { + const int max_h = AOMMIN(h - y_o, block_size); + const int max_w = AOMMIN(w - x_o, block_size); + double block_mean = 0; + for (int y = 0; y < max_h; ++y) { + for (int x = 0; x < max_w; ++x) { + block_mean += data[(y_o + y) * stride + x_o + x]; + } + } + return block_mean / (max_w * max_h); +} + +static void equation_system_clear(aom_equation_system_t *eqns) { + const int n = eqns->n; + memset(eqns->A, 0, sizeof(*eqns->A) * n * n); + memset(eqns->x, 0, sizeof(*eqns->x) * n); + memset(eqns->b, 0, sizeof(*eqns->b) * n); +} + +static int equation_system_init(aom_equation_system_t *eqns, int n) { + eqns->A = (double *)aom_malloc(sizeof(*eqns->A) * n * n); + eqns->b = (double *)aom_malloc(sizeof(*eqns->b) * n); + eqns->x = (double *)aom_malloc(sizeof(*eqns->x) * n); + eqns->n = n; + if (!eqns->A || !eqns->b || !eqns->x) { + fprintf(stderr, "Failed to allocate system of equations of size %d\n", n); + aom_free(eqns->A); + aom_free(eqns->b); + aom_free(eqns->x); + memset(eqns, 0, sizeof(*eqns)); + return 0; + } + equation_system_clear(eqns); + return 1; +} + +static int equation_system_solve(aom_equation_system_t *eqns) { + const int n = eqns->n; + double *b = (double *)aom_malloc(sizeof(*b) * n); + double *A = (double *)aom_malloc(sizeof(*A) * n * n); + int ret = 0; + if (A == NULL || b == NULL) { + fprintf(stderr, "Unable to allocate temp values of size %dx%d\n", n, n); + aom_free(b); + aom_free(A); + return 0; + } + memcpy(A, eqns->A, sizeof(*eqns->A) * n * n); + memcpy(b, eqns->b, sizeof(*eqns->b) * n); + ret = linsolve(n, A, eqns->n, b, eqns->x); + aom_free(b); + aom_free(A); + + if (ret == 0) { + fprintf(stderr, "Solving %dx%d system failed!\n", n, n); + return 0; + } + return 1; +} + +static void equation_system_add(aom_equation_system_t *dest, + aom_equation_system_t *src) { + const int n = dest->n; + int i, j; + for (i = 0; i < n; ++i) { + for (j = 0; j < n; ++j) { + dest->A[i * n + j] += src->A[i * n + j]; + } + dest->b[i] += src->b[i]; + } +} + +static void equation_system_free(aom_equation_system_t *eqns) { + if (!eqns) return; + aom_free(eqns->A); + aom_free(eqns->b); + aom_free(eqns->x); + memset(eqns, 0, sizeof(*eqns)); +} + +static void noise_strength_solver_add(aom_noise_strength_solver_t *dest, + aom_noise_strength_solver_t *src) { + equation_system_add(&dest->eqns, &src->eqns); + dest->num_equations += src->num_equations; + dest->total += src->total; +} + +// Return the number of coefficients required for the given parameters +static int num_coeffs(const aom_noise_model_params_t params) { + const int n = 2 * params.lag + 1; + switch (params.shape) { + case AOM_NOISE_SHAPE_DIAMOND: return params.lag * (params.lag + 1); + case AOM_NOISE_SHAPE_SQUARE: return (n * n) / 2; + } + return 0; +} + +static int noise_state_init(aom_noise_state_t *state, int n) { + const int kNumBins = 20; + if (!equation_system_init(&state->eqns, n)) { + fprintf(stderr, "Failed initialization noise state with size %d\n", n); + return 0; + } + return aom_noise_strength_solver_init(&state->strength_solver, kNumBins); +} + +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)); + if (!lut->points) return 0; + lut->num_points = num_points; + memset(lut->points, 0, sizeof(*lut->points) * num_points); + return 1; +} + +void aom_noise_strength_lut_free(aom_noise_strength_lut_t *lut) { + if (!lut) return; + aom_free(lut->points); + memset(lut, 0, sizeof(*lut)); +} + +double aom_noise_strength_lut_eval(const aom_noise_strength_lut_t *lut, + double x) { + int i = 0; + // Constant extrapolation for x < x_0. + if (x < lut->points[0][0]) return lut->points[0][1]; + for (i = 0; i < lut->num_points - 1; ++i) { + if (x >= lut->points[i][0] && x <= lut->points[i + 1][0]) { + const double a = + (x - lut->points[i][0]) / (lut->points[i + 1][0] - lut->points[i][0]); + return lut->points[i + 1][1] * a + lut->points[i][1] * (1.0 - a); + } + } + // Constant extrapolation for x > x_{n-1} + return lut->points[lut->num_points - 1][1]; +} + +void aom_noise_strength_solver_add_measurement( + aom_noise_strength_solver_t *solver, double block_mean, double noise_std) { + const double val = + AOMMIN(AOMMAX(block_mean, solver->min_intensity), solver->max_intensity); + const double range = solver->max_intensity - solver->min_intensity; + const double bin = + (solver->num_bins - 1) * (val - solver->min_intensity) / range; + const int bin_i0 = (int)floor(bin); + const int bin_i1 = AOMMIN(solver->num_bins - 1, bin_i0 + 1); + const double a = bin - bin_i0; + const int n = solver->num_bins; + solver->eqns.A[bin_i0 * n + bin_i0] += (1.0 - a) * (1.0 - a); + solver->eqns.A[bin_i1 * n + bin_i0] += a * (1.0 - a); + solver->eqns.A[bin_i1 * n + bin_i1] += a * a; + solver->eqns.A[bin_i0 * n + bin_i1] += a * (1.0 - a); + solver->eqns.b[bin_i0] += (1.0 - a) * noise_std; + solver->eqns.b[bin_i1] += a * noise_std; + solver->total += noise_std; + solver->num_equations++; +} + +int aom_noise_strength_solver_solve(aom_noise_strength_solver_t *solver) { + // Add regularization proportional to the number of constraints + const int n = solver->num_bins; + const double kAlpha = 2.0 * (double)(solver->num_equations) / n; + int result = 0; + double mean = 0; + + // Do this in a non-destructive manner so it is not confusing to the caller + double *old_A = solver->eqns.A; + double *A = (double *)aom_malloc(sizeof(*A) * n * n); + if (!A) { + fprintf(stderr, "Unable to allocate copy of A\n"); + return 0; + } + memcpy(A, old_A, sizeof(*A) * n * n); + + for (int i = 0; i < n; ++i) { + const int i_lo = AOMMAX(0, i - 1); + const int i_hi = AOMMIN(n - 1, i + 1); + A[i * n + i_lo] -= kAlpha; + A[i * n + i] += 2 * kAlpha; + A[i * n + i_hi] -= kAlpha; + } + + // Small regularization to give average noise strength + mean = solver->total / solver->num_equations; + for (int i = 0; i < n; ++i) { + A[i * n + i] += 1.0 / 8192.; + solver->eqns.b[i] += mean / 8192.; + } + solver->eqns.A = A; + result = equation_system_solve(&solver->eqns); + solver->eqns.A = old_A; + + aom_free(A); + return result; +} + +int aom_noise_strength_solver_init(aom_noise_strength_solver_t *solver, + int num_bins) { + if (!solver) return 0; + memset(solver, 0, sizeof(*solver)); + solver->num_bins = num_bins; + solver->min_intensity = 0; + solver->max_intensity = 255; + return equation_system_init(&solver->eqns, num_bins); +} + +void aom_noise_strength_solver_free(aom_noise_strength_solver_t *solver) { + if (!solver) return; + equation_system_free(&solver->eqns); +} + +double aom_noise_strength_solver_get_center( + const aom_noise_strength_solver_t *solver, int i) { + const double range = solver->max_intensity - solver->min_intensity; + const int n = solver->num_bins; + return ((double)i) / (n - 1) * range + solver->min_intensity; +} + +int aom_noise_strength_solver_fit_piecewise( + const aom_noise_strength_solver_t *solver, aom_noise_strength_lut_t *lut) { + const double kTolerance = 0.1; + int low_point = 0; + aom_equation_system_t sys; + if (!equation_system_init(&sys, 2)) { + fprintf(stderr, "Failed to init equation system\n"); + return 0; + } + + if (!aom_noise_strength_lut_init(lut, solver->num_bins + 1)) { + fprintf(stderr, "Failed to init lut\n"); + return 0; + } + + lut->points[0][0] = aom_noise_strength_solver_get_center(solver, 0); + lut->points[0][1] = solver->eqns.x[0]; + lut->num_points = 1; + + while (low_point < solver->num_bins - 1) { + int i = low_point; + equation_system_clear(&sys); + for (; i < solver->num_bins; ++i) { + int x = i - low_point; + double b = 1; + sys.A[0 * 2 + 0] += x * x; + sys.A[0 * 2 + 1] += x * b; + sys.A[1 * 2 + 0] += x * b; + sys.A[1 * 2 + 1] += b * b; + sys.b[0] += x * solver->eqns.x[i]; + sys.b[1] += b * solver->eqns.x[i]; + + if (x > 1) { + double res = 0; + int k; + equation_system_solve(&sys); + + for (k = low_point; k <= i; ++k) { + double y; + x = k - low_point; + y = sys.x[0] * x + sys.x[1]; + y -= solver->eqns.x[k]; + res += y * y; + } + const int n = i - low_point + 1; + if (sqrt(res / n) > kTolerance) { + low_point = i - 1; + + lut->points[lut->num_points][0] = + aom_noise_strength_solver_get_center(solver, i - 1); + lut->points[lut->num_points][1] = solver->eqns.x[i - 1]; + lut->num_points++; + } + } + } + if (i == solver->num_bins) { + lut->points[lut->num_points][0] = + aom_noise_strength_solver_get_center(solver, i - 1); + lut->points[lut->num_points][1] = solver->eqns.x[i - 1]; + lut->num_points++; + break; + } + } + equation_system_free(&sys); + return 1; +} + +int aom_flat_block_finder_init(aom_flat_block_finder_t *block_finder, + int block_size) { + const int n = block_size * block_size; + aom_equation_system_t eqns; + double *AtA_inv = 0; + double *A = 0; + int x = 0, y = 0, i = 0, j = 0; + if (!equation_system_init(&eqns, kLowPolyNumParams)) { + fprintf(stderr, "Failed to init equation system for block_size=%d\n", + block_size); + return 0; + } + + AtA_inv = (double *)aom_malloc(kLowPolyNumParams * kLowPolyNumParams * + sizeof(*AtA_inv)); + A = (double *)aom_malloc(kLowPolyNumParams * n * sizeof(*A)); + if (AtA_inv == NULL || A == NULL) { + fprintf(stderr, "Failed to alloc A or AtA_inv for block_size=%d\n", + block_size); + aom_free(AtA_inv); + aom_free(A); + equation_system_free(&eqns); + return 0; + } + + block_finder->A = A; + block_finder->AtA_inv = AtA_inv; + block_finder->block_size = block_size; + + for (y = 0; y < block_size; ++y) { + const double yd = ((double)y - block_size / 2.) / (block_size / 2.); + for (x = 0; x < block_size; ++x) { + const double xd = ((double)x - block_size / 2.) / (block_size / 2.); + const double coords[3] = { yd, xd, 1 }; + const int row = y * block_size + x; + A[kLowPolyNumParams * row + 0] = yd; + A[kLowPolyNumParams * row + 1] = xd; + A[kLowPolyNumParams * row + 2] = 1; + + for (i = 0; i < kLowPolyNumParams; ++i) { + for (j = 0; j < kLowPolyNumParams; ++j) { + eqns.A[kLowPolyNumParams * i + j] += coords[i] * coords[j]; + } + } + } + } + + // Lazy inverse using existing equation solver. + for (i = 0; i < kLowPolyNumParams; ++i) { + memset(eqns.b, 0, sizeof(*eqns.b) * kLowPolyNumParams); + eqns.b[i] = 1; + equation_system_solve(&eqns); + + for (j = 0; j < kLowPolyNumParams; ++j) { + AtA_inv[j * kLowPolyNumParams + i] = eqns.x[j]; + } + } + equation_system_free(&eqns); + return 1; +} + +void aom_flat_block_finder_free(aom_flat_block_finder_t *block_finder) { + if (!block_finder) return; + aom_free(block_finder->A); + aom_free(block_finder->AtA_inv); + memset(block_finder, 0, sizeof(*block_finder)); +} + +void aom_flat_block_finder_extract_block( + const aom_flat_block_finder_t *block_finder, const uint8_t *const data, + int w, int h, int stride, int offsx, int offsy, double *plane, + double *block) { + const int block_size = block_finder->block_size; + const int n = block_size * block_size; + const double *A = block_finder->A; + const double *AtA_inv = block_finder->AtA_inv; + double plane_coords[kLowPolyNumParams]; + double AtA_inv_b[kLowPolyNumParams]; + int xi, yi, i; + + for (yi = 0; yi < block_size; ++yi) { + const int y = AOMMAX(0, AOMMIN(h - 1, offsy + yi)); + for (xi = 0; xi < block_size; ++xi) { + const int x = AOMMAX(0, AOMMIN(w - 1, offsx + xi)); + block[yi * block_size + xi] = + ((double)data[y * stride + x]) / kBlockNormalization; + } + } + + multiply_mat(block, A, AtA_inv_b, 1, n, kLowPolyNumParams); + multiply_mat(AtA_inv, AtA_inv_b, plane_coords, kLowPolyNumParams, + kLowPolyNumParams, 1); + multiply_mat(A, plane_coords, plane, n, kLowPolyNumParams, 1); + + for (i = 0; i < n; ++i) { + block[i] -= plane[i]; + } +} + +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) { + 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 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; + int num_flat = 0; + 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) { + fprintf(stderr, "Failed to allocate memory for block of size %d\n", n); + aom_free(plane); + aom_free(block); + return -1; + } + + for (by = 0; by < num_blocks_h; ++by) { + for (bx = 0; bx < num_blocks_w; ++bx) { + // Compute gradient covariance matrix. + double Gxx = 0, Gxy = 0, Gyy = 0; + double var = 0; + double mean = 0; + int xi, yi; + aom_flat_block_finder_extract_block(block_finder, data, w, h, stride, + bx * block_size, by * block_size, + plane, block); + + for (yi = 1; yi < block_size - 1; ++yi) { + for (xi = 1; xi < block_size - 1; ++xi) { + const double gx = (block[yi * block_size + xi + 1] - + block[yi * block_size + xi - 1]) / + 2; + const double gy = (block[yi * block_size + xi + block_size] - + block[yi * block_size + xi - block_size]) / + 2; + Gxx += gx * gx; + Gxy += gx * gy; + Gyy += gy * gy; + + mean += block[yi * block_size + xi]; + var += block[yi * block_size + xi] * block[yi * block_size + xi]; + } + } + 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; + + { + 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 int is_flat = (trace < kTraceThreshold) && + (e1 / AOMMAX(e2, 1e-8) < kRatioThreshold) && + norm < kNormThreshold && var > kVarThreshold; + flat_blocks[by * num_blocks_w + bx] = is_flat ? 255 : 0; + num_flat += is_flat; + } + } + } + + aom_free(block); + aom_free(plane); + return num_flat; +} + +int aom_noise_model_init(aom_noise_model_t *model, + const aom_noise_model_params_t params) { + const int n = num_coeffs(params); + const int lag = params.lag; + int x = 0, y = 0, i = 0, c = 0; + + memset(model, 0, sizeof(*model)); + if (params.lag < 1) { + fprintf(stderr, "Invalid noise param: lag = %d must be >= 1\n", params.lag); + return 0; + } + if (params.lag > kMaxLag) { + fprintf(stderr, "Invalid noise param: lag = %d must be <= %d\n", params.lag, + kMaxLag); + return 0; + } + + memcpy(&model->params, ¶ms, sizeof(params)); + for (c = 0; c < 3; ++c) { + if (!noise_state_init(&model->combined_state[c], n + (c > 0))) { + fprintf(stderr, "Failed to allocate noise state for channel %d\n", c); + aom_noise_model_free(model); + return 0; + } + if (!noise_state_init(&model->latest_state[c], n + (c > 0))) { + fprintf(stderr, "Failed to allocate noise state for channel %d\n", c); + aom_noise_model_free(model); + return 0; + } + } + model->n = n; + model->coords = (int(*)[2])aom_malloc(sizeof(*model->coords) * n); + + for (y = -lag; y <= 0; ++y) { + const int max_x = y == 0 ? -1 : lag; + for (x = -lag; x <= max_x; ++x) { + switch (params.shape) { + case AOM_NOISE_SHAPE_DIAMOND: + if (abs(x) <= y + lag) { + model->coords[i][0] = x; + model->coords[i][1] = y; + ++i; + } + break; + case AOM_NOISE_SHAPE_SQUARE: + model->coords[i][0] = x; + model->coords[i][1] = y; + ++i; + break; + default: + fprintf(stderr, "Invalid shape\n"); + aom_noise_model_free(model); + return 0; + } + } + } + assert(i == n); + return 1; +} + +void aom_noise_model_free(aom_noise_model_t *model) { + int c = 0; + if (!model) return; + + aom_free(model->coords); + for (c = 0; c < 3; ++c) { + equation_system_free(&model->latest_state[c].eqns); + equation_system_free(&model->combined_state[c].eqns); + + equation_system_free(&model->latest_state[c].strength_solver.eqns); + equation_system_free(&model->combined_state[c].strength_solver.eqns); + } + memset(model, 0, sizeof(*model)); +} + +static int add_block_observations( + aom_noise_model_t *noise_model, int c, 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; + double *A = noise_model->latest_state[c].eqns.A; + double *b = noise_model->latest_state[c].eqns.b; + 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; + for (bx = 0; bx < num_blocks_w; ++bx) { + const int x_o = bx * block_size; + 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]) { + continue; + } + 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); + x_end = AOMMIN( + w - bx * block_size - lag, + (bx + 1 < num_blocks_w && flat_blocks[by * num_blocks_w + bx + 1]) + ? block_size + : block_size - lag); + for (y = y_start; y < y_end; ++y) { + for (x = x_start; x < x_end; ++x) { + double val = 0; + for (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; + assert(x_i < w && y_i < h); + buffer[i] = ((double)(data[y_i * stride + x_i]) - + (double)(denoised[y_i * stride + x_i])); + } + val = ((double)data[(y_o + y) * stride + (x_o + x)]) - + ((double)denoised[(y_o + y) * stride + (x_o + x)]); + + // For the color channels we must also consider the correlation with + // 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)]); + } + + for (i = 0; i < n; ++i) { + for (j = 0; j < n; ++j) { + A[i * n + j] += (buffer[i] * buffer[j]) / + (kBlockNormalization * kBlockNormalization); + } + b[i] += + (buffer[i] * val) / (kBlockNormalization * kBlockNormalization); + } + } + } + } + } + aom_free(buffer); + 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) { + 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; + 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); + double noise_var = 0; + int num_samples_in_block = 0; + int y_start = + (by > 0 && flat_blocks[(by - 1) * num_blocks_w + bx]) ? 0 : lag; + 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 = + (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) { + const double actual = + ((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; + 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)])); + } + 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); + } + } +} + +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) { + 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; + int num_blocks = 0; + int i = 0, channel = 0; + + if (block_size <= 1) { + fprintf(stderr, "block_size = %d must be > 1\n", block_size); + return AOM_NOISE_STATUS_INVALID_ARGUMENT; + } + + if (block_size < noise_model->params.lag * 2 + 1) { + fprintf(stderr, "block_size = %d must be >= %d\n", block_size, + noise_model->params.lag * 2 + 1); + return AOM_NOISE_STATUS_INVALID_ARGUMENT; + } + + // Clear the latest equation system + for (i = 0; i < 3; ++i) { + equation_system_clear(&noise_model->latest_state[i].eqns); + } + + // Check that we have enough flat blocks + for (i = 0; i < num_blocks_h * num_blocks_w; ++i) { + if (flat_blocks[i]) { + num_blocks++; + } + } + + if (num_blocks <= 1) { + fprintf(stderr, "Not enough flat blocks to update noise estimate\n"); + return AOM_NOISE_STATUS_INSUFFICIENT_FLAT_BLOCKS; + } + + for (channel = 0; channel < 3; ++channel) { + 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; + 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, + block_size, num_blocks_w, num_blocks_h)) { + fprintf(stderr, "Adding block observation failed\n"); + return AOM_NOISE_STATUS_INTERNAL_ERROR; + } + + 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; + } + + add_noise_std_observations( + noise_model, channel, noise_model->latest_state[channel].eqns.x, + data[channel], denoised[channel], w, h, stride[channel], alt_data, + alt_denoised, stride[0], flat_blocks, block_size, num_blocks_w, + num_blocks_h); + + if (!aom_noise_strength_solver_solve( + &noise_model->latest_state[channel].strength_solver)) { + fprintf(stderr, "Solving latest noise strength failed!\n"); + return AOM_NOISE_STATUS_INTERNAL_ERROR; + } + + // Check noise characteristics and return if error. + if (channel == 0 && + noise_model->combined_state[channel].strength_solver.num_equations > + 0) { + y_model_different = 1; + } + + // Don't update the combined stats if the y model is different. + if (y_model_different) continue; + + 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; + } + + noise_strength_solver_add( + &noise_model->combined_state[channel].strength_solver, + &noise_model->latest_state[channel].strength_solver); + + if (!aom_noise_strength_solver_solve( + &noise_model->combined_state[channel].strength_solver)) { + fprintf(stderr, "Solving combined noise strength failed!\n"); + return AOM_NOISE_STATUS_INTERNAL_ERROR; + } + } + + return y_model_different ? AOM_NOISE_STATUS_DIFFERENT_NOISE_TYPE + : AOM_NOISE_STATUS_OK; +} diff --git a/aom_dsp/noise_model.h b/aom_dsp/noise_model.h new file mode 100644 index 000000000..c373f4540 --- /dev/null +++ b/aom_dsp/noise_model.h @@ -0,0 +1,229 @@ +/* + * Copyright (c) 2017, Alliance for Open Media. All rights reserved + * + * This source code is subject to the terms of the BSD 2 Clause License and + * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License + * was not distributed with this source code in the LICENSE file, you can + * obtain it at www.aomedia.org/license/software. If the Alliance for Open + * Media Patent License 1.0 was not distributed with this source code in the + * PATENTS file, you can obtain it at www.aomedia.org/license/patent. + */ + +#ifndef AOM_DSP_NOISE_MODEL_H_ +#define AOM_DSP_NOISE_MODEL_H_ + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +#include <stdint.h> + +/*!\brief Wrapper of data required to represent linear system of eqns and soln. + */ +typedef struct { + double *A; + double *b; + double *x; + int n; +} aom_equation_system_t; + +/*!\brief Representation of a piecewise linear curve + * + * Holds n points as (x, y) pairs, that store the curve. + */ +typedef struct { + double (*points)[2]; + int num_points; +} aom_noise_strength_lut_t; + +/*!\brief Init the noise strength lut with the given number of points*/ +int aom_noise_strength_lut_init(aom_noise_strength_lut_t *lut, int num_points); + +/*!\brief Frees the noise strength lut. */ +void aom_noise_strength_lut_free(aom_noise_strength_lut_t *lut); + +/*!\brief Evaluate the lut at the point x. + * + * \param[in] lut The lut data. + * \param[in] x The coordinate to evaluate the lut. + */ +double aom_noise_strength_lut_eval(const aom_noise_strength_lut_t *lut, + double x); + +/*!\brief Helper struct to model noise strength as a function of intensity. + * + * Internally, this structure holds a representation of a linear system + * of equations that models noise strength (standard deviation) as a + * function of intensity. The mapping is initially stored using a + * piecewise representation with evenly spaced bins that cover the entire + * domain from [min_intensity, max_intensity]. Each observation (x,y) gives a + * constraint of the form: + * y_{i} (1 - a) + y_{i+1} a = y + * where y_{i} is the value of bin i and x_{i} <= x <= x_{i+1} and + * a = x/(x_{i+1} - x{i}). The equation system holds the corresponding + * normal equations. + * + * As there may be missing data, the solution is regularized to get a + * complete set of values for the bins. A reduced representation after + * solving can be obtained by getting the corresponding noise_strength_lut_t. + */ +typedef struct { + aom_equation_system_t eqns; + double min_intensity; + double max_intensity; + int num_bins; + int num_equations; + double total; +} aom_noise_strength_solver_t; + +/*!\brief Initializes the noise solver with the given number of bins. + * + * Returns 0 if initialization fails. + * + * \param[in] solver The noise solver to be initialized. + * \param[in] num_bins Number of bins to use in the internal representation. + */ +int aom_noise_strength_solver_init(aom_noise_strength_solver_t *solver, + int num_bins); +void aom_noise_strength_solver_free(aom_noise_strength_solver_t *solver); + +/*!\brief Gets the x coordinate of bin i. + * + * \param[in] i The bin whose coordinate to query. + */ +double aom_noise_strength_solver_get_center( + const aom_noise_strength_solver_t *solver, int i); + +/*!\brief Add an observation of the block mean intensity to its noise strength. + * + * \param[in] block_mean The average block intensity, + * \param[in] noise_std The observed noise strength. + */ +void aom_noise_strength_solver_add_measurement( + aom_noise_strength_solver_t *solver, double block_mean, double noise_std); + +/*!\brief Solves the current set of equations for the noise strength. */ +int aom_noise_strength_solver_solve(aom_noise_strength_solver_t *solver); + +/*!\brief Fits a reduced piecewise linear lut to the internal solution + * + * \param[out] lut The output piecewise linear lut. + */ +int aom_noise_strength_solver_fit_piecewise( + const aom_noise_strength_solver_t *solver, aom_noise_strength_lut_t *lut); + +/*!\brief Helper for holding precomputed data for finding flat blocks. + * + * Internally a block is modeled with a low-order polynomial model. A + * planar model would be a bunch of equations like: + * <[y_i x_i 1], [a_1, a_2, a_3]> = b_i + * for each point in the block. The system matrix A with row i as [y_i x_i 1] + * is maintained as is the inverse, inv(A'*A), so that the plane parameters + * can be fit for each block. + */ +typedef struct { + double *AtA_inv; + double *A; + int num_params; // The number of parameters used for internal low-order model + int block_size; // The block size the finder was initialized with +} aom_flat_block_finder_t; + +/*!\brief Init the block_finder with the given block size */ +int aom_flat_block_finder_init(aom_flat_block_finder_t *block_finder, + int block_size); +void aom_flat_block_finder_free(aom_flat_block_finder_t *block_finder); + +/*!\brief Helper to extract a block and low order "planar" model. */ +void aom_flat_block_finder_extract_block( + const aom_flat_block_finder_t *block_finder, const uint8_t *const data, + int w, int h, int stride, int offsx, int offsy, double *plane, + double *block); + +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 noise shape indicates the allowed coefficients in the AR model. +typedef enum { + AOM_NOISE_SHAPE_DIAMOND = 0, + AOM_NOISE_SHAPE_SQUARE = 1 +} aom_noise_shape; + +// The parameters of the noise model include the shape type and the lag. +typedef struct { + aom_noise_shape shape; + int lag; +} aom_noise_model_params_t; + +/*!\brief State of a noise model estimate for a single channel. + * + * This contains a system of equations that can be used to solve + * for the auto-regressive coefficients as well as a noise strength + * solver that can be used to model noise strength as a function of + * intensity. + */ +typedef struct { + aom_equation_system_t eqns; + aom_noise_strength_solver_t strength_solver; +} aom_noise_state_t; + +/*!\brief Complete model of noise for a planar video + * + * This includes a noise model for the latest frame and an aggregated + * estimate over all previous frames that had similar parameters. + */ +typedef struct { + aom_noise_model_params_t params; + aom_noise_state_t combined_state[3]; // Combined state per channel + aom_noise_state_t latest_state[3]; // Latest state per channel + int (*coords)[2]; // Offsets (x,y) of the coefficient samples + int n; // Number of parameters (size of coords) +} aom_noise_model_t; + +/*!\brief Result of a noise model update. */ +typedef enum { + AOM_NOISE_STATUS_OK = 0, + AOM_NOISE_STATUS_INVALID_ARGUMENT, + AOM_NOISE_STATUS_INSUFFICIENT_FLAT_BLOCKS, + AOM_NOISE_STATUS_DIFFERENT_NOISE_TYPE, + AOM_NOISE_STATUS_INTERNAL_ERROR, +} aom_noise_status_t; + +/*!\brief Initializes a noise model with the given parameters. + * + * Returns 0 on failure. + */ +int aom_noise_model_init(aom_noise_model_t *model, + const aom_noise_model_params_t params); +void aom_noise_model_free(aom_noise_model_t *model); + +/*!\brief Updates the noise model with a new frame observation. + * + * Updates the noise model with measurements from the given input frame and a + * denoised variant of it. Noise is sampled from flat blocks using the flat + * block map. + * + * Returns a noise_status indicating if the update was successful. If the + * Update was successful, the combined_state is updated with measurements from + * the provided frame. If status is OK or DIFFERENT_NOISE_TYPE, the latest noise + * state will be updated with measurements from the provided frame. + * + * \param[in,out] noise_model The noise model to be updated + * \param[in] data Raw frame data + * \param[in] denoised Denoised frame data. + * \param[in] w Frame width + * \param[in] h Frame height + * \param[in] strides Stride of the planes + * \param[in] chroma_sub_log2 Chroma subsampling for planes != 0. + * \param[in] flat_blocks A map to blocks that have been determined flat + * \param[in] block_size The size of blocks. + */ +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 strides[3], + int chroma_sub_log2[2], const uint8_t *const flat_blocks, int block_size); + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus +#endif // AOM_DSP_NOISE_MODEL_H_ diff --git a/aom_dsp/noise_util.c b/aom_dsp/noise_util.c new file mode 100644 index 000000000..d1a0ccf90 --- /dev/null +++ b/aom_dsp/noise_util.c @@ -0,0 +1,149 @@ +/* + * Copyright (c) 2017, Alliance for Open Media. All rights reserved + * + * This source code is subject to the terms of the BSD 2 Clause License and + * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License + * was not distributed with this source code in the LICENSE file, you can + * obtain it at www.aomedia.org/license/software. If the Alliance for Open + * Media Patent License 1.0 was not distributed with this source code in the + * PATENTS file, you can obtain it at www.aomedia.org/license/patent. + */ + +#include <math.h> + +#include <stdio.h> +#include <stdlib.h> + +#include "aom_dsp/noise_util.h" +#include "aom_mem/aom_mem.h" + +// Return normally distrbuted values with standard deviation of sigma. +double aom_randn(double sigma) { + while (1) { + const double u = 2.0 * ((double)rand()) / RAND_MAX - 1.0; + const double v = 2.0 * ((double)rand()) / RAND_MAX - 1.0; + const double s = u * u + v * v; + if (s > 0 && s < 1) { + return sigma * (u * sqrt(-2.0 * log(s) / s)); + } + } + return 0; +} + +double aom_normalized_cross_correlation(const double *a, const double *b, + int n) { + double c = 0; + double a_len = 0; + double b_len = 0; + for (int i = 0; i < n; ++i) { + a_len += a[i] * a[i]; + b_len += b[i] * b[i]; + c += a[i] * b[i]; + } + return c / (sqrt(a_len) * sqrt(b_len)); +} + +void aom_noise_synth(int lag, int n, const int (*coords)[2], + const double *coeffs, double *data, int w, int h) { + const int pad_size = 3 * lag; + const int padded_w = w + pad_size; + const int padded_h = h + pad_size; + int x = 0, y = 0; + double *padded = (double *)aom_malloc(padded_w * padded_h * sizeof(*padded)); + + for (y = 0; y < padded_h; ++y) { + for (x = 0; x < padded_w; ++x) { + padded[y * padded_w + x] = aom_randn(1.0); + } + } + for (y = lag; y < padded_h; ++y) { + for (x = lag; x < padded_w; ++x) { + double sum = 0; + int i = 0; + for (i = 0; i < n; ++i) { + const int dx = coords[i][0]; + const int dy = coords[i][1]; + sum += padded[(y + dy) * padded_w + (x + dx)] * coeffs[i]; + } + padded[y * padded_w + x] += sum; + } + } + // Copy over the padded rows to the output + for (y = 0; y < h; ++y) { + memcpy(data + y * w, padded + y * padded_w, sizeof(*data) * w); + } + aom_free(padded); +} + +int aom_noise_data_validate(const double *data, int w, int h) { + const double kVarianceThreshold = 2; + const double kMeanThreshold = 2; + + int x = 0, y = 0; + int ret_value = 1; + double var = 0, mean = 0; + double *mean_x, *mean_y, *var_x, *var_y; + + // Check that noise variance is not increasing in x or y + // and that the data is zero mean. + mean_x = (double *)aom_malloc(sizeof(*mean_x) * w); + var_x = (double *)aom_malloc(sizeof(*var_x) * w); + mean_y = (double *)aom_malloc(sizeof(*mean_x) * h); + var_y = (double *)aom_malloc(sizeof(*var_y) * h); + + memset(mean_x, 0, sizeof(*mean_x) * w); + memset(var_x, 0, sizeof(*var_x) * w); + memset(mean_y, 0, sizeof(*mean_y) * h); + memset(var_y, 0, sizeof(*var_y) * h); + + for (y = 0; y < h; ++y) { + for (x = 0; x < w; ++x) { + const double d = data[y * w + x]; + var_x[x] += d * d; + var_y[y] += d * d; + mean_x[x] += d; + mean_y[y] += d; + var += d * d; + mean += d; + } + } + mean /= (w * h); + var = var / (w * h) - mean * mean; + + for (y = 0; y < h; ++y) { + mean_y[y] /= h; + var_y[y] = var_y[y] / h - mean_y[y] * mean_y[y]; + if (fabs(var_y[y] - var) >= kVarianceThreshold) { + fprintf(stderr, "Variance distance too large %f %f\n", var_y[y], var); + ret_value = 0; + break; + } + if (fabs(mean_y[y] - mean) >= kMeanThreshold) { + fprintf(stderr, "Mean distance too large %f %f\n", mean_y[y], mean); + ret_value = 0; + break; + } + } + + for (x = 0; x < w; ++x) { + mean_x[x] /= w; + var_x[x] = var_x[x] / w - mean_x[x] * mean_x[x]; + if (fabs(var_x[x] - var) >= kVarianceThreshold) { + fprintf(stderr, "Variance distance too large %f %f\n", var_x[x], var); + ret_value = 0; + break; + } + if (fabs(mean_x[x] - mean) >= kMeanThreshold) { + fprintf(stderr, "Mean distance too large %f %f\n", mean_x[x], mean); + ret_value = 0; + break; + } + } + + aom_free(mean_x); + aom_free(mean_y); + aom_free(var_x); + aom_free(var_y); + + return ret_value; +} diff --git a/aom_dsp/noise_util.h b/aom_dsp/noise_util.h new file mode 100644 index 000000000..beff9886b --- /dev/null +++ b/aom_dsp/noise_util.h @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2017, Alliance for Open Media. All rights reserved + * + * This source code is subject to the terms of the BSD 2 Clause License and + * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License + * was not distributed with this source code in the LICENSE file, you can + * obtain it at www.aomedia.org/license/software. If the Alliance for Open + * Media Patent License 1.0 was not distributed with this source code in the + * PATENTS file, you can obtain it at www.aomedia.org/license/patent. + */ + +#ifndef AOM_DSP_NOISE_UTIL_H_ +#define AOM_DSP_NOISE_UTIL_H_ + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +// Computes normalized cross correlation of two vectors a and b of length n. +double aom_normalized_cross_correlation(const double *a, const double *b, + int n); + +// Synthesizes noise using the auto-regressive filter of the given lag, +// with the provided n coefficients sampled at the given coords. +void aom_noise_synth(int lag, int n, const int (*coords)[2], + const double *coeffs, double *data, int w, int h); + +// Validates the correlated noise in the data buffer of size (w, h). +int aom_noise_data_validate(const double *data, int w, int h); + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus + +#endif // AOM_DSP_NOISE_UTIL_H_ diff --git a/test/noise_model_test.cc b/test/noise_model_test.cc new file mode 100644 index 000000000..fb67762eb --- /dev/null +++ b/test/noise_model_test.cc @@ -0,0 +1,563 @@ +#include <algorithm> +#include <vector> + +#include "./aom_dsp/noise_model.h" +#include "./aom_dsp/noise_util.h" +#include "third_party/googletest/src/googletest/include/gtest/gtest.h" + +extern "C" double aom_randn(double sigma); + +TEST(NoiseStrengthSolver, GetCentersTwoBins) { + aom_noise_strength_solver_t solver; + aom_noise_strength_solver_init(&solver, 2); + EXPECT_NEAR(0, aom_noise_strength_solver_get_center(&solver, 0), 1e-5); + EXPECT_NEAR(255, aom_noise_strength_solver_get_center(&solver, 1), 1e-5); + aom_noise_strength_solver_free(&solver); +} + +TEST(NoiseStrengthSolver, GetCenters256Bins) { + const int num_bins = 256; + aom_noise_strength_solver_t solver; + aom_noise_strength_solver_init(&solver, num_bins); + + for (int i = 0; i < 256; ++i) { + EXPECT_NEAR(i, aom_noise_strength_solver_get_center(&solver, i), 1e-5); + } + aom_noise_strength_solver_free(&solver); +} + +// Tests that the noise strength solver returns the identity transform when +// given identity-like constraints. +TEST(NoiseStrengthSolver, ObserveIdentity) { + const int num_bins = 256; + aom_noise_strength_solver_t solver; + EXPECT_EQ(1, aom_noise_strength_solver_init(&solver, num_bins)); + + // We have to add a big more strength to constraints at the boundary to + // overcome any regularization. + for (int j = 0; j < 5; ++j) { + aom_noise_strength_solver_add_measurement(&solver, 0, 0); + aom_noise_strength_solver_add_measurement(&solver, 255, 255); + } + for (int i = 0; i < 256; ++i) { + aom_noise_strength_solver_add_measurement(&solver, i, i); + } + EXPECT_EQ(1, aom_noise_strength_solver_solve(&solver)); + for (int i = 2; i < num_bins - 2; ++i) { + EXPECT_NEAR(i, solver.eqns.x[i], 0.1); + } + + aom_noise_strength_lut_t lut; + EXPECT_EQ(1, aom_noise_strength_solver_fit_piecewise(&solver, &lut)); + + ASSERT_EQ(2, lut.num_points); + EXPECT_NEAR(0.0, lut.points[0][0], 1e-5); + EXPECT_NEAR(0.0, lut.points[0][1], 0.5); + EXPECT_NEAR(255.0, lut.points[1][0], 1e-5); + EXPECT_NEAR(255.0, lut.points[1][1], 0.5); + + aom_noise_strength_lut_free(&lut); + aom_noise_strength_solver_free(&solver); +} + +TEST(NoiseStrengthLut, LutEvalSinglePoint) { + aom_noise_strength_lut_t lut; + ASSERT_TRUE(aom_noise_strength_lut_init(&lut, 1)); + ASSERT_EQ(1, lut.num_points); + lut.points[0][0] = 0; + lut.points[0][1] = 1; + EXPECT_EQ(1, aom_noise_strength_lut_eval(&lut, -1)); + EXPECT_EQ(1, aom_noise_strength_lut_eval(&lut, 0)); + EXPECT_EQ(1, aom_noise_strength_lut_eval(&lut, 1)); + aom_noise_strength_lut_free(&lut); +} + +TEST(NoiseStrengthLut, LutEvalMultiPointInterp) { + const double kEps = 1e-5; + aom_noise_strength_lut_t lut; + ASSERT_TRUE(aom_noise_strength_lut_init(&lut, 4)); + ASSERT_EQ(4, lut.num_points); + + lut.points[0][0] = 0; + lut.points[0][1] = 0; + + lut.points[1][0] = 1; + lut.points[1][1] = 1; + + lut.points[2][0] = 2; + lut.points[2][1] = 1; + + lut.points[3][0] = 100; + lut.points[3][1] = 1001; + + // Test lower boundary + EXPECT_EQ(0, aom_noise_strength_lut_eval(&lut, -1)); + EXPECT_EQ(0, aom_noise_strength_lut_eval(&lut, 0)); + + // Test first part that should be identity + EXPECT_NEAR(0.25, aom_noise_strength_lut_eval(&lut, 0.25), kEps); + EXPECT_NEAR(0.75, aom_noise_strength_lut_eval(&lut, 0.75), kEps); + + // This is a constant section (should evaluate to 1) + EXPECT_NEAR(1.0, aom_noise_strength_lut_eval(&lut, 1.25), kEps); + EXPECT_NEAR(1.0, aom_noise_strength_lut_eval(&lut, 1.75), kEps); + + // Test interpolation between to non-zero y coords. + EXPECT_NEAR(1, aom_noise_strength_lut_eval(&lut, 2), kEps); + EXPECT_NEAR(251, aom_noise_strength_lut_eval(&lut, 26.5), kEps); + EXPECT_NEAR(751, aom_noise_strength_lut_eval(&lut, 75.5), kEps); + + // Test upper boundary + EXPECT_EQ(1001, aom_noise_strength_lut_eval(&lut, 100)); + EXPECT_EQ(1001, aom_noise_strength_lut_eval(&lut, 101)); + + aom_noise_strength_lut_free(&lut); +} + +TEST(NoiseModel, InitSuccessWithValidSquareShape) { + aom_noise_model_params_t params = {.shape = AOM_NOISE_SHAPE_SQUARE, + .lag = 2 }; + aom_noise_model_t model; + + EXPECT_TRUE(aom_noise_model_init(&model, params)); + + const int kNumCoords = 12; + const int kCoords[][2] = { { -2, -2 }, { -1, -2 }, { 0, -2 }, { 1, -2 }, + { 2, -2 }, { -2, -1 }, { -1, -1 }, { 0, -1 }, + { 1, -1 }, { 2, -1 }, { -2, 0 }, { -1, 0 } }; + EXPECT_EQ(kNumCoords, model.n); + for (int i = 0; i < kNumCoords; ++i) { + const int *coord = kCoords[i]; + EXPECT_EQ(coord[0], model.coords[i][0]); + EXPECT_EQ(coord[1], model.coords[i][1]); + } + aom_noise_model_free(&model); +} + +TEST(NoiseModel, InitSuccessWithValidDiamondShape) { + aom_noise_model_t model; + aom_noise_model_params_t params = {.shape = AOM_NOISE_SHAPE_DIAMOND, + .lag = 2 }; + EXPECT_TRUE(aom_noise_model_init(&model, params)); + EXPECT_EQ(6, model.n); + const int kNumCoords = 6; + const int kCoords[][2] = { { 0, -2 }, { -1, -1 }, { 0, -1 }, + { 1, -1 }, { -2, 0 }, { -1, 0 } }; + EXPECT_EQ(kNumCoords, model.n); + for (int i = 0; i < kNumCoords; ++i) { + const int *coord = kCoords[i]; + EXPECT_EQ(coord[0], model.coords[i][0]); + EXPECT_EQ(coord[1], model.coords[i][1]); + } + aom_noise_model_free(&model); +} + +TEST(NoiseModel, InitFailsWithTooLargeLag) { + aom_noise_model_t model; + aom_noise_model_params_t params = {.shape = AOM_NOISE_SHAPE_SQUARE, + .lag = 10 }; + EXPECT_FALSE(aom_noise_model_init(&model, params)); + aom_noise_model_free(&model); +} + +TEST(NoiseModel, InitFailsWithTooSmallLag) { + aom_noise_model_t model; + aom_noise_model_params_t params = {.shape = AOM_NOISE_SHAPE_SQUARE, + .lag = 0 }; + EXPECT_FALSE(aom_noise_model_init(&model, params)); + aom_noise_model_free(&model); +} + +TEST(NoiseModel, InitFailsWithInvalidShape) { + aom_noise_model_t model; + aom_noise_model_params_t params = {.shape = aom_noise_shape(100), .lag = 3 }; + EXPECT_FALSE(aom_noise_model_init(&model, params)); + aom_noise_model_free(&model); +} + +TEST(FlatBlockEstimator, ExtractBlock) { + const int kBlockSize = 16; + aom_flat_block_finder_t flat_block_finder; + ASSERT_EQ(1, aom_flat_block_finder_init(&flat_block_finder, kBlockSize)); + + // Test with an image of more than one block. + const int h = 2 * kBlockSize; + const int w = 2 * kBlockSize; + const int stride = 2 * kBlockSize; + std::vector<uint8_t> data(h * stride, 128); + + // Set up the (0,0) block to be a plane and the (0,1) block to be a + // checkerboard + for (int y = 0; y < kBlockSize; ++y) { + for (int x = 0; x < kBlockSize; ++x) { + data[y * stride + x] = -y + x + 128; + data[y * stride + x + kBlockSize] = + (x % 2 + y % 2) % 2 ? 128 - 20 : 128 + 20; + } + } + std::vector<double> block(kBlockSize * kBlockSize, 1); + std::vector<double> plane(kBlockSize * kBlockSize, 1); + + // The block data should be a constant (zero) and the rest of the plane + // trend is covered in the plane data. + aom_flat_block_finder_extract_block(&flat_block_finder, &data[0], w, h, + stride, 0, 0, &plane[0], &block[0]); + for (int y = 0; y < kBlockSize; ++y) { + for (int x = 0; x < kBlockSize; ++x) { + EXPECT_NEAR(0, block[y * kBlockSize + x], 1e-5); + EXPECT_NEAR((double)(data[y * stride + x]) / 255, + plane[y * kBlockSize + x], 1e-5); + } + } + + // The plane trend is a constant, and the block is a zero mean checkerboard. + aom_flat_block_finder_extract_block(&flat_block_finder, &data[0], w, h, + stride, kBlockSize, 0, &plane[0], + &block[0]); + for (int y = 0; y < kBlockSize; ++y) { + for (int x = 0; x < kBlockSize; ++x) { + EXPECT_NEAR(((double)data[y * stride + x + kBlockSize] - 128.0) / 255, + block[y * kBlockSize + x], 1e-5); + EXPECT_NEAR(128.0 / 255.0, plane[y * kBlockSize + x], 1e-5); + } + } + aom_flat_block_finder_free(&flat_block_finder); +} + +TEST(FlatBlockEstimator, FindFlatBlocks) { + const int kBlockSize = 32; + aom_flat_block_finder_t flat_block_finder; + ASSERT_EQ(1, aom_flat_block_finder_init(&flat_block_finder, kBlockSize)); + + const int num_blocks_w = 8; + const int h = kBlockSize; + const int w = kBlockSize * num_blocks_w; + const int stride = w; + std::vector<uint8_t> data(h * stride, 128); + std::vector<uint8_t> flat_blocks(num_blocks_w, 0); + + for (int y = 0; y < kBlockSize; ++y) { + for (int x = 0; x < kBlockSize; ++x) { + // Block 0 (not flat): constant doesn't have enough variance to qualify + data[y * stride + x + 0 * kBlockSize] = 128; + + // Block 1 (not flat): too high of variance is hard to validate as flat + data[y * stride + x + 1 * kBlockSize] = (uint8_t)(128 + aom_randn(5)); + + // Block 2 (flat): slight checkerboard added to constant + const int check = (x % 2 + y % 2) % 2 ? -2 : 2; + data[y * stride + x + 2 * kBlockSize] = 128 + check; + + // Block 3 (flat): planar block with checkerboard pattern is also flat + data[y * stride + x + 3 * kBlockSize] = y * 2 - x / 2 + 128 + check; + + // Block 4 (flat): gaussian random with standard deviation 1. + data[y * stride + x + 4 * kBlockSize] = + (uint8_t)(aom_randn(1) + x + 128.0); + + // Block 5 (flat): gaussian random with standard deviation 2. + data[y * stride + x + 5 * kBlockSize] = + (uint8_t)(aom_randn(2) + y + 128.0); + + // Block 6 (not flat): too high of directional gradient. + const int strong_edge = x > kBlockSize / 2 ? 64 : 0; + data[y * stride + x + 6 * kBlockSize] = + (uint8_t)(aom_randn(1) + strong_edge + 128.0); + + // Block 7 (not flat): too high gradient. + const int big_check = ((x >> 2) % 2 + (y >> 2) % 2) % 2 ? -16 : 16; + data[y * stride + x + 7 * kBlockSize] = + (uint8_t)(aom_randn(1) + big_check + 128.0); + } + } + + EXPECT_EQ(4, aom_flat_block_finder_run(&flat_block_finder, &data[0], w, h, + stride, &flat_blocks[0])); + + // First two blocks are not flat + EXPECT_EQ(0, flat_blocks[0]); + 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]); + + // Last 2 are not. + EXPECT_EQ(0, flat_blocks[6]); + EXPECT_EQ(0, flat_blocks[7]); + + aom_flat_block_finder_free(&flat_block_finder); +} + +class NoiseModelUpdateTest : public ::testing::Test { + public: + static const int kWidth = 128; + static const int kHeight = 128; + static const int kBlockSize = 16; + static const int kNumBlocksX = kWidth / kBlockSize; + static const int kNumBlocksY = kHeight / kBlockSize; + + void SetUp() { + const aom_noise_model_params_t params = {.shape = AOM_NOISE_SHAPE_SQUARE, + .lag = 3 }; + ASSERT_TRUE(aom_noise_model_init(&model_, params)); + + data_.resize(kWidth * kHeight * 3); + denoised_.resize(kWidth * kHeight * 3); + noise_.resize(kWidth * kHeight); + 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; + chroma_sub_[0] = 0; + chroma_sub_[1] = 0; + } + + void TearDown() { aom_noise_model_free(&model_); } + + protected: + aom_noise_model_t model_; + std::vector<uint8_t> data_; + std::vector<uint8_t> denoised_; + + std::vector<double> noise_; + std::vector<double> renoise_; + std::vector<uint8_t> flat_blocks_; + + uint8_t *data_ptr_[3]; + uint8_t *denoised_ptr_[3]; + int strides_[3]; + int chroma_sub_[2]; +}; + +TEST_F(NoiseModelUpdateTest, UpdateFailsNoFlatBlocks) { + EXPECT_EQ(AOM_NOISE_STATUS_INSUFFICIENT_FLAT_BLOCKS, + aom_noise_model_update(&model_, data_ptr_, denoised_ptr_, kWidth, + kHeight, strides_, chroma_sub_, + &flat_blocks_[0], kBlockSize)); +} + +TEST_F(NoiseModelUpdateTest, UpdateSuccessForZeroNoiseAllFlat) { + flat_blocks_.assign(flat_blocks_.size(), 1); + denoised_.assign(denoised_.size(), 128); + data_.assign(denoised_.size(), 128); + EXPECT_EQ(AOM_NOISE_STATUS_INTERNAL_ERROR, + aom_noise_model_update(&model_, data_ptr_, denoised_ptr_, kWidth, + kHeight, strides_, chroma_sub_, + &flat_blocks_[0], kBlockSize)); +} + +TEST_F(NoiseModelUpdateTest, UpdateFailsBlockSizeTooSmall) { + flat_blocks_.assign(flat_blocks_.size(), 1); + denoised_.assign(denoised_.size(), 128); + data_.assign(denoised_.size(), 128); + EXPECT_EQ( + AOM_NOISE_STATUS_INVALID_ARGUMENT, + aom_noise_model_update(&model_, data_ptr_, denoised_ptr_, kWidth, kHeight, + strides_, chroma_sub_, &flat_blocks_[0], + 6 /* block_size=2 is too small*/)); +} + +TEST_F(NoiseModelUpdateTest, UpdateSuccessForWhiteRandomNoise) { + for (int y = 0; y < kHeight; ++y) { + for (int x = 0; x < kWidth; ++x) { + data_ptr_[0][y * kWidth + x] = int(64 + y + aom_randn(1)); + denoised_ptr_[0][y * kWidth + x] = 64 + y; + // Make the chroma planes completely correlated with the Y plane + for (int c = 1; c < 3; ++c) { + data_ptr_[c][y * kWidth + x] = data_ptr_[0][y * kWidth + x]; + denoised_ptr_[c][y * kWidth + x] = denoised_ptr_[0][y * kWidth + x]; + } + } + } + flat_blocks_.assign(flat_blocks_.size(), 1); + EXPECT_EQ(AOM_NOISE_STATUS_OK, + aom_noise_model_update(&model_, data_ptr_, denoised_ptr_, kWidth, + kHeight, strides_, chroma_sub_, + &flat_blocks_[0], kBlockSize)); + + const double kCoeffEps = 0.075; + const int n = model_.n; + for (int c = 0; c < 3; ++c) { + for (int i = 0; i < n; ++i) { + EXPECT_NEAR(0, model_.latest_state[c].eqns.x[i], kCoeffEps); + EXPECT_NEAR(0, model_.combined_state[c].eqns.x[i], kCoeffEps); + } + // The second and third channels are highly correlated with the first. + if (c > 0) { + ASSERT_EQ(n + 1, model_.latest_state[c].eqns.n); + ASSERT_EQ(n + 1, model_.combined_state[c].eqns.n); + + EXPECT_NEAR(1, model_.latest_state[c].eqns.x[n], kCoeffEps); + EXPECT_NEAR(1, model_.combined_state[c].eqns.x[n], kCoeffEps); + } + } + + // The fitted noise strength should be close to the standard deviation + // for all intensity bins. + const double kStdEps = 0.1; + for (int i = 0; i < model_.latest_state[0].strength_solver.eqns.n; ++i) { + EXPECT_NEAR(1.0, model_.latest_state[0].strength_solver.eqns.x[i], kStdEps); + EXPECT_NEAR(1.0, model_.combined_state[0].strength_solver.eqns.x[i], + kStdEps); + } + + aom_noise_strength_lut_t lut; + aom_noise_strength_solver_fit_piecewise( + &model_.latest_state[0].strength_solver, &lut); + ASSERT_EQ(2, lut.num_points); + EXPECT_NEAR(0.0, lut.points[0][0], 1e-5); + EXPECT_NEAR(1.0, lut.points[0][1], kStdEps); + EXPECT_NEAR(255.0, lut.points[1][0], 1e-5); + EXPECT_NEAR(1.0, lut.points[1][1], kStdEps); + aom_noise_strength_lut_free(&lut); +} + +TEST_F(NoiseModelUpdateTest, UpdateSuccessForScaledWhiteNoise) { + const double kCoeffEps = 0.05; + const double kLowStd = 1; + const double kHighStd = 4; + for (int y = 0; y < kHeight; ++y) { + for (int x = 0; x < kWidth; ++x) { + for (int c = 0; c < 3; ++c) { + // The image data is bimodal: + // Bottom half has low intensity and low noise strength + // Top half has high intensity and high noise strength + const int avg = (y < kHeight / 2) ? 4 : 245; + const double std = (y < kHeight / 2) ? kLowStd : kHighStd; + data_ptr_[c][y * kWidth + x] = + (uint8_t)std::min((int)255, (int)(2 + avg + aom_randn(std))); + denoised_ptr_[c][y * kWidth + x] = 2 + avg; + } + } + } + // Label all blocks as flat for the update + flat_blocks_.assign(flat_blocks_.size(), 1); + EXPECT_EQ(AOM_NOISE_STATUS_OK, + aom_noise_model_update(&model_, data_ptr_, denoised_ptr_, kWidth, + kHeight, strides_, chroma_sub_, + &flat_blocks_[0], kBlockSize)); + + const int n = model_.n; + // The noise is uncorrelated spatially and with the y channel. + // All coefficients should be reasonably close to zero. + for (int c = 0; c < 3; ++c) { + for (int i = 0; i < n; ++i) { + EXPECT_NEAR(0, model_.latest_state[c].eqns.x[i], kCoeffEps); + EXPECT_NEAR(0, model_.combined_state[c].eqns.x[i], kCoeffEps); + } + if (c > 0) { + ASSERT_EQ(n + 1, model_.latest_state[c].eqns.n); + ASSERT_EQ(n + 1, model_.combined_state[c].eqns.n); + + // The correlation to the y channel should be low (near zero) + EXPECT_NEAR(0, model_.latest_state[c].eqns.x[n], kCoeffEps); + EXPECT_NEAR(0, model_.combined_state[c].eqns.x[n], kCoeffEps); + } + } + + // Noise strength should vary between kLowStd and kHighStd. + const double kStdEps = 0.15; + ASSERT_EQ(20, model_.latest_state[0].strength_solver.eqns.n); + for (int i = 0; i < model_.latest_state[0].strength_solver.eqns.n; ++i) { + const double a = i / 19.0; + const double expected = (kLowStd * (1.0 - a) + kHighStd * a); + EXPECT_NEAR(expected, model_.latest_state[0].strength_solver.eqns.x[i], + kStdEps); + EXPECT_NEAR(expected, model_.combined_state[0].strength_solver.eqns.x[i], + kStdEps); + } + + // If we fit a piecewise linear model, there should be two points: + // one near kLowStd at 0, and the other near kHighStd and 255. + aom_noise_strength_lut_t lut; + aom_noise_strength_solver_fit_piecewise( + &model_.latest_state[0].strength_solver, &lut); + ASSERT_EQ(2, lut.num_points); + EXPECT_NEAR(0, lut.points[0][0], 1e-4); + EXPECT_NEAR(kLowStd, lut.points[0][1], kStdEps); + EXPECT_NEAR(255.0, lut.points[1][0], 1e-5); + EXPECT_NEAR(kHighStd, lut.points[1][1], kStdEps); + aom_noise_strength_lut_free(&lut); +} + +TEST_F(NoiseModelUpdateTest, UpdateSuccessForCorrelatedNoise) { + const int kNumCoeffs = 24; + 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, + }; + ASSERT_EQ(model_.n, kNumCoeffs); + aom_noise_synth(model_.params.lag, model_.n, model_.coords, kCoeffs, + &noise_[0], kWidth, kHeight); + 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) { + 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); + } + } + } + EXPECT_EQ(AOM_NOISE_STATUS_OK, + aom_noise_model_update(&model_, data_ptr_, denoised_ptr_, kWidth, + kHeight, strides_, chroma_sub_, + &flat_blocks_[0], kBlockSize)); + + // 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); + } + // 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); + } + + // 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); + } +} diff --git a/test/test.cmake b/test/test.cmake index 002f03c9d..152243611 100644 --- a/test/test.cmake +++ b/test/test.cmake @@ -225,6 +225,7 @@ if (CONFIG_AV1_ENCODER) "${AOM_ROOT}/test/masked_sad_test.cc" "${AOM_ROOT}/test/masked_variance_test.cc" "${AOM_ROOT}/test/minmax_test.cc" + "${AOM_ROOT}/test/noise_model_test.cc" "${AOM_ROOT}/test/subtract_test.cc" "${AOM_ROOT}/test/sum_squares_test.cc" "${AOM_ROOT}/test/variance_test.cc") |