Skip to main content

ferraus_denoise/
profile.rs

1//! Noise profiles: learning, automatic noise-region detection, the `.noiseprof` JSON
2//! format, and merging.
3//!
4//! A profile stores, per frequency bin, the mean and standard deviation of the STFT
5//! magnitude in dB across the profile frames (per channel, plus the statistics of the
6//! channel-mean power used when channels are linked), together with the FFT size, hop,
7//! sample rate, channel count, and the ferraus version that wrote it. Levels use the
8//! `dbfs_sine` scale: a full-scale sine reads 0 dB at its bin.
9
10use 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/// Mean and standard deviation of the magnitude in dB, per bin.
21#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
22pub struct BinStats {
23    pub mean_db: Vec<f32>,
24    pub std_db: Vec<f32>,
25}
26
27/// Where a profile came from (informational).
28#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
29pub struct ProfileSource {
30    /// `range`, `auto`, `from`, or `merge-max` / `merge-mean`.
31    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/// A learned noise fingerprint (the `.noiseprof` file).
43#[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    /// Number of STFT frames the statistics were computed from.
55    pub frames: usize,
56    pub bin_hz: f64,
57    #[serde(default)]
58    pub source: ProfileSource,
59    /// `mean_db[channel][bin]`.
60    pub mean_db: Vec<Vec<f32>>,
61    /// `std_db[channel][bin]`.
62    pub std_db: Vec<Vec<f32>>,
63    /// Statistics of the channel-mean power (used with `--link on`).
64    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    /// Statistics to use for mask group `group`: the linked statistics when channels are
77    /// linked or when the profile's channel count does not match the input.
78    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    /// Mean level of the profile across bins, in dB (for reports).
88    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    /// Validate structure (used after parsing a file).
94    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    /// JSON with one top-level key per line and arrays kept on one line, so the file stays
140    /// small and diffable while remaining plain JSON.
141    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
162/// Streaming accumulator of per-bin dB statistics (Welford), constant memory.
163pub 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    /// Add one frame (`frame[c]` has `fft_size` samples per channel).
204    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    /// Add every frame that lies fully inside `[start, end)` of `buf`.
230    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
286/// Learn a profile from `[start, end)` of `buf`.
287pub 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/// The region chosen by `--auto`.
294#[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    /// Mean level of the region (dBFS-sine scale, averaged over bins).
301    pub level_db: f32,
302    /// Mean over bands of the level's standard deviation across frames (dB); lower is steadier.
303    pub stability_db: f32,
304}
305
306const AUTO_BANDS: usize = 16;
307/// Weight of spectral instability against loudness when ranking candidate windows.
308const STABILITY_WEIGHT: f64 = 1.0;
309/// Weight of the loudest frame above the window mean (penalizes windows that touch the
310/// start or end of a sound).
311const PEAK_WEIGHT: f64 = 0.5;
312/// Windows quieter than this are digital silence and useless as a noise sample.
313const SILENCE_DB: f64 = -130.0;
314
315/// Find the quietest spectrally stable window of `window_s` seconds (or a shorter window
316/// when the input is short). Streams over the input with constant memory.
317pub 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    // log-spaced band edges between 60 Hz and min(16 kHz, Nyquist)
333    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    // ring of per-frame feature vectors: [total_db, band_db...]
345    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; // (score, frame index, level, stab)
349    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/// How `denoise profile merge` combines profiles.
424#[derive(Clone, Copy, Debug, PartialEq, Eq)]
425pub enum MergeMethod {
426    /// Element-wise maximum of means and of standard deviations (the most conservative
427    /// fingerprint: removes the loudest noise seen in any take).
428    Max,
429    /// Pooled statistics weighted by frame count (as if all frames had been profiled together).
430    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
479/// Merge two or more profiles. Returns the merged profile and any warnings.
480pub 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}