git.sesse.net Git - ffmpeg/blob - 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->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
 170     s->nb_taps = av_audio_fifo_size(s->fifo[1]);
 171     if (s->nb_taps <= 0)
 172         return AVERROR(EINVAL);
 173
 174     for (n = 4; (1 << n) < s->nb_taps; n++);
 175     N = FFMIN(n, 16);
 176     s->ir_length = 1 << n;
 177     s->fft_length = (1 << (N + 1)) + 1;
 178     s->part_size = 1 << (N - 1);
 179     s->block_size = FFALIGN(s->fft_length, 32);
 180     s->coeff_size = FFALIGN(s->part_size + 1, 32);
 181     s->nb_partitions = (s->nb_taps + s->part_size - 1) / s->part_size;
 182     s->nb_coeffs = s->ir_length + s->nb_partitions;
 183
 184     for (ch = 0; ch < ctx->inputs[0]->channels; ch++) {
 185         s->sum[ch] = av_calloc(s->fft_length, sizeof(**s->sum));
 186         if (!s->sum[ch])
 187             return AVERROR(ENOMEM);
 188     }
 189
 190     for (ch = 0; ch < ctx->inputs[1]->channels; ch++) {
 191         s->coeff[ch] = av_calloc(s->nb_partitions * s->coeff_size, sizeof(**s->coeff));
 192         if (!s->coeff[ch])
 193             return AVERROR(ENOMEM);
 194     }
 195
 196     for (ch = 0; ch < ctx->inputs[0]->channels; ch++) {
 197         s->block[ch] = av_calloc(s->nb_partitions * s->block_size, sizeof(**s->block));
 198         if (!s->block[ch])
 199             return AVERROR(ENOMEM);
 200     }
 201
 202     for (ch = 0; ch < ctx->inputs[0]->channels; ch++) {
 203         s->rdft[ch]  = av_rdft_init(N, DFT_R2C);
 204         s->irdft[ch] = av_rdft_init(N, IDFT_C2R);
 205         if (!s->rdft[ch] || !s->irdft[ch])
 206             return AVERROR(ENOMEM);
 207     }
 208
 209     s->in[1] = ff_get_audio_buffer(ctx->inputs[1], s->nb_taps);
 210     if (!s->in[1])
 211         return AVERROR(ENOMEM);
 212
 213     s->buffer = ff_get_audio_buffer(ctx->inputs[0], s->part_size * 3);
 214     if (!s->buffer)
 215         return AVERROR(ENOMEM);
 216
 217     av_audio_fifo_read(s->fifo[1], (void **)s->in[1]->extended_data, s->nb_taps);
 218
 219     if (s->again) {
 220         float power = 0;
 221
 222         for (ch = 0; ch < ctx->inputs[1]->channels; ch++) {
 223             float *time = (float *)s->in[1]->extended_data[!s->one2many * ch];
 224
 225             for (i = 0; i < s->nb_taps; i++)
 226                 power += FFABS(time[i]);
 227         }
 228
 229         s->gain = sqrtf(1.f / (ctx->inputs[1]->channels * power)) / (sqrtf(ctx->inputs[1]->channels));
 230         for (ch = 0; ch < ctx->inputs[1]->channels; ch++) {
 231             float *time = (float *)s->in[1]->extended_data[!s->one2many * ch];
 232
 233             s->fdsp->vector_fmul_scalar(time, time, s->gain, FFALIGN(s->nb_taps, 4));
 234         }
 235     }
 236
 237     for (ch = 0; ch < ctx->inputs[1]->channels; ch++) {
 238         float *time = (float *)s->in[1]->extended_data[!s->one2many * ch];
 239         float *block = s->block[ch];
 240         FFTComplex *coeff = s->coeff[ch];
 241
 242         for (i = FFMAX(1, s->length * s->nb_taps); i < s->nb_taps; i++)
 243             time[i] = 0;
 244
 245         for (i = 0; i < s->nb_partitions; i++) {
 246             const float scale = 1.f / s->part_size;
 247             const int toffset = i * s->part_size;
 248             const int coffset = i * s->coeff_size;
 249             const int boffset = s->part_size;
 250             const int remaining = s->nb_taps - (i * s->part_size);
 251             const int size = remaining >= s->part_size ? s->part_size : remaining;
 252
 253             memset(block, 0, sizeof(*block) * s->fft_length);
 254             memcpy(block + boffset, time + toffset, size * sizeof(*block));
 255
 256             av_rdft_calc(s->rdft[0], block);
 257
 258             coeff[coffset].re = block[0] * scale;
 259             coeff[coffset].im = 0;
 260             for (n = 1; n < s->part_size; n++) {
 261                 coeff[coffset + n].re = block[2 * n] * scale;
 262                 coeff[coffset + n].im = block[2 * n + 1] * scale;
 263             }
 264             coeff[coffset + s->part_size].re = block[1] * scale;
 265             coeff[coffset + s->part_size].im = 0;
 266         }
 267     }
 268
 269     av_frame_free(&s->in[1]);
 270     av_log(ctx, AV_LOG_DEBUG, "nb_taps: %d\n", s->nb_taps);
 271     av_log(ctx, AV_LOG_DEBUG, "nb_partitions: %d\n", s->nb_partitions);
 272     av_log(ctx, AV_LOG_DEBUG, "partition size: %d\n", s->part_size);
 273     av_log(ctx, AV_LOG_DEBUG, "ir_length: %d\n", s->ir_length);
 274
 275     s->have_coeffs = 1;
 276
 277     return 0;
 278 }
 279
 280 static int read_ir(AVFilterLink *link, AVFrame *frame)
 281 {
 282     AVFilterContext *ctx = link->dst;
 283     AudioFIRContext *s = ctx->priv;
 284     int nb_taps, max_nb_taps, ret;
 285
 286     ret = av_audio_fifo_write(s->fifo[1], (void **)frame->extended_data,
 287                              frame->nb_samples);
 288     av_frame_free(&frame);
 289     if (ret < 0)
 290         return ret;
 291
 292     nb_taps = av_audio_fifo_size(s->fifo[1]);
 293     max_nb_taps = s->max_ir_len * ctx->outputs[0]->sample_rate;
 294     if (nb_taps > max_nb_taps) {
 295         av_log(ctx, AV_LOG_ERROR, "Too big number of coefficients: %d > %d.\n", nb_taps, max_nb_taps);
 296         return AVERROR(EINVAL);
 297     }
 298
 299     return 0;
 300 }
 301
 302 static int filter_frame(AVFilterLink *link, AVFrame *frame)
 303 {
 304     AVFilterContext *ctx = link->dst;
 305     AudioFIRContext *s = ctx->priv;
 306     AVFilterLink *outlink = ctx->outputs[0];
 307     int ret;
 308
 309     ret = av_audio_fifo_write(s->fifo[0], (void **)frame->extended_data,
 310                               frame->nb_samples);
 311     if (ret > 0 && s->pts == AV_NOPTS_VALUE)
 312         s->pts = frame->pts;
 313
 314     av_frame_free(&frame);
 315
 316     if (ret < 0)
 317         return ret;
 318
 319     if (!s->have_coeffs && s->eof_coeffs) {
 320         ret = convert_coeffs(ctx);
 321         if (ret < 0)
 322             return ret;
 323     }
 324
 325     if (s->have_coeffs) {
 326         while (av_audio_fifo_size(s->fifo[0]) >= s->part_size) {
 327             ret = fir_frame(s, outlink);
 328             if (ret < 0)
 329                 return ret;
 330         }
 331     }
 332     return 0;
 333 }
 334
 335 static int request_frame(AVFilterLink *outlink)
 336 {
 337     AVFilterContext *ctx = outlink->src;
 338     AudioFIRContext *s = ctx->priv;
 339     int ret;
 340
 341     if (!s->eof_coeffs) {
 342         ret = ff_request_frame(ctx->inputs[1]);
 343         if (ret == AVERROR_EOF) {
 344             s->eof_coeffs = 1;
 345             ret = 0;
 346         }
 347         return ret;
 348     }
 349     ret = ff_request_frame(ctx->inputs[0]);
 350     if (ret == AVERROR_EOF && s->have_coeffs) {
 351         if (s->need_padding) {
 352             AVFrame *silence = ff_get_audio_buffer(outlink, s->part_size);
 353
 354             if (!silence)
 355                 return AVERROR(ENOMEM);
 356             ret = av_audio_fifo_write(s->fifo[0], (void **)silence->extended_data,
 357                                       silence->nb_samples);
 358             av_frame_free(&silence);
 359             if (ret < 0)
 360                 return ret;
 361             s->need_padding = 0;
 362         }
 363
 364         while (av_audio_fifo_size(s->fifo[0]) > 0) {
 365             ret = fir_frame(s, outlink);
 366             if (ret < 0)
 367                 return ret;
 368         }
 369         ret = AVERROR_EOF;
 370     }
 371     return ret;
 372 }
 373
 374 static int query_formats(AVFilterContext *ctx)
 375 {
 376     AVFilterFormats *formats;
 377     AVFilterChannelLayouts *layouts;
 378     static const enum AVSampleFormat sample_fmts[] = {
 379         AV_SAMPLE_FMT_FLTP,
 380         AV_SAMPLE_FMT_NONE
 381     };
 382     int ret, i;
 383
 384     layouts = ff_all_channel_counts();
 385     if ((ret = ff_channel_layouts_ref(layouts, &ctx->outputs[0]->in_channel_layouts)) < 0)
 386         return ret;
 387
 388     for (i = 0; i < 2; i++) {
 389         layouts = ff_all_channel_counts();
 390         if ((ret = ff_channel_layouts_ref(layouts, &ctx->inputs[i]->out_channel_layouts)) < 0)
 391             return ret;
 392     }
 393
 394     formats = ff_make_format_list(sample_fmts);
 395     if ((ret = ff_set_common_formats(ctx, formats)) < 0)
 396         return ret;
 397
 398     formats = ff_all_samplerates();
 399     return ff_set_common_samplerates(ctx, formats);
 400 }
 401
 402 static int config_output(AVFilterLink *outlink)
 403 {
 404     AVFilterContext *ctx = outlink->src;
 405     AudioFIRContext *s = ctx->priv;
 406
 407     if (ctx->inputs[0]->channels != ctx->inputs[1]->channels &&
 408         ctx->inputs[1]->channels != 1) {
 409         av_log(ctx, AV_LOG_ERROR,
 410                "Second input must have same number of channels as first input or "
 411                "exactly 1 channel.\n");
 412         return AVERROR(EINVAL);
 413     }
 414
 415     s->one2many = ctx->inputs[1]->channels == 1;
 416     outlink->sample_rate = ctx->inputs[0]->sample_rate;
 417     outlink->time_base   = ctx->inputs[0]->time_base;
 418     outlink->channel_layout = ctx->inputs[0]->channel_layout;
 419     outlink->channels = ctx->inputs[0]->channels;
 420
 421     s->fifo[0] = av_audio_fifo_alloc(ctx->inputs[0]->format, ctx->inputs[0]->channels, 1024);
 422     s->fifo[1] = av_audio_fifo_alloc(ctx->inputs[1]->format, ctx->inputs[1]->channels, 1024);
 423     if (!s->fifo[0] || !s->fifo[1])
 424         return AVERROR(ENOMEM);
 425
 426     s->sum = av_calloc(outlink->channels, sizeof(*s->sum));
 427     s->coeff = av_calloc(ctx->inputs[1]->channels, sizeof(*s->coeff));
 428     s->block = av_calloc(ctx->inputs[0]->channels, sizeof(*s->block));
 429     s->rdft = av_calloc(outlink->channels, sizeof(*s->rdft));
 430     s->irdft = av_calloc(outlink->channels, sizeof(*s->irdft));
 431     if (!s->sum || !s->coeff || !s->block || !s->rdft || !s->irdft)
 432         return AVERROR(ENOMEM);
 433
 434     s->nb_channels = outlink->channels;
 435     s->nb_coef_channels = ctx->inputs[1]->channels;
 436     s->want_skip = 1;
 437     s->need_padding = 1;
 438     s->pts = AV_NOPTS_VALUE;
 439
 440     return 0;
 441 }
 442
 443 static av_cold void uninit(AVFilterContext *ctx)
 444 {
 445     AudioFIRContext *s = ctx->priv;
 446     int ch;
 447
 448     if (s->sum) {
 449         for (ch = 0; ch < s->nb_channels; ch++) {
 450             av_freep(&s->sum[ch]);
 451         }
 452     }
 453     av_freep(&s->sum);
 454
 455     if (s->coeff) {
 456         for (ch = 0; ch < s->nb_coef_channels; ch++) {
 457             av_freep(&s->coeff[ch]);
 458         }
 459     }
 460     av_freep(&s->coeff);
 461
 462     if (s->block) {
 463         for (ch = 0; ch < s->nb_channels; ch++) {
 464             av_freep(&s->block[ch]);
 465         }
 466     }
 467     av_freep(&s->block);
 468
 469     if (s->rdft) {
 470         for (ch = 0; ch < s->nb_channels; ch++) {
 471             av_rdft_end(s->rdft[ch]);
 472         }
 473     }
 474     av_freep(&s->rdft);
 475
 476     if (s->irdft) {
 477         for (ch = 0; ch < s->nb_channels; ch++) {
 478             av_rdft_end(s->irdft[ch]);
 479         }
 480     }
 481     av_freep(&s->irdft);
 482
 483     av_frame_free(&s->in[0]);
 484     av_frame_free(&s->in[1]);
 485     av_frame_free(&s->buffer);
 486
 487     av_audio_fifo_free(s->fifo[0]);
 488     av_audio_fifo_free(s->fifo[1]);
 489
 490     av_freep(&s->fdsp);
 491 }
 492
 493 static av_cold int init(AVFilterContext *ctx)
 494 {
 495     AudioFIRContext *s = ctx->priv;
 496
 497     s->fcmul_add = fcmul_add_c;
 498
 499     s->fdsp = avpriv_float_dsp_alloc(0);
 500     if (!s->fdsp)
 501         return AVERROR(ENOMEM);
 502
 503     if (ARCH_X86)
 504         ff_afir_init_x86(s);
 505
 506     return 0;
 507 }
 508
 509 static const AVFilterPad afir_inputs[] = {
 510     {
 511         .name           = "main",
 512         .type           = AVMEDIA_TYPE_AUDIO,
 513         .filter_frame   = filter_frame,
 514     },{
 515         .name           = "ir",
 516         .type           = AVMEDIA_TYPE_AUDIO,
 517         .filter_frame   = read_ir,
 518     },
 519     { NULL }
 520 };
 521
 522 static const AVFilterPad afir_outputs[] = {
 523     {
 524         .name          = "default",
 525         .type          = AVMEDIA_TYPE_AUDIO,
 526         .config_props  = config_output,
 527         .request_frame = request_frame,
 528     },
 529     { NULL }
 530 };
 531
 532 #define AF AV_OPT_FLAG_AUDIO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
 533 #define OFFSET(x) offsetof(AudioFIRContext, x)
 534
 535 static const AVOption afir_options[] = {
 536     { "dry",    "set dry gain",      OFFSET(dry_gain),   AV_OPT_TYPE_FLOAT, {.dbl=1},    0, 10, AF },
 537     { "wet",    "set wet gain",      OFFSET(wet_gain),   AV_OPT_TYPE_FLOAT, {.dbl=1},    0, 10, AF },
 538     { "length", "set IR length",     OFFSET(length),     AV_OPT_TYPE_FLOAT, {.dbl=1},    0,  1, AF },
 539     { "again",  "enable auto gain",  OFFSET(again),      AV_OPT_TYPE_BOOL,  {.i64=1},    0,  1, AF },
 540     { "maxir",  "set max ir length", OFFSET(max_ir_len), AV_OPT_TYPE_FLOAT, {.dbl=30}, 0.1, 60, AF },
 541     { NULL }
 542 };
 543
 544 AVFILTER_DEFINE_CLASS(afir);
 545
 546 AVFilter ff_af_afir = {
 547     .name          = "afir",
 548     .description   = NULL_IF_CONFIG_SMALL("Apply Finite Impulse Response filter with supplied coefficients in 2nd stream."),
 549     .priv_size     = sizeof(AudioFIRContext),
 550     .priv_class    = &afir_class,
 551     .query_formats = query_formats,
 552     .init          = init,
 553     .uninit        = uninit,
 554     .inputs        = afir_inputs,
 555     .outputs       = afir_outputs,
 556     .flags         = AVFILTER_FLAG_SLICE_THREADS,
 557 };