]> git.sesse.net Git - ffmpeg/blob - libavfilter/vf_dctdnoiz.c
a89f2631c829aa8afa8ea90f5ededf60e810f950
[ffmpeg] / libavfilter / vf_dctdnoiz.c
1 /*
2  * Copyright (c) 2013-2014 Clément Bœsch
3  *
4  * This file is part of FFmpeg.
5  *
6  * FFmpeg is free software; you can redistribute it and/or
7  * modify it under the terms of the GNU Lesser General Public
8  * License as published by the Free Software Foundation; either
9  * version 2.1 of the License, or (at your option) any later version.
10  *
11  * FFmpeg is distributed in the hope that it will be useful,
12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14  * Lesser General Public License for more details.
15  *
16  * You should have received a copy of the GNU Lesser General Public
17  * License along with FFmpeg; if not, write to the Free Software
18  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
19  */
20
21 /**
22  * A simple, relatively efficient and slow DCT image denoiser.
23  *
24  * @see http://www.ipol.im/pub/art/2011/ys-dct/
25  *
26  * The DCT factorization used is based on "Fast and numerically stable
27  * algorithms for discrete cosine transforms" from Gerlind Plonkaa & Manfred
28  * Tasche (DOI: 10.1016/j.laa.2004.07.015).
29  */
30
31 #include "libavutil/avassert.h"
32 #include "libavutil/eval.h"
33 #include "libavutil/mem_internal.h"
34 #include "libavutil/opt.h"
35 #include "internal.h"
36
37 static const char *const var_names[] = { "c", NULL };
38 enum { VAR_C, VAR_VARS_NB };
39
40 #define MAX_THREADS 8
41
42 typedef struct DCTdnoizContext {
43     const AVClass *class;
44
45     /* coefficient factor expression */
46     char *expr_str;
47     AVExpr *expr[MAX_THREADS];
48     double var_values[MAX_THREADS][VAR_VARS_NB];
49
50     int nb_threads;
51     int pr_width, pr_height;    // width and height to process
52     float sigma;                // used when no expression are st
53     float th;                   // threshold (3*sigma)
54     float *cbuf[2][3];          // two planar rgb color buffers
55     float *slices[MAX_THREADS]; // slices buffers (1 slice buffer per thread)
56     float *weights;             // dct coeff are cumulated with overlapping; these values are used for averaging
57     int p_linesize;             // line sizes for color and weights
58     int overlap;                // number of block overlapping pixels
59     int step;                   // block step increment (blocksize - overlap)
60     int n;                      // 1<<n is the block size
61     int bsize;                  // block size, 1<<n
62     void (*filter_freq_func)(struct DCTdnoizContext *s,
63                              const float *src, int src_linesize,
64                              float *dst, int dst_linesize,
65                              int thread_id);
66     void (*color_decorrelation)(float **dst, int dst_linesize,
67                                 const uint8_t **src, int src_linesize,
68                                 int w, int h);
69     void (*color_correlation)(uint8_t **dst, int dst_linesize,
70                               float **src, int src_linesize,
71                               int w, int h);
72 } DCTdnoizContext;
73
74 #define MIN_NBITS 3 /* blocksize = 1<<3 =  8 */
75 #define MAX_NBITS 4 /* blocksize = 1<<4 = 16 */
76 #define DEFAULT_NBITS 3
77
78 #define OFFSET(x) offsetof(DCTdnoizContext, x)
79 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
80 static const AVOption dctdnoiz_options[] = {
81     { "sigma",   "set noise sigma constant",               OFFSET(sigma),    AV_OPT_TYPE_FLOAT,  {.dbl=0},            0, 999,          .flags = FLAGS },
82     { "s",       "set noise sigma constant",               OFFSET(sigma),    AV_OPT_TYPE_FLOAT,  {.dbl=0},            0, 999,          .flags = FLAGS },
83     { "overlap", "set number of block overlapping pixels", OFFSET(overlap),  AV_OPT_TYPE_INT,    {.i64=-1}, -1, (1<<MAX_NBITS)-1, .flags = FLAGS },
84     { "expr",    "set coefficient factor expression",      OFFSET(expr_str), AV_OPT_TYPE_STRING, {.str=NULL},                          .flags = FLAGS },
85     { "e",       "set coefficient factor expression",      OFFSET(expr_str), AV_OPT_TYPE_STRING, {.str=NULL},                          .flags = FLAGS },
86     { "n",       "set the block size, expressed in bits",  OFFSET(n),        AV_OPT_TYPE_INT,    {.i64=DEFAULT_NBITS}, MIN_NBITS, MAX_NBITS, .flags = FLAGS },
87     { NULL }
88 };
89
90 AVFILTER_DEFINE_CLASS(dctdnoiz);
91
92 static void av_always_inline fdct8_1d(float *dst, const float *src,
93                                       int dst_stridea, int dst_strideb,
94                                       int src_stridea, int src_strideb)
95 {
96     int i;
97
98     for (i = 0; i < 8; i++) {
99         const float x00 = src[0*src_stridea] + src[7*src_stridea];
100         const float x01 = src[1*src_stridea] + src[6*src_stridea];
101         const float x02 = src[2*src_stridea] + src[5*src_stridea];
102         const float x03 = src[3*src_stridea] + src[4*src_stridea];
103         const float x04 = src[0*src_stridea] - src[7*src_stridea];
104         const float x05 = src[1*src_stridea] - src[6*src_stridea];
105         const float x06 = src[2*src_stridea] - src[5*src_stridea];
106         const float x07 = src[3*src_stridea] - src[4*src_stridea];
107         const float x08 = x00 + x03;
108         const float x09 = x01 + x02;
109         const float x0a = x00 - x03;
110         const float x0b = x01 - x02;
111         const float x0c = 1.38703984532215f*x04 + 0.275899379282943f*x07;
112         const float x0d = 1.17587560241936f*x05 + 0.785694958387102f*x06;
113         const float x0e = -0.785694958387102f*x05 + 1.17587560241936f*x06;
114         const float x0f = 0.275899379282943f*x04 - 1.38703984532215f*x07;
115         const float x10 = 0.353553390593274f * (x0c - x0d);
116         const float x11 = 0.353553390593274f * (x0e - x0f);
117         dst[0*dst_stridea] = 0.353553390593274f * (x08 + x09);
118         dst[1*dst_stridea] = 0.353553390593274f * (x0c + x0d);
119         dst[2*dst_stridea] = 0.461939766255643f*x0a + 0.191341716182545f*x0b;
120         dst[3*dst_stridea] = 0.707106781186547f * (x10 - x11);
121         dst[4*dst_stridea] = 0.353553390593274f * (x08 - x09);
122         dst[5*dst_stridea] = 0.707106781186547f * (x10 + x11);
123         dst[6*dst_stridea] = 0.191341716182545f*x0a - 0.461939766255643f*x0b;
124         dst[7*dst_stridea] = 0.353553390593274f * (x0e + x0f);
125         dst += dst_strideb;
126         src += src_strideb;
127     }
128 }
129
130 static void av_always_inline idct8_1d(float *dst, const float *src,
131                                       int dst_stridea, int dst_strideb,
132                                       int src_stridea, int src_strideb,
133                                       int add)
134 {
135     int i;
136
137     for (i = 0; i < 8; i++) {
138         const float x00 =  1.4142135623731f  *src[0*src_stridea];
139         const float x01 =  1.38703984532215f *src[1*src_stridea] + 0.275899379282943f*src[7*src_stridea];
140         const float x02 =  1.30656296487638f *src[2*src_stridea] + 0.541196100146197f*src[6*src_stridea];
141         const float x03 =  1.17587560241936f *src[3*src_stridea] + 0.785694958387102f*src[5*src_stridea];
142         const float x04 =  1.4142135623731f  *src[4*src_stridea];
143         const float x05 = -0.785694958387102f*src[3*src_stridea] + 1.17587560241936f*src[5*src_stridea];
144         const float x06 =  0.541196100146197f*src[2*src_stridea] - 1.30656296487638f*src[6*src_stridea];
145         const float x07 = -0.275899379282943f*src[1*src_stridea] + 1.38703984532215f*src[7*src_stridea];
146         const float x09 = x00 + x04;
147         const float x0a = x01 + x03;
148         const float x0b = 1.4142135623731f*x02;
149         const float x0c = x00 - x04;
150         const float x0d = x01 - x03;
151         const float x0e = 0.353553390593274f * (x09 - x0b);
152         const float x0f = 0.353553390593274f * (x0c + x0d);
153         const float x10 = 0.353553390593274f * (x0c - x0d);
154         const float x11 = 1.4142135623731f*x06;
155         const float x12 = x05 + x07;
156         const float x13 = x05 - x07;
157         const float x14 = 0.353553390593274f * (x11 + x12);
158         const float x15 = 0.353553390593274f * (x11 - x12);
159         const float x16 = 0.5f*x13;
160         dst[0*dst_stridea] = (add ? dst[ 0*dst_stridea] : 0) + 0.25f * (x09 + x0b) + 0.353553390593274f*x0a;
161         dst[1*dst_stridea] = (add ? dst[ 1*dst_stridea] : 0) + 0.707106781186547f * (x0f + x15);
162         dst[2*dst_stridea] = (add ? dst[ 2*dst_stridea] : 0) + 0.707106781186547f * (x0f - x15);
163         dst[3*dst_stridea] = (add ? dst[ 3*dst_stridea] : 0) + 0.707106781186547f * (x0e + x16);
164         dst[4*dst_stridea] = (add ? dst[ 4*dst_stridea] : 0) + 0.707106781186547f * (x0e - x16);
165         dst[5*dst_stridea] = (add ? dst[ 5*dst_stridea] : 0) + 0.707106781186547f * (x10 - x14);
166         dst[6*dst_stridea] = (add ? dst[ 6*dst_stridea] : 0) + 0.707106781186547f * (x10 + x14);
167         dst[7*dst_stridea] = (add ? dst[ 7*dst_stridea] : 0) + 0.25f * (x09 + x0b) - 0.353553390593274f*x0a;
168         dst += dst_strideb;
169         src += src_strideb;
170     }
171 }
172
173
174 static void av_always_inline fdct16_1d(float *dst, const float *src,
175                                        int dst_stridea, int dst_strideb,
176                                        int src_stridea, int src_strideb)
177 {
178     int i;
179
180     for (i = 0; i < 16; i++) {
181         const float x00 = src[ 0*src_stridea] + src[15*src_stridea];
182         const float x01 = src[ 1*src_stridea] + src[14*src_stridea];
183         const float x02 = src[ 2*src_stridea] + src[13*src_stridea];
184         const float x03 = src[ 3*src_stridea] + src[12*src_stridea];
185         const float x04 = src[ 4*src_stridea] + src[11*src_stridea];
186         const float x05 = src[ 5*src_stridea] + src[10*src_stridea];
187         const float x06 = src[ 6*src_stridea] + src[ 9*src_stridea];
188         const float x07 = src[ 7*src_stridea] + src[ 8*src_stridea];
189         const float x08 = src[ 0*src_stridea] - src[15*src_stridea];
190         const float x09 = src[ 1*src_stridea] - src[14*src_stridea];
191         const float x0a = src[ 2*src_stridea] - src[13*src_stridea];
192         const float x0b = src[ 3*src_stridea] - src[12*src_stridea];
193         const float x0c = src[ 4*src_stridea] - src[11*src_stridea];
194         const float x0d = src[ 5*src_stridea] - src[10*src_stridea];
195         const float x0e = src[ 6*src_stridea] - src[ 9*src_stridea];
196         const float x0f = src[ 7*src_stridea] - src[ 8*src_stridea];
197         const float x10 = x00 + x07;
198         const float x11 = x01 + x06;
199         const float x12 = x02 + x05;
200         const float x13 = x03 + x04;
201         const float x14 = x00 - x07;
202         const float x15 = x01 - x06;
203         const float x16 = x02 - x05;
204         const float x17 = x03 - x04;
205         const float x18 = x10 + x13;
206         const float x19 = x11 + x12;
207         const float x1a = x10 - x13;
208         const float x1b = x11 - x12;
209         const float x1c =   1.38703984532215f*x14 + 0.275899379282943f*x17;
210         const float x1d =   1.17587560241936f*x15 + 0.785694958387102f*x16;
211         const float x1e = -0.785694958387102f*x15 + 1.17587560241936f *x16;
212         const float x1f =  0.275899379282943f*x14 - 1.38703984532215f *x17;
213         const float x20 = 0.25f * (x1c - x1d);
214         const float x21 = 0.25f * (x1e - x1f);
215         const float x22 =  1.40740373752638f *x08 + 0.138617169199091f*x0f;
216         const float x23 =  1.35331800117435f *x09 + 0.410524527522357f*x0e;
217         const float x24 =  1.24722501298667f *x0a + 0.666655658477747f*x0d;
218         const float x25 =  1.09320186700176f *x0b + 0.897167586342636f*x0c;
219         const float x26 = -0.897167586342636f*x0b + 1.09320186700176f *x0c;
220         const float x27 =  0.666655658477747f*x0a - 1.24722501298667f *x0d;
221         const float x28 = -0.410524527522357f*x09 + 1.35331800117435f *x0e;
222         const float x29 =  0.138617169199091f*x08 - 1.40740373752638f *x0f;
223         const float x2a = x22 + x25;
224         const float x2b = x23 + x24;
225         const float x2c = x22 - x25;
226         const float x2d = x23 - x24;
227         const float x2e = 0.25f * (x2a - x2b);
228         const float x2f = 0.326640741219094f*x2c + 0.135299025036549f*x2d;
229         const float x30 = 0.135299025036549f*x2c - 0.326640741219094f*x2d;
230         const float x31 = x26 + x29;
231         const float x32 = x27 + x28;
232         const float x33 = x26 - x29;
233         const float x34 = x27 - x28;
234         const float x35 = 0.25f * (x31 - x32);
235         const float x36 = 0.326640741219094f*x33 + 0.135299025036549f*x34;
236         const float x37 = 0.135299025036549f*x33 - 0.326640741219094f*x34;
237         dst[ 0*dst_stridea] = 0.25f * (x18 + x19);
238         dst[ 1*dst_stridea] = 0.25f * (x2a + x2b);
239         dst[ 2*dst_stridea] = 0.25f * (x1c + x1d);
240         dst[ 3*dst_stridea] = 0.707106781186547f * (x2f - x37);
241         dst[ 4*dst_stridea] = 0.326640741219094f*x1a + 0.135299025036549f*x1b;
242         dst[ 5*dst_stridea] = 0.707106781186547f * (x2f + x37);
243         dst[ 6*dst_stridea] = 0.707106781186547f * (x20 - x21);
244         dst[ 7*dst_stridea] = 0.707106781186547f * (x2e + x35);
245         dst[ 8*dst_stridea] = 0.25f * (x18 - x19);
246         dst[ 9*dst_stridea] = 0.707106781186547f * (x2e - x35);
247         dst[10*dst_stridea] = 0.707106781186547f * (x20 + x21);
248         dst[11*dst_stridea] = 0.707106781186547f * (x30 - x36);
249         dst[12*dst_stridea] = 0.135299025036549f*x1a - 0.326640741219094f*x1b;
250         dst[13*dst_stridea] = 0.707106781186547f * (x30 + x36);
251         dst[14*dst_stridea] = 0.25f * (x1e + x1f);
252         dst[15*dst_stridea] = 0.25f * (x31 + x32);
253         dst += dst_strideb;
254         src += src_strideb;
255     }
256 }
257
258 static void av_always_inline idct16_1d(float *dst, const float *src,
259                                        int dst_stridea, int dst_strideb,
260                                        int src_stridea, int src_strideb,
261                                        int add)
262 {
263     int i;
264
265     for (i = 0; i < 16; i++) {
266         const float x00 =  1.4142135623731f  *src[ 0*src_stridea];
267         const float x01 =  1.40740373752638f *src[ 1*src_stridea] + 0.138617169199091f*src[15*src_stridea];
268         const float x02 =  1.38703984532215f *src[ 2*src_stridea] + 0.275899379282943f*src[14*src_stridea];
269         const float x03 =  1.35331800117435f *src[ 3*src_stridea] + 0.410524527522357f*src[13*src_stridea];
270         const float x04 =  1.30656296487638f *src[ 4*src_stridea] + 0.541196100146197f*src[12*src_stridea];
271         const float x05 =  1.24722501298667f *src[ 5*src_stridea] + 0.666655658477747f*src[11*src_stridea];
272         const float x06 =  1.17587560241936f *src[ 6*src_stridea] + 0.785694958387102f*src[10*src_stridea];
273         const float x07 =  1.09320186700176f *src[ 7*src_stridea] + 0.897167586342636f*src[ 9*src_stridea];
274         const float x08 =  1.4142135623731f  *src[ 8*src_stridea];
275         const float x09 = -0.897167586342636f*src[ 7*src_stridea] + 1.09320186700176f*src[ 9*src_stridea];
276         const float x0a =  0.785694958387102f*src[ 6*src_stridea] - 1.17587560241936f*src[10*src_stridea];
277         const float x0b = -0.666655658477747f*src[ 5*src_stridea] + 1.24722501298667f*src[11*src_stridea];
278         const float x0c =  0.541196100146197f*src[ 4*src_stridea] - 1.30656296487638f*src[12*src_stridea];
279         const float x0d = -0.410524527522357f*src[ 3*src_stridea] + 1.35331800117435f*src[13*src_stridea];
280         const float x0e =  0.275899379282943f*src[ 2*src_stridea] - 1.38703984532215f*src[14*src_stridea];
281         const float x0f = -0.138617169199091f*src[ 1*src_stridea] + 1.40740373752638f*src[15*src_stridea];
282         const float x12 = x00 + x08;
283         const float x13 = x01 + x07;
284         const float x14 = x02 + x06;
285         const float x15 = x03 + x05;
286         const float x16 = 1.4142135623731f*x04;
287         const float x17 = x00 - x08;
288         const float x18 = x01 - x07;
289         const float x19 = x02 - x06;
290         const float x1a = x03 - x05;
291         const float x1d = x12 + x16;
292         const float x1e = x13 + x15;
293         const float x1f = 1.4142135623731f*x14;
294         const float x20 = x12 - x16;
295         const float x21 = x13 - x15;
296         const float x22 = 0.25f * (x1d - x1f);
297         const float x23 = 0.25f * (x20 + x21);
298         const float x24 = 0.25f * (x20 - x21);
299         const float x25 = 1.4142135623731f*x17;
300         const float x26 = 1.30656296487638f*x18 + 0.541196100146197f*x1a;
301         const float x27 = 1.4142135623731f*x19;
302         const float x28 = -0.541196100146197f*x18 + 1.30656296487638f*x1a;
303         const float x29 = 0.176776695296637f * (x25 + x27) + 0.25f*x26;
304         const float x2a = 0.25f * (x25 - x27);
305         const float x2b = 0.176776695296637f * (x25 + x27) - 0.25f*x26;
306         const float x2c = 0.353553390593274f*x28;
307         const float x1b = 0.707106781186547f * (x2a - x2c);
308         const float x1c = 0.707106781186547f * (x2a + x2c);
309         const float x2d = 1.4142135623731f*x0c;
310         const float x2e = x0b + x0d;
311         const float x2f = x0a + x0e;
312         const float x30 = x09 + x0f;
313         const float x31 = x09 - x0f;
314         const float x32 = x0a - x0e;
315         const float x33 = x0b - x0d;
316         const float x37 = 1.4142135623731f*x2d;
317         const float x38 = 1.30656296487638f*x2e + 0.541196100146197f*x30;
318         const float x39 = 1.4142135623731f*x2f;
319         const float x3a = -0.541196100146197f*x2e + 1.30656296487638f*x30;
320         const float x3b = 0.176776695296637f * (x37 + x39) + 0.25f*x38;
321         const float x3c = 0.25f * (x37 - x39);
322         const float x3d = 0.176776695296637f * (x37 + x39) - 0.25f*x38;
323         const float x3e = 0.353553390593274f*x3a;
324         const float x34 = 0.707106781186547f * (x3c - x3e);
325         const float x35 = 0.707106781186547f * (x3c + x3e);
326         const float x3f = 1.4142135623731f*x32;
327         const float x40 = x31 + x33;
328         const float x41 = x31 - x33;
329         const float x42 = 0.25f * (x3f + x40);
330         const float x43 = 0.25f * (x3f - x40);
331         const float x44 = 0.353553390593274f*x41;
332         dst[ 0*dst_stridea] = (add ? dst[ 0*dst_stridea] : 0) + 0.176776695296637f * (x1d + x1f) + 0.25f*x1e;
333         dst[ 1*dst_stridea] = (add ? dst[ 1*dst_stridea] : 0) + 0.707106781186547f * (x29 + x3d);
334         dst[ 2*dst_stridea] = (add ? dst[ 2*dst_stridea] : 0) + 0.707106781186547f * (x29 - x3d);
335         dst[ 3*dst_stridea] = (add ? dst[ 3*dst_stridea] : 0) + 0.707106781186547f * (x23 - x43);
336         dst[ 4*dst_stridea] = (add ? dst[ 4*dst_stridea] : 0) + 0.707106781186547f * (x23 + x43);
337         dst[ 5*dst_stridea] = (add ? dst[ 5*dst_stridea] : 0) + 0.707106781186547f * (x1b - x35);
338         dst[ 6*dst_stridea] = (add ? dst[ 6*dst_stridea] : 0) + 0.707106781186547f * (x1b + x35);
339         dst[ 7*dst_stridea] = (add ? dst[ 7*dst_stridea] : 0) + 0.707106781186547f * (x22 + x44);
340         dst[ 8*dst_stridea] = (add ? dst[ 8*dst_stridea] : 0) + 0.707106781186547f * (x22 - x44);
341         dst[ 9*dst_stridea] = (add ? dst[ 9*dst_stridea] : 0) + 0.707106781186547f * (x1c + x34);
342         dst[10*dst_stridea] = (add ? dst[10*dst_stridea] : 0) + 0.707106781186547f * (x1c - x34);
343         dst[11*dst_stridea] = (add ? dst[11*dst_stridea] : 0) + 0.707106781186547f * (x24 + x42);
344         dst[12*dst_stridea] = (add ? dst[12*dst_stridea] : 0) + 0.707106781186547f * (x24 - x42);
345         dst[13*dst_stridea] = (add ? dst[13*dst_stridea] : 0) + 0.707106781186547f * (x2b - x3b);
346         dst[14*dst_stridea] = (add ? dst[14*dst_stridea] : 0) + 0.707106781186547f * (x2b + x3b);
347         dst[15*dst_stridea] = (add ? dst[15*dst_stridea] : 0) + 0.176776695296637f * (x1d + x1f) - 0.25f*x1e;
348         dst += dst_strideb;
349         src += src_strideb;
350     }
351 }
352
353 #define DEF_FILTER_FREQ_FUNCS(bsize)                                                        \
354 static av_always_inline void filter_freq_##bsize(const float *src, int src_linesize,        \
355                                                  float *dst, int dst_linesize,              \
356                                                  AVExpr *expr, double *var_values,          \
357                                                  int sigma_th)                              \
358 {                                                                                           \
359     unsigned i;                                                                             \
360     DECLARE_ALIGNED(32, float, tmp_block1)[bsize * bsize];                                  \
361     DECLARE_ALIGNED(32, float, tmp_block2)[bsize * bsize];                                  \
362                                                                                             \
363     /* forward DCT */                                                                       \
364     fdct##bsize##_1d(tmp_block1, src, 1, bsize, 1, src_linesize);                           \
365     fdct##bsize##_1d(tmp_block2, tmp_block1, bsize, 1, bsize, 1);                           \
366                                                                                             \
367     for (i = 0; i < bsize*bsize; i++) {                                                     \
368         float *b = &tmp_block2[i];                                                          \
369         /* frequency filtering */                                                           \
370         if (expr) {                                                                         \
371             var_values[VAR_C] = fabsf(*b);                                                  \
372             *b *= av_expr_eval(expr, var_values, NULL);                                     \
373         } else {                                                                            \
374             if (fabsf(*b) < sigma_th)                                                       \
375                 *b = 0;                                                                     \
376         }                                                                                   \
377     }                                                                                       \
378                                                                                             \
379     /* inverse DCT */                                                                       \
380     idct##bsize##_1d(tmp_block1, tmp_block2, 1, bsize, 1, bsize, 0);                        \
381     idct##bsize##_1d(dst, tmp_block1, dst_linesize, 1, bsize, 1, 1);                        \
382 }                                                                                           \
383                                                                                             \
384 static void filter_freq_sigma_##bsize(DCTdnoizContext *s,                                   \
385                                       const float *src, int src_linesize,                   \
386                                       float *dst, int dst_linesize, int thread_id)          \
387 {                                                                                           \
388     filter_freq_##bsize(src, src_linesize, dst, dst_linesize, NULL, NULL, s->th);           \
389 }                                                                                           \
390                                                                                             \
391 static void filter_freq_expr_##bsize(DCTdnoizContext *s,                                    \
392                                      const float *src, int src_linesize,                    \
393                                      float *dst, int dst_linesize, int thread_id)           \
394 {                                                                                           \
395     filter_freq_##bsize(src, src_linesize, dst, dst_linesize,                               \
396                         s->expr[thread_id], s->var_values[thread_id], 0);                   \
397 }
398
399 DEF_FILTER_FREQ_FUNCS(8)
400 DEF_FILTER_FREQ_FUNCS(16)
401
402 #define DCT3X3_0_0  0.5773502691896258f /*  1/sqrt(3) */
403 #define DCT3X3_0_1  0.5773502691896258f /*  1/sqrt(3) */
404 #define DCT3X3_0_2  0.5773502691896258f /*  1/sqrt(3) */
405 #define DCT3X3_1_0  0.7071067811865475f /*  1/sqrt(2) */
406 #define DCT3X3_1_2 -0.7071067811865475f /* -1/sqrt(2) */
407 #define DCT3X3_2_0  0.4082482904638631f /*  1/sqrt(6) */
408 #define DCT3X3_2_1 -0.8164965809277261f /* -2/sqrt(6) */
409 #define DCT3X3_2_2  0.4082482904638631f /*  1/sqrt(6) */
410
411 static av_always_inline void color_decorrelation(float **dst, int dst_linesize,
412                                                  const uint8_t **src, int src_linesize,
413                                                  int w, int h,
414                                                  int r, int g, int b)
415 {
416     int x, y;
417     float *dstp_r = dst[0];
418     float *dstp_g = dst[1];
419     float *dstp_b = dst[2];
420     const uint8_t *srcp = src[0];
421
422     for (y = 0; y < h; y++) {
423         for (x = 0; x < w; x++) {
424             dstp_r[x] = srcp[r] * DCT3X3_0_0 + srcp[g] * DCT3X3_0_1 + srcp[b] * DCT3X3_0_2;
425             dstp_g[x] = srcp[r] * DCT3X3_1_0 +                        srcp[b] * DCT3X3_1_2;
426             dstp_b[x] = srcp[r] * DCT3X3_2_0 + srcp[g] * DCT3X3_2_1 + srcp[b] * DCT3X3_2_2;
427             srcp += 3;
428         }
429         srcp   += src_linesize - w * 3;
430         dstp_r += dst_linesize;
431         dstp_g += dst_linesize;
432         dstp_b += dst_linesize;
433     }
434 }
435
436 static av_always_inline void color_correlation(uint8_t **dst, int dst_linesize,
437                                                float **src, int src_linesize,
438                                                int w, int h,
439                                                int r, int g, int b)
440 {
441     int x, y;
442     const float *src_r = src[0];
443     const float *src_g = src[1];
444     const float *src_b = src[2];
445     uint8_t *dstp = dst[0];
446
447     for (y = 0; y < h; y++) {
448         for (x = 0; x < w; x++) {
449             dstp[r] = av_clip_uint8(src_r[x] * DCT3X3_0_0 + src_g[x] * DCT3X3_1_0 + src_b[x] * DCT3X3_2_0);
450             dstp[g] = av_clip_uint8(src_r[x] * DCT3X3_0_1 +                         src_b[x] * DCT3X3_2_1);
451             dstp[b] = av_clip_uint8(src_r[x] * DCT3X3_0_2 + src_g[x] * DCT3X3_1_2 + src_b[x] * DCT3X3_2_2);
452             dstp += 3;
453         }
454         dstp  += dst_linesize - w * 3;
455         src_r += src_linesize;
456         src_g += src_linesize;
457         src_b += src_linesize;
458     }
459 }
460
461 #define DECLARE_COLOR_FUNCS(name, r, g, b)                                          \
462 static void color_decorrelation_##name(float **dst, int dst_linesize,               \
463                                        const uint8_t **src, int src_linesize,       \
464                                        int w, int h)                                \
465 {                                                                                   \
466     color_decorrelation(dst, dst_linesize, src, src_linesize, w, h, r, g, b);       \
467 }                                                                                   \
468                                                                                     \
469 static void color_correlation_##name(uint8_t **dst, int dst_linesize,               \
470                                      float **src, int src_linesize,                 \
471                                      int w, int h)                                  \
472 {                                                                                   \
473     color_correlation(dst, dst_linesize, src, src_linesize, w, h, r, g, b);         \
474 }
475
476 DECLARE_COLOR_FUNCS(rgb, 0, 1, 2)
477 DECLARE_COLOR_FUNCS(bgr, 2, 1, 0)
478
479 static av_always_inline void color_decorrelation_gbrp(float **dst, int dst_linesize,
480                                                       const uint8_t **src, int src_linesize,
481                                                       int w, int h)
482 {
483     int x, y;
484     float *dstp_r = dst[0];
485     float *dstp_g = dst[1];
486     float *dstp_b = dst[2];
487     const uint8_t *srcp_r = src[2];
488     const uint8_t *srcp_g = src[0];
489     const uint8_t *srcp_b = src[1];
490
491     for (y = 0; y < h; y++) {
492         for (x = 0; x < w; x++) {
493             dstp_r[x] = srcp_r[x] * DCT3X3_0_0 + srcp_g[x] * DCT3X3_0_1 + srcp_b[x] * DCT3X3_0_2;
494             dstp_g[x] = srcp_r[x] * DCT3X3_1_0 +                          srcp_b[x] * DCT3X3_1_2;
495             dstp_b[x] = srcp_r[x] * DCT3X3_2_0 + srcp_g[x] * DCT3X3_2_1 + srcp_b[x] * DCT3X3_2_2;
496         }
497         srcp_r += src_linesize;
498         srcp_g += src_linesize;
499         srcp_b += src_linesize;
500         dstp_r += dst_linesize;
501         dstp_g += dst_linesize;
502         dstp_b += dst_linesize;
503     }
504 }
505
506 static av_always_inline void color_correlation_gbrp(uint8_t **dst, int dst_linesize,
507                                                     float **src, int src_linesize,
508                                                     int w, int h)
509 {
510     int x, y;
511     const float *src_r = src[0];
512     const float *src_g = src[1];
513     const float *src_b = src[2];
514     uint8_t *dstp_r = dst[2];
515     uint8_t *dstp_g = dst[0];
516     uint8_t *dstp_b = dst[1];
517
518     for (y = 0; y < h; y++) {
519         for (x = 0; x < w; x++) {
520             dstp_r[x] = av_clip_uint8(src_r[x] * DCT3X3_0_0 + src_g[x] * DCT3X3_1_0 + src_b[x] * DCT3X3_2_0);
521             dstp_g[x] = av_clip_uint8(src_r[x] * DCT3X3_0_1 +                         src_b[x] * DCT3X3_2_1);
522             dstp_b[x] = av_clip_uint8(src_r[x] * DCT3X3_0_2 + src_g[x] * DCT3X3_1_2 + src_b[x] * DCT3X3_2_2);
523         }
524         dstp_r += dst_linesize;
525         dstp_g += dst_linesize;
526         dstp_b += dst_linesize;
527         src_r += src_linesize;
528         src_g += src_linesize;
529         src_b += src_linesize;
530     }
531 }
532
533 static int config_input(AVFilterLink *inlink)
534 {
535     AVFilterContext *ctx = inlink->dst;
536     DCTdnoizContext *s = ctx->priv;
537     int i, x, y, bx, by, linesize, *iweights, max_slice_h, slice_h;
538     const int bsize = 1 << s->n;
539
540     switch (inlink->format) {
541     case AV_PIX_FMT_BGR24:
542         s->color_decorrelation = color_decorrelation_bgr;
543         s->color_correlation   = color_correlation_bgr;
544         break;
545     case AV_PIX_FMT_RGB24:
546         s->color_decorrelation = color_decorrelation_rgb;
547         s->color_correlation   = color_correlation_rgb;
548         break;
549     case AV_PIX_FMT_GBRP:
550         s->color_decorrelation = color_decorrelation_gbrp;
551         s->color_correlation   = color_correlation_gbrp;
552         break;
553     default:
554         av_assert0(0);
555     }
556
557     s->pr_width  = inlink->w - (inlink->w - bsize) % s->step;
558     s->pr_height = inlink->h - (inlink->h - bsize) % s->step;
559     if (s->pr_width != inlink->w)
560         av_log(ctx, AV_LOG_WARNING, "The last %d horizontal pixels won't be denoised\n",
561                inlink->w - s->pr_width);
562     if (s->pr_height != inlink->h)
563         av_log(ctx, AV_LOG_WARNING, "The last %d vertical pixels won't be denoised\n",
564                inlink->h - s->pr_height);
565
566     max_slice_h = s->pr_height / ((s->bsize - 1) * 2);
567     s->nb_threads = FFMIN3(MAX_THREADS, ff_filter_get_nb_threads(ctx), max_slice_h);
568     av_log(ctx, AV_LOG_DEBUG, "threads: [max=%d hmax=%d user=%d] => %d\n",
569            MAX_THREADS, max_slice_h, ff_filter_get_nb_threads(ctx), s->nb_threads);
570
571     s->p_linesize = linesize = FFALIGN(s->pr_width, 32);
572     for (i = 0; i < 2; i++) {
573         s->cbuf[i][0] = av_malloc_array(linesize * s->pr_height, sizeof(*s->cbuf[i][0]));
574         s->cbuf[i][1] = av_malloc_array(linesize * s->pr_height, sizeof(*s->cbuf[i][1]));
575         s->cbuf[i][2] = av_malloc_array(linesize * s->pr_height, sizeof(*s->cbuf[i][2]));
576         if (!s->cbuf[i][0] || !s->cbuf[i][1] || !s->cbuf[i][2])
577             return AVERROR(ENOMEM);
578     }
579
580     /* eval expressions are probably not thread safe when the eval internal
581      * state can be changed (typically through load & store operations) */
582     if (s->expr_str) {
583         for (i = 0; i < s->nb_threads; i++) {
584             int ret = av_expr_parse(&s->expr[i], s->expr_str, var_names,
585                                     NULL, NULL, NULL, NULL, 0, ctx);
586             if (ret < 0)
587                 return ret;
588         }
589     }
590
591     /* each slice will need to (pre & re)process the top and bottom block of
592      * the previous one in in addition to its processing area. This is because
593      * each pixel is averaged by all the surrounding blocks */
594     slice_h = (int)ceilf(s->pr_height / (float)s->nb_threads) + (s->bsize - 1) * 2;
595     for (i = 0; i < s->nb_threads; i++) {
596         s->slices[i] = av_malloc_array(linesize, slice_h * sizeof(*s->slices[i]));
597         if (!s->slices[i])
598             return AVERROR(ENOMEM);
599     }
600
601     s->weights = av_malloc(s->pr_height * linesize * sizeof(*s->weights));
602     if (!s->weights)
603         return AVERROR(ENOMEM);
604     iweights = av_calloc(s->pr_height, linesize * sizeof(*iweights));
605     if (!iweights)
606         return AVERROR(ENOMEM);
607     for (y = 0; y < s->pr_height - bsize + 1; y += s->step)
608         for (x = 0; x < s->pr_width - bsize + 1; x += s->step)
609             for (by = 0; by < bsize; by++)
610                 for (bx = 0; bx < bsize; bx++)
611                     iweights[(y + by)*linesize + x + bx]++;
612     for (y = 0; y < s->pr_height; y++)
613         for (x = 0; x < s->pr_width; x++)
614             s->weights[y*linesize + x] = 1. / iweights[y*linesize + x];
615     av_free(iweights);
616
617     return 0;
618 }
619
620 static av_cold int init(AVFilterContext *ctx)
621 {
622     DCTdnoizContext *s = ctx->priv;
623
624     s->bsize = 1 << s->n;
625     if (s->overlap == -1)
626         s->overlap = s->bsize - 1;
627
628     if (s->overlap > s->bsize - 1) {
629         av_log(s, AV_LOG_ERROR, "Overlap value can not except %d "
630                "with a block size of %dx%d\n",
631                s->bsize - 1, s->bsize, s->bsize);
632         return AVERROR(EINVAL);
633     }
634
635     if (s->expr_str) {
636         switch (s->n) {
637         case 3: s->filter_freq_func = filter_freq_expr_8;  break;
638         case 4: s->filter_freq_func = filter_freq_expr_16; break;
639         default: av_assert0(0);
640         }
641     } else {
642         switch (s->n) {
643         case 3: s->filter_freq_func = filter_freq_sigma_8;  break;
644         case 4: s->filter_freq_func = filter_freq_sigma_16; break;
645         default: av_assert0(0);
646         }
647     }
648
649     s->th   = s->sigma * 3.;
650     s->step = s->bsize - s->overlap;
651     return 0;
652 }
653
654 static int query_formats(AVFilterContext *ctx)
655 {
656     static const enum AVPixelFormat pix_fmts[] = {
657         AV_PIX_FMT_BGR24, AV_PIX_FMT_RGB24,
658         AV_PIX_FMT_GBRP,
659         AV_PIX_FMT_NONE
660     };
661     AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts);
662     if (!fmts_list)
663         return AVERROR(ENOMEM);
664     return ff_set_common_formats(ctx, fmts_list);
665 }
666
667 typedef struct ThreadData {
668     float *src, *dst;
669 } ThreadData;
670
671 static int filter_slice(AVFilterContext *ctx,
672                         void *arg, int jobnr, int nb_jobs)
673 {
674     int x, y;
675     DCTdnoizContext *s = ctx->priv;
676     const ThreadData *td = arg;
677     const int w = s->pr_width;
678     const int h = s->pr_height;
679     const int slice_start = (h *  jobnr   ) / nb_jobs;
680     const int slice_end   = (h * (jobnr+1)) / nb_jobs;
681     const int slice_start_ctx = FFMAX(slice_start - s->bsize + 1, 0);
682     const int slice_end_ctx   = FFMIN(slice_end, h - s->bsize + 1);
683     const int slice_h = slice_end_ctx - slice_start_ctx;
684     const int src_linesize   = s->p_linesize;
685     const int dst_linesize   = s->p_linesize;
686     const int slice_linesize = s->p_linesize;
687     float *dst;
688     const float *src = td->src + slice_start_ctx * src_linesize;
689     const float *weights = s->weights + slice_start * dst_linesize;
690     float *slice = s->slices[jobnr];
691
692     // reset block sums
693     memset(slice, 0, (slice_h + s->bsize - 1) * dst_linesize * sizeof(*slice));
694
695     // block dct sums
696     for (y = 0; y < slice_h; y += s->step) {
697         for (x = 0; x < w - s->bsize + 1; x += s->step)
698             s->filter_freq_func(s, src + x, src_linesize,
699                                 slice + x, slice_linesize,
700                                 jobnr);
701         src += s->step * src_linesize;
702         slice += s->step * slice_linesize;
703     }
704
705     // average blocks
706     slice = s->slices[jobnr] + (slice_start - slice_start_ctx) * slice_linesize;
707     dst = td->dst + slice_start * dst_linesize;
708     for (y = slice_start; y < slice_end; y++) {
709         for (x = 0; x < w; x++)
710             dst[x] = slice[x] * weights[x];
711         slice += slice_linesize;
712         dst += dst_linesize;
713         weights += dst_linesize;
714     }
715
716     return 0;
717 }
718
719 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
720 {
721     AVFilterContext *ctx = inlink->dst;
722     DCTdnoizContext *s = ctx->priv;
723     AVFilterLink *outlink = inlink->dst->outputs[0];
724     int direct, plane;
725     AVFrame *out;
726
727     if (av_frame_is_writable(in)) {
728         direct = 1;
729         out = in;
730     } else {
731         direct = 0;
732         out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
733         if (!out) {
734             av_frame_free(&in);
735             return AVERROR(ENOMEM);
736         }
737         av_frame_copy_props(out, in);
738     }
739
740     s->color_decorrelation(s->cbuf[0], s->p_linesize,
741                            (const uint8_t **)in->data, in->linesize[0],
742                            s->pr_width, s->pr_height);
743     for (plane = 0; plane < 3; plane++) {
744         ThreadData td = {
745             .src = s->cbuf[0][plane],
746             .dst = s->cbuf[1][plane],
747         };
748         ctx->internal->execute(ctx, filter_slice, &td, NULL, s->nb_threads);
749     }
750     s->color_correlation(out->data, out->linesize[0],
751                          s->cbuf[1], s->p_linesize,
752                          s->pr_width, s->pr_height);
753
754     if (!direct) {
755         int y;
756         uint8_t *dst = out->data[0];
757         const uint8_t *src = in->data[0];
758         const int dst_linesize = out->linesize[0];
759         const int src_linesize = in->linesize[0];
760         const int hpad = (inlink->w - s->pr_width) * 3;
761         const int vpad = (inlink->h - s->pr_height);
762
763         if (hpad) {
764             uint8_t       *dstp = dst + s->pr_width * 3;
765             const uint8_t *srcp = src + s->pr_width * 3;
766
767             for (y = 0; y < s->pr_height; y++) {
768                 memcpy(dstp, srcp, hpad);
769                 dstp += dst_linesize;
770                 srcp += src_linesize;
771             }
772         }
773         if (vpad) {
774             uint8_t       *dstp = dst + s->pr_height * dst_linesize;
775             const uint8_t *srcp = src + s->pr_height * src_linesize;
776
777             for (y = 0; y < vpad; y++) {
778                 memcpy(dstp, srcp, inlink->w * 3);
779                 dstp += dst_linesize;
780                 srcp += src_linesize;
781             }
782         }
783
784         av_frame_free(&in);
785     }
786
787     return ff_filter_frame(outlink, out);
788 }
789
790 static av_cold void uninit(AVFilterContext *ctx)
791 {
792     int i;
793     DCTdnoizContext *s = ctx->priv;
794
795     av_freep(&s->weights);
796     for (i = 0; i < 2; i++) {
797         av_freep(&s->cbuf[i][0]);
798         av_freep(&s->cbuf[i][1]);
799         av_freep(&s->cbuf[i][2]);
800     }
801     for (i = 0; i < s->nb_threads; i++) {
802         av_freep(&s->slices[i]);
803         av_expr_free(s->expr[i]);
804     }
805 }
806
807 static const AVFilterPad dctdnoiz_inputs[] = {
808     {
809         .name         = "default",
810         .type         = AVMEDIA_TYPE_VIDEO,
811         .filter_frame = filter_frame,
812         .config_props = config_input,
813     },
814     { NULL }
815 };
816
817 static const AVFilterPad dctdnoiz_outputs[] = {
818     {
819         .name = "default",
820         .type = AVMEDIA_TYPE_VIDEO,
821     },
822     { NULL }
823 };
824
825 AVFilter ff_vf_dctdnoiz = {
826     .name          = "dctdnoiz",
827     .description   = NULL_IF_CONFIG_SMALL("Denoise frames using 2D DCT."),
828     .priv_size     = sizeof(DCTdnoizContext),
829     .init          = init,
830     .uninit        = uninit,
831     .query_formats = query_formats,
832     .inputs        = dctdnoiz_inputs,
833     .outputs       = dctdnoiz_outputs,
834     .priv_class    = &dctdnoiz_class,
835     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS,
836 };