|
|
import torch |
|
|
import torch.nn as nn |
|
|
from torch import Tensor |
|
|
from torchaudio.transforms import MelScale |
|
|
|
|
|
|
|
|
class LinearSpectrogram(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
n_fft=2048, |
|
|
win_length=2048, |
|
|
hop_length=512, |
|
|
center=False, |
|
|
mode="pow2_sqrt", |
|
|
): |
|
|
super().__init__() |
|
|
|
|
|
self.n_fft = n_fft |
|
|
self.win_length = win_length |
|
|
self.hop_length = hop_length |
|
|
self.center = center |
|
|
self.mode = mode |
|
|
|
|
|
self.register_buffer("window", torch.hann_window(win_length)) |
|
|
|
|
|
def forward(self, y: Tensor) -> Tensor: |
|
|
if y.ndim == 3: |
|
|
y = y.squeeze(1) |
|
|
|
|
|
y = torch.nn.functional.pad( |
|
|
y.unsqueeze(1), |
|
|
( |
|
|
(self.win_length - self.hop_length) // 2, |
|
|
(self.win_length - self.hop_length + 1) // 2, |
|
|
), |
|
|
mode="reflect", |
|
|
).squeeze(1) |
|
|
dtype = y.dtype |
|
|
spec = torch.stft( |
|
|
y.float(), |
|
|
self.n_fft, |
|
|
hop_length=self.hop_length, |
|
|
win_length=self.win_length, |
|
|
window=self.window, |
|
|
center=self.center, |
|
|
pad_mode="reflect", |
|
|
normalized=False, |
|
|
onesided=True, |
|
|
return_complex=True, |
|
|
) |
|
|
spec = torch.view_as_real(spec) |
|
|
|
|
|
if self.mode == "pow2_sqrt": |
|
|
spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6) |
|
|
spec = spec.to(dtype) |
|
|
return spec |
|
|
|
|
|
|
|
|
class LogMelSpectrogram(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
sample_rate=44100, |
|
|
n_fft=2048, |
|
|
win_length=2048, |
|
|
hop_length=512, |
|
|
n_mels=128, |
|
|
center=False, |
|
|
f_min=0.0, |
|
|
f_max=None, |
|
|
): |
|
|
super().__init__() |
|
|
|
|
|
self.sample_rate = sample_rate |
|
|
self.n_fft = n_fft |
|
|
self.win_length = win_length |
|
|
self.hop_length = hop_length |
|
|
self.center = center |
|
|
self.n_mels = n_mels |
|
|
self.f_min = f_min |
|
|
self.f_max = f_max or sample_rate // 2 |
|
|
|
|
|
self.spectrogram = LinearSpectrogram(n_fft, win_length, hop_length, center) |
|
|
self.mel_scale = MelScale( |
|
|
self.n_mels, |
|
|
self.sample_rate, |
|
|
self.f_min, |
|
|
self.f_max, |
|
|
self.n_fft // 2 + 1, |
|
|
"slaney", |
|
|
"slaney", |
|
|
) |
|
|
|
|
|
def compress(self, x: Tensor) -> Tensor: |
|
|
return torch.log(torch.clamp(x, min=1e-5)) |
|
|
|
|
|
def decompress(self, x: Tensor) -> Tensor: |
|
|
return torch.exp(x) |
|
|
|
|
|
def forward(self, x: Tensor, return_linear: bool = False) -> Tensor: |
|
|
linear = self.spectrogram(x) |
|
|
x = self.mel_scale(linear) |
|
|
x = self.compress(x) |
|
|
|
|
|
if return_linear: |
|
|
return x, self.compress(linear) |
|
|
|
|
|
return x |
|
|
|
