| """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) |
| 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() |
|
|
| |
| 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) |
|
|
| |
| lo, hi = float(arr.min()), float(arr.max()) |
| rng = max(hi - lo, 1e-6) |
| arr = (arr - lo) / rng |
|
|
| return arr.astype(np.float32).tolist() |
|
|
