Match URGENT's SNR definition and 4 kHz band limit

configs/data/multiasr.yaml

@@ -882,8 +882,9 @@ noise_snr_db: [-5.0, 20.0]
 noise_corpus: FluidInference/musan
 noise_corpus_glob: "noise/**/*.wav"
 bandlimit_prob: 0.1
-# Uniform over the list. 4 kHz is URGENT's band limit (resampling to 8 kHz);
-# 5.5 and 7 kHz mimic band-limited source recordings, which URGENT inherits
+# Uniform over the list. 4 kHz is URGENT's band limit, done as URGENT does it
+# (resample to 8 kHz and back); 5.5 and 7 kHz are brickwalls mimicking
+# band-limited source recordings, which URGENT inherits
 # from its speech corpora (cutoffs measured on its test clips: 121 at 4 kHz,
 # 16 at 5 kHz, 51 at 7 kHz).
 bandlimit_hz: [4000, 4000, 4000, 5512, 7000]

scripts/train_augment.py

@@ -69,12 +69,43 @@ def _peak_safe(audio: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
     return audio / peak if peak > 1.0 else audio
 
 
+def non_silent(
+    x: npt.NDArray[Any], threshold: float = 0.01, frame_length: int = 1024, frame_shift: int = 512
+) -> npt.NDArray[np.bool_]:
+    """ESPnet's `detect_non_silence` (espnet2/train/preprocessor.py), for one channel.
+
+    A sample is non-silent when its frame's power exceeds `threshold` of the
+    clip's mean frame power. URGENT measures both sides of an SNR over these
+    samples only, so pauses in the speech or gaps in a bursty noise do not
+    shift the mix.
+    """
+    if x.size < frame_length:
+        return np.ones(x.size, dtype=bool)
+    nadd = (-(x.size - frame_length) % frame_shift) % frame_length
+    padded = np.pad(x.astype(np.float64), (0, nadd))
+    frames = np.lib.stride_tricks.sliding_window_view(padded, frame_length)[::frame_shift]
+    power = np.mean(frames**2, axis=-1)
+    mean_power = power.mean()
+    if mean_power == 0:
+        return np.ones(x.size, dtype=bool)
+    detects = np.repeat(power / mean_power > threshold, frame_shift)[: x.size]
+    return np.pad(detects, (0, x.size - detects.size), mode="edge")
+
+
+def _active_power(x: npt.NDArray[Any]) -> float:
+    return float(np.mean(np.square(x[non_silent(x)], dtype=np.float64)))
+
+
 def mix_at_snr(
     speech: npt.NDArray[Any], noise: npt.NDArray[Any], snr_db: float
 ) -> npt.NDArray[np.float32]:
-    """`speech` with `noise` (tiled or cropped to its length) mixed in `snr_db` below it."""
+    """`speech` with `noise` (tiled or cropped to its length) mixed in `snr_db` below it.
+
+    URGENT's `mix_noise`: speech and noise power are taken over their
+    non-silent samples (`non_silent`), with the same 1e-10 floor on the noise.
+    """
     noise = _fit(noise, speech.size)
-    scale = _rms(speech) / (_rms(noise) * 10 ** (snr_db / 20))
+    scale = 10 ** (-snr_db / 20) * np.sqrt(_active_power(speech) / max(_active_power(noise), 1e-10))
     return _peak_safe(speech.astype(np.float32) + np.float32(scale) * noise)
 
 
@@ -87,13 +118,27 @@ def colored_noise(size: int, exponent: float, rng: np.random.Generator) -> npt.N
     return noise / _rms(noise)
 
 
+# URGENT's bandwidth limitation resamples to one of these lower rates and back.
+URGENT_SAMPLE_RATES = (8000, 16000, 22050, 24000, 32000, 44100, 48000)
+
+
 def bandlimit(
     audio: npt.NDArray[Any], sample_rate: int, cutoff_hz: float
 ) -> npt.NDArray[np.float32]:
-    """`audio` with every frequency above `cutoff_hz` removed (a brickwall, like
-    URGENT's cutoffs: a 30+ dB drop within 500 Hz)."""
+    """`audio` with every frequency above `cutoff_hz` removed.
+
+    When `2 * cutoff_hz` is one of URGENT's rates (at 16 kHz input only 8 kHz,
+    the 4 kHz limit), this is URGENT's own augmentation: resample down and back
+    with `resample_poly`, i.e. librosa's "polyphase", one of URGENT's resamplers.
+    Other cutoffs mimic band-limited source recordings, whose edges measure
+    sharp (a 30+ dB drop within 500 Hz), so they are a brickwall.
+    """
     if 2 * cutoff_hz >= sample_rate:
         return audio.astype(np.float32, copy=False)
+    if int(2 * cutoff_hz) in URGENT_SAMPLE_RATES:
+        ratio = Fraction(int(2 * cutoff_hz), sample_rate).limit_denominator(1000)
+        down = resample_poly(audio, ratio.numerator, ratio.denominator)
+        return _fit(resample_poly(down, ratio.denominator, ratio.numerator), audio.size)
     spectrum = np.fft.rfft(audio)
     spectrum[np.fft.rfftfreq(audio.size, d=1 / sample_rate) > cutoff_hz] = 0
     return np.fft.irfft(spectrum, n=audio.size).astype(np.float32)

tests/test_train_augment.py

@@ -17,6 +17,7 @@
     clip,
     colored_noise,
     mix_at_snr,
+    non_silent,
     phone,
     reverberate,
     room_response,
@@ -55,6 +56,38 @@ def test_never_clips(self) -> None:
         assert np.abs(mix_at_snr(loud, loud, -5.0)).max() <= 1.0
 
 
+class TestNonSilent:
+    """ESPnet's detect_non_silence: frames above 1% of the mean frame power are active."""
+
+    def test_burst_is_active_and_silence_is_not(self) -> None:
+        x = np.zeros(16000, dtype=np.float32)
+        x[4096:8192] = _tone(4096 / 16000, amp=0.3)
+        mask = non_silent(x)
+        assert mask.shape == x.shape
+        assert mask[5000:7000].all()
+        assert not mask[:3000].any()
+        assert not mask[10000:].any()
+
+    def test_short_or_silent_clips_count_as_all_active(self) -> None:
+        assert non_silent(np.ones(500, dtype=np.float32)).all()
+        assert non_silent(np.zeros(4000, dtype=np.float32)).all()
+
+
+class TestUrgentSnr:
+    """mix_at_snr measures SNR over non-silent samples, as URGENT's mix_noise does."""
+
+    def test_snr_ignores_pauses_and_noise_gaps(self) -> None:
+        speech = np.zeros(32000, dtype=np.float32)
+        speech[:16000] = _tone(1.0, amp=0.1)  # a 1 s pause after the speech
+        noise = np.random.default_rng(0).standard_normal(32000).astype(np.float32)
+        noise[8000:24000] = 0  # a bursty noise
+        added = mix_at_snr(speech, noise, 5.0) - speech
+        snr = 10 * np.log10(
+            np.mean(speech[non_silent(speech)] ** 2) / np.mean(added[non_silent(added)] ** 2)
+        )
+        assert snr == pytest.approx(5.0, abs=0.05)
+
+
 class TestColoredNoise:
     """colored_noise is unit RMS and tilts toward low frequencies as the exponent grows."""
 
@@ -85,6 +118,25 @@ def test_cutoff_at_or_above_nyquist_is_a_no_op(self) -> None:
         np.testing.assert_array_equal(bandlimit(tone, 16000, 8000.0), tone)
 
 
+class TestUrgentBandlimit:
+    """The 4 kHz limit is URGENT's resampling to 8 kHz and back, not a brickwall."""
+
+    def test_resampled_limit_cuts_the_top_band_and_keeps_the_passband(self) -> None:
+        white = colored_noise(32000, 0.0, np.random.default_rng(0)) * 0.1
+        out = bandlimit(white, 16000, 4000.0)
+        assert out.shape == white.shape
+        assert _high_share(out, 4500.0) < 0.01
+
+        def passband(a: npt.NDArray[np.float32]) -> float:  # ~100 Hz - 3 kHz
+            return float((np.abs(np.fft.rfft(a))[200:6000] ** 2).sum())
+
+        assert passband(out) / passband(white) == pytest.approx(1.0, rel=0.05)
+
+    def test_other_cutoffs_stay_brickwall(self) -> None:
+        white = colored_noise(32000, 0.0, np.random.default_rng(1)) * 0.1
+        assert _high_share(bandlimit(white, 16000, 7000.0), 7000.0 + 1) < 1e-6
+
+
 class TestClip:
     """clip flattens everything outside two quantiles (URGENT's scheme) and keeps the level."""