"""Audio Loader Module. Loads and preprocesses WAV audio files for the similarity pipeline. Notes: - This pipeline uses a strict single-backend mechanism (torchaudio) for reading audio. This methodological decision ensures strict experimental reproducibility by avoiding backend-dependent variations in decoding or numerical representation. - Optional WebRTC VAD-based endpoint trimming removes leading/trailing silence without cutting internal pauses. """ from __future__ import annotations import struct from typing import Tuple import numpy as np import torch import torchaudio try: import soundfile as sf # type: ignore _SOUNDFILE_AVAILABLE = True except ImportError: _SOUNDFILE_AVAILABLE = False try: import webrtcvad # type: ignore _WEBRTCVAD_AVAILABLE = True except ImportError: _WEBRTCVAD_AVAILABLE = False class AudioLoader: """Loads and preprocesses audio files.""" def __init__( self, target_sr: int = 16000, use_vad: bool = True, vad_mode: int = 1, vad_frame_ms: int = 10, vad_onset_frames: int = 2, vad_offset_frames: int = 4, energy_trim_threshold: float = 0.06, ): """ Initialize AudioLoader. Args: target_sr: Target sampling rate (default: 16000 Hz for WavLM). use_vad: If True, apply WebRTC VAD endpoint trimming to strip leading/trailing silence. vad_mode: WebRTC VAD aggressiveness (0–3). 1–2 recommended. vad_frame_ms: Frame length in ms for VAD (10, 20, or 30). vad_onset_frames: Minimum consecutive voiced frames to detect speech onset (hysteresis). vad_offset_frames: Minimum consecutive unvoiced frames to detect speech offset (hysteresis). energy_trim_threshold: Fraction of peak absolute amplitude used as threshold for energy-based refinement after VAD (0.0–1.0). Frames with amplitude below this fraction of the local peak are trimmed from the head and tail. """ self.target_sr = target_sr self.use_vad = use_vad self.vad_mode = vad_mode self.vad_frame_ms = vad_frame_ms self.vad_onset_frames = vad_onset_frames self.vad_offset_frames = vad_offset_frames self.energy_trim_threshold = energy_trim_threshold @staticmethod def _load_with_soundfile(audio_path: str) -> Tuple[torch.Tensor, int]: """Fallback loader using soundfile (libsndfile). Used when torchaudio's default backend (torchcodec) is unavailable, e.g. on systems without FFmpeg DLLs. Produces identical float32 waveform tensors for WAV files. """ if not _SOUNDFILE_AVAILABLE: raise ImportError( "soundfile is not installed. Install it with: pip install soundfile" ) data, sample_rate = sf.read(audio_path, dtype="float32", always_2d=True) # soundfile returns (num_samples, channels) → transpose to (channels, num_samples) waveform = torch.from_numpy(data.T) return waveform, sample_rate def load_audio(self, audio_path: str) -> Tuple[torch.Tensor, int]: """ Load audio file and resample if necessary. Uses soundfile (libsndfile) as the primary loader because torchaudio ≥2.11 requires torchcodec + FFmpeg DLLs which are often unavailable on Windows. soundfile handles WAV natively and produces identical float32 results. Falls back to torchaudio.load() if soundfile is unavailable. Args: audio_path: Path to WAV audio file Returns: Tuple of (waveform tensor, sample_rate) waveform shape: (1, num_samples) for mono Raises: RuntimeError: If both soundfile and torchaudio fail to read the audio file. """ # Primary: soundfile (no FFmpeg dependency, handles WAV natively) try: waveform, sample_rate = self._load_with_soundfile(audio_path) except Exception as sf_err: # Fallback: torchaudio (needs torchcodec/FFmpeg on ≥2.11) try: waveform, sample_rate = torchaudio.load(audio_path) except Exception as torchaudio_err: raise RuntimeError( f"Failed to read audio file '{audio_path}'. " f"soundfile error: {sf_err} | " f"torchaudio error: {torchaudio_err}" ) # Convert to mono if stereo if waveform.shape[0] > 1: waveform = torch.mean(waveform, dim=0, keepdim=True) # Resample if necessary if sample_rate != self.target_sr: resampler = torchaudio.transforms.Resample( orig_freq=sample_rate, new_freq=self.target_sr ) waveform = resampler(waveform) sample_rate = self.target_sr return waveform, sample_rate # ------------------------------------------------------------------ # VAD-based endpoint trimming # ------------------------------------------------------------------ def vad_trim_endpoints( self, waveform: torch.Tensor, sample_rate: int ) -> torch.Tensor: """Trim leading and trailing silence using WebRTC VAD. Only the outermost silent segments are removed; internal pauses (e.g. between words or ayat) are preserved intact. The algorithm uses a hysteresis mechanism: * **Onset**: the first sample of the first window in a run of at least ``vad_onset_frames`` consecutive *voiced* frames is taken as the speech start. * **Offset**: the last sample of the last *voiced* frame before a run of at least ``vad_offset_frames`` consecutive *unvoiced* frames that extends to the end of the signal is taken as the speech end. Args: waveform: Mono waveform tensor of shape ``(1, num_samples)`` in float32 (values in roughly [-1, 1]). sample_rate: Must be 8000, 16000, 32000, or 48000. Returns: Trimmed waveform tensor ``(1, trimmed_samples)``, float32. """ if not _WEBRTCVAD_AVAILABLE: raise ImportError( "webrtcvad is not installed but VAD trimming was requested. " "Please install it with: pip install webrtcvad" ) # --- prepare int16 PCM bytes for webrtcvad ----------------------- audio_np: np.ndarray = waveform.squeeze(0).numpy() # (num_samples,) # Clip and convert float32 -> int16 pcm_int16 = np.clip(audio_np * 32767, -32768, 32767).astype(np.int16) frame_len = int(sample_rate * self.vad_frame_ms / 1000) # samples per frame num_frames = len(pcm_int16) // frame_len if num_frames == 0: return waveform # too short to analyse vad = webrtcvad.Vad(self.vad_mode) # --- per-frame voiced/unvoiced labels ---------------------------- is_voiced: list[bool] = [] for i in range(num_frames): start = i * frame_len end = start + frame_len frame_bytes = struct.pack(f"<{frame_len}h", *pcm_int16[start:end]) is_voiced.append(vad.is_speech(frame_bytes, sample_rate)) # --- onset detection (left-to-right) ----------------------------- speech_start_frame: int = 0 consecutive_voiced = 0 for idx, voiced in enumerate(is_voiced): if voiced: consecutive_voiced += 1 if consecutive_voiced >= self.vad_onset_frames: speech_start_frame = idx - self.vad_onset_frames + 1 break else: consecutive_voiced = 0 else: # No onset found – return original waveform (all silence?) return waveform # --- offset detection (right-to-left) ---------------------------- trailing_unvoiced = 0 for idx in range(num_frames - 1, -1, -1): if not is_voiced[idx]: trailing_unvoiced += 1 else: break if trailing_unvoiced >= self.vad_offset_frames: speech_end_frame = num_frames - trailing_unvoiced - 1 else: speech_end_frame = num_frames - 1 # no significant trailing silence # Sanity: ensure start <= end if speech_start_frame > speech_end_frame: return waveform # --- reconstruct trimmed waveform -------------------------------- start_sample = speech_start_frame * frame_len end_sample = (speech_end_frame + 1) * frame_len end_sample = min(end_sample, waveform.shape[1]) trimmed = waveform[:, start_sample:end_sample] # Guard against empty result if trimmed.shape[1] == 0: return waveform # --- energy-based refinement ------------------------------------- trimmed = self._energy_refine(trimmed) return trimmed def _energy_refine(self, waveform: torch.Tensor) -> torch.Tensor: """Refine VAD boundaries by trimming low-energy head/tail. Uses a short-window RMS envelope to find the first and last sample whose energy exceeds ``energy_trim_threshold`` of the peak RMS value. This removes residual breath or resonance that WebRTC VAD may leave behind. Args: waveform: Mono waveform ``(1, num_samples)``. Returns: Refined waveform ``(1, refined_samples)``. """ if self.energy_trim_threshold <= 0.0: return waveform audio = waveform.squeeze(0).numpy() # (num_samples,) abs_env = np.abs(audio) # Short-window RMS (window ≈ 10 ms) win = max(int(self.target_sr * 0.01), 1) # Cumulative sum trick for fast moving average of squared signal sq = audio.astype(np.float64) ** 2 cs = np.concatenate(([0.0], np.cumsum(sq))) rms = np.sqrt((cs[win:] - cs[:-win]) / win).astype(np.float32) if len(rms) == 0: return waveform peak_rms = rms.max() if peak_rms == 0: return waveform threshold = self.energy_trim_threshold * peak_rms # Find first and last frame above threshold above = np.where(rms >= threshold)[0] if len(above) == 0: return waveform start = int(above[0]) end = int(above[-1]) + win # include the window trailing edge end = min(end, len(audio)) refined = waveform[:, start:end] if refined.shape[1] == 0: return waveform return refined def normalize_audio(self, waveform: torch.Tensor) -> torch.Tensor: """ Normalize audio waveform to [-1, 1] range. Args: waveform: Audio waveform tensor Returns: Normalized waveform """ # Normalize to [-1, 1] max_val = torch.max(torch.abs(waveform)) if max_val > 0: waveform = waveform / max_val return waveform def load_and_preprocess(self, audio_path: str) -> torch.Tensor: """ Load and preprocess audio in one step. Args: audio_path: Path to WAV audio file Returns: Preprocessed waveform tensor (1, num_samples) """ waveform, sr = self.load_audio(audio_path) # VAD endpoint trimming (after resampling, before normalisation) if self.use_vad: waveform = self.vad_trim_endpoints(waveform, sr) waveform = self.normalize_audio(waveform) return waveform