Home | History | Annotate | Download | only in aom_dsp
      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