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