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 #ifndef AOM_AOM_DSP_NOISE_MODEL_H_ 13 #define AOM_AOM_DSP_NOISE_MODEL_H_ 14 15 #ifdef __cplusplus 16 extern "C" { 17 #endif // __cplusplus 18 19 #include <stdint.h> 20 #include "aom_dsp/grain_synthesis.h" 21 #include "aom_scale/yv12config.h" 22 23 /*!\brief Wrapper of data required to represent linear system of eqns and soln. 24 */ 25 typedef struct { 26 double *A; 27 double *b; 28 double *x; 29 int n; 30 } aom_equation_system_t; 31 32 /*!\brief Representation of a piecewise linear curve 33 * 34 * Holds n points as (x, y) pairs, that store the curve. 35 */ 36 typedef struct { 37 double (*points)[2]; 38 int num_points; 39 } aom_noise_strength_lut_t; 40 41 /*!\brief Init the noise strength lut with the given number of points*/ 42 int aom_noise_strength_lut_init(aom_noise_strength_lut_t *lut, int num_points); 43 44 /*!\brief Frees the noise strength lut. */ 45 void aom_noise_strength_lut_free(aom_noise_strength_lut_t *lut); 46 47 /*!\brief Evaluate the lut at the point x. 48 * 49 * \param[in] lut The lut data. 50 * \param[in] x The coordinate to evaluate the lut. 51 */ 52 double aom_noise_strength_lut_eval(const aom_noise_strength_lut_t *lut, 53 double x); 54 55 /*!\brief Helper struct to model noise strength as a function of intensity. 56 * 57 * Internally, this structure holds a representation of a linear system 58 * of equations that models noise strength (standard deviation) as a 59 * function of intensity. The mapping is initially stored using a 60 * piecewise representation with evenly spaced bins that cover the entire 61 * domain from [min_intensity, max_intensity]. Each observation (x,y) gives a 62 * constraint of the form: 63 * y_{i} (1 - a) + y_{i+1} a = y 64 * where y_{i} is the value of bin i and x_{i} <= x <= x_{i+1} and 65 * a = x/(x_{i+1} - x{i}). The equation system holds the corresponding 66 * normal equations. 67 * 68 * As there may be missing data, the solution is regularized to get a 69 * complete set of values for the bins. A reduced representation after 70 * solving can be obtained by getting the corresponding noise_strength_lut_t. 71 */ 72 typedef struct { 73 aom_equation_system_t eqns; 74 double min_intensity; 75 double max_intensity; 76 int num_bins; 77 int num_equations; 78 double total; 79 } aom_noise_strength_solver_t; 80 81 /*!\brief Initializes the noise solver with the given number of bins. 82 * 83 * Returns 0 if initialization fails. 84 * 85 * \param[in] solver The noise solver to be initialized. 86 * \param[in] num_bins Number of bins to use in the internal representation. 87 * \param[in] bit_depth The bit depth used to derive {min,max}_intensity. 88 */ 89 int aom_noise_strength_solver_init(aom_noise_strength_solver_t *solver, 90 int num_bins, int bit_depth); 91 void aom_noise_strength_solver_free(aom_noise_strength_solver_t *solver); 92 93 /*!\brief Gets the x coordinate of bin i. 94 * 95 * \param[in] i The bin whose coordinate to query. 96 */ 97 double aom_noise_strength_solver_get_center( 98 const aom_noise_strength_solver_t *solver, int i); 99 100 /*!\brief Add an observation of the block mean intensity to its noise strength. 101 * 102 * \param[in] block_mean The average block intensity, 103 * \param[in] noise_std The observed noise strength. 104 */ 105 void aom_noise_strength_solver_add_measurement( 106 aom_noise_strength_solver_t *solver, double block_mean, double noise_std); 107 108 /*!\brief Solves the current set of equations for the noise strength. */ 109 int aom_noise_strength_solver_solve(aom_noise_strength_solver_t *solver); 110 111 /*!\brief Fits a reduced piecewise linear lut to the internal solution 112 * 113 * \param[in] max_num_points The maximum number of output points 114 * \param[out] lut The output piecewise linear lut. 115 */ 116 int aom_noise_strength_solver_fit_piecewise( 117 const aom_noise_strength_solver_t *solver, int max_num_points, 118 aom_noise_strength_lut_t *lut); 119 120 /*!\brief Helper for holding precomputed data for finding flat blocks. 121 * 122 * Internally a block is modeled with a low-order polynomial model. A 123 * planar model would be a bunch of equations like: 124 * <[y_i x_i 1], [a_1, a_2, a_3]> = b_i 125 * for each point in the block. The system matrix A with row i as [y_i x_i 1] 126 * is maintained as is the inverse, inv(A'*A), so that the plane parameters 127 * can be fit for each block. 128 */ 129 typedef struct { 130 double *AtA_inv; 131 double *A; 132 int num_params; // The number of parameters used for internal low-order model 133 int block_size; // The block size the finder was initialized with 134 double normalization; // Normalization factor (1 / (2^(bit_depth) - 1)) 135 int use_highbd; // Whether input data should be interpreted as uint16 136 } aom_flat_block_finder_t; 137 138 /*!\brief Init the block_finder with the given block size, bit_depth */ 139 int aom_flat_block_finder_init(aom_flat_block_finder_t *block_finder, 140 int block_size, int bit_depth, int use_highbd); 141 void aom_flat_block_finder_free(aom_flat_block_finder_t *block_finder); 142 143 /*!\brief Helper to extract a block and low order "planar" model. */ 144 void aom_flat_block_finder_extract_block( 145 const aom_flat_block_finder_t *block_finder, const uint8_t *const data, 146 int w, int h, int stride, int offsx, int offsy, double *plane, 147 double *block); 148 149 /*!\brief Runs the flat block finder on the input data. 150 * 151 * Find flat blocks in the input image data. Returns a map of 152 * flat_blocks, where the value of flat_blocks map will be non-zero 153 * when a block is determined to be flat. A higher value indicates a bigger 154 * confidence in the decision. 155 */ 156 int aom_flat_block_finder_run(const aom_flat_block_finder_t *block_finder, 157 const uint8_t *const data, int w, int h, 158 int stride, uint8_t *flat_blocks); 159 160 // The noise shape indicates the allowed coefficients in the AR model. 161 enum { 162 AOM_NOISE_SHAPE_DIAMOND = 0, 163 AOM_NOISE_SHAPE_SQUARE = 1 164 } UENUM1BYTE(aom_noise_shape); 165 166 // The parameters of the noise model include the shape type, lag, the 167 // bit depth of the input images provided, and whether the input images 168 // will be using uint16 (or uint8) representation. 169 typedef struct { 170 aom_noise_shape shape; 171 int lag; 172 int bit_depth; 173 int use_highbd; 174 } aom_noise_model_params_t; 175 176 /*!\brief State of a noise model estimate for a single channel. 177 * 178 * This contains a system of equations that can be used to solve 179 * for the auto-regressive coefficients as well as a noise strength 180 * solver that can be used to model noise strength as a function of 181 * intensity. 182 */ 183 typedef struct { 184 aom_equation_system_t eqns; 185 aom_noise_strength_solver_t strength_solver; 186 int num_observations; // The number of observations in the eqn system 187 double ar_gain; // The gain of the current AR filter 188 } aom_noise_state_t; 189 190 /*!\brief Complete model of noise for a planar video 191 * 192 * This includes a noise model for the latest frame and an aggregated 193 * estimate over all previous frames that had similar parameters. 194 */ 195 typedef struct { 196 aom_noise_model_params_t params; 197 aom_noise_state_t combined_state[3]; // Combined state per channel 198 aom_noise_state_t latest_state[3]; // Latest state per channel 199 int (*coords)[2]; // Offsets (x,y) of the coefficient samples 200 int n; // Number of parameters (size of coords) 201 int bit_depth; 202 } aom_noise_model_t; 203 204 /*!\brief Result of a noise model update. */ 205 enum { 206 AOM_NOISE_STATUS_OK = 0, 207 AOM_NOISE_STATUS_INVALID_ARGUMENT, 208 AOM_NOISE_STATUS_INSUFFICIENT_FLAT_BLOCKS, 209 AOM_NOISE_STATUS_DIFFERENT_NOISE_TYPE, 210 AOM_NOISE_STATUS_INTERNAL_ERROR, 211 } UENUM1BYTE(aom_noise_status_t); 212 213 /*!\brief Initializes a noise model with the given parameters. 214 * 215 * Returns 0 on failure. 216 */ 217 int aom_noise_model_init(aom_noise_model_t *model, 218 const aom_noise_model_params_t params); 219 void aom_noise_model_free(aom_noise_model_t *model); 220 221 /*!\brief Updates the noise model with a new frame observation. 222 * 223 * Updates the noise model with measurements from the given input frame and a 224 * denoised variant of it. Noise is sampled from flat blocks using the flat 225 * block map. 226 * 227 * Returns a noise_status indicating if the update was successful. If the 228 * Update was successful, the combined_state is updated with measurements from 229 * the provided frame. If status is OK or DIFFERENT_NOISE_TYPE, the latest noise 230 * state will be updated with measurements from the provided frame. 231 * 232 * \param[in,out] noise_model The noise model to be updated 233 * \param[in] data Raw frame data 234 * \param[in] denoised Denoised frame data. 235 * \param[in] w Frame width 236 * \param[in] h Frame height 237 * \param[in] strides Stride of the planes 238 * \param[in] chroma_sub_log2 Chroma subsampling for planes != 0. 239 * \param[in] flat_blocks A map to blocks that have been determined flat 240 * \param[in] block_size The size of blocks. 241 */ 242 aom_noise_status_t aom_noise_model_update( 243 aom_noise_model_t *const noise_model, const uint8_t *const data[3], 244 const uint8_t *const denoised[3], int w, int h, int strides[3], 245 int chroma_sub_log2[2], const uint8_t *const flat_blocks, int block_size); 246 247 /*\brief Save the "latest" estimate into the "combined" estimate. 248 * 249 * This is meant to be called when the noise modeling detected a change 250 * in parameters (or for example, if a user wanted to reset estimation at 251 * a shot boundary). 252 */ 253 void aom_noise_model_save_latest(aom_noise_model_t *noise_model); 254 255 /*!\brief Converts the noise_model parameters to the corresponding 256 * grain_parameters. 257 * 258 * The noise structs in this file are suitable for estimation (e.g., using 259 * floats), but the grain parameters in the bitstream are quantized. This 260 * function does the conversion by selecting the correct quantization levels. 261 */ 262 int aom_noise_model_get_grain_parameters(aom_noise_model_t *const noise_model, 263 aom_film_grain_t *film_grain); 264 265 /*!\brief Perform a Wiener filter denoising in 2D using the provided noise psd. 266 * 267 * \param[in] data Raw frame data 268 * \param[out] denoised Denoised frame data 269 * \param[in] w Frame width 270 * \param[in] h Frame height 271 * \param[in] stride Stride of the planes 272 * \param[in] chroma_sub_log2 Chroma subsampling for planes != 0. 273 * \param[in] noise_psd The power spectral density of the noise 274 * \param[in] block_size The size of blocks 275 * \param[in] bit_depth Bit depth of the image 276 * \param[in] use_highbd If true, uint8 pointers are interpreted as 277 * uint16 and stride is measured in uint16. 278 * This must be true when bit_depth >= 10. 279 */ 280 int aom_wiener_denoise_2d(const uint8_t *const data[3], uint8_t *denoised[3], 281 int w, int h, int stride[3], int chroma_sub_log2[2], 282 float *noise_psd[3], int block_size, int bit_depth, 283 int use_highbd); 284 285 struct aom_denoise_and_model_t; 286 287 /*!\brief Denoise the buffer and model the residual noise. 288 * 289 * This is meant to be called sequentially on input frames. The input buffer 290 * is denoised and the residual noise is modelled. The current noise estimate 291 * is populated in film_grain. Returns true on success. The grain.apply_grain 292 * parameter will be true when the input buffer was successfully denoised and 293 * grain was modelled. Returns false on error. 294 * 295 * \param[in] ctx Struct allocated with aom_denoise_and_model_alloc 296 * that holds some buffers for denoising and the current 297 * noise estimate. 298 * \param[in/out] buf The raw input buffer to be denoised. 299 * \param[out] grain Output film grain parameters 300 */ 301 int aom_denoise_and_model_run(struct aom_denoise_and_model_t *ctx, 302 YV12_BUFFER_CONFIG *buf, aom_film_grain_t *grain); 303 304 /*!\brief Allocates a context that can be used for denoising and noise modeling. 305 * 306 * \param[in] bit_depth Bit depth of buffers this will be run on. 307 * \param[in] block_size Block size for noise modeling and flat block 308 * estimation 309 * \param[in] noise_level The noise_level (2.5 for moderate noise, and 5 for 310 * higher levels of noise) 311 */ 312 struct aom_denoise_and_model_t *aom_denoise_and_model_alloc(int bit_depth, 313 int block_size, 314 float noise_level); 315 316 /*!\brief Frees the denoise context allocated with aom_denoise_and_model_alloc 317 */ 318 void aom_denoise_and_model_free(struct aom_denoise_and_model_t *denoise_model); 319 320 #ifdef __cplusplus 321 } // extern "C" 322 #endif // __cplusplus 323 #endif // AOM_AOM_DSP_NOISE_MODEL_H_ 324