diff --git a/ml/params.yaml b/ml/params.yaml index 60bca4c5..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: @@ -20,3 +21,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..e548c093 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,18 @@ 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_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 +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..a0f5ac56 --- /dev/null +++ b/ml/pipeline/speech/background_noise_stage.py @@ -0,0 +1,151 @@ +"""BackgroundNoiseAugmentor: mix environmental noise into WAV audio samples.""" + +from __future__ import annotations + +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 supplying pre-loaded noise audio data. + + 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[tuple[str, AudioData]]: ... + + +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_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). + + 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]. + """ + + 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, even when not applied. + # 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) + else: + noise_volume = 0.0 + + return { + "noise_file": noise_file, + "noise_start_s": 0.0, # Always zero — noise is mixed from the start of the file. + "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_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"] + + input_path = self._input_dir / input_sample.path + audio = await self._audio_reader.read(input_path) + + # 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") + 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..11f53045 --- /dev/null +++ b/ml/pipeline/speech/mic_noise_stage.py @@ -0,0 +1,109 @@ +"""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. +""" + +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. 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. + """ + + 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"] + 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( + 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"] + + # 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) + + 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") + 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..8153dc77 100644 --- a/ml/pipeline/stages/params.py +++ b/ml/pipeline/stages/params.py @@ -37,13 +37,30 @@ 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 subsample_rate: int + sample_rate: int 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,10 +71,13 @@ 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"]), 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"]), @@ -77,4 +97,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..116e6f89 --- /dev/null +++ b/ml/pipeline/stages/speech_05_add_background_noise.py @@ -0,0 +1,89 @@ +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 AudioData, 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. + + 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. + """ + + 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) + + @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: + 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.sample_rate), + 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..8dafe8ff --- /dev/null +++ b/ml/test/pipeline/speech/test_background_noise_stage.py @@ -0,0 +1,498 @@ +"""Unit tests for BackgroundNoiseAugmentor.""" + +from __future__ import annotations + +import asyncio +import hashlib +import sys +from pathlib import Path + +import numpy as np +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 records read() calls and returns silence.""" + + def __init__( + self, + sample_rate: int = 16000, + audio_duration_s: float = 0.5, + ) -> None: + self.calls: list[Path] = [] + self._sample_rate = sample_rate + self._audio_duration_s = audio_duration_s + + async def read(self, path: Path) -> AudioData: + self.calls.append(path) + 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 (name, AudioData) tuples.""" + + 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[tuple[str, AudioData]]: + return list(self._items) + + +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, + *, + 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, + sample_rate: int = 16000, +) -> tuple[BackgroundNoiseAugmentor, _RecordingAudioReader, _RecordingAudioWriter]: + if audio_reader is None: + 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 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_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_provider=noise_provider, + params=params, + ) + return stage, audio_reader, audio_writer + + +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.zeros(n_samples, dtype=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) + + +# --------------------------------------------------------------------------- +# TestNoiseStartS +# --------------------------------------------------------------------------- + + +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, + vary_probability=1.0, + volume_min=0.1, + volume_max=0.1, + ) + sample = _make_audio_sample() + 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_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 + + 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, + 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")) + # 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, + 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")) + 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 + _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, + 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_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, + 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) + result = asyncio.run( + stage.transform(Manifest([sample]), tmp_path / "manifest.json") + ) + 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 +# --------------------------------------------------------------------------- + + +class TestSkipPath: + def test_skip_path_does_not_call_audio_writer_on_second_run(self, tmp_path: Path) -> None: + """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, + 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")) + 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: + """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, + vary_probability=1.0, + volume_min=0.1, + volume_max=0.1, + 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..9ed99325 --- /dev/null +++ b/ml/test/pipeline/speech/test_mic_noise_stage.py @@ -0,0 +1,354 @@ +"""Unit tests for MicrophoneNoiseAugmentor.""" + +from __future__ import annotations + +import asyncio +import hashlib +import sys +from pathlib import Path + +import numpy as np +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)} 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_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" + + 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.01}) + 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_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 +# --------------------------------------------------------------------------- + + +class TestGenerateOutput: + 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_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 + + 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_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_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 +# --------------------------------------------------------------------------- + + +class TestSkipPath: + def test_skip_path_does_not_call_audio_writer_on_second_run(self, tmp_path: Path) -> None: + """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")) + 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: + """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")) + 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..39a09a8d 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: @@ -23,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": { @@ -44,6 +52,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, + }, }, } @@ -55,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) @@ -153,3 +185,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)