avfilter/dctdnoiz: add 8x8 dct and make it the default
[ffmpeg.git] / 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/opt.h"
34 #include "drawutils.h"
35 #include "internal.h"
36
37 static const char *const var_names[] = { "c", NULL };
38 enum { VAR_C, VAR_VARS_NB };
39
40 typedef struct DCTdnoizContext {
41     const AVClass *class;
42
43     /* coefficient factor expression */
44     char *expr_str;
45     AVExpr *expr;
46     double var_values[VAR_VARS_NB];
47
48     int pr_width, pr_height;    // width and height to process
49     float sigma;                // used when no expression are st
50     float th;                   // threshold (3*sigma)
51     float color_dct[3][3];      // 3x3 DCT for color decorrelation
52     float *cbuf[2][3];          // two planar rgb color buffers
53     float *weights;             // dct coeff are cumulated with overlapping; these values are used for averaging
54     int p_linesize;             // line sizes for color and weights
55     int overlap;                // number of block overlapping pixels
56     int step;                   // block step increment (blocksize - overlap)
57     int n;                      // 1<<n is the block size
58     int bsize;                  // block size, 1<<n
59     void (*filter_freq_func)(struct DCTdnoizContext *s,
60                              const float *src, int src_linesize,
61                              float *dst, int dst_linesize);
62 } DCTdnoizContext;
63
64 #define MIN_NBITS 3 /* blocksize = 1<<3 =  8 */
65 #define MAX_NBITS 4 /* blocksize = 1<<4 = 16 */
66 #define DEFAULT_NBITS 3
67
68 #define OFFSET(x) offsetof(DCTdnoizContext, x)
69 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
70 static const AVOption dctdnoiz_options[] = {
71     { "sigma",   "set noise sigma constant",               OFFSET(sigma),    AV_OPT_TYPE_FLOAT,  {.dbl=0},            0, 999,          .flags = FLAGS },
72     { "s",       "set noise sigma constant",               OFFSET(sigma),    AV_OPT_TYPE_FLOAT,  {.dbl=0},            0, 999,          .flags = FLAGS },
73     { "overlap", "set number of block overlapping pixels", OFFSET(overlap),  AV_OPT_TYPE_INT,    {.i64=-1}, -1, (1<<MAX_NBITS)-1, .flags = FLAGS },
74     { "expr",    "set coefficient factor expression",      OFFSET(expr_str), AV_OPT_TYPE_STRING, {.str=NULL},                          .flags = FLAGS },
75     { "e",       "set coefficient factor expression",      OFFSET(expr_str), AV_OPT_TYPE_STRING, {.str=NULL},                          .flags = FLAGS },
76     { "n",       "set the block size, expressed in bits",  OFFSET(n),        AV_OPT_TYPE_INT,    {.i64=DEFAULT_NBITS}, MIN_NBITS, MAX_NBITS, .flags = FLAGS },
77     { NULL }
78 };
79
80 AVFILTER_DEFINE_CLASS(dctdnoiz);
81
82 static void av_always_inline fdct8_1d(float *dst, const float *src,
83                                       int dst_stridea, int dst_strideb,
84                                       int src_stridea, int src_strideb)
85 {
86     int i;
87
88     for (i = 0; i < 8; i++) {
89         const float x00 = src[0*src_stridea] + src[7*src_stridea];
90         const float x01 = src[1*src_stridea] + src[6*src_stridea];
91         const float x02 = src[2*src_stridea] + src[5*src_stridea];
92         const float x03 = src[3*src_stridea] + src[4*src_stridea];
93         const float x04 = src[0*src_stridea] - src[7*src_stridea];
94         const float x05 = src[1*src_stridea] - src[6*src_stridea];
95         const float x06 = src[2*src_stridea] - src[5*src_stridea];
96         const float x07 = src[3*src_stridea] - src[4*src_stridea];
97         const float x08 = x00 + x03;
98         const float x09 = x01 + x02;
99         const float x0a = x00 - x03;
100         const float x0b = x01 - x02;
101         const float x0c = 1.38703984532215*x04 + 0.275899379282943*x07;
102         const float x0d = 1.17587560241936*x05 + 0.785694958387102*x06;
103         const float x0e = -0.785694958387102*x05 + 1.17587560241936*x06;
104         const float x0f = 0.275899379282943*x04 - 1.38703984532215*x07;
105         const float x10 = 0.353553390593274 * (x0c - x0d);
106         const float x11 = 0.353553390593274 * (x0e - x0f);
107         dst[0*dst_stridea] = 0.353553390593274 * (x08 + x09);
108         dst[1*dst_stridea] = 0.353553390593274 * (x0c + x0d);
109         dst[2*dst_stridea] = 0.461939766255643*x0a + 0.191341716182545*x0b;
110         dst[3*dst_stridea] = 0.707106781186547 * (x10 - x11);
111         dst[4*dst_stridea] = 0.353553390593274 * (x08 - x09);
112         dst[5*dst_stridea] = 0.707106781186547 * (x10 + x11);
113         dst[6*dst_stridea] = 0.191341716182545*x0a - 0.461939766255643*x0b;
114         dst[7*dst_stridea] = 0.353553390593274 * (x0e + x0f);
115         dst += dst_strideb;
116         src += src_strideb;
117     }
118 }
119
120 static void av_always_inline idct8_1d(float *dst, const float *src,
121                                       int dst_stridea, int dst_strideb,
122                                       int src_stridea, int src_strideb,
123                                       int add)
124 {
125     int i;
126
127     for (i = 0; i < 8; i++) {
128         const float x00 = 1.4142135623731*src[0*src_stridea];
129         const float x01 = 1.38703984532215*src[1*src_stridea] + 0.275899379282943*src[7*src_stridea];
130         const float x02 = 1.30656296487638*src[2*src_stridea] + 0.541196100146197*src[6*src_stridea];
131         const float x03 = 1.17587560241936*src[3*src_stridea] + 0.785694958387102*src[5*src_stridea];
132         const float x04 = 1.4142135623731*src[4*src_stridea];
133         const float x05 = -0.785694958387102*src[3*src_stridea] + 1.17587560241936*src[5*src_stridea];
134         const float x06 = 0.541196100146197*src[2*src_stridea] - 1.30656296487638*src[6*src_stridea];
135         const float x07 = -0.275899379282943*src[1*src_stridea] + 1.38703984532215*src[7*src_stridea];
136         const float x09 = x00 + x04;
137         const float x0a = x01 + x03;
138         const float x0b = 1.4142135623731*x02;
139         const float x0c = x00 - x04;
140         const float x0d = x01 - x03;
141         const float x0e = 0.353553390593274 * (x09 - x0b);
142         const float x0f = 0.353553390593274 * (x0c + x0d);
143         const float x10 = 0.353553390593274 * (x0c - x0d);
144         const float x11 = 1.4142135623731*x06;
145         const float x12 = x05 + x07;
146         const float x13 = x05 - x07;
147         const float x14 = 0.353553390593274 * (x11 + x12);
148         const float x15 = 0.353553390593274 * (x11 - x12);
149         const float x16 = 0.5*x13;
150         const float x08 = -x15;
151         dst[0*dst_stridea] = (add ? dst[ 0*dst_stridea] : 0) + 0.25 * (x09 + x0b) + 0.353553390593274*x0a;
152         dst[1*dst_stridea] = (add ? dst[ 1*dst_stridea] : 0) + 0.707106781186547 * (x0f - x08);
153         dst[2*dst_stridea] = (add ? dst[ 2*dst_stridea] : 0) + 0.707106781186547 * (x0f + x08);
154         dst[3*dst_stridea] = (add ? dst[ 3*dst_stridea] : 0) + 0.707106781186547 * (x0e + x16);
155         dst[4*dst_stridea] = (add ? dst[ 4*dst_stridea] : 0) + 0.707106781186547 * (x0e - x16);
156         dst[5*dst_stridea] = (add ? dst[ 5*dst_stridea] : 0) + 0.707106781186547 * (x10 - x14);
157         dst[6*dst_stridea] = (add ? dst[ 6*dst_stridea] : 0) + 0.707106781186547 * (x10 + x14);
158         dst[7*dst_stridea] = (add ? dst[ 7*dst_stridea] : 0) + 0.25 * (x09 + x0b) - 0.353553390593274*x0a;
159         dst += dst_strideb;
160         src += src_strideb;
161     }
162 }
163
164
165 static void av_always_inline fdct16_1d(float *dst, const float *src,
166                                        int dst_stridea, int dst_strideb,
167                                        int src_stridea, int src_strideb)
168 {
169     int i;
170
171     for (i = 0; i < 16; i++) {
172         const float x00 = src[ 0*src_stridea] + src[15*src_stridea];
173         const float x01 = src[ 1*src_stridea] + src[14*src_stridea];
174         const float x02 = src[ 2*src_stridea] + src[13*src_stridea];
175         const float x03 = src[ 3*src_stridea] + src[12*src_stridea];
176         const float x04 = src[ 4*src_stridea] + src[11*src_stridea];
177         const float x05 = src[ 5*src_stridea] + src[10*src_stridea];
178         const float x06 = src[ 6*src_stridea] + src[ 9*src_stridea];
179         const float x07 = src[ 7*src_stridea] + src[ 8*src_stridea];
180         const float x08 = src[ 0*src_stridea] - src[15*src_stridea];
181         const float x09 = src[ 1*src_stridea] - src[14*src_stridea];
182         const float x0a = src[ 2*src_stridea] - src[13*src_stridea];
183         const float x0b = src[ 3*src_stridea] - src[12*src_stridea];
184         const float x0c = src[ 4*src_stridea] - src[11*src_stridea];
185         const float x0d = src[ 5*src_stridea] - src[10*src_stridea];
186         const float x0e = src[ 6*src_stridea] - src[ 9*src_stridea];
187         const float x0f = src[ 7*src_stridea] - src[ 8*src_stridea];
188         const float x10 = x00 + x07;
189         const float x11 = x01 + x06;
190         const float x12 = x02 + x05;
191         const float x13 = x03 + x04;
192         const float x14 = x00 - x07;
193         const float x15 = x01 - x06;
194         const float x16 = x02 - x05;
195         const float x17 = x03 - x04;
196         const float x18 = x10 + x13;
197         const float x19 = x11 + x12;
198         const float x1a = x10 - x13;
199         const float x1b = x11 - x12;
200         const float x1c = 1.38703984532215*x14 + 0.275899379282943*x17;
201         const float x1d = 1.17587560241936*x15 + 0.785694958387102*x16;
202         const float x1e = -0.785694958387102*x15 + 1.17587560241936*x16;
203         const float x1f = 0.275899379282943*x14 - 1.38703984532215*x17;
204         const float x20 = 0.25 * (x1c - x1d);
205         const float x21 = 0.25 * (x1e - x1f);
206         const float x22 = 1.40740373752638*x08 + 0.138617169199091*x0f;
207         const float x23 = 1.35331800117435*x09 + 0.410524527522357*x0e;
208         const float x24 = 1.24722501298667*x0a + 0.666655658477747*x0d;
209         const float x25 = 1.09320186700176*x0b + 0.897167586342636*x0c;
210         const float x26 = -0.897167586342636*x0b + 1.09320186700176*x0c;
211         const float x27 = 0.666655658477747*x0a - 1.24722501298667*x0d;
212         const float x28 = -0.410524527522357*x09 + 1.35331800117435*x0e;
213         const float x29 = 0.138617169199091*x08 - 1.40740373752638*x0f;
214         const float x2a = x22 + x25;
215         const float x2b = x23 + x24;
216         const float x2c = x22 - x25;
217         const float x2d = x23 - x24;
218         const float x2e = 0.25 * (x2a - x2b);
219         const float x2f = 0.326640741219094*x2c + 0.135299025036549*x2d;
220         const float x30 = 0.135299025036549*x2c - 0.326640741219094*x2d;
221         const float x31 = x26 + x29;
222         const float x32 = x27 + x28;
223         const float x33 = x26 - x29;
224         const float x34 = x27 - x28;
225         const float x35 = 0.25 * (x31 - x32);
226         const float x36 = 0.326640741219094*x33 + 0.135299025036549*x34;
227         const float x37 = 0.135299025036549*x33 - 0.326640741219094*x34;
228         dst[ 0*dst_stridea] = 0.25 * (x18 + x19);
229         dst[ 1*dst_stridea] = 0.25 * (x2a + x2b);
230         dst[ 2*dst_stridea] = 0.25 * (x1c + x1d);
231         dst[ 3*dst_stridea] = 0.707106781186547 * (x2f - x37);
232         dst[ 4*dst_stridea] = 0.326640741219094*x1a + 0.135299025036549*x1b;
233         dst[ 5*dst_stridea] = 0.707106781186547 * (x2f + x37);
234         dst[ 6*dst_stridea] = 0.707106781186547 * (x20 - x21);
235         dst[ 7*dst_stridea] = 0.707106781186547 * (x2e + x35);
236         dst[ 8*dst_stridea] = 0.25 * (x18 - x19);
237         dst[ 9*dst_stridea] = 0.707106781186547 * (x2e - x35);
238         dst[10*dst_stridea] = 0.707106781186547 * (x20 + x21);
239         dst[11*dst_stridea] = 0.707106781186547 * (x30 - x36);
240         dst[12*dst_stridea] = 0.135299025036549*x1a - 0.326640741219094*x1b;
241         dst[13*dst_stridea] = 0.707106781186547 * (x30 + x36);
242         dst[14*dst_stridea] = 0.25 * (x1e + x1f);
243         dst[15*dst_stridea] = 0.25 * (x31 + x32);
244         dst += dst_strideb;
245         src += src_strideb;
246     }
247 }
248
249 static void av_always_inline idct16_1d(float *dst, const float *src,
250                                        int dst_stridea, int dst_strideb,
251                                        int src_stridea, int src_strideb,
252                                        int add)
253 {
254     int i;
255
256     for (i = 0; i < 16; i++) {
257         const float x00 =  1.4142135623731  *src[ 0*src_stridea];
258         const float x01 =  1.40740373752638 *src[ 1*src_stridea] + 0.138617169199091*src[15*src_stridea];
259         const float x02 =  1.38703984532215 *src[ 2*src_stridea] + 0.275899379282943*src[14*src_stridea];
260         const float x03 =  1.35331800117435 *src[ 3*src_stridea] + 0.410524527522357*src[13*src_stridea];
261         const float x04 =  1.30656296487638 *src[ 4*src_stridea] + 0.541196100146197*src[12*src_stridea];
262         const float x05 =  1.24722501298667 *src[ 5*src_stridea] + 0.666655658477747*src[11*src_stridea];
263         const float x06 =  1.17587560241936 *src[ 6*src_stridea] + 0.785694958387102*src[10*src_stridea];
264         const float x07 =  1.09320186700176 *src[ 7*src_stridea] + 0.897167586342636*src[ 9*src_stridea];
265         const float x08 =  1.4142135623731  *src[ 8*src_stridea];
266         const float x09 = -0.897167586342636*src[ 7*src_stridea] + 1.09320186700176*src[ 9*src_stridea];
267         const float x0a =  0.785694958387102*src[ 6*src_stridea] - 1.17587560241936*src[10*src_stridea];
268         const float x0b = -0.666655658477747*src[ 5*src_stridea] + 1.24722501298667*src[11*src_stridea];
269         const float x0c =  0.541196100146197*src[ 4*src_stridea] - 1.30656296487638*src[12*src_stridea];
270         const float x0d = -0.410524527522357*src[ 3*src_stridea] + 1.35331800117435*src[13*src_stridea];
271         const float x0e =  0.275899379282943*src[ 2*src_stridea] - 1.38703984532215*src[14*src_stridea];
272         const float x0f = -0.138617169199091*src[ 1*src_stridea] + 1.40740373752638*src[15*src_stridea];
273         const float x12 = x00 + x08;
274         const float x13 = x01 + x07;
275         const float x14 = x02 + x06;
276         const float x15 = x03 + x05;
277         const float x16 = 1.4142135623731*x04;
278         const float x17 = x00 - x08;
279         const float x18 = x01 - x07;
280         const float x19 = x02 - x06;
281         const float x1a = x03 - x05;
282         const float x1d = x12 + x16;
283         const float x1e = x13 + x15;
284         const float x1f = 1.4142135623731*x14;
285         const float x20 = x12 - x16;
286         const float x21 = x13 - x15;
287         const float x22 = 0.25 * (x1d - x1f);
288         const float x23 = 0.25 * (x20 + x21);
289         const float x24 = 0.25 * (x20 - x21);
290         const float x25 = 1.4142135623731*x17;
291         const float x26 = 1.30656296487638*x18 + 0.541196100146197*x1a;
292         const float x27 = 1.4142135623731*x19;
293         const float x28 = -0.541196100146197*x18 + 1.30656296487638*x1a;
294         const float x29 = 0.176776695296637 * (x25 + x27) + 0.25*x26;
295         const float x2a = 0.25 * (x25 - x27);
296         const float x2b = 0.176776695296637 * (x25 + x27) - 0.25*x26;
297         const float x2c = 0.353553390593274*x28;
298         const float x1b = 0.707106781186547 * (x2a - x2c);
299         const float x1c = 0.707106781186547 * (x2a + x2c);
300         const float x2d = 1.4142135623731*x0c;
301         const float x2e = x0b + x0d;
302         const float x2f = x0a + x0e;
303         const float x30 = x09 + x0f;
304         const float x31 = x09 - x0f;
305         const float x32 = x0a - x0e;
306         const float x33 = x0b - x0d;
307         const float x37 = 1.4142135623731*x2d;
308         const float x38 = 1.30656296487638*x2e + 0.541196100146197*x30;
309         const float x39 = 1.4142135623731*x2f;
310         const float x3a = -0.541196100146197*x2e + 1.30656296487638*x30;
311         const float x3b = 0.176776695296637 * (x37 + x39) + 0.25*x38;
312         const float x3c = 0.25 * (x37 - x39);
313         const float x3d = 0.176776695296637 * (x37 + x39) - 0.25*x38;
314         const float x3e = 0.353553390593274*x3a;
315         const float x34 = 0.707106781186547 * (x3c - x3e);
316         const float x35 = 0.707106781186547 * (x3c + x3e);
317         const float x3f = 1.4142135623731*x32;
318         const float x40 = x31 + x33;
319         const float x41 = x31 - x33;
320         const float x42 = 0.25 * (x3f + x40);
321         const float x43 = 0.25 * (x3f - x40);
322         const float x44 = 0.353553390593274*x41;
323         const float x36 = -x43;
324         const float x10 = -x34;
325         const float x11 = -x3d;
326         dst[ 0*dst_stridea] = (add ? dst[ 0*dst_stridea] : 0) + 0.176776695296637 * (x1d + x1f) + 0.25*x1e;
327         dst[ 1*dst_stridea] = (add ? dst[ 1*dst_stridea] : 0) + 0.707106781186547 * (x29 - x11);
328         dst[ 2*dst_stridea] = (add ? dst[ 2*dst_stridea] : 0) + 0.707106781186547 * (x29 + x11);
329         dst[ 3*dst_stridea] = (add ? dst[ 3*dst_stridea] : 0) + 0.707106781186547 * (x23 + x36);
330         dst[ 4*dst_stridea] = (add ? dst[ 4*dst_stridea] : 0) + 0.707106781186547 * (x23 - x36);
331         dst[ 5*dst_stridea] = (add ? dst[ 5*dst_stridea] : 0) + 0.707106781186547 * (x1b - x35);
332         dst[ 6*dst_stridea] = (add ? dst[ 6*dst_stridea] : 0) + 0.707106781186547 * (x1b + x35);
333         dst[ 7*dst_stridea] = (add ? dst[ 7*dst_stridea] : 0) + 0.707106781186547 * (x22 + x44);
334         dst[ 8*dst_stridea] = (add ? dst[ 8*dst_stridea] : 0) + 0.707106781186547 * (x22 - x44);
335         dst[ 9*dst_stridea] = (add ? dst[ 9*dst_stridea] : 0) + 0.707106781186547 * (x1c - x10);
336         dst[10*dst_stridea] = (add ? dst[10*dst_stridea] : 0) + 0.707106781186547 * (x1c + x10);
337         dst[11*dst_stridea] = (add ? dst[11*dst_stridea] : 0) + 0.707106781186547 * (x24 + x42);
338         dst[12*dst_stridea] = (add ? dst[12*dst_stridea] : 0) + 0.707106781186547 * (x24 - x42);
339         dst[13*dst_stridea] = (add ? dst[13*dst_stridea] : 0) + 0.707106781186547 * (x2b - x3b);
340         dst[14*dst_stridea] = (add ? dst[14*dst_stridea] : 0) + 0.707106781186547 * (x2b + x3b);
341         dst[15*dst_stridea] = (add ? dst[15*dst_stridea] : 0) + 0.176776695296637 * (x1d + x1f) - 0.25*x1e;
342         dst += dst_strideb;
343         src += src_strideb;
344     }
345 }
346
347 #define DEF_FILTER_FREQ_FUNCS(bsize)                                                        \
348 static av_always_inline void filter_freq_##bsize(const float *src, int src_linesize,        \
349                                                  float *dst, int dst_linesize,              \
350                                                  AVExpr *expr, double *var_values,          \
351                                                  int sigma_th)                              \
352 {                                                                                           \
353     unsigned i;                                                                             \
354     DECLARE_ALIGNED(32, float, tmp_block1)[bsize * bsize];                                  \
355     DECLARE_ALIGNED(32, float, tmp_block2)[bsize * bsize];                                  \
356                                                                                             \
357     /* forward DCT */                                                                       \
358     fdct##bsize##_1d(tmp_block1, src, 1, bsize, 1, src_linesize);                           \
359     fdct##bsize##_1d(tmp_block2, tmp_block1, bsize, 1, bsize, 1);                           \
360                                                                                             \
361     for (i = 0; i < bsize*bsize; i++) {                                                     \
362         float *b = &tmp_block2[i];                                                          \
363         /* frequency filtering */                                                           \
364         if (expr) {                                                                         \
365             var_values[VAR_C] = FFABS(*b);                                                  \
366             *b *= av_expr_eval(expr, var_values, NULL);                                     \
367         } else {                                                                            \
368             if (FFABS(*b) < sigma_th)                                                       \
369                 *b = 0;                                                                     \
370         }                                                                                   \
371     }                                                                                       \
372                                                                                             \
373     /* inverse DCT */                                                                       \
374     idct##bsize##_1d(tmp_block1, tmp_block2, 1, bsize, 1, bsize, 0);                        \
375     idct##bsize##_1d(dst, tmp_block1, dst_linesize, 1, bsize, 1, 1);                        \
376 }                                                                                           \
377                                                                                             \
378 static void filter_freq_sigma_##bsize(DCTdnoizContext *s,                                   \
379                                       const float *src, int src_linesize,                   \
380                                       float *dst, int dst_linesize)                         \
381 {                                                                                           \
382     filter_freq_##bsize(src, src_linesize, dst, dst_linesize, NULL, NULL, s->th);           \
383 }                                                                                           \
384                                                                                             \
385 static void filter_freq_expr_##bsize(DCTdnoizContext *s,                                    \
386                                      const float *src, int src_linesize,                    \
387                                      float *dst, int dst_linesize)                          \
388 {                                                                                           \
389     filter_freq_##bsize(src, src_linesize, dst, dst_linesize, s->expr, s->var_values, 0);   \
390 }
391
392 DEF_FILTER_FREQ_FUNCS(8)
393 DEF_FILTER_FREQ_FUNCS(16)
394
395 static int config_input(AVFilterLink *inlink)
396 {
397     AVFilterContext *ctx = inlink->dst;
398     DCTdnoizContext *s = ctx->priv;
399     int i, x, y, bx, by, linesize, *iweights;
400     const int bsize = 1 << s->n;
401     const float dct_3x3[3][3] = {
402         { 1./sqrt(3),  1./sqrt(3),  1./sqrt(3) },
403         { 1./sqrt(2),           0, -1./sqrt(2) },
404         { 1./sqrt(6), -2./sqrt(6),  1./sqrt(6) },
405     };
406     uint8_t rgba_map[4];
407
408     ff_fill_rgba_map(rgba_map, inlink->format);
409     for (y = 0; y < 3; y++)
410         for (x = 0; x < 3; x++)
411             s->color_dct[y][x] = dct_3x3[rgba_map[y]][rgba_map[x]];
412
413     s->pr_width  = inlink->w - (inlink->w - bsize) % s->step;
414     s->pr_height = inlink->h - (inlink->h - bsize) % s->step;
415     if (s->pr_width != inlink->w)
416         av_log(ctx, AV_LOG_WARNING, "The last %d horizontal pixels won't be denoised\n",
417                inlink->w - s->pr_width);
418     if (s->pr_height != inlink->h)
419         av_log(ctx, AV_LOG_WARNING, "The last %d vertical pixels won't be denoised\n",
420                inlink->h - s->pr_height);
421
422     s->p_linesize = linesize = FFALIGN(s->pr_width, 32);
423     for (i = 0; i < 2; i++) {
424         s->cbuf[i][0] = av_malloc(linesize * s->pr_height * sizeof(*s->cbuf[i][0]));
425         s->cbuf[i][1] = av_malloc(linesize * s->pr_height * sizeof(*s->cbuf[i][1]));
426         s->cbuf[i][2] = av_malloc(linesize * s->pr_height * sizeof(*s->cbuf[i][2]));
427         if (!s->cbuf[i][0] || !s->cbuf[i][1] || !s->cbuf[i][2])
428             return AVERROR(ENOMEM);
429     }
430
431     s->weights = av_malloc(s->pr_height * linesize * sizeof(*s->weights));
432     if (!s->weights)
433         return AVERROR(ENOMEM);
434     iweights = av_calloc(s->pr_height, linesize * sizeof(*iweights));
435     if (!iweights)
436         return AVERROR(ENOMEM);
437     for (y = 0; y < s->pr_height - bsize + 1; y += s->step)
438         for (x = 0; x < s->pr_width - bsize + 1; x += s->step)
439             for (by = 0; by < bsize; by++)
440                 for (bx = 0; bx < bsize; bx++)
441                     iweights[(y + by)*linesize + x + bx]++;
442     for (y = 0; y < s->pr_height; y++)
443         for (x = 0; x < s->pr_width; x++)
444             s->weights[y*linesize + x] = 1. / iweights[y*linesize + x];
445     av_free(iweights);
446
447     return 0;
448 }
449
450 static av_cold int init(AVFilterContext *ctx)
451 {
452     DCTdnoizContext *s = ctx->priv;
453
454     s->bsize = 1 << s->n;
455     if (s->overlap == -1)
456         s->overlap = s->bsize - 1;
457
458     if (s->overlap > s->bsize - 1) {
459         av_log(s, AV_LOG_ERROR, "Overlap value can not except %d "
460                "with a block size of %dx%d\n",
461                s->bsize - 1, s->bsize, s->bsize);
462         return AVERROR(EINVAL);
463     }
464
465     if (s->expr_str) {
466         int ret = av_expr_parse(&s->expr, s->expr_str, var_names,
467                                 NULL, NULL, NULL, NULL, 0, ctx);
468         if (ret < 0)
469             return ret;
470         switch (s->n) {
471         case 3: s->filter_freq_func = filter_freq_expr_8;  break;
472         case 4: s->filter_freq_func = filter_freq_expr_16; break;
473         default: av_assert0(0);
474         }
475     } else {
476         switch (s->n) {
477         case 3: s->filter_freq_func = filter_freq_sigma_8;  break;
478         case 4: s->filter_freq_func = filter_freq_sigma_16; break;
479         default: av_assert0(0);
480         }
481     }
482
483     s->th   = s->sigma * 3.;
484     s->step = s->bsize - s->overlap;
485     return 0;
486 }
487
488 static int query_formats(AVFilterContext *ctx)
489 {
490     static const enum AVPixelFormat pix_fmts[] = {
491         AV_PIX_FMT_BGR24, AV_PIX_FMT_RGB24,
492         AV_PIX_FMT_NONE
493     };
494     ff_set_common_formats(ctx, ff_make_format_list(pix_fmts));
495     return 0;
496 }
497
498 static void color_decorrelation(float dct3ch[3][3], float **dst, int dst_linesize,
499                                 const uint8_t *src, int src_linesize, int w, int h)
500 {
501     int x, y;
502     float *dstp_r = dst[0];
503     float *dstp_g = dst[1];
504     float *dstp_b = dst[2];
505
506     for (y = 0; y < h; y++) {
507         const uint8_t *srcp = src;
508
509         for (x = 0; x < w; x++) {
510             dstp_r[x] = srcp[0] * dct3ch[0][0] + srcp[1] * dct3ch[0][1] + srcp[2] * dct3ch[0][2];
511             dstp_g[x] = srcp[0] * dct3ch[1][0] + srcp[1] * dct3ch[1][1] + srcp[2] * dct3ch[1][2];
512             dstp_b[x] = srcp[0] * dct3ch[2][0] + srcp[1] * dct3ch[2][1] + srcp[2] * dct3ch[2][2];
513             srcp += 3;
514         }
515         src += src_linesize;
516         dstp_r += dst_linesize;
517         dstp_g += dst_linesize;
518         dstp_b += dst_linesize;
519     }
520 }
521
522 static void color_correlation(float dct3ch[3][3], uint8_t *dst, int dst_linesize,
523                               float **src, int src_linesize, int w, int h)
524 {
525     int x, y;
526     const float *src_r = src[0];
527     const float *src_g = src[1];
528     const float *src_b = src[2];
529
530     for (y = 0; y < h; y++) {
531         uint8_t *dstp = dst;
532
533         for (x = 0; x < w; x++) {
534             dstp[0] = av_clip_uint8(src_r[x] * dct3ch[0][0] + src_g[x] * dct3ch[1][0] + src_b[x] * dct3ch[2][0]);
535             dstp[1] = av_clip_uint8(src_r[x] * dct3ch[0][1] + src_g[x] * dct3ch[1][1] + src_b[x] * dct3ch[2][1]);
536             dstp[2] = av_clip_uint8(src_r[x] * dct3ch[0][2] + src_g[x] * dct3ch[1][2] + src_b[x] * dct3ch[2][2]);
537             dstp += 3;
538         }
539         dst += dst_linesize;
540         src_r += src_linesize;
541         src_g += src_linesize;
542         src_b += src_linesize;
543     }
544 }
545
546 static void filter_plane(AVFilterContext *ctx,
547                          float *dst, int dst_linesize,
548                          const float *src, int src_linesize,
549                          int w, int h)
550 {
551     int x, y;
552     DCTdnoizContext *s = ctx->priv;
553     float *dst0 = dst;
554     const float *weights = s->weights;
555
556     // reset block sums
557     memset(dst, 0, h * dst_linesize * sizeof(*dst));
558
559     // block dct sums
560     for (y = 0; y < h - s->bsize + 1; y += s->step) {
561         for (x = 0; x < w - s->bsize + 1; x += s->step)
562             s->filter_freq_func(s, src + x, src_linesize,
563                                    dst + x, dst_linesize);
564         src += s->step * src_linesize;
565         dst += s->step * dst_linesize;
566     }
567
568     // average blocks
569     dst = dst0;
570     for (y = 0; y < h; y++) {
571         for (x = 0; x < w; x++)
572             dst[x] *= weights[x];
573         dst += dst_linesize;
574         weights += dst_linesize;
575     }
576 }
577
578 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
579 {
580     AVFilterContext *ctx = inlink->dst;
581     DCTdnoizContext *s = ctx->priv;
582     AVFilterLink *outlink = inlink->dst->outputs[0];
583     int direct, plane;
584     AVFrame *out;
585
586     if (av_frame_is_writable(in)) {
587         direct = 1;
588         out = in;
589     } else {
590         direct = 0;
591         out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
592         if (!out) {
593             av_frame_free(&in);
594             return AVERROR(ENOMEM);
595         }
596         av_frame_copy_props(out, in);
597     }
598
599     color_decorrelation(s->color_dct, s->cbuf[0], s->p_linesize,
600                         in->data[0], in->linesize[0], s->pr_width, s->pr_height);
601     for (plane = 0; plane < 3; plane++)
602         filter_plane(ctx, s->cbuf[1][plane], s->p_linesize,
603                           s->cbuf[0][plane], s->p_linesize,
604                           s->pr_width, s->pr_height);
605     color_correlation(s->color_dct, out->data[0], out->linesize[0],
606                       s->cbuf[1], s->p_linesize, s->pr_width, s->pr_height);
607
608     if (!direct) {
609         int y;
610         uint8_t *dst = out->data[0];
611         const uint8_t *src = in->data[0];
612         const int dst_linesize = out->linesize[0];
613         const int src_linesize = in->linesize[0];
614         const int hpad = (inlink->w - s->pr_width) * 3;
615         const int vpad = (inlink->h - s->pr_height);
616
617         if (hpad) {
618             uint8_t       *dstp = dst + s->pr_width * 3;
619             const uint8_t *srcp = src + s->pr_width * 3;
620
621             for (y = 0; y < s->pr_height; y++) {
622                 memcpy(dstp, srcp, hpad);
623                 dstp += dst_linesize;
624                 srcp += src_linesize;
625             }
626         }
627         if (vpad) {
628             uint8_t       *dstp = dst + s->pr_height * dst_linesize;
629             const uint8_t *srcp = src + s->pr_height * src_linesize;
630
631             for (y = 0; y < vpad; y++) {
632                 memcpy(dstp, srcp, inlink->w * 3);
633                 dstp += dst_linesize;
634                 srcp += src_linesize;
635             }
636         }
637
638         av_frame_free(&in);
639     }
640
641     return ff_filter_frame(outlink, out);
642 }
643
644 static av_cold void uninit(AVFilterContext *ctx)
645 {
646     int i;
647     DCTdnoizContext *s = ctx->priv;
648
649     av_free(s->weights);
650     for (i = 0; i < 2; i++) {
651         av_free(s->cbuf[i][0]);
652         av_free(s->cbuf[i][1]);
653         av_free(s->cbuf[i][2]);
654     }
655     av_expr_free(s->expr);
656 }
657
658 static const AVFilterPad dctdnoiz_inputs[] = {
659     {
660         .name         = "default",
661         .type         = AVMEDIA_TYPE_VIDEO,
662         .filter_frame = filter_frame,
663         .config_props = config_input,
664     },
665     { NULL }
666 };
667
668 static const AVFilterPad dctdnoiz_outputs[] = {
669     {
670         .name = "default",
671         .type = AVMEDIA_TYPE_VIDEO,
672     },
673     { NULL }
674 };
675
676 AVFilter ff_vf_dctdnoiz = {
677     .name          = "dctdnoiz",
678     .description   = NULL_IF_CONFIG_SMALL("Denoise frames using 2D DCT."),
679     .priv_size     = sizeof(DCTdnoizContext),
680     .init          = init,
681     .uninit        = uninit,
682     .query_formats = query_formats,
683     .inputs        = dctdnoiz_inputs,
684     .outputs       = dctdnoiz_outputs,
685     .priv_class    = &dctdnoiz_class,
686     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC,
687 };