From 58a177c5ebf00342fe0b966b65dfffbedea3c404 Mon Sep 17 00:00:00 2001 From: Joe Davis Date: Wed, 17 Jun 2026 18:40:32 -0700 Subject: [PATCH 1/7] ADR-227: Add BackgroundNoiseAugmentor and MicrophoneNoiseAugmentor - background_noise_stage.py: BackgroundNoiseAugmentor with NoiseProvider protocol; noise_file always chosen for hash stability; noise_start_s bounds derived from file durations at runtime; all three keys always present in applied_values - mic_noise_stage.py: MicrophoneNoiseAugmentor with Gaussian noise seeded from output_seed - speech_05_add_background_noise.py, speech_06_add_mic_noise.py: entry-point scripts - params.py: AddBackgroundNoiseParams, AddMicNoiseParams; PipelineParams.load() extended - params.yaml: add_background_noise and add_mic_noise sections - _doc_speech.md: updated stages table and key design decisions --- ml/params.yaml | 8 + ml/pipeline/speech/_doc_speech.md | 24 +- ml/pipeline/speech/background_noise_stage.py | 181 ++++++ ml/pipeline/speech/mic_noise_stage.py | 100 ++++ ml/pipeline/stages/params.py | 28 + .../stages/speech_05_add_background_noise.py | 68 +++ ml/pipeline/stages/speech_06_add_mic_noise.py | 52 ++ .../speech/test_background_noise_stage.py | 517 ++++++++++++++++++ .../pipeline/speech/test_mic_noise_stage.py | 330 +++++++++++ ml/test/pipeline/stages/test_params.py | 87 ++- 10 files changed, 1392 insertions(+), 3 deletions(-) create mode 100644 ml/pipeline/speech/background_noise_stage.py create mode 100644 ml/pipeline/speech/mic_noise_stage.py create mode 100644 ml/pipeline/stages/speech_05_add_background_noise.py create mode 100644 ml/pipeline/stages/speech_06_add_mic_noise.py create mode 100644 ml/test/pipeline/speech/test_background_noise_stage.py create mode 100644 ml/test/pipeline/speech/test_mic_noise_stage.py diff --git a/ml/params.yaml b/ml/params.yaml index 60bca4c5..b7326429 100644 --- a/ml/params.yaml +++ b/ml/params.yaml @@ -20,3 +20,11 @@ stages: suffix_vary_probability: 0.333 suffix_min_s: 0.1 suffix_max_s: 0.5 + add_background_noise: + vary_probability: 0.5 + volume_min: 0.0 + volume_max: 0.3 + add_mic_noise: + vary_probability: 0.5 + amplitude_min: 0.001 + amplitude_max: 0.05 diff --git a/ml/pipeline/speech/_doc_speech.md b/ml/pipeline/speech/_doc_speech.md index 493c998d..9b27553c 100644 --- a/ml/pipeline/speech/_doc_speech.md +++ b/ml/pipeline/speech/_doc_speech.md @@ -10,9 +10,10 @@ produces a derived `AudioSample` manifest with per-sample variation applied. |--------|-------------|-----------| | [`tts_stage.py`](tts_stage.py) | `TtsSampleGenerator` | `TextSample → AudioSample` via edge_tts | | [`delay_stage.py`](delay_stage.py) | `DelayAugmentor` | `AudioSample → AudioSample` (silence padding) | +| [`background_noise_stage.py`](background_noise_stage.py) | `BackgroundNoiseAugmentor` | `AudioSample → AudioSample` (environmental noise mix) | +| [`mic_noise_stage.py`](mic_noise_stage.py) | `MicrophoneNoiseAugmentor` | `AudioSample → AudioSample` (Gaussian mic noise) | -Planned stages (not yet implemented): `background_noise_stage.py`, -`mic_noise_stage.py`, `token_stage.py`, `spectrogram_stage.py`. +Planned stages (not yet implemented): `token_stage.py`, `spectrogram_stage.py`. ## Key design decisions @@ -36,3 +37,22 @@ prefix for the same reason. accepts `list[str]` voices directly. The entry-point runs `asyncio.run(edge_tts.list_voices())` and filters: `Gender == 'Female'`, `Locale == 'en-US'`, `':' not in ShortName`, `'DragonHD' not in ShortName`, `'Turbo' not in ShortName`. + +**`NoiseProvider` protocol as the noise-file seam for `BackgroundNoiseAugmentor`.** +Consistent with `TtsProvider` in `tts_stage.py` — the protocol lives in the same +module as the stage that uses it. Unit tests supply `_FakeNoiseProvider`; the +entry-point supplies `_DirectoryNoiseProvider` (globbing `*.wav` from `--noise-dir`). + +**`noise_file` is always chosen (hash stability), `noise_start_s`/`noise_volume` are +0.0 when not applied.** `VariationGenerator.choose()` runs on the sorted filename list +before the `should_vary` check so the content hash does not change if `vary_probability` +is toggled. All three keys are always present in `applied_values`. + +**`noise_start_s` bounds are derived from file durations at runtime.** Both the noise +file and the audio sample file are read in `_get_applied_values` when noise is applied. +`max_start_s = noise_duration_s - audio_duration_s`; clamped to 0.0 when negative +(i.e. noise file is shorter than audio). + +**Gaussian noise in `MicrophoneNoiseAugmentor` is seeded from `output_seed`.** Uses +`np.random.default_rng(output_seed).normal(0, amplitude, len(samples))` for +reproducibility. The amplitude is stored as 0.0 when not applied. diff --git a/ml/pipeline/speech/background_noise_stage.py b/ml/pipeline/speech/background_noise_stage.py new file mode 100644 index 00000000..d6fd14d6 --- /dev/null +++ b/ml/pipeline/speech/background_noise_stage.py @@ -0,0 +1,181 @@ +"""BackgroundNoiseAugmentor: mix environmental noise into WAV audio samples. + +noise_file is always chosen (choose() always called) for hash stability, even when +not applied. noise_start_s and noise_volume are stored as 0.0 when should_vary +returns False. All three keys (noise_file, noise_start_s, noise_volume) are always +present in every applied_values dict. + +noise_start_s bounds are derived from file durations: both the noise file and the +audio sample file are read to compute max_start_s = noise_duration_s - audio_duration_s. +When the noise file is shorter than the audio sample, max_start_s is negative and +clamped to 0.0 (min and max both 0.0). +""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +from pathlib import Path +from typing import Any, Protocol + +import numpy as np + +from pipeline.core.manifest import ManifestStore +from pipeline.core.modifier_stage import ModifierStage +from pipeline.core.randomization import MinMaxFilter, VariationGenerator +from pipeline.core.sample import AudioSample +from pipeline.io.audio_io import AudioData, AudioReader, AudioWriter +from pipeline.stages import conventions +from pipeline.stages.params import AddBackgroundNoiseParams + + +class NoiseProvider(Protocol): + """Protocol for listing available noise files. + + Consistent with TtsProvider living in tts_stage.py — the protocol belongs + alongside the stage that uses it. + """ + + def list_files(self) -> list[Path]: ... + + +def _read_duration_s(audio_reader: AudioReader, path: Path) -> float: + """Read a file and return its duration in seconds. + + Handles both a running event loop (uses a thread executor) and no loop + (uses asyncio.run()). + """ + async def _read() -> float: + data = await audio_reader.read(path) + return len(data.samples) / data.sample_rate + + try: + loop = asyncio.get_running_loop() + except RuntimeError: + # No running loop — safe to call asyncio.run() + return asyncio.run(_read()) + + # Inside a running loop — run in a thread to avoid blocking the loop + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(asyncio.run, _read()) + return future.result() + + +class BackgroundNoiseAugmentor(ModifierStage[AudioSample, AudioSample]): + """Augmentation stage that mixes environmental noise into WAV samples. + + The noise file is always chosen (even when not applied) so that the + content hash remains stable across configuration changes. Only + noise_start_s and noise_volume are set to 0.0 when the noise is not applied. + + Noise mixing algorithm: slice noise at noise_start_s for len(audio) samples, + zero-pad if noise is shorter after the slice, multiply by noise_volume, add to + audio samples, then clip to [-1.0, 1.0]. + + input_dir is the directory containing the input WAV files referenced by + AudioSample.path. This is typically the output directory of the preceding stage. + """ + + def __init__( + self, + output_dir: Path, + manifest_store: ManifestStore, + audio_reader: AudioReader, + audio_writer: AudioWriter, + input_dir: Path, + noise_provider: NoiseProvider, + params: AddBackgroundNoiseParams, + ) -> None: + super().__init__(output_dir, manifest_store) + self._audio_reader = audio_reader + self._audio_writer = audio_writer + self._input_dir = input_dir + self._noise_provider = noise_provider + self._vary_probability = params.vary_probability + self._volume_filter = MinMaxFilter(params.volume_min, params.volume_max, precision=2) + + def _get_applied_values( + self, sample: AudioSample, generator: VariationGenerator + ) -> dict[str, Any]: + # Always choose a noise file for hash stability + noise_files = sorted([p.name for p in self._noise_provider.list_files()]) + noise_file: str = generator.choose("noise_file", noise_files) + + if generator.should_vary("noise", self._vary_probability): + # Derive noise_start_s bounds from file durations + noise_path = self._noise_provider.list_files()[0].parent / noise_file + audio_path = self._input_dir / sample.path + + noise_duration_s = _read_duration_s(self._audio_reader, noise_path) + audio_duration_s = _read_duration_s(self._audio_reader, audio_path) + + max_start_s = noise_duration_s - audio_duration_s + clamped_max = max(0.0, max_start_s) + start_filter = MinMaxFilter(0.0, clamped_max, precision=2) + noise_start_s = generator.generate("noise_start_s", start_filter) + noise_volume = generator.generate("noise_volume", self._volume_filter) + else: + noise_start_s = 0.0 + noise_volume = 0.0 + + return { + "noise_file": noise_file, + "noise_start_s": float(noise_start_s), + "noise_volume": float(noise_volume), + } + + def _derive_id(self, input_sample: AudioSample, applied_values: dict[str, Any]) -> str: + noise_file: str = applied_values["noise_file"] + noise_volume: float = applied_values["noise_volume"] + noise_filestem = Path(noise_file).stem + return f"{input_sample.id}_{noise_filestem}_v{int(noise_volume * 100)}" + + async def _generate_output( + self, + input_sample: AudioSample, + output_id: str, + output_seed: int, + applied_values: dict[str, Any], + parent_content_hash: str, + ) -> AudioSample: + noise_file: str = applied_values["noise_file"] + noise_start_s: float = applied_values["noise_start_s"] + noise_volume: float = applied_values["noise_volume"] + + input_path = self._input_dir / input_sample.path + audio = await self._audio_reader.read(input_path) + output_samples = audio.samples.copy() + + if noise_volume > 0.0: + # Resolve noise file path via provider + noise_path = self._noise_provider.list_files()[0].parent / noise_file + noise_audio = await self._audio_reader.read(noise_path) + + start_sample = int(noise_start_s * noise_audio.sample_rate) + n_needed = len(output_samples) + noise_slice = noise_audio.samples[start_sample: start_sample + n_needed] + + # Zero-pad if noise slice is shorter than audio + if len(noise_slice) < n_needed: + noise_slice = np.pad(noise_slice, (0, n_needed - len(noise_slice))) + + output_samples = output_samples + noise_volume * noise_slice + output_samples = np.clip(output_samples, -1.0, 1.0).astype(np.float32) + + self._output_dir.mkdir(parents=True, exist_ok=True) + output_path = conventions.sample_file_path(self._output_dir, output_id, "wav") + await self._audio_writer.write(output_path, AudioData(samples=output_samples, sample_rate=audio.sample_rate)) + + content_hash = self._compute_content_hash( + parent_content_hash, output_seed, applied_values + ) + + return AudioSample( + id=output_id, + seed=output_seed, + content_hash=content_hash, + path=Path(f"{output_id}.wav"), + parent_content_hash=parent_content_hash, + transcript=input_sample.transcript, + applied_values=applied_values, + ) diff --git a/ml/pipeline/speech/mic_noise_stage.py b/ml/pipeline/speech/mic_noise_stage.py new file mode 100644 index 00000000..7fc0eeb2 --- /dev/null +++ b/ml/pipeline/speech/mic_noise_stage.py @@ -0,0 +1,100 @@ +"""MicrophoneNoiseAugmentor: add Gaussian microphone noise to WAV audio samples. + +mic_noise_amplitude is stored as 0.0 when should_vary returns False. +Gaussian noise is seeded from output_seed for reproducibility. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np + +from pipeline.core.manifest import ManifestStore +from pipeline.core.modifier_stage import ModifierStage +from pipeline.core.randomization import MinMaxFilter, VariationGenerator +from pipeline.core.sample import AudioSample +from pipeline.io.audio_io import AudioData, AudioReader, AudioWriter +from pipeline.stages import conventions +from pipeline.stages.params import AddMicNoiseParams + + +class MicrophoneNoiseAugmentor(ModifierStage[AudioSample, AudioSample]): + """Augmentation stage that adds Gaussian microphone noise to WAV samples. + + mic_noise_amplitude is stored as 0.0 when the noise is not applied. Gaussian + noise uses output_seed for reproducibility via numpy's default_rng. + + input_dir is the directory containing the input WAV files referenced by + AudioSample.path. This is typically the output directory of the preceding stage. + """ + + def __init__( + self, + output_dir: Path, + manifest_store: ManifestStore, + audio_reader: AudioReader, + audio_writer: AudioWriter, + input_dir: Path, + params: AddMicNoiseParams, + ) -> None: + super().__init__(output_dir, manifest_store) + self._audio_reader = audio_reader + self._audio_writer = audio_writer + self._input_dir = input_dir + self._vary_probability = params.vary_probability + self._amplitude_filter = MinMaxFilter(params.amplitude_min, params.amplitude_max, precision=3) + + def _get_applied_values( + self, sample: AudioSample, generator: VariationGenerator + ) -> dict[str, Any]: + if generator.should_vary("mic_noise_amplitude", self._vary_probability): + amplitude = generator.generate("mic_noise_amplitude", self._amplitude_filter) + else: + amplitude = 0.0 + + return { + "mic_noise_amplitude": float(amplitude), + } + + def _derive_id(self, input_sample: AudioSample, applied_values: dict[str, Any]) -> str: + amplitude: float = applied_values["mic_noise_amplitude"] + return f"{input_sample.id}_mic{int(amplitude * 1000)}" + + async def _generate_output( + self, + input_sample: AudioSample, + output_id: str, + output_seed: int, + applied_values: dict[str, Any], + parent_content_hash: str, + ) -> AudioSample: + amplitude: float = applied_values["mic_noise_amplitude"] + + input_path = self._input_dir / input_sample.path + audio = await self._audio_reader.read(input_path) + output_samples = audio.samples.copy() + + if amplitude > 0.0: + rng = np.random.default_rng(output_seed) + noise = rng.normal(0, amplitude, len(output_samples)).astype(np.float32) + output_samples = np.clip(output_samples + noise, -1.0, 1.0).astype(np.float32) + + self._output_dir.mkdir(parents=True, exist_ok=True) + output_path = conventions.sample_file_path(self._output_dir, output_id, "wav") + await self._audio_writer.write(output_path, AudioData(samples=output_samples, sample_rate=audio.sample_rate)) + + content_hash = self._compute_content_hash( + parent_content_hash, output_seed, applied_values + ) + + return AudioSample( + id=output_id, + seed=output_seed, + content_hash=content_hash, + path=Path(f"{output_id}.wav"), + parent_content_hash=parent_content_hash, + transcript=input_sample.transcript, + applied_values=applied_values, + ) diff --git a/ml/pipeline/stages/params.py b/ml/pipeline/stages/params.py index 69d0ac29..04e72a2e 100644 --- a/ml/pipeline/stages/params.py +++ b/ml/pipeline/stages/params.py @@ -37,6 +37,20 @@ class AddDelaysParams: suffix_max_s: float +@dataclass +class AddBackgroundNoiseParams: + vary_probability: float + volume_min: float + volume_max: float + + +@dataclass +class AddMicNoiseParams: + vary_probability: float + amplitude_min: float + amplitude_max: float + + @dataclass class PipelineParams: variations_per_phrase: int @@ -44,6 +58,8 @@ class PipelineParams: generate_phrases: GeneratePhraseParams generate_samples: GenerateSamplesParams add_delays: AddDelaysParams + add_background_noise: AddBackgroundNoiseParams + add_mic_noise: AddMicNoiseParams @classmethod def load(cls, path: Path) -> "PipelineParams": @@ -54,6 +70,8 @@ def load(cls, path: Path) -> "PipelineParams": stage_phrases = raw["stages"]["generate_phrases"] stage_samples = raw["stages"]["generate_speech_samples"] stage_delays = raw["stages"]["add_delays"] + stage_bg_noise = raw["stages"]["add_background_noise"] + stage_mic_noise = raw["stages"]["add_mic_noise"] return cls( variations_per_phrase=int(pipeline["variations_per_phrase"]), @@ -77,4 +95,14 @@ def load(cls, path: Path) -> "PipelineParams": suffix_min_s=float(stage_delays["suffix_min_s"]), suffix_max_s=float(stage_delays["suffix_max_s"]), ), + add_background_noise=AddBackgroundNoiseParams( + vary_probability=float(stage_bg_noise["vary_probability"]), + volume_min=float(stage_bg_noise["volume_min"]), + volume_max=float(stage_bg_noise["volume_max"]), + ), + add_mic_noise=AddMicNoiseParams( + vary_probability=float(stage_mic_noise["vary_probability"]), + amplitude_min=float(stage_mic_noise["amplitude_min"]), + amplitude_max=float(stage_mic_noise["amplitude_max"]), + ), ) diff --git a/ml/pipeline/stages/speech_05_add_background_noise.py b/ml/pipeline/stages/speech_05_add_background_noise.py new file mode 100644 index 00000000..c833468e --- /dev/null +++ b/ml/pipeline/stages/speech_05_add_background_noise.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import argparse +import asyncio +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +from pipeline.core.manifest import ManifestStore +from pipeline.io.audio_io import LibrosaAudioReader, SoundfileAudioWriter +from pipeline.speech.background_noise_stage import BackgroundNoiseAugmentor +from pipeline.stages import conventions +from pipeline.stages.params import PipelineParams + +_PROJECT_ROOT = Path(__file__).parents[2] + + +class _DirectoryNoiseProvider: + """NoiseProvider backed by a filesystem directory. + + Lists all WAV files in the given directory. DVC wiring points this at the + appropriate data/ path; this class requires no knowledge of the DVC layout. + """ + + def __init__(self, noise_dir: Path) -> None: + self._noise_dir = noise_dir + + def list_files(self) -> list[Path]: + return list(self._noise_dir.glob("*.wav")) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Add environmental background noise to WAV audio samples" + ) + parser.add_argument("--input-manifest-dir", required=True, type=Path) + parser.add_argument("--noise-dir", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + args = parser.parse_args() + + params = PipelineParams.load(conventions.params_path(_PROJECT_ROOT)) + + store = ManifestStore() + input_manifest = store.read(conventions.manifest_path(args.input_manifest_dir)) + + args.output_dir.mkdir(parents=True, exist_ok=True) + + stage = BackgroundNoiseAugmentor( + output_dir=args.output_dir, + manifest_store=store, + audio_reader=LibrosaAudioReader(), + audio_writer=SoundfileAudioWriter(), + input_dir=args.input_manifest_dir, + noise_provider=_DirectoryNoiseProvider(args.noise_dir), + params=params.add_background_noise, + ) + + asyncio.run( + stage.transform( + input_manifest, + conventions.manifest_path(args.output_dir), + ) + ) + + +if __name__ == "__main__": + main() diff --git a/ml/pipeline/stages/speech_06_add_mic_noise.py b/ml/pipeline/stages/speech_06_add_mic_noise.py new file mode 100644 index 00000000..aeb5cbfd --- /dev/null +++ b/ml/pipeline/stages/speech_06_add_mic_noise.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import argparse +import asyncio +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +from pipeline.core.manifest import ManifestStore +from pipeline.io.audio_io import LibrosaAudioReader, SoundfileAudioWriter +from pipeline.speech.mic_noise_stage import MicrophoneNoiseAugmentor +from pipeline.stages import conventions +from pipeline.stages.params import PipelineParams + +_PROJECT_ROOT = Path(__file__).parents[2] + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Add Gaussian microphone noise to WAV audio samples" + ) + parser.add_argument("--input-manifest-dir", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + args = parser.parse_args() + + params = PipelineParams.load(conventions.params_path(_PROJECT_ROOT)) + + store = ManifestStore() + input_manifest = store.read(conventions.manifest_path(args.input_manifest_dir)) + + args.output_dir.mkdir(parents=True, exist_ok=True) + + stage = MicrophoneNoiseAugmentor( + output_dir=args.output_dir, + manifest_store=store, + audio_reader=LibrosaAudioReader(), + audio_writer=SoundfileAudioWriter(), + input_dir=args.input_manifest_dir, + params=params.add_mic_noise, + ) + + asyncio.run( + stage.transform( + input_manifest, + conventions.manifest_path(args.output_dir), + ) + ) + + +if __name__ == "__main__": + main() diff --git a/ml/test/pipeline/speech/test_background_noise_stage.py b/ml/test/pipeline/speech/test_background_noise_stage.py new file mode 100644 index 00000000..62ea27cd --- /dev/null +++ b/ml/test/pipeline/speech/test_background_noise_stage.py @@ -0,0 +1,517 @@ +"""Unit tests for BackgroundNoiseAugmentor.""" + +from __future__ import annotations + +import asyncio +import hashlib +import sys +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import soundfile as sf + +sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent)) + +from pipeline.core.manifest import Manifest, ManifestStore +from pipeline.core.randomization import VariationGenerator +from pipeline.core.sample import AudioSample +from pipeline.io.audio_io import AudioData +from pipeline.speech.background_noise_stage import BackgroundNoiseAugmentor, NoiseProvider +from pipeline.stages.params import AddBackgroundNoiseParams + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_audio_sample( + sample_id: str = "TV_ON_Jenny_r100", + transcript: str = "TV_ON", + sample_rate: int = 16000, + duration_s: float = 0.5, + wav_dir: Path | None = None, +) -> AudioSample: + content = f"{sample_id}:audio" + content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() + path = Path(f"{sample_id}.wav") + if wav_dir is not None: + wav_path = wav_dir / path + n_samples = int(sample_rate * duration_s) + data = np.zeros(n_samples, dtype=np.float32) + sf.write(str(wav_path), data, sample_rate, format="WAV", subtype="PCM_16") + return AudioSample( + id=sample_id, + seed=0, + content_hash=content_hash, + path=path, + parent_content_hash="parent_hash", + transcript=transcript, + applied_values={}, + ) + + +class _RecordingAudioReader: + """Stub AudioReader that returns configured durations per path.""" + + def __init__( + self, + sample_rate: int = 16000, + audio_duration_s: float = 0.5, + noise_duration_s: float = 2.0, + ) -> None: + self.calls: list[Path] = [] + self._sample_rate = sample_rate + self._audio_duration_s = audio_duration_s + self._noise_duration_s = noise_duration_s + # If path stem ends with "noise", return noise duration, else audio duration + self._noise_paths: set[Path] = set() + + def register_noise_path(self, path: Path) -> None: + self._noise_paths.add(path) + + async def read(self, path: Path) -> AudioData: + self.calls.append(path) + if path in self._noise_paths or path.parent.name == "noise": + n_samples = int(self._sample_rate * self._noise_duration_s) + else: + n_samples = int(self._sample_rate * self._audio_duration_s) + samples = np.zeros(n_samples, dtype=np.float32) + return AudioData(samples=samples, sample_rate=self._sample_rate) + + +class _RecordingAudioWriter: + """Stub AudioWriter that records write() calls.""" + + def __init__(self) -> None: + self.calls: list[tuple[Path, AudioData]] = [] + + async def write(self, path: Path, audio: AudioData) -> None: + self.calls.append((path, audio)) + sf.write(str(path), audio.samples, audio.sample_rate, format="WAV", subtype="PCM_16") + + +class _FakeNoiseProvider: + """Fake NoiseProvider that returns a fixed list of noise file paths.""" + + def __init__(self, noise_dir: Path, filenames: list[str]) -> None: + self._noise_dir = noise_dir + self._filenames = filenames + + def list_files(self) -> list[Path]: + return [self._noise_dir / name for name in self._filenames] + + +def _make_params( + vary_probability: float = 0.0, + volume_min: float = 0.0, + volume_max: float = 0.3, +) -> AddBackgroundNoiseParams: + return AddBackgroundNoiseParams( + vary_probability=vary_probability, + volume_min=volume_min, + volume_max=volume_max, + ) + + +def _make_stage( + output_dir: Path, + noise_dir: Path | None = None, + *, + audio_reader: _RecordingAudioReader | None = None, + audio_writer: _RecordingAudioWriter | None = None, + input_dir: Path | None = None, + noise_provider: _FakeNoiseProvider | None = None, + params: AddBackgroundNoiseParams | None = None, + vary_probability: float = 0.0, + volume_min: float = 0.0, + volume_max: float = 0.3, + noise_filenames: list[str] | None = None, +) -> tuple[BackgroundNoiseAugmentor, _RecordingAudioReader, _RecordingAudioWriter]: + if audio_reader is None: + audio_reader = _RecordingAudioReader() + if audio_writer is None: + audio_writer = _RecordingAudioWriter() + if input_dir is None: + input_dir = output_dir + if noise_dir is None: + noise_dir = output_dir / "noise" + noise_dir.mkdir(parents=True, exist_ok=True) + if params is None: + params = _make_params( + vary_probability=vary_probability, + volume_min=volume_min, + volume_max=volume_max, + ) + if noise_provider is None: + if noise_filenames is None: + noise_filenames = ["traffic.wav"] + noise_provider = _FakeNoiseProvider(noise_dir, noise_filenames) + stage = BackgroundNoiseAugmentor( + output_dir=output_dir, + manifest_store=ManifestStore(), + audio_reader=audio_reader, + audio_writer=audio_writer, + input_dir=input_dir, + noise_provider=noise_provider, + params=params, + ) + return stage, audio_reader, audio_writer + + +def _write_noise_wav(noise_dir: Path, filename: str, sample_rate: int = 16000, duration_s: float = 2.0) -> Path: + noise_dir.mkdir(parents=True, exist_ok=True) + path = noise_dir / filename + n_samples = int(sample_rate * duration_s) + data = (np.random.default_rng(42).random(n_samples) * 2 - 1).astype(np.float32) + sf.write(str(path), data, sample_rate, format="WAV", subtype="PCM_16") + return path + + +# --------------------------------------------------------------------------- +# TestAppliedValues +# --------------------------------------------------------------------------- + + +class TestAppliedValues: + def test_applied_values_has_noise_file_key(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert "noise_file" in av + + def test_applied_values_has_noise_start_s_key(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert "noise_start_s" in av + + def test_applied_values_has_noise_volume_key(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert "noise_volume" in av + + def test_applied_values_has_exactly_three_keys(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert set(av.keys()) == {"noise_file", "noise_start_s", "noise_volume"} + + def test_noise_file_always_chosen_when_not_applied(self, tmp_path: Path) -> None: + """noise_file is always selected even when should_vary returns False (vary_probability=0).""" + stage, _, _ = _make_stage( + tmp_path, + vary_probability=0.0, + noise_filenames=["rain.wav", "traffic.wav"], + ) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + # noise_file must be one of the provided filenames + assert av["noise_file"] in ["rain.wav", "traffic.wav"] + + def test_noise_file_stored_even_when_volume_is_zero(self, tmp_path: Path) -> None: + """noise_file is stored regardless of whether noise is applied.""" + stage, _, _ = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert av["noise_file"] == "traffic.wav" + assert av["noise_volume"] == 0.0 + + def test_noise_start_s_zero_when_not_applied(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert av["noise_start_s"] == 0.0 + + def test_noise_volume_zero_when_not_applied(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert av["noise_volume"] == 0.0 + + def test_noise_volume_nonzero_when_applied(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage( + tmp_path, + vary_probability=1.0, + volume_min=0.2, + volume_max=0.2, + ) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert av["noise_volume"] > 0.0 + + def test_noise_file_chosen_from_sorted_list(self, tmp_path: Path) -> None: + """choose() operates on sorted filenames for OS-independent determinism.""" + filenames = ["z_noise.wav", "a_noise.wav", "m_noise.wav"] + stage, _, _ = _make_stage(tmp_path, noise_filenames=filenames) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(0)) + # Must be one of the provided filenames + assert av["noise_file"] in filenames + + def test_noise_start_s_stored_as_float(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert isinstance(av["noise_start_s"], float) + + def test_noise_volume_stored_as_float(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert isinstance(av["noise_volume"], float) + + +# --------------------------------------------------------------------------- +# TestAppliedValuesNoiseShorterThanAudio +# --------------------------------------------------------------------------- + + +class TestNoiseShorterThanAudio: + def test_noise_start_s_clamped_to_zero_when_noise_shorter_than_audio( + self, tmp_path: Path + ) -> None: + """When noise file is shorter than audio, max_start_s < 0 → clamped to 0.0.""" + # noise is 0.1s, audio is 0.5s → max_start_s = -0.4s → clamp to 0.0 + reader = _RecordingAudioReader( + sample_rate=16000, + audio_duration_s=0.5, + noise_duration_s=0.1, + ) + noise_dir = tmp_path / "noise" + noise_dir.mkdir() + reader.register_noise_path(noise_dir / "short.wav") + + stage, _, _ = _make_stage( + tmp_path, + noise_dir=noise_dir, + audio_reader=reader, + vary_probability=1.0, + volume_min=0.1, + volume_max=0.1, + noise_filenames=["short.wav"], + ) + sample = _make_audio_sample(wav_dir=tmp_path) + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert av["noise_start_s"] == 0.0 + + +# --------------------------------------------------------------------------- +# TestDeriveId +# --------------------------------------------------------------------------- + + +class TestDeriveId: + def test_derive_id_format_with_noise_applied(self, tmp_path: Path) -> None: + """Format: {input.id}_{noise_filestem}_v{int(volume*100)}""" + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(sample_id="TV_ON_Jenny_r100") + result = stage._derive_id(sample, {"noise_file": "traffic.wav", "noise_start_s": 0.5, "noise_volume": 0.25}) + assert result == "TV_ON_Jenny_r100_traffic_v25" + + def test_derive_id_format_with_zero_volume(self, tmp_path: Path) -> None: + """Even when volume=0 (not applied), the filestem is still in the id.""" + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(sample_id="TV_ON_Jenny_r100") + result = stage._derive_id(sample, {"noise_file": "rain.wav", "noise_start_s": 0.0, "noise_volume": 0.0}) + assert result == "TV_ON_Jenny_r100_rain_v0" + + def test_derive_id_uses_stem_not_full_filename(self, tmp_path: Path) -> None: + """noise_filestem = Path(noise_file).stem — extension excluded.""" + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(sample_id="X") + result = stage._derive_id(sample, {"noise_file": "background_noise.wav", "noise_start_s": 0.0, "noise_volume": 0.1}) + assert "background_noise" in result + assert ".wav" not in result + + def test_derive_id_volume_int_conversion(self, tmp_path: Path) -> None: + """int(volume * 100): 0.3 → 30, 0.25 → 25.""" + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(sample_id="S") + result = stage._derive_id(sample, {"noise_file": "t.wav", "noise_start_s": 0.0, "noise_volume": 0.3}) + assert result == "S_t_v30" + + def test_derive_id_uses_input_id_as_prefix(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(sample_id="VOLUME_UP_Aria_r110") + result = stage._derive_id(sample, {"noise_file": "cafe.wav", "noise_start_s": 0.0, "noise_volume": 0.0}) + assert result.startswith("VOLUME_UP_Aria_r110_") + + +# --------------------------------------------------------------------------- +# TestGenerateOutput +# --------------------------------------------------------------------------- + + +class TestGenerateOutput: + def test_generate_output_calls_audio_reader_for_input_and_noise(self, tmp_path: Path) -> None: + """When noise is applied (vary_probability=1.0), reader is called for both noise and audio.""" + noise_dir = tmp_path / "noise" + _write_noise_wav(noise_dir, "traffic.wav") + + reader = _RecordingAudioReader(audio_duration_s=0.5, noise_duration_s=2.0) + reader.register_noise_path(noise_dir / "traffic.wav") + + stage, _, _ = _make_stage( + tmp_path, + noise_dir=noise_dir, + audio_reader=reader, + vary_probability=1.0, + volume_min=0.1, + volume_max=0.1, + noise_filenames=["traffic.wav"], + ) + sample = _make_audio_sample(wav_dir=tmp_path) + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + # Should have called reader at least 3 times: + # - noise file duration (in _get_applied_values) + # - audio duration (in _get_applied_values) + # - audio file (in _generate_output) + # - noise file (in _generate_output) + assert len(reader.calls) >= 3 + + def test_generate_output_calls_audio_writer(self, tmp_path: Path) -> None: + noise_dir = tmp_path / "noise" + _write_noise_wav(noise_dir, "traffic.wav") + + reader = _RecordingAudioReader(audio_duration_s=0.5, noise_duration_s=2.0) + reader.register_noise_path(noise_dir / "traffic.wav") + + stage, _, writer = _make_stage( + tmp_path, + noise_dir=noise_dir, + audio_reader=reader, + noise_filenames=["traffic.wav"], + ) + sample = _make_audio_sample(wav_dir=tmp_path) + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert len(writer.calls) == 1 + + def test_output_audio_same_length_as_input(self, tmp_path: Path) -> None: + """Background noise does not change the length of the audio.""" + sample_rate = 16000 + duration_s = 0.5 + noise_dir = tmp_path / "noise" + _write_noise_wav(noise_dir, "traffic.wav", sample_rate=sample_rate, duration_s=2.0) + + reader = _RecordingAudioReader( + sample_rate=sample_rate, audio_duration_s=duration_s, noise_duration_s=2.0 + ) + reader.register_noise_path(noise_dir / "traffic.wav") + + stage, _, writer = _make_stage( + tmp_path, + noise_dir=noise_dir, + audio_reader=reader, + vary_probability=1.0, + volume_min=0.1, + volume_max=0.1, + noise_filenames=["traffic.wav"], + ) + sample = _make_audio_sample(wav_dir=tmp_path, sample_rate=sample_rate, duration_s=duration_s) + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + _path, written_audio = writer.calls[0] + assert len(written_audio.samples) == int(sample_rate * duration_s) + + def test_output_audio_unchanged_when_volume_is_zero(self, tmp_path: Path) -> None: + """When noise not applied (volume=0), output samples equal input samples.""" + sample_rate = 16000 + duration_s = 0.1 + noise_dir = tmp_path / "noise" + _write_noise_wav(noise_dir, "traffic.wav", sample_rate=sample_rate, duration_s=2.0) + + # Use a special reader that returns a known nonzero signal for audio + class _ConstantReader: + async def read(self, path: Path) -> AudioData: + n = int(sample_rate * duration_s) + if path.parent.name == "noise": + samples = np.ones(int(sample_rate * 2.0), dtype=np.float32) * 0.9 + else: + samples = np.full(n, 0.5, dtype=np.float32) + return AudioData(samples=samples, sample_rate=sample_rate) + + writer = _RecordingAudioWriter() + stage = BackgroundNoiseAugmentor( + output_dir=tmp_path, + manifest_store=ManifestStore(), + audio_reader=_ConstantReader(), + audio_writer=writer, + input_dir=tmp_path, + noise_provider=_FakeNoiseProvider(noise_dir, ["traffic.wav"]), + params=_make_params(vary_probability=0.0), + ) + sample = _make_audio_sample(wav_dir=tmp_path, sample_rate=sample_rate, duration_s=duration_s) + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + _path, written_audio = writer.calls[0] + # All samples should still be 0.5 (unchanged) + assert np.allclose(written_audio.samples, 0.5) + + def test_transcript_preserved_in_output_sample(self, tmp_path: Path) -> None: + noise_dir = tmp_path / "noise" + _write_noise_wav(noise_dir, "traffic.wav") + + reader = _RecordingAudioReader(audio_duration_s=0.5, noise_duration_s=2.0) + reader.register_noise_path(noise_dir / "traffic.wav") + + stage, _, _ = _make_stage( + tmp_path, + noise_dir=noise_dir, + audio_reader=reader, + noise_filenames=["traffic.wav"], + ) + sample = _make_audio_sample(transcript="VOLUME_UP", wav_dir=tmp_path) + result = asyncio.run( + stage.transform(Manifest([sample]), tmp_path / "manifest.json") + ) + assert result.samples[0].transcript == "VOLUME_UP" + + +# --------------------------------------------------------------------------- +# TestSkipPath +# --------------------------------------------------------------------------- + + +class TestSkipPath: + def test_skip_path_does_not_call_audio_writer_on_second_run(self, tmp_path: Path) -> None: + noise_dir = tmp_path / "noise" + _write_noise_wav(noise_dir, "traffic.wav") + + reader = _RecordingAudioReader(audio_duration_s=0.5, noise_duration_s=2.0) + reader.register_noise_path(noise_dir / "traffic.wav") + + stage, _, writer = _make_stage( + tmp_path, + noise_dir=noise_dir, + audio_reader=reader, + noise_filenames=["traffic.wav"], + ) + sample = _make_audio_sample(wav_dir=tmp_path) + + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + initial_count = len(writer.calls) + assert initial_count == 1 + + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert len(writer.calls) == initial_count + + def test_skip_path_preserves_output_sample_id(self, tmp_path: Path) -> None: + noise_dir = tmp_path / "noise" + _write_noise_wav(noise_dir, "traffic.wav") + + reader = _RecordingAudioReader(audio_duration_s=0.5, noise_duration_s=2.0) + reader.register_noise_path(noise_dir / "traffic.wav") + + stage, _, _ = _make_stage( + tmp_path, + noise_dir=noise_dir, + audio_reader=reader, + noise_filenames=["traffic.wav"], + ) + sample = _make_audio_sample(wav_dir=tmp_path) + + result1 = asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + result2 = asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert result2.samples[0].id == result1.samples[0].id diff --git a/ml/test/pipeline/speech/test_mic_noise_stage.py b/ml/test/pipeline/speech/test_mic_noise_stage.py new file mode 100644 index 00000000..dd06e644 --- /dev/null +++ b/ml/test/pipeline/speech/test_mic_noise_stage.py @@ -0,0 +1,330 @@ +"""Unit tests for MicrophoneNoiseAugmentor.""" + +from __future__ import annotations + +import asyncio +import hashlib +import sys +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import soundfile as sf + +sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent)) + +from pipeline.core.manifest import Manifest, ManifestStore +from pipeline.core.randomization import VariationGenerator +from pipeline.core.sample import AudioSample +from pipeline.io.audio_io import AudioData +from pipeline.speech.mic_noise_stage import MicrophoneNoiseAugmentor +from pipeline.stages.params import AddMicNoiseParams + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_audio_sample( + sample_id: str = "TV_ON_Jenny_r100", + transcript: str = "TV_ON", + sample_rate: int = 16000, + duration_s: float = 0.1, + wav_dir: Path | None = None, +) -> AudioSample: + content = f"{sample_id}:audio" + content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() + path = Path(f"{sample_id}.wav") + if wav_dir is not None: + wav_path = wav_dir / path + n_samples = int(sample_rate * duration_s) + data = np.zeros(n_samples, dtype=np.float32) + sf.write(str(wav_path), data, sample_rate, format="WAV", subtype="PCM_16") + return AudioSample( + id=sample_id, + seed=0, + content_hash=content_hash, + path=path, + parent_content_hash="parent_hash", + transcript=transcript, + applied_values={}, + ) + + +class _RecordingAudioReader: + """Stub AudioReader that records read() calls and returns silence.""" + + def __init__(self, sample_rate: int = 16000, duration_s: float = 0.1) -> None: + self.calls: list[Path] = [] + self._sample_rate = sample_rate + self._duration_s = duration_s + + async def read(self, path: Path) -> AudioData: + self.calls.append(path) + n_samples = int(self._sample_rate * self._duration_s) + samples = np.zeros(n_samples, dtype=np.float32) + return AudioData(samples=samples, sample_rate=self._sample_rate) + + +class _RecordingAudioWriter: + """Stub AudioWriter that records write() calls.""" + + def __init__(self) -> None: + self.calls: list[tuple[Path, AudioData]] = [] + + async def write(self, path: Path, audio: AudioData) -> None: + self.calls.append((path, audio)) + sf.write(str(path), audio.samples, audio.sample_rate, format="WAV", subtype="PCM_16") + + +def _make_params( + vary_probability: float = 0.0, + amplitude_min: float = 0.001, + amplitude_max: float = 0.05, +) -> AddMicNoiseParams: + return AddMicNoiseParams( + vary_probability=vary_probability, + amplitude_min=amplitude_min, + amplitude_max=amplitude_max, + ) + + +def _make_stage( + output_dir: Path, + *, + audio_reader: _RecordingAudioReader | None = None, + audio_writer: _RecordingAudioWriter | None = None, + input_dir: Path | None = None, + params: AddMicNoiseParams | None = None, + vary_probability: float = 0.0, + amplitude_min: float = 0.001, + amplitude_max: float = 0.05, +) -> tuple[MicrophoneNoiseAugmentor, _RecordingAudioReader, _RecordingAudioWriter]: + if audio_reader is None: + audio_reader = _RecordingAudioReader() + if audio_writer is None: + audio_writer = _RecordingAudioWriter() + if input_dir is None: + input_dir = output_dir + if params is None: + params = _make_params( + vary_probability=vary_probability, + amplitude_min=amplitude_min, + amplitude_max=amplitude_max, + ) + stage = MicrophoneNoiseAugmentor( + output_dir=output_dir, + manifest_store=ManifestStore(), + audio_reader=audio_reader, + audio_writer=audio_writer, + input_dir=input_dir, + params=params, + ) + return stage, audio_reader, audio_writer + + +# --------------------------------------------------------------------------- +# TestAppliedValues +# --------------------------------------------------------------------------- + + +class TestAppliedValues: + def test_applied_values_has_mic_noise_amplitude_key(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert "mic_noise_amplitude" in av + + def test_applied_values_has_exactly_one_key(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert set(av.keys()) == {"mic_noise_amplitude"} + + def test_amplitude_zero_when_not_applied(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path, vary_probability=0.0, amplitude_min=0.01, amplitude_max=0.05) + sample = _make_audio_sample() + for seed in range(10): + av = stage._get_applied_values(sample, VariationGenerator(seed)) + assert av["mic_noise_amplitude"] == 0.0 + + def test_amplitude_nonzero_when_applied(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage( + tmp_path, vary_probability=1.0, amplitude_min=0.01, amplitude_max=0.01 + ) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert av["mic_noise_amplitude"] > 0.0 + + def test_amplitude_stored_as_float(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert isinstance(av["mic_noise_amplitude"], float) + + def test_amplitude_zero_stored_as_float_not_int(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample() + av = stage._get_applied_values(sample, VariationGenerator(42)) + assert av["mic_noise_amplitude"] == 0.0 + assert isinstance(av["mic_noise_amplitude"], float) + + +# --------------------------------------------------------------------------- +# TestDeriveId +# --------------------------------------------------------------------------- + + +class TestDeriveId: + def test_derive_id_format_with_noise_applied(self, tmp_path: Path) -> None: + """Format: {input.id}_mic{int(amplitude*1000)}""" + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(sample_id="TV_ON_Jenny_r100") + result = stage._derive_id(sample, {"mic_noise_amplitude": 0.025}) + assert result == "TV_ON_Jenny_r100_mic25" + + def test_derive_id_format_with_zero_amplitude(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(sample_id="TV_ON_Jenny_r100") + result = stage._derive_id(sample, {"mic_noise_amplitude": 0.0}) + assert result == "TV_ON_Jenny_r100_mic0" + + def test_derive_id_uses_input_id_as_prefix(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(sample_id="VOLUME_UP_Aria_r110") + result = stage._derive_id(sample, {"mic_noise_amplitude": 0.0}) + assert result.startswith("VOLUME_UP_Aria_r110_") + + def test_derive_id_amplitude_int_conversion(self, tmp_path: Path) -> None: + """int(amplitude * 1000): 0.05 → 50, 0.001 → 1.""" + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(sample_id="S") + result = stage._derive_id(sample, {"mic_noise_amplitude": 0.05}) + assert result == "S_mic50" + + def test_derive_id_always_has_mic_suffix(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(sample_id="X") + result = stage._derive_id(sample, {"mic_noise_amplitude": 0.01}) + assert "_mic" in result + + +# --------------------------------------------------------------------------- +# TestGenerateOutput +# --------------------------------------------------------------------------- + + +class TestGenerateOutput: + def test_generate_output_calls_audio_reader(self, tmp_path: Path) -> None: + stage, reader, _ = _make_stage(tmp_path) + sample = _make_audio_sample(wav_dir=tmp_path) + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert len(reader.calls) == 1 + + def test_generate_output_calls_audio_writer(self, tmp_path: Path) -> None: + stage, _, writer = _make_stage(tmp_path) + sample = _make_audio_sample(wav_dir=tmp_path) + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert len(writer.calls) == 1 + + def test_output_audio_same_length_as_input(self, tmp_path: Path) -> None: + """Mic noise does not change the length of the audio.""" + sample_rate = 16000 + duration_s = 0.1 + stage, _, writer = _make_stage( + tmp_path, + audio_reader=_RecordingAudioReader(sample_rate=sample_rate, duration_s=duration_s), + vary_probability=1.0, + amplitude_min=0.01, + amplitude_max=0.01, + ) + sample = _make_audio_sample(wav_dir=tmp_path, sample_rate=sample_rate, duration_s=duration_s) + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + _path, written_audio = writer.calls[0] + assert len(written_audio.samples) == int(sample_rate * duration_s) + + def test_output_audio_unchanged_when_amplitude_is_zero(self, tmp_path: Path) -> None: + """When noise not applied (amplitude=0), output samples equal input samples.""" + sample_rate = 16000 + duration_s = 0.1 + + class _ConstantReader: + async def read(self, path: Path) -> AudioData: + n = int(sample_rate * duration_s) + return AudioData(samples=np.full(n, 0.5, dtype=np.float32), sample_rate=sample_rate) + + writer = _RecordingAudioWriter() + stage = MicrophoneNoiseAugmentor( + output_dir=tmp_path, + manifest_store=ManifestStore(), + audio_reader=_ConstantReader(), + audio_writer=writer, + input_dir=tmp_path, + params=_make_params(vary_probability=0.0), + ) + sample = _make_audio_sample(wav_dir=tmp_path, sample_rate=sample_rate, duration_s=duration_s) + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + _path, written_audio = writer.calls[0] + assert np.allclose(written_audio.samples, 0.5) + + def test_gaussian_noise_added_when_amplitude_nonzero(self, tmp_path: Path) -> None: + """When amplitude > 0, output should differ from input (noise was added).""" + sample_rate = 16000 + duration_s = 0.5 # Long enough for statistical significance + amplitude = 0.1 + + class _ZeroReader: + async def read(self, path: Path) -> AudioData: + n = int(sample_rate * duration_s) + return AudioData(samples=np.zeros(n, dtype=np.float32), sample_rate=sample_rate) + + writer = _RecordingAudioWriter() + stage = MicrophoneNoiseAugmentor( + output_dir=tmp_path, + manifest_store=ManifestStore(), + audio_reader=_ZeroReader(), + audio_writer=writer, + input_dir=tmp_path, + params=_make_params(vary_probability=1.0, amplitude_min=amplitude, amplitude_max=amplitude), + ) + sample = _make_audio_sample(wav_dir=tmp_path, sample_rate=sample_rate, duration_s=duration_s) + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + _path, written_audio = writer.calls[0] + # Noise was added — not all zeros anymore + assert not np.all(written_audio.samples == 0.0) + + def test_transcript_preserved_in_output_sample(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(transcript="CHANNEL_UP", wav_dir=tmp_path) + result = asyncio.run( + stage.transform(Manifest([sample]), tmp_path / "manifest.json") + ) + assert result.samples[0].transcript == "CHANNEL_UP" + + +# --------------------------------------------------------------------------- +# TestSkipPath +# --------------------------------------------------------------------------- + + +class TestSkipPath: + def test_skip_path_does_not_call_audio_writer_on_second_run(self, tmp_path: Path) -> None: + stage, _, writer = _make_stage(tmp_path) + sample = _make_audio_sample(wav_dir=tmp_path) + + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + initial_count = len(writer.calls) + assert initial_count == 1 + + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert len(writer.calls) == initial_count + + def test_skip_path_preserves_output_sample_id(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(wav_dir=tmp_path) + + result1 = asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + result2 = asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert result2.samples[0].id == result1.samples[0].id diff --git a/ml/test/pipeline/stages/test_params.py b/ml/test/pipeline/stages/test_params.py index 18bcef67..3a4cce92 100644 --- a/ml/test/pipeline/stages/test_params.py +++ b/ml/test/pipeline/stages/test_params.py @@ -10,7 +10,14 @@ sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent)) -from pipeline.stages.params import AddDelaysParams, GeneratePhraseParams, GenerateSamplesParams, PipelineParams +from pipeline.stages.params import ( + AddBackgroundNoiseParams, + AddDelaysParams, + AddMicNoiseParams, + GeneratePhraseParams, + GenerateSamplesParams, + PipelineParams, +) def _write_params(path: Path, data: dict) -> Path: @@ -44,6 +51,16 @@ def _write_params(path: Path, data: dict) -> Path: "suffix_min_s": 0.02, "suffix_max_s": 0.15, }, + "add_background_noise": { + "vary_probability": 0.5, + "volume_min": 0.0, + "volume_max": 0.3, + }, + "add_mic_noise": { + "vary_probability": 0.5, + "amplitude_min": 0.001, + "amplitude_max": 0.05, + }, }, } @@ -153,3 +170,71 @@ def test_add_delays_fields_are_floats(self, tmp_path: Path) -> None: assert isinstance(ad.suffix_vary_probability, float) assert isinstance(ad.suffix_min_s, float) assert isinstance(ad.suffix_max_s, float) + + +class TestAddBackgroundNoiseParamsLoad: + def test_loads_add_background_noise_fields(self, tmp_path: Path) -> None: + params_file = _write_params(tmp_path, _VALID_DATA) + params = PipelineParams.load(params_file) + + bn = params.add_background_noise + assert bn.vary_probability == pytest.approx(0.5) + assert bn.volume_min == pytest.approx(0.0) + assert bn.volume_max == pytest.approx(0.3) + + def test_add_background_noise_is_correct_type(self, tmp_path: Path) -> None: + params_file = _write_params(tmp_path, _VALID_DATA) + params = PipelineParams.load(params_file) + + assert isinstance(params.add_background_noise, AddBackgroundNoiseParams) + + def test_add_background_noise_fields_are_floats(self, tmp_path: Path) -> None: + params_file = _write_params(tmp_path, _VALID_DATA) + params = PipelineParams.load(params_file) + + bn = params.add_background_noise + assert isinstance(bn.vary_probability, float) + assert isinstance(bn.volume_min, float) + assert isinstance(bn.volume_max, float) + + def test_missing_add_background_noise_stage_raises(self, tmp_path: Path) -> None: + import copy + data = copy.deepcopy(_VALID_DATA) + del data["stages"]["add_background_noise"] + params_file = _write_params(tmp_path, data) + with pytest.raises(KeyError): + PipelineParams.load(params_file) + + +class TestAddMicNoiseParamsLoad: + def test_loads_add_mic_noise_fields(self, tmp_path: Path) -> None: + params_file = _write_params(tmp_path, _VALID_DATA) + params = PipelineParams.load(params_file) + + mn = params.add_mic_noise + assert mn.vary_probability == pytest.approx(0.5) + assert mn.amplitude_min == pytest.approx(0.001) + assert mn.amplitude_max == pytest.approx(0.05) + + def test_add_mic_noise_is_correct_type(self, tmp_path: Path) -> None: + params_file = _write_params(tmp_path, _VALID_DATA) + params = PipelineParams.load(params_file) + + assert isinstance(params.add_mic_noise, AddMicNoiseParams) + + def test_add_mic_noise_fields_are_floats(self, tmp_path: Path) -> None: + params_file = _write_params(tmp_path, _VALID_DATA) + params = PipelineParams.load(params_file) + + mn = params.add_mic_noise + assert isinstance(mn.vary_probability, float) + assert isinstance(mn.amplitude_min, float) + assert isinstance(mn.amplitude_max, float) + + def test_missing_add_mic_noise_stage_raises(self, tmp_path: Path) -> None: + import copy + data = copy.deepcopy(_VALID_DATA) + del data["stages"]["add_mic_noise"] + params_file = _write_params(tmp_path, data) + with pytest.raises(KeyError): + PipelineParams.load(params_file) From fdfa5b73a782c277f24ac461ead0caed0031046d Mon Sep 17 00:00:00 2001 From: Joe Davis Date: Wed, 17 Jun 2026 19:10:36 -0700 Subject: [PATCH 2/7] ADR-227: Fix executor.submit coroutine form; store _noise_dir to remove redundant list_files() calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Issue 1: Changed executor.submit(asyncio.run, _read()) to executor.submit(lambda: asyncio.run(_read())) so the coroutine is created inside the worker thread. The old form passed a pre-created coroutine object across threads, which is fragile (the coroutine is bound to the originating thread's context). The lambda form creates a fresh coroutine in the worker thread per asyncio.run() semantics. Issues 2 & 3: Added noise_dir: Path constructor parameter to BackgroundNoiseAugmentor and stored it as _noise_dir. Both _get_applied_values (line 106) and _generate_output (line 151) previously called self._noise_provider.list_files()[0].parent to recover the noise directory — re-globbing the filesystem on each call and raising an opaque IndexError on an empty directory. _noise_dir eliminates both calls. --- ml/pipeline/speech/background_noise_stage.py | 46 +++++++++++++++----- 1 file changed, 35 insertions(+), 11 deletions(-) diff --git a/ml/pipeline/speech/background_noise_stage.py b/ml/pipeline/speech/background_noise_stage.py index d6fd14d6..96f876d1 100644 --- a/ml/pipeline/speech/background_noise_stage.py +++ b/ml/pipeline/speech/background_noise_stage.py @@ -9,6 +9,14 @@ audio sample file are read to compute max_start_s = noise_duration_s - audio_duration_s. When the noise file is shorter than the audio sample, max_start_s is negative and clamped to 0.0 (min and max both 0.0). + +_noise_dir is stored directly in the constructor and used to resolve noise file +paths, eliminating redundant list_files() calls and the IndexError risk from an +empty noise directory. + +_read_duration_s runs the async read in a thread executor when called from a running +event loop, using the lambda form (executor.submit(lambda: asyncio.run(...))) to avoid +passing a pre-created coroutine to asyncio.run in another thread. """ from __future__ import annotations @@ -42,22 +50,29 @@ def list_files(self) -> list[Path]: ... def _read_duration_s(audio_reader: AudioReader, path: Path) -> float: """Read a file and return its duration in seconds. - Handles both a running event loop (uses a thread executor) and no loop - (uses asyncio.run()). + When called outside a running event loop, uses asyncio.run() directly. + When called from within a running event loop (e.g. from a synchronous + method invoked by an async caller), runs asyncio.run() in a thread executor + so the loop thread is not blocked. Uses the lambda form + (executor.submit(lambda: asyncio.run(_read()))) to create the coroutine + inside the worker thread, avoiding passing a pre-created coroutine across + threads. """ async def _read() -> float: data = await audio_reader.read(path) return len(data.samples) / data.sample_rate try: - loop = asyncio.get_running_loop() + asyncio.get_running_loop() except RuntimeError: - # No running loop — safe to call asyncio.run() + # No running loop — safe to call asyncio.run() directly return asyncio.run(_read()) - # Inside a running loop — run in a thread to avoid blocking the loop + # Inside a running loop — run in a thread to avoid blocking the loop. + # The lambda creates the coroutine inside the worker thread so a fresh + # event loop is used per asyncio.run() call. with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(asyncio.run, _read()) + future = executor.submit(lambda: asyncio.run(_read())) return future.result() @@ -74,6 +89,10 @@ class BackgroundNoiseAugmentor(ModifierStage[AudioSample, AudioSample]): input_dir is the directory containing the input WAV files referenced by AudioSample.path. This is typically the output directory of the preceding stage. + + noise_dir is the directory containing noise WAV files. Stored directly rather + than derived from NoiseProvider.list_files()[0].parent to avoid redundant + filesystem calls and IndexError on an empty directory. """ def __init__( @@ -83,6 +102,7 @@ def __init__( audio_reader: AudioReader, audio_writer: AudioWriter, input_dir: Path, + noise_dir: Path, noise_provider: NoiseProvider, params: AddBackgroundNoiseParams, ) -> None: @@ -90,6 +110,7 @@ def __init__( self._audio_reader = audio_reader self._audio_writer = audio_writer self._input_dir = input_dir + self._noise_dir = noise_dir self._noise_provider = noise_provider self._vary_probability = params.vary_probability self._volume_filter = MinMaxFilter(params.volume_min, params.volume_max, precision=2) @@ -97,13 +118,17 @@ def __init__( def _get_applied_values( self, sample: AudioSample, generator: VariationGenerator ) -> dict[str, Any]: - # Always choose a noise file for hash stability + # Always choose a noise file for hash stability. + # list_files() is called once; _noise_dir resolves the full path. noise_files = sorted([p.name for p in self._noise_provider.list_files()]) noise_file: str = generator.choose("noise_file", noise_files) if generator.should_vary("noise", self._vary_probability): - # Derive noise_start_s bounds from file durations - noise_path = self._noise_provider.list_files()[0].parent / noise_file + # noise_start_s bounds are derived from file durations. Both files are read + # synchronously here because _get_applied_values is a synchronous method. + # The durations are needed to clamp the start so the noise slice contains + # actual noise data rather than silence from an out-of-range offset. + noise_path = self._noise_dir / noise_file audio_path = self._input_dir / sample.path noise_duration_s = _read_duration_s(self._audio_reader, noise_path) @@ -147,8 +172,7 @@ async def _generate_output( output_samples = audio.samples.copy() if noise_volume > 0.0: - # Resolve noise file path via provider - noise_path = self._noise_provider.list_files()[0].parent / noise_file + noise_path = self._noise_dir / noise_file noise_audio = await self._audio_reader.read(noise_path) start_sample = int(noise_start_s * noise_audio.sample_rate) From c729bd32df2973d748da962e7e7e98ce6c796e48 Mon Sep 17 00:00:00 2001 From: Joe Davis Date: Wed, 17 Jun 2026 19:10:42 -0700 Subject: [PATCH 3/7] ADR-227: Remove unused imports from test_background_noise_stage.py and test_mic_noise_stage.py Issues 4 & 5: 'from typing import Any' and 'import pytest' were imported but never referenced in either test file. Also updates test_background_noise_stage.py to pass the new noise_dir parameter required by the BackgroundNoiseAugmentor constructor. --- ml/test/pipeline/speech/test_background_noise_stage.py | 4 ++-- ml/test/pipeline/speech/test_mic_noise_stage.py | 2 -- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/ml/test/pipeline/speech/test_background_noise_stage.py b/ml/test/pipeline/speech/test_background_noise_stage.py index 62ea27cd..7baee819 100644 --- a/ml/test/pipeline/speech/test_background_noise_stage.py +++ b/ml/test/pipeline/speech/test_background_noise_stage.py @@ -6,10 +6,8 @@ import hashlib import sys from pathlib import Path -from typing import Any import numpy as np -import pytest import soundfile as sf sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent)) @@ -155,6 +153,7 @@ def _make_stage( audio_reader=audio_reader, audio_writer=audio_writer, input_dir=input_dir, + noise_dir=noise_dir, noise_provider=noise_provider, params=params, ) @@ -440,6 +439,7 @@ async def read(self, path: Path) -> AudioData: audio_reader=_ConstantReader(), audio_writer=writer, input_dir=tmp_path, + noise_dir=noise_dir, noise_provider=_FakeNoiseProvider(noise_dir, ["traffic.wav"]), params=_make_params(vary_probability=0.0), ) diff --git a/ml/test/pipeline/speech/test_mic_noise_stage.py b/ml/test/pipeline/speech/test_mic_noise_stage.py index dd06e644..0e2ecb4f 100644 --- a/ml/test/pipeline/speech/test_mic_noise_stage.py +++ b/ml/test/pipeline/speech/test_mic_noise_stage.py @@ -6,10 +6,8 @@ import hashlib import sys from pathlib import Path -from typing import Any import numpy as np -import pytest import soundfile as sf sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent)) From 905e7f365d14cc55a2d8785bb20837fc824b2fe0 Mon Sep 17 00:00:00 2001 From: Joe Davis Date: Wed, 17 Jun 2026 19:10:49 -0700 Subject: [PATCH 4/7] ADR-227: Pass noise_dir to BackgroundNoiseAugmentor; add validation for empty noise directory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Issue 6: _DirectoryNoiseProvider.list_files() now raises ValueError with the directory path if no WAV files are found. Previously it returned [] which caused VariationGenerator.choose() to raise ValueError("options must be non-empty") — an opaque message that didn't identify the misconfigured path. The new guard surfaces the root cause immediately. Also passes args.noise_dir as the new noise_dir constructor parameter added to BackgroundNoiseAugmentor. --- ml/pipeline/stages/speech_05_add_background_noise.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/ml/pipeline/stages/speech_05_add_background_noise.py b/ml/pipeline/stages/speech_05_add_background_noise.py index c833468e..3dd91b84 100644 --- a/ml/pipeline/stages/speech_05_add_background_noise.py +++ b/ml/pipeline/stages/speech_05_add_background_noise.py @@ -21,13 +21,21 @@ class _DirectoryNoiseProvider: Lists all WAV files in the given directory. DVC wiring points this at the appropriate data/ path; this class requires no knowledge of the DVC layout. + + Raises ValueError with the directory path if the directory contains no WAV + files, so that a misconfigured --noise-dir produces an actionable error + rather than the opaque ValueError("options must be non-empty") from + VariationGenerator.choose(). """ def __init__(self, noise_dir: Path) -> None: self._noise_dir = noise_dir def list_files(self) -> list[Path]: - return list(self._noise_dir.glob("*.wav")) + files = list(self._noise_dir.glob("*.wav")) + if not files: + raise ValueError(f"No WAV files found in noise directory: {self._noise_dir}") + return files def main() -> None: @@ -52,6 +60,7 @@ def main() -> None: audio_reader=LibrosaAudioReader(), audio_writer=SoundfileAudioWriter(), input_dir=args.input_manifest_dir, + noise_dir=args.noise_dir, noise_provider=_DirectoryNoiseProvider(args.noise_dir), params=params.add_background_noise, ) From 0dc95ecfb9d1cd1f7864808e2a81f7b6a0a81dee Mon Sep 17 00:00:00 2001 From: Joe Davis Date: Thu, 18 Jun 2026 06:38:40 -0700 Subject: [PATCH 5/7] ADR-227: Remove noise_start_s variation and _read_duration_s MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per user decision (PRRT_kwDOQYcoLM6KbINz, PRRT_kwDOQYcoLM6KbIOC, PRRT_kwDOQYcoLM6KjMkD): noise_start_s is no longer varied — it is always 0.0. This eliminates the hacky synchronous _read_duration_s helper that called asyncio.run inside a ThreadPoolExecutor to avoid blocking the event loop. noise_start_s is still stored in applied_values for hash stability. Update _doc_speech.md to remove the now-obsolete bounds-derivation design note. --- ml/pipeline/speech/_doc_speech.md | 14 +-- ml/pipeline/speech/background_noise_stage.py | 96 +++---------------- .../speech/test_background_noise_stage.py | 35 ++----- 3 files changed, 26 insertions(+), 119 deletions(-) diff --git a/ml/pipeline/speech/_doc_speech.md b/ml/pipeline/speech/_doc_speech.md index 9b27553c..e548c093 100644 --- a/ml/pipeline/speech/_doc_speech.md +++ b/ml/pipeline/speech/_doc_speech.md @@ -43,15 +43,11 @@ Consistent with `TtsProvider` in `tts_stage.py` — the protocol lives in the sa module as the stage that uses it. Unit tests supply `_FakeNoiseProvider`; the entry-point supplies `_DirectoryNoiseProvider` (globbing `*.wav` from `--noise-dir`). -**`noise_file` is always chosen (hash stability), `noise_start_s`/`noise_volume` are -0.0 when not applied.** `VariationGenerator.choose()` runs on the sorted filename list -before the `should_vary` check so the content hash does not change if `vary_probability` -is toggled. All three keys are always present in `applied_values`. - -**`noise_start_s` bounds are derived from file durations at runtime.** Both the noise -file and the audio sample file are read in `_get_applied_values` when noise is applied. -`max_start_s = noise_duration_s - audio_duration_s`; clamped to 0.0 when negative -(i.e. noise file is shorter than audio). +**`noise_file` is always chosen (hash stability), `noise_volume` is 0.0 when not applied.** +`VariationGenerator.choose()` runs on the sorted filename list before the `should_vary` +check so the content hash does not change if `vary_probability` is toggled. All three +keys (`noise_file`, `noise_start_s`, `noise_volume`) are always present in `applied_values`. +`noise_start_s` is always 0.0 — noise is mixed from the beginning of the noise file. **Gaussian noise in `MicrophoneNoiseAugmentor` is seeded from `output_seed`.** Uses `np.random.default_rng(output_seed).normal(0, amplitude, len(samples))` for diff --git a/ml/pipeline/speech/background_noise_stage.py b/ml/pipeline/speech/background_noise_stage.py index 96f876d1..ca212106 100644 --- a/ml/pipeline/speech/background_noise_stage.py +++ b/ml/pipeline/speech/background_noise_stage.py @@ -1,28 +1,7 @@ -"""BackgroundNoiseAugmentor: mix environmental noise into WAV audio samples. - -noise_file is always chosen (choose() always called) for hash stability, even when -not applied. noise_start_s and noise_volume are stored as 0.0 when should_vary -returns False. All three keys (noise_file, noise_start_s, noise_volume) are always -present in every applied_values dict. - -noise_start_s bounds are derived from file durations: both the noise file and the -audio sample file are read to compute max_start_s = noise_duration_s - audio_duration_s. -When the noise file is shorter than the audio sample, max_start_s is negative and -clamped to 0.0 (min and max both 0.0). - -_noise_dir is stored directly in the constructor and used to resolve noise file -paths, eliminating redundant list_files() calls and the IndexError risk from an -empty noise directory. - -_read_duration_s runs the async read in a thread executor when called from a running -event loop, using the lambda form (executor.submit(lambda: asyncio.run(...))) to avoid -passing a pre-created coroutine to asyncio.run in another thread. -""" +"""BackgroundNoiseAugmentor: mix environmental noise into WAV audio samples.""" from __future__ import annotations -import asyncio -import concurrent.futures from pathlib import Path from typing import Any, Protocol @@ -47,52 +26,16 @@ class NoiseProvider(Protocol): def list_files(self) -> list[Path]: ... -def _read_duration_s(audio_reader: AudioReader, path: Path) -> float: - """Read a file and return its duration in seconds. - - When called outside a running event loop, uses asyncio.run() directly. - When called from within a running event loop (e.g. from a synchronous - method invoked by an async caller), runs asyncio.run() in a thread executor - so the loop thread is not blocked. Uses the lambda form - (executor.submit(lambda: asyncio.run(_read()))) to create the coroutine - inside the worker thread, avoiding passing a pre-created coroutine across - threads. - """ - async def _read() -> float: - data = await audio_reader.read(path) - return len(data.samples) / data.sample_rate - - try: - asyncio.get_running_loop() - except RuntimeError: - # No running loop — safe to call asyncio.run() directly - return asyncio.run(_read()) - - # Inside a running loop — run in a thread to avoid blocking the loop. - # The lambda creates the coroutine inside the worker thread so a fresh - # event loop is used per asyncio.run() call. - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(lambda: asyncio.run(_read())) - return future.result() - - class BackgroundNoiseAugmentor(ModifierStage[AudioSample, AudioSample]): """Augmentation stage that mixes environmental noise into WAV samples. The noise file is always chosen (even when not applied) so that the content hash remains stable across configuration changes. Only - noise_start_s and noise_volume are set to 0.0 when the noise is not applied. - - Noise mixing algorithm: slice noise at noise_start_s for len(audio) samples, - zero-pad if noise is shorter after the slice, multiply by noise_volume, add to - audio samples, then clip to [-1.0, 1.0]. - - input_dir is the directory containing the input WAV files referenced by - AudioSample.path. This is typically the output directory of the preceding stage. + noise_volume is set to 0.0 when the noise is not applied. noise_start_s + is always 0.0 (noise always mixed from the start of the noise file). - noise_dir is the directory containing noise WAV files. Stored directly rather - than derived from NoiseProvider.list_files()[0].parent to avoid redundant - filesystem calls and IndexError on an empty directory. + Noise mixing: slice noise at noise_start_s for len(audio) samples, + zero-pad if noise is shorter, multiply by noise_volume, add, clip to [-1.0, 1.0]. """ def __init__( @@ -118,34 +61,19 @@ def __init__( def _get_applied_values( self, sample: AudioSample, generator: VariationGenerator ) -> dict[str, Any]: - # Always choose a noise file for hash stability. + # Always choose a noise file for hash stability, even when not applied. # list_files() is called once; _noise_dir resolves the full path. noise_files = sorted([p.name for p in self._noise_provider.list_files()]) noise_file: str = generator.choose("noise_file", noise_files) if generator.should_vary("noise", self._vary_probability): - # noise_start_s bounds are derived from file durations. Both files are read - # synchronously here because _get_applied_values is a synchronous method. - # The durations are needed to clamp the start so the noise slice contains - # actual noise data rather than silence from an out-of-range offset. - noise_path = self._noise_dir / noise_file - audio_path = self._input_dir / sample.path - - noise_duration_s = _read_duration_s(self._audio_reader, noise_path) - audio_duration_s = _read_duration_s(self._audio_reader, audio_path) - - max_start_s = noise_duration_s - audio_duration_s - clamped_max = max(0.0, max_start_s) - start_filter = MinMaxFilter(0.0, clamped_max, precision=2) - noise_start_s = generator.generate("noise_start_s", start_filter) noise_volume = generator.generate("noise_volume", self._volume_filter) else: - noise_start_s = 0.0 noise_volume = 0.0 return { "noise_file": noise_file, - "noise_start_s": float(noise_start_s), + "noise_start_s": 0.0, # Always zero — noise is mixed from the start of the file. "noise_volume": float(noise_volume), } @@ -169,22 +97,24 @@ async def _generate_output( input_path = self._input_dir / input_sample.path audio = await self._audio_reader.read(input_path) - output_samples = audio.samples.copy() if noise_volume > 0.0: noise_path = self._noise_dir / noise_file noise_audio = await self._audio_reader.read(noise_path) start_sample = int(noise_start_s * noise_audio.sample_rate) - n_needed = len(output_samples) + n_needed = len(audio.samples) noise_slice = noise_audio.samples[start_sample: start_sample + n_needed] # Zero-pad if noise slice is shorter than audio if len(noise_slice) < n_needed: noise_slice = np.pad(noise_slice, (0, n_needed - len(noise_slice))) - output_samples = output_samples + noise_volume * noise_slice - output_samples = np.clip(output_samples, -1.0, 1.0).astype(np.float32) + output_samples: np.ndarray = np.clip( + audio.samples + noise_volume * noise_slice, -1.0, 1.0 + ).astype(np.float32) + else: + output_samples = audio.samples self._output_dir.mkdir(parents=True, exist_ok=True) output_path = conventions.sample_file_path(self._output_dir, output_id, "wav") diff --git a/ml/test/pipeline/speech/test_background_noise_stage.py b/ml/test/pipeline/speech/test_background_noise_stage.py index 7baee819..c6b1f8c8 100644 --- a/ml/test/pipeline/speech/test_background_noise_stage.py +++ b/ml/test/pipeline/speech/test_background_noise_stage.py @@ -265,35 +265,20 @@ def test_noise_volume_stored_as_float(self, tmp_path: Path) -> None: # --------------------------------------------------------------------------- -# TestAppliedValuesNoiseShorterThanAudio +# TestNoiseStartS # --------------------------------------------------------------------------- -class TestNoiseShorterThanAudio: - def test_noise_start_s_clamped_to_zero_when_noise_shorter_than_audio( - self, tmp_path: Path - ) -> None: - """When noise file is shorter than audio, max_start_s < 0 → clamped to 0.0.""" - # noise is 0.1s, audio is 0.5s → max_start_s = -0.4s → clamp to 0.0 - reader = _RecordingAudioReader( - sample_rate=16000, - audio_duration_s=0.5, - noise_duration_s=0.1, - ) - noise_dir = tmp_path / "noise" - noise_dir.mkdir() - reader.register_noise_path(noise_dir / "short.wav") - +class TestNoiseStartS: + def test_noise_start_s_is_always_zero(self, tmp_path: Path) -> None: + """noise_start_s is always 0.0 — noise is mixed from the start of the file.""" stage, _, _ = _make_stage( tmp_path, - noise_dir=noise_dir, - audio_reader=reader, vary_probability=1.0, volume_min=0.1, volume_max=0.1, - noise_filenames=["short.wav"], ) - sample = _make_audio_sample(wav_dir=tmp_path) + sample = _make_audio_sample() av = stage._get_applied_values(sample, VariationGenerator(42)) assert av["noise_start_s"] == 0.0 @@ -347,7 +332,7 @@ def test_derive_id_uses_input_id_as_prefix(self, tmp_path: Path) -> None: class TestGenerateOutput: def test_generate_output_calls_audio_reader_for_input_and_noise(self, tmp_path: Path) -> None: - """When noise is applied (vary_probability=1.0), reader is called for both noise and audio.""" + """When noise is applied (vary_probability=1.0), reader is called for both audio and noise.""" noise_dir = tmp_path / "noise" _write_noise_wav(noise_dir, "traffic.wav") @@ -365,12 +350,8 @@ def test_generate_output_calls_audio_reader_for_input_and_noise(self, tmp_path: ) sample = _make_audio_sample(wav_dir=tmp_path) asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) - # Should have called reader at least 3 times: - # - noise file duration (in _get_applied_values) - # - audio duration (in _get_applied_values) - # - audio file (in _generate_output) - # - noise file (in _generate_output) - assert len(reader.calls) >= 3 + # _generate_output reads the input audio file and the noise file — 2 calls total. + assert len(reader.calls) >= 2 def test_generate_output_calls_audio_writer(self, tmp_path: Path) -> None: noise_dir = tmp_path / "noise" From b85d2b24356e809d6a13a108f2dafd11a71860d3 Mon Sep 17 00:00:00 2001 From: Joe Davis Date: Thu, 18 Jun 2026 06:38:51 -0700 Subject: [PATCH 6/7] ADR-227: Avoid output_samples copy when amplitude is 0 in MicrophoneNoiseAugmentor Per Copilot review (PRRT_kwDOQYcoLM6KbINm) and user decision (PRRT_kwDOQYcoLM6KjhqB): only add Gaussian noise if amplitude > 0. Use audio.samples directly on the no-noise path instead of creating an unnecessary copy. New array is only allocated when mixing occurs. --- ml/pipeline/speech/mic_noise_stage.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/ml/pipeline/speech/mic_noise_stage.py b/ml/pipeline/speech/mic_noise_stage.py index 7fc0eeb2..68c5b29a 100644 --- a/ml/pipeline/speech/mic_noise_stage.py +++ b/ml/pipeline/speech/mic_noise_stage.py @@ -74,12 +74,13 @@ async def _generate_output( input_path = self._input_dir / input_sample.path audio = await self._audio_reader.read(input_path) - output_samples = audio.samples.copy() if amplitude > 0.0: rng = np.random.default_rng(output_seed) - noise = rng.normal(0, amplitude, len(output_samples)).astype(np.float32) - output_samples = np.clip(output_samples + noise, -1.0, 1.0).astype(np.float32) + noise = rng.normal(0, amplitude, len(audio.samples)).astype(np.float32) + output_samples: np.ndarray = np.clip(audio.samples + noise, -1.0, 1.0).astype(np.float32) + else: + output_samples = audio.samples self._output_dir.mkdir(parents=True, exist_ok=True) output_path = conventions.sample_file_path(self._output_dir, output_id, "wav") From 852e879df6bb2b91942d2f13e974829f435e029d Mon Sep 17 00:00:00 2001 From: Joe Davis Date: Thu, 18 Jun 2026 16:46:51 -0700 Subject: [PATCH 7/7] ADR-227: Redesign NoiseProvider; return input when noise not applied MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - NoiseProvider.list_files() now returns list[tuple[str, AudioData]] so that implementations can load and resample files once at construction time, avoiding per-sample I/O and ensuring sample rate consistency. - _DirectoryNoiseProvider gains a sample_rate: int constructor parameter; it loads all WAV files eagerly via librosa (resampling to sample_rate), satisfying the sample-rate equality contract. - BackgroundNoiseAugmentor removes noise_dir from its constructor (no longer needed; noise audio is retrieved directly from the provider). - _generate_output in both BackgroundNoiseAugmentor and MicrophoneNoiseAugmentor returns input_sample unchanged when noise_volume / amplitude == 0.0 — no file is written and no I/O occurs, eliminating unnecessary file copies. - MicrophoneNoiseAugmentor._derive_id returns input_sample.id when amplitude == 0.0 (no _mic### suffix added for unmodified samples). - A sample_rate: int field is added to PipelineParams (pipeline-level param) with a sample_rate: 16000 default in params.yaml. - Added ValueError guard in _generate_output when noise and audio sample rates differ, with a diagnostic message explaining the fix. Co-Authored-By: Claude Sonnet 4.6 --- ml/params.yaml | 1 + ml/pipeline/speech/background_noise_stage.py | 72 +++--- ml/pipeline/speech/mic_noise_stage.py | 24 +- ml/pipeline/stages/params.py | 2 + .../stages/speech_05_add_background_noise.py | 40 +-- .../speech/test_background_noise_stage.py | 228 +++++++++--------- .../pipeline/speech/test_mic_noise_stage.py | 100 +++++--- ml/test/pipeline/stages/test_params.py | 15 ++ 8 files changed, 281 insertions(+), 201 deletions(-) diff --git a/ml/params.yaml b/ml/params.yaml index b7326429..e57d7eaa 100644 --- a/ml/params.yaml +++ b/ml/params.yaml @@ -2,6 +2,7 @@ pipeline: input_phrases_path: scripts/intent_prediction/01_input_phrases.csv variations_per_phrase: 20 subsample_rate: 200 + sample_rate: 16000 stages: generate_phrases: diff --git a/ml/pipeline/speech/background_noise_stage.py b/ml/pipeline/speech/background_noise_stage.py index ca212106..a0f5ac56 100644 --- a/ml/pipeline/speech/background_noise_stage.py +++ b/ml/pipeline/speech/background_noise_stage.py @@ -17,13 +17,17 @@ class NoiseProvider(Protocol): - """Protocol for listing available noise files. + """Protocol for supplying pre-loaded noise audio data. - Consistent with TtsProvider living in tts_stage.py — the protocol belongs - alongside the stage that uses it. + Implementations load and resample noise files at construction time so that + each call to list_files() is cheap and all returned audio uses the same + sample rate. + + Returns a list of (filename, AudioData) pairs. The filename is used for + hash-stable variation selection; AudioData is used for mixing. """ - def list_files(self) -> list[Path]: ... + def list_files(self) -> list[tuple[str, AudioData]]: ... class BackgroundNoiseAugmentor(ModifierStage[AudioSample, AudioSample]): @@ -34,7 +38,10 @@ class BackgroundNoiseAugmentor(ModifierStage[AudioSample, AudioSample]): noise_volume is set to 0.0 when the noise is not applied. noise_start_s is always 0.0 (noise always mixed from the start of the noise file). - Noise mixing: slice noise at noise_start_s for len(audio) samples, + When noise_volume == 0.0, the input sample is returned unchanged — no + file is written and no I/O occurs. + + Noise mixing (when applied): slice noise for len(audio) samples, zero-pad if noise is shorter, multiply by noise_volume, add, clip to [-1.0, 1.0]. """ @@ -45,7 +52,6 @@ def __init__( audio_reader: AudioReader, audio_writer: AudioWriter, input_dir: Path, - noise_dir: Path, noise_provider: NoiseProvider, params: AddBackgroundNoiseParams, ) -> None: @@ -53,7 +59,6 @@ def __init__( self._audio_reader = audio_reader self._audio_writer = audio_writer self._input_dir = input_dir - self._noise_dir = noise_dir self._noise_provider = noise_provider self._vary_probability = params.vary_probability self._volume_filter = MinMaxFilter(params.volume_min, params.volume_max, precision=2) @@ -62,9 +67,9 @@ def _get_applied_values( self, sample: AudioSample, generator: VariationGenerator ) -> dict[str, Any]: # Always choose a noise file for hash stability, even when not applied. - # list_files() is called once; _noise_dir resolves the full path. - noise_files = sorted([p.name for p in self._noise_provider.list_files()]) - noise_file: str = generator.choose("noise_file", noise_files) + # list_files() returns pre-loaded (name, AudioData) pairs. + noise_items = self._noise_provider.list_files() + noise_file: str = generator.choose("noise_file", sorted([name for name, _ in noise_items])) if generator.should_vary("noise", self._vary_probability): noise_volume = generator.generate("noise_volume", self._volume_filter) @@ -91,30 +96,41 @@ async def _generate_output( applied_values: dict[str, Any], parent_content_hash: str, ) -> AudioSample: + noise_volume: float = applied_values["noise_volume"] + + # When not applied, return the input sample unchanged — no I/O. + if noise_volume == 0.0: + return input_sample + noise_file: str = applied_values["noise_file"] noise_start_s: float = applied_values["noise_start_s"] - noise_volume: float = applied_values["noise_volume"] input_path = self._input_dir / input_sample.path audio = await self._audio_reader.read(input_path) - if noise_volume > 0.0: - noise_path = self._noise_dir / noise_file - noise_audio = await self._audio_reader.read(noise_path) - - start_sample = int(noise_start_s * noise_audio.sample_rate) - n_needed = len(audio.samples) - noise_slice = noise_audio.samples[start_sample: start_sample + n_needed] - - # Zero-pad if noise slice is shorter than audio - if len(noise_slice) < n_needed: - noise_slice = np.pad(noise_slice, (0, n_needed - len(noise_slice))) - - output_samples: np.ndarray = np.clip( - audio.samples + noise_volume * noise_slice, -1.0, 1.0 - ).astype(np.float32) - else: - output_samples = audio.samples + # Look up the pre-loaded noise audio from the provider. + noise_items = self._noise_provider.list_files() + noise_audio = next(data for name, data in noise_items if name == noise_file) + + if noise_audio.sample_rate != audio.sample_rate: + raise ValueError( + f"Noise file '{noise_file}' has sample_rate {noise_audio.sample_rate} Hz " + f"but input audio has sample_rate {audio.sample_rate} Hz. " + f"Configure _DirectoryNoiseProvider with the pipeline sample_rate so that " + f"noise files are resampled at load time." + ) + + start_sample = int(noise_start_s * noise_audio.sample_rate) + n_needed = len(audio.samples) + noise_slice = noise_audio.samples[start_sample: start_sample + n_needed] + + # Zero-pad if noise slice is shorter than audio + if len(noise_slice) < n_needed: + noise_slice = np.pad(noise_slice, (0, n_needed - len(noise_slice))) + + output_samples: np.ndarray = np.clip( + audio.samples + noise_volume * noise_slice, -1.0, 1.0 + ).astype(np.float32) self._output_dir.mkdir(parents=True, exist_ok=True) output_path = conventions.sample_file_path(self._output_dir, output_id, "wav") diff --git a/ml/pipeline/speech/mic_noise_stage.py b/ml/pipeline/speech/mic_noise_stage.py index 68c5b29a..11f53045 100644 --- a/ml/pipeline/speech/mic_noise_stage.py +++ b/ml/pipeline/speech/mic_noise_stage.py @@ -1,6 +1,7 @@ """MicrophoneNoiseAugmentor: add Gaussian microphone noise to WAV audio samples. mic_noise_amplitude is stored as 0.0 when should_vary returns False. +When amplitude == 0.0, the input sample is returned unchanged — no file is written. Gaussian noise is seeded from output_seed for reproducibility. """ @@ -23,8 +24,11 @@ class MicrophoneNoiseAugmentor(ModifierStage[AudioSample, AudioSample]): """Augmentation stage that adds Gaussian microphone noise to WAV samples. - mic_noise_amplitude is stored as 0.0 when the noise is not applied. Gaussian - noise uses output_seed for reproducibility via numpy's default_rng. + mic_noise_amplitude is stored as 0.0 when the noise is not applied. When + amplitude == 0.0, the input sample is returned unchanged — no file is + written and no I/O occurs. + + Gaussian noise uses output_seed for reproducibility via numpy's default_rng. input_dir is the directory containing the input WAV files referenced by AudioSample.path. This is typically the output directory of the preceding stage. @@ -60,6 +64,9 @@ def _get_applied_values( def _derive_id(self, input_sample: AudioSample, applied_values: dict[str, Any]) -> str: amplitude: float = applied_values["mic_noise_amplitude"] + if amplitude == 0.0: + # No modification — return the input id unchanged. + return input_sample.id return f"{input_sample.id}_mic{int(amplitude * 1000)}" async def _generate_output( @@ -72,15 +79,16 @@ async def _generate_output( ) -> AudioSample: amplitude: float = applied_values["mic_noise_amplitude"] + # When not applied, return the input sample unchanged — no I/O. + if amplitude == 0.0: + return input_sample + input_path = self._input_dir / input_sample.path audio = await self._audio_reader.read(input_path) - if amplitude > 0.0: - rng = np.random.default_rng(output_seed) - noise = rng.normal(0, amplitude, len(audio.samples)).astype(np.float32) - output_samples: np.ndarray = np.clip(audio.samples + noise, -1.0, 1.0).astype(np.float32) - else: - output_samples = audio.samples + rng = np.random.default_rng(output_seed) + noise = rng.normal(0, amplitude, len(audio.samples)).astype(np.float32) + output_samples: np.ndarray = np.clip(audio.samples + noise, -1.0, 1.0).astype(np.float32) self._output_dir.mkdir(parents=True, exist_ok=True) output_path = conventions.sample_file_path(self._output_dir, output_id, "wav") diff --git a/ml/pipeline/stages/params.py b/ml/pipeline/stages/params.py index 04e72a2e..8153dc77 100644 --- a/ml/pipeline/stages/params.py +++ b/ml/pipeline/stages/params.py @@ -55,6 +55,7 @@ class AddMicNoiseParams: class PipelineParams: variations_per_phrase: int subsample_rate: int + sample_rate: int generate_phrases: GeneratePhraseParams generate_samples: GenerateSamplesParams add_delays: AddDelaysParams @@ -76,6 +77,7 @@ def load(cls, path: Path) -> "PipelineParams": return cls( variations_per_phrase=int(pipeline["variations_per_phrase"]), subsample_rate=int(pipeline["subsample_rate"]), + sample_rate=int(pipeline["sample_rate"]), generate_phrases=GeneratePhraseParams( repeat_modifier_chance=float(stage_phrases["repeat_modifier_chance"]), pleasantry_chance=float(stage_phrases["pleasantry_chance"]), diff --git a/ml/pipeline/stages/speech_05_add_background_noise.py b/ml/pipeline/stages/speech_05_add_background_noise.py index 3dd91b84..116e6f89 100644 --- a/ml/pipeline/stages/speech_05_add_background_noise.py +++ b/ml/pipeline/stages/speech_05_add_background_noise.py @@ -8,7 +8,7 @@ sys.path.insert(0, str(Path(__file__).parent.parent.parent)) from pipeline.core.manifest import ManifestStore -from pipeline.io.audio_io import LibrosaAudioReader, SoundfileAudioWriter +from pipeline.io.audio_io import AudioData, LibrosaAudioReader, SoundfileAudioWriter from pipeline.speech.background_noise_stage import BackgroundNoiseAugmentor from pipeline.stages import conventions from pipeline.stages.params import PipelineParams @@ -19,23 +19,36 @@ class _DirectoryNoiseProvider: """NoiseProvider backed by a filesystem directory. - Lists all WAV files in the given directory. DVC wiring points this at the - appropriate data/ path; this class requires no knowledge of the DVC layout. + Loads all WAV files in the given directory at construction time, resampling + them to ``sample_rate`` so that all noise audio matches the pipeline sample + rate. This avoids per-sample I/O at augmentation time and ensures noise + files are at the correct sample rate before mixing. Raises ValueError with the directory path if the directory contains no WAV - files, so that a misconfigured --noise-dir produces an actionable error - rather than the opaque ValueError("options must be non-empty") from - VariationGenerator.choose(). + files, so that a misconfigured --noise-dir produces an actionable error. """ - def __init__(self, noise_dir: Path) -> None: + def __init__(self, noise_dir: Path, sample_rate: int) -> None: self._noise_dir = noise_dir + self._items: list[tuple[str, AudioData]] = self._load(noise_dir, sample_rate) - def list_files(self) -> list[Path]: - files = list(self._noise_dir.glob("*.wav")) - if not files: - raise ValueError(f"No WAV files found in noise directory: {self._noise_dir}") - return files + @staticmethod + def _load(noise_dir: Path, sample_rate: int) -> list[tuple[str, AudioData]]: + import librosa + + wav_paths = sorted(noise_dir.glob("*.wav")) + if not wav_paths: + raise ValueError(f"No WAV files found in noise directory: {noise_dir}") + + items: list[tuple[str, AudioData]] = [] + for path in wav_paths: + # librosa.load with sr=sample_rate resamples on load if needed. + samples, sr = librosa.load(str(path), sr=sample_rate, mono=True, dtype="float32") + items.append((path.name, AudioData(samples=samples, sample_rate=int(sr)))) + return items + + def list_files(self) -> list[tuple[str, AudioData]]: + return list(self._items) def main() -> None: @@ -60,8 +73,7 @@ def main() -> None: audio_reader=LibrosaAudioReader(), audio_writer=SoundfileAudioWriter(), input_dir=args.input_manifest_dir, - noise_dir=args.noise_dir, - noise_provider=_DirectoryNoiseProvider(args.noise_dir), + noise_provider=_DirectoryNoiseProvider(args.noise_dir, params.sample_rate), params=params.add_background_noise, ) diff --git a/ml/test/pipeline/speech/test_background_noise_stage.py b/ml/test/pipeline/speech/test_background_noise_stage.py index c6b1f8c8..8dafe8ff 100644 --- a/ml/test/pipeline/speech/test_background_noise_stage.py +++ b/ml/test/pipeline/speech/test_background_noise_stage.py @@ -52,30 +52,20 @@ def _make_audio_sample( class _RecordingAudioReader: - """Stub AudioReader that returns configured durations per path.""" + """Stub AudioReader that records read() calls and returns silence.""" def __init__( self, sample_rate: int = 16000, audio_duration_s: float = 0.5, - noise_duration_s: float = 2.0, ) -> None: self.calls: list[Path] = [] self._sample_rate = sample_rate self._audio_duration_s = audio_duration_s - self._noise_duration_s = noise_duration_s - # If path stem ends with "noise", return noise duration, else audio duration - self._noise_paths: set[Path] = set() - - def register_noise_path(self, path: Path) -> None: - self._noise_paths.add(path) async def read(self, path: Path) -> AudioData: self.calls.append(path) - if path in self._noise_paths or path.parent.name == "noise": - n_samples = int(self._sample_rate * self._noise_duration_s) - else: - n_samples = int(self._sample_rate * self._audio_duration_s) + n_samples = int(self._sample_rate * self._audio_duration_s) samples = np.zeros(n_samples, dtype=np.float32) return AudioData(samples=samples, sample_rate=self._sample_rate) @@ -92,14 +82,27 @@ async def write(self, path: Path, audio: AudioData) -> None: class _FakeNoiseProvider: - """Fake NoiseProvider that returns a fixed list of noise file paths.""" + """Fake NoiseProvider that returns a fixed list of (name, AudioData) tuples.""" - def __init__(self, noise_dir: Path, filenames: list[str]) -> None: - self._noise_dir = noise_dir - self._filenames = filenames + def __init__( + self, + filenames: list[str], + sample_rate: int = 16000, + duration_s: float = 2.0, + ) -> None: + self._items = [ + ( + name, + AudioData( + samples=(np.random.default_rng(i).random(int(sample_rate * duration_s)) * 2 - 1).astype(np.float32), + sample_rate=sample_rate, + ), + ) + for i, name in enumerate(filenames) + ] - def list_files(self) -> list[Path]: - return [self._noise_dir / name for name in self._filenames] + def list_files(self) -> list[tuple[str, AudioData]]: + return list(self._items) def _make_params( @@ -116,7 +119,6 @@ def _make_params( def _make_stage( output_dir: Path, - noise_dir: Path | None = None, *, audio_reader: _RecordingAudioReader | None = None, audio_writer: _RecordingAudioWriter | None = None, @@ -127,16 +129,14 @@ def _make_stage( volume_min: float = 0.0, volume_max: float = 0.3, noise_filenames: list[str] | None = None, + sample_rate: int = 16000, ) -> tuple[BackgroundNoiseAugmentor, _RecordingAudioReader, _RecordingAudioWriter]: if audio_reader is None: - audio_reader = _RecordingAudioReader() + audio_reader = _RecordingAudioReader(sample_rate=sample_rate) if audio_writer is None: audio_writer = _RecordingAudioWriter() if input_dir is None: input_dir = output_dir - if noise_dir is None: - noise_dir = output_dir / "noise" - noise_dir.mkdir(parents=True, exist_ok=True) if params is None: params = _make_params( vary_probability=vary_probability, @@ -146,25 +146,24 @@ def _make_stage( if noise_provider is None: if noise_filenames is None: noise_filenames = ["traffic.wav"] - noise_provider = _FakeNoiseProvider(noise_dir, noise_filenames) + noise_provider = _FakeNoiseProvider(noise_filenames, sample_rate=sample_rate) stage = BackgroundNoiseAugmentor( output_dir=output_dir, manifest_store=ManifestStore(), audio_reader=audio_reader, audio_writer=audio_writer, input_dir=input_dir, - noise_dir=noise_dir, noise_provider=noise_provider, params=params, ) return stage, audio_reader, audio_writer -def _write_noise_wav(noise_dir: Path, filename: str, sample_rate: int = 16000, duration_s: float = 2.0) -> Path: - noise_dir.mkdir(parents=True, exist_ok=True) - path = noise_dir / filename +def _write_audio_wav(dir: Path, filename: str, sample_rate: int = 16000, duration_s: float = 0.5) -> Path: + dir.mkdir(parents=True, exist_ok=True) + path = dir / filename n_samples = int(sample_rate * duration_s) - data = (np.random.default_rng(42).random(n_samples) * 2 - 1).astype(np.float32) + data = np.zeros(n_samples, dtype=np.float32) sf.write(str(path), data, sample_rate, format="WAV", subtype="PCM_16") return path @@ -331,17 +330,42 @@ def test_derive_id_uses_input_id_as_prefix(self, tmp_path: Path) -> None: class TestGenerateOutput: - def test_generate_output_calls_audio_reader_for_input_and_noise(self, tmp_path: Path) -> None: - """When noise is applied (vary_probability=1.0), reader is called for both audio and noise.""" - noise_dir = tmp_path / "noise" - _write_noise_wav(noise_dir, "traffic.wav") + def test_generate_output_does_not_call_reader_when_volume_is_zero(self, tmp_path: Path) -> None: + """When noise not applied (volume=0), no audio reading occurs.""" + reader = _RecordingAudioReader() + stage, _, _ = _make_stage( + tmp_path, + audio_reader=reader, + vary_probability=0.0, + ) + sample = _make_audio_sample() + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert len(reader.calls) == 0 - reader = _RecordingAudioReader(audio_duration_s=0.5, noise_duration_s=2.0) - reader.register_noise_path(noise_dir / "traffic.wav") + def test_generate_output_does_not_call_writer_when_volume_is_zero(self, tmp_path: Path) -> None: + """When noise not applied (volume=0), no file is written.""" + stage, _, writer = _make_stage( + tmp_path, + vary_probability=0.0, + ) + sample = _make_audio_sample() + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert len(writer.calls) == 0 + def test_generate_output_returns_input_when_volume_is_zero(self, tmp_path: Path) -> None: + """When noise not applied, transform returns the original input sample.""" + stage, _, _ = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample(sample_id="MY_SAMPLE") + result = asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert result.samples[0].id == sample.id + assert result.samples[0].path == sample.path + + def test_generate_output_calls_audio_reader_for_input_and_does_not_call_for_noise(self, tmp_path: Path) -> None: + """When noise is applied (vary_probability=1.0), reader is called for input only; noise comes from provider.""" + _write_audio_wav(tmp_path, "TV_ON_Jenny_r100.wav") + reader = _RecordingAudioReader() stage, _, _ = _make_stage( tmp_path, - noise_dir=noise_dir, audio_reader=reader, vary_probability=1.0, volume_min=0.1, @@ -350,20 +374,16 @@ def test_generate_output_calls_audio_reader_for_input_and_noise(self, tmp_path: ) sample = _make_audio_sample(wav_dir=tmp_path) asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) - # _generate_output reads the input audio file and the noise file — 2 calls total. - assert len(reader.calls) >= 2 - - def test_generate_output_calls_audio_writer(self, tmp_path: Path) -> None: - noise_dir = tmp_path / "noise" - _write_noise_wav(noise_dir, "traffic.wav") - - reader = _RecordingAudioReader(audio_duration_s=0.5, noise_duration_s=2.0) - reader.register_noise_path(noise_dir / "traffic.wav") + # Only 1 call: the input audio file. Noise audio comes from the provider directly. + assert len(reader.calls) == 1 + def test_generate_output_calls_audio_writer_when_noise_applied(self, tmp_path: Path) -> None: + _write_audio_wav(tmp_path, "TV_ON_Jenny_r100.wav") stage, _, writer = _make_stage( tmp_path, - noise_dir=noise_dir, - audio_reader=reader, + vary_probability=1.0, + volume_min=0.1, + volume_max=0.1, noise_filenames=["traffic.wav"], ) sample = _make_audio_sample(wav_dir=tmp_path) @@ -374,73 +394,30 @@ def test_output_audio_same_length_as_input(self, tmp_path: Path) -> None: """Background noise does not change the length of the audio.""" sample_rate = 16000 duration_s = 0.5 - noise_dir = tmp_path / "noise" - _write_noise_wav(noise_dir, "traffic.wav", sample_rate=sample_rate, duration_s=2.0) - - reader = _RecordingAudioReader( - sample_rate=sample_rate, audio_duration_s=duration_s, noise_duration_s=2.0 - ) - reader.register_noise_path(noise_dir / "traffic.wav") + _write_audio_wav(tmp_path, "TV_ON_Jenny_r100.wav", sample_rate=sample_rate, duration_s=duration_s) + reader = _RecordingAudioReader(sample_rate=sample_rate, audio_duration_s=duration_s) stage, _, writer = _make_stage( tmp_path, - noise_dir=noise_dir, audio_reader=reader, vary_probability=1.0, volume_min=0.1, volume_max=0.1, noise_filenames=["traffic.wav"], + sample_rate=sample_rate, ) sample = _make_audio_sample(wav_dir=tmp_path, sample_rate=sample_rate, duration_s=duration_s) asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) _path, written_audio = writer.calls[0] assert len(written_audio.samples) == int(sample_rate * duration_s) - def test_output_audio_unchanged_when_volume_is_zero(self, tmp_path: Path) -> None: - """When noise not applied (volume=0), output samples equal input samples.""" - sample_rate = 16000 - duration_s = 0.1 - noise_dir = tmp_path / "noise" - _write_noise_wav(noise_dir, "traffic.wav", sample_rate=sample_rate, duration_s=2.0) - - # Use a special reader that returns a known nonzero signal for audio - class _ConstantReader: - async def read(self, path: Path) -> AudioData: - n = int(sample_rate * duration_s) - if path.parent.name == "noise": - samples = np.ones(int(sample_rate * 2.0), dtype=np.float32) * 0.9 - else: - samples = np.full(n, 0.5, dtype=np.float32) - return AudioData(samples=samples, sample_rate=sample_rate) - - writer = _RecordingAudioWriter() - stage = BackgroundNoiseAugmentor( - output_dir=tmp_path, - manifest_store=ManifestStore(), - audio_reader=_ConstantReader(), - audio_writer=writer, - input_dir=tmp_path, - noise_dir=noise_dir, - noise_provider=_FakeNoiseProvider(noise_dir, ["traffic.wav"]), - params=_make_params(vary_probability=0.0), - ) - sample = _make_audio_sample(wav_dir=tmp_path, sample_rate=sample_rate, duration_s=duration_s) - asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) - _path, written_audio = writer.calls[0] - # All samples should still be 0.5 (unchanged) - assert np.allclose(written_audio.samples, 0.5) - - def test_transcript_preserved_in_output_sample(self, tmp_path: Path) -> None: - noise_dir = tmp_path / "noise" - _write_noise_wav(noise_dir, "traffic.wav") - - reader = _RecordingAudioReader(audio_duration_s=0.5, noise_duration_s=2.0) - reader.register_noise_path(noise_dir / "traffic.wav") - + def test_transcript_preserved_in_output_sample_when_noise_applied(self, tmp_path: Path) -> None: + _write_audio_wav(tmp_path, "TV_ON_Jenny_r100.wav") stage, _, _ = _make_stage( tmp_path, - noise_dir=noise_dir, - audio_reader=reader, + vary_probability=1.0, + volume_min=0.1, + volume_max=0.1, noise_filenames=["traffic.wav"], ) sample = _make_audio_sample(transcript="VOLUME_UP", wav_dir=tmp_path) @@ -449,6 +426,35 @@ def test_transcript_preserved_in_output_sample(self, tmp_path: Path) -> None: ) assert result.samples[0].transcript == "VOLUME_UP" + def test_transcript_preserved_in_output_sample_when_not_applied(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample(transcript="CHANNEL_DOWN") + result = asyncio.run( + stage.transform(Manifest([sample]), tmp_path / "manifest.json") + ) + assert result.samples[0].transcript == "CHANNEL_DOWN" + + def test_sample_rate_mismatch_raises(self, tmp_path: Path) -> None: + """Raises ValueError when noise audio has a different sample rate than input audio.""" + _write_audio_wav(tmp_path, "TV_ON_Jenny_r100.wav", sample_rate=16000) + # Provide noise at a different sample rate + noise_provider = _FakeNoiseProvider(["traffic.wav"], sample_rate=8000) + reader = _RecordingAudioReader(sample_rate=16000) # Input is 16000 + writer = _RecordingAudioWriter() + stage = BackgroundNoiseAugmentor( + output_dir=tmp_path, + manifest_store=ManifestStore(), + audio_reader=reader, + audio_writer=writer, + input_dir=tmp_path, + noise_provider=noise_provider, + params=_make_params(vary_probability=1.0, volume_min=0.1, volume_max=0.1), + ) + sample = _make_audio_sample(wav_dir=tmp_path, sample_rate=16000) + import pytest + with pytest.raises(ValueError, match="sample.rate"): + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + # --------------------------------------------------------------------------- # TestSkipPath @@ -457,16 +463,13 @@ def test_transcript_preserved_in_output_sample(self, tmp_path: Path) -> None: class TestSkipPath: def test_skip_path_does_not_call_audio_writer_on_second_run(self, tmp_path: Path) -> None: - noise_dir = tmp_path / "noise" - _write_noise_wav(noise_dir, "traffic.wav") - - reader = _RecordingAudioReader(audio_duration_s=0.5, noise_duration_s=2.0) - reader.register_noise_path(noise_dir / "traffic.wav") - + """When noise IS applied, second run skips writing (output unchanged).""" + _write_audio_wav(tmp_path, "TV_ON_Jenny_r100.wav") stage, _, writer = _make_stage( tmp_path, - noise_dir=noise_dir, - audio_reader=reader, + vary_probability=1.0, + volume_min=0.1, + volume_max=0.1, noise_filenames=["traffic.wav"], ) sample = _make_audio_sample(wav_dir=tmp_path) @@ -479,16 +482,13 @@ def test_skip_path_does_not_call_audio_writer_on_second_run(self, tmp_path: Path assert len(writer.calls) == initial_count def test_skip_path_preserves_output_sample_id(self, tmp_path: Path) -> None: - noise_dir = tmp_path / "noise" - _write_noise_wav(noise_dir, "traffic.wav") - - reader = _RecordingAudioReader(audio_duration_s=0.5, noise_duration_s=2.0) - reader.register_noise_path(noise_dir / "traffic.wav") - + """When noise IS applied, second run returns the same output id.""" + _write_audio_wav(tmp_path, "TV_ON_Jenny_r100.wav") stage, _, _ = _make_stage( tmp_path, - noise_dir=noise_dir, - audio_reader=reader, + vary_probability=1.0, + volume_min=0.1, + volume_max=0.1, noise_filenames=["traffic.wav"], ) sample = _make_audio_sample(wav_dir=tmp_path) diff --git a/ml/test/pipeline/speech/test_mic_noise_stage.py b/ml/test/pipeline/speech/test_mic_noise_stage.py index 0e2ecb4f..9ed99325 100644 --- a/ml/test/pipeline/speech/test_mic_noise_stage.py +++ b/ml/test/pipeline/speech/test_mic_noise_stage.py @@ -177,22 +177,23 @@ def test_amplitude_zero_stored_as_float_not_int(self, tmp_path: Path) -> None: class TestDeriveId: def test_derive_id_format_with_noise_applied(self, tmp_path: Path) -> None: - """Format: {input.id}_mic{int(amplitude*1000)}""" + """Format: {input.id}_mic{int(amplitude*1000)} when amplitude > 0.""" stage, _, _ = _make_stage(tmp_path) sample = _make_audio_sample(sample_id="TV_ON_Jenny_r100") result = stage._derive_id(sample, {"mic_noise_amplitude": 0.025}) assert result == "TV_ON_Jenny_r100_mic25" - def test_derive_id_format_with_zero_amplitude(self, tmp_path: Path) -> None: + def test_derive_id_format_with_zero_amplitude_is_input_id(self, tmp_path: Path) -> None: + """When amplitude == 0 (no noise), id equals input id — no suffix added.""" stage, _, _ = _make_stage(tmp_path) sample = _make_audio_sample(sample_id="TV_ON_Jenny_r100") result = stage._derive_id(sample, {"mic_noise_amplitude": 0.0}) - assert result == "TV_ON_Jenny_r100_mic0" + assert result == "TV_ON_Jenny_r100" def test_derive_id_uses_input_id_as_prefix(self, tmp_path: Path) -> None: stage, _, _ = _make_stage(tmp_path) sample = _make_audio_sample(sample_id="VOLUME_UP_Aria_r110") - result = stage._derive_id(sample, {"mic_noise_amplitude": 0.0}) + result = stage._derive_id(sample, {"mic_noise_amplitude": 0.01}) assert result.startswith("VOLUME_UP_Aria_r110_") def test_derive_id_amplitude_int_conversion(self, tmp_path: Path) -> None: @@ -202,12 +203,19 @@ def test_derive_id_amplitude_int_conversion(self, tmp_path: Path) -> None: result = stage._derive_id(sample, {"mic_noise_amplitude": 0.05}) assert result == "S_mic50" - def test_derive_id_always_has_mic_suffix(self, tmp_path: Path) -> None: + def test_derive_id_always_has_mic_suffix_when_nonzero(self, tmp_path: Path) -> None: stage, _, _ = _make_stage(tmp_path) sample = _make_audio_sample(sample_id="X") result = stage._derive_id(sample, {"mic_noise_amplitude": 0.01}) assert "_mic" in result + def test_derive_id_no_mic_suffix_when_zero(self, tmp_path: Path) -> None: + """No _mic suffix when amplitude is zero.""" + stage, _, _ = _make_stage(tmp_path) + sample = _make_audio_sample(sample_id="X") + result = stage._derive_id(sample, {"mic_noise_amplitude": 0.0}) + assert "_mic" not in result + # --------------------------------------------------------------------------- # TestGenerateOutput @@ -215,14 +223,40 @@ def test_derive_id_always_has_mic_suffix(self, tmp_path: Path) -> None: class TestGenerateOutput: - def test_generate_output_calls_audio_reader(self, tmp_path: Path) -> None: - stage, reader, _ = _make_stage(tmp_path) + def test_generate_output_does_not_call_reader_when_amplitude_is_zero(self, tmp_path: Path) -> None: + """When amplitude == 0, no audio reading occurs.""" + stage, reader, _ = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample() + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert len(reader.calls) == 0 + + def test_generate_output_does_not_call_writer_when_amplitude_is_zero(self, tmp_path: Path) -> None: + """When amplitude == 0, no file is written.""" + stage, _, writer = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample() + asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert len(writer.calls) == 0 + + def test_generate_output_returns_input_when_amplitude_is_zero(self, tmp_path: Path) -> None: + """When amplitude == 0, transform returns the original input sample.""" + stage, _, _ = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample(sample_id="MY_SAMPLE") + result = asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) + assert result.samples[0].id == sample.id + assert result.samples[0].path == sample.path + + def test_generate_output_calls_audio_reader_when_amplitude_nonzero(self, tmp_path: Path) -> None: + stage, reader, _ = _make_stage( + tmp_path, vary_probability=1.0, amplitude_min=0.01, amplitude_max=0.01 + ) sample = _make_audio_sample(wav_dir=tmp_path) asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) assert len(reader.calls) == 1 - def test_generate_output_calls_audio_writer(self, tmp_path: Path) -> None: - stage, _, writer = _make_stage(tmp_path) + def test_generate_output_calls_audio_writer_when_amplitude_nonzero(self, tmp_path: Path) -> None: + stage, _, writer = _make_stage( + tmp_path, vary_probability=1.0, amplitude_min=0.01, amplitude_max=0.01 + ) sample = _make_audio_sample(wav_dir=tmp_path) asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) assert len(writer.calls) == 1 @@ -243,30 +277,6 @@ def test_output_audio_same_length_as_input(self, tmp_path: Path) -> None: _path, written_audio = writer.calls[0] assert len(written_audio.samples) == int(sample_rate * duration_s) - def test_output_audio_unchanged_when_amplitude_is_zero(self, tmp_path: Path) -> None: - """When noise not applied (amplitude=0), output samples equal input samples.""" - sample_rate = 16000 - duration_s = 0.1 - - class _ConstantReader: - async def read(self, path: Path) -> AudioData: - n = int(sample_rate * duration_s) - return AudioData(samples=np.full(n, 0.5, dtype=np.float32), sample_rate=sample_rate) - - writer = _RecordingAudioWriter() - stage = MicrophoneNoiseAugmentor( - output_dir=tmp_path, - manifest_store=ManifestStore(), - audio_reader=_ConstantReader(), - audio_writer=writer, - input_dir=tmp_path, - params=_make_params(vary_probability=0.0), - ) - sample = _make_audio_sample(wav_dir=tmp_path, sample_rate=sample_rate, duration_s=duration_s) - asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) - _path, written_audio = writer.calls[0] - assert np.allclose(written_audio.samples, 0.5) - def test_gaussian_noise_added_when_amplitude_nonzero(self, tmp_path: Path) -> None: """When amplitude > 0, output should differ from input (noise was added).""" sample_rate = 16000 @@ -293,14 +303,24 @@ async def read(self, path: Path) -> AudioData: # Noise was added — not all zeros anymore assert not np.all(written_audio.samples == 0.0) - def test_transcript_preserved_in_output_sample(self, tmp_path: Path) -> None: - stage, _, _ = _make_stage(tmp_path) + def test_transcript_preserved_when_noise_applied(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage( + tmp_path, vary_probability=1.0, amplitude_min=0.01, amplitude_max=0.01 + ) sample = _make_audio_sample(transcript="CHANNEL_UP", wav_dir=tmp_path) result = asyncio.run( stage.transform(Manifest([sample]), tmp_path / "manifest.json") ) assert result.samples[0].transcript == "CHANNEL_UP" + def test_transcript_preserved_when_not_applied(self, tmp_path: Path) -> None: + stage, _, _ = _make_stage(tmp_path, vary_probability=0.0) + sample = _make_audio_sample(transcript="CHANNEL_DOWN") + result = asyncio.run( + stage.transform(Manifest([sample]), tmp_path / "manifest.json") + ) + assert result.samples[0].transcript == "CHANNEL_DOWN" + # --------------------------------------------------------------------------- # TestSkipPath @@ -309,7 +329,10 @@ def test_transcript_preserved_in_output_sample(self, tmp_path: Path) -> None: class TestSkipPath: def test_skip_path_does_not_call_audio_writer_on_second_run(self, tmp_path: Path) -> None: - stage, _, writer = _make_stage(tmp_path) + """When noise IS applied, second run skips writing (output unchanged).""" + stage, _, writer = _make_stage( + tmp_path, vary_probability=1.0, amplitude_min=0.01, amplitude_max=0.01 + ) sample = _make_audio_sample(wav_dir=tmp_path) asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) @@ -320,7 +343,10 @@ def test_skip_path_does_not_call_audio_writer_on_second_run(self, tmp_path: Path assert len(writer.calls) == initial_count def test_skip_path_preserves_output_sample_id(self, tmp_path: Path) -> None: - stage, _, _ = _make_stage(tmp_path) + """When noise IS applied, second run returns the same output id.""" + stage, _, _ = _make_stage( + tmp_path, vary_probability=1.0, amplitude_min=0.01, amplitude_max=0.01 + ) sample = _make_audio_sample(wav_dir=tmp_path) result1 = asyncio.run(stage.transform(Manifest([sample]), tmp_path / "manifest.json")) diff --git a/ml/test/pipeline/stages/test_params.py b/ml/test/pipeline/stages/test_params.py index 3a4cce92..39a09a8d 100644 --- a/ml/test/pipeline/stages/test_params.py +++ b/ml/test/pipeline/stages/test_params.py @@ -30,6 +30,7 @@ def _write_params(path: Path, data: dict) -> Path: "pipeline": { "variations_per_phrase": 20, "subsample_rate": 200, + "sample_rate": 16000, }, "stages": { "generate_phrases": { @@ -72,6 +73,20 @@ def test_loads_pipeline_fields(self, tmp_path: Path) -> None: assert params.variations_per_phrase == 20 assert params.subsample_rate == 200 + assert params.sample_rate == 16000 + + def test_sample_rate_is_int(self, tmp_path: Path) -> None: + params_file = _write_params(tmp_path, _VALID_DATA) + params = PipelineParams.load(params_file) + assert isinstance(params.sample_rate, int) + + def test_missing_sample_rate_raises(self, tmp_path: Path) -> None: + import copy + data = copy.deepcopy(_VALID_DATA) + del data["pipeline"]["sample_rate"] + params_file = _write_params(tmp_path, data) + with pytest.raises(KeyError): + PipelineParams.load(params_file) def test_loads_generate_phrases_fields(self, tmp_path: Path) -> None: params_file = _write_params(tmp_path, _VALID_DATA)