src/data/augmentation.py

"""
Field noise augmentation for West African farm environments.
Mixes clean speech with tractor, wind, and livestock audio samples.
Degrades gracefully to Gaussian noise when no .wav files are present.
"""
from __future__ import annotations

import logging
from pathlib import Path

import numpy as np

logger = logging.getLogger(__name__)


class FieldNoiseAugmenter:
    """
    Applies audiomentations transforms that simulate noisy field conditions.
    If the noise_dir has no .wav files, falls back to Gaussian noise only.
    """

    def __init__(self, noise_dir: str, config: dict) -> None:
        self.noise_dir = Path(noise_dir)
        self.config = config
        self._compose = None
        self._gaussian_only = False
        self._build_pipeline()

    def _build_pipeline(self) -> None:
        try:
            from audiomentations import (
                AddBackgroundNoise,
                AddGaussianNoise,
                Compose,
                RoomSimulator,
                TimeStretch,
            )
        except ImportError:
            logger.warning("audiomentations not installed — augmentation disabled.")
            self._compose = None
            return

        snr_range = self.config.get("audio", {}).get("noise_snr_db_range", [5, 20])
        prob = self.config.get("audio", {}).get("augmentation_prob", 0.6)

        wav_files = list(self.noise_dir.glob("*.wav")) if self.noise_dir.exists() else []

        transforms = []

        if wav_files:
            transforms.append(
                AddBackgroundNoise(
                    sounds_path=str(self.noise_dir),
                    min_snr_db=float(snr_range[0]),
                    max_snr_db=float(snr_range[1]),
                    p=prob,
                )
            )
            logger.info("FieldNoiseAugmenter: loaded %d noise files from %s", len(wav_files), self.noise_dir)
        else:
            logger.warning(
                "FieldNoiseAugmenter: no .wav files found in %s — using Gaussian noise only. "
                "Populate noise_samples/ for realistic field augmentation.",
                self.noise_dir,
            )
            self._gaussian_only = True

        transforms += [
            AddGaussianNoise(min_amplitude=0.001, max_amplitude=0.015, p=0.3),
            TimeStretch(min_rate=0.9, max_rate=1.1, leave_length_unchanged=True, p=0.2),
            RoomSimulator(p=0.3),
        ]

        self._compose = Compose(transforms)

    def augment(self, audio: np.ndarray, sr: int) -> np.ndarray:
        """Apply augmentation pipeline to a float32 audio array."""
        if self._compose is None:
            return audio
        return self._compose(samples=audio, sample_rate=sr)

    def is_ready(self) -> bool:
        """Returns True if augmentation is available (even Gaussian-only)."""
        return self._compose is not None