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