File size: 2,526 Bytes
eca55dc | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import torch
import torch.nn as nn
import torchaudio
from typing import Optional
class Spectrogram(nn.Module):
"""
Mel Spectrogram module with AmplitudeToDB conversion.
Args:
sample_rate (int): Sample rate of the audio.
n_fft (int): Size of FFT.
win_length (Optional[int]): Window length. Defaults to n_fft.
win_length_ms (Optional[float]): Window length in milliseconds. Overrides win_length if provided.
hop_length (Optional[int]): Hop length. Defaults to win_length // 2.
hop_length_ms (Optional[float]): Hop length in milliseconds. Overrides hop_length if provided.
n_mels (int): Number of mel filterbanks.
f_min (float): Minimum frequency.
f_max (Optional[float]): Maximum frequency.
power (float): Power of the magnitude.
"""
def __init__(
self,
sample_rate: int = 32000,
n_fft: int = 4096,
win_length: Optional[int] = None,
win_length_ms: Optional[float] = None,
hop_length: Optional[int] = None,
hop_length_ms: Optional[float] = None,
n_mels: int = 128,
f_min: float = 0.0,
f_max: Optional[float] = None,
power: float = 2.0,
):
super().__init__()
if win_length is None:
if win_length_ms is None:
win_length = n_fft
else:
win_length = int(sample_rate * win_length_ms / 1000)
if hop_length is None:
if hop_length_ms is None:
hop_length = win_length // 2
else:
hop_length = int(sample_rate * hop_length_ms / 1000)
self.mel_spec = torchaudio.transforms.MelSpectrogram(
sample_rate=sample_rate,
n_fft=n_fft,
win_length=win_length,
hop_length=hop_length,
n_mels=n_mels,
f_min=f_min,
f_max=f_max,
power=power,
normalized=True,
)
self.amplitude_to_db = torchaudio.transforms.AmplitudeToDB()
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass.
Args:
x (torch.Tensor): Input waveform [B, C, T] or [B, T].
Returns:
torch.Tensor: Log-Mel Spectrogram [B, C, F, T].
"""
# x: [B, C, T]
# MelSpectrogram expects [..., T]
# Output will be [..., n_mels, time]
spec = self.mel_spec(x)
spec = self.amplitude_to_db(spec)
return spec
|