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 |
|
|
|
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 |
|
|
|
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 |
|
|
|
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 |
|
|
|
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 |
|
|
|
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 |
|
|
|
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 |
|
|
|
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 |
|
|
|
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 |
|
|
|
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 |
|
|
|
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 |
|
|
|
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 |
|
|
}; |