| import os |
| import sys |
|
|
| import torch |
| import torchaudio |
| from torch import nn |
| current_dir = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.append(current_dir) |
| print(sys.path) |
| from common import safe_log |
|
|
|
|
| class FeatureExtractor(nn.Module): |
| """Base class for feature extractors.""" |
|
|
| def forward(self, audio: torch.Tensor, **kwargs) -> torch.Tensor: |
| """ |
| Extract features from the given audio. |
| |
| Args: |
| audio (Tensor): Input audio waveform. |
| |
| Returns: |
| Tensor: Extracted features of shape (B, C, L), where B is the batch size, |
| C denotes output features, and L is the sequence length. |
| """ |
| raise NotImplementedError("Subclasses must implement the forward method.") |
|
|
|
|
| class MelSpectrogramFeatures(FeatureExtractor): |
| def __init__(self, sample_rate=24000, n_fft=1024, hop_length=256, win_length=None, |
| n_mels=100, mel_fmin=0, mel_fmax=None, normalize=False, padding="center"): |
| super().__init__() |
| if padding not in ["center", "same"]: |
| raise ValueError("Padding must be 'center' or 'same'.") |
| self.padding = padding |
| self.mel_spec = torchaudio.transforms.MelSpectrogram( |
| sample_rate=sample_rate, |
| n_fft=n_fft, |
| hop_length=hop_length, |
| win_length=win_length, |
| power=1, |
| normalized=normalize, |
| f_min=mel_fmin, |
| f_max=mel_fmax, |
| n_mels=n_mels, |
| center=padding == "center", |
| ) |
|
|
| def forward(self, audio, **kwargs): |
| if self.padding == "same": |
| pad = self.mel_spec.win_length - self.mel_spec.hop_length |
| audio = torch.nn.functional.pad(audio, (pad // 2, pad // 2), mode="reflect") |
| mel = self.mel_spec(audio) |
| mel = safe_log(mel) |
| return mel |
