Extend feature extraction module to allow for RASR compatible logmel features

i6_models/primitives/feature_extraction.py

@@ -1,7 +1,7 @@
 __all__ = ["LogMelFeatureExtractionV1", "LogMelFeatureExtractionV1Config"]
 
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Any, Dict
 
 from librosa import filters
 import torch
@@ -22,6 +22,9 @@ class LogMelFeatureExtractionV1Config(ModelConfiguration):
         min_amp: minimum amplitude for safe log
         num_filters: number of mel windows
         center: centered STFT with automatic padding
+        periodic: whether the window is assumed to be periodic
+        mel_options: extra options for mel filters
+        rasr_compatible: apply FFT to make features compatible to RASR's
     """
 
     sample_rate: int
@@ -33,6 +36,9 @@ class LogMelFeatureExtractionV1Config(ModelConfiguration):
     num_filters: int
     center: bool
     n_fft: Optional[int] = None
+    periodic: Optional[bool] = True
+    mel_options: Optional[Dict[str, Any]] = None
+    rasr_compatible: Optional[bool] = False
 
     def __post_init__(self) -> None:
         super().__post_init__()
@@ -62,6 +68,8 @@ def __init__(self, cfg: LogMelFeatureExtractionV1Config):
         self.min_amp = cfg.min_amp
         self.n_fft = cfg.n_fft
         self.win_length = int(cfg.win_size * cfg.sample_rate)
+        self.mel_options = cfg.mel_options or {}
+        self.rasr_compatible = cfg.rasr_compatible
 
         self.register_buffer(
             "mel_basis",
@@ -72,32 +80,42 @@ def __init__(self, cfg: LogMelFeatureExtractionV1Config):
                     n_mels=cfg.num_filters,
                     fmin=cfg.f_min,
                     fmax=cfg.f_max,
+                    **self.mel_options,
                 )
             ),
         )
-        self.register_buffer("window", torch.hann_window(self.win_length))
+        self.register_buffer("window", torch.hann_window(self.win_length, periodic=cfg.periodic))
 
     def forward(self, raw_audio, length) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         :param raw_audio: [B, T]
         :param length in samples: [B]
         :return features as [B,T,F] and length in frames [B]
         """
-        power_spectrum = (
-            torch.abs(
-                torch.stft(
-                    raw_audio,
-                    n_fft=self.n_fft,
-                    hop_length=self.hop_length,
-                    win_length=self.win_length,
-                    window=self.window,
-                    center=self.center,
-                    pad_mode="constant",
-                    return_complex=True,
+        if self.rasr_compatible:
+            windowed = raw_audio.unfold(1, size=self.win_length, step=self.hop_length)
+            smoothed = windowed * self.window.unsqueeze(0)
+
+            # Compute power spectrum using torch.fft.rfftn
+            power_spectrum = torch.abs(torch.fft.rfftn(smoothed, s=self.n_fft)) ** 2  # [B, F, T]
+            power_spectrum = power_spectrum.transpose(1, 2)  # [B, T, F]
+        else:
+            power_spectrum = (
+                torch.abs(
+                    torch.stft(
+                        raw_audio,
+                        n_fft=self.n_fft,
+                        hop_length=self.hop_length,
+                        win_length=self.win_length,
+                        window=self.window,
+                        center=self.center,
+                        pad_mode="constant",
+                        return_complex=True,
+                    )
                 )
+                ** 2
             )
-            ** 2
-        )
+
         if len(power_spectrum.size()) == 2:
             # For some reason torch.stft removes the batch axis for batch sizes of 1, so we need to add it again
             power_spectrum = torch.unsqueeze(power_spectrum, 0)