File size: 1,675 Bytes
f6aec75 | 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 | """Log-mel feature frontend.
40 log-mel bands, 16 kHz, 25 ms window / 10 ms hop. Deliberately matches
what is cheap to compute on the ESP32-S3 later (esp-dsp / TFLite-Micro
audio frontend), so the on-device features can be made to line up with
training.
"""
import torch
import torchaudio
SAMPLE_RATE = 16000
N_MELS = 40
WIN_LENGTH = 400 # 25 ms
HOP_LENGTH = 160 # 10 ms
N_FFT = 512
class LogMel(torch.nn.Module):
def __init__(self):
super().__init__()
self.mel = torchaudio.transforms.MelSpectrogram(
sample_rate=SAMPLE_RATE,
n_fft=N_FFT,
win_length=WIN_LENGTH,
hop_length=HOP_LENGTH,
n_mels=N_MELS,
center=True,
power=2.0,
)
EMA_ALPHA = 0.02 # ~0.5 s time constant at 10 ms hop
def forward(self, waveform):
"""waveform (B, samples) -> features (B, T, N_MELS).
Normalization is a causal per-channel EMA mean subtraction so the
exact same computation can run frame-by-frame on the device.
"""
mel = self.mel(waveform) # (B, n_mels, T)
logmel = torch.log(mel + 1e-6)
a = self.EMA_ALPHA
ema = torchaudio.functional.lfilter(
logmel,
a_coeffs=torch.tensor([1.0, -(1.0 - a)], device=logmel.device),
b_coeffs=torch.tensor([a, 0.0], device=logmel.device),
clamp=False,
)
logmel = logmel - ema
return logmel.transpose(1, 2) # (B, T, n_mels)
def num_frames(num_samples):
"""Frame count produced for a waveform length (center=True)."""
return num_samples // HOP_LENGTH + 1
|