File size: 1,982 Bytes
e6a1f55 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | """Mel spectrogram utilities for the UI and BiCrossMamba-ST input."""
from __future__ import annotations
from typing import List
import numpy as np
import torch
import torchaudio
_MEL_TRANSFORM_CACHE: dict = {}
def _get_mel(sr: int, n_mels: int, n_fft: int, hop_length: int) -> torchaudio.transforms.MelSpectrogram:
key = (sr, n_mels, n_fft, hop_length)
if key not in _MEL_TRANSFORM_CACHE:
_MEL_TRANSFORM_CACHE[key] = torchaudio.transforms.MelSpectrogram(
sample_rate=sr,
n_fft=n_fft,
hop_length=hop_length,
n_mels=n_mels,
f_min=0.0,
f_max=sr // 2,
power=2.0,
)
return _MEL_TRANSFORM_CACHE[key]
def mel_spectrogram(
waveform: torch.Tensor,
sample_rate: int = 16000,
n_mels: int = 64,
n_fft: int = 1024,
hop_length: int = 256,
) -> torch.Tensor:
"""Compute a log-mel spectrogram. Returns shape [n_mels, T]."""
mel = _get_mel(sample_rate, n_mels, n_fft, hop_length)
spec = mel(waveform) # [1, n_mels, T]
log_spec = torch.log(spec + 1e-6)
return log_spec.squeeze(0)
def mel_for_ui(
waveform: torch.Tensor,
sample_rate: int = 16000,
n_mels: int = 64,
n_fft: int = 1024,
hop_length: int = 256,
max_time_steps: int = 256,
) -> List[List[float]]:
"""Return a normalised [0,1] list-of-lists suitable for canvas rendering."""
log_spec = mel_spectrogram(waveform, sample_rate, n_mels, n_fft, hop_length)
arr = log_spec.detach().cpu().numpy()
# Down-sample time axis if needed
T = arr.shape[1]
if T > max_time_steps:
bucket = T // max_time_steps
trimmed = arr[:, : bucket * max_time_steps]
arr = trimmed.reshape(arr.shape[0], max_time_steps, bucket).mean(axis=2)
# Per-clip min-max normalize
lo, hi = float(arr.min()), float(arr.max())
rng = max(hi - lo, 1e-6)
arr = (arr - lo) / rng
return arr.astype(np.float32).tolist()
|