vp9_denoiser.c 16.8 KB
Newer Older
Tim Kopp's avatar
Tim Kopp committed
1 2 3 4 5 6 7 8 9 10
/*
 *  Copyright (c) 2012 The WebM project authors. All Rights Reserved.
 *
 *  Use of this source code is governed by a BSD-style license
 *  that can be found in the LICENSE file in the root of the source
 *  tree. An additional intellectual property rights grant can be found
 *  in the file PATENTS.  All contributing project authors may
 *  be found in the AUTHORS file in the root of the source tree.
 */

Tim Kopp's avatar
Tim Kopp committed
11
#include <assert.h>
12
#include <limits.h>
Tim Kopp's avatar
Tim Kopp committed
13
#include "vpx_scale/yv12config.h"
Tim Kopp's avatar
Tim Kopp committed
14
#include "vpx/vpx_integer.h"
15
#include "vp9/common/vp9_reconinter.h"
16
#include "vp9/encoder/vp9_context_tree.h"
Tim Kopp's avatar
Tim Kopp committed
17
#include "vp9/encoder/vp9_denoiser.h"
Tim Kopp's avatar
Tim Kopp committed
18

19 20 21 22 23 24 25 26 27 28 29
/* The VP9 denoiser is a work-in-progress. It currently is only designed to work
 * with speed 6, though it (inexplicably) seems to also work with speed 5 (one
 * would need to modify the source code in vp9_pickmode.c and vp9_encoder.c to
 * make the calls to the vp9_denoiser_* functions when in speed 5).
 *
 * The implementation is very similar to that of the VP8 denoiser. While
 * choosing the motion vectors / reference frames, the denoiser is run, and if
 * it did not modify the signal to much, the denoised block is copied to the
 * signal.
 */

Tim Kopp's avatar
Tim Kopp committed
30 31 32 33
#ifdef OUTPUT_YUV_DENOISED
static void make_grayscale(YV12_BUFFER_CONFIG *yuv);
#endif

34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51
static int absdiff_thresh(BLOCK_SIZE bs, int increase_denoising) {
  (void)bs;
  return 3 + (increase_denoising ? 1 : 0);
}

static int delta_thresh(BLOCK_SIZE bs, int increase_denoising) {
  (void)bs;
  (void)increase_denoising;
  return 4;
}

static int noise_motion_thresh(BLOCK_SIZE bs, int increase_denoising) {
  (void)bs;
  (void)increase_denoising;
  return 25 * 25;
}

static unsigned int sse_thresh(BLOCK_SIZE bs, int increase_denoising) {
52 53 54
  return (4 << b_width_log2_lookup[bs]) *
         (4 << b_height_log2_lookup[bs]) *
         (increase_denoising ? 60 : 40);
55 56 57 58 59 60 61 62
}

static int sse_diff_thresh(BLOCK_SIZE bs, int increase_denoising,
                           int mv_row, int mv_col) {
  if (mv_row * mv_row + mv_col * mv_col >
      noise_motion_thresh(bs, increase_denoising)) {
    return 0;
  } else {
63 64
    return (4 << b_width_log2_lookup[bs]) *
           (4 << b_height_log2_lookup[bs]) * 20;
65 66 67
  }
}

68 69 70
int total_adj_strong_thresh(BLOCK_SIZE bs, int increase_denoising) {
  return (4 << b_width_log2_lookup[bs]) *
         (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 75
  return (4 << b_width_log2_lookup[bs]) *
         (4 << b_height_log2_lookup[bs]) * (increase_denoising ? 3 : 2);
76 77
}

78 79 80 81 82 83 84 85 86 87
// TODO(jackychen): If increase_denoising is enabled in the future,
// we might need to update the code for calculating 'total_adj' in
// case the C code is not bit-exact with corresponding sse2 code.
int vp9_denoiser_filter_c(const uint8_t *sig, int sig_stride,
                          const uint8_t *mc_avg,
                          int mc_avg_stride,
                          uint8_t *avg, int avg_stride,
                          int increase_denoising,
                          BLOCK_SIZE bs,
                          int motion_magnitude) {
Tim Kopp's avatar
Tim Kopp committed
88
  int r, c;
89 90 91 92
  const uint8_t *sig_start = sig;
  const uint8_t *mc_avg_start = mc_avg;
  uint8_t *avg_start = avg;
  int diff, adj, absdiff, delta;
93
  int adj_val[] = {3, 4, 6};
Tim Kopp's avatar
Tim Kopp committed
94
  int total_adj = 0;
95 96
  int shift_inc = 1;

97 98 99
  // If motion_magnitude is small, making the denoiser more aggressive by
  // increasing the adjustment for each level. Add another increment for
  // blocks that are labeled for increase denoising.
100 101 102 103 104 105 106 107
  if (motion_magnitude <= MOTION_MAGNITUDE_THRESHOLD) {
    if (increase_denoising) {
      shift_inc = 2;
    }
    adj_val[0] += shift_inc;
    adj_val[1] += shift_inc;
    adj_val[2] += shift_inc;
  }
Tim Kopp's avatar
Tim Kopp committed
108

109
  // First attempt to apply a strong temporal denoising filter.
110 111
  for (r = 0; r < (4 << b_height_log2_lookup[bs]); ++r) {
    for (c = 0; c < (4 << b_width_log2_lookup[bs]); ++c) {
Tim Kopp's avatar
Tim Kopp committed
112
      diff = mc_avg[c] - sig[c];
113
      absdiff = abs(diff);
Tim Kopp's avatar
Tim Kopp committed
114

115
      if (absdiff <= absdiff_thresh(bs, increase_denoising)) {
116 117
        avg[c] = mc_avg[c];
        total_adj += diff;
Tim Kopp's avatar
Tim Kopp committed
118
      } else {
119 120 121
        switch (absdiff) {
          case 4: case 5: case 6: case 7:
            adj = adj_val[0];
Tim Kopp's avatar
Tim Kopp committed
122
            break;
123 124 125
          case 8: case 9: case 10: case 11:
          case 12: case 13: case 14: case 15:
            adj = adj_val[1];
Tim Kopp's avatar
Tim Kopp committed
126
            break;
127 128 129 130 131 132 133 134 135 136
          default:
            adj = adj_val[2];
        }
        if (diff > 0) {
          avg[c] = MIN(UINT8_MAX, sig[c] + adj);
          total_adj += adj;
        } else {
          avg[c] = MAX(0, sig[c] - adj);
          total_adj -= adj;
        }
Tim Kopp's avatar
Tim Kopp committed
137 138 139 140 141 142
      }
    }
    sig += sig_stride;
    avg += avg_stride;
    mc_avg += mc_avg_stride;
  }
143 144 145 146 147 148 149 150

  // If the strong filter did not modify the signal too much, we're all set.
  if (abs(total_adj) <= total_adj_strong_thresh(bs, increase_denoising)) {
    return FILTER_BLOCK;
  }

  // Otherwise, we try to dampen the filter if the delta is not too high.
  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 155 156 157 158 159
    return COPY_BLOCK;
  }

  mc_avg =  mc_avg_start;
  avg = avg_start;
  sig = sig_start;
160 161
  for (r = 0; r < (4 << b_height_log2_lookup[bs]); ++r) {
    for (c = 0; c < (4 << b_width_log2_lookup[bs]); ++c) {
162 163 164 165 166 167
      diff = mc_avg[c] - sig[c];
      adj = abs(diff);
      if (adj > delta) {
        adj = delta;
      }
      if (diff > 0) {
JackyChen's avatar
JackyChen committed
168 169 170
        // Diff positive means we made positive adjustment above
        // (in first try/attempt), so now make negative adjustment to bring
        // denoised signal down.
171
        avg[c] = MAX(0, avg[c] - adj);
JackyChen's avatar
JackyChen committed
172
        total_adj -= adj;
173
      } else {
JackyChen's avatar
JackyChen committed
174 175 176
        // Diff negative means we made negative adjustment above
        // (in first try/attempt), so now make positive adjustment to bring
        // denoised signal up.
177
        avg[c] = MIN(UINT8_MAX, avg[c] + adj);
JackyChen's avatar
JackyChen committed
178
        total_adj += adj;
179 180 181 182 183 184 185 186 187 188 189 190
      }
    }
    sig += sig_stride;
    avg += avg_stride;
    mc_avg += mc_avg_stride;
  }

  // We can use the filter if it has been sufficiently dampened
  if (abs(total_adj) <= total_adj_weak_thresh(bs, increase_denoising)) {
    return FILTER_BLOCK;
  }
  return COPY_BLOCK;
Tim Kopp's avatar
Tim Kopp committed
191 192
}

Tim Kopp's avatar
Tim Kopp committed
193 194
static uint8_t *block_start(uint8_t *framebuf, int stride,
                            int mi_row, int mi_col) {
195
  return framebuf + (stride * mi_row * 8) + (mi_col * 8);
196 197 198
}


199 200 201 202 203
static VP9_DENOISER_DECISION perform_motion_compensation(VP9_DENOISER *denoiser,
                                                         MACROBLOCK *mb,
                                                         BLOCK_SIZE bs,
                                                         int increase_denoising,
                                                         int mi_row,
204
                                                         int mi_col,
205 206
                                                         PICK_MODE_CONTEXT *ctx,
                                                         int *motion_magnitude
207
                                                         ) {
208
  int mv_col, mv_row;
209
  int sse_diff = ctx->zeromv_sse - ctx->newmv_sse;
210 211
  MV_REFERENCE_FRAME frame;
  MACROBLOCKD *filter_mbd = &mb->e_mbd;
hkuang's avatar
hkuang committed
212
  MB_MODE_INFO *mbmi = &filter_mbd->mi[0].src_mi->mbmi;
213

Tim Kopp's avatar
Tim Kopp committed
214 215 216 217 218
  MB_MODE_INFO saved_mbmi;
  int i, j;
  struct buf_2d saved_dst[MAX_MB_PLANE];
  struct buf_2d saved_pre[MAX_MB_PLANE][2];  // 2 pre buffers

219
  // We will restore these after motion compensation.
Tim Kopp's avatar
Tim Kopp committed
220 221 222 223 224 225 226
  saved_mbmi = *mbmi;
  for (i = 0; i < MAX_MB_PLANE; ++i) {
    for (j = 0; j < 2; ++j) {
      saved_pre[i][j] = filter_mbd->plane[i].pre[j];
    }
    saved_dst[i] = filter_mbd->plane[i].dst;
  }
227

228 229
  mv_col = ctx->best_sse_mv.as_mv.col;
  mv_row = ctx->best_sse_mv.as_mv.row;
230

231 232
  *motion_magnitude = mv_row * mv_row + mv_col * mv_col;

233
  frame = ctx->best_reference_frame;
234 235 236

  // If the best reference frame uses inter-prediction and there is enough of a
  // difference in sum-squared-error, use it.
237 238
  if (frame != INTRA_FRAME &&
      sse_diff > sse_diff_thresh(bs, increase_denoising, mv_row, mv_col)) {
239 240 241
    mbmi->ref_frame[0] = ctx->best_reference_frame;
    mbmi->mode = ctx->best_sse_inter_mode;
    mbmi->mv[0] = ctx->best_sse_mv;
242 243
  } else {
    // Otherwise, use the zero reference frame.
244
    frame = ctx->best_zeromv_reference_frame;
245

246
    mbmi->ref_frame[0] = ctx->best_zeromv_reference_frame;
247 248 249
    mbmi->mode = ZEROMV;
    mbmi->mv[0].as_int = 0;

250 251 252
    ctx->best_sse_inter_mode = ZEROMV;
    ctx->best_sse_mv.as_int = 0;
    ctx->newmv_sse = ctx->zeromv_sse;
253 254 255 256
  }

  // Set the pointers in the MACROBLOCKD to point to the buffers in the denoiser
  // struct.
Tim Kopp's avatar
Tim Kopp committed
257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276
  for (j = 0; j < 2; ++j) {
    filter_mbd->plane[0].pre[j].buf =
        block_start(denoiser->running_avg_y[frame].y_buffer,
                    denoiser->running_avg_y[frame].y_stride,
                    mi_row, mi_col);
    filter_mbd->plane[0].pre[j].stride =
        denoiser->running_avg_y[frame].y_stride;
    filter_mbd->plane[1].pre[j].buf =
        block_start(denoiser->running_avg_y[frame].u_buffer,
                    denoiser->running_avg_y[frame].uv_stride,
                    mi_row, mi_col);
    filter_mbd->plane[1].pre[j].stride =
        denoiser->running_avg_y[frame].uv_stride;
    filter_mbd->plane[2].pre[j].buf =
        block_start(denoiser->running_avg_y[frame].v_buffer,
                    denoiser->running_avg_y[frame].uv_stride,
                    mi_row, mi_col);
    filter_mbd->plane[2].pre[j].stride =
        denoiser->running_avg_y[frame].uv_stride;
  }
277 278 279 280 281 282 283 284 285
  filter_mbd->plane[0].dst.buf =
      block_start(denoiser->mc_running_avg_y.y_buffer,
                  denoiser->mc_running_avg_y.y_stride,
                  mi_row, mi_col);
  filter_mbd->plane[0].dst.stride = denoiser->mc_running_avg_y.y_stride;
  filter_mbd->plane[1].dst.buf =
      block_start(denoiser->mc_running_avg_y.u_buffer,
                  denoiser->mc_running_avg_y.uv_stride,
                  mi_row, mi_col);
Tim Kopp's avatar
Tim Kopp committed
286
  filter_mbd->plane[1].dst.stride = denoiser->mc_running_avg_y.uv_stride;
287 288 289 290
  filter_mbd->plane[2].dst.buf =
      block_start(denoiser->mc_running_avg_y.v_buffer,
                  denoiser->mc_running_avg_y.uv_stride,
                  mi_row, mi_col);
Tim Kopp's avatar
Tim Kopp committed
291
  filter_mbd->plane[2].dst.stride = denoiser->mc_running_avg_y.uv_stride;
292 293 294 295 296

  vp9_build_inter_predictors_sby(filter_mbd, mv_row, mv_col, bs);

  // Restore everything to its original state
  *mbmi = saved_mbmi;
Tim Kopp's avatar
Tim Kopp committed
297 298 299 300 301 302
  for (i = 0; i < MAX_MB_PLANE; ++i) {
    for (j = 0; j < 2; ++j) {
      filter_mbd->plane[i].pre[j] = saved_pre[i][j];
    }
    filter_mbd->plane[i].dst = saved_dst[i];
  }
303

304 305
  mv_row = ctx->best_sse_mv.as_mv.row;
  mv_col = ctx->best_sse_mv.as_mv.col;
306

307
  if (ctx->newmv_sse > sse_thresh(bs, increase_denoising)) {
308 309 310 311
    return COPY_BLOCK;
  }
  if (mv_row * mv_row + mv_col * mv_col >
      8 * noise_motion_thresh(bs, increase_denoising)) {
312 313 314 315 316
    return COPY_BLOCK;
  }
  return FILTER_BLOCK;
}

317
void vp9_denoiser_denoise(VP9_DENOISER *denoiser, MACROBLOCK *mb,
318 319
                          int mi_row, int mi_col, BLOCK_SIZE bs,
                          PICK_MODE_CONTEXT *ctx) {
320
  int motion_magnitude = 0;
321
  VP9_DENOISER_DECISION decision = FILTER_BLOCK;
322
  YV12_BUFFER_CONFIG avg = denoiser->running_avg_y[INTRA_FRAME];
Tim Kopp's avatar
Tim Kopp committed
323 324 325 326
  YV12_BUFFER_CONFIG mc_avg = denoiser->mc_running_avg_y;
  uint8_t *avg_start = block_start(avg.y_buffer, avg.y_stride, mi_row, mi_col);
  uint8_t *mc_avg_start = block_start(mc_avg.y_buffer, mc_avg.y_stride,
                                          mi_row, mi_col);
327 328
  struct buf_2d src = mb->plane[0].src;

329 330
  decision = perform_motion_compensation(denoiser, mb, bs,
                                         denoiser->increase_denoising,
331 332
                                         mi_row, mi_col, ctx,
                                         &motion_magnitude);
333 334

  if (decision == FILTER_BLOCK) {
335 336 337 338
    decision = vp9_denoiser_filter(src.buf, src.stride,
                                 mc_avg_start, mc_avg.y_stride,
                                 avg_start, avg.y_stride,
                                 0, bs, motion_magnitude);
339
  }
340 341

  if (decision == FILTER_BLOCK) {
342 343 344 345
    vp9_convolve_copy(avg_start, avg.y_stride, src.buf, src.stride,
                      NULL, 0, NULL, 0,
                      num_4x4_blocks_wide_lookup[bs] << 2,
                      num_4x4_blocks_high_lookup[bs] << 2);
346
  } else {  // COPY_BLOCK
347 348 349 350
    vp9_convolve_copy(src.buf, src.stride, avg_start, avg.y_stride,
                      NULL, 0, NULL, 0,
                      num_4x4_blocks_wide_lookup[bs] << 2,
                      num_4x4_blocks_high_lookup[bs] << 2);
351
  }
Tim Kopp's avatar
Tim Kopp committed
352 353
}

Tim Kopp's avatar
Tim Kopp committed
354
static void copy_frame(YV12_BUFFER_CONFIG dest, const YV12_BUFFER_CONFIG src) {
Tim Kopp's avatar
Tim Kopp committed
355
  int r;
Tim Kopp's avatar
Tim Kopp committed
356
  const uint8_t *srcbuf = src.y_buffer;
357 358 359 360 361
  uint8_t *destbuf = dest.y_buffer;
  assert(dest.y_width == src.y_width);
  assert(dest.y_height == src.y_height);

  for (r = 0; r < dest.y_height; ++r) {
Tim Kopp's avatar
Tim Kopp committed
362
    vpx_memcpy(destbuf, srcbuf, dest.y_width);
363 364 365
    destbuf += dest.y_stride;
    srcbuf += src.y_stride;
  }
366 367
}

368 369 370 371 372 373 374 375 376
static void swap_frame_buffer(YV12_BUFFER_CONFIG dest,
                              YV12_BUFFER_CONFIG src) {
  uint8_t *tmp_buf = dest.y_buffer;
  assert(dest.y_width == src.y_width);
  assert(dest.y_height == src.y_height);
  dest.y_buffer = src.y_buffer;
  src.y_buffer = tmp_buf;
}

Tim Kopp's avatar
Tim Kopp committed
377
void vp9_denoiser_update_frame_info(VP9_DENOISER *denoiser,
378
                                    YV12_BUFFER_CONFIG src,
Tim Kopp's avatar
Tim Kopp committed
379 380 381 382
                                    FRAME_TYPE frame_type,
                                    int refresh_alt_ref_frame,
                                    int refresh_golden_frame,
                                    int refresh_last_frame) {
383
  if (frame_type == KEY_FRAME) {
Tim Kopp's avatar
Tim Kopp committed
384
    int i;
385
    // Start at 1 so as not to overwrite the INTRA_FRAME
386
    for (i = 1; i < MAX_REF_FRAMES; ++i)
387
      copy_frame(denoiser->running_avg_y[i], src);
388 389 390 391 392 393 394 395 396 397 398 399 400 401 402
    return;
  }

  /* For non key frames */
  if (refresh_alt_ref_frame) {
    swap_frame_buffer(denoiser->running_avg_y[ALTREF_FRAME],
                      denoiser->running_avg_y[INTRA_FRAME]);
  }
  if (refresh_golden_frame) {
    swap_frame_buffer(denoiser->running_avg_y[GOLDEN_FRAME],
                      denoiser->running_avg_y[INTRA_FRAME]);
  }
  if (refresh_last_frame) {
    swap_frame_buffer(denoiser->running_avg_y[LAST_FRAME],
                      denoiser->running_avg_y[INTRA_FRAME]);
403
  }
Tim Kopp's avatar
Tim Kopp committed
404 405
}

406 407 408
void vp9_denoiser_reset_frame_stats(PICK_MODE_CONTEXT *ctx) {
  ctx->zeromv_sse = UINT_MAX;
  ctx->newmv_sse = UINT_MAX;
409 410
}

411 412
void vp9_denoiser_update_frame_stats(MB_MODE_INFO *mbmi, unsigned int sse,
                                     PREDICTION_MODE mode,
413
                                     PICK_MODE_CONTEXT *ctx) {
414
  // TODO(tkopp): Use both MVs if possible
415 416 417
  if (mbmi->mv[0].as_int == 0 && sse < ctx->zeromv_sse) {
    ctx->zeromv_sse = sse;
    ctx->best_zeromv_reference_frame = mbmi->ref_frame[0];
418 419
  }

420 421 422 423 424
  if (mode == NEWMV) {
    ctx->newmv_sse = sse;
    ctx->best_sse_inter_mode = mode;
    ctx->best_sse_mv = mbmi->mv[0];
    ctx->best_reference_frame = mbmi->ref_frame[0];
425
  }
Tim Kopp's avatar
Tim Kopp committed
426 427 428
}

int vp9_denoiser_alloc(VP9_DENOISER *denoiser, int width, int height,
429 430 431 432 433
                       int ssx, int ssy,
#if CONFIG_VP9_HIGHBITDEPTH
                       int use_highbitdepth,
#endif
                       int border) {
434
  int i, fail;
435
  const int legacy_byte_alignment = 0;
Tim Kopp's avatar
Tim Kopp committed
436
  assert(denoiser != NULL);
437

438 439
  for (i = 0; i < MAX_REF_FRAMES; ++i) {
    fail = vp9_alloc_frame_buffer(&denoiser->running_avg_y[i], width, height,
440 441 442 443
                                  ssx, ssy,
#if CONFIG_VP9_HIGHBITDEPTH
                                  use_highbitdepth,
#endif
444
                                  border, legacy_byte_alignment);
445 446 447 448
    if (fail) {
      vp9_denoiser_free(denoiser);
      return 1;
    }
Tim Kopp's avatar
Tim Kopp committed
449 450 451
#ifdef OUTPUT_YUV_DENOISED
    make_grayscale(&denoiser->running_avg_y[i]);
#endif
452 453
  }

454
  fail = vp9_alloc_frame_buffer(&denoiser->mc_running_avg_y, width, height,
455 456 457 458
                                ssx, ssy,
#if CONFIG_VP9_HIGHBITDEPTH
                                use_highbitdepth,
#endif
459
                                border, legacy_byte_alignment);
460
  if (fail) {
461 462 463
    vp9_denoiser_free(denoiser);
    return 1;
  }
Tim Kopp's avatar
Tim Kopp committed
464 465 466
#ifdef OUTPUT_YUV_DENOISED
  make_grayscale(&denoiser->running_avg_y[i]);
#endif
467
  denoiser->increase_denoising = 0;
468
  denoiser->frame_buffer_initialized = 1;
469

Tim Kopp's avatar
Tim Kopp committed
470 471 472 473
  return 0;
}

void vp9_denoiser_free(VP9_DENOISER *denoiser) {
474
  int i;
475
  denoiser->frame_buffer_initialized = 0;
Tim Kopp's avatar
Tim Kopp committed
476 477 478
  if (denoiser == NULL) {
    return;
  }
479 480 481 482
  for (i = 0; i < MAX_REF_FRAMES; ++i) {
    if (&denoiser->running_avg_y[i] != NULL) {
      vp9_free_frame_buffer(&denoiser->running_avg_y[i]);
    }
483
  }
484 485
  if (&denoiser->mc_running_avg_y != NULL) {
    vp9_free_frame_buffer(&denoiser->mc_running_avg_y);
486
  }
Tim Kopp's avatar
Tim Kopp committed
487
}
Tim Kopp's avatar
Tim Kopp committed
488 489 490 491 492 493 494

#ifdef OUTPUT_YUV_DENOISED
static void make_grayscale(YV12_BUFFER_CONFIG *yuv) {
  int r, c;
  uint8_t *u = yuv->u_buffer;
  uint8_t *v = yuv->v_buffer;

495 496
  for (r = 0; r < yuv->uv_height; ++r) {
    for (c = 0; c < yuv->uv_width; ++c) {
Tim Kopp's avatar
Tim Kopp committed
497 498 499
      u[c] = UINT8_MAX / 2;
      v[c] = UINT8_MAX / 2;
    }
500 501
    u += yuv->uv_stride;
    v += yuv->uv_stride;
Tim Kopp's avatar
Tim Kopp committed
502 503 504
  }
}
#endif