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