libavfilter/af_anlms.c

   1 /*
   2  * Copyright (c) 2019 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 "libavutil/channel_layout.h"
  22 #include "libavutil/common.h"
  23 #include "libavutil/float_dsp.h"
  24 #include "libavutil/opt.h"
  25
  26 #include "audio.h"
  27 #include "avfilter.h"
  28 #include "formats.h"
  29 #include "filters.h"
  30 #include "internal.h"
  31
  32 enum OutModes {
  33     IN_MODE,
  34     DESIRED_MODE,
  35     OUT_MODE,
  36     NOISE_MODE,
  37     NB_OMODES
  38 };
  39
  40 typedef struct AudioNLMSContext {
  41     const AVClass *class;
  42
  43     int order;
  44     float mu;
  45     float eps;
  46     float leakage;
  47     int output_mode;
  48
  49     int kernel_size;
  50     AVFrame *offset;
  51     AVFrame *delay;
  52     AVFrame *coeffs;
  53     AVFrame *tmp;
  54
  55     AVFrame *frame[2];
  56
  57     int anlmf;
  58
  59     AVFloatDSPContext *fdsp;
  60 } AudioNLMSContext;
  61
  62 #define OFFSET(x) offsetof(AudioNLMSContext, x)
  63 #define A AV_OPT_FLAG_AUDIO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
  64 #define AT AV_OPT_FLAG_AUDIO_PARAM|AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_RUNTIME_PARAM
  65
  66 static const AVOption anlms_options[] = {
  67     { "order",   "set the filter order",   OFFSET(order),   AV_OPT_TYPE_INT,   {.i64=256},  1, INT16_MAX, A },
  68     { "mu",      "set the filter mu",      OFFSET(mu),      AV_OPT_TYPE_FLOAT, {.dbl=0.75}, 0, 2, AT },
  69     { "eps",     "set the filter eps",     OFFSET(eps),     AV_OPT_TYPE_FLOAT, {.dbl=1},    0, 1, AT },
  70     { "leakage", "set the filter leakage", OFFSET(leakage), AV_OPT_TYPE_FLOAT, {.dbl=0},    0, 1, AT },
  71     { "out_mode", "set output mode",       OFFSET(output_mode), AV_OPT_TYPE_INT, {.i64=OUT_MODE}, 0, NB_OMODES-1, AT, "mode" },
  72     {  "i", "input",                 0,          AV_OPT_TYPE_CONST,    {.i64=IN_MODE},      0, 0, AT, "mode" },
  73     {  "d", "desired",               0,          AV_OPT_TYPE_CONST,    {.i64=DESIRED_MODE}, 0, 0, AT, "mode" },
  74     {  "o", "output",                0,          AV_OPT_TYPE_CONST,    {.i64=OUT_MODE},     0, 0, AT, "mode" },
  75     {  "n", "noise",                 0,          AV_OPT_TYPE_CONST,    {.i64=NOISE_MODE},   0, 0, AT, "mode" },
  76     { NULL }
  77 };
  78
  79 AVFILTER_DEFINE_CLASS_EXT(anlms, "anlm(f|s)", anlms_options);
  80
  81 static int query_formats(AVFilterContext *ctx)
  82 {
  83     static const enum AVSampleFormat sample_fmts[] = {
  84         AV_SAMPLE_FMT_FLTP,
  85         AV_SAMPLE_FMT_NONE
  86     };
  87     int ret = ff_set_common_all_channel_counts(ctx);
  88     if (ret < 0)
  89         return ret;
  90
  91     ret = ff_set_common_formats_from_list(ctx, sample_fmts);
  92     if (ret < 0)
  93         return ret;
  94
  95     return ff_set_common_all_samplerates(ctx);
  96 }
  97
  98 static float fir_sample(AudioNLMSContext *s, float sample, float *delay,
  99                         float *coeffs, float *tmp, int *offset)
 100 {
 101     const int order = s->order;
 102     float output;
 103
 104     delay[*offset] = sample;
 105
 106     memcpy(tmp, coeffs + order - *offset, order * sizeof(float));
 107
 108     output = s->fdsp->scalarproduct_float(delay, tmp, s->kernel_size);
 109
 110     if (--(*offset) < 0)
 111         *offset = order - 1;
 112
 113     return output;
 114 }
 115
 116 static float process_sample(AudioNLMSContext *s, float input, float desired,
 117                             float *delay, float *coeffs, float *tmp, int *offsetp)
 118 {
 119     const int order = s->order;
 120     const float leakage = s->leakage;
 121     const float mu = s->mu;
 122     const float a = 1.f - leakage * mu;
 123     float sum, output, e, norm, b;
 124     int offset = *offsetp;
 125
 126     delay[offset + order] = input;
 127
 128     output = fir_sample(s, input, delay, coeffs, tmp, offsetp);
 129     e = desired - output;
 130
 131     sum = s->fdsp->scalarproduct_float(delay, delay, s->kernel_size);
 132
 133     norm = s->eps + sum;
 134     b = mu * e / norm;
 135     if (s->anlmf)
 136         b *= 4.f * e * e;
 137
 138     memcpy(tmp, delay + offset, order * sizeof(float));
 139
 140     s->fdsp->vector_fmul_scalar(coeffs, coeffs, a, s->kernel_size);
 141
 142     s->fdsp->vector_fmac_scalar(coeffs, tmp, b, s->kernel_size);
 143
 144     memcpy(coeffs + order, coeffs, order * sizeof(float));
 145
 146     switch (s->output_mode) {
 147     case IN_MODE:       output = input;         break;
 148     case DESIRED_MODE:  output = desired;       break;
 149     case OUT_MODE: /*output = output;*/         break;
 150     case NOISE_MODE: output = desired - output; break;
 151     }
 152     return output;
 153 }
 154
 155 static int process_channels(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
 156 {
 157     AudioNLMSContext *s = ctx->priv;
 158     AVFrame *out = arg;
 159     const int start = (out->ch_layout.nb_channels * jobnr) / nb_jobs;
 160     const int end = (out->ch_layout.nb_channels * (jobnr+1)) / nb_jobs;
 161
 162     for (int c = start; c < end; c++) {
 163         const float *input = (const float *)s->frame[0]->extended_data[c];
 164         const float *desired = (const float *)s->frame[1]->extended_data[c];
 165         float *delay = (float *)s->delay->extended_data[c];
 166         float *coeffs = (float *)s->coeffs->extended_data[c];
 167         float *tmp = (float *)s->tmp->extended_data[c];
 168         int *offset = (int *)s->offset->extended_data[c];
 169         float *output = (float *)out->extended_data[c];
 170
 171         for (int n = 0; n < out->nb_samples; n++) {
 172             output[n] = process_sample(s, input[n], desired[n], delay, coeffs, tmp, offset);
 173             if (ctx->is_disabled)
 174                 output[n] = input[n];
 175         }
 176     }
 177
 178     return 0;
 179 }
 180
 181 static int activate(AVFilterContext *ctx)
 182 {
 183     AudioNLMSContext *s = ctx->priv;
 184     int i, ret, status;
 185     int nb_samples;
 186     int64_t pts;
 187
 188     FF_FILTER_FORWARD_STATUS_BACK_ALL(ctx->outputs[0], ctx);
 189
 190     nb_samples = FFMIN(ff_inlink_queued_samples(ctx->inputs[0]),
 191                        ff_inlink_queued_samples(ctx->inputs[1]));
 192     for (i = 0; i < ctx->nb_inputs && nb_samples > 0; i++) {
 193         if (s->frame[i])
 194             continue;
 195
 196         if (ff_inlink_check_available_samples(ctx->inputs[i], nb_samples) > 0) {
 197             ret = ff_inlink_consume_samples(ctx->inputs[i], nb_samples, nb_samples, &s->frame[i]);
 198             if (ret < 0)
 199                 return ret;
 200         }
 201     }
 202
 203     if (s->frame[0] && s->frame[1]) {
 204         AVFrame *out;
 205
 206         out = ff_get_audio_buffer(ctx->outputs[0], s->frame[0]->nb_samples);
 207         if (!out) {
 208             av_frame_free(&s->frame[0]);
 209             av_frame_free(&s->frame[1]);
 210             return AVERROR(ENOMEM);
 211         }
 212
 213         ff_filter_execute(ctx, process_channels, out, NULL,
 214                           FFMIN(ctx->outputs[0]->ch_layout.nb_channels, ff_filter_get_nb_threads(ctx)));
 215
 216         out->pts = s->frame[0]->pts;
 217
 218         av_frame_free(&s->frame[0]);
 219         av_frame_free(&s->frame[1]);
 220
 221         ret = ff_filter_frame(ctx->outputs[0], out);
 222         if (ret < 0)
 223             return ret;
 224     }
 225
 226     if (!nb_samples) {
 227         for (i = 0; i < 2; i++) {
 228             if (ff_inlink_acknowledge_status(ctx->inputs[i], &status, &pts)) {
 229                 ff_outlink_set_status(ctx->outputs[0], status, pts);
 230                 return 0;
 231             }
 232         }
 233     }
 234
 235     if (ff_outlink_frame_wanted(ctx->outputs[0])) {
 236         for (i = 0; i < 2; i++) {
 237             if (ff_inlink_queued_samples(ctx->inputs[i]) > 0)
 238                 continue;
 239             ff_inlink_request_frame(ctx->inputs[i]);
 240             return 0;
 241         }
 242     }
 243     return 0;
 244 }
 245
 246 static int config_output(AVFilterLink *outlink)
 247 {
 248     AVFilterContext *ctx = outlink->src;
 249     AudioNLMSContext *s = ctx->priv;
 250
 251     s->anlmf = !strcmp(ctx->filter->name, "anlmf");
 252     s->kernel_size = FFALIGN(s->order, 16);
 253
 254     if (!s->offset)
 255         s->offset = ff_get_audio_buffer(outlink, 1);
 256     if (!s->delay)
 257         s->delay = ff_get_audio_buffer(outlink, 2 * s->kernel_size);
 258     if (!s->coeffs)
 259         s->coeffs = ff_get_audio_buffer(outlink, 2 * s->kernel_size);
 260     if (!s->tmp)
 261         s->tmp = ff_get_audio_buffer(outlink, s->kernel_size);
 262     if (!s->delay || !s->coeffs || !s->offset || !s->tmp)
 263         return AVERROR(ENOMEM);
 264
 265     return 0;
 266 }
 267
 268 static av_cold int init(AVFilterContext *ctx)
 269 {
 270     AudioNLMSContext *s = ctx->priv;
 271
 272     s->fdsp = avpriv_float_dsp_alloc(0);
 273     if (!s->fdsp)
 274         return AVERROR(ENOMEM);
 275
 276     return 0;
 277 }
 278
 279 static av_cold void uninit(AVFilterContext *ctx)
 280 {
 281     AudioNLMSContext *s = ctx->priv;
 282
 283     av_freep(&s->fdsp);
 284     av_frame_free(&s->delay);
 285     av_frame_free(&s->coeffs);
 286     av_frame_free(&s->offset);
 287     av_frame_free(&s->tmp);
 288 }
 289
 290 static const AVFilterPad inputs[] = {
 291     {
 292         .name = "input",
 293         .type = AVMEDIA_TYPE_AUDIO,
 294     },
 295     {
 296         .name = "desired",
 297         .type = AVMEDIA_TYPE_AUDIO,
 298     },
 299 };
 300
 301 static const AVFilterPad outputs[] = {
 302     {
 303         .name         = "default",
 304         .type         = AVMEDIA_TYPE_AUDIO,
 305         .config_props = config_output,
 306     },
 307 };
 308
 309 const AVFilter ff_af_anlms = {
 310     .name           = "anlms",
 311     .description    = NULL_IF_CONFIG_SMALL("Apply Normalized Least-Mean-Squares algorithm to first audio stream."),
 312     .priv_size      = sizeof(AudioNLMSContext),
 313     .priv_class     = &anlms_class,
 314     .init           = init,
 315     .uninit         = uninit,
 316     .activate       = activate,
 317     FILTER_INPUTS(inputs),
 318     FILTER_OUTPUTS(outputs),
 319     FILTER_QUERY_FUNC(query_formats),
 320     .flags          = AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL |
 321                       AVFILTER_FLAG_SLICE_THREADS,
 322     .process_command = ff_filter_process_command,
 323 };
 324
 325 const AVFilter ff_af_anlmf = {
 326     .name           = "anlmf",
 327     .description    = NULL_IF_CONFIG_SMALL("Apply Normalized Least-Mean-Fourth algorithm to first audio stream."),
 328     .priv_size      = sizeof(AudioNLMSContext),
 329     .priv_class     = &anlms_class,
 330     .init           = init,
 331     .uninit         = uninit,
 332     .activate       = activate,
 333     FILTER_INPUTS(inputs),
 334     FILTER_OUTPUTS(outputs),
 335     FILTER_QUERY_FUNC(query_formats),
 336     .flags          = AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL |
 337                       AVFILTER_FLAG_SLICE_THREADS,
 338     .process_command = ff_filter_process_command,
 339 };