melspec.py

"""Log-mel spectrograms, in one place.

Training and inference must compute this identically. When they do not, the
model sees a slightly different picture in production than it was taught on and
degrades in a way that looks like a bad model rather than a bad pipeline — so
there is one implementation and both sides import it.

Plain numpy, on purpose. The alternative is librosa, which pulls in a large
dependency tree to compute forty lines of arithmetic, and this has to install
cleanly on a small droplet.
"""

from __future__ import annotations

import subprocess

import numpy as np

SR = 16_000
N_FFT = 400          # 25 ms
HOP = 320            # 20 ms, so 50 frames a second
N_MELS = 64
FMIN, FMAX = 50.0, 7600.0
FPS = SR / HOP


def _filterbank(sr=SR, n_fft=N_FFT, n_mels=N_MELS, fmin=FMIN, fmax=FMAX):
    def to_mel(f):
        return 2595.0 * np.log10(1.0 + f / 700.0)

    def to_hz(m):
        return 700.0 * (10.0 ** (m / 2595.0) - 1.0)

    points = to_hz(np.linspace(to_mel(fmin), to_mel(fmax), n_mels + 2))
    bins = np.floor((n_fft + 1) * points / sr).astype(int)
    bank = np.zeros((n_mels, n_fft // 2 + 1), dtype=np.float32)
    for i in range(n_mels):
        left, centre, right = bins[i], bins[i + 1], bins[i + 2]
        centre = max(centre, left + 1)
        right = max(right, centre + 1)
        bank[i, left:centre] = np.linspace(0, 1, centre - left, endpoint=False)
        bank[i, centre:right] = np.linspace(1, 0, right - centre, endpoint=False)
    return bank


BANK = _filterbank()


def logmel(samples: np.ndarray) -> np.ndarray:
    """(frames, 64) log-mel from mono float samples at 16 kHz."""
    samples = np.ascontiguousarray(samples, dtype=np.float32)
    if len(samples) < N_FFT:
        return np.zeros((0, N_MELS), dtype=np.float32)
    frames = 1 + (len(samples) - N_FFT) // HOP
    window = np.hanning(N_FFT).astype(np.float32)
    strided = np.lib.stride_tricks.as_strided(
        samples, shape=(frames, N_FFT),
        strides=(samples.strides[0] * HOP, samples.strides[0]))
    spectrum = np.fft.rfft(strided * window, axis=1)
    power = (spectrum.real ** 2 + spectrum.imag ** 2).astype(np.float32)
    # The float32 matmul raises divide-by-zero, overflow and invalid on macOS
    # regardless of the values — Accelerate's vectorised kernel sets the FP
    # status flags from lanes past the end of the data. Checked on real audio:
    # the largest power in a chunk is around 27 and every output is finite.
    with np.errstate(all="ignore"):
        energy = power @ BANK.T
    return np.log(energy + 1e-6).astype(np.float32)


def decode(path: str) -> np.ndarray:
    """Any audio file to mono float32 at 16 kHz, via ffmpeg."""
    raw = subprocess.run(
        ["ffmpeg", "-v", "error", "-i", str(path), "-f", "f32le",
         "-ac", "1", "-ar", str(SR), "-"],
        capture_output=True, check=True).stdout
    return np.frombuffer(raw, dtype=np.float32)