summaryrefslogtreecommitdiffstats
path: root/third_party/aom/aom_dsp/noise_model.h
blob: 8228aeacfc085f69519968614bea7f3514f616ed (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
/*
 * 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_AOM_DSP_NOISE_MODEL_H_
#define AOM_AOM_DSP_NOISE_MODEL_H_

#ifdef __cplusplus
extern "C" {
#endif  // __cplusplus

#include <stdint.h>
#include "aom_dsp/grain_params.h"
#include "aom_ports/mem.h"
#include "aom_scale/yv12config.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.
 * \param[in]  bit_depth The bit depth used to derive {min,max}_intensity.
 */
int aom_noise_strength_solver_init(aom_noise_strength_solver_t *solver,
                                   int num_bins, int bit_depth);
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[in] max_num_points  The maximum number of output points
 * \param[out] lut  The output piecewise linear lut.
 */
int aom_noise_strength_solver_fit_piecewise(
    const aom_noise_strength_solver_t *solver, int max_num_points,
    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
  double normalization;  // Normalization factor (1 / (2^(bit_depth) - 1))
  int use_highbd;        // Whether input data should be interpreted as uint16
} aom_flat_block_finder_t;

/*!\brief Init the block_finder with the given block size, bit_depth */
int aom_flat_block_finder_init(aom_flat_block_finder_t *block_finder,
                               int block_size, int bit_depth, int use_highbd);
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);

/*!\brief Runs the flat block finder on the input data.
 *
 * Find flat blocks in the input image data. Returns a map of
 * flat_blocks, where the value of flat_blocks map will be non-zero
 * when a block is determined to be flat. A higher value indicates a bigger
 * confidence in the decision.
 */
int aom_flat_block_finder_run(const aom_flat_block_finder_t *block_finder,
                              const uint8_t *const data, int w, int h,
                              int stride, uint8_t *flat_blocks);

// The noise shape indicates the allowed coefficients in the AR model.
enum {
  AOM_NOISE_SHAPE_DIAMOND = 0,
  AOM_NOISE_SHAPE_SQUARE = 1
} UENUM1BYTE(aom_noise_shape);

// The parameters of the noise model include the shape type, lag, the
// bit depth of the input images provided, and whether the input images
// will be using uint16 (or uint8) representation.
typedef struct {
  aom_noise_shape shape;
  int lag;
  int bit_depth;
  int use_highbd;
} 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;
  int num_observations;  // The number of observations in the eqn system
  double ar_gain;        // The gain of the current AR filter
} 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)
  int bit_depth;
} aom_noise_model_t;

/*!\brief Result of a noise model update. */
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,
} UENUM1BYTE(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);

/*\brief Save the "latest" estimate into the "combined" estimate.
 *
 * This is meant to be called when the noise modeling detected a change
 * in parameters (or for example, if a user wanted to reset estimation at
 * a shot boundary).
 */
void aom_noise_model_save_latest(aom_noise_model_t *noise_model);

/*!\brief Converts the noise_model parameters to the corresponding
 *    grain_parameters.
 *
 * The noise structs in this file are suitable for estimation (e.g., using
 * floats), but the grain parameters in the bitstream are quantized. This
 * function does the conversion by selecting the correct quantization levels.
 */
int aom_noise_model_get_grain_parameters(aom_noise_model_t *const noise_model,
                                         aom_film_grain_t *film_grain);

/*!\brief Perform a Wiener filter denoising in 2D using the provided noise psd.
 *
 * \param[in]     data            Raw frame data
 * \param[out]    denoised        Denoised frame data
 * \param[in]     w               Frame width
 * \param[in]     h               Frame height
 * \param[in]     stride          Stride of the planes
 * \param[in]     chroma_sub_log2 Chroma subsampling for planes != 0.
 * \param[in]     noise_psd       The power spectral density of the noise
 * \param[in]     block_size      The size of blocks
 * \param[in]     bit_depth       Bit depth of the image
 * \param[in]     use_highbd      If true, uint8 pointers are interpreted as
 *                                uint16 and stride is measured in uint16.
 *                                This must be true when bit_depth >= 10.
 */
int aom_wiener_denoise_2d(const uint8_t *const data[3], uint8_t *denoised[3],
                          int w, int h, int stride[3], int chroma_sub_log2[2],
                          float *noise_psd[3], int block_size, int bit_depth,
                          int use_highbd);

struct aom_denoise_and_model_t;

/*!\brief Denoise the buffer and model the residual noise.
 *
 * This is meant to be called sequentially on input frames. The input buffer
 * is denoised and the residual noise is modelled. The current noise estimate
 * is populated in film_grain. Returns true on success. The grain.apply_grain
 * parameter will be true when the input buffer was successfully denoised and
 * grain was modelled. Returns false on error.
 *
 * \param[in]     ctx           Struct allocated with
 *                              aom_denoise_and_model_alloc that holds some
 *                              buffers for denoising and the current noise
 *                              estimate.
 * \param[in,out] buf           The raw input buffer to be denoised.
 * \param[out]    grain         Output film grain parameters
 * \param[in]     apply_denoise Whether or not to apply the denoising to the
 *                              frame that will be encoded
 */
int aom_denoise_and_model_run(struct aom_denoise_and_model_t *ctx,
                              YV12_BUFFER_CONFIG *buf, aom_film_grain_t *grain,
                              int apply_denoise);

/*!\brief Allocates a context that can be used for denoising and noise modeling.
 *
 * \param[in]  bit_depth   Bit depth of buffers this will be run on.
 * \param[in]  block_size  Block size for noise modeling and flat block
 *                         estimation
 * \param[in]  noise_level The noise_level (2.5 for moderate noise, and 5 for
 *                         higher levels of noise)
 */
struct aom_denoise_and_model_t *aom_denoise_and_model_alloc(int bit_depth,
                                                            int block_size,
                                                            float noise_level);

/*!\brief Frees the denoise context allocated with aom_denoise_and_model_alloc
 */
void aom_denoise_and_model_free(struct aom_denoise_and_model_t *denoise_model);

#ifdef __cplusplus
}  // extern "C"
#endif  // __cplusplus
#endif  // AOM_AOM_DSP_NOISE_MODEL_H_