| """Recording-chain augmentation — teaches the model real-world audio.
|
|
|
| An earlier model trained on clean LibriSpeech generalized to LibriSpeech but
|
| failed on real phone/laptop recordings: it never saw mic coloration,
|
| limited bandwidth, compression artifacts, or room tone. This module degrades the
|
| synthetic mixtures to *look like* real recordings, closing that domain gap.
|
|
|
| Two kinds of degradation, applied differently on purpose:
|
|
|
| LINEAR (EQ tilt, band-limiting): these color the whole recording, target voice
|
| included. We apply the SAME linear filter to the mixture AND the clean target,
|
| so the model isn't asked to invent bandwidth the "mic" never captured — it only
|
| has to remove the interferer + noise within the recorded band.
|
|
|
| MIXTURE-ONLY (soft clipping, mu-law codec sim, broadband noise floor): these
|
| simulate junk in the captured signal. Applied to the mixture only; the target
|
| stays the clean reference, so the model learns to clean them off.
|
|
|
| All ops are numpy/scipy so they run in DataLoader workers without ffmpeg overhead.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| from dataclasses import dataclass
|
|
|
| import numpy as np
|
| from scipy.signal import butter, sosfilt
|
|
|
| EPS = 1e-8
|
|
|
|
|
| @dataclass
|
| class AugConfig:
|
| enabled: bool = False
|
|
|
| eq_prob: float = 0.6
|
| bandlimit_prob: float = 0.5
|
| clip_prob: float = 0.3
|
| mulaw_prob: float = 0.3
|
| noise_floor_prob: float = 0.5
|
|
|
| lowpass_hz: tuple[float, float] = (3400.0, 7800.0)
|
|
|
| eq_tilt_db: tuple[float, float] = (-9.0, 9.0)
|
|
|
| clip_drive: tuple[float, float] = (2.0, 8.0)
|
|
|
| mulaw_levels: tuple[int, int] = (64, 256)
|
|
|
| noise_floor_snr_db: tuple[float, float] = (20.0, 40.0)
|
|
|
|
|
| class RecordingAugment:
|
| """Sample a recording chain, then apply its linear/mixture-only parts."""
|
|
|
| def __init__(self, cfg: AugConfig, sr: int):
|
| self.cfg = cfg
|
| self.sr = sr
|
|
|
|
|
| def sample_params(self, rng: np.random.Generator) -> dict:
|
| c = self.cfg
|
| p: dict = {}
|
| p["eq"] = (
|
| float(rng.uniform(*c.eq_tilt_db)) if rng.random() < c.eq_prob else None
|
| )
|
| p["lowpass"] = (
|
| float(rng.uniform(*c.lowpass_hz)) if rng.random() < c.bandlimit_prob else None
|
| )
|
| p["clip"] = (
|
| float(rng.uniform(*c.clip_drive)) if rng.random() < c.clip_prob else None
|
| )
|
| p["mulaw"] = (
|
| int(rng.integers(c.mulaw_levels[0], c.mulaw_levels[1] + 1))
|
| if rng.random() < c.mulaw_prob else None
|
| )
|
| p["noise_floor"] = (
|
| float(rng.uniform(*c.noise_floor_snr_db))
|
| if rng.random() < c.noise_floor_prob else None
|
| )
|
| return p
|
|
|
|
|
| def apply_linear(self, wav: np.ndarray, params: dict) -> np.ndarray:
|
| out = wav
|
| if params.get("eq") is not None:
|
| out = self._spectral_tilt(out, params["eq"])
|
| if params.get("lowpass") is not None:
|
| out = self._lowpass(out, params["lowpass"])
|
| return out.astype(np.float32, copy=False)
|
|
|
|
|
| def apply_mixture_only(
|
| self, wav: np.ndarray, params: dict, rng: np.random.Generator
|
| ) -> np.ndarray:
|
| out = wav
|
| if params.get("clip") is not None:
|
| out = self._soft_clip(out, params["clip"])
|
| if params.get("mulaw") is not None:
|
| out = self._mulaw(out, params["mulaw"])
|
| if params.get("noise_floor") is not None:
|
| out = self._add_noise_floor(out, params["noise_floor"], rng)
|
| return out.astype(np.float32, copy=False)
|
|
|
|
|
| def _spectral_tilt(self, wav: np.ndarray, tilt_db: float) -> np.ndarray:
|
|
|
|
|
| n = len(wav)
|
| if n < 8:
|
| return wav
|
| spec = np.fft.rfft(wav)
|
| ramp = np.linspace(-0.5, 0.5, spec.shape[0])
|
| gain = 10 ** ((tilt_db * ramp) / 20.0)
|
| return np.fft.irfft(spec * gain, n=n)
|
|
|
| def _lowpass(self, wav: np.ndarray, cutoff_hz: float) -> np.ndarray:
|
| nyq = self.sr / 2.0
|
| cutoff = min(max(cutoff_hz, 500.0), nyq * 0.99)
|
| sos = butter(6, cutoff / nyq, btype="low", output="sos")
|
| return sosfilt(sos, wav)
|
|
|
| def _soft_clip(self, wav: np.ndarray, drive: float) -> np.ndarray:
|
| peak = np.max(np.abs(wav)) + EPS
|
| x = wav / peak
|
| y = np.tanh(drive * x) / np.tanh(drive)
|
| return y * peak
|
|
|
| def _mulaw(self, wav: np.ndarray, levels: int) -> np.ndarray:
|
| peak = np.max(np.abs(wav)) + EPS
|
| x = np.clip(wav / peak, -1.0, 1.0)
|
| mu = float(levels - 1)
|
| comp = np.sign(x) * np.log1p(mu * np.abs(x)) / np.log1p(mu)
|
| q = np.round((comp + 1) / 2 * (levels - 1)) / (levels - 1) * 2 - 1
|
| dec = np.sign(q) * (1.0 / mu) * (np.power(1 + mu, np.abs(q)) - 1)
|
| return dec * peak
|
|
|
| def _add_noise_floor(
|
| self, wav: np.ndarray, snr_db: float, rng: np.random.Generator
|
| ) -> np.ndarray:
|
| sig_rms = np.sqrt(np.mean(wav**2) + EPS)
|
| noise = rng.standard_normal(len(wav)).astype(np.float32)
|
| noise_rms = np.sqrt(np.mean(noise**2) + EPS)
|
| target_noise_rms = sig_rms / (10 ** (snr_db / 20.0))
|
| return wav + noise * (target_noise_rms / noise_rms)
|
|
|
