src/third_party/libvpx/source/libvpx/vp9/encoder/vp9_denoiser.c

   1 /*
   2  *  Copyright (c) 2012 The WebM project authors. All Rights Reserved.
   3  *
   4  *  Use of this source code is governed by a BSD-style license
   5  *  that can be found in the LICENSE file in the root of the source
   6  *  tree. An additional intellectual property rights grant can be found
   7  *  in the file PATENTS.  All contributing project authors may
   8  *  be found in the AUTHORS file in the root of the source tree.
   9  */
  10
  11 #include <assert.h>
  12 #include <limits.h>
  13 #include "vpx_scale/yv12config.h"
  14 #include "vpx/vpx_integer.h"
  15 #include "vp9/common/vp9_reconinter.h"
  16 #include "vp9/encoder/vp9_context_tree.h"
  17 #include "vp9/encoder/vp9_denoiser.h"
  18
  19 /* The VP9 denoiser is a work-in-progress. It currently is only designed to work
  20  * with speed 6, though it (inexplicably) seems to also work with speed 5 (one
  21  * would need to modify the source code in vp9_pickmode.c and vp9_encoder.c to
  22  * make the calls to the vp9_denoiser_* functions when in speed 5).
  23  *
  24  * The implementation is very similar to that of the VP8 denoiser. While
  25  * choosing the motion vectors / reference frames, the denoiser is run, and if
  26  * it did not modify the signal to much, the denoised block is copied to the
  27  * signal.
  28  */
  29
  30 #ifdef OUTPUT_YUV_DENOISED
  31 static void make_grayscale(YV12_BUFFER_CONFIG *yuv);
  32 #endif
  33
  34 static int absdiff_thresh(BLOCK_SIZE bs, int increase_denoising) {
  35   (void)bs;
  36   return 3 + (increase_denoising ? 1 : 0);
  37 }
  38
  39 static int delta_thresh(BLOCK_SIZE bs, int increase_denoising) {
  40   (void)bs;
  41   (void)increase_denoising;
  42   return 4;
  43 }
  44
  45 static int noise_motion_thresh(BLOCK_SIZE bs, int increase_denoising) {
  46   (void)bs;
  47   (void)increase_denoising;
  48   return 25 * 25;
  49 }
  50
  51 static unsigned int sse_thresh(BLOCK_SIZE bs, int increase_denoising) {
  52   return (4 << b_width_log2_lookup[bs]) *
  53          (4 << b_height_log2_lookup[bs]) *
  54          (increase_denoising ? 60 : 40);
  55 }
  56
  57 static int sse_diff_thresh(BLOCK_SIZE bs, int increase_denoising,
  58                            int mv_row, int mv_col) {
  59   if (mv_row * mv_row + mv_col * mv_col >
  60       noise_motion_thresh(bs, increase_denoising)) {
  61     return 0;
  62   } else {
  63     return (4 << b_width_log2_lookup[bs]) *
  64            (4 << b_height_log2_lookup[bs]) * 20;
  65   }
  66 }
  67
  68 int total_adj_strong_thresh(BLOCK_SIZE bs, int increase_denoising) {
  69   return (4 << b_width_log2_lookup[bs]) *
  70          (4 << b_height_log2_lookup[bs]) * (increase_denoising ? 3 : 2);
  71 }
  72
  73 static int total_adj_weak_thresh(BLOCK_SIZE bs, int increase_denoising) {
  74   return (4 << b_width_log2_lookup[bs]) *
  75          (4 << b_height_log2_lookup[bs]) * (increase_denoising ? 3 : 2);
  76 }
  77
  78 // TODO(jackychen): If increase_denoising is enabled in the future,
  79 // we might need to update the code for calculating 'total_adj' in
  80 // case the C code is not bit-exact with corresponding sse2 code.
  81 int vp9_denoiser_filter_c(const uint8_t *sig, int sig_stride,
  82                           const uint8_t *mc_avg,
  83                           int mc_avg_stride,
  84                           uint8_t *avg, int avg_stride,
  85                           int increase_denoising,
  86                           BLOCK_SIZE bs,
  87                           int motion_magnitude) {
  88   int r, c;
  89   const uint8_t *sig_start = sig;
  90   const uint8_t *mc_avg_start = mc_avg;
  91   uint8_t *avg_start = avg;
  92   int diff, adj, absdiff, delta;
  93   int adj_val[] = {3, 4, 6};
  94   int total_adj = 0;
  95   int shift_inc = 1;
  96
  97   // If motion_magnitude is small, making the denoiser more aggressive by
  98   // increasing the adjustment for each level. Add another increment for
  99   // blocks that are labeled for increase denoising.
 100   if (motion_magnitude <= MOTION_MAGNITUDE_THRESHOLD) {
 101     if (increase_denoising) {
 102       shift_inc = 2;
 103     }
 104     adj_val[0] += shift_inc;
 105     adj_val[1] += shift_inc;
 106     adj_val[2] += shift_inc;
 107   }
 108
 109   // First attempt to apply a strong temporal denoising filter.
 110   for (r = 0; r < (4 << b_height_log2_lookup[bs]); ++r) {
 111     for (c = 0; c < (4 << b_width_log2_lookup[bs]); ++c) {
 112       diff = mc_avg[c] - sig[c];
 113       absdiff = abs(diff);
 114
 115       if (absdiff <= absdiff_thresh(bs, increase_denoising)) {
 116         avg[c] = mc_avg[c];
 117         total_adj += diff;
 118       } else {
 119         switch (absdiff) {
 120           case 4: case 5: case 6: case 7:
 121             adj = adj_val[0];
 122             break;
 123           case 8: case 9: case 10: case 11:
 124           case 12: case 13: case 14: case 15:
 125             adj = adj_val[1];
 126             break;
 127           default:
 128             adj = adj_val[2];
 129         }
 130         if (diff > 0) {
 131           avg[c] = MIN(UINT8_MAX, sig[c] + adj);
 132           total_adj += adj;
 133         } else {
 134           avg[c] = MAX(0, sig[c] - adj);
 135           total_adj -= adj;
 136         }
 137       }
 138     }
 139     sig += sig_stride;
 140     avg += avg_stride;
 141     mc_avg += mc_avg_stride;
 142   }
 143
 144   // If the strong filter did not modify the signal too much, we're all set.
 145   if (abs(total_adj) <= total_adj_strong_thresh(bs, increase_denoising)) {
 146     return FILTER_BLOCK;
 147   }
 148
 149   // Otherwise, we try to dampen the filter if the delta is not too high.
 150   delta = ((abs(total_adj) - total_adj_strong_thresh(bs, increase_denoising))
 151            >> num_pels_log2_lookup[bs]) + 1;
 152
 153   if (delta >= delta_thresh(bs, increase_denoising)) {
 154     return COPY_BLOCK;
 155   }
 156
 157   mc_avg =  mc_avg_start;
 158   avg = avg_start;
 159   sig = sig_start;
 160   for (r = 0; r < (4 << b_height_log2_lookup[bs]); ++r) {
 161     for (c = 0; c < (4 << b_width_log2_lookup[bs]); ++c) {
 162       diff = mc_avg[c] - sig[c];
 163       adj = abs(diff);
 164       if (adj > delta) {
 165         adj = delta;
 166       }
 167       if (diff > 0) {
 168         // Diff positive means we made positive adjustment above
 169         // (in first try/attempt), so now make negative adjustment to bring
 170         // denoised signal down.
 171         avg[c] = MAX(0, avg[c] - adj);
 172         total_adj -= adj;
 173       } else {
 174         // Diff negative means we made negative adjustment above
 175         // (in first try/attempt), so now make positive adjustment to bring
 176         // denoised signal up.
 177         avg[c] = MIN(UINT8_MAX, avg[c] + adj);
 178         total_adj += adj;
 179       }
 180     }
 181     sig += sig_stride;
 182     avg += avg_stride;
 183     mc_avg += mc_avg_stride;
 184   }
 185
 186   // We can use the filter if it has been sufficiently dampened
 187   if (abs(total_adj) <= total_adj_weak_thresh(bs, increase_denoising)) {
 188     return FILTER_BLOCK;
 189   }
 190   return COPY_BLOCK;
 191 }
 192
 193 static uint8_t *block_start(uint8_t *framebuf, int stride,
 194                             int mi_row, int mi_col) {
 195   return framebuf + (stride * mi_row * 8) + (mi_col * 8);
 196 }
 197
 198 static void copy_block(uint8_t *dest, int dest_stride,
 199                        const uint8_t *src, int src_stride, BLOCK_SIZE bs) {
 200   int r;
 201   for (r = 0; r < (4 << b_height_log2_lookup[bs]); ++r) {
 202     vpx_memcpy(dest, src, (4 << b_width_log2_lookup[bs]));
 203     dest += dest_stride;
 204     src += src_stride;
 205   }
 206 }
 207
 208 static VP9_DENOISER_DECISION perform_motion_compensation(VP9_DENOISER *denoiser,
 209                                                          MACROBLOCK *mb,
 210                                                          BLOCK_SIZE bs,
 211                                                          int increase_denoising,
 212                                                          int mi_row,
 213                                                          int mi_col,
 214                                                          PICK_MODE_CONTEXT *ctx,
 215                                                          int *motion_magnitude
 216                                                          ) {
 217   int mv_col, mv_row;
 218   int sse_diff = ctx->zeromv_sse - ctx->newmv_sse;
 219   MV_REFERENCE_FRAME frame;
 220   MACROBLOCKD *filter_mbd = &mb->e_mbd;
 221   MB_MODE_INFO *mbmi = &filter_mbd->mi[0].src_mi->mbmi;
 222
 223   MB_MODE_INFO saved_mbmi;
 224   int i, j;
 225   struct buf_2d saved_dst[MAX_MB_PLANE];
 226   struct buf_2d saved_pre[MAX_MB_PLANE][2];  // 2 pre buffers
 227
 228   // We will restore these after motion compensation.
 229   saved_mbmi = *mbmi;
 230   for (i = 0; i < MAX_MB_PLANE; ++i) {
 231     for (j = 0; j < 2; ++j) {
 232       saved_pre[i][j] = filter_mbd->plane[i].pre[j];
 233     }
 234     saved_dst[i] = filter_mbd->plane[i].dst;
 235   }
 236
 237   mv_col = ctx->best_sse_mv.as_mv.col;
 238   mv_row = ctx->best_sse_mv.as_mv.row;
 239
 240   *motion_magnitude = mv_row * mv_row + mv_col * mv_col;
 241
 242   frame = ctx->best_reference_frame;
 243
 244   // If the best reference frame uses inter-prediction and there is enough of a
 245   // difference in sum-squared-error, use it.
 246   if (frame != INTRA_FRAME &&
 247       sse_diff > sse_diff_thresh(bs, increase_denoising, mv_row, mv_col)) {
 248     mbmi->ref_frame[0] = ctx->best_reference_frame;
 249     mbmi->mode = ctx->best_sse_inter_mode;
 250     mbmi->mv[0] = ctx->best_sse_mv;
 251   } else {
 252     // Otherwise, use the zero reference frame.
 253     frame = ctx->best_zeromv_reference_frame;
 254
 255     mbmi->ref_frame[0] = ctx->best_zeromv_reference_frame;
 256     mbmi->mode = ZEROMV;
 257     mbmi->mv[0].as_int = 0;
 258
 259     ctx->best_sse_inter_mode = ZEROMV;
 260     ctx->best_sse_mv.as_int = 0;
 261     ctx->newmv_sse = ctx->zeromv_sse;
 262   }
 263
 264   // Set the pointers in the MACROBLOCKD to point to the buffers in the denoiser
 265   // struct.
 266   for (j = 0; j < 2; ++j) {
 267     filter_mbd->plane[0].pre[j].buf =
 268         block_start(denoiser->running_avg_y[frame].y_buffer,
 269                     denoiser->running_avg_y[frame].y_stride,
 270                     mi_row, mi_col);
 271     filter_mbd->plane[0].pre[j].stride =
 272         denoiser->running_avg_y[frame].y_stride;
 273     filter_mbd->plane[1].pre[j].buf =
 274         block_start(denoiser->running_avg_y[frame].u_buffer,
 275                     denoiser->running_avg_y[frame].uv_stride,
 276                     mi_row, mi_col);
 277     filter_mbd->plane[1].pre[j].stride =
 278         denoiser->running_avg_y[frame].uv_stride;
 279     filter_mbd->plane[2].pre[j].buf =
 280         block_start(denoiser->running_avg_y[frame].v_buffer,
 281                     denoiser->running_avg_y[frame].uv_stride,
 282                     mi_row, mi_col);
 283     filter_mbd->plane[2].pre[j].stride =
 284         denoiser->running_avg_y[frame].uv_stride;
 285   }
 286   filter_mbd->plane[0].dst.buf =
 287       block_start(denoiser->mc_running_avg_y.y_buffer,
 288                   denoiser->mc_running_avg_y.y_stride,
 289                   mi_row, mi_col);
 290   filter_mbd->plane[0].dst.stride = denoiser->mc_running_avg_y.y_stride;
 291   filter_mbd->plane[1].dst.buf =
 292       block_start(denoiser->mc_running_avg_y.u_buffer,
 293                   denoiser->mc_running_avg_y.uv_stride,
 294                   mi_row, mi_col);
 295   filter_mbd->plane[1].dst.stride = denoiser->mc_running_avg_y.uv_stride;
 296   filter_mbd->plane[2].dst.buf =
 297       block_start(denoiser->mc_running_avg_y.v_buffer,
 298                   denoiser->mc_running_avg_y.uv_stride,
 299                   mi_row, mi_col);
 300   filter_mbd->plane[2].dst.stride = denoiser->mc_running_avg_y.uv_stride;
 301
 302   vp9_build_inter_predictors_sby(filter_mbd, mv_row, mv_col, bs);
 303
 304   // Restore everything to its original state
 305   *mbmi = saved_mbmi;
 306   for (i = 0; i < MAX_MB_PLANE; ++i) {
 307     for (j = 0; j < 2; ++j) {
 308       filter_mbd->plane[i].pre[j] = saved_pre[i][j];
 309     }
 310     filter_mbd->plane[i].dst = saved_dst[i];
 311   }
 312
 313   mv_row = ctx->best_sse_mv.as_mv.row;
 314   mv_col = ctx->best_sse_mv.as_mv.col;
 315
 316   if (ctx->newmv_sse > sse_thresh(bs, increase_denoising)) {
 317     return COPY_BLOCK;
 318   }
 319   if (mv_row * mv_row + mv_col * mv_col >
 320       8 * noise_motion_thresh(bs, increase_denoising)) {
 321     return COPY_BLOCK;
 322   }
 323   return FILTER_BLOCK;
 324 }
 325
 326 void vp9_denoiser_denoise(VP9_DENOISER *denoiser, MACROBLOCK *mb,
 327                           int mi_row, int mi_col, BLOCK_SIZE bs,
 328                           PICK_MODE_CONTEXT *ctx) {
 329   int motion_magnitude = 0;
 330   VP9_DENOISER_DECISION decision = FILTER_BLOCK;
 331   YV12_BUFFER_CONFIG avg = denoiser->running_avg_y[INTRA_FRAME];
 332   YV12_BUFFER_CONFIG mc_avg = denoiser->mc_running_avg_y;
 333   uint8_t *avg_start = block_start(avg.y_buffer, avg.y_stride, mi_row, mi_col);
 334   uint8_t *mc_avg_start = block_start(mc_avg.y_buffer, mc_avg.y_stride,
 335                                           mi_row, mi_col);
 336   struct buf_2d src = mb->plane[0].src;
 337
 338   decision = perform_motion_compensation(denoiser, mb, bs,
 339                                          denoiser->increase_denoising,
 340                                          mi_row, mi_col, ctx,
 341                                          &motion_magnitude);
 342
 343   if (decision == FILTER_BLOCK) {
 344     decision = vp9_denoiser_filter(src.buf, src.stride,
 345                                  mc_avg_start, mc_avg.y_stride,
 346                                  avg_start, avg.y_stride,
 347                                  0, bs, motion_magnitude);
 348   }
 349
 350   if (decision == FILTER_BLOCK) {
 351     copy_block(src.buf, src.stride, avg_start, avg.y_stride, bs);
 352   } else {  // COPY_BLOCK
 353     copy_block(avg_start, avg.y_stride, src.buf, src.stride, bs);
 354   }
 355 }
 356
 357 static void copy_frame(YV12_BUFFER_CONFIG dest, const YV12_BUFFER_CONFIG src) {
 358   int r;
 359   const uint8_t *srcbuf = src.y_buffer;
 360   uint8_t *destbuf = dest.y_buffer;
 361   assert(dest.y_width == src.y_width);
 362   assert(dest.y_height == src.y_height);
 363
 364   for (r = 0; r < dest.y_height; ++r) {
 365     vpx_memcpy(destbuf, srcbuf, dest.y_width);
 366     destbuf += dest.y_stride;
 367     srcbuf += src.y_stride;
 368   }
 369 }
 370
 371 void vp9_denoiser_update_frame_info(VP9_DENOISER *denoiser,
 372                                     YV12_BUFFER_CONFIG src,
 373                                     FRAME_TYPE frame_type,
 374                                     int refresh_alt_ref_frame,
 375                                     int refresh_golden_frame,
 376                                     int refresh_last_frame) {
 377   if (frame_type == KEY_FRAME) {
 378     int i;
 379     // Start at 1 so as not to overwrite the INTRA_FRAME
 380     for (i = 1; i < MAX_REF_FRAMES; ++i) {
 381       copy_frame(denoiser->running_avg_y[i], src);
 382     }
 383   } else {  /* For non key frames */
 384     if (refresh_alt_ref_frame) {
 385       copy_frame(denoiser->running_avg_y[ALTREF_FRAME],
 386                  denoiser->running_avg_y[INTRA_FRAME]);
 387     }
 388     if (refresh_golden_frame) {
 389       copy_frame(denoiser->running_avg_y[GOLDEN_FRAME],
 390                  denoiser->running_avg_y[INTRA_FRAME]);
 391     }
 392     if (refresh_last_frame) {
 393       copy_frame(denoiser->running_avg_y[LAST_FRAME],
 394                  denoiser->running_avg_y[INTRA_FRAME]);
 395     }
 396   }
 397 }
 398
 399 void vp9_denoiser_reset_frame_stats(PICK_MODE_CONTEXT *ctx) {
 400   ctx->zeromv_sse = UINT_MAX;
 401   ctx->newmv_sse = UINT_MAX;
 402 }
 403
 404 void vp9_denoiser_update_frame_stats(MB_MODE_INFO *mbmi, unsigned int sse,
 405                                      PREDICTION_MODE mode,
 406                                      PICK_MODE_CONTEXT *ctx) {
 407   // TODO(tkopp): Use both MVs if possible
 408   if (mbmi->mv[0].as_int == 0 && sse < ctx->zeromv_sse) {
 409     ctx->zeromv_sse = sse;
 410     ctx->best_zeromv_reference_frame = mbmi->ref_frame[0];
 411   }
 412
 413   if (mode == NEWMV) {
 414     ctx->newmv_sse = sse;
 415     ctx->best_sse_inter_mode = mode;
 416     ctx->best_sse_mv = mbmi->mv[0];
 417     ctx->best_reference_frame = mbmi->ref_frame[0];
 418   }
 419 }
 420
 421 int vp9_denoiser_alloc(VP9_DENOISER *denoiser, int width, int height,
 422                        int ssx, int ssy,
 423 #if CONFIG_VP9_HIGHBITDEPTH
 424                        int use_highbitdepth,
 425 #endif
 426                        int border) {
 427   int i, fail;
 428   assert(denoiser != NULL);
 429
 430   for (i = 0; i < MAX_REF_FRAMES; ++i) {
 431     fail = vp9_alloc_frame_buffer(&denoiser->running_avg_y[i], width, height,
 432                                   ssx, ssy,
 433 #if CONFIG_VP9_HIGHBITDEPTH
 434                                   use_highbitdepth,
 435 #endif
 436                                   border);
 437     if (fail) {
 438       vp9_denoiser_free(denoiser);
 439       return 1;
 440     }
 441 #ifdef OUTPUT_YUV_DENOISED
 442     make_grayscale(&denoiser->running_avg_y[i]);
 443 #endif
 444   }
 445
 446   fail = vp9_alloc_frame_buffer(&denoiser->mc_running_avg_y, width, height,
 447                                 ssx, ssy,
 448 #if CONFIG_VP9_HIGHBITDEPTH
 449                                 use_highbitdepth,
 450 #endif
 451                                 border);
 452   if (fail) {
 453     vp9_denoiser_free(denoiser);
 454     return 1;
 455   }
 456 #ifdef OUTPUT_YUV_DENOISED
 457   make_grayscale(&denoiser->running_avg_y[i]);
 458 #endif
 459   denoiser->increase_denoising = 0;
 460
 461   return 0;
 462 }
 463
 464 void vp9_denoiser_free(VP9_DENOISER *denoiser) {
 465   int i;
 466   if (denoiser == NULL) {
 467     return;
 468   }
 469   for (i = 0; i < MAX_REF_FRAMES; ++i) {
 470     if (&denoiser->running_avg_y[i] != NULL) {
 471       vp9_free_frame_buffer(&denoiser->running_avg_y[i]);
 472     }
 473   }
 474   if (&denoiser->mc_running_avg_y != NULL) {
 475     vp9_free_frame_buffer(&denoiser->mc_running_avg_y);
 476   }
 477 }
 478
 479 #ifdef OUTPUT_YUV_DENOISED
 480 static void make_grayscale(YV12_BUFFER_CONFIG *yuv) {
 481   int r, c;
 482   uint8_t *u = yuv->u_buffer;
 483   uint8_t *v = yuv->v_buffer;
 484
 485   // The '/2's are there because we have a 440 buffer, but we want to output
 486   // 420.
 487   for (r = 0; r < yuv->uv_height / 2; ++r) {
 488     for (c = 0; c < yuv->uv_width / 2; ++c) {
 489       u[c] = UINT8_MAX / 2;
 490       v[c] = UINT8_MAX / 2;
 491     }
 492     u += yuv->uv_stride + yuv->uv_width / 2;
 493     v += yuv->uv_stride + yuv->uv_width / 2;
 494   }
 495 }
 496 #endif