• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 <float.h>
22 
23 #include "libavutil/avassert.h"
24 #include "libavutil/audio_fifo.h"
25 #include "libavutil/avstring.h"
26 #include "libavutil/opt.h"
27 #include "avfilter.h"
28 #include "audio.h"
29 #include "formats.h"
30 
31 #include "af_anlmdndsp.h"
32 
33 #define WEIGHT_LUT_NBITS 20
34 #define WEIGHT_LUT_SIZE  (1<<WEIGHT_LUT_NBITS)
35 
36 #define SQR(x) ((x) * (x))
37 
38 typedef struct AudioNLMeansContext {
39     const AVClass *class;
40 
41     float a;
42     int64_t pd;
43     int64_t rd;
44     float m;
45     int om;
46 
47     float pdiff_lut_scale;
48     float weight_lut[WEIGHT_LUT_SIZE];
49 
50     int K;
51     int S;
52     int N;
53     int H;
54 
55     int offset;
56     AVFrame *in;
57     AVFrame *cache;
58 
59     int64_t pts;
60 
61     AVAudioFifo *fifo;
62     int eof_left;
63 
64     AudioNLMDNDSPContext dsp;
65 } AudioNLMeansContext;
66 
67 enum OutModes {
68     IN_MODE,
69     OUT_MODE,
70     NOISE_MODE,
71     NB_MODES
72 };
73 
74 #define OFFSET(x) offsetof(AudioNLMeansContext, x)
75 #define AFT AV_OPT_FLAG_AUDIO_PARAM|AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_RUNTIME_PARAM
76 
77 static const AVOption anlmdn_options[] = {
78     { "s", "set denoising strength", OFFSET(a),  AV_OPT_TYPE_FLOAT,    {.dbl=0.00001},0.00001, 10, AFT },
79     { "p", "set patch duration",     OFFSET(pd), AV_OPT_TYPE_DURATION, {.i64=2000}, 1000, 100000, AFT },
80     { "r", "set research duration",  OFFSET(rd), AV_OPT_TYPE_DURATION, {.i64=6000}, 2000, 300000, AFT },
81     { "o", "set output mode",        OFFSET(om), AV_OPT_TYPE_INT,      {.i64=OUT_MODE},  0, NB_MODES-1, AFT, "mode" },
82     {  "i", "input",                 0,          AV_OPT_TYPE_CONST,    {.i64=IN_MODE},   0,  0, AFT, "mode" },
83     {  "o", "output",                0,          AV_OPT_TYPE_CONST,    {.i64=OUT_MODE},  0,  0, AFT, "mode" },
84     {  "n", "noise",                 0,          AV_OPT_TYPE_CONST,    {.i64=NOISE_MODE},0,  0, AFT, "mode" },
85     { "m", "set smooth factor",      OFFSET(m),  AV_OPT_TYPE_FLOAT,    {.dbl=11.},       1, 15, AFT },
86     { NULL }
87 };
88 
89 AVFILTER_DEFINE_CLASS(anlmdn);
90 
query_formats(AVFilterContext * ctx)91 static int query_formats(AVFilterContext *ctx)
92 {
93     AVFilterFormats *formats = NULL;
94     AVFilterChannelLayouts *layouts = NULL;
95     static const enum AVSampleFormat sample_fmts[] = {
96         AV_SAMPLE_FMT_FLTP,
97         AV_SAMPLE_FMT_NONE
98     };
99     int ret;
100 
101     formats = ff_make_format_list(sample_fmts);
102     if (!formats)
103         return AVERROR(ENOMEM);
104     ret = ff_set_common_formats(ctx, formats);
105     if (ret < 0)
106         return ret;
107 
108     layouts = ff_all_channel_counts();
109     if (!layouts)
110         return AVERROR(ENOMEM);
111 
112     ret = ff_set_common_channel_layouts(ctx, layouts);
113     if (ret < 0)
114         return ret;
115 
116     formats = ff_all_samplerates();
117     return ff_set_common_samplerates(ctx, formats);
118 }
119 
compute_distance_ssd_c(const float * f1,const float * f2,ptrdiff_t K)120 static float compute_distance_ssd_c(const float *f1, const float *f2, ptrdiff_t K)
121 {
122     float distance = 0.;
123 
124     for (int k = -K; k <= K; k++)
125         distance += SQR(f1[k] - f2[k]);
126 
127     return distance;
128 }
129 
compute_cache_c(float * cache,const float * f,ptrdiff_t S,ptrdiff_t K,ptrdiff_t i,ptrdiff_t jj)130 static void compute_cache_c(float *cache, const float *f,
131                             ptrdiff_t S, ptrdiff_t K,
132                             ptrdiff_t i, ptrdiff_t jj)
133 {
134     int v = 0;
135 
136     for (int j = jj; j < jj + S; j++, v++)
137         cache[v] += -SQR(f[i - K - 1] - f[j - K - 1]) + SQR(f[i + K] - f[j + K]);
138 }
139 
ff_anlmdn_init(AudioNLMDNDSPContext * dsp)140 void ff_anlmdn_init(AudioNLMDNDSPContext *dsp)
141 {
142     dsp->compute_distance_ssd = compute_distance_ssd_c;
143     dsp->compute_cache        = compute_cache_c;
144 
145     if (ARCH_X86)
146         ff_anlmdn_init_x86(dsp);
147 }
148 
config_filter(AVFilterContext * ctx)149 static int config_filter(AVFilterContext *ctx)
150 {
151     AudioNLMeansContext *s = ctx->priv;
152     AVFilterLink *outlink = ctx->outputs[0];
153     int newK, newS, newH, newN;
154     AVFrame *new_in, *new_cache;
155 
156     newK = av_rescale(s->pd, outlink->sample_rate, AV_TIME_BASE);
157     newS = av_rescale(s->rd, outlink->sample_rate, AV_TIME_BASE);
158 
159     newH = newK * 2 + 1;
160     newN = newH + (newK + newS) * 2;
161 
162     av_log(ctx, AV_LOG_DEBUG, "K:%d S:%d H:%d N:%d\n", newK, newS, newH, newN);
163 
164     if (!s->cache || s->cache->nb_samples < newS * 2) {
165         new_cache = ff_get_audio_buffer(outlink, newS * 2);
166         if (new_cache) {
167             av_frame_free(&s->cache);
168             s->cache = new_cache;
169         } else {
170             return AVERROR(ENOMEM);
171         }
172     }
173     if (!s->cache)
174         return AVERROR(ENOMEM);
175 
176     s->pdiff_lut_scale = 1.f / s->m * WEIGHT_LUT_SIZE;
177     for (int i = 0; i < WEIGHT_LUT_SIZE; i++) {
178         float w = -i / s->pdiff_lut_scale;
179 
180         s->weight_lut[i] = expf(w);
181     }
182 
183     if (!s->in || s->in->nb_samples < newN) {
184         new_in = ff_get_audio_buffer(outlink, newN);
185         if (new_in) {
186             av_frame_free(&s->in);
187             s->in = new_in;
188         } else {
189             return AVERROR(ENOMEM);
190         }
191     }
192     if (!s->in)
193         return AVERROR(ENOMEM);
194 
195     s->K = newK;
196     s->S = newS;
197     s->H = newH;
198     s->N = newN;
199 
200     return 0;
201 }
202 
config_output(AVFilterLink * outlink)203 static int config_output(AVFilterLink *outlink)
204 {
205     AVFilterContext *ctx = outlink->src;
206     AudioNLMeansContext *s = ctx->priv;
207     int ret;
208 
209     s->eof_left = -1;
210     s->pts = AV_NOPTS_VALUE;
211 
212     ret = config_filter(ctx);
213     if (ret < 0)
214         return ret;
215 
216     s->fifo = av_audio_fifo_alloc(outlink->format, outlink->channels, s->N);
217     if (!s->fifo)
218         return AVERROR(ENOMEM);
219 
220     ret = av_audio_fifo_write(s->fifo, (void **)s->in->extended_data, s->K + s->S);
221     if (ret < 0)
222         return ret;
223 
224     ff_anlmdn_init(&s->dsp);
225 
226     return 0;
227 }
228 
filter_channel(AVFilterContext * ctx,void * arg,int ch,int nb_jobs)229 static int filter_channel(AVFilterContext *ctx, void *arg, int ch, int nb_jobs)
230 {
231     AudioNLMeansContext *s = ctx->priv;
232     AVFrame *out = arg;
233     const int S = s->S;
234     const int K = s->K;
235     const int om = s->om;
236     const float *f = (const float *)(s->in->extended_data[ch]) + K;
237     float *cache = (float *)s->cache->extended_data[ch];
238     const float sw = (65536.f / (4 * K + 2)) / sqrtf(s->a);
239     float *dst = (float *)out->extended_data[ch] + s->offset;
240     const float smooth = s->m;
241 
242     for (int i = S; i < s->H + S; i++) {
243         float P = 0.f, Q = 0.f;
244         int v = 0;
245 
246         if (i == S) {
247             for (int j = i - S; j <= i + S; j++) {
248                 if (i == j)
249                     continue;
250                 cache[v++] = s->dsp.compute_distance_ssd(f + i, f + j, K);
251             }
252         } else {
253             s->dsp.compute_cache(cache, f, S, K, i, i - S);
254             s->dsp.compute_cache(cache + S, f, S, K, i, i + 1);
255         }
256 
257         for (int j = 0; j < 2 * S && !ctx->is_disabled; j++) {
258             const float distance = cache[j];
259             unsigned weight_lut_idx;
260             float w;
261 
262             if (distance < 0.f) {
263                 cache[j] = 0.f;
264                 continue;
265             }
266             w = distance * sw;
267             if (w >= smooth)
268                 continue;
269             weight_lut_idx = w * s->pdiff_lut_scale;
270             av_assert2(weight_lut_idx < WEIGHT_LUT_SIZE);
271             w = s->weight_lut[weight_lut_idx];
272             P += w * f[i - S + j + (j >= S)];
273             Q += w;
274         }
275 
276         P += f[i];
277         Q += 1;
278 
279         switch (om) {
280         case IN_MODE:    dst[i - S] = f[i];           break;
281         case OUT_MODE:   dst[i - S] = P / Q;          break;
282         case NOISE_MODE: dst[i - S] = f[i] - (P / Q); break;
283         }
284     }
285 
286     return 0;
287 }
288 
filter_frame(AVFilterLink * inlink,AVFrame * in)289 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
290 {
291     AVFilterContext *ctx = inlink->dst;
292     AVFilterLink *outlink = ctx->outputs[0];
293     AudioNLMeansContext *s = ctx->priv;
294     AVFrame *out = NULL;
295     int available, wanted, ret;
296 
297     if (s->pts == AV_NOPTS_VALUE)
298         s->pts = in->pts;
299 
300     ret = av_audio_fifo_write(s->fifo, (void **)in->extended_data,
301                               in->nb_samples);
302     av_frame_free(&in);
303 
304     s->offset = 0;
305     available = av_audio_fifo_size(s->fifo);
306     wanted = (available / s->H) * s->H;
307 
308     if (wanted >= s->H && available >= s->N) {
309         out = ff_get_audio_buffer(outlink, wanted);
310         if (!out)
311             return AVERROR(ENOMEM);
312     }
313 
314     while (available >= s->N) {
315         ret = av_audio_fifo_peek(s->fifo, (void **)s->in->extended_data, s->N);
316         if (ret < 0)
317             break;
318 
319         ctx->internal->execute(ctx, filter_channel, out, NULL, inlink->channels);
320 
321         av_audio_fifo_drain(s->fifo, s->H);
322 
323         s->offset += s->H;
324         available -= s->H;
325     }
326 
327     if (out) {
328         out->pts = s->pts;
329         out->nb_samples = s->offset;
330         if (s->eof_left >= 0) {
331             out->nb_samples = FFMIN(s->eof_left, s->offset);
332             s->eof_left -= out->nb_samples;
333         }
334         s->pts += av_rescale_q(s->offset, (AVRational){1, outlink->sample_rate}, outlink->time_base);
335 
336         return ff_filter_frame(outlink, out);
337     }
338 
339     return ret;
340 }
341 
request_frame(AVFilterLink * outlink)342 static int request_frame(AVFilterLink *outlink)
343 {
344     AVFilterContext *ctx = outlink->src;
345     AudioNLMeansContext *s = ctx->priv;
346     int ret;
347 
348     ret = ff_request_frame(ctx->inputs[0]);
349 
350     if (ret == AVERROR_EOF && s->eof_left != 0) {
351         AVFrame *in;
352 
353         if (s->eof_left < 0)
354             s->eof_left = av_audio_fifo_size(s->fifo) - (s->S + s->K);
355         if (s->eof_left <= 0)
356             return AVERROR_EOF;
357         in = ff_get_audio_buffer(outlink, s->H);
358         if (!in)
359             return AVERROR(ENOMEM);
360 
361         return filter_frame(ctx->inputs[0], in);
362     }
363 
364     return ret;
365 }
366 
process_command(AVFilterContext * ctx,const char * cmd,const char * args,char * res,int res_len,int flags)367 static int process_command(AVFilterContext *ctx, const char *cmd, const char *args,
368                            char *res, int res_len, int flags)
369 {
370     int ret;
371 
372     ret = ff_filter_process_command(ctx, cmd, args, res, res_len, flags);
373     if (ret < 0)
374         return ret;
375 
376     ret = config_filter(ctx);
377     if (ret < 0)
378         return ret;
379 
380     return 0;
381 }
382 
uninit(AVFilterContext * ctx)383 static av_cold void uninit(AVFilterContext *ctx)
384 {
385     AudioNLMeansContext *s = ctx->priv;
386 
387     av_audio_fifo_free(s->fifo);
388     av_frame_free(&s->in);
389     av_frame_free(&s->cache);
390 }
391 
392 static const AVFilterPad inputs[] = {
393     {
394         .name         = "default",
395         .type         = AVMEDIA_TYPE_AUDIO,
396         .filter_frame = filter_frame,
397     },
398     { NULL }
399 };
400 
401 static const AVFilterPad outputs[] = {
402     {
403         .name          = "default",
404         .type          = AVMEDIA_TYPE_AUDIO,
405         .config_props  = config_output,
406         .request_frame = request_frame,
407     },
408     { NULL }
409 };
410 
411 AVFilter ff_af_anlmdn = {
412     .name          = "anlmdn",
413     .description   = NULL_IF_CONFIG_SMALL("Reduce broadband noise from stream using Non-Local Means."),
414     .query_formats = query_formats,
415     .priv_size     = sizeof(AudioNLMeansContext),
416     .priv_class    = &anlmdn_class,
417     .uninit        = uninit,
418     .inputs        = inputs,
419     .outputs       = outputs,
420     .process_command = process_command,
421     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL |
422                      AVFILTER_FLAG_SLICE_THREADS,
423 };
424