avfilter/af_aiir: add support for alternative coefficients format
[ffmpeg.git] / libavfilter / af_aiir.c
1 /*
2  * Copyright (c) 2018 Paul B Mahol
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 #include <float.h>
22
23 #include "libavutil/avassert.h"
24 #include "libavutil/avstring.h"
25 #include "libavutil/opt.h"
26 #include "audio.h"
27 #include "avfilter.h"
28 #include "internal.h"
29
30 typedef struct AudioIIRContext {
31     const AVClass *class;
32     char *a_str, *b_str;
33     double dry_gain, wet_gain;
34     int format;
35
36     int *nb_a, *nb_b;
37     double **a, **b;
38     double **input, **output;
39     int clippings;
40     int channels;
41
42     void (*iir_frame)(AVFilterContext *ctx, AVFrame *in, AVFrame *out);
43 } AudioIIRContext;
44
45 static int query_formats(AVFilterContext *ctx)
46 {
47     AVFilterFormats *formats;
48     AVFilterChannelLayouts *layouts;
49     static const enum AVSampleFormat sample_fmts[] = {
50         AV_SAMPLE_FMT_DBLP,
51         AV_SAMPLE_FMT_FLTP,
52         AV_SAMPLE_FMT_S32P,
53         AV_SAMPLE_FMT_S16P,
54         AV_SAMPLE_FMT_NONE
55     };
56     int ret;
57
58     layouts = ff_all_channel_counts();
59     if (!layouts)
60         return AVERROR(ENOMEM);
61     ret = ff_set_common_channel_layouts(ctx, layouts);
62     if (ret < 0)
63         return ret;
64
65     formats = ff_make_format_list(sample_fmts);
66     if (!formats)
67         return AVERROR(ENOMEM);
68     ret = ff_set_common_formats(ctx, formats);
69     if (ret < 0)
70         return ret;
71
72     formats = ff_all_samplerates();
73     if (!formats)
74         return AVERROR(ENOMEM);
75     return ff_set_common_samplerates(ctx, formats);
76 }
77
78 #define IIR_FRAME(name, type, min, max, need_clipping)                  \
79 static void iir_frame_## name(AVFilterContext *ctx, AVFrame *in, AVFrame *out)  \
80 {                                                                       \
81     AudioIIRContext *s = ctx->priv;                                     \
82     const double ig = s->dry_gain;                                      \
83     const double og = s->wet_gain;                                      \
84     int ch, n;                                                          \
85                                                                         \
86     for (ch = 0; ch < out->channels; ch++) {                            \
87         const type *src = (const type *)in->extended_data[ch];          \
88         double *ic = (double *)s->input[ch];                            \
89         double *oc = (double *)s->output[ch];                           \
90         const int nb_a = s->nb_a[ch];                                   \
91         const int nb_b = s->nb_b[ch];                                   \
92         const double *a = s->a[ch];                                     \
93         const double *b = s->b[ch];                                     \
94         type *dst = (type *)out->extended_data[ch];                     \
95                                                                         \
96         for (n = 0; n < in->nb_samples; n++) {                          \
97             double sample = 0.;                                         \
98             int x;                                                      \
99                                                                         \
100             memmove(&ic[1], &ic[0], (nb_b - 1) * sizeof(*ic));          \
101             memmove(&oc[1], &oc[0], (nb_a - 1) * sizeof(*oc));          \
102             ic[0] = src[n] * ig;                                        \
103             for (x = 0; x < nb_b; x++)                                  \
104                 sample += b[x] * ic[x];                                 \
105                                                                         \
106             for (x = 1; x < nb_a; x++)                                  \
107                 sample -= a[x] * oc[x];                                 \
108                                                                         \
109             oc[0] = sample;                                             \
110             sample *= og;                                               \
111             if (need_clipping && sample < min) {                        \
112                 s->clippings++;                                         \
113                 dst[n] = min;                                           \
114             } else if (need_clipping && sample > max) {                 \
115                 s->clippings++;                                         \
116                 dst[n] = max;                                           \
117             } else {                                                    \
118                 dst[n] = sample;                                        \
119             }                                                           \
120         }                                                               \
121     }                                                                   \
122 }
123
124 IIR_FRAME(s16p, int16_t, INT16_MIN, INT16_MAX, 1)
125 IIR_FRAME(s32p, int32_t, INT32_MIN, INT32_MAX, 1)
126 IIR_FRAME(fltp, float,         -1.,        1., 0)
127 IIR_FRAME(dblp, double,        -1.,        1., 0)
128
129 static void count_coefficients(char *item_str, int *nb_items)
130 {
131     char *p;
132
133     if (!item_str)
134         return;
135
136     *nb_items = 1;
137     for (p = item_str; *p && *p != '|'; p++) {
138         if (*p == ' ')
139             (*nb_items)++;
140     }
141 }
142
143 static int read_tf_coefficients(AVFilterContext *ctx, char *item_str, int nb_items, double *dst)
144 {
145     char *p, *arg, *old_str, *saveptr = NULL;
146     int i;
147
148     p = old_str = av_strdup(item_str);
149     if (!p)
150         return AVERROR(ENOMEM);
151     for (i = 0; i < nb_items; i++) {
152         if (!(arg = av_strtok(p, " ", &saveptr)))
153             break;
154
155         p = NULL;
156         if (sscanf(arg, "%lf", &dst[i]) != 1) {
157             av_log(ctx, AV_LOG_ERROR, "Invalid coefficients supplied: %s\n", arg);
158             return AVERROR(EINVAL);
159         }
160     }
161
162     av_freep(&old_str);
163
164     return 0;
165 }
166
167 static int read_zp_coefficients(AVFilterContext *ctx, char *item_str, int nb_items, double *dst, int is_zeros)
168 {
169     char *p, *arg, *old_str, *saveptr = NULL;
170     int i;
171
172     p = old_str = av_strdup(item_str);
173     if (!p)
174         return AVERROR(ENOMEM);
175     for (i = 0; i < nb_items; i++) {
176         if (!(arg = av_strtok(p, " ", &saveptr)))
177             break;
178
179         p = NULL;
180         if (i == 0 && is_zeros) {
181             if (sscanf(arg, "%lf", &dst[i]) != 1) {
182                 av_log(ctx, AV_LOG_ERROR, "Invalid gain supplied: %s\n", arg);
183                 return AVERROR(EINVAL);
184             }
185         } else {
186             if (sscanf(arg, "%lf %lfi", &dst[i*2], &dst[i*2+1]) != 2) {
187                 av_log(ctx, AV_LOG_ERROR, "Invalid coefficients supplied: %s\n", arg);
188                 return AVERROR(EINVAL);
189             }
190         }
191     }
192
193     av_freep(&old_str);
194
195     return 0;
196 }
197
198 static int read_channels(AVFilterContext *ctx, int channels, uint8_t *item_str, int *nb, double **c, double **cache, int is_zeros)
199 {
200     AudioIIRContext *s = ctx->priv;
201     char *p, *arg, *old_str, *prev_arg = NULL, *saveptr = NULL;
202     int i, ret;
203
204     p = old_str = av_strdup(item_str);
205     if (!p)
206         return AVERROR(ENOMEM);
207     for (i = 0; i < channels; i++) {
208         if (!(arg = av_strtok(p, "|", &saveptr)))
209             arg = prev_arg;
210
211         if (!arg)
212             return AVERROR(EINVAL);
213
214         count_coefficients(arg, &nb[i]);
215
216         p = NULL;
217         cache[i] = av_calloc(nb[i] + 1, sizeof(double));
218         c[i] = av_calloc(nb[i] * (s->format + 1), sizeof(double));
219         if (!c[i] || !cache[i])
220             return AVERROR(ENOMEM);
221
222         if (s->format) {
223             ret = read_zp_coefficients(ctx, arg, nb[i], c[i], is_zeros);
224             if (is_zeros)
225                 nb[i]--;
226         } else {
227             ret = read_tf_coefficients(ctx, arg, nb[i], c[i]);
228         }
229         if (ret < 0)
230             return ret;
231         prev_arg = arg;
232     }
233
234     av_freep(&old_str);
235
236     return 0;
237 }
238
239 static void multiply(double wre, double wim, int npz, double *coeffs)
240 {
241     double nwre = -wre, nwim = -wim;
242     double cre, cim;
243     int i;
244
245     for (i = npz; i >= 1; i--) {
246         cre = coeffs[2 * i + 0];
247         cim = coeffs[2 * i + 1];
248
249         coeffs[2 * i + 0] = (nwre * cre - nwim * cim) + coeffs[2 * (i - 1) + 0];
250         coeffs[2 * i + 1] = (nwre * cim + nwim * cre) + coeffs[2 * (i - 1) + 1];
251     }
252
253     cre = coeffs[0];
254     cim = coeffs[1];
255     coeffs[0] = nwre * cre - nwim * cim;
256     coeffs[1] = nwre * cim + nwim * cre;
257 }
258
259 static int expand(AVFilterContext *ctx, double *pz, int nb, double *coeffs)
260 {
261     int i;
262
263     coeffs[0] = 1.0;
264     coeffs[1] = 0.0;
265
266     for (i = 0; i < nb; i++) {
267         coeffs[2 * (i + 1)    ] = 0.0;
268         coeffs[2 * (i + 1) + 1] = 0.0;
269     }
270
271     for (i = 0; i < nb; i++)
272         multiply(pz[2 * i], pz[2 * i + 1], nb, coeffs);
273
274     for (i = 0; i < nb + 1; i++) {
275         if (fabs(coeffs[2 * i + 1]) > DBL_EPSILON) {
276             av_log(ctx, AV_LOG_ERROR, "coeff: %lf of z^%d is not real; poles/zeros are not complex conjugates.\n",
277                    coeffs[2 * i + i], i);
278             return AVERROR(EINVAL);
279         }
280     }
281
282     return 0;
283 }
284
285 static int convert_zp2tf(AVFilterContext *ctx, int channels)
286 {
287     AudioIIRContext *s = ctx->priv;
288     int ch, i, j, ret;
289
290     for (ch = 0; ch < channels; ch++) {
291         double *topc, *botc, gain;
292
293         topc = av_calloc((s->nb_b[ch] + 1) * 2, sizeof(*topc));
294         botc = av_calloc((s->nb_a[ch] + 1) * 2, sizeof(*botc));
295         if (!topc || !botc)
296             return AVERROR(ENOMEM);
297
298         ret = expand(ctx, s->a[ch], s->nb_a[ch], botc);
299         if (ret < 0) {
300             av_free(topc);
301             av_free(botc);
302             return ret;
303         }
304
305         ret = expand(ctx, &s->b[ch][2], s->nb_b[ch], topc);
306         if (ret < 0) {
307             av_free(topc);
308             av_free(botc);
309             return ret;
310         }
311
312         gain = s->b[ch][0];
313         for (j = 0, i = s->nb_b[ch]; i >= 0; j++, i--) {
314             s->b[ch][j] = topc[2 * i] * gain;
315         }
316         s->nb_b[ch]++;
317
318         for (j = 0, i = s->nb_a[ch]; i >= 0; j++, i--) {
319             s->a[ch][j] = botc[2 * i];
320         }
321         s->nb_a[ch]++;
322
323         av_free(topc);
324         av_free(botc);
325     }
326
327     return 0;
328 }
329
330 static int config_output(AVFilterLink *outlink)
331 {
332     AVFilterContext *ctx = outlink->src;
333     AudioIIRContext *s = ctx->priv;
334     AVFilterLink *inlink = ctx->inputs[0];
335     int ch, ret, i;
336
337     s->channels = inlink->channels;
338     s->a = av_calloc(inlink->channels, sizeof(*s->a));
339     s->b = av_calloc(inlink->channels, sizeof(*s->b));
340     s->nb_a = av_calloc(inlink->channels, sizeof(*s->nb_a));
341     s->nb_b = av_calloc(inlink->channels, sizeof(*s->nb_b));
342     s->input = av_calloc(inlink->channels, sizeof(*s->input));
343     s->output = av_calloc(inlink->channels, sizeof(*s->output));
344     if (!s->a || !s->b || !s->nb_a || !s->nb_b || !s->input || !s->output)
345         return AVERROR(ENOMEM);
346
347     ret = read_channels(ctx, inlink->channels, s->a_str, s->nb_a, s->a, s->output, 0);
348     if (ret < 0)
349         return ret;
350
351     ret = read_channels(ctx, inlink->channels, s->b_str, s->nb_b, s->b, s->input, 1);
352     if (ret < 0)
353         return ret;
354
355     if (s->format) {
356         ret = convert_zp2tf(ctx, inlink->channels);
357         if (ret < 0)
358             return ret;
359     }
360
361     for (ch = 0; ch < inlink->channels; ch++) {
362         for (i = 1; i < s->nb_a[ch]; i++) {
363             s->a[ch][i] /= s->a[ch][0];
364         }
365
366         for (i = 0; i < s->nb_b[ch]; i++) {
367             s->b[ch][i] /= s->a[ch][0];
368         }
369     }
370
371     switch (inlink->format) {
372     case AV_SAMPLE_FMT_DBLP: s->iir_frame = iir_frame_dblp; break;
373     case AV_SAMPLE_FMT_FLTP: s->iir_frame = iir_frame_fltp; break;
374     case AV_SAMPLE_FMT_S32P: s->iir_frame = iir_frame_s32p; break;
375     case AV_SAMPLE_FMT_S16P: s->iir_frame = iir_frame_s16p; break;
376     }
377
378     return 0;
379 }
380
381 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
382 {
383     AVFilterContext *ctx = inlink->dst;
384     AudioIIRContext *s = ctx->priv;
385     AVFilterLink *outlink = ctx->outputs[0];
386     AVFrame *out;
387
388     if (av_frame_is_writable(in)) {
389         out = in;
390     } else {
391         out = ff_get_audio_buffer(outlink, in->nb_samples);
392         if (!out) {
393             av_frame_free(&in);
394             return AVERROR(ENOMEM);
395         }
396         av_frame_copy_props(out, in);
397     }
398
399     s->iir_frame(ctx, in, out);
400
401     if (s->clippings > 0)
402         av_log(ctx, AV_LOG_WARNING, "clipping %d times. Please reduce gain.\n", s->clippings);
403     s->clippings = 0;
404
405     if (in != out)
406         av_frame_free(&in);
407
408     return ff_filter_frame(outlink, out);
409 }
410
411 static av_cold int init(AVFilterContext *ctx)
412 {
413     AudioIIRContext *s = ctx->priv;
414
415     if (!s->a_str || !s->b_str) {
416         av_log(ctx, AV_LOG_ERROR, "Valid coefficients are mandatory.\n");
417         return AVERROR(EINVAL);
418     }
419
420     return 0;
421 }
422
423 static av_cold void uninit(AVFilterContext *ctx)
424 {
425     AudioIIRContext *s = ctx->priv;
426     int ch;
427
428     if (s->a) {
429         for (ch = 0; ch < s->channels; ch++) {
430             av_freep(&s->a[ch]);
431             av_freep(&s->output[ch]);
432         }
433     }
434     av_freep(&s->a);
435
436     if (s->b) {
437         for (ch = 0; ch < s->channels; ch++) {
438             av_freep(&s->b[ch]);
439             av_freep(&s->input[ch]);
440         }
441     }
442     av_freep(&s->b);
443
444     av_freep(&s->input);
445     av_freep(&s->output);
446
447     av_freep(&s->nb_a);
448     av_freep(&s->nb_b);
449 }
450
451 static const AVFilterPad inputs[] = {
452     {
453         .name         = "default",
454         .type         = AVMEDIA_TYPE_AUDIO,
455         .filter_frame = filter_frame,
456     },
457     { NULL }
458 };
459
460 static const AVFilterPad outputs[] = {
461     {
462         .name         = "default",
463         .type         = AVMEDIA_TYPE_AUDIO,
464         .config_props = config_output,
465     },
466     { NULL }
467 };
468
469 #define OFFSET(x) offsetof(AudioIIRContext, x)
470 #define AF AV_OPT_FLAG_AUDIO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
471
472 static const AVOption aiir_options[] = {
473     { "a", "set A/denominator/poles coefficients", OFFSET(a_str),    AV_OPT_TYPE_STRING, {.str="1 1"}, 0, 0, AF },
474     { "b", "set B/numerator/zeros coefficients",   OFFSET(b_str),    AV_OPT_TYPE_STRING, {.str="1 1"}, 0, 0, AF },
475     { "dry", "set dry gain",                       OFFSET(dry_gain), AV_OPT_TYPE_DOUBLE, {.dbl=1},     0, 1, AF },
476     { "wet", "set wet gain",                       OFFSET(wet_gain), AV_OPT_TYPE_DOUBLE, {.dbl=1},     0, 1, AF },
477     { "f", "set coefficients format",              OFFSET(format),   AV_OPT_TYPE_INT,    {.i64=0},     0, 1, AF, "format" },
478     { "tf", "transfer function",                   0,                AV_OPT_TYPE_CONST,  {.i64=0},     0, 0, AF, "format" },
479     { "zp", "Z-plane zeros/poles",                 0,                AV_OPT_TYPE_CONST,  {.i64=1},     0, 0, AF, "format" },
480     { NULL },
481 };
482
483 AVFILTER_DEFINE_CLASS(aiir);
484
485 AVFilter ff_af_aiir = {
486     .name          = "aiir",
487     .description   = NULL_IF_CONFIG_SMALL("Apply Infinite Impulse Response filter with supplied coefficients."),
488     .priv_size     = sizeof(AudioIIRContext),
489     .init          = init,
490     .uninit        = uninit,
491     .query_formats = query_formats,
492     .inputs        = inputs,
493     .outputs       = outputs,
494     .priv_class    = &aiir_class,
495 };