Coverage Report

Created: 2022-08-24 06:15

/src/aom/aom_dsp/noise_model.c
Line
Count
Source (jump to first uncovered line)
1
/*
2
 * Copyright (c) 2017, Alliance for Open Media. All rights reserved
3
 *
4
 * This source code is subject to the terms of the BSD 2 Clause License and
5
 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
6
 * was not distributed with this source code in the LICENSE file, you can
7
 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
8
 * Media Patent License 1.0 was not distributed with this source code in the
9
 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
10
 */
11
12
#include <math.h>
13
#include <stdio.h>
14
#include <stdlib.h>
15
#include <string.h>
16
17
#include "aom_dsp/aom_dsp_common.h"
18
#include "aom_dsp/mathutils.h"
19
#include "aom_dsp/noise_model.h"
20
#include "aom_dsp/noise_util.h"
21
#include "aom_mem/aom_mem.h"
22
23
0
#define kLowPolyNumParams 3
24
25
static const int kMaxLag = 4;
26
27
// Defines a function that can be used to obtain the mean of a block for the
28
// provided data type (uint8_t, or uint16_t)
29
#define GET_BLOCK_MEAN(INT_TYPE, suffix)                                    \
30
  static double get_block_mean_##suffix(const INT_TYPE *data, int w, int h, \
31
                                        int stride, int x_o, int y_o,       \
32
0
                                        int block_size) {                   \
33
0
    const int max_h = AOMMIN(h - y_o, block_size);                          \
34
0
    const int max_w = AOMMIN(w - x_o, block_size);                          \
35
0
    double block_mean = 0;                                                  \
36
0
    for (int y = 0; y < max_h; ++y) {                                       \
37
0
      for (int x = 0; x < max_w; ++x) {                                     \
38
0
        block_mean += data[(y_o + y) * stride + x_o + x];                   \
39
0
      }                                                                     \
40
0
    }                                                                       \
41
0
    return block_mean / (max_w * max_h);                                    \
42
0
  }
Unexecuted instantiation: noise_model.c:get_block_mean_highbd
Unexecuted instantiation: noise_model.c:get_block_mean_lowbd
43
44
GET_BLOCK_MEAN(uint8_t, lowbd);
45
GET_BLOCK_MEAN(uint16_t, highbd);
46
47
static INLINE double get_block_mean(const uint8_t *data, int w, int h,
48
                                    int stride, int x_o, int y_o,
49
0
                                    int block_size, int use_highbd) {
50
0
  if (use_highbd)
51
0
    return get_block_mean_highbd((const uint16_t *)data, w, h, stride, x_o, y_o,
52
0
                                 block_size);
53
0
  return get_block_mean_lowbd(data, w, h, stride, x_o, y_o, block_size);
54
0
}
55
56
// Defines a function that can be used to obtain the variance of a block
57
// for the provided data type (uint8_t, or uint16_t)
58
#define GET_NOISE_VAR(INT_TYPE, suffix)                                  \
59
  static double get_noise_var_##suffix(                                  \
60
      const INT_TYPE *data, const INT_TYPE *denoised, int stride, int w, \
61
0
      int h, int x_o, int y_o, int block_size_x, int block_size_y) {     \
62
0
    const int max_h = AOMMIN(h - y_o, block_size_y);                     \
63
0
    const int max_w = AOMMIN(w - x_o, block_size_x);                     \
64
0
    double noise_var = 0;                                                \
65
0
    double noise_mean = 0;                                               \
66
0
    for (int y = 0; y < max_h; ++y) {                                    \
67
0
      for (int x = 0; x < max_w; ++x) {                                  \
68
0
        double noise = (double)data[(y_o + y) * stride + x_o + x] -      \
69
0
                       denoised[(y_o + y) * stride + x_o + x];           \
70
0
        noise_mean += noise;                                             \
71
0
        noise_var += noise * noise;                                      \
72
0
      }                                                                  \
73
0
    }                                                                    \
74
0
    noise_mean /= (max_w * max_h);                                       \
75
0
    return noise_var / (max_w * max_h) - noise_mean * noise_mean;        \
76
0
  }
Unexecuted instantiation: noise_model.c:get_noise_var_highbd
Unexecuted instantiation: noise_model.c:get_noise_var_lowbd
77
78
GET_NOISE_VAR(uint8_t, lowbd);
79
GET_NOISE_VAR(uint16_t, highbd);
80
81
static INLINE double get_noise_var(const uint8_t *data, const uint8_t *denoised,
82
                                   int w, int h, int stride, int x_o, int y_o,
83
                                   int block_size_x, int block_size_y,
84
0
                                   int use_highbd) {
85
0
  if (use_highbd)
86
0
    return get_noise_var_highbd((const uint16_t *)data,
87
0
                                (const uint16_t *)denoised, w, h, stride, x_o,
88
0
                                y_o, block_size_x, block_size_y);
89
0
  return get_noise_var_lowbd(data, denoised, w, h, stride, x_o, y_o,
90
0
                             block_size_x, block_size_y);
91
0
}
92
93
0
static void equation_system_clear(aom_equation_system_t *eqns) {
94
0
  const int n = eqns->n;
95
0
  memset(eqns->A, 0, sizeof(*eqns->A) * n * n);
96
0
  memset(eqns->x, 0, sizeof(*eqns->x) * n);
97
0
  memset(eqns->b, 0, sizeof(*eqns->b) * n);
98
0
}
99
100
static void equation_system_copy(aom_equation_system_t *dst,
101
0
                                 const aom_equation_system_t *src) {
102
0
  const int n = dst->n;
103
0
  memcpy(dst->A, src->A, sizeof(*dst->A) * n * n);
104
0
  memcpy(dst->x, src->x, sizeof(*dst->x) * n);
105
0
  memcpy(dst->b, src->b, sizeof(*dst->b) * n);
106
0
}
107
108
0
static int equation_system_init(aom_equation_system_t *eqns, int n) {
109
0
  eqns->A = (double *)aom_malloc(sizeof(*eqns->A) * n * n);
110
0
  eqns->b = (double *)aom_malloc(sizeof(*eqns->b) * n);
111
0
  eqns->x = (double *)aom_malloc(sizeof(*eqns->x) * n);
112
0
  eqns->n = n;
113
0
  if (!eqns->A || !eqns->b || !eqns->x) {
114
0
    fprintf(stderr, "Failed to allocate system of equations of size %d\n", n);
115
0
    aom_free(eqns->A);
116
0
    aom_free(eqns->b);
117
0
    aom_free(eqns->x);
118
0
    memset(eqns, 0, sizeof(*eqns));
119
0
    return 0;
120
0
  }
121
0
  equation_system_clear(eqns);
122
0
  return 1;
123
0
}
124
125
0
static int equation_system_solve(aom_equation_system_t *eqns) {
126
0
  const int n = eqns->n;
127
0
  double *b = (double *)aom_malloc(sizeof(*b) * n);
128
0
  double *A = (double *)aom_malloc(sizeof(*A) * n * n);
129
0
  int ret = 0;
130
0
  if (A == NULL || b == NULL) {
131
0
    fprintf(stderr, "Unable to allocate temp values of size %dx%d\n", n, n);
132
0
    aom_free(b);
133
0
    aom_free(A);
134
0
    return 0;
135
0
  }
136
0
  memcpy(A, eqns->A, sizeof(*eqns->A) * n * n);
137
0
  memcpy(b, eqns->b, sizeof(*eqns->b) * n);
138
0
  ret = linsolve(n, A, eqns->n, b, eqns->x);
139
0
  aom_free(b);
140
0
  aom_free(A);
141
142
0
  if (ret == 0) {
143
0
    return 0;
144
0
  }
145
0
  return 1;
146
0
}
147
148
static void equation_system_add(aom_equation_system_t *dest,
149
0
                                aom_equation_system_t *src) {
150
0
  const int n = dest->n;
151
0
  int i, j;
152
0
  for (i = 0; i < n; ++i) {
153
0
    for (j = 0; j < n; ++j) {
154
0
      dest->A[i * n + j] += src->A[i * n + j];
155
0
    }
156
0
    dest->b[i] += src->b[i];
157
0
  }
158
0
}
159
160
0
static void equation_system_free(aom_equation_system_t *eqns) {
161
0
  if (!eqns) return;
162
0
  aom_free(eqns->A);
163
0
  aom_free(eqns->b);
164
0
  aom_free(eqns->x);
165
0
  memset(eqns, 0, sizeof(*eqns));
166
0
}
167
168
0
static void noise_strength_solver_clear(aom_noise_strength_solver_t *solver) {
169
0
  equation_system_clear(&solver->eqns);
170
0
  solver->num_equations = 0;
171
0
  solver->total = 0;
172
0
}
173
174
static void noise_strength_solver_add(aom_noise_strength_solver_t *dest,
175
0
                                      aom_noise_strength_solver_t *src) {
176
0
  equation_system_add(&dest->eqns, &src->eqns);
177
0
  dest->num_equations += src->num_equations;
178
0
  dest->total += src->total;
179
0
}
180
181
// Return the number of coefficients required for the given parameters
182
0
static int num_coeffs(const aom_noise_model_params_t params) {
183
0
  const int n = 2 * params.lag + 1;
184
0
  switch (params.shape) {
185
0
    case AOM_NOISE_SHAPE_DIAMOND: return params.lag * (params.lag + 1);
186
0
    case AOM_NOISE_SHAPE_SQUARE: return (n * n) / 2;
187
0
  }
188
0
  return 0;
189
0
}
190
191
0
static int noise_state_init(aom_noise_state_t *state, int n, int bit_depth) {
192
0
  const int kNumBins = 20;
193
0
  if (!equation_system_init(&state->eqns, n)) {
194
0
    fprintf(stderr, "Failed initialization noise state with size %d\n", n);
195
0
    return 0;
196
0
  }
197
0
  state->ar_gain = 1.0;
198
0
  state->num_observations = 0;
199
0
  return aom_noise_strength_solver_init(&state->strength_solver, kNumBins,
200
0
                                        bit_depth);
201
0
}
202
203
0
static void set_chroma_coefficient_fallback_soln(aom_equation_system_t *eqns) {
204
0
  const double kTolerance = 1e-6;
205
0
  const int last = eqns->n - 1;
206
  // Set all of the AR coefficients to zero, but try to solve for correlation
207
  // with the luma channel
208
0
  memset(eqns->x, 0, sizeof(*eqns->x) * eqns->n);
209
0
  if (fabs(eqns->A[last * eqns->n + last]) > kTolerance) {
210
0
    eqns->x[last] = eqns->b[last] / eqns->A[last * eqns->n + last];
211
0
  }
212
0
}
213
214
0
int aom_noise_strength_lut_init(aom_noise_strength_lut_t *lut, int num_points) {
215
0
  if (!lut) return 0;
216
0
  if (num_points <= 0) return 0;
217
0
  lut->num_points = 0;
218
0
  lut->points = (double(*)[2])aom_malloc(num_points * sizeof(*lut->points));
219
0
  if (!lut->points) return 0;
220
0
  lut->num_points = num_points;
221
0
  memset(lut->points, 0, sizeof(*lut->points) * num_points);
222
0
  return 1;
223
0
}
224
225
0
void aom_noise_strength_lut_free(aom_noise_strength_lut_t *lut) {
226
0
  if (!lut) return;
227
0
  aom_free(lut->points);
228
0
  memset(lut, 0, sizeof(*lut));
229
0
}
230
231
double aom_noise_strength_lut_eval(const aom_noise_strength_lut_t *lut,
232
0
                                   double x) {
233
0
  int i = 0;
234
  // Constant extrapolation for x <  x_0.
235
0
  if (x < lut->points[0][0]) return lut->points[0][1];
236
0
  for (i = 0; i < lut->num_points - 1; ++i) {
237
0
    if (x >= lut->points[i][0] && x <= lut->points[i + 1][0]) {
238
0
      const double a =
239
0
          (x - lut->points[i][0]) / (lut->points[i + 1][0] - lut->points[i][0]);
240
0
      return lut->points[i + 1][1] * a + lut->points[i][1] * (1.0 - a);
241
0
    }
242
0
  }
243
  // Constant extrapolation for x > x_{n-1}
244
0
  return lut->points[lut->num_points - 1][1];
245
0
}
246
247
static double noise_strength_solver_get_bin_index(
248
0
    const aom_noise_strength_solver_t *solver, double value) {
249
0
  const double val =
250
0
      fclamp(value, solver->min_intensity, solver->max_intensity);
251
0
  const double range = solver->max_intensity - solver->min_intensity;
252
0
  return (solver->num_bins - 1) * (val - solver->min_intensity) / range;
253
0
}
254
255
static double noise_strength_solver_get_value(
256
0
    const aom_noise_strength_solver_t *solver, double x) {
257
0
  const double bin = noise_strength_solver_get_bin_index(solver, x);
258
0
  const int bin_i0 = (int)floor(bin);
259
0
  const int bin_i1 = AOMMIN(solver->num_bins - 1, bin_i0 + 1);
260
0
  const double a = bin - bin_i0;
261
0
  return (1.0 - a) * solver->eqns.x[bin_i0] + a * solver->eqns.x[bin_i1];
262
0
}
263
264
void aom_noise_strength_solver_add_measurement(
265
0
    aom_noise_strength_solver_t *solver, double block_mean, double noise_std) {
266
0
  const double bin = noise_strength_solver_get_bin_index(solver, block_mean);
267
0
  const int bin_i0 = (int)floor(bin);
268
0
  const int bin_i1 = AOMMIN(solver->num_bins - 1, bin_i0 + 1);
269
0
  const double a = bin - bin_i0;
270
0
  const int n = solver->num_bins;
271
0
  solver->eqns.A[bin_i0 * n + bin_i0] += (1.0 - a) * (1.0 - a);
272
0
  solver->eqns.A[bin_i1 * n + bin_i0] += a * (1.0 - a);
273
0
  solver->eqns.A[bin_i1 * n + bin_i1] += a * a;
274
0
  solver->eqns.A[bin_i0 * n + bin_i1] += a * (1.0 - a);
275
0
  solver->eqns.b[bin_i0] += (1.0 - a) * noise_std;
276
0
  solver->eqns.b[bin_i1] += a * noise_std;
277
0
  solver->total += noise_std;
278
0
  solver->num_equations++;
279
0
}
280
281
0
int aom_noise_strength_solver_solve(aom_noise_strength_solver_t *solver) {
282
  // Add regularization proportional to the number of constraints
283
0
  const int n = solver->num_bins;
284
0
  const double kAlpha = 2.0 * (double)(solver->num_equations) / n;
285
0
  int result = 0;
286
0
  double mean = 0;
287
288
  // Do this in a non-destructive manner so it is not confusing to the caller
289
0
  double *old_A = solver->eqns.A;
290
0
  double *A = (double *)aom_malloc(sizeof(*A) * n * n);
291
0
  if (!A) {
292
0
    fprintf(stderr, "Unable to allocate copy of A\n");
293
0
    return 0;
294
0
  }
295
0
  memcpy(A, old_A, sizeof(*A) * n * n);
296
297
0
  for (int i = 0; i < n; ++i) {
298
0
    const int i_lo = AOMMAX(0, i - 1);
299
0
    const int i_hi = AOMMIN(n - 1, i + 1);
300
0
    A[i * n + i_lo] -= kAlpha;
301
0
    A[i * n + i] += 2 * kAlpha;
302
0
    A[i * n + i_hi] -= kAlpha;
303
0
  }
304
305
  // Small regularization to give average noise strength
306
0
  mean = solver->total / solver->num_equations;
307
0
  for (int i = 0; i < n; ++i) {
308
0
    A[i * n + i] += 1.0 / 8192.;
309
0
    solver->eqns.b[i] += mean / 8192.;
310
0
  }
311
0
  solver->eqns.A = A;
312
0
  result = equation_system_solve(&solver->eqns);
313
0
  solver->eqns.A = old_A;
314
315
0
  aom_free(A);
316
0
  return result;
317
0
}
318
319
int aom_noise_strength_solver_init(aom_noise_strength_solver_t *solver,
320
0
                                   int num_bins, int bit_depth) {
321
0
  if (!solver) return 0;
322
0
  memset(solver, 0, sizeof(*solver));
323
0
  solver->num_bins = num_bins;
324
0
  solver->min_intensity = 0;
325
0
  solver->max_intensity = (1 << bit_depth) - 1;
326
0
  solver->total = 0;
327
0
  solver->num_equations = 0;
328
0
  return equation_system_init(&solver->eqns, num_bins);
329
0
}
330
331
0
void aom_noise_strength_solver_free(aom_noise_strength_solver_t *solver) {
332
0
  if (!solver) return;
333
0
  equation_system_free(&solver->eqns);
334
0
}
335
336
double aom_noise_strength_solver_get_center(
337
0
    const aom_noise_strength_solver_t *solver, int i) {
338
0
  const double range = solver->max_intensity - solver->min_intensity;
339
0
  const int n = solver->num_bins;
340
0
  return ((double)i) / (n - 1) * range + solver->min_intensity;
341
0
}
342
343
// Computes the residual if a point were to be removed from the lut. This is
344
// calculated as the area between the output of the solver and the line segment
345
// that would be formed between [x_{i - 1}, x_{i + 1}).
346
static void update_piecewise_linear_residual(
347
    const aom_noise_strength_solver_t *solver,
348
0
    const aom_noise_strength_lut_t *lut, double *residual, int start, int end) {
349
0
  const double dx = 255. / solver->num_bins;
350
0
  for (int i = AOMMAX(start, 1); i < AOMMIN(end, lut->num_points - 1); ++i) {
351
0
    const int lower = AOMMAX(0, (int)floor(noise_strength_solver_get_bin_index(
352
0
                                    solver, lut->points[i - 1][0])));
353
0
    const int upper = AOMMIN(solver->num_bins - 1,
354
0
                             (int)ceil(noise_strength_solver_get_bin_index(
355
0
                                 solver, lut->points[i + 1][0])));
356
0
    double r = 0;
357
0
    for (int j = lower; j <= upper; ++j) {
358
0
      const double x = aom_noise_strength_solver_get_center(solver, j);
359
0
      if (x < lut->points[i - 1][0]) continue;
360
0
      if (x >= lut->points[i + 1][0]) continue;
361
0
      const double y = solver->eqns.x[j];
362
0
      const double a = (x - lut->points[i - 1][0]) /
363
0
                       (lut->points[i + 1][0] - lut->points[i - 1][0]);
364
0
      const double estimate_y =
365
0
          lut->points[i - 1][1] * (1.0 - a) + lut->points[i + 1][1] * a;
366
0
      r += fabs(y - estimate_y);
367
0
    }
368
0
    residual[i] = r * dx;
369
0
  }
370
0
}
371
372
int aom_noise_strength_solver_fit_piecewise(
373
    const aom_noise_strength_solver_t *solver, int max_output_points,
374
0
    aom_noise_strength_lut_t *lut) {
375
  // The tolerance is normalized to be give consistent results between
376
  // different bit-depths.
377
0
  const double kTolerance = solver->max_intensity * 0.00625 / 255.0;
378
0
  if (!aom_noise_strength_lut_init(lut, solver->num_bins)) {
379
0
    fprintf(stderr, "Failed to init lut\n");
380
0
    return 0;
381
0
  }
382
0
  for (int i = 0; i < solver->num_bins; ++i) {
383
0
    lut->points[i][0] = aom_noise_strength_solver_get_center(solver, i);
384
0
    lut->points[i][1] = solver->eqns.x[i];
385
0
  }
386
0
  if (max_output_points < 0) {
387
0
    max_output_points = solver->num_bins;
388
0
  }
389
390
0
  double *residual = aom_malloc(solver->num_bins * sizeof(*residual));
391
0
  memset(residual, 0, sizeof(*residual) * solver->num_bins);
392
393
0
  update_piecewise_linear_residual(solver, lut, residual, 0, solver->num_bins);
394
395
  // Greedily remove points if there are too many or if it doesn't hurt local
396
  // approximation (never remove the end points)
397
0
  while (lut->num_points > 2) {
398
0
    int min_index = 1;
399
0
    for (int j = 1; j < lut->num_points - 1; ++j) {
400
0
      if (residual[j] < residual[min_index]) {
401
0
        min_index = j;
402
0
      }
403
0
    }
404
0
    const double dx =
405
0
        lut->points[min_index + 1][0] - lut->points[min_index - 1][0];
406
0
    const double avg_residual = residual[min_index] / dx;
407
0
    if (lut->num_points <= max_output_points && avg_residual > kTolerance) {
408
0
      break;
409
0
    }
410
411
0
    const int num_remaining = lut->num_points - min_index - 1;
412
0
    memmove(lut->points + min_index, lut->points + min_index + 1,
413
0
            sizeof(lut->points[0]) * num_remaining);
414
0
    lut->num_points--;
415
416
0
    update_piecewise_linear_residual(solver, lut, residual, min_index - 1,
417
0
                                     min_index + 1);
418
0
  }
419
0
  aom_free(residual);
420
0
  return 1;
421
0
}
422
423
int aom_flat_block_finder_init(aom_flat_block_finder_t *block_finder,
424
0
                               int block_size, int bit_depth, int use_highbd) {
425
0
  const int n = block_size * block_size;
426
0
  aom_equation_system_t eqns;
427
0
  double *AtA_inv = 0;
428
0
  double *A = 0;
429
0
  int x = 0, y = 0, i = 0, j = 0;
430
0
  block_finder->A = NULL;
431
0
  block_finder->AtA_inv = NULL;
432
433
0
  if (!equation_system_init(&eqns, kLowPolyNumParams)) {
434
0
    fprintf(stderr, "Failed to init equation system for block_size=%d\n",
435
0
            block_size);
436
0
    return 0;
437
0
  }
438
439
0
  AtA_inv = (double *)aom_malloc(kLowPolyNumParams * kLowPolyNumParams *
440
0
                                 sizeof(*AtA_inv));
441
0
  A = (double *)aom_malloc(kLowPolyNumParams * n * sizeof(*A));
442
0
  if (AtA_inv == NULL || A == NULL) {
443
0
    fprintf(stderr, "Failed to alloc A or AtA_inv for block_size=%d\n",
444
0
            block_size);
445
0
    aom_free(AtA_inv);
446
0
    aom_free(A);
447
0
    equation_system_free(&eqns);
448
0
    return 0;
449
0
  }
450
451
0
  block_finder->A = A;
452
0
  block_finder->AtA_inv = AtA_inv;
453
0
  block_finder->block_size = block_size;
454
0
  block_finder->normalization = (1 << bit_depth) - 1;
455
0
  block_finder->use_highbd = use_highbd;
456
457
0
  for (y = 0; y < block_size; ++y) {
458
0
    const double yd = ((double)y - block_size / 2.) / (block_size / 2.);
459
0
    for (x = 0; x < block_size; ++x) {
460
0
      const double xd = ((double)x - block_size / 2.) / (block_size / 2.);
461
0
      const double coords[3] = { yd, xd, 1 };
462
0
      const int row = y * block_size + x;
463
0
      A[kLowPolyNumParams * row + 0] = yd;
464
0
      A[kLowPolyNumParams * row + 1] = xd;
465
0
      A[kLowPolyNumParams * row + 2] = 1;
466
467
0
      for (i = 0; i < kLowPolyNumParams; ++i) {
468
0
        for (j = 0; j < kLowPolyNumParams; ++j) {
469
0
          eqns.A[kLowPolyNumParams * i + j] += coords[i] * coords[j];
470
0
        }
471
0
      }
472
0
    }
473
0
  }
474
475
  // Lazy inverse using existing equation solver.
476
0
  for (i = 0; i < kLowPolyNumParams; ++i) {
477
0
    memset(eqns.b, 0, sizeof(*eqns.b) * kLowPolyNumParams);
478
0
    eqns.b[i] = 1;
479
0
    equation_system_solve(&eqns);
480
481
0
    for (j = 0; j < kLowPolyNumParams; ++j) {
482
0
      AtA_inv[j * kLowPolyNumParams + i] = eqns.x[j];
483
0
    }
484
0
  }
485
0
  equation_system_free(&eqns);
486
0
  return 1;
487
0
}
488
489
0
void aom_flat_block_finder_free(aom_flat_block_finder_t *block_finder) {
490
0
  if (!block_finder) return;
491
0
  aom_free(block_finder->A);
492
0
  aom_free(block_finder->AtA_inv);
493
0
  memset(block_finder, 0, sizeof(*block_finder));
494
0
}
495
496
void aom_flat_block_finder_extract_block(
497
    const aom_flat_block_finder_t *block_finder, const uint8_t *const data,
498
    int w, int h, int stride, int offsx, int offsy, double *plane,
499
0
    double *block) {
500
0
  const int block_size = block_finder->block_size;
501
0
  const int n = block_size * block_size;
502
0
  const double *A = block_finder->A;
503
0
  const double *AtA_inv = block_finder->AtA_inv;
504
0
  double plane_coords[kLowPolyNumParams];
505
0
  double AtA_inv_b[kLowPolyNumParams];
506
0
  int xi, yi, i;
507
508
0
  if (block_finder->use_highbd) {
509
0
    const uint16_t *const data16 = (const uint16_t *const)data;
510
0
    for (yi = 0; yi < block_size; ++yi) {
511
0
      const int y = clamp(offsy + yi, 0, h - 1);
512
0
      for (xi = 0; xi < block_size; ++xi) {
513
0
        const int x = clamp(offsx + xi, 0, w - 1);
514
0
        block[yi * block_size + xi] =
515
0
            ((double)data16[y * stride + x]) / block_finder->normalization;
516
0
      }
517
0
    }
518
0
  } else {
519
0
    for (yi = 0; yi < block_size; ++yi) {
520
0
      const int y = clamp(offsy + yi, 0, h - 1);
521
0
      for (xi = 0; xi < block_size; ++xi) {
522
0
        const int x = clamp(offsx + xi, 0, w - 1);
523
0
        block[yi * block_size + xi] =
524
0
            ((double)data[y * stride + x]) / block_finder->normalization;
525
0
      }
526
0
    }
527
0
  }
528
0
  multiply_mat(block, A, AtA_inv_b, 1, n, kLowPolyNumParams);
529
0
  multiply_mat(AtA_inv, AtA_inv_b, plane_coords, kLowPolyNumParams,
530
0
               kLowPolyNumParams, 1);
531
0
  multiply_mat(A, plane_coords, plane, n, kLowPolyNumParams, 1);
532
533
0
  for (i = 0; i < n; ++i) {
534
0
    block[i] -= plane[i];
535
0
  }
536
0
}
537
538
typedef struct {
539
  int index;
540
  float score;
541
} index_and_score_t;
542
543
0
static int compare_scores(const void *a, const void *b) {
544
0
  const float diff =
545
0
      ((index_and_score_t *)a)->score - ((index_and_score_t *)b)->score;
546
0
  if (diff < 0)
547
0
    return -1;
548
0
  else if (diff > 0)
549
0
    return 1;
550
0
  return 0;
551
0
}
552
553
int aom_flat_block_finder_run(const aom_flat_block_finder_t *block_finder,
554
                              const uint8_t *const data, int w, int h,
555
0
                              int stride, uint8_t *flat_blocks) {
556
  // The gradient-based features used in this code are based on:
557
  //  A. Kokaram, D. Kelly, H. Denman and A. Crawford, "Measuring noise
558
  //  correlation for improved video denoising," 2012 19th, ICIP.
559
  // The thresholds are more lenient to allow for correct grain modeling
560
  // if extreme cases.
561
0
  const int block_size = block_finder->block_size;
562
0
  const int n = block_size * block_size;
563
0
  const double kTraceThreshold = 0.15 / (32 * 32);
564
0
  const double kRatioThreshold = 1.25;
565
0
  const double kNormThreshold = 0.08 / (32 * 32);
566
0
  const double kVarThreshold = 0.005 / (double)n;
567
0
  const int num_blocks_w = (w + block_size - 1) / block_size;
568
0
  const int num_blocks_h = (h + block_size - 1) / block_size;
569
0
  int num_flat = 0;
570
0
  int bx = 0, by = 0;
571
0
  double *plane = (double *)aom_malloc(n * sizeof(*plane));
572
0
  double *block = (double *)aom_malloc(n * sizeof(*block));
573
0
  index_and_score_t *scores = (index_and_score_t *)aom_malloc(
574
0
      num_blocks_w * num_blocks_h * sizeof(*scores));
575
0
  if (plane == NULL || block == NULL || scores == NULL) {
576
0
    fprintf(stderr, "Failed to allocate memory for block of size %d\n", n);
577
0
    aom_free(plane);
578
0
    aom_free(block);
579
0
    aom_free(scores);
580
0
    return -1;
581
0
  }
582
583
#ifdef NOISE_MODEL_LOG_SCORE
584
  fprintf(stderr, "score = [");
585
#endif
586
0
  for (by = 0; by < num_blocks_h; ++by) {
587
0
    for (bx = 0; bx < num_blocks_w; ++bx) {
588
      // Compute gradient covariance matrix.
589
0
      double Gxx = 0, Gxy = 0, Gyy = 0;
590
0
      double var = 0;
591
0
      double mean = 0;
592
0
      int xi, yi;
593
0
      aom_flat_block_finder_extract_block(block_finder, data, w, h, stride,
594
0
                                          bx * block_size, by * block_size,
595
0
                                          plane, block);
596
597
0
      for (yi = 1; yi < block_size - 1; ++yi) {
598
0
        for (xi = 1; xi < block_size - 1; ++xi) {
599
0
          const double gx = (block[yi * block_size + xi + 1] -
600
0
                             block[yi * block_size + xi - 1]) /
601
0
                            2;
602
0
          const double gy = (block[yi * block_size + xi + block_size] -
603
0
                             block[yi * block_size + xi - block_size]) /
604
0
                            2;
605
0
          Gxx += gx * gx;
606
0
          Gxy += gx * gy;
607
0
          Gyy += gy * gy;
608
609
0
          mean += block[yi * block_size + xi];
610
0
          var += block[yi * block_size + xi] * block[yi * block_size + xi];
611
0
        }
612
0
      }
613
0
      mean /= (block_size - 2) * (block_size - 2);
614
615
      // Normalize gradients by block_size.
616
0
      Gxx /= ((block_size - 2) * (block_size - 2));
617
0
      Gxy /= ((block_size - 2) * (block_size - 2));
618
0
      Gyy /= ((block_size - 2) * (block_size - 2));
619
0
      var = var / ((block_size - 2) * (block_size - 2)) - mean * mean;
620
621
0
      {
622
0
        const double trace = Gxx + Gyy;
623
0
        const double det = Gxx * Gyy - Gxy * Gxy;
624
0
        const double e1 = (trace + sqrt(trace * trace - 4 * det)) / 2.;
625
0
        const double e2 = (trace - sqrt(trace * trace - 4 * det)) / 2.;
626
0
        const double norm = e1;  // Spectral norm
627
0
        const double ratio = (e1 / AOMMAX(e2, 1e-6));
628
0
        const int is_flat = (trace < kTraceThreshold) &&
629
0
                            (ratio < kRatioThreshold) &&
630
0
                            (norm < kNormThreshold) && (var > kVarThreshold);
631
        // The following weights are used to combine the above features to give
632
        // a sigmoid score for flatness. If the input was normalized to [0,100]
633
        // the magnitude of these values would be close to 1 (e.g., weights
634
        // corresponding to variance would be a factor of 10000x smaller).
635
        // The weights are given in the following order:
636
        //    [{var}, {ratio}, {trace}, {norm}, offset]
637
        // with one of the most discriminative being simply the variance.
638
0
        const double weights[5] = { -6682, -0.2056, 13087, -12434, 2.5694 };
639
0
        double sum_weights = weights[0] * var + weights[1] * ratio +
640
0
                             weights[2] * trace + weights[3] * norm +
641
0
                             weights[4];
642
        // clamp the value to [-25.0, 100.0] to prevent overflow
643
0
        sum_weights = fclamp(sum_weights, -25.0, 100.0);
644
0
        const float score = (float)(1.0 / (1 + exp(-sum_weights)));
645
0
        flat_blocks[by * num_blocks_w + bx] = is_flat ? 255 : 0;
646
0
        scores[by * num_blocks_w + bx].score = var > kVarThreshold ? score : 0;
647
0
        scores[by * num_blocks_w + bx].index = by * num_blocks_w + bx;
648
#ifdef NOISE_MODEL_LOG_SCORE
649
        fprintf(stderr, "%g %g %g %g %g %d ", score, var, ratio, trace, norm,
650
                is_flat);
651
#endif
652
0
        num_flat += is_flat;
653
0
      }
654
0
    }
655
#ifdef NOISE_MODEL_LOG_SCORE
656
    fprintf(stderr, "\n");
657
#endif
658
0
  }
659
#ifdef NOISE_MODEL_LOG_SCORE
660
  fprintf(stderr, "];\n");
661
#endif
662
  // Find the top-scored blocks (most likely to be flat) and set the flat blocks
663
  // be the union of the thresholded results and the top 10th percentile of the
664
  // scored results.
665
0
  qsort(scores, num_blocks_w * num_blocks_h, sizeof(*scores), &compare_scores);
666
0
  const int top_nth_percentile = num_blocks_w * num_blocks_h * 90 / 100;
667
0
  const float score_threshold = scores[top_nth_percentile].score;
668
0
  for (int i = 0; i < num_blocks_w * num_blocks_h; ++i) {
669
0
    if (scores[i].score >= score_threshold) {
670
0
      num_flat += flat_blocks[scores[i].index] == 0;
671
0
      flat_blocks[scores[i].index] |= 1;
672
0
    }
673
0
  }
674
0
  aom_free(block);
675
0
  aom_free(plane);
676
0
  aom_free(scores);
677
0
  return num_flat;
678
0
}
679
680
int aom_noise_model_init(aom_noise_model_t *model,
681
0
                         const aom_noise_model_params_t params) {
682
0
  const int n = num_coeffs(params);
683
0
  const int lag = params.lag;
684
0
  const int bit_depth = params.bit_depth;
685
0
  int x = 0, y = 0, i = 0, c = 0;
686
687
0
  memset(model, 0, sizeof(*model));
688
0
  if (params.lag < 1) {
689
0
    fprintf(stderr, "Invalid noise param: lag = %d must be >= 1\n", params.lag);
690
0
    return 0;
691
0
  }
692
0
  if (params.lag > kMaxLag) {
693
0
    fprintf(stderr, "Invalid noise param: lag = %d must be <= %d\n", params.lag,
694
0
            kMaxLag);
695
0
    return 0;
696
0
  }
697
698
0
  memcpy(&model->params, &params, sizeof(params));
699
0
  for (c = 0; c < 3; ++c) {
700
0
    if (!noise_state_init(&model->combined_state[c], n + (c > 0), bit_depth)) {
701
0
      fprintf(stderr, "Failed to allocate noise state for channel %d\n", c);
702
0
      aom_noise_model_free(model);
703
0
      return 0;
704
0
    }
705
0
    if (!noise_state_init(&model->latest_state[c], n + (c > 0), bit_depth)) {
706
0
      fprintf(stderr, "Failed to allocate noise state for channel %d\n", c);
707
0
      aom_noise_model_free(model);
708
0
      return 0;
709
0
    }
710
0
  }
711
0
  model->n = n;
712
0
  model->coords = (int(*)[2])aom_malloc(sizeof(*model->coords) * n);
713
714
0
  for (y = -lag; y <= 0; ++y) {
715
0
    const int max_x = y == 0 ? -1 : lag;
716
0
    for (x = -lag; x <= max_x; ++x) {
717
0
      switch (params.shape) {
718
0
        case AOM_NOISE_SHAPE_DIAMOND:
719
0
          if (abs(x) <= y + lag) {
720
0
            model->coords[i][0] = x;
721
0
            model->coords[i][1] = y;
722
0
            ++i;
723
0
          }
724
0
          break;
725
0
        case AOM_NOISE_SHAPE_SQUARE:
726
0
          model->coords[i][0] = x;
727
0
          model->coords[i][1] = y;
728
0
          ++i;
729
0
          break;
730
0
        default:
731
0
          fprintf(stderr, "Invalid shape\n");
732
0
          aom_noise_model_free(model);
733
0
          return 0;
734
0
      }
735
0
    }
736
0
  }
737
0
  assert(i == n);
738
0
  return 1;
739
0
}
740
741
0
void aom_noise_model_free(aom_noise_model_t *model) {
742
0
  int c = 0;
743
0
  if (!model) return;
744
745
0
  aom_free(model->coords);
746
0
  for (c = 0; c < 3; ++c) {
747
0
    equation_system_free(&model->latest_state[c].eqns);
748
0
    equation_system_free(&model->combined_state[c].eqns);
749
750
0
    equation_system_free(&model->latest_state[c].strength_solver.eqns);
751
0
    equation_system_free(&model->combined_state[c].strength_solver.eqns);
752
0
  }
753
0
  memset(model, 0, sizeof(*model));
754
0
}
755
756
// Extracts the neighborhood defined by coords around point (x, y) from
757
// the difference between the data and denoised images. Also extracts the
758
// entry (possibly downsampled) for (x, y) in the alt_data (e.g., luma).
759
#define EXTRACT_AR_ROW(INT_TYPE, suffix)                                   \
760
  static double extract_ar_row_##suffix(                                   \
761
      int(*coords)[2], int num_coords, const INT_TYPE *const data,         \
762
      const INT_TYPE *const denoised, int stride, int sub_log2[2],         \
763
      const INT_TYPE *const alt_data, const INT_TYPE *const alt_denoised,  \
764
0
      int alt_stride, int x, int y, double *buffer) {                      \
765
0
    for (int i = 0; i < num_coords; ++i) {                                 \
766
0
      const int x_i = x + coords[i][0], y_i = y + coords[i][1];            \
767
0
      buffer[i] =                                                          \
768
0
          (double)data[y_i * stride + x_i] - denoised[y_i * stride + x_i]; \
769
0
    }                                                                      \
770
0
    const double val =                                                     \
771
0
        (double)data[y * stride + x] - denoised[y * stride + x];           \
772
0
                                                                           \
773
0
    if (alt_data && alt_denoised) {                                        \
774
0
      double avg_data = 0, avg_denoised = 0;                               \
775
0
      int num_samples = 0;                                                 \
776
0
      for (int dy_i = 0; dy_i < (1 << sub_log2[1]); dy_i++) {              \
777
0
        const int y_up = (y << sub_log2[1]) + dy_i;                        \
778
0
        for (int dx_i = 0; dx_i < (1 << sub_log2[0]); dx_i++) {            \
779
0
          const int x_up = (x << sub_log2[0]) + dx_i;                      \
780
0
          avg_data += alt_data[y_up * alt_stride + x_up];                  \
781
0
          avg_denoised += alt_denoised[y_up * alt_stride + x_up];          \
782
0
          num_samples++;                                                   \
783
0
        }                                                                  \
784
0
      }                                                                    \
785
0
      buffer[num_coords] = (avg_data - avg_denoised) / num_samples;        \
786
0
    }                                                                      \
787
0
    return val;                                                            \
788
0
  }
Unexecuted instantiation: noise_model.c:extract_ar_row_highbd
Unexecuted instantiation: noise_model.c:extract_ar_row_lowbd
789
790
EXTRACT_AR_ROW(uint8_t, lowbd);
791
EXTRACT_AR_ROW(uint16_t, highbd);
792
793
static int add_block_observations(
794
    aom_noise_model_t *noise_model, int c, const uint8_t *const data,
795
    const uint8_t *const denoised, int w, int h, int stride, int sub_log2[2],
796
    const uint8_t *const alt_data, const uint8_t *const alt_denoised,
797
    int alt_stride, const uint8_t *const flat_blocks, int block_size,
798
0
    int num_blocks_w, int num_blocks_h) {
799
0
  const int lag = noise_model->params.lag;
800
0
  const int num_coords = noise_model->n;
801
0
  const double normalization = (1 << noise_model->params.bit_depth) - 1;
802
0
  double *A = noise_model->latest_state[c].eqns.A;
803
0
  double *b = noise_model->latest_state[c].eqns.b;
804
0
  double *buffer = (double *)aom_malloc(sizeof(*buffer) * (num_coords + 1));
805
0
  const int n = noise_model->latest_state[c].eqns.n;
806
807
0
  if (!buffer) {
808
0
    fprintf(stderr, "Unable to allocate buffer of size %d\n", num_coords + 1);
809
0
    return 0;
810
0
  }
811
0
  for (int by = 0; by < num_blocks_h; ++by) {
812
0
    const int y_o = by * (block_size >> sub_log2[1]);
813
0
    for (int bx = 0; bx < num_blocks_w; ++bx) {
814
0
      const int x_o = bx * (block_size >> sub_log2[0]);
815
0
      if (!flat_blocks[by * num_blocks_w + bx]) {
816
0
        continue;
817
0
      }
818
0
      int y_start =
819
0
          (by > 0 && flat_blocks[(by - 1) * num_blocks_w + bx]) ? 0 : lag;
820
0
      int x_start =
821
0
          (bx > 0 && flat_blocks[by * num_blocks_w + bx - 1]) ? 0 : lag;
822
0
      int y_end = AOMMIN((h >> sub_log2[1]) - by * (block_size >> sub_log2[1]),
823
0
                         block_size >> sub_log2[1]);
824
0
      int x_end = AOMMIN(
825
0
          (w >> sub_log2[0]) - bx * (block_size >> sub_log2[0]) - lag,
826
0
          (bx + 1 < num_blocks_w && flat_blocks[by * num_blocks_w + bx + 1])
827
0
              ? (block_size >> sub_log2[0])
828
0
              : ((block_size >> sub_log2[0]) - lag));
829
0
      for (int y = y_start; y < y_end; ++y) {
830
0
        for (int x = x_start; x < x_end; ++x) {
831
0
          const double val =
832
0
              noise_model->params.use_highbd
833
0
                  ? extract_ar_row_highbd(noise_model->coords, num_coords,
834
0
                                          (const uint16_t *const)data,
835
0
                                          (const uint16_t *const)denoised,
836
0
                                          stride, sub_log2,
837
0
                                          (const uint16_t *const)alt_data,
838
0
                                          (const uint16_t *const)alt_denoised,
839
0
                                          alt_stride, x + x_o, y + y_o, buffer)
840
0
                  : extract_ar_row_lowbd(noise_model->coords, num_coords, data,
841
0
                                         denoised, stride, sub_log2, alt_data,
842
0
                                         alt_denoised, alt_stride, x + x_o,
843
0
                                         y + y_o, buffer);
844
0
          for (int i = 0; i < n; ++i) {
845
0
            for (int j = 0; j < n; ++j) {
846
0
              A[i * n + j] +=
847
0
                  (buffer[i] * buffer[j]) / (normalization * normalization);
848
0
            }
849
0
            b[i] += (buffer[i] * val) / (normalization * normalization);
850
0
          }
851
0
          noise_model->latest_state[c].num_observations++;
852
0
        }
853
0
      }
854
0
    }
855
0
  }
856
0
  aom_free(buffer);
857
0
  return 1;
858
0
}
859
860
static void add_noise_std_observations(
861
    aom_noise_model_t *noise_model, int c, const double *coeffs,
862
    const uint8_t *const data, const uint8_t *const denoised, int w, int h,
863
    int stride, int sub_log2[2], const uint8_t *const alt_data, int alt_stride,
864
    const uint8_t *const flat_blocks, int block_size, int num_blocks_w,
865
0
    int num_blocks_h) {
866
0
  const int num_coords = noise_model->n;
867
0
  aom_noise_strength_solver_t *noise_strength_solver =
868
0
      &noise_model->latest_state[c].strength_solver;
869
870
0
  const aom_noise_strength_solver_t *noise_strength_luma =
871
0
      &noise_model->latest_state[0].strength_solver;
872
0
  const double luma_gain = noise_model->latest_state[0].ar_gain;
873
0
  const double noise_gain = noise_model->latest_state[c].ar_gain;
874
0
  for (int by = 0; by < num_blocks_h; ++by) {
875
0
    const int y_o = by * (block_size >> sub_log2[1]);
876
0
    for (int bx = 0; bx < num_blocks_w; ++bx) {
877
0
      const int x_o = bx * (block_size >> sub_log2[0]);
878
0
      if (!flat_blocks[by * num_blocks_w + bx]) {
879
0
        continue;
880
0
      }
881
0
      const int num_samples_h =
882
0
          AOMMIN((h >> sub_log2[1]) - by * (block_size >> sub_log2[1]),
883
0
                 block_size >> sub_log2[1]);
884
0
      const int num_samples_w =
885
0
          AOMMIN((w >> sub_log2[0]) - bx * (block_size >> sub_log2[0]),
886
0
                 (block_size >> sub_log2[0]));
887
      // Make sure that we have a reasonable amount of samples to consider the
888
      // block
889
0
      if (num_samples_w * num_samples_h > block_size) {
890
0
        const double block_mean = get_block_mean(
891
0
            alt_data ? alt_data : data, w, h, alt_data ? alt_stride : stride,
892
0
            x_o << sub_log2[0], y_o << sub_log2[1], block_size,
893
0
            noise_model->params.use_highbd);
894
0
        const double noise_var = get_noise_var(
895
0
            data, denoised, stride, w >> sub_log2[0], h >> sub_log2[1], x_o,
896
0
            y_o, block_size >> sub_log2[0], block_size >> sub_log2[1],
897
0
            noise_model->params.use_highbd);
898
        // We want to remove the part of the noise that came from being
899
        // correlated with luma. Note that the noise solver for luma must
900
        // have already been run.
901
0
        const double luma_strength =
902
0
            c > 0 ? luma_gain * noise_strength_solver_get_value(
903
0
                                    noise_strength_luma, block_mean)
904
0
                  : 0;
905
0
        const double corr = c > 0 ? coeffs[num_coords] : 0;
906
        // Chroma noise:
907
        //    N(0, noise_var) = N(0, uncorr_var) + corr * N(0, luma_strength^2)
908
        // The uncorrelated component:
909
        //   uncorr_var = noise_var - (corr * luma_strength)^2
910
        // But don't allow fully correlated noise (hence the max), since the
911
        // synthesis cannot model it.
912
0
        const double uncorr_std = sqrt(
913
0
            AOMMAX(noise_var / 16, noise_var - pow(corr * luma_strength, 2)));
914
        // After we've removed correlation with luma, undo the gain that will
915
        // come from running the IIR filter.
916
0
        const double adjusted_strength = uncorr_std / noise_gain;
917
0
        aom_noise_strength_solver_add_measurement(
918
0
            noise_strength_solver, block_mean, adjusted_strength);
919
0
      }
920
0
    }
921
0
  }
922
0
}
923
924
// Return true if the noise estimate appears to be different from the combined
925
// (multi-frame) estimate. The difference is measured by checking whether the
926
// AR coefficients have diverged (using a threshold on normalized cross
927
// correlation), or whether the noise strength has changed.
928
0
static int is_noise_model_different(aom_noise_model_t *const noise_model) {
929
  // These thresholds are kind of arbitrary and will likely need further tuning
930
  // (or exported as parameters). The threshold on noise strength is a weighted
931
  // difference between the noise strength histograms
932
0
  const double kCoeffThreshold = 0.9;
933
0
  const double kStrengthThreshold =
934
0
      0.005 * (1 << (noise_model->params.bit_depth - 8));
935
0
  for (int c = 0; c < 1; ++c) {
936
0
    const double corr =
937
0
        aom_normalized_cross_correlation(noise_model->latest_state[c].eqns.x,
938
0
                                         noise_model->combined_state[c].eqns.x,
939
0
                                         noise_model->combined_state[c].eqns.n);
940
0
    if (corr < kCoeffThreshold) return 1;
941
942
0
    const double dx =
943
0
        1.0 / noise_model->latest_state[c].strength_solver.num_bins;
944
945
0
    const aom_equation_system_t *latest_eqns =
946
0
        &noise_model->latest_state[c].strength_solver.eqns;
947
0
    const aom_equation_system_t *combined_eqns =
948
0
        &noise_model->combined_state[c].strength_solver.eqns;
949
0
    double diff = 0;
950
0
    double total_weight = 0;
951
0
    for (int j = 0; j < latest_eqns->n; ++j) {
952
0
      double weight = 0;
953
0
      for (int i = 0; i < latest_eqns->n; ++i) {
954
0
        weight += latest_eqns->A[i * latest_eqns->n + j];
955
0
      }
956
0
      weight = sqrt(weight);
957
0
      diff += weight * fabs(latest_eqns->x[j] - combined_eqns->x[j]);
958
0
      total_weight += weight;
959
0
    }
960
0
    if (diff * dx / total_weight > kStrengthThreshold) return 1;
961
0
  }
962
0
  return 0;
963
0
}
964
965
0
static int ar_equation_system_solve(aom_noise_state_t *state, int is_chroma) {
966
0
  const int ret = equation_system_solve(&state->eqns);
967
0
  state->ar_gain = 1.0;
968
0
  if (!ret) return ret;
969
970
  // Update the AR gain from the equation system as it will be used to fit
971
  // the noise strength as a function of intensity.  In the Yule-Walker
972
  // equations, the diagonal should be the variance of the correlated noise.
973
  // In the case of the least squares estimate, there will be some variability
974
  // in the diagonal. So use the mean of the diagonal as the estimate of
975
  // overall variance (this works for least squares or Yule-Walker formulation).
976
0
  double var = 0;
977
0
  const int n = state->eqns.n;
978
0
  for (int i = 0; i < (state->eqns.n - is_chroma); ++i) {
979
0
    var += state->eqns.A[i * n + i] / state->num_observations;
980
0
  }
981
0
  var /= (n - is_chroma);
982
983
  // Keep track of E(Y^2) = <b, x> + E(X^2)
984
  // In the case that we are using chroma and have an estimate of correlation
985
  // with luma we adjust that estimate slightly to remove the correlated bits by
986
  // subtracting out the last column of a scaled by our correlation estimate
987
  // from b. E(y^2) = <b - A(:, end)*x(end), x>
988
0
  double sum_covar = 0;
989
0
  for (int i = 0; i < state->eqns.n - is_chroma; ++i) {
990
0
    double bi = state->eqns.b[i];
991
0
    if (is_chroma) {
992
0
      bi -= state->eqns.A[i * n + (n - 1)] * state->eqns.x[n - 1];
993
0
    }
994
0
    sum_covar += (bi * state->eqns.x[i]) / state->num_observations;
995
0
  }
996
  // Now, get an estimate of the variance of uncorrelated noise signal and use
997
  // it to determine the gain of the AR filter.
998
0
  const double noise_var = AOMMAX(var - sum_covar, 1e-6);
999
0
  state->ar_gain = AOMMAX(1, sqrt(AOMMAX(var / noise_var, 1e-6)));
1000
0
  return ret;
1001
0
}
1002
1003
aom_noise_status_t aom_noise_model_update(
1004
    aom_noise_model_t *const noise_model, const uint8_t *const data[3],
1005
    const uint8_t *const denoised[3], int w, int h, int stride[3],
1006
0
    int chroma_sub_log2[2], const uint8_t *const flat_blocks, int block_size) {
1007
0
  const int num_blocks_w = (w + block_size - 1) / block_size;
1008
0
  const int num_blocks_h = (h + block_size - 1) / block_size;
1009
0
  int y_model_different = 0;
1010
0
  int num_blocks = 0;
1011
0
  int i = 0, channel = 0;
1012
1013
0
  if (block_size <= 1) {
1014
0
    fprintf(stderr, "block_size = %d must be > 1\n", block_size);
1015
0
    return AOM_NOISE_STATUS_INVALID_ARGUMENT;
1016
0
  }
1017
1018
0
  if (block_size < noise_model->params.lag * 2 + 1) {
1019
0
    fprintf(stderr, "block_size = %d must be >= %d\n", block_size,
1020
0
            noise_model->params.lag * 2 + 1);
1021
0
    return AOM_NOISE_STATUS_INVALID_ARGUMENT;
1022
0
  }
1023
1024
  // Clear the latest equation system
1025
0
  for (i = 0; i < 3; ++i) {
1026
0
    equation_system_clear(&noise_model->latest_state[i].eqns);
1027
0
    noise_model->latest_state[i].num_observations = 0;
1028
0
    noise_strength_solver_clear(&noise_model->latest_state[i].strength_solver);
1029
0
  }
1030
1031
  // Check that we have enough flat blocks
1032
0
  for (i = 0; i < num_blocks_h * num_blocks_w; ++i) {
1033
0
    if (flat_blocks[i]) {
1034
0
      num_blocks++;
1035
0
    }
1036
0
  }
1037
1038
0
  if (num_blocks <= 1) {
1039
0
    fprintf(stderr, "Not enough flat blocks to update noise estimate\n");
1040
0
    return AOM_NOISE_STATUS_INSUFFICIENT_FLAT_BLOCKS;
1041
0
  }
1042
1043
0
  for (channel = 0; channel < 3; ++channel) {
1044
0
    int no_subsampling[2] = { 0, 0 };
1045
0
    const uint8_t *alt_data = channel > 0 ? data[0] : 0;
1046
0
    const uint8_t *alt_denoised = channel > 0 ? denoised[0] : 0;
1047
0
    int *sub = channel > 0 ? chroma_sub_log2 : no_subsampling;
1048
0
    const int is_chroma = channel != 0;
1049
0
    if (!data[channel] || !denoised[channel]) break;
1050
0
    if (!add_block_observations(noise_model, channel, data[channel],
1051
0
                                denoised[channel], w, h, stride[channel], sub,
1052
0
                                alt_data, alt_denoised, stride[0], flat_blocks,
1053
0
                                block_size, num_blocks_w, num_blocks_h)) {
1054
0
      fprintf(stderr, "Adding block observation failed\n");
1055
0
      return AOM_NOISE_STATUS_INTERNAL_ERROR;
1056
0
    }
1057
1058
0
    if (!ar_equation_system_solve(&noise_model->latest_state[channel],
1059
0
                                  is_chroma)) {
1060
0
      if (is_chroma) {
1061
0
        set_chroma_coefficient_fallback_soln(
1062
0
            &noise_model->latest_state[channel].eqns);
1063
0
      } else {
1064
0
        fprintf(stderr, "Solving latest noise equation system failed %d!\n",
1065
0
                channel);
1066
0
        return AOM_NOISE_STATUS_INTERNAL_ERROR;
1067
0
      }
1068
0
    }
1069
1070
0
    add_noise_std_observations(
1071
0
        noise_model, channel, noise_model->latest_state[channel].eqns.x,
1072
0
        data[channel], denoised[channel], w, h, stride[channel], sub, alt_data,
1073
0
        stride[0], flat_blocks, block_size, num_blocks_w, num_blocks_h);
1074
1075
0
    if (!aom_noise_strength_solver_solve(
1076
0
            &noise_model->latest_state[channel].strength_solver)) {
1077
0
      fprintf(stderr, "Solving latest noise strength failed!\n");
1078
0
      return AOM_NOISE_STATUS_INTERNAL_ERROR;
1079
0
    }
1080
1081
    // Check noise characteristics and return if error.
1082
0
    if (channel == 0 &&
1083
0
        noise_model->combined_state[channel].strength_solver.num_equations >
1084
0
            0 &&
1085
0
        is_noise_model_different(noise_model)) {
1086
0
      y_model_different = 1;
1087
0
    }
1088
1089
    // Don't update the combined stats if the y model is different.
1090
0
    if (y_model_different) continue;
1091
1092
0
    noise_model->combined_state[channel].num_observations +=
1093
0
        noise_model->latest_state[channel].num_observations;
1094
0
    equation_system_add(&noise_model->combined_state[channel].eqns,
1095
0
                        &noise_model->latest_state[channel].eqns);
1096
0
    if (!ar_equation_system_solve(&noise_model->combined_state[channel],
1097
0
                                  is_chroma)) {
1098
0
      if (is_chroma) {
1099
0
        set_chroma_coefficient_fallback_soln(
1100
0
            &noise_model->combined_state[channel].eqns);
1101
0
      } else {
1102
0
        fprintf(stderr, "Solving combined noise equation system failed %d!\n",
1103
0
                channel);
1104
0
        return AOM_NOISE_STATUS_INTERNAL_ERROR;
1105
0
      }
1106
0
    }
1107
1108
0
    noise_strength_solver_add(
1109
0
        &noise_model->combined_state[channel].strength_solver,
1110
0
        &noise_model->latest_state[channel].strength_solver);
1111
1112
0
    if (!aom_noise_strength_solver_solve(
1113
0
            &noise_model->combined_state[channel].strength_solver)) {
1114
0
      fprintf(stderr, "Solving combined noise strength failed!\n");
1115
0
      return AOM_NOISE_STATUS_INTERNAL_ERROR;
1116
0
    }
1117
0
  }
1118
1119
0
  return y_model_different ? AOM_NOISE_STATUS_DIFFERENT_NOISE_TYPE
1120
0
                           : AOM_NOISE_STATUS_OK;
1121
0
}
1122
1123
0
void aom_noise_model_save_latest(aom_noise_model_t *noise_model) {
1124
0
  for (int c = 0; c < 3; c++) {
1125
0
    equation_system_copy(&noise_model->combined_state[c].eqns,
1126
0
                         &noise_model->latest_state[c].eqns);
1127
0
    equation_system_copy(&noise_model->combined_state[c].strength_solver.eqns,
1128
0
                         &noise_model->latest_state[c].strength_solver.eqns);
1129
0
    noise_model->combined_state[c].strength_solver.num_equations =
1130
0
        noise_model->latest_state[c].strength_solver.num_equations;
1131
0
    noise_model->combined_state[c].num_observations =
1132
0
        noise_model->latest_state[c].num_observations;
1133
0
    noise_model->combined_state[c].ar_gain =
1134
0
        noise_model->latest_state[c].ar_gain;
1135
0
  }
1136
0
}
1137
1138
int aom_noise_model_get_grain_parameters(aom_noise_model_t *const noise_model,
1139
0
                                         aom_film_grain_t *film_grain) {
1140
0
  if (noise_model->params.lag > 3) {
1141
0
    fprintf(stderr, "params.lag = %d > 3\n", noise_model->params.lag);
1142
0
    return 0;
1143
0
  }
1144
0
  uint16_t random_seed = film_grain->random_seed;
1145
0
  memset(film_grain, 0, sizeof(*film_grain));
1146
0
  film_grain->random_seed = random_seed;
1147
1148
0
  film_grain->apply_grain = 1;
1149
0
  film_grain->update_parameters = 1;
1150
1151
0
  film_grain->ar_coeff_lag = noise_model->params.lag;
1152
1153
  // Convert the scaling functions to 8 bit values
1154
0
  aom_noise_strength_lut_t scaling_points[3];
1155
0
  if (!aom_noise_strength_solver_fit_piecewise(
1156
0
          &noise_model->combined_state[0].strength_solver, 14,
1157
0
          scaling_points + 0)) {
1158
0
    return 0;
1159
0
  }
1160
0
  if (!aom_noise_strength_solver_fit_piecewise(
1161
0
          &noise_model->combined_state[1].strength_solver, 10,
1162
0
          scaling_points + 1)) {
1163
0
    aom_noise_strength_lut_free(scaling_points + 0);
1164
0
    return 0;
1165
0
  }
1166
0
  if (!aom_noise_strength_solver_fit_piecewise(
1167
0
          &noise_model->combined_state[2].strength_solver, 10,
1168
0
          scaling_points + 2)) {
1169
0
    aom_noise_strength_lut_free(scaling_points + 0);
1170
0
    aom_noise_strength_lut_free(scaling_points + 1);
1171
0
    return 0;
1172
0
  }
1173
1174
  // Both the domain and the range of the scaling functions in the film_grain
1175
  // are normalized to 8-bit (e.g., they are implicitly scaled during grain
1176
  // synthesis).
1177
0
  const double strength_divisor = 1 << (noise_model->params.bit_depth - 8);
1178
0
  double max_scaling_value = 1e-4;
1179
0
  for (int c = 0; c < 3; ++c) {
1180
0
    for (int i = 0; i < scaling_points[c].num_points; ++i) {
1181
0
      scaling_points[c].points[i][0] =
1182
0
          AOMMIN(255, scaling_points[c].points[i][0] / strength_divisor);
1183
0
      scaling_points[c].points[i][1] =
1184
0
          AOMMIN(255, scaling_points[c].points[i][1] / strength_divisor);
1185
0
      max_scaling_value =
1186
0
          AOMMAX(scaling_points[c].points[i][1], max_scaling_value);
1187
0
    }
1188
0
  }
1189
1190
  // Scaling_shift values are in the range [8,11]
1191
0
  const int max_scaling_value_log2 =
1192
0
      clamp((int)floor(log2(max_scaling_value) + 1), 2, 5);
1193
0
  film_grain->scaling_shift = 5 + (8 - max_scaling_value_log2);
1194
1195
0
  const double scale_factor = 1 << (8 - max_scaling_value_log2);
1196
0
  film_grain->num_y_points = scaling_points[0].num_points;
1197
0
  film_grain->num_cb_points = scaling_points[1].num_points;
1198
0
  film_grain->num_cr_points = scaling_points[2].num_points;
1199
1200
0
  int(*film_grain_scaling[3])[2] = {
1201
0
    film_grain->scaling_points_y,
1202
0
    film_grain->scaling_points_cb,
1203
0
    film_grain->scaling_points_cr,
1204
0
  };
1205
0
  for (int c = 0; c < 3; c++) {
1206
0
    for (int i = 0; i < scaling_points[c].num_points; ++i) {
1207
0
      film_grain_scaling[c][i][0] = (int)(scaling_points[c].points[i][0] + 0.5);
1208
0
      film_grain_scaling[c][i][1] = clamp(
1209
0
          (int)(scale_factor * scaling_points[c].points[i][1] + 0.5), 0, 255);
1210
0
    }
1211
0
  }
1212
0
  aom_noise_strength_lut_free(scaling_points + 0);
1213
0
  aom_noise_strength_lut_free(scaling_points + 1);
1214
0
  aom_noise_strength_lut_free(scaling_points + 2);
1215
1216
  // Convert the ar_coeffs into 8-bit values
1217
0
  const int n_coeff = noise_model->combined_state[0].eqns.n;
1218
0
  double max_coeff = 1e-4, min_coeff = -1e-4;
1219
0
  double y_corr[2] = { 0, 0 };
1220
0
  double avg_luma_strength = 0;
1221
0
  for (int c = 0; c < 3; c++) {
1222
0
    aom_equation_system_t *eqns = &noise_model->combined_state[c].eqns;
1223
0
    for (int i = 0; i < n_coeff; ++i) {
1224
0
      max_coeff = AOMMAX(max_coeff, eqns->x[i]);
1225
0
      min_coeff = AOMMIN(min_coeff, eqns->x[i]);
1226
0
    }
1227
    // Since the correlation between luma/chroma was computed in an already
1228
    // scaled space, we adjust it in the un-scaled space.
1229
0
    aom_noise_strength_solver_t *solver =
1230
0
        &noise_model->combined_state[c].strength_solver;
1231
    // Compute a weighted average of the strength for the channel.
1232
0
    double average_strength = 0, total_weight = 0;
1233
0
    for (int i = 0; i < solver->eqns.n; ++i) {
1234
0
      double w = 0;
1235
0
      for (int j = 0; j < solver->eqns.n; ++j) {
1236
0
        w += solver->eqns.A[i * solver->eqns.n + j];
1237
0
      }
1238
0
      w = sqrt(w);
1239
0
      average_strength += solver->eqns.x[i] * w;
1240
0
      total_weight += w;
1241
0
    }
1242
0
    if (total_weight == 0)
1243
0
      average_strength = 1;
1244
0
    else
1245
0
      average_strength /= total_weight;
1246
0
    if (c == 0) {
1247
0
      avg_luma_strength = average_strength;
1248
0
    } else {
1249
0
      y_corr[c - 1] = avg_luma_strength * eqns->x[n_coeff] / average_strength;
1250
0
      max_coeff = AOMMAX(max_coeff, y_corr[c - 1]);
1251
0
      min_coeff = AOMMIN(min_coeff, y_corr[c - 1]);
1252
0
    }
1253
0
  }
1254
  // Shift value: AR coeffs range (values 6-9)
1255
  // 6: [-2, 2),  7: [-1, 1), 8: [-0.5, 0.5), 9: [-0.25, 0.25)
1256
0
  film_grain->ar_coeff_shift =
1257
0
      clamp(7 - (int)AOMMAX(1 + floor(log2(max_coeff)), ceil(log2(-min_coeff))),
1258
0
            6, 9);
1259
0
  double scale_ar_coeff = 1 << film_grain->ar_coeff_shift;
1260
0
  int *ar_coeffs[3] = {
1261
0
    film_grain->ar_coeffs_y,
1262
0
    film_grain->ar_coeffs_cb,
1263
0
    film_grain->ar_coeffs_cr,
1264
0
  };
1265
0
  for (int c = 0; c < 3; ++c) {
1266
0
    aom_equation_system_t *eqns = &noise_model->combined_state[c].eqns;
1267
0
    for (int i = 0; i < n_coeff; ++i) {
1268
0
      ar_coeffs[c][i] =
1269
0
          clamp((int)round(scale_ar_coeff * eqns->x[i]), -128, 127);
1270
0
    }
1271
0
    if (c > 0) {
1272
0
      ar_coeffs[c][n_coeff] =
1273
0
          clamp((int)round(scale_ar_coeff * y_corr[c - 1]), -128, 127);
1274
0
    }
1275
0
  }
1276
1277
  // At the moment, the noise modeling code assumes that the chroma scaling
1278
  // functions are a function of luma.
1279
0
  film_grain->cb_mult = 128;       // 8 bits
1280
0
  film_grain->cb_luma_mult = 192;  // 8 bits
1281
0
  film_grain->cb_offset = 256;     // 9 bits
1282
1283
0
  film_grain->cr_mult = 128;       // 8 bits
1284
0
  film_grain->cr_luma_mult = 192;  // 8 bits
1285
0
  film_grain->cr_offset = 256;     // 9 bits
1286
1287
0
  film_grain->chroma_scaling_from_luma = 0;
1288
0
  film_grain->grain_scale_shift = 0;
1289
0
  film_grain->overlap_flag = 1;
1290
0
  return 1;
1291
0
}
1292
1293
0
static void pointwise_multiply(const float *a, float *b, int n) {
1294
0
  for (int i = 0; i < n; ++i) {
1295
0
    b[i] *= a[i];
1296
0
  }
1297
0
}
1298
1299
0
static float *get_half_cos_window(int block_size) {
1300
0
  float *window_function =
1301
0
      (float *)aom_malloc(block_size * block_size * sizeof(*window_function));
1302
0
  for (int y = 0; y < block_size; ++y) {
1303
0
    const double cos_yd = cos((.5 + y) * PI / block_size - PI / 2);
1304
0
    for (int x = 0; x < block_size; ++x) {
1305
0
      const double cos_xd = cos((.5 + x) * PI / block_size - PI / 2);
1306
0
      window_function[y * block_size + x] = (float)(cos_yd * cos_xd);
1307
0
    }
1308
0
  }
1309
0
  return window_function;
1310
0
}
1311
1312
#define DITHER_AND_QUANTIZE(INT_TYPE, suffix)                               \
1313
  static void dither_and_quantize_##suffix(                                 \
1314
      float *result, int result_stride, INT_TYPE *denoised, int w, int h,   \
1315
      int stride, int chroma_sub_w, int chroma_sub_h, int block_size,       \
1316
0
      float block_normalization) {                                          \
1317
0
    for (int y = 0; y < (h >> chroma_sub_h); ++y) {                         \
1318
0
      for (int x = 0; x < (w >> chroma_sub_w); ++x) {                       \
1319
0
        const int result_idx =                                              \
1320
0
            (y + (block_size >> chroma_sub_h)) * result_stride + x +        \
1321
0
            (block_size >> chroma_sub_w);                                   \
1322
0
        INT_TYPE new_val = (INT_TYPE)AOMMIN(                                \
1323
0
            AOMMAX(result[result_idx] * block_normalization + 0.5f, 0),     \
1324
0
            block_normalization);                                           \
1325
0
        const float err =                                                   \
1326
0
            -(((float)new_val) / block_normalization - result[result_idx]); \
1327
0
        denoised[y * stride + x] = new_val;                                 \
1328
0
        if (x + 1 < (w >> chroma_sub_w)) {                                  \
1329
0
          result[result_idx + 1] += err * 7.0f / 16.0f;                     \
1330
0
        }                                                                   \
1331
0
        if (y + 1 < (h >> chroma_sub_h)) {                                  \
1332
0
          if (x > 0) {                                                      \
1333
0
            result[result_idx + result_stride - 1] += err * 3.0f / 16.0f;   \
1334
0
          }                                                                 \
1335
0
          result[result_idx + result_stride] += err * 5.0f / 16.0f;         \
1336
0
          if (x + 1 < (w >> chroma_sub_w)) {                                \
1337
0
            result[result_idx + result_stride + 1] += err * 1.0f / 16.0f;   \
1338
0
          }                                                                 \
1339
0
        }                                                                   \
1340
0
      }                                                                     \
1341
0
    }                                                                       \
1342
0
  }
Unexecuted instantiation: noise_model.c:dither_and_quantize_highbd
Unexecuted instantiation: noise_model.c:dither_and_quantize_lowbd
1343
1344
DITHER_AND_QUANTIZE(uint8_t, lowbd);
1345
DITHER_AND_QUANTIZE(uint16_t, highbd);
1346
1347
int aom_wiener_denoise_2d(const uint8_t *const data[3], uint8_t *denoised[3],
1348
                          int w, int h, int stride[3], int chroma_sub[2],
1349
                          float *noise_psd[3], int block_size, int bit_depth,
1350
0
                          int use_highbd) {
1351
0
  float *plane = NULL, *block = NULL, *window_full = NULL,
1352
0
        *window_chroma = NULL;
1353
0
  double *block_d = NULL, *plane_d = NULL;
1354
0
  struct aom_noise_tx_t *tx_full = NULL;
1355
0
  struct aom_noise_tx_t *tx_chroma = NULL;
1356
0
  const int num_blocks_w = (w + block_size - 1) / block_size;
1357
0
  const int num_blocks_h = (h + block_size - 1) / block_size;
1358
0
  const int result_stride = (num_blocks_w + 2) * block_size;
1359
0
  const int result_height = (num_blocks_h + 2) * block_size;
1360
0
  float *result = NULL;
1361
0
  int init_success = 1;
1362
0
  aom_flat_block_finder_t block_finder_full;
1363
0
  aom_flat_block_finder_t block_finder_chroma;
1364
0
  const float kBlockNormalization = (float)((1 << bit_depth) - 1);
1365
0
  if (chroma_sub[0] != chroma_sub[1]) {
1366
0
    fprintf(stderr,
1367
0
            "aom_wiener_denoise_2d doesn't handle different chroma "
1368
0
            "subsampling\n");
1369
0
    return 0;
1370
0
  }
1371
0
  init_success &= aom_flat_block_finder_init(&block_finder_full, block_size,
1372
0
                                             bit_depth, use_highbd);
1373
0
  result = (float *)aom_malloc((num_blocks_h + 2) * block_size * result_stride *
1374
0
                               sizeof(*result));
1375
0
  plane = (float *)aom_malloc(block_size * block_size * sizeof(*plane));
1376
0
  block =
1377
0
      (float *)aom_memalign(32, 2 * block_size * block_size * sizeof(*block));
1378
0
  block_d = (double *)aom_malloc(block_size * block_size * sizeof(*block_d));
1379
0
  plane_d = (double *)aom_malloc(block_size * block_size * sizeof(*plane_d));
1380
0
  window_full = get_half_cos_window(block_size);
1381
0
  tx_full = aom_noise_tx_malloc(block_size);
1382
1383
0
  if (chroma_sub[0] != 0) {
1384
0
    init_success &= aom_flat_block_finder_init(&block_finder_chroma,
1385
0
                                               block_size >> chroma_sub[0],
1386
0
                                               bit_depth, use_highbd);
1387
0
    window_chroma = get_half_cos_window(block_size >> chroma_sub[0]);
1388
0
    tx_chroma = aom_noise_tx_malloc(block_size >> chroma_sub[0]);
1389
0
  } else {
1390
0
    window_chroma = window_full;
1391
0
    tx_chroma = tx_full;
1392
0
  }
1393
1394
0
  init_success &= (tx_full != NULL) && (tx_chroma != NULL) && (plane != NULL) &&
1395
0
                  (plane_d != NULL) && (block != NULL) && (block_d != NULL) &&
1396
0
                  (window_full != NULL) && (window_chroma != NULL) &&
1397
0
                  (result != NULL);
1398
0
  for (int c = init_success ? 0 : 3; c < 3; ++c) {
1399
0
    float *window_function = c == 0 ? window_full : window_chroma;
1400
0
    aom_flat_block_finder_t *block_finder = &block_finder_full;
1401
0
    const int chroma_sub_h = c > 0 ? chroma_sub[1] : 0;
1402
0
    const int chroma_sub_w = c > 0 ? chroma_sub[0] : 0;
1403
0
    struct aom_noise_tx_t *tx =
1404
0
        (c > 0 && chroma_sub[0] > 0) ? tx_chroma : tx_full;
1405
0
    if (!data[c] || !denoised[c]) continue;
1406
0
    if (c > 0 && chroma_sub[0] != 0) {
1407
0
      block_finder = &block_finder_chroma;
1408
0
    }
1409
0
    memset(result, 0, sizeof(*result) * result_stride * result_height);
1410
    // Do overlapped block processing (half overlapped). The block rows can
1411
    // easily be done in parallel
1412
0
    for (int offsy = 0; offsy < (block_size >> chroma_sub_h);
1413
0
         offsy += (block_size >> chroma_sub_h) / 2) {
1414
0
      for (int offsx = 0; offsx < (block_size >> chroma_sub_w);
1415
0
           offsx += (block_size >> chroma_sub_w) / 2) {
1416
        // Pad the boundary when processing each block-set.
1417
0
        for (int by = -1; by < num_blocks_h; ++by) {
1418
0
          for (int bx = -1; bx < num_blocks_w; ++bx) {
1419
0
            const int pixels_per_block =
1420
0
                (block_size >> chroma_sub_w) * (block_size >> chroma_sub_h);
1421
0
            aom_flat_block_finder_extract_block(
1422
0
                block_finder, data[c], w >> chroma_sub_w, h >> chroma_sub_h,
1423
0
                stride[c], bx * (block_size >> chroma_sub_w) + offsx,
1424
0
                by * (block_size >> chroma_sub_h) + offsy, plane_d, block_d);
1425
0
            for (int j = 0; j < pixels_per_block; ++j) {
1426
0
              block[j] = (float)block_d[j];
1427
0
              plane[j] = (float)plane_d[j];
1428
0
            }
1429
0
            pointwise_multiply(window_function, block, pixels_per_block);
1430
0
            aom_noise_tx_forward(tx, block);
1431
0
            aom_noise_tx_filter(tx, noise_psd[c]);
1432
0
            aom_noise_tx_inverse(tx, block);
1433
1434
            // Apply window function to the plane approximation (we will apply
1435
            // it to the sum of plane + block when composing the results).
1436
0
            pointwise_multiply(window_function, plane, pixels_per_block);
1437
1438
0
            for (int y = 0; y < (block_size >> chroma_sub_h); ++y) {
1439
0
              const int y_result =
1440
0
                  y + (by + 1) * (block_size >> chroma_sub_h) + offsy;
1441
0
              for (int x = 0; x < (block_size >> chroma_sub_w); ++x) {
1442
0
                const int x_result =
1443
0
                    x + (bx + 1) * (block_size >> chroma_sub_w) + offsx;
1444
0
                result[y_result * result_stride + x_result] +=
1445
0
                    (block[y * (block_size >> chroma_sub_w) + x] +
1446
0
                     plane[y * (block_size >> chroma_sub_w) + x]) *
1447
0
                    window_function[y * (block_size >> chroma_sub_w) + x];
1448
0
              }
1449
0
            }
1450
0
          }
1451
0
        }
1452
0
      }
1453
0
    }
1454
0
    if (use_highbd) {
1455
0
      dither_and_quantize_highbd(result, result_stride, (uint16_t *)denoised[c],
1456
0
                                 w, h, stride[c], chroma_sub_w, chroma_sub_h,
1457
0
                                 block_size, kBlockNormalization);
1458
0
    } else {
1459
0
      dither_and_quantize_lowbd(result, result_stride, denoised[c], w, h,
1460
0
                                stride[c], chroma_sub_w, chroma_sub_h,
1461
0
                                block_size, kBlockNormalization);
1462
0
    }
1463
0
  }
1464
0
  aom_free(result);
1465
0
  aom_free(plane);
1466
0
  aom_free(block);
1467
0
  aom_free(plane_d);
1468
0
  aom_free(block_d);
1469
0
  aom_free(window_full);
1470
1471
0
  aom_noise_tx_free(tx_full);
1472
1473
0
  aom_flat_block_finder_free(&block_finder_full);
1474
0
  if (chroma_sub[0] != 0) {
1475
0
    aom_flat_block_finder_free(&block_finder_chroma);
1476
0
    aom_free(window_chroma);
1477
0
    aom_noise_tx_free(tx_chroma);
1478
0
  }
1479
0
  return init_success;
1480
0
}
1481
1482
struct aom_denoise_and_model_t {
1483
  int block_size;
1484
  int bit_depth;
1485
  float noise_level;
1486
1487
  // Size of current denoised buffer and flat_block buffer
1488
  int width;
1489
  int height;
1490
  int y_stride;
1491
  int uv_stride;
1492
  int num_blocks_w;
1493
  int num_blocks_h;
1494
1495
  // Buffers for image and noise_psd allocated on the fly
1496
  float *noise_psd[3];
1497
  uint8_t *denoised[3];
1498
  uint8_t *flat_blocks;
1499
1500
  aom_flat_block_finder_t flat_block_finder;
1501
  aom_noise_model_t noise_model;
1502
};
1503
1504
struct aom_denoise_and_model_t *aom_denoise_and_model_alloc(int bit_depth,
1505
                                                            int block_size,
1506
0
                                                            float noise_level) {
1507
0
  struct aom_denoise_and_model_t *ctx =
1508
0
      (struct aom_denoise_and_model_t *)aom_malloc(
1509
0
          sizeof(struct aom_denoise_and_model_t));
1510
0
  if (!ctx) {
1511
0
    fprintf(stderr, "Unable to allocate denoise_and_model struct\n");
1512
0
    return NULL;
1513
0
  }
1514
0
  memset(ctx, 0, sizeof(*ctx));
1515
1516
0
  ctx->block_size = block_size;
1517
0
  ctx->noise_level = noise_level;
1518
0
  ctx->bit_depth = bit_depth;
1519
1520
0
  ctx->noise_psd[0] =
1521
0
      aom_malloc(sizeof(*ctx->noise_psd[0]) * block_size * block_size);
1522
0
  ctx->noise_psd[1] =
1523
0
      aom_malloc(sizeof(*ctx->noise_psd[1]) * block_size * block_size);
1524
0
  ctx->noise_psd[2] =
1525
0
      aom_malloc(sizeof(*ctx->noise_psd[2]) * block_size * block_size);
1526
0
  if (!ctx->noise_psd[0] || !ctx->noise_psd[1] || !ctx->noise_psd[2]) {
1527
0
    fprintf(stderr, "Unable to allocate noise PSD buffers\n");
1528
0
    aom_denoise_and_model_free(ctx);
1529
0
    return NULL;
1530
0
  }
1531
0
  return ctx;
1532
0
}
1533
1534
0
void aom_denoise_and_model_free(struct aom_denoise_and_model_t *ctx) {
1535
0
  aom_free(ctx->flat_blocks);
1536
0
  for (int i = 0; i < 3; ++i) {
1537
0
    aom_free(ctx->denoised[i]);
1538
0
    aom_free(ctx->noise_psd[i]);
1539
0
  }
1540
0
  aom_noise_model_free(&ctx->noise_model);
1541
0
  aom_flat_block_finder_free(&ctx->flat_block_finder);
1542
0
  aom_free(ctx);
1543
0
}
1544
1545
static int denoise_and_model_realloc_if_necessary(
1546
0
    struct aom_denoise_and_model_t *ctx, YV12_BUFFER_CONFIG *sd) {
1547
0
  if (ctx->width == sd->y_width && ctx->height == sd->y_height &&
1548
0
      ctx->y_stride == sd->y_stride && ctx->uv_stride == sd->uv_stride)
1549
0
    return 1;
1550
0
  const int use_highbd = (sd->flags & YV12_FLAG_HIGHBITDEPTH) != 0;
1551
0
  const int block_size = ctx->block_size;
1552
1553
0
  ctx->width = sd->y_width;
1554
0
  ctx->height = sd->y_height;
1555
0
  ctx->y_stride = sd->y_stride;
1556
0
  ctx->uv_stride = sd->uv_stride;
1557
1558
0
  for (int i = 0; i < 3; ++i) {
1559
0
    aom_free(ctx->denoised[i]);
1560
0
    ctx->denoised[i] = NULL;
1561
0
  }
1562
0
  aom_free(ctx->flat_blocks);
1563
0
  ctx->flat_blocks = NULL;
1564
1565
0
  ctx->denoised[0] = aom_malloc((sd->y_stride * sd->y_height) << use_highbd);
1566
0
  ctx->denoised[1] = aom_malloc((sd->uv_stride * sd->uv_height) << use_highbd);
1567
0
  ctx->denoised[2] = aom_malloc((sd->uv_stride * sd->uv_height) << use_highbd);
1568
0
  if (!ctx->denoised[0] || !ctx->denoised[1] || !ctx->denoised[2]) {
1569
0
    fprintf(stderr, "Unable to allocate denoise buffers\n");
1570
0
    return 0;
1571
0
  }
1572
0
  ctx->num_blocks_w = (sd->y_width + ctx->block_size - 1) / ctx->block_size;
1573
0
  ctx->num_blocks_h = (sd->y_height + ctx->block_size - 1) / ctx->block_size;
1574
0
  ctx->flat_blocks = aom_malloc(ctx->num_blocks_w * ctx->num_blocks_h);
1575
1576
0
  aom_flat_block_finder_free(&ctx->flat_block_finder);
1577
0
  if (!aom_flat_block_finder_init(&ctx->flat_block_finder, ctx->block_size,
1578
0
                                  ctx->bit_depth, use_highbd)) {
1579
0
    fprintf(stderr, "Unable to init flat block finder\n");
1580
0
    return 0;
1581
0
  }
1582
1583
0
  const aom_noise_model_params_t params = { AOM_NOISE_SHAPE_SQUARE, 3,
1584
0
                                            ctx->bit_depth, use_highbd };
1585
0
  aom_noise_model_free(&ctx->noise_model);
1586
0
  if (!aom_noise_model_init(&ctx->noise_model, params)) {
1587
0
    fprintf(stderr, "Unable to init noise model\n");
1588
0
    return 0;
1589
0
  }
1590
1591
  // Simply use a flat PSD (although we could use the flat blocks to estimate
1592
  // PSD) those to estimate an actual noise PSD)
1593
0
  const float y_noise_level =
1594
0
      aom_noise_psd_get_default_value(ctx->block_size, ctx->noise_level);
1595
0
  const float uv_noise_level = aom_noise_psd_get_default_value(
1596
0
      ctx->block_size >> sd->subsampling_x, ctx->noise_level);
1597
0
  for (int i = 0; i < block_size * block_size; ++i) {
1598
0
    ctx->noise_psd[0][i] = y_noise_level;
1599
0
    ctx->noise_psd[1][i] = ctx->noise_psd[2][i] = uv_noise_level;
1600
0
  }
1601
0
  return 1;
1602
0
}
1603
1604
int aom_denoise_and_model_run(struct aom_denoise_and_model_t *ctx,
1605
                              YV12_BUFFER_CONFIG *sd,
1606
0
                              aom_film_grain_t *film_grain, int apply_denoise) {
1607
0
  const int block_size = ctx->block_size;
1608
0
  const int use_highbd = (sd->flags & YV12_FLAG_HIGHBITDEPTH) != 0;
1609
0
  uint8_t *raw_data[3] = {
1610
0
    use_highbd ? (uint8_t *)CONVERT_TO_SHORTPTR(sd->y_buffer) : sd->y_buffer,
1611
0
    use_highbd ? (uint8_t *)CONVERT_TO_SHORTPTR(sd->u_buffer) : sd->u_buffer,
1612
0
    use_highbd ? (uint8_t *)CONVERT_TO_SHORTPTR(sd->v_buffer) : sd->v_buffer,
1613
0
  };
1614
0
  const uint8_t *const data[3] = { raw_data[0], raw_data[1], raw_data[2] };
1615
0
  int strides[3] = { sd->y_stride, sd->uv_stride, sd->uv_stride };
1616
0
  int chroma_sub_log2[2] = { sd->subsampling_x, sd->subsampling_y };
1617
1618
0
  if (!denoise_and_model_realloc_if_necessary(ctx, sd)) {
1619
0
    fprintf(stderr, "Unable to realloc buffers\n");
1620
0
    return 0;
1621
0
  }
1622
1623
0
  aom_flat_block_finder_run(&ctx->flat_block_finder, data[0], sd->y_width,
1624
0
                            sd->y_height, strides[0], ctx->flat_blocks);
1625
1626
0
  if (!aom_wiener_denoise_2d(data, ctx->denoised, sd->y_width, sd->y_height,
1627
0
                             strides, chroma_sub_log2, ctx->noise_psd,
1628
0
                             block_size, ctx->bit_depth, use_highbd)) {
1629
0
    fprintf(stderr, "Unable to denoise image\n");
1630
0
    return 0;
1631
0
  }
1632
1633
0
  const aom_noise_status_t status = aom_noise_model_update(
1634
0
      &ctx->noise_model, data, (const uint8_t *const *)ctx->denoised,
1635
0
      sd->y_width, sd->y_height, strides, chroma_sub_log2, ctx->flat_blocks,
1636
0
      block_size);
1637
0
  int have_noise_estimate = 0;
1638
0
  if (status == AOM_NOISE_STATUS_OK) {
1639
0
    have_noise_estimate = 1;
1640
0
  } else if (status == AOM_NOISE_STATUS_DIFFERENT_NOISE_TYPE) {
1641
0
    aom_noise_model_save_latest(&ctx->noise_model);
1642
0
    have_noise_estimate = 1;
1643
0
  } else {
1644
    // Unable to update noise model; proceed if we have a previous estimate.
1645
0
    have_noise_estimate =
1646
0
        (ctx->noise_model.combined_state[0].strength_solver.num_equations > 0);
1647
0
  }
1648
1649
0
  film_grain->apply_grain = 0;
1650
0
  if (have_noise_estimate) {
1651
0
    if (!aom_noise_model_get_grain_parameters(&ctx->noise_model, film_grain)) {
1652
0
      fprintf(stderr, "Unable to get grain parameters.\n");
1653
0
      return 0;
1654
0
    }
1655
0
    if (!film_grain->random_seed) {
1656
0
      film_grain->random_seed = 7391;
1657
0
    }
1658
0
    if (apply_denoise) {
1659
0
      memcpy(raw_data[0], ctx->denoised[0],
1660
0
             (strides[0] * sd->y_height) << use_highbd);
1661
0
      memcpy(raw_data[1], ctx->denoised[1],
1662
0
             (strides[1] * sd->uv_height) << use_highbd);
1663
0
      memcpy(raw_data[2], ctx->denoised[2],
1664
0
             (strides[2] * sd->uv_height) << use_highbd);
1665
0
    }
1666
0
  }
1667
0
  return 1;
1668
0
}