| import torch |
| import torchaudio.functional as F |
| from torch import Tensor, nn |
| 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.return_complex = True |
|
|
| self.register_buffer("window", torch.hann_window(win_length), persistent=False) |
|
|
| 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) |
|
|
| spec = torch.stft( |
| y, |
| 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=self.return_complex, |
| ) |
|
|
| if self.return_complex: |
| spec = torch.view_as_real(spec) |
|
|
| if self.mode == "pow2_sqrt": |
| spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6) |
|
|
| 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 float(sample_rate // 2) |
|
|
| self.spectrogram = LinearSpectrogram(n_fft, win_length, hop_length, center) |
|
|
| fb = F.melscale_fbanks( |
| n_freqs=self.n_fft // 2 + 1, |
| f_min=self.f_min, |
| f_max=self.f_max, |
| n_mels=self.n_mels, |
| sample_rate=self.sample_rate, |
| norm="slaney", |
| mel_scale="slaney", |
| ) |
| self.register_buffer( |
| "fb", |
| fb, |
| persistent=False, |
| ) |
|
|
| 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 apply_mel_scale(self, x: Tensor) -> Tensor: |
| return torch.matmul(x.transpose(-1, -2), self.fb).transpose(-1, -2) |
|
|
| def forward( |
| self, x: Tensor, return_linear: bool = False, sample_rate: int = None |
| ) -> Tensor: |
| if sample_rate is not None and sample_rate != self.sample_rate: |
| x = F.resample(x, orig_freq=sample_rate, new_freq=self.sample_rate) |
|
|
| linear = self.spectrogram(x) |
| x = self.apply_mel_scale(linear) |
| x = self.compress(x) |
|
|
| if return_linear: |
| return x, self.compress(linear) |
|
|
| return x |
|
|
