File size: 6,198 Bytes
d67c04d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """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
# Probability each op fires per sample.
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
# Band-limit cutoff range (Hz). 3.4 kHz ~ telephone; 7.8 kHz ~ near full band.
lowpass_hz: tuple[float, float] = (3400.0, 7800.0)
# First-order spectral tilt in dB across the band (+ = brighter, - = darker).
eq_tilt_db: tuple[float, float] = (-9.0, 9.0)
# Soft-clip drive (higher = more distortion). tanh(drive * x) / tanh(drive).
clip_drive: tuple[float, float] = (2.0, 8.0)
# mu-law companding quantization levels (lower = more codec-like crunch).
mulaw_levels: tuple[int, int] = (64, 256)
# Broadband noise floor SNR (dB) relative to signal RMS.
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
# ---- parameter sampling -------------------------------------------------
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
# ---- linear ops (apply identically to mixture AND target) ---------------
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)
# ---- mixture-only ops (degradations to be cleaned off) ------------------
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)
# ---- individual effects -------------------------------------------------
def _spectral_tilt(self, wav: np.ndarray, tilt_db: float) -> np.ndarray:
# Frequency-domain first-order tilt: linear gain ramp from -tilt/2 at DC
# to +tilt/2 at Nyquist (dB), applied to the magnitude spectrum.
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) # encode
q = np.round((comp + 1) / 2 * (levels - 1)) / (levels - 1) * 2 - 1 # quantize
dec = np.sign(q) * (1.0 / mu) * (np.power(1 + mu, np.abs(q)) - 1) # decode
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)
|