libavfilter/af_afir.c

   1 /*
   2  * Copyright (c) 2017 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 /**
  22  * @file
  23  * An arbitrary audio FIR filter
  24  */
  25
  26 #include "libavutil/audio_fifo.h"
  27 #include "libavutil/common.h"
  28 #include "libavutil/float_dsp.h"
  29 #include "libavutil/opt.h"
  30 #include "libavcodec/avfft.h"
  31
  32 #include "audio.h"
  33 #include "avfilter.h"
  34 #include "formats.h"
  35 #include "internal.h"
  36 #include "af_afir.h"
  37
  38 static void fcmul_add_c(float *sum, const float *t, const float *c, ptrdiff_t len)
  39 {
  40     int n;
  41
  42     for (n = 0; n < len; n++) {
  43         const float cre = c[2 * n    ];
  44         const float cim = c[2 * n + 1];
  45         const float tre = t[2 * n    ];
  46         const float tim = t[2 * n + 1];
  47
  48         sum[2 * n    ] += tre * cre - tim * cim;
  49         sum[2 * n + 1] += tre * cim + tim * cre;
  50     }
  51
  52     sum[2 * n] += t[2 * n] * c[2 * n];
  53 }
  54
  55 static int fir_channel(AVFilterContext *ctx, void *arg, int ch, int nb_jobs)
  56 {
  57     AudioFIRContext *s = ctx->priv;
  58     const float *src = (const float *)s->in[0]->extended_data[ch];
  59     int index1 = (s->index + 1) % 3;
  60     int index2 = (s->index + 2) % 3;
  61     float *sum = s->sum[ch];
  62     AVFrame *out = arg;
  63     float *block;
  64     float *dst;
  65     int n, i, j;
  66
  67     memset(sum, 0, sizeof(*sum) * s->fft_length);
  68     block = s->block[ch] + s->part_index * s->block_size;
  69     memset(block, 0, sizeof(*block) * s->fft_length);
  70
  71     s->fdsp->vector_fmul_scalar(block + s->part_size, src, s->dry_gain, FFALIGN(s->nb_samples, 4));
  72     emms_c();
  73
  74     av_rdft_calc(s->rdft[ch], block);
  75     block[2 * s->part_size] = block[1];
  76     block[1] = 0;
  77
  78     j = s->part_index;
  79
  80     for (i = 0; i < s->nb_partitions; i++) {
  81         const int coffset = i * s->coeff_size;
  82         const FFTComplex *coeff = s->coeff[ch * !s->one2many] + coffset;
  83
  84         block = s->block[ch] + j * s->block_size;
  85         s->fcmul_add(sum, block, (const float *)coeff, s->part_size);
  86
  87         if (j == 0)
  88             j = s->nb_partitions;
  89         j--;
  90     }
  91
  92     sum[1] = sum[2 * s->part_size];
  93     av_rdft_calc(s->irdft[ch], sum);
  94
  95     dst = (float *)s->buffer->extended_data[ch] + index1 * s->part_size;
  96     for (n = 0; n < s->part_size; n++) {
  97         dst[n] += sum[n];
  98     }
  99
 100     dst = (float *)s->buffer->extended_data[ch] + index2 * s->part_size;
 101
 102     memcpy(dst, sum + s->part_size, s->part_size * sizeof(*dst));
 103
 104     dst = (float *)s->buffer->extended_data[ch] + s->index * s->part_size;
 105
 106     if (out) {
 107         float *ptr = (float *)out->extended_data[ch];
 108         s->fdsp->vector_fmul_scalar(ptr, dst, s->gain * s->wet_gain, FFALIGN(out->nb_samples, 4));
 109         emms_c();
 110     }
 111
 112     return 0;
 113 }
 114
 115 static int fir_frame(AudioFIRContext *s, AVFilterLink *outlink)
 116 {
 117     AVFilterContext *ctx = outlink->src;
 118     AVFrame *out = NULL;
 119     int ret;
 120
 121     s->nb_samples = FFMIN(s->part_size, av_audio_fifo_size(s->fifo[0]));
 122
 123     if (!s->want_skip) {
 124         out = ff_get_audio_buffer(outlink, s->nb_samples);
 125         if (!out)
 126             return AVERROR(ENOMEM);
 127     }
 128
 129     s->in[0] = ff_get_audio_buffer(ctx->inputs[0], s->nb_samples);
 130     if (!s->in[0]) {
 131         av_frame_free(&out);
 132         return AVERROR(ENOMEM);
 133     }
 134
 135     av_audio_fifo_peek(s->fifo[0], (void **)s->in[0]->extended_data, s->nb_samples);
 136
 137     ctx->internal->execute(ctx, fir_channel, out, NULL, outlink->channels);
 138
 139     s->part_index = (s->part_index + 1) % s->nb_partitions;
 140
 141     av_audio_fifo_drain(s->fifo[0], s->nb_samples);
 142
 143     if (!s->want_skip) {
 144         out->pts = s->pts;
 145         if (s->pts != AV_NOPTS_VALUE)
 146             s->pts += av_rescale_q(out->nb_samples, (AVRational){1, outlink->sample_rate}, outlink->time_base);
 147     }
 148
 149     s->index++;
 150     if (s->index == 3)
 151         s->index = 0;
 152
 153     av_frame_free(&s->in[0]);
 154
 155     if (s->want_skip == 1) {
 156         s->want_skip = 0;
 157         ret = 0;
 158     } else {
 159         ret = ff_filter_frame(outlink, out);
 160     }
 161
 162     return ret;
 163 }
 164
 165 static int convert_coeffs(AVFilterContext *ctx)
 166 {
 167     AudioFIRContext *s = ctx->priv;
 168     int i, ch, n, N;
 169     float power = 0;
 170
 171     s->nb_taps = av_audio_fifo_size(s->fifo[1]);
 172     if (s->nb_taps <= 0)
 173         return AVERROR(EINVAL);
 174
 175     for (n = 4; (1 << n) < s->nb_taps; n++);
 176     N = FFMIN(n, 16);
 177     s->ir_length = 1 << n;
 178     s->fft_length = (1 << (N + 1)) + 1;
 179     s->part_size = 1 << (N - 1);
 180     s->block_size = FFALIGN(s->fft_length, 32);
 181     s->coeff_size = FFALIGN(s->part_size + 1, 32);
 182     s->nb_partitions = (s->nb_taps + s->part_size - 1) / s->part_size;
 183     s->nb_coeffs = s->ir_length + s->nb_partitions;
 184
 185     for (ch = 0; ch < ctx->inputs[0]->channels; ch++) {
 186         s->sum[ch] = av_calloc(s->fft_length, sizeof(**s->sum));
 187         if (!s->sum[ch])
 188             return AVERROR(ENOMEM);
 189     }
 190
 191     for (ch = 0; ch < ctx->inputs[1]->channels; ch++) {
 192         s->coeff[ch] = av_calloc(s->nb_partitions * s->coeff_size, sizeof(**s->coeff));
 193         if (!s->coeff[ch])
 194             return AVERROR(ENOMEM);
 195     }
 196
 197     for (ch = 0; ch < ctx->inputs[0]->channels; ch++) {
 198         s->block[ch] = av_calloc(s->nb_partitions * s->block_size, sizeof(**s->block));
 199         if (!s->block[ch])
 200             return AVERROR(ENOMEM);
 201     }
 202
 203     for (ch = 0; ch < ctx->inputs[0]->channels; ch++) {
 204         s->rdft[ch]  = av_rdft_init(N, DFT_R2C);
 205         s->irdft[ch] = av_rdft_init(N, IDFT_C2R);
 206         if (!s->rdft[ch] || !s->irdft[ch])
 207             return AVERROR(ENOMEM);
 208     }
 209
 210     s->in[1] = ff_get_audio_buffer(ctx->inputs[1], s->nb_taps);
 211     if (!s->in[1])
 212         return AVERROR(ENOMEM);
 213
 214     s->buffer = ff_get_audio_buffer(ctx->inputs[0], s->part_size * 3);
 215     if (!s->buffer)
 216         return AVERROR(ENOMEM);
 217
 218     av_audio_fifo_read(s->fifo[1], (void **)s->in[1]->extended_data, s->nb_taps);
 219
 220     for (ch = 0; ch < ctx->inputs[1]->channels; ch++) {
 221         float *time = (float *)s->in[1]->extended_data[!s->one2many * ch];
 222         float *block = s->block[ch];
 223         FFTComplex *coeff = s->coeff[ch];
 224
 225         power += s->fdsp->scalarproduct_float(time, time, s->nb_taps);
 226
 227         for (i = FFMAX(1, s->length * s->nb_taps); i < s->nb_taps; i++)
 228             time[i] = 0;
 229
 230         for (i = 0; i < s->nb_partitions; i++) {
 231             const float scale = 1.f / s->part_size;
 232             const int toffset = i * s->part_size;
 233             const int coffset = i * s->coeff_size;
 234             const int boffset = s->part_size;
 235             const int remaining = s->nb_taps - (i * s->part_size);
 236             const int size = remaining >= s->part_size ? s->part_size : remaining;
 237
 238             memset(block, 0, sizeof(*block) * s->fft_length);
 239             memcpy(block + boffset, time + toffset, size * sizeof(*block));
 240
 241             av_rdft_calc(s->rdft[0], block);
 242
 243             coeff[coffset].re = block[0] * scale;
 244             coeff[coffset].im = 0;
 245             for (n = 1; n < s->part_size; n++) {
 246                 coeff[coffset + n].re = block[2 * n] * scale;
 247                 coeff[coffset + n].im = block[2 * n + 1] * scale;
 248             }
 249             coeff[coffset + s->part_size].re = block[1] * scale;
 250             coeff[coffset + s->part_size].im = 0;
 251         }
 252     }
 253
 254     av_frame_free(&s->in[1]);
 255     s->gain = s->again ? 1.f / sqrtf(power / ctx->inputs[1]->channels) : 1.f;
 256     av_log(ctx, AV_LOG_DEBUG, "nb_taps: %d\n", s->nb_taps);
 257     av_log(ctx, AV_LOG_DEBUG, "nb_partitions: %d\n", s->nb_partitions);
 258     av_log(ctx, AV_LOG_DEBUG, "partition size: %d\n", s->part_size);
 259     av_log(ctx, AV_LOG_DEBUG, "ir_length: %d\n", s->ir_length);
 260
 261     s->have_coeffs = 1;
 262
 263     return 0;
 264 }
 265
 266 static int read_ir(AVFilterLink *link, AVFrame *frame)
 267 {
 268     AVFilterContext *ctx = link->dst;
 269     AudioFIRContext *s = ctx->priv;
 270     int nb_taps, max_nb_taps;
 271
 272     av_audio_fifo_write(s->fifo[1], (void **)frame->extended_data,
 273                         frame->nb_samples);
 274     av_frame_free(&frame);
 275
 276     nb_taps = av_audio_fifo_size(s->fifo[1]);
 277     max_nb_taps = MAX_IR_DURATION * ctx->outputs[0]->sample_rate;
 278     if (nb_taps > max_nb_taps) {
 279         av_log(ctx, AV_LOG_ERROR, "Too big number of coefficients: %d > %d.\n", nb_taps, max_nb_taps);
 280         return AVERROR(EINVAL);
 281     }
 282
 283     return 0;
 284 }
 285
 286 static int filter_frame(AVFilterLink *link, AVFrame *frame)
 287 {
 288     AVFilterContext *ctx = link->dst;
 289     AudioFIRContext *s = ctx->priv;
 290     AVFilterLink *outlink = ctx->outputs[0];
 291     int ret = 0;
 292
 293     av_audio_fifo_write(s->fifo[0], (void **)frame->extended_data,
 294                         frame->nb_samples);
 295     if (s->pts == AV_NOPTS_VALUE)
 296         s->pts = frame->pts;
 297
 298     av_frame_free(&frame);
 299
 300     if (!s->have_coeffs && s->eof_coeffs) {
 301         ret = convert_coeffs(ctx);
 302         if (ret < 0)
 303             return ret;
 304     }
 305
 306     if (s->have_coeffs) {
 307         while (av_audio_fifo_size(s->fifo[0]) >= s->part_size) {
 308             ret = fir_frame(s, outlink);
 309             if (ret < 0)
 310                 break;
 311         }
 312     }
 313     return ret;
 314 }
 315
 316 static int request_frame(AVFilterLink *outlink)
 317 {
 318     AVFilterContext *ctx = outlink->src;
 319     AudioFIRContext *s = ctx->priv;
 320     int ret;
 321
 322     if (!s->eof_coeffs) {
 323         ret = ff_request_frame(ctx->inputs[1]);
 324         if (ret == AVERROR_EOF) {
 325             s->eof_coeffs = 1;
 326             ret = 0;
 327         }
 328         return ret;
 329     }
 330     ret = ff_request_frame(ctx->inputs[0]);
 331     if (ret == AVERROR_EOF && s->have_coeffs) {
 332         if (s->need_padding) {
 333             AVFrame *silence = ff_get_audio_buffer(outlink, s->part_size);
 334
 335             if (!silence)
 336                 return AVERROR(ENOMEM);
 337             av_audio_fifo_write(s->fifo[0], (void **)silence->extended_data,
 338                         silence->nb_samples);
 339             av_frame_free(&silence);
 340             s->need_padding = 0;
 341         }
 342
 343         while (av_audio_fifo_size(s->fifo[0]) > 0) {
 344             ret = fir_frame(s, outlink);
 345             if (ret < 0)
 346                 return ret;
 347         }
 348         ret = AVERROR_EOF;
 349     }
 350     return ret;
 351 }
 352
 353 static int query_formats(AVFilterContext *ctx)
 354 {
 355     AVFilterFormats *formats;
 356     AVFilterChannelLayouts *layouts;
 357     static const enum AVSampleFormat sample_fmts[] = {
 358         AV_SAMPLE_FMT_FLTP,
 359         AV_SAMPLE_FMT_NONE
 360     };
 361     int ret, i;
 362
 363     layouts = ff_all_channel_counts();
 364     if ((ret = ff_channel_layouts_ref(layouts, &ctx->outputs[0]->in_channel_layouts)) < 0)
 365         return ret;
 366
 367     for (i = 0; i < 2; i++) {
 368         layouts = ff_all_channel_counts();
 369         if ((ret = ff_channel_layouts_ref(layouts, &ctx->inputs[i]->out_channel_layouts)) < 0)
 370             return ret;
 371     }
 372
 373     formats = ff_make_format_list(sample_fmts);
 374     if ((ret = ff_set_common_formats(ctx, formats)) < 0)
 375         return ret;
 376
 377     formats = ff_all_samplerates();
 378     return ff_set_common_samplerates(ctx, formats);
 379 }
 380
 381 static int config_output(AVFilterLink *outlink)
 382 {
 383     AVFilterContext *ctx = outlink->src;
 384     AudioFIRContext *s = ctx->priv;
 385
 386     if (ctx->inputs[0]->channels != ctx->inputs[1]->channels &&
 387         ctx->inputs[1]->channels != 1) {
 388         av_log(ctx, AV_LOG_ERROR,
 389                "Second input must have same number of channels as first input or "
 390                "exactly 1 channel.\n");
 391         return AVERROR(EINVAL);
 392     }
 393
 394     s->one2many = ctx->inputs[1]->channels == 1;
 395     outlink->sample_rate = ctx->inputs[0]->sample_rate;
 396     outlink->time_base   = ctx->inputs[0]->time_base;
 397     outlink->channel_layout = ctx->inputs[0]->channel_layout;
 398     outlink->channels = ctx->inputs[0]->channels;
 399
 400     s->fifo[0] = av_audio_fifo_alloc(ctx->inputs[0]->format, ctx->inputs[0]->channels, 1024);
 401     s->fifo[1] = av_audio_fifo_alloc(ctx->inputs[1]->format, ctx->inputs[1]->channels, 1024);
 402     if (!s->fifo[0] || !s->fifo[1])
 403         return AVERROR(ENOMEM);
 404
 405     s->sum = av_calloc(outlink->channels, sizeof(*s->sum));
 406     s->coeff = av_calloc(ctx->inputs[1]->channels, sizeof(*s->coeff));
 407     s->block = av_calloc(ctx->inputs[0]->channels, sizeof(*s->block));
 408     s->rdft = av_calloc(outlink->channels, sizeof(*s->rdft));
 409     s->irdft = av_calloc(outlink->channels, sizeof(*s->irdft));
 410     if (!s->sum || !s->coeff || !s->block || !s->rdft || !s->irdft)
 411         return AVERROR(ENOMEM);
 412
 413     s->nb_channels = outlink->channels;
 414     s->nb_coef_channels = ctx->inputs[1]->channels;
 415     s->want_skip = 1;
 416     s->need_padding = 1;
 417     s->pts = AV_NOPTS_VALUE;
 418
 419     return 0;
 420 }
 421
 422 static av_cold void uninit(AVFilterContext *ctx)
 423 {
 424     AudioFIRContext *s = ctx->priv;
 425     int ch;
 426
 427     if (s->sum) {
 428         for (ch = 0; ch < s->nb_channels; ch++) {
 429             av_freep(&s->sum[ch]);
 430         }
 431     }
 432     av_freep(&s->sum);
 433
 434     if (s->coeff) {
 435         for (ch = 0; ch < s->nb_coef_channels; ch++) {
 436             av_freep(&s->coeff[ch]);
 437         }
 438     }
 439     av_freep(&s->coeff);
 440
 441     if (s->block) {
 442         for (ch = 0; ch < s->nb_channels; ch++) {
 443             av_freep(&s->block[ch]);
 444         }
 445     }
 446     av_freep(&s->block);
 447
 448     if (s->rdft) {
 449         for (ch = 0; ch < s->nb_channels; ch++) {
 450             av_rdft_end(s->rdft[ch]);
 451         }
 452     }
 453     av_freep(&s->rdft);
 454
 455     if (s->irdft) {
 456         for (ch = 0; ch < s->nb_channels; ch++) {
 457             av_rdft_end(s->irdft[ch]);
 458         }
 459     }
 460     av_freep(&s->irdft);
 461
 462     av_frame_free(&s->in[0]);
 463     av_frame_free(&s->in[1]);
 464     av_frame_free(&s->buffer);
 465
 466     av_audio_fifo_free(s->fifo[0]);
 467     av_audio_fifo_free(s->fifo[1]);
 468
 469     av_freep(&s->fdsp);
 470 }
 471
 472 static av_cold int init(AVFilterContext *ctx)
 473 {
 474     AudioFIRContext *s = ctx->priv;
 475
 476     s->fcmul_add = fcmul_add_c;
 477
 478     s->fdsp = avpriv_float_dsp_alloc(0);
 479     if (!s->fdsp)
 480         return AVERROR(ENOMEM);
 481
 482     if (ARCH_X86)
 483         ff_afir_init_x86(s);
 484
 485     return 0;
 486 }
 487
 488 static const AVFilterPad afir_inputs[] = {
 489     {
 490         .name           = "main",
 491         .type           = AVMEDIA_TYPE_AUDIO,
 492         .filter_frame   = filter_frame,
 493     },{
 494         .name           = "ir",
 495         .type           = AVMEDIA_TYPE_AUDIO,
 496         .filter_frame   = read_ir,
 497     },
 498     { NULL }
 499 };
 500
 501 static const AVFilterPad afir_outputs[] = {
 502     {
 503         .name          = "default",
 504         .type          = AVMEDIA_TYPE_AUDIO,
 505         .config_props  = config_output,
 506         .request_frame = request_frame,
 507     },
 508     { NULL }
 509 };
 510
 511 #define AF AV_OPT_FLAG_AUDIO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
 512 #define OFFSET(x) offsetof(AudioFIRContext, x)
 513
 514 static const AVOption afir_options[] = {
 515     { "dry",    "set dry gain",     OFFSET(dry_gain), AV_OPT_TYPE_FLOAT, {.dbl=1}, 0, 1, AF },
 516     { "wet",    "set wet gain",     OFFSET(wet_gain), AV_OPT_TYPE_FLOAT, {.dbl=1}, 0, 1, AF },
 517     { "length", "set IR length",    OFFSET(length),   AV_OPT_TYPE_FLOAT, {.dbl=1}, 0, 1, AF },
 518     { "again",  "enable auto gain", OFFSET(again),    AV_OPT_TYPE_BOOL,  {.i64=1}, 0, 1, AF },
 519     { NULL }
 520 };
 521
 522 AVFILTER_DEFINE_CLASS(afir);
 523
 524 AVFilter ff_af_afir = {
 525     .name          = "afir",
 526     .description   = NULL_IF_CONFIG_SMALL("Apply Finite Impulse Response filter with supplied coefficients in 2nd stream."),
 527     .priv_size     = sizeof(AudioFIRContext),
 528     .priv_class    = &afir_class,
 529     .query_formats = query_formats,
 530     .init          = init,
 531     .uninit        = uninit,
 532     .inputs        = afir_inputs,
 533     .outputs       = afir_outputs,
 534     .flags         = AVFILTER_FLAG_SLICE_THREADS,
 535 };