| """Torch-native waveform to log-mel feature extraction. |
| |
| The frontend intentionally avoids torchaudio so the deployable student has one |
| fewer binary dependency. Exported production models normally accept log-mel |
| features; keeping this implementation in the repository gives training, demo, |
| and parity tests one canonical preprocessing contract. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Mapping |
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import Tensor, nn |
|
|
|
|
| @dataclass(frozen=True) |
| class LogMelConfig: |
| sample_rate: int = 16_000 |
| n_fft: int = 400 |
| hop_length: int = 160 |
| win_length: int = 400 |
| n_mels: int = 80 |
| f_min: float = 0.0 |
| f_max: float | None = 8_000.0 |
| log_floor: float = 1e-10 |
| normalize: bool = True |
| mel_scale: str = "htk" |
| log_scale: str = "standard" |
| center: bool = False |
| drop_last_frame: bool = False |
| pad_side: str = "right" |
|
|
| @classmethod |
| def from_mapping(cls, values: Mapping[str, object]) -> LogMelConfig: |
| known = {field.name for field in cls.__dataclass_fields__.values()} |
| return cls(**{key: value for key, value in values.items() if key in known}) |
|
|
|
|
| def _hz_to_mel(freq: Tensor) -> Tensor: |
| |
| |
| return 2595.0 * torch.log10(1.0 + freq / 700.0) |
|
|
|
|
| def _mel_to_hz(mels: Tensor) -> Tensor: |
| return 700.0 * (torch.pow(10.0, mels / 2595.0) - 1.0) |
|
|
|
|
| def _hz_to_slaney_mel(freq: Tensor) -> Tensor: |
| linear_spacing = 200.0 / 3.0 |
| min_log_hz = 1_000.0 |
| min_log_mel = min_log_hz / linear_spacing |
| log_step = torch.log(torch.tensor(6.4, dtype=freq.dtype, device=freq.device)) / 27.0 |
| linear = freq / linear_spacing |
| logarithmic = min_log_mel + torch.log((freq / min_log_hz).clamp_min(1e-12)) / log_step |
| return torch.where(freq >= min_log_hz, logarithmic, linear) |
|
|
|
|
| def _slaney_mel_to_hz(mels: Tensor) -> Tensor: |
| linear_spacing = 200.0 / 3.0 |
| min_log_hz = 1_000.0 |
| min_log_mel = min_log_hz / linear_spacing |
| log_step = torch.log(torch.tensor(6.4, dtype=mels.dtype, device=mels.device)) / 27.0 |
| linear = mels * linear_spacing |
| logarithmic = min_log_hz * torch.exp(log_step * (mels - min_log_mel)) |
| return torch.where(mels >= min_log_mel, logarithmic, linear) |
|
|
|
|
| def make_mel_filterbank(config: LogMelConfig) -> Tensor: |
| """Create a ``[n_mels, n_fft // 2 + 1]`` triangular filterbank.""" |
|
|
| max_hz = float(config.sample_rate / 2 if config.f_max is None else config.f_max) |
| if not 0.0 <= config.f_min < max_hz <= config.sample_rate / 2: |
| raise ValueError("expected 0 <= f_min < f_max <= sample_rate / 2") |
|
|
| fft_freqs = torch.linspace(0.0, config.sample_rate / 2, config.n_fft // 2 + 1) |
| if config.mel_scale == "htk": |
| to_mel, to_hz = _hz_to_mel, _mel_to_hz |
| elif config.mel_scale == "slaney": |
| to_mel, to_hz = _hz_to_slaney_mel, _slaney_mel_to_hz |
| else: |
| raise ValueError("mel_scale must be 'htk' or 'slaney'") |
| mel_edges = torch.linspace( |
| to_mel(torch.tensor(float(config.f_min))), |
| to_mel(torch.tensor(max_hz)), |
| config.n_mels + 2, |
| ) |
| hz_edges = to_hz(mel_edges) |
| lower = hz_edges[:-2, None] |
| center = hz_edges[1:-1, None] |
| upper = hz_edges[2:, None] |
| rising = (fft_freqs[None, :] - lower) / (center - lower).clamp_min(1e-12) |
| falling = (upper - fft_freqs[None, :]) / (upper - center).clamp_min(1e-12) |
| filters = torch.minimum(rising, falling).clamp_min(0.0) |
|
|
| |
| enorm = 2.0 / (upper - lower).clamp_min(1e-12) |
| return filters * enorm |
|
|
|
|
| class LogMelFrontend(nn.Module): |
| """Convert padded mono waveforms to normalized log-mel features. |
| |
| Parameters |
| ---------- |
| waveforms: |
| Float tensor shaped ``[batch, samples]`` (or ``[samples]``). |
| lengths: |
| Optional valid sample counts. The returned mask is ``[batch, frames]``. |
| """ |
|
|
| def __init__(self, config: LogMelConfig | None = None) -> None: |
| super().__init__() |
| config = config or LogMelConfig() |
| self.config = config |
| if config.pad_side not in {"left", "right"}: |
| raise ValueError("pad_side must be 'left' or 'right'") |
| |
| self.register_buffer( |
| "window", torch.hann_window(config.win_length, periodic=False), persistent=False |
| ) |
| self.register_buffer("mel_filters", make_mel_filterbank(config), persistent=True) |
|
|
| def forward(self, waveforms: Tensor, lengths: Tensor | None = None) -> tuple[Tensor, Tensor]: |
| if waveforms.ndim == 1: |
| waveforms = waveforms.unsqueeze(0) |
| if waveforms.ndim != 2: |
| raise ValueError("waveforms must have shape [batch, samples]") |
|
|
| batch, original_samples = waveforms.shape |
| if lengths is None: |
| lengths = torch.full( |
| (batch,), original_samples, dtype=torch.long, device=waveforms.device |
| ) |
| else: |
| lengths = lengths.to(device=waveforms.device, dtype=torch.long).clamp( |
| min=0, max=original_samples |
| ) |
| if original_samples < self.config.n_fft: |
| waveforms = F.pad(waveforms, (0, self.config.n_fft - original_samples)) |
|
|
| spectrum = torch.stft( |
| waveforms, |
| n_fft=self.config.n_fft, |
| hop_length=self.config.hop_length, |
| win_length=self.config.win_length, |
| window=self.window.to(dtype=waveforms.dtype), |
| center=self.config.center, |
| return_complex=True, |
| ) |
| if self.config.drop_last_frame: |
| spectrum = spectrum[..., :-1] |
| power = spectrum.abs().square() |
| mel = torch.matmul(self.mel_filters.to(dtype=power.dtype), power) |
| log_mel = torch.log10(mel.clamp_min(self.config.log_floor)) |
| if self.config.log_scale == "whisper": |
| dynamic_floor = log_mel.amax(dim=(-2, -1), keepdim=True) - 8.0 |
| log_mel = torch.maximum(log_mel, dynamic_floor) |
| log_mel = (log_mel + 4.0) / 4.0 |
| elif self.config.log_scale != "standard": |
| raise ValueError("log_scale must be 'standard' or 'whisper'") |
|
|
| if self.config.center: |
| if self.config.drop_last_frame: |
| frame_lengths = torch.div( |
| lengths + self.config.hop_length - 1, |
| self.config.hop_length, |
| rounding_mode="floor", |
| ) |
| else: |
| frame_lengths = 1 + torch.div( |
| lengths, self.config.hop_length, rounding_mode="floor" |
| ) |
| else: |
| padded_lengths = lengths.clamp_min(self.config.n_fft) |
| frame_lengths = 1 + torch.div( |
| padded_lengths - self.config.n_fft, |
| self.config.hop_length, |
| rounding_mode="floor", |
| ) |
| if self.config.drop_last_frame and not self.config.center: |
| frame_lengths = (frame_lengths - 1).clamp_min(0) |
| frame_lengths = torch.where(lengths > 0, frame_lengths, torch.zeros_like(frame_lengths)) |
| frame_lengths = frame_lengths.clamp(max=log_mel.shape[-1]) |
| positions = torch.arange(log_mel.shape[-1], device=waveforms.device) |
| if self.config.pad_side == "left": |
| mask = positions.unsqueeze(0) >= (log_mel.shape[-1] - frame_lengths).unsqueeze(1) |
| else: |
| mask = positions.unsqueeze(0) < frame_lengths.unsqueeze(1) |
|
|
| if self.config.normalize: |
| valid = mask.unsqueeze(1).to(log_mel.dtype) |
| denominator = valid.sum(dim=-1, keepdim=True).clamp_min(1.0) |
| mean = (log_mel * valid).sum(dim=-1, keepdim=True) / denominator |
| variance = ((log_mel - mean).square() * valid).sum(dim=-1, keepdim=True) |
| variance = variance / denominator |
| log_mel = (log_mel - mean) * torch.rsqrt(variance + 1e-5) |
| log_mel = log_mel * valid |
|
|
| return log_mel, mask |
|
|
