avfilter/af_sofalizer: speed up fast convolution
[ffmpeg.git] / libavfilter / af_afftfilt.c
1 /*
2  * Copyright (c) 2016 Paul B Mahol
3  *
4  * This file is part of FFmpeg.
5  *
6  * FFmpeg is free software; you can redistribute it and/or modify it
7  * under the terms of the GNU Lesser General Public License as published
8  * by the Free Software Foundation; either version 2.1 of the License,
9  * 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 #include "libavutil/audio_fifo.h"
22 #include "libavutil/avstring.h"
23 #include "libavfilter/internal.h"
24 #include "libavutil/common.h"
25 #include "libavutil/opt.h"
26 #include "libavcodec/avfft.h"
27 #include "libavutil/eval.h"
28 #include "audio.h"
29 #include "window_func.h"
30
31 typedef struct AFFTFiltContext {
32     const AVClass *class;
33     char *real_str;
34     char *img_str;
35     int fft_bits;
36
37     FFTContext *fft, *ifft;
38     FFTComplex **fft_data;
39     FFTComplex **fft_temp;
40     int nb_exprs;
41     int window_size;
42     AVExpr **real;
43     AVExpr **imag;
44     AVAudioFifo *fifo;
45     int64_t pts;
46     int hop_size;
47     float overlap;
48     AVFrame *buffer;
49     int start, end;
50     int win_func;
51     float win_scale;
52     float *window_func_lut;
53 } AFFTFiltContext;
54
55 static const char *const var_names[] = {            "sr",     "b",       "nb",        "ch",        "chs",   "pts",     "re",     "im", NULL };
56 enum                                   { VAR_SAMPLE_RATE, VAR_BIN, VAR_NBBINS, VAR_CHANNEL, VAR_CHANNELS, VAR_PTS, VAR_REAL, VAR_IMAG, VAR_VARS_NB };
57
58 #define OFFSET(x) offsetof(AFFTFiltContext, x)
59 #define A AV_OPT_FLAG_AUDIO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
60
61 static const AVOption afftfilt_options[] = {
62     { "real", "set channels real expressions",       OFFSET(real_str), AV_OPT_TYPE_STRING, {.str = "re" }, 0, 0, A },
63     { "imag", "set channels imaginary expressions",  OFFSET(img_str),  AV_OPT_TYPE_STRING, {.str = "im" }, 0, 0, A },
64     { "win_size", "set window size", OFFSET(fft_bits), AV_OPT_TYPE_INT, {.i64=12}, 4, 17, A, "fft" },
65         { "w16",    0, 0, AV_OPT_TYPE_CONST, {.i64=4},  0, 0, A, "fft" },
66         { "w32",    0, 0, AV_OPT_TYPE_CONST, {.i64=5},  0, 0, A, "fft" },
67         { "w64",    0, 0, AV_OPT_TYPE_CONST, {.i64=6},  0, 0, A, "fft" },
68         { "w128",   0, 0, AV_OPT_TYPE_CONST, {.i64=7},  0, 0, A, "fft" },
69         { "w256",   0, 0, AV_OPT_TYPE_CONST, {.i64=8},  0, 0, A, "fft" },
70         { "w512",   0, 0, AV_OPT_TYPE_CONST, {.i64=9},  0, 0, A, "fft" },
71         { "w1024",  0, 0, AV_OPT_TYPE_CONST, {.i64=10}, 0, 0, A, "fft" },
72         { "w2048",  0, 0, AV_OPT_TYPE_CONST, {.i64=11}, 0, 0, A, "fft" },
73         { "w4096",  0, 0, AV_OPT_TYPE_CONST, {.i64=12}, 0, 0, A, "fft" },
74         { "w8192",  0, 0, AV_OPT_TYPE_CONST, {.i64=13}, 0, 0, A, "fft" },
75         { "w16384", 0, 0, AV_OPT_TYPE_CONST, {.i64=14}, 0, 0, A, "fft" },
76         { "w32768", 0, 0, AV_OPT_TYPE_CONST, {.i64=15}, 0, 0, A, "fft" },
77         { "w65536", 0, 0, AV_OPT_TYPE_CONST, {.i64=16}, 0, 0, A, "fft" },
78         { "w131072",0, 0, AV_OPT_TYPE_CONST, {.i64=17}, 0, 0, A, "fft" },
79     { "win_func", "set window function", OFFSET(win_func), AV_OPT_TYPE_INT, {.i64 = WFUNC_HANNING}, 0, NB_WFUNC-1, A, "win_func" },
80         { "rect",     "Rectangular",      0, AV_OPT_TYPE_CONST, {.i64=WFUNC_RECT},     0, 0, A, "win_func" },
81         { "bartlett", "Bartlett",         0, AV_OPT_TYPE_CONST, {.i64=WFUNC_BARTLETT}, 0, 0, A, "win_func" },
82         { "hann",     "Hann",             0, AV_OPT_TYPE_CONST, {.i64=WFUNC_HANNING},  0, 0, A, "win_func" },
83         { "hanning",  "Hanning",          0, AV_OPT_TYPE_CONST, {.i64=WFUNC_HANNING},  0, 0, A, "win_func" },
84         { "hamming",  "Hamming",          0, AV_OPT_TYPE_CONST, {.i64=WFUNC_HAMMING},  0, 0, A, "win_func" },
85         { "blackman", "Blackman",         0, AV_OPT_TYPE_CONST, {.i64=WFUNC_BLACKMAN}, 0, 0, A, "win_func" },
86         { "welch",    "Welch",            0, AV_OPT_TYPE_CONST, {.i64=WFUNC_WELCH},    0, 0, A, "win_func" },
87         { "flattop",  "Flat-top",         0, AV_OPT_TYPE_CONST, {.i64=WFUNC_FLATTOP},  0, 0, A, "win_func" },
88         { "bharris",  "Blackman-Harris",  0, AV_OPT_TYPE_CONST, {.i64=WFUNC_BHARRIS},  0, 0, A, "win_func" },
89         { "bnuttall", "Blackman-Nuttall", 0, AV_OPT_TYPE_CONST, {.i64=WFUNC_BNUTTALL}, 0, 0, A, "win_func" },
90         { "bhann",    "Bartlett-Hann",    0, AV_OPT_TYPE_CONST, {.i64=WFUNC_BHANN},    0, 0, A, "win_func" },
91         { "sine",     "Sine",             0, AV_OPT_TYPE_CONST, {.i64=WFUNC_SINE},     0, 0, A, "win_func" },
92         { "nuttall",  "Nuttall",          0, AV_OPT_TYPE_CONST, {.i64=WFUNC_NUTTALL},  0, 0, A, "win_func" },
93         { "lanczos",  "Lanczos",          0, AV_OPT_TYPE_CONST, {.i64=WFUNC_LANCZOS},  0, 0, A, "win_func" },
94         { "gauss",    "Gauss",            0, AV_OPT_TYPE_CONST, {.i64=WFUNC_GAUSS},    0, 0, A, "win_func" },
95         { "tukey",    "Tukey",            0, AV_OPT_TYPE_CONST, {.i64=WFUNC_TUKEY},    0, 0, A, "win_func" },
96         { "dolph",    "Dolph-Chebyshev",  0, AV_OPT_TYPE_CONST, {.i64=WFUNC_DOLPH},    0, 0, A, "win_func" },
97         { "cauchy",   "Cauchy",           0, AV_OPT_TYPE_CONST, {.i64=WFUNC_CAUCHY},   0, 0, A, "win_func" },
98         { "parzen",   "Parzen",           0, AV_OPT_TYPE_CONST, {.i64=WFUNC_PARZEN},   0, 0, A, "win_func" },
99         { "poisson",  "Poisson",          0, AV_OPT_TYPE_CONST, {.i64=WFUNC_POISSON},  0, 0, A, "win_func" },
100         { "bohman",   "Bohman",           0, AV_OPT_TYPE_CONST, {.i64=WFUNC_BOHMAN},   0, 0, A, "win_func" },
101     { "overlap", "set window overlap", OFFSET(overlap), AV_OPT_TYPE_FLOAT, {.dbl=0.75}, 0,  1, A },
102     { NULL },
103 };
104
105 AVFILTER_DEFINE_CLASS(afftfilt);
106
107 static inline double getreal(void *priv, double x, double ch)
108 {
109     AFFTFiltContext *s = priv;
110     int ich, ix;
111
112     ich = av_clip(ch, 0, s->nb_exprs - 1);
113     ix = av_clip(x, 0, s->window_size / 2);
114
115     return s->fft_data[ich][ix].re;
116 }
117
118 static inline double getimag(void *priv, double x, double ch)
119 {
120     AFFTFiltContext *s = priv;
121     int ich, ix;
122
123     ich = av_clip(ch, 0, s->nb_exprs - 1);
124     ix = av_clip(x, 0, s->window_size / 2);
125
126     return s->fft_data[ich][ix].im;
127 }
128
129 static double realf(void *priv, double x, double ch) { return getreal(priv, x, ch); }
130 static double imagf(void *priv, double x, double ch) { return getimag(priv, x, ch); }
131
132 static const char *const func2_names[]    = { "real", "imag", NULL };
133 double (*func2[])(void *, double, double) = {  realf,  imagf, NULL };
134
135 static int config_input(AVFilterLink *inlink)
136 {
137     AVFilterContext *ctx = inlink->dst;
138     AFFTFiltContext *s = ctx->priv;
139     char *saveptr = NULL;
140     int ret = 0, ch, i;
141     float overlap;
142     char *args;
143     const char *last_expr = "1";
144
145     s->pts  = AV_NOPTS_VALUE;
146     s->fft  = av_fft_init(s->fft_bits, 0);
147     s->ifft = av_fft_init(s->fft_bits, 1);
148     if (!s->fft || !s->ifft)
149         return AVERROR(ENOMEM);
150
151     s->window_size = 1 << s->fft_bits;
152
153     s->fft_data = av_calloc(inlink->channels, sizeof(*s->fft_data));
154     if (!s->fft_data)
155         return AVERROR(ENOMEM);
156
157     s->fft_temp = av_calloc(inlink->channels, sizeof(*s->fft_temp));
158     if (!s->fft_temp)
159         return AVERROR(ENOMEM);
160
161     for (ch = 0; ch < inlink->channels; ch++) {
162         s->fft_data[ch] = av_calloc(s->window_size, sizeof(**s->fft_data));
163         if (!s->fft_data[ch])
164             return AVERROR(ENOMEM);
165     }
166
167     for (ch = 0; ch < inlink->channels; ch++) {
168         s->fft_temp[ch] = av_calloc(s->window_size, sizeof(**s->fft_temp));
169         if (!s->fft_temp[ch])
170             return AVERROR(ENOMEM);
171     }
172
173     s->real = av_calloc(inlink->channels, sizeof(*s->real));
174     if (!s->real)
175         return AVERROR(ENOMEM);
176
177     s->imag = av_calloc(inlink->channels, sizeof(*s->imag));
178     if (!s->imag)
179         return AVERROR(ENOMEM);
180
181     args = av_strdup(s->real_str);
182     if (!args)
183         return AVERROR(ENOMEM);
184
185     for (ch = 0; ch < inlink->channels; ch++) {
186         char *arg = av_strtok(ch == 0 ? args : NULL, "|", &saveptr);
187
188         ret = av_expr_parse(&s->real[ch], arg ? arg : last_expr, var_names,
189                             NULL, NULL, func2_names, func2, 0, ctx);
190         if (ret < 0)
191             break;
192         if (arg)
193             last_expr = arg;
194         s->nb_exprs++;
195     }
196
197     av_free(args);
198
199     args = av_strdup(s->img_str ? s->img_str : s->real_str);
200     if (!args)
201         return AVERROR(ENOMEM);
202
203     for (ch = 0; ch < inlink->channels; ch++) {
204         char *arg = av_strtok(ch == 0 ? args : NULL, "|", &saveptr);
205
206         ret = av_expr_parse(&s->imag[ch], arg ? arg : last_expr, var_names,
207                             NULL, NULL, func2_names, func2, 0, ctx);
208         if (ret < 0)
209             break;
210         if (arg)
211             last_expr = arg;
212     }
213
214     av_free(args);
215
216     s->fifo = av_audio_fifo_alloc(inlink->format, inlink->channels, s->window_size);
217     if (!s->fifo)
218         return AVERROR(ENOMEM);
219
220     s->window_func_lut = av_realloc_f(s->window_func_lut, s->window_size,
221                                       sizeof(*s->window_func_lut));
222     if (!s->window_func_lut)
223         return AVERROR(ENOMEM);
224     generate_window_func(s->window_func_lut, s->window_size, s->win_func, &overlap);
225     if (s->overlap == 1)
226         s->overlap = overlap;
227
228     for (s->win_scale = 0, i = 0; i < s->window_size; i++) {
229         s->win_scale += s->window_func_lut[i] * s->window_func_lut[i];
230     }
231
232     s->hop_size = s->window_size * (1 - s->overlap);
233     if (s->hop_size <= 0)
234         return AVERROR(EINVAL);
235
236     s->buffer = ff_get_audio_buffer(inlink, s->window_size * 2);
237     if (!s->buffer)
238         return AVERROR(ENOMEM);
239
240     return ret;
241 }
242
243 static int filter_frame(AVFilterLink *inlink, AVFrame *frame)
244 {
245     AVFilterContext *ctx = inlink->dst;
246     AVFilterLink *outlink = ctx->outputs[0];
247     AFFTFiltContext *s = ctx->priv;
248     const int window_size = s->window_size;
249     const float f = 1. / s->win_scale;
250     double values[VAR_VARS_NB];
251     AVFrame *out, *in = NULL;
252     int ch, n, ret, i, j, k;
253     int start = s->start, end = s->end;
254
255     if (s->pts == AV_NOPTS_VALUE)
256         s->pts = frame->pts;
257
258     ret = av_audio_fifo_write(s->fifo, (void **)frame->extended_data, frame->nb_samples);
259     av_frame_free(&frame);
260     if (ret < 0)
261         return ret;
262
263     while (av_audio_fifo_size(s->fifo) >= window_size) {
264         if (!in) {
265             in = ff_get_audio_buffer(outlink, window_size);
266             if (!in)
267                 return AVERROR(ENOMEM);
268         }
269
270         ret = av_audio_fifo_peek(s->fifo, (void **)in->extended_data, window_size);
271         if (ret < 0)
272             break;
273
274         for (ch = 0; ch < inlink->channels; ch++) {
275             const float *src = (float *)in->extended_data[ch];
276             FFTComplex *fft_data = s->fft_data[ch];
277
278             for (n = 0; n < in->nb_samples; n++) {
279                 fft_data[n].re = src[n] * s->window_func_lut[n];
280                 fft_data[n].im = 0;
281             }
282
283             for (; n < window_size; n++) {
284                 fft_data[n].re = 0;
285                 fft_data[n].im = 0;
286             }
287         }
288
289         values[VAR_PTS]         = s->pts;
290         values[VAR_SAMPLE_RATE] = inlink->sample_rate;
291         values[VAR_NBBINS]      = window_size / 2;
292         values[VAR_CHANNELS]    = inlink->channels;
293
294         for (ch = 0; ch < inlink->channels; ch++) {
295             FFTComplex *fft_data = s->fft_data[ch];
296
297             av_fft_permute(s->fft, fft_data);
298             av_fft_calc(s->fft, fft_data);
299         }
300
301         for (ch = 0; ch < inlink->channels; ch++) {
302             FFTComplex *fft_data = s->fft_data[ch];
303             FFTComplex *fft_temp = s->fft_temp[ch];
304             float *buf = (float *)s->buffer->extended_data[ch];
305             int x;
306             values[VAR_CHANNEL] = ch;
307
308             for (n = 0; n <= window_size / 2; n++) {
309                 float fr, fi;
310
311                 values[VAR_BIN] = n;
312                 values[VAR_REAL] = fft_data[n].re;
313                 values[VAR_IMAG] = fft_data[n].im;
314
315                 fr = av_expr_eval(s->real[ch], values, s);
316                 fi = av_expr_eval(s->imag[ch], values, s);
317
318                 fft_temp[n].re = fr;
319                 fft_temp[n].im = fi;
320             }
321
322             for (n = window_size / 2 + 1, x = window_size / 2 - 1; n < window_size; n++, x--) {
323                 fft_temp[n].re =  fft_temp[x].re;
324                 fft_temp[n].im = -fft_temp[x].im;
325             }
326
327             av_fft_permute(s->ifft, fft_temp);
328             av_fft_calc(s->ifft, fft_temp);
329
330             start = s->start;
331             end = s->end;
332             k = end;
333             for (i = 0, j = start; j < k && i < window_size; i++, j++) {
334                 buf[j] += s->fft_temp[ch][i].re * f;
335             }
336
337             for (; i < window_size; i++, j++) {
338                 buf[j] = s->fft_temp[ch][i].re * f;
339             }
340
341             start += s->hop_size;
342             end = j;
343         }
344
345         s->start = start;
346         s->end = end;
347
348         if (start >= window_size) {
349             float *dst, *buf;
350
351             start -= window_size;
352             end   -= window_size;
353
354             s->start = start;
355             s->end = end;
356
357             out = ff_get_audio_buffer(outlink, window_size);
358             if (!out) {
359                 ret = AVERROR(ENOMEM);
360                 break;
361             }
362
363             out->pts = s->pts;
364             s->pts += window_size;
365
366             for (ch = 0; ch < inlink->channels; ch++) {
367                 dst = (float *)out->extended_data[ch];
368                 buf = (float *)s->buffer->extended_data[ch];
369
370                 for (n = 0; n < window_size; n++) {
371                     dst[n] = buf[n] * (1 - s->overlap);
372                 }
373                 memmove(buf, buf + window_size, window_size * 4);
374             }
375
376             ret = ff_filter_frame(outlink, out);
377             if (ret < 0)
378                 break;
379         }
380
381         av_audio_fifo_drain(s->fifo, s->hop_size);
382     }
383
384     av_frame_free(&in);
385     return ret < 0 ? ret : 0;
386 }
387
388 static int query_formats(AVFilterContext *ctx)
389 {
390     AVFilterFormats *formats;
391     AVFilterChannelLayouts *layouts;
392     static const enum AVSampleFormat sample_fmts[] = {
393         AV_SAMPLE_FMT_FLTP,
394         AV_SAMPLE_FMT_NONE
395     };
396     int ret;
397
398     layouts = ff_all_channel_counts();
399     if (!layouts)
400         return AVERROR(ENOMEM);
401     ret = ff_set_common_channel_layouts(ctx, layouts);
402     if (ret < 0)
403         return ret;
404
405     formats = ff_make_format_list(sample_fmts);
406     if (!formats)
407         return AVERROR(ENOMEM);
408     ret = ff_set_common_formats(ctx, formats);
409     if (ret < 0)
410         return ret;
411
412     formats = ff_all_samplerates();
413     if (!formats)
414         return AVERROR(ENOMEM);
415     return ff_set_common_samplerates(ctx, formats);
416 }
417
418 static av_cold void uninit(AVFilterContext *ctx)
419 {
420     AFFTFiltContext *s = ctx->priv;
421     int i;
422
423     av_fft_end(s->fft);
424     av_fft_end(s->ifft);
425
426     for (i = 0; i < s->nb_exprs; i++) {
427         if (s->fft_data)
428             av_freep(&s->fft_data[i]);
429         if (s->fft_temp)
430             av_freep(&s->fft_temp[i]);
431     }
432     av_freep(&s->fft_data);
433     av_freep(&s->fft_temp);
434
435     for (i = 0; i < s->nb_exprs; i++) {
436         av_expr_free(s->real[i]);
437         av_expr_free(s->imag[i]);
438     }
439
440     av_freep(&s->real);
441     av_freep(&s->imag);
442     av_frame_free(&s->buffer);
443     av_freep(&s->window_func_lut);
444
445     av_audio_fifo_free(s->fifo);
446 }
447
448 static const AVFilterPad inputs[] = {
449     {
450         .name         = "default",
451         .type         = AVMEDIA_TYPE_AUDIO,
452         .config_props = config_input,
453         .filter_frame = filter_frame,
454     },
455     { NULL }
456 };
457
458 static const AVFilterPad outputs[] = {
459     {
460         .name = "default",
461         .type = AVMEDIA_TYPE_AUDIO,
462     },
463     { NULL }
464 };
465
466 AVFilter ff_af_afftfilt = {
467     .name            = "afftfilt",
468     .description     = NULL_IF_CONFIG_SMALL("Apply arbitrary expressions to samples in frequency domain."),
469     .priv_size       = sizeof(AFFTFiltContext),
470     .priv_class      = &afftfilt_class,
471     .inputs          = inputs,
472     .outputs         = outputs,
473     .query_formats   = query_formats,
474     .uninit          = uninit,
475 };