1use crate::stft::{fft_size_for, hop_for, pow_db, Fft, FLOOR_DB};
11use ferraus_core::{AudioBuffer, Error, Result, FERRAUS_VERSION};
12use realfft::num_complex::Complex32;
13use serde::{Deserialize, Serialize};
14use serde_json::Value;
15use std::collections::VecDeque;
16
17pub const PROFILE_FORMAT: &str = "ferraus-noiseprof";
18pub const PROFILE_FORMAT_VERSION: u32 = 1;
19
20#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
22pub struct BinStats {
23 pub mean_db: Vec<f32>,
24 pub std_db: Vec<f32>,
25}
26
27#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
29pub struct ProfileSource {
30 pub method: String,
32 #[serde(default, skip_serializing_if = "Option::is_none")]
33 pub path: Option<String>,
34 #[serde(default, skip_serializing_if = "Option::is_none")]
35 pub start_s: Option<f64>,
36 #[serde(default, skip_serializing_if = "Option::is_none")]
37 pub end_s: Option<f64>,
38 #[serde(default, skip_serializing_if = "Vec::is_empty")]
39 pub merged_from: Vec<String>,
40}
41
42#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
44pub struct NoiseProfile {
45 pub format: String,
46 pub format_version: u32,
47 pub ferraus_version: String,
48 pub sample_rate: u32,
49 pub fft_size: usize,
50 pub hop_size: usize,
51 pub window: String,
52 pub magnitude_scale: String,
53 pub channels: usize,
54 pub frames: usize,
56 pub bin_hz: f64,
57 #[serde(default)]
58 pub source: ProfileSource,
59 pub mean_db: Vec<Vec<f32>>,
61 pub std_db: Vec<Vec<f32>>,
63 pub linked: BinStats,
65}
66
67fn round3(x: f64) -> f32 {
68 ((x * 1000.0).round() / 1000.0) as f32
69}
70
71impl NoiseProfile {
72 pub fn bins(&self) -> usize {
73 self.fft_size / 2 + 1
74 }
75
76 pub fn stats_for(&self, group: usize, link: bool, input_channels: usize) -> (&[f32], &[f32]) {
79 if link || self.channels != input_channels {
80 (&self.linked.mean_db, &self.linked.std_db)
81 } else {
82 let c = group.min(self.channels - 1);
83 (&self.mean_db[c], &self.std_db[c])
84 }
85 }
86
87 pub fn mean_level_db(&self) -> f32 {
89 let v = &self.linked.mean_db;
90 v.iter().sum::<f32>() / v.len().max(1) as f32
91 }
92
93 pub fn validate(&self) -> Result<()> {
95 let bad = |m: String| {
96 Error::input(format!("invalid noise profile: {m}"))
97 .with_hint("re-learn it with 'ferraus denoise profile <input> --auto -o room.noiseprof'")
98 };
99 if self.format != PROFILE_FORMAT {
100 return Err(bad(format!("format is '{}', expected '{PROFILE_FORMAT}'", self.format)));
101 }
102 if self.format_version > PROFILE_FORMAT_VERSION {
103 return Err(bad(format!(
104 "format_version {} is newer than this ferraus understands ({PROFILE_FORMAT_VERSION})",
105 self.format_version
106 )));
107 }
108 if self.fft_size < 16 || !self.fft_size.is_power_of_two() {
109 return Err(bad(format!("fft_size {} is not a power of two", self.fft_size)));
110 }
111 if self.sample_rate == 0 || self.channels == 0 {
112 return Err(bad("sample_rate and channels must be positive".into()));
113 }
114 let bins = self.bins();
115 if self.mean_db.len() != self.channels || self.std_db.len() != self.channels {
116 return Err(bad(format!("expected {} channel arrays in mean_db/std_db", self.channels)));
117 }
118 let all = self.mean_db.iter().chain(&self.std_db).chain([&self.linked.mean_db, &self.linked.std_db]);
119 for a in all {
120 if a.len() != bins {
121 return Err(bad(format!("expected {bins} bins (fft_size/2+1), found {}", a.len())));
122 }
123 if a.iter().any(|v| !v.is_finite()) {
124 return Err(bad("non-finite values".into()));
125 }
126 }
127 Ok(())
128 }
129
130 pub fn from_json(bytes: &[u8]) -> Result<NoiseProfile> {
131 let p: NoiseProfile = serde_json::from_slice(bytes).map_err(|e| {
132 Error::input(format!("cannot parse noise profile: {e}"))
133 .with_hint("profiles are JSON written by 'ferraus denoise profile'")
134 })?;
135 p.validate()?;
136 Ok(p)
137 }
138
139 pub fn to_json(&self) -> Vec<u8> {
142 let v = serde_json::to_value(self).expect("profile serializes");
143 let mut out = String::from("{\n");
144 if let Value::Object(m) = v {
145 let n = m.len();
146 for (i, (k, val)) in m.iter().enumerate() {
147 out.push_str(" ");
148 out.push_str(&serde_json::to_string(k).unwrap());
149 out.push_str(": ");
150 out.push_str(&serde_json::to_string(val).unwrap());
151 if i + 1 < n {
152 out.push(',');
153 }
154 out.push('\n');
155 }
156 }
157 out.push_str("}\n");
158 out.into_bytes()
159 }
160}
161
162pub struct ProfileLearner {
164 fft: Fft,
165 sample_rate: u32,
166 channels: usize,
167 frames: usize,
168 mean: Vec<Vec<f64>>,
169 m2: Vec<Vec<f64>>,
170 lmean: Vec<f64>,
171 lm2: Vec<f64>,
172 spec: Vec<Complex32>,
173 lpow: Vec<f32>,
174}
175
176impl ProfileLearner {
177 pub fn new(sample_rate: u32, channels: usize) -> ProfileLearner {
178 let fft = Fft::new(fft_size_for(sample_rate));
179 let bins = fft.bins();
180 let channels = channels.max(1);
181 ProfileLearner {
182 sample_rate,
183 channels,
184 frames: 0,
185 mean: vec![vec![0.0; bins]; channels],
186 m2: vec![vec![0.0; bins]; channels],
187 lmean: vec![0.0; bins],
188 lm2: vec![0.0; bins],
189 spec: vec![Complex32::new(0.0, 0.0); bins],
190 lpow: vec![0.0; bins],
191 fft,
192 }
193 }
194
195 pub fn fft_size(&self) -> usize {
196 self.fft.n
197 }
198
199 pub fn frames(&self) -> usize {
200 self.frames
201 }
202
203 pub fn add_frame(&mut self, frame: &[&[f32]]) {
205 self.frames += 1;
206 let nf = self.frames as f64;
207 let scale2 = self.fft.mag_scale * self.fft.mag_scale;
208 self.lpow.iter_mut().for_each(|p| *p = 0.0);
209 for c in 0..self.channels {
210 let x = frame[c.min(frame.len() - 1)];
211 self.fft.forward(x, &mut self.spec);
212 for (k, z) in self.spec.iter().enumerate() {
213 let p = z.norm_sqr() * scale2;
214 self.lpow[k] += p / self.channels as f32;
215 let db = pow_db(p) as f64;
216 let d = db - self.mean[c][k];
217 self.mean[c][k] += d / nf;
218 self.m2[c][k] += d * (db - self.mean[c][k]);
219 }
220 }
221 for k in 0..self.lpow.len() {
222 let db = pow_db(self.lpow[k]) as f64;
223 let d = db - self.lmean[k];
224 self.lmean[k] += d / nf;
225 self.lm2[k] += d * (db - self.lmean[k]);
226 }
227 }
228
229 pub fn add_region(&mut self, buf: &AudioBuffer, start: usize, end: usize) {
231 let n = self.fft.n;
232 let hop = hop_for(n);
233 let end = end.min(buf.frames());
234 let mut s = start;
235 while s + n <= end {
236 let frame: Vec<&[f32]> = buf.data.iter().map(|c| &c[s..s + n]).collect();
237 self.add_frame(&frame);
238 s += hop;
239 }
240 }
241
242 pub fn finish(self, source: ProfileSource) -> Result<NoiseProfile> {
243 if self.frames == 0 {
244 let ms = self.fft.n as f64 * 1000.0 / self.sample_rate as f64;
245 return Err(Error::usage(format!(
246 "the noise region is too short: it needs at least {ms:.0} ms ({} samples)",
247 self.fft.n
248 )));
249 }
250 let nf = self.frames as f64;
251 let std = |m2: f64| {
252 if self.frames > 1 {
253 (m2 / (nf - 1.0)).max(0.0).sqrt()
254 } else {
255 0.0
256 }
257 };
258 let mean_db: Vec<Vec<f32>> =
259 self.mean.iter().map(|c| c.iter().map(|v| round3(v.max(FLOOR_DB as f64))).collect()).collect();
260 let std_db: Vec<Vec<f32>> = self.m2.iter().map(|c| c.iter().map(|v| round3(std(*v))).collect()).collect();
261 let linked = BinStats {
262 mean_db: self.lmean.iter().map(|v| round3(v.max(FLOOR_DB as f64))).collect(),
263 std_db: self.lm2.iter().map(|v| round3(std(*v))).collect(),
264 };
265 let n = self.fft.n;
266 Ok(NoiseProfile {
267 format: PROFILE_FORMAT.into(),
268 format_version: PROFILE_FORMAT_VERSION,
269 ferraus_version: FERRAUS_VERSION.into(),
270 sample_rate: self.sample_rate,
271 fft_size: n,
272 hop_size: hop_for(n),
273 window: "hann".into(),
274 magnitude_scale: "dbfs_sine".into(),
275 channels: self.channels,
276 frames: self.frames,
277 bin_hz: self.sample_rate as f64 / n as f64,
278 source,
279 mean_db,
280 std_db,
281 linked,
282 })
283 }
284}
285
286pub fn learn(buf: &AudioBuffer, start: usize, end: usize, source: ProfileSource) -> Result<NoiseProfile> {
288 let mut l = ProfileLearner::new(buf.sample_rate, buf.channels());
289 l.add_region(buf, start, end);
290 l.finish(source)
291}
292
293#[derive(Clone, Debug, PartialEq, Serialize)]
295pub struct QuietRegion {
296 pub start: usize,
297 pub end: usize,
298 pub start_s: f64,
299 pub end_s: f64,
300 pub level_db: f32,
302 pub stability_db: f32,
304}
305
306const AUTO_BANDS: usize = 16;
307const STABILITY_WEIGHT: f64 = 1.0;
309const PEAK_WEIGHT: f64 = 0.5;
312const SILENCE_DB: f64 = -130.0;
314
315pub fn find_quiet_region(buf: &AudioBuffer, window_s: f64) -> Option<QuietRegion> {
318 let sr = buf.sample_rate;
319 let n = fft_size_for(sr);
320 let hop = hop_for(n);
321 let total = buf.frames();
322 if total < n {
323 return None;
324 }
325 let available = (total - n) / hop + 1;
326 let want = ((window_s * sr as f64 - n as f64) / hop as f64).floor() as isize + 1;
327 let w = (want.max(1) as usize).min(available);
328
329 let mut fft = Fft::new(n);
330 let bins = fft.bins();
331 let scale2 = fft.mag_scale * fft.mag_scale;
332 let lo = 60.0f64;
334 let hi = (sr as f64 / 2.0).min(16_000.0).max(lo * 2.0);
335 let bin_hz = sr as f64 / n as f64;
336 let edges: Vec<usize> = (0..=AUTO_BANDS)
337 .map(|b| {
338 let f = lo * (hi / lo).powf(b as f64 / AUTO_BANDS as f64);
339 ((f / bin_hz).round() as usize).clamp(1, bins - 1)
340 })
341 .collect();
342 let mut mono = vec![0f32; n];
343 let mut spec = vec![Complex32::new(0.0, 0.0); bins];
344 let mut ring: VecDeque<[f64; AUTO_BANDS + 1]> = VecDeque::with_capacity(w + 1);
346 let mut sum = [0f64; AUTO_BANDS + 1];
347 let mut sum2 = [0f64; AUTO_BANDS + 1];
348 let mut best: Option<(f64, usize, f64, f64)> = None; let mut best_silent: Option<(usize, f64)> = None;
350 let chans = buf.channels() as f32;
351 for f in 0..available {
352 let s = f * hop;
353 for (i, m) in mono.iter_mut().enumerate() {
354 *m = buf.data.iter().map(|c| c[s + i]).sum::<f32>() / chans;
355 }
356 fft.forward(&mono, &mut spec);
357 let mut feat = [0f64; AUTO_BANDS + 1];
358 let mut tot = 0f64;
359 for z in &spec {
360 tot += (z.norm_sqr() * scale2) as f64;
361 }
362 feat[0] = pow_db((tot / bins as f64) as f32) as f64;
363 for b in 0..AUTO_BANDS {
364 let (a, e) = (edges[b], edges[b + 1].max(edges[b] + 1));
365 let mut p = 0f64;
366 for z in &spec[a..e.min(bins)] {
367 p += (z.norm_sqr() * scale2) as f64;
368 }
369 feat[b + 1] = pow_db((p / (e - a) as f64) as f32) as f64;
370 }
371 for i in 0..=AUTO_BANDS {
372 sum[i] += feat[i];
373 sum2[i] += feat[i] * feat[i];
374 }
375 ring.push_back(feat);
376 if ring.len() > w {
377 let old = ring.pop_front().unwrap();
378 for i in 0..=AUTO_BANDS {
379 sum[i] -= old[i];
380 sum2[i] -= old[i] * old[i];
381 }
382 }
383 if ring.len() == w {
384 let wn = w as f64;
385 let level = sum[0] / wn;
386 let mut stab = 0f64;
387 for i in 1..=AUTO_BANDS {
388 let m = sum[i] / wn;
389 stab += (sum2[i] / wn - m * m).max(0.0).sqrt();
390 }
391 stab /= AUTO_BANDS as f64;
392 let first = f + 1 - w;
393 if level < SILENCE_DB {
394 if best_silent.is_none() {
395 best_silent = Some((first, level));
396 }
397 continue;
398 }
399 let peak = ring.iter().map(|f| f[0]).fold(f64::NEG_INFINITY, f64::max);
400 let score = level + STABILITY_WEIGHT * stab + PEAK_WEIGHT * (peak - level);
401 if best.map(|b| score < b.0).unwrap_or(true) {
402 best = Some((score, first, level, stab));
403 }
404 }
405 }
406 let (first, level, stab) = match (best, best_silent) {
407 (Some((_, f, l, s)), _) => (f, l, s),
408 (None, Some((f, l))) => (f, l, 0.0),
409 (None, None) => return None,
410 };
411 let start = first * hop;
412 let end = (start + (w - 1) * hop + n).min(total);
413 Some(QuietRegion {
414 start,
415 end,
416 start_s: start as f64 / sr as f64,
417 end_s: end as f64 / sr as f64,
418 level_db: level as f32,
419 stability_db: stab as f32,
420 })
421}
422
423#[derive(Clone, Copy, Debug, PartialEq, Eq)]
425pub enum MergeMethod {
426 Max,
429 Mean,
431}
432
433impl MergeMethod {
434 pub fn parse(s: &str) -> Result<MergeMethod> {
435 match s {
436 "max" => Ok(MergeMethod::Max),
437 "mean" => Ok(MergeMethod::Mean),
438 other => Err(Error::usage(format!("--method must be max or mean, not '{other}'"))),
439 }
440 }
441 pub fn name(self) -> &'static str {
442 match self {
443 MergeMethod::Max => "max",
444 MergeMethod::Mean => "mean",
445 }
446 }
447}
448
449fn merge_stats(parts: &[(&[f32], &[f32], usize)], method: MergeMethod) -> BinStats {
450 let bins = parts[0].0.len();
451 let mut mean_db = vec![0f32; bins];
452 let mut std_db = vec![0f32; bins];
453 match method {
454 MergeMethod::Max => {
455 for k in 0..bins {
456 mean_db[k] = parts.iter().map(|p| p.0[k]).fold(f32::MIN, f32::max);
457 std_db[k] = parts.iter().map(|p| p.1[k]).fold(f32::MIN, f32::max);
458 }
459 }
460 MergeMethod::Mean => {
461 let total: f64 = parts.iter().map(|p| p.2.max(1) as f64).sum();
462 for k in 0..bins {
463 let mut m = 0f64;
464 let mut s2 = 0f64;
465 for p in parts {
466 let wt = p.2.max(1) as f64 / total;
467 let (mu, sd) = (p.0[k] as f64, p.1[k] as f64);
468 m += wt * mu;
469 s2 += wt * (sd * sd + mu * mu);
470 }
471 mean_db[k] = round3(m);
472 std_db[k] = round3((s2 - m * m).max(0.0).sqrt());
473 }
474 }
475 }
476 BinStats { mean_db, std_db }
477}
478
479pub fn merge(profiles: &[NoiseProfile], names: &[String], method: MergeMethod) -> Result<(NoiseProfile, Vec<String>)> {
481 if profiles.len() < 2 {
482 return Err(Error::usage("merge needs at least two profiles"));
483 }
484 let first = &profiles[0];
485 for (p, name) in profiles.iter().zip(names).skip(1) {
486 if p.sample_rate != first.sample_rate || p.fft_size != first.fft_size {
487 return Err(Error::usage(format!(
488 "{name} was learned at {} Hz with FFT size {}, but the first profile uses {} Hz / {}; profiles can only be merged when both match",
489 p.sample_rate, p.fft_size, first.sample_rate, first.fft_size
490 ))
491 .with_hint("convert the takes to one sample rate ('ferraus convert in.wav --sample-rate 48000 -o out.wav') and re-learn the profiles"));
492 }
493 }
494 let mut warnings = Vec::new();
495 let same_channels = profiles.iter().all(|p| p.channels == first.channels);
496 let channels = profiles.iter().map(|p| p.channels).max().unwrap_or(1);
497 if !same_channels {
498 warnings.push(format!(
499 "profiles have different channel counts; per-channel statistics were built from each profile's linked statistics ({channels} channels)"
500 ));
501 }
502 let mut mean_db = Vec::with_capacity(channels);
503 let mut std_db = Vec::with_capacity(channels);
504 for c in 0..channels {
505 let parts: Vec<(&[f32], &[f32], usize)> = profiles
506 .iter()
507 .map(|p| {
508 if same_channels {
509 (p.mean_db[c].as_slice(), p.std_db[c].as_slice(), p.frames)
510 } else {
511 (p.linked.mean_db.as_slice(), p.linked.std_db.as_slice(), p.frames)
512 }
513 })
514 .collect();
515 let s = merge_stats(&parts, method);
516 mean_db.push(s.mean_db);
517 std_db.push(s.std_db);
518 }
519 let lparts: Vec<(&[f32], &[f32], usize)> =
520 profiles.iter().map(|p| (p.linked.mean_db.as_slice(), p.linked.std_db.as_slice(), p.frames)).collect();
521 let linked = merge_stats(&lparts, method);
522 let merged = NoiseProfile {
523 format: PROFILE_FORMAT.into(),
524 format_version: PROFILE_FORMAT_VERSION,
525 ferraus_version: FERRAUS_VERSION.into(),
526 sample_rate: first.sample_rate,
527 fft_size: first.fft_size,
528 hop_size: first.hop_size,
529 window: first.window.clone(),
530 magnitude_scale: first.magnitude_scale.clone(),
531 channels,
532 frames: profiles.iter().map(|p| p.frames).sum(),
533 bin_hz: first.bin_hz,
534 source: ProfileSource {
535 method: format!("merge-{}", method.name()),
536 merged_from: names.to_vec(),
537 ..Default::default()
538 },
539 mean_db,
540 std_db,
541 linked,
542 };
543 Ok((merged, warnings))
544}