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