| """NumPy implementation of MuScriptor's 16 kHz log-mel frontend.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import numpy as np |
| import soundfile as sf |
| from scipy.signal import resample_poly |
|
|
|
|
| SAMPLE_RATE = 16_000 |
| N_FFT = 2_048 |
| HOP_LENGTH = 160 |
| N_MELS = 512 |
| SEGMENT_SAMPLES = 5 * SAMPLE_RATE |
|
|
|
|
| def load_audio_16k(path: str | Path) -> np.ndarray: |
| """Load an audio file as mono float32 at 16 kHz.""" |
| audio, sample_rate = sf.read(path, dtype="float32", always_2d=True) |
| audio = audio.mean(axis=1) |
| if sample_rate != SAMPLE_RATE: |
| divisor = np.gcd(sample_rate, SAMPLE_RATE) |
| audio = resample_poly( |
| audio, SAMPLE_RATE // divisor, sample_rate // divisor |
| ).astype(np.float32) |
| return audio.astype(np.float32, copy=False) |
|
|
|
|
| def first_five_second_chunk(audio: np.ndarray) -> np.ndarray: |
| """Crop/pad audio to the five-second chunk used by MuScriptor.""" |
| audio = np.asarray(audio, dtype=np.float32).reshape(-1) |
| if audio.size >= SEGMENT_SAMPLES: |
| return audio[:SEGMENT_SAMPLES] |
| return np.pad(audio, (0, SEGMENT_SAMPLES - audio.size)) |
|
|
|
|
| def _hz_to_mel_htk(freq: np.ndarray) -> np.ndarray: |
| return 2595.0 * np.log10(1.0 + freq / 700.0) |
|
|
|
|
| def _mel_to_hz_htk(mel: np.ndarray) -> np.ndarray: |
| return 700.0 * (10.0 ** (mel / 2595.0) - 1.0) |
|
|
|
|
| def mel_filterbank() -> np.ndarray: |
| """Match the pure-torch HTK filterbank bundled with MuScriptor.""" |
| all_freqs = np.linspace(0, SAMPLE_RATE // 2, N_FFT // 2 + 1, dtype=np.float32) |
| mel_min = _hz_to_mel_htk(np.asarray(0.0, dtype=np.float32)) |
| mel_max = _hz_to_mel_htk(np.asarray(SAMPLE_RATE / 2, dtype=np.float32)) |
| mel_points = np.linspace(mel_min, mel_max, N_MELS + 2, dtype=np.float32) |
| freq_points = _mel_to_hz_htk(mel_points) |
| freq_diff = freq_points[1:] - freq_points[:-1] |
| slopes = freq_points[None, :] - all_freqs[:, None] |
| down = -slopes[:, :-2] / freq_diff[:-1] |
| up = slopes[:, 2:] / freq_diff[1:] |
| return np.maximum(0.0, np.minimum(down, up)).astype(np.float32) |
|
|
|
|
| def log_mel_spectrogram(audio: np.ndarray) -> np.ndarray: |
| """Return `[1, 501, 512]` log-mel features for a five-second chunk. |
| |
| The reflection padding, periodic Hann window, FFT, HTK mel bank, and log |
| epsilon mirror ``muscriptor.modules.mel_spectrogram``. |
| """ |
| audio = first_five_second_chunk(audio) |
| padded = np.pad(audio, (N_FFT // 2, N_FFT // 2), mode="reflect") |
| frames = np.lib.stride_tricks.sliding_window_view(padded, N_FFT)[::HOP_LENGTH] |
| window = np.hanning(N_FFT + 1)[:-1].astype(np.float32) |
| spectrum = np.abs(np.fft.rfft(frames * window, n=N_FFT, axis=-1)).astype(np.float32) |
| mel = spectrum @ mel_filterbank() |
| return np.log(mel + np.float32(1e-6))[None, ...].astype(np.float32) |
|
|
|
|
