libavfilter/vf_fftdnoiz.c

   1 /*
   2  * This file is part of FFmpeg.
   3  *
   4  * FFmpeg is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU Lesser General Public
   6  * License as published by the Free Software Foundation; either
   7  * version 2.1 of the License, or (at your option) any later version.
   8  *
   9  * FFmpeg is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12  * Lesser General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU Lesser General Public
  15  * License along with FFmpeg; if not, write to the Free Software
  16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  17  */
  18
  19 #include <float.h>
  20
  21 #include "libavutil/common.h"
  22 #include "libavutil/imgutils.h"
  23 #include "libavutil/mem.h"
  24 #include "libavutil/opt.h"
  25 #include "libavutil/pixdesc.h"
  26 #include "libavutil/tx.h"
  27
  28 #include "filters.h"
  29 #include "video.h"
  30 #include "window_func.h"
  31
  32 #define MAX_BLOCK 256
  33 #define MAX_THREADS 32
  34
  35 enum BufferTypes {
  36     CURRENT,
  37     PREV,
  38     NEXT,
  39     BSIZE
  40 };
  41
  42 typedef struct PlaneContext {
  43     int planewidth, planeheight;
  44     int nox, noy;
  45     int b;
  46     int o;
  47     float n;
  48
  49     float *buffer[MAX_THREADS][BSIZE];
  50     AVComplexFloat *hdata[MAX_THREADS], *vdata[MAX_THREADS];
  51     AVComplexFloat *hdata_out[MAX_THREADS], *vdata_out[MAX_THREADS];
  52     int data_linesize;
  53     int buffer_linesize;
  54 } PlaneContext;
  55
  56 typedef struct FFTdnoizContext {
  57     const AVClass *class;
  58
  59     float sigma;
  60     float amount;
  61     int   block_size;
  62     float overlap;
  63     int   method;
  64     int   window;
  65     int   nb_prev;
  66     int   nb_next;
  67     int   planesf;
  68
  69     AVFrame *prev, *cur, *next;
  70
  71     int depth;
  72     int nb_planes;
  73     int nb_threads;
  74     PlaneContext planes[4];
  75     float win[MAX_BLOCK][MAX_BLOCK];
  76
  77     AVTXContext *fft[MAX_THREADS], *ifft[MAX_THREADS];
  78     AVTXContext *fft_r[MAX_THREADS], *ifft_r[MAX_THREADS];
  79
  80     av_tx_fn tx_fn, itx_fn;
  81     av_tx_fn tx_r_fn, itx_r_fn;
  82
  83     void (*import_row)(AVComplexFloat *dst, uint8_t *src, int rw, float scale, float *win, int off);
  84     void (*export_row)(AVComplexFloat *src, uint8_t *dst, int rw, int depth, float *win);
  85 } FFTdnoizContext;
  86
  87 #define OFFSET(x) offsetof(FFTdnoizContext, x)
  88 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
  89 #define TFLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_RUNTIME_PARAM
  90 static const AVOption fftdnoiz_options[] = {
  91     { "sigma",   "set denoise strength",
  92         OFFSET(sigma),      AV_OPT_TYPE_FLOAT, {.dbl=1},        0, 100, .flags = TFLAGS },
  93     { "amount",  "set amount of denoising",
  94         OFFSET(amount),     AV_OPT_TYPE_FLOAT, {.dbl=1},     0.01,   1, .flags = TFLAGS },
  95     { "block",   "set block size",
  96         OFFSET(block_size), AV_OPT_TYPE_INT,   {.i64=32}, 8, MAX_BLOCK, .flags = FLAGS },
  97     { "overlap", "set block overlap",
  98         OFFSET(overlap),    AV_OPT_TYPE_FLOAT, {.dbl=0.5},    0.2, 0.8, .flags = FLAGS },
  99     { "method",  "set method of denoising",
 100         OFFSET(method),     AV_OPT_TYPE_INT,   {.i64=0},        0,   1, .flags = TFLAGS, .unit = "method" },
 101     { "wiener", "wiener method",
 102         0,                  AV_OPT_TYPE_CONST, {.i64=0},        0,   0, .flags = TFLAGS, .unit = "method" },
 103     { "hard",   "hard thresholding",
 104         0,                  AV_OPT_TYPE_CONST, {.i64=1},        0,   0, .flags = TFLAGS, .unit = "method" },
 105     { "prev",    "set number of previous frames for temporal denoising",
 106         OFFSET(nb_prev),    AV_OPT_TYPE_INT,   {.i64=0},        0,   1, .flags = FLAGS },
 107     { "next",    "set number of next frames for temporal denoising",
 108         OFFSET(nb_next),    AV_OPT_TYPE_INT,   {.i64=0},        0,   1, .flags = FLAGS },
 109     { "planes",  "set planes to filter",
 110         OFFSET(planesf),    AV_OPT_TYPE_INT,   {.i64=7},        0,  15, .flags = TFLAGS },
 111     WIN_FUNC_OPTION("window", OFFSET(window), FLAGS, WFUNC_HANNING),
 112     { NULL }
 113 };
 114
 115 AVFILTER_DEFINE_CLASS(fftdnoiz);
 116
 117 static const enum AVPixelFormat pix_fmts[] = {
 118     AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAY9,
 119     AV_PIX_FMT_GRAY10, AV_PIX_FMT_GRAY12,
 120     AV_PIX_FMT_GRAY14, AV_PIX_FMT_GRAY16,
 121     AV_PIX_FMT_YUV410P, AV_PIX_FMT_YUV411P,
 122     AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P,
 123     AV_PIX_FMT_YUV440P, AV_PIX_FMT_YUV444P,
 124     AV_PIX_FMT_YUVJ420P, AV_PIX_FMT_YUVJ422P,
 125     AV_PIX_FMT_YUVJ440P, AV_PIX_FMT_YUVJ444P,
 126     AV_PIX_FMT_YUVJ411P,
 127     AV_PIX_FMT_YUV420P9, AV_PIX_FMT_YUV422P9, AV_PIX_FMT_YUV444P9,
 128     AV_PIX_FMT_YUV420P10, AV_PIX_FMT_YUV422P10, AV_PIX_FMT_YUV444P10,
 129     AV_PIX_FMT_YUV440P10,
 130     AV_PIX_FMT_YUV444P12, AV_PIX_FMT_YUV422P12, AV_PIX_FMT_YUV420P12,
 131     AV_PIX_FMT_YUV440P12,
 132     AV_PIX_FMT_YUV444P14, AV_PIX_FMT_YUV422P14, AV_PIX_FMT_YUV420P14,
 133     AV_PIX_FMT_YUV420P16, AV_PIX_FMT_YUV422P16, AV_PIX_FMT_YUV444P16,
 134     AV_PIX_FMT_GBRP, AV_PIX_FMT_GBRP9, AV_PIX_FMT_GBRP10,
 135     AV_PIX_FMT_GBRP12, AV_PIX_FMT_GBRP14, AV_PIX_FMT_GBRP16,
 136     AV_PIX_FMT_YUVA420P,  AV_PIX_FMT_YUVA422P,   AV_PIX_FMT_YUVA444P,
 137     AV_PIX_FMT_YUVA444P9, AV_PIX_FMT_YUVA444P10, AV_PIX_FMT_YUVA444P12, AV_PIX_FMT_YUVA444P16,
 138     AV_PIX_FMT_YUVA422P9, AV_PIX_FMT_YUVA422P10, AV_PIX_FMT_YUVA422P12, AV_PIX_FMT_YUVA422P16,
 139     AV_PIX_FMT_YUVA420P9, AV_PIX_FMT_YUVA420P10, AV_PIX_FMT_YUVA420P16,
 140     AV_PIX_FMT_GBRAP,     AV_PIX_FMT_GBRAP10,    AV_PIX_FMT_GBRAP12,    AV_PIX_FMT_GBRAP16,
 141     AV_PIX_FMT_NONE
 142 };
 143
 144 typedef struct ThreadData {
 145     float *src, *dst;
 146 } ThreadData;
 147
 148 static void import_row8(AVComplexFloat *dst, uint8_t *src, int rw,
 149                         float scale, float *win, int off)
 150 {
 151     for (int j = 0; j < rw; j++) {
 152         const int i = abs(j + off);
 153         dst[j].re = src[i] * scale * win[j];
 154         dst[j].im = 0.f;
 155     }
 156 }
 157
 158 static void export_row8(AVComplexFloat *src, uint8_t *dst, int rw, int depth, float *win)
 159 {
 160     for (int j = 0; j < rw; j++)
 161         dst[j] = av_clip_uint8(lrintf(src[j].re / win[j]));
 162 }
 163
 164 static void import_row16(AVComplexFloat *dst, uint8_t *srcp, int rw,
 165                          float scale, float *win, int off)
 166 {
 167     uint16_t *src = (uint16_t *)srcp;
 168
 169     for (int j = 0; j < rw; j++) {
 170         const int i = abs(j + off);
 171         dst[j].re = src[i] * scale * win[j];
 172         dst[j].im = 0;
 173     }
 174 }
 175
 176 static void export_row16(AVComplexFloat *src, uint8_t *dstp, int rw, int depth, float *win)
 177 {
 178     uint16_t *dst = (uint16_t *)dstp;
 179
 180     for (int j = 0; j < rw; j++)
 181         dst[j] = av_clip_uintp2_c(lrintf(src[j].re / win[j]), depth);
 182 }
 183
 184 static int config_input(AVFilterLink *inlink)
 185 {
 186     AVFilterContext *ctx = inlink->dst;
 187     const AVPixFmtDescriptor *desc;
 188     FFTdnoizContext *s = ctx->priv;
 189     float lut[MAX_BLOCK + 1];
 190     float overlap;
 191     int i;
 192
 193     desc = av_pix_fmt_desc_get(inlink->format);
 194     s->depth = desc->comp[0].depth;
 195
 196     if (s->depth <= 8) {
 197         s->import_row = import_row8;
 198         s->export_row = export_row8;
 199     } else {
 200         s->import_row = import_row16;
 201         s->export_row = export_row16;
 202     }
 203
 204     s->planes[1].planewidth = s->planes[2].planewidth = AV_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
 205     s->planes[0].planewidth = s->planes[3].planewidth = inlink->w;
 206     s->planes[1].planeheight = s->planes[2].planeheight = AV_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h);
 207     s->planes[0].planeheight = s->planes[3].planeheight = inlink->h;
 208
 209     s->nb_planes = av_pix_fmt_count_planes(inlink->format);
 210     s->nb_threads = FFMIN(ff_filter_get_nb_threads(ctx), MAX_THREADS);
 211
 212     for (int i = 0; i < s->nb_threads; i++) {
 213         float scale = 1.f, iscale = 1.f;
 214         int ret;
 215
 216         if ((ret = av_tx_init(&s->fft[i],    &s->tx_fn,    AV_TX_FLOAT_FFT,
 217                               0, s->block_size,               &scale,  0)) < 0 ||
 218             (ret = av_tx_init(&s->ifft[i],   &s->itx_fn,   AV_TX_FLOAT_FFT,
 219                               1, s->block_size,               &iscale, 0)) < 0 ||
 220             (ret = av_tx_init(&s->fft_r[i],  &s->tx_r_fn,  AV_TX_FLOAT_FFT,
 221                               0, 1 + s->nb_prev + s->nb_next, &scale,  0)) < 0 ||
 222             (ret = av_tx_init(&s->ifft_r[i], &s->itx_r_fn, AV_TX_FLOAT_FFT,
 223                               1, 1 + s->nb_prev + s->nb_next, &iscale, 0)) < 0)
 224             return ret;
 225     }
 226
 227     for (i = 0; i < s->nb_planes; i++) {
 228         PlaneContext *p = &s->planes[i];
 229         int size;
 230
 231         p->b = s->block_size;
 232         p->n = 1.f / (p->b * p->b);
 233         p->o = lrintf(p->b * s->overlap);
 234         size = p->b - p->o;
 235         p->nox = (p->planewidth  + (size - 1)) / size;
 236         p->noy = (p->planeheight + (size - 1)) / size;
 237
 238         av_log(ctx, AV_LOG_DEBUG, "nox:%d noy:%d size:%d\n", p->nox, p->noy, size);
 239
 240         p->buffer_linesize = p->b * sizeof(AVComplexFloat);
 241         p->data_linesize = 2 * p->b * sizeof(float);
 242         for (int j = 0; j < s->nb_threads; j++) {
 243             p->hdata[j] = av_calloc(p->b, p->data_linesize);
 244             p->hdata_out[j] = av_calloc(p->b, p->data_linesize);
 245             p->vdata[j] = av_calloc(p->b, p->data_linesize);
 246             p->vdata_out[j] = av_calloc(p->b, p->data_linesize);
 247             p->buffer[j][CURRENT] = av_calloc(p->b, p->buffer_linesize);
 248             if (!p->buffer[j][CURRENT])
 249                 return AVERROR(ENOMEM);
 250             if (s->nb_prev > 0) {
 251                 p->buffer[j][PREV] = av_calloc(p->b, p->buffer_linesize);
 252                 if (!p->buffer[j][PREV])
 253                     return AVERROR(ENOMEM);
 254             }
 255             if (s->nb_next > 0) {
 256                 p->buffer[j][NEXT] = av_calloc(p->b, p->buffer_linesize);
 257                 if (!p->buffer[j][NEXT])
 258                     return AVERROR(ENOMEM);
 259             }
 260             if (!p->hdata[j] || !p->vdata[j] ||
 261                 !p->hdata_out[j] || !p->vdata_out[j])
 262                 return AVERROR(ENOMEM);
 263         }
 264     }
 265
 266     generate_window_func(lut, s->block_size + 1, s->window, &overlap);
 267
 268     for (int y = 0; y < s->block_size; y++) {
 269         for (int x = 0; x < s->block_size; x++)
 270             s->win[y][x] = lut[y] * lut[x];
 271     }
 272
 273     return 0;
 274 }
 275
 276 static void import_block(FFTdnoizContext *s,
 277                          uint8_t *srcp, int src_linesize,
 278                          float *buffer, int buffer_linesize, int plane,
 279                          int jobnr, int y, int x)
 280 {
 281     PlaneContext *p = &s->planes[plane];
 282     const int width = p->planewidth;
 283     const int height = p->planeheight;
 284     const int block = p->b;
 285     const int overlap = p->o;
 286     const int hoverlap = overlap / 2;
 287     const int size = block - overlap;
 288     const int bpp = (s->depth + 7) / 8;
 289     const int data_linesize = p->data_linesize / sizeof(AVComplexFloat);
 290     const float scale = 1.f / ((1.f + s->nb_prev + s->nb_next) * s->block_size * s->block_size);
 291     AVComplexFloat *hdata = p->hdata[jobnr];
 292     AVComplexFloat *hdata_out = p->hdata_out[jobnr];
 293     AVComplexFloat *vdata_out = p->vdata_out[jobnr];
 294     const int woff = -hoverlap;
 295     const int hoff = -hoverlap;
 296     const int rh = FFMIN(block, height - y * size + hoverlap);
 297     const int rw = FFMIN(block, width  - x * size + hoverlap);
 298     AVComplexFloat *ssrc, *ddst, *dst = hdata, *dst_out = hdata_out;
 299     float *bdst = buffer;
 300
 301     buffer_linesize /= sizeof(float);
 302
 303     for (int i = 0; i < rh; i++) {
 304         uint8_t *src = srcp + src_linesize * abs(y * size + i + hoff) + x * size * bpp;
 305
 306         s->import_row(dst, src, rw, scale, s->win[i], woff);
 307         for (int j = rw; j < block; j++) {
 308             dst[j].re = dst[rw - 1].re;
 309             dst[j].im = 0.f;
 310         }
 311         s->tx_fn(s->fft[jobnr], dst_out, dst, sizeof(AVComplexFloat));
 312
 313         ddst = dst_out;
 314         dst += data_linesize;
 315         dst_out += data_linesize;
 316     }
 317
 318     dst = dst_out;
 319     for (int i = rh; i < block; i++) {
 320         for (int j = 0; j < block; j++) {
 321             dst[j].re = ddst[j].re;
 322             dst[j].im = ddst[j].im;
 323         }
 324
 325         dst += data_linesize;
 326     }
 327
 328     ssrc = hdata_out;
 329     dst = vdata_out;
 330     for (int i = 0; i < block; i++) {
 331         for (int j = 0; j < block; j++)
 332             dst[j] = ssrc[j * data_linesize + i];
 333         s->tx_fn(s->fft[jobnr], bdst, dst, sizeof(AVComplexFloat));
 334
 335         dst += data_linesize;
 336         bdst += buffer_linesize;
 337     }
 338 }
 339
 340 static void export_block(FFTdnoizContext *s,
 341                          uint8_t *dstp, int dst_linesize,
 342                          float *buffer, int buffer_linesize, int plane,
 343                          int jobnr, int y, int x)
 344 {
 345     PlaneContext *p = &s->planes[plane];
 346     const int depth = s->depth;
 347     const int bpp = (depth + 7) / 8;
 348     const int width = p->planewidth;
 349     const int height = p->planeheight;
 350     const int block = p->b;
 351     const int overlap = p->o;
 352     const int hoverlap = overlap / 2;
 353     const int size = block - overlap;
 354     const int data_linesize = p->data_linesize / sizeof(AVComplexFloat);
 355     AVComplexFloat *hdata = p->hdata[jobnr];
 356     AVComplexFloat *hdata_out = p->hdata_out[jobnr];
 357     AVComplexFloat *vdata_out = p->vdata_out[jobnr];
 358     const int rw = FFMIN(size, width  - x * size);
 359     const int rh = FFMIN(size, height - y * size);
 360     AVComplexFloat *hdst, *vdst = vdata_out, *hdst_out = hdata_out;
 361     float *bsrc = buffer;
 362
 363     hdst = hdata;
 364     buffer_linesize /= sizeof(float);
 365
 366     for (int i = 0; i < block; i++) {
 367         s->itx_fn(s->ifft[jobnr], vdst, bsrc, sizeof(AVComplexFloat));
 368         for (int j = 0; j < block; j++)
 369             hdst[j * data_linesize + i] = vdst[j];
 370
 371         vdst += data_linesize;
 372         bsrc += buffer_linesize;
 373     }
 374
 375     hdst = hdata + hoverlap * data_linesize;
 376     for (int i = 0; i < rh && (y * size + i) < height; i++) {
 377         uint8_t *dst = dstp + dst_linesize * (y * size + i) + x * size * bpp;
 378
 379         s->itx_fn(s->ifft[jobnr], hdst_out, hdst, sizeof(AVComplexFloat));
 380         s->export_row(hdst_out + hoverlap, dst, rw, depth, s->win[i + hoverlap] + hoverlap);
 381
 382         hdst += data_linesize;
 383         hdst_out += data_linesize;
 384     }
 385 }
 386
 387 static void filter_block3d2(FFTdnoizContext *s, int plane, float *pbuffer, float *nbuffer,
 388                             int jobnr)
 389 {
 390     PlaneContext *p = &s->planes[plane];
 391     const int block = p->b;
 392     const int buffer_linesize = p->buffer_linesize / sizeof(float);
 393     const float depthx = (1 << (s->depth - 8)) * (1 << (s->depth - 8));
 394     const float sigma = s->sigma * depthx / (3.f * s->block_size * s->block_size);
 395     const float limit = 1.f - s->amount;
 396     float *cbuffer = p->buffer[jobnr][CURRENT];
 397     const int method = s->method;
 398     float *cbuff = cbuffer;
 399     float *pbuff = pbuffer;
 400     float *nbuff = nbuffer;
 401
 402     for (int i = 0; i < block; i++) {
 403         for (int j = 0; j < block; j++) {
 404             AVComplexFloat buffer[BSIZE];
 405             AVComplexFloat outbuffer[BSIZE];
 406
 407             buffer[0].re = pbuff[2 * j    ];
 408             buffer[0].im = pbuff[2 * j + 1];
 409
 410             buffer[1].re = cbuff[2 * j    ];
 411             buffer[1].im = cbuff[2 * j + 1];
 412
 413             buffer[2].re = nbuff[2 * j    ];
 414             buffer[2].im = nbuff[2 * j + 1];
 415
 416             s->tx_r_fn(s->fft_r[jobnr], outbuffer, buffer, sizeof(AVComplexFloat));
 417
 418             for (int z = 0; z < 3; z++) {
 419                 const float re = outbuffer[z].re;
 420                 const float im = outbuffer[z].im;
 421                 const float power = re * re + im * im;
 422                 float factor;
 423
 424                 switch (method) {
 425                 case 0:
 426                     factor = fmaxf(limit, (power - sigma) / (power + 1e-15f));
 427                     break;
 428                 case 1:
 429                     factor = power < sigma ? limit : 1.f;
 430                     break;
 431                 }
 432
 433                 outbuffer[z].re *= factor;
 434                 outbuffer[z].im *= factor;
 435             }
 436
 437             s->itx_r_fn(s->ifft_r[jobnr], buffer, outbuffer, sizeof(AVComplexFloat));
 438
 439             cbuff[2 * j + 0] = buffer[1].re;
 440             cbuff[2 * j + 1] = buffer[1].im;
 441         }
 442
 443         cbuff += buffer_linesize;
 444         pbuff += buffer_linesize;
 445         nbuff += buffer_linesize;
 446     }
 447 }
 448
 449 static void filter_block3d1(FFTdnoizContext *s, int plane, float *pbuffer,
 450                             int jobnr)
 451 {
 452     PlaneContext *p = &s->planes[plane];
 453     const int block = p->b;
 454     const int buffer_linesize = p->buffer_linesize / sizeof(float);
 455     const float depthx = (1 << (s->depth - 8)) * (1 << (s->depth - 8));
 456     const float sigma = s->sigma * depthx / (2.f * s->block_size * s->block_size);
 457     const float limit = 1.f - s->amount;
 458     float *cbuffer = p->buffer[jobnr][CURRENT];
 459     const int method = s->method;
 460     float *cbuff = cbuffer;
 461     float *pbuff = pbuffer;
 462
 463     for (int i = 0; i < block; i++) {
 464         for (int j = 0; j < block; j++) {
 465             AVComplexFloat buffer[BSIZE];
 466             AVComplexFloat outbuffer[BSIZE];
 467
 468             buffer[0].re = pbuff[2 * j    ];
 469             buffer[0].im = pbuff[2 * j + 1];
 470
 471             buffer[1].re = cbuff[2 * j    ];
 472             buffer[1].im = cbuff[2 * j + 1];
 473
 474             s->tx_r_fn(s->fft_r[jobnr], outbuffer, buffer, sizeof(AVComplexFloat));
 475
 476             for (int z = 0; z < 2; z++) {
 477                 const float re = outbuffer[z].re;
 478                 const float im = outbuffer[z].im;
 479                 const float power = re * re + im * im;
 480                 float factor;
 481
 482                 switch (method) {
 483                 case 0:
 484                     factor = fmaxf(limit, (power - sigma) / (power + 1e-15f));
 485                     break;
 486                 case 1:
 487                     factor = power < sigma ? limit : 1.f;
 488                     break;
 489                 }
 490
 491                 outbuffer[z].re *= factor;
 492                 outbuffer[z].im *= factor;
 493             }
 494
 495             s->itx_r_fn(s->ifft_r[jobnr], buffer, outbuffer, sizeof(AVComplexFloat));
 496
 497             cbuff[2 * j + 0] = buffer[1].re;
 498             cbuff[2 * j + 1] = buffer[1].im;
 499         }
 500
 501         cbuff += buffer_linesize;
 502         pbuff += buffer_linesize;
 503     }
 504 }
 505
 506 static void filter_block2d(FFTdnoizContext *s, int plane,
 507                            int jobnr)
 508 {
 509     PlaneContext *p = &s->planes[plane];
 510     const int block = p->b;
 511     const int method = s->method;
 512     const int buffer_linesize = p->buffer_linesize / sizeof(float);
 513     const float depthx = (1 << (s->depth - 8)) * (1 << (s->depth - 8));
 514     const float sigma = s->sigma * depthx / (s->block_size * s->block_size);
 515     const float limit = 1.f - s->amount;
 516     float *buff = p->buffer[jobnr][CURRENT];
 517
 518     for (int i = 0; i < block; i++) {
 519         for (int j = 0; j < block; j++) {
 520             float factor, power, re, im;
 521
 522             re = buff[j * 2    ];
 523             im = buff[j * 2 + 1];
 524             power = re * re + im * im;
 525             switch (method) {
 526             case 0:
 527                 factor = fmaxf(limit, (power - sigma) / (power + 1e-15f));
 528                 break;
 529             case 1:
 530                 factor = power < sigma ? limit : 1.f;
 531                 break;
 532             }
 533
 534             buff[j * 2    ] *= factor;
 535             buff[j * 2 + 1] *= factor;
 536         }
 537
 538         buff += buffer_linesize;
 539     }
 540 }
 541
 542 static int denoise(AVFilterContext *ctx, void *arg,
 543                    int jobnr, int nb_jobs)
 544 {
 545     FFTdnoizContext *s = ctx->priv;
 546     AVFrame *out = arg;
 547
 548     for (int plane = 0; plane < s->nb_planes; plane++) {
 549         PlaneContext *p = &s->planes[plane];
 550         const int nox = p->nox;
 551         const int noy = p->noy;
 552         const int slice_start = (noy * jobnr) / nb_jobs;
 553         const int slice_end = (noy * (jobnr+1)) / nb_jobs;
 554
 555         if (!((1 << plane) & s->planesf) || ctx->is_disabled)
 556             continue;
 557
 558         for (int y = slice_start; y < slice_end; y++) {
 559             for (int x = 0; x < nox; x++) {
 560                 if (s->next) {
 561                     import_block(s, s->next->data[plane], s->next->linesize[plane],
 562                                  p->buffer[jobnr][NEXT], p->buffer_linesize, plane,
 563                                  jobnr, y, x);
 564                 }
 565
 566                 if (s->prev) {
 567                     import_block(s, s->prev->data[plane], s->prev->linesize[plane],
 568                                  p->buffer[jobnr][PREV], p->buffer_linesize, plane,
 569                                  jobnr, y, x);
 570                 }
 571
 572                 import_block(s, s->cur->data[plane], s->cur->linesize[plane],
 573                              p->buffer[jobnr][CURRENT], p->buffer_linesize, plane,
 574                              jobnr, y, x);
 575
 576                 if (s->next && s->prev) {
 577                     filter_block3d2(s, plane, p->buffer[jobnr][PREV], p->buffer[jobnr][NEXT], jobnr);
 578                 } else if (s->next) {
 579                     filter_block3d1(s, plane, p->buffer[jobnr][NEXT], jobnr);
 580                 } else  if (s->prev) {
 581                     filter_block3d1(s, plane, p->buffer[jobnr][PREV], jobnr);
 582                 } else {
 583                     filter_block2d(s, plane, jobnr);
 584                 }
 585
 586                 export_block(s, out->data[plane], out->linesize[plane],
 587                              p->buffer[jobnr][CURRENT], p->buffer_linesize, plane,
 588                              jobnr, y, x);
 589             }
 590         }
 591     }
 592
 593     return 0;
 594 }
 595
 596 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
 597 {
 598     AVFilterContext *ctx = inlink->dst;
 599     FFTdnoizContext *s = ctx->priv;
 600     AVFilterLink *outlink = ctx->outputs[0];
 601     int direct, plane;
 602     AVFrame *out;
 603
 604     if (s->nb_next > 0 && s->nb_prev > 0) {
 605         av_frame_free(&s->prev);
 606         s->prev = s->cur;
 607         s->cur = s->next;
 608         s->next = in;
 609
 610         if (!s->prev && s->cur) {
 611             s->prev = av_frame_clone(s->cur);
 612             if (!s->prev)
 613                 return AVERROR(ENOMEM);
 614         }
 615         if (!s->cur)
 616             return 0;
 617     } else if (s->nb_next > 0) {
 618         av_frame_free(&s->cur);
 619         s->cur = s->next;
 620         s->next = in;
 621
 622         if (!s->cur)
 623             return 0;
 624     } else if (s->nb_prev > 0) {
 625         av_frame_free(&s->prev);
 626         s->prev = s->cur;
 627         s->cur = in;
 628
 629         if (!s->prev)
 630             s->prev = av_frame_clone(s->cur);
 631         if (!s->prev)
 632             return AVERROR(ENOMEM);
 633     } else {
 634         s->cur = in;
 635     }
 636
 637     if (av_frame_is_writable(in) && s->nb_next == 0 && s->nb_prev == 0) {
 638         direct = 1;
 639         out = in;
 640     } else {
 641         direct = 0;
 642         out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
 643         if (!out)
 644             return AVERROR(ENOMEM);
 645         av_frame_copy_props(out, s->cur);
 646     }
 647
 648     ff_filter_execute(ctx, denoise, out, NULL,
 649                       FFMIN(s->planes[0].noy, s->nb_threads));
 650
 651     for (plane = 0; plane < s->nb_planes; plane++) {
 652         PlaneContext *p = &s->planes[plane];
 653
 654         if (!((1 << plane) & s->planesf) || ctx->is_disabled) {
 655             if (!direct)
 656                 av_image_copy_plane(out->data[plane], out->linesize[plane],
 657                                     s->cur->data[plane], s->cur->linesize[plane],
 658                                     p->planewidth * (1 + (s->depth > 8)), p->planeheight);
 659             continue;
 660         }
 661     }
 662
 663     if (s->nb_next == 0 && s->nb_prev == 0) {
 664         if (direct) {
 665             s->cur = NULL;
 666         } else {
 667             av_frame_free(&s->cur);
 668         }
 669     }
 670     return ff_filter_frame(outlink, out);
 671 }
 672
 673 static int request_frame(AVFilterLink *outlink)
 674 {
 675     AVFilterContext *ctx = outlink->src;
 676     FFTdnoizContext *s = ctx->priv;
 677     int ret = 0;
 678
 679     ret = ff_request_frame(ctx->inputs[0]);
 680
 681     if (ret == AVERROR_EOF && (s->nb_next > 0)) {
 682         AVFrame *buf;
 683
 684         if (s->next && s->nb_next > 0)
 685             buf = av_frame_clone(s->next);
 686         else if (s->cur)
 687             buf = av_frame_clone(s->cur);
 688         else
 689             buf = av_frame_clone(s->prev);
 690         if (!buf)
 691             return AVERROR(ENOMEM);
 692
 693         ret = filter_frame(ctx->inputs[0], buf);
 694         if (ret < 0)
 695             return ret;
 696         ret = AVERROR_EOF;
 697     }
 698
 699     return ret;
 700 }
 701
 702 static av_cold void uninit(AVFilterContext *ctx)
 703 {
 704     FFTdnoizContext *s = ctx->priv;
 705     int i;
 706
 707     for (i = 0; i < 4; i++) {
 708         PlaneContext *p = &s->planes[i];
 709
 710         for (int j = 0; j < s->nb_threads; j++) {
 711             av_freep(&p->hdata[j]);
 712             av_freep(&p->vdata[j]);
 713             av_freep(&p->hdata_out[j]);
 714             av_freep(&p->vdata_out[j]);
 715             av_freep(&p->buffer[j][PREV]);
 716             av_freep(&p->buffer[j][CURRENT]);
 717             av_freep(&p->buffer[j][NEXT]);
 718         }
 719     }
 720
 721     for (i = 0; i < s->nb_threads; i++) {
 722         av_tx_uninit(&s->fft[i]);
 723         av_tx_uninit(&s->ifft[i]);
 724         av_tx_uninit(&s->fft_r[i]);
 725         av_tx_uninit(&s->ifft_r[i]);
 726     }
 727
 728     av_frame_free(&s->prev);
 729     av_frame_free(&s->cur);
 730     av_frame_free(&s->next);
 731 }
 732
 733 static const AVFilterPad fftdnoiz_inputs[] = {
 734     {
 735         .name         = "default",
 736         .type         = AVMEDIA_TYPE_VIDEO,
 737         .filter_frame = filter_frame,
 738         .config_props = config_input,
 739     },
 740 };
 741
 742 static const AVFilterPad fftdnoiz_outputs[] = {
 743     {
 744         .name          = "default",
 745         .type          = AVMEDIA_TYPE_VIDEO,
 746         .request_frame = request_frame,
 747     },
 748 };
 749
 750 const FFFilter ff_vf_fftdnoiz = {
 751     .p.name        = "fftdnoiz",
 752     .p.description = NULL_IF_CONFIG_SMALL("Denoise frames using 3D FFT."),
 753     .p.priv_class  = &fftdnoiz_class,
 754     .p.flags       = AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL |
 755                      AVFILTER_FLAG_SLICE_THREADS,
 756     .priv_size     = sizeof(FFTdnoizContext),
 757     .uninit        = uninit,
 758     FILTER_INPUTS(fftdnoiz_inputs),
 759     FILTER_OUTPUTS(fftdnoiz_outputs),
 760     FILTER_PIXFMTS_ARRAY(pix_fmts),
 761     .process_command = ff_filter_process_command,
 762 };