spectrogram.py

# MIT License
# 
# Copyright (c) 2026 audio-embeddings contributors
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from typing import Optional

import torch
import torch.nn as nn
import torchaudio


class Spectrogram(nn.Module):
    """
    Mel-frequency audio representation with optional temporal derivatives.

    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.
        representation (str): ``log_mel`` or phase-aware ``complex_mel``.
        complex_log_dynamic_range_db (float): Retained complex-mel log range.
        add_delta (bool): Append one temporal derivative per base channel.
        add_delta_delta (bool): Append a second derivative per base channel.
        delta_win_length (int): Odd regression window used for derivatives.
    """

    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,
        representation: str = "log_mel",
        complex_log_dynamic_range_db: float = 80.0,
        add_delta: bool = False,
        add_delta_delta: bool = False,
        delta_win_length: int = 5,
    ):
        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)

        representation = representation.strip().lower().replace("-", "_")
        if representation not in {"log_mel", "complex_mel"}:
            raise ValueError(
                "representation must be 'log_mel' or 'complex_mel', "
                f"got {representation!r}"
            )
        if complex_log_dynamic_range_db <= 0.0:
            raise ValueError(
                "complex_log_dynamic_range_db must be positive, "
                f"got {complex_log_dynamic_range_db}"
            )
        if delta_win_length < 3 or delta_win_length % 2 == 0:
            raise ValueError(
                f"delta_win_length must be an odd integer >= 3, got {delta_win_length}"
            )

        self.representation = representation
        self.complex_log_dynamic_range_db = float(complex_log_dynamic_range_db)
        self.add_delta = bool(add_delta)
        self.add_delta_delta = bool(add_delta_delta)
        self.delta_win_length = int(delta_win_length)
        self.base_output_channels = 1 if representation == "log_mel" else 2
        derivative_orders = int(self.add_delta) + int(self.add_delta_delta)
        self.output_channels = self.base_output_channels * (1 + derivative_orders)

        if representation == "log_mel":
            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()
        else:
            # Keep the historical ``mel_spec`` attribute for sample-rate and
            # hop-length discovery in callbacks and HEAR adapters.
            self.mel_spec = torchaudio.transforms.Spectrogram(
                n_fft=n_fft,
                win_length=win_length,
                hop_length=hop_length,
                power=None,
                normalized=True,
            )
            self.mel_spec.sample_rate = sample_rate
            self.amplitude_to_db = None
            mel_fb = torchaudio.functional.melscale_fbanks(
                n_freqs=n_fft // 2 + 1,
                f_min=f_min,
                f_max=float(sample_rate / 2 if f_max is None else f_max),
                n_mels=n_mels,
                sample_rate=sample_rate,
                norm=None,
                mel_scale="htk",
            )
            self.register_buffer("mel_fb", mel_fb)

    def _complex_mel(self, x: torch.Tensor) -> torch.Tensor:
        if x.shape[1] != 1:
            raise ValueError(
                "complex_mel expects mono waveform input [B, 1, T], "
                f"got {tuple(x.shape)}"
            )

        complex_spec = self.mel_spec(x[:, 0])
        mel_fb = self.mel_fb.to(dtype=complex_spec.dtype)
        complex_mel = torch.matmul(
            complex_spec.transpose(-1, -2),
            mel_fb,
        ).transpose(-1, -2)

        magnitude = complex_mel.abs()
        eps = torch.finfo(magnitude.dtype).eps
        reference = magnitude.amax(dim=(-2, -1), keepdim=True)
        relative_magnitude = magnitude / reference.clamp_min(eps)
        floor_ratio = 10.0 ** (-self.complex_log_dynamic_range_db / 20.0)
        log_magnitude_db = 20.0 * torch.log10(relative_magnitude.clamp_min(floor_ratio))
        log_radius = (
            log_magnitude_db + self.complex_log_dynamic_range_db
        ) / self.complex_log_dynamic_range_db
        log_radius = log_radius.clamp_(0.0, 1.0)
        log_radius = torch.where(
            reference > eps, log_radius, torch.zeros_like(log_radius)
        )

        unit_phase = complex_mel / magnitude.clamp_min(eps)
        compressed = unit_phase * log_radius
        return torch.stack([compressed.real, compressed.imag], dim=1)

    def _append_deltas(self, spec: torch.Tensor) -> torch.Tensor:
        if not self.add_delta and not self.add_delta_delta:
            return spec

        delta = torchaudio.functional.compute_deltas(
            spec,
            win_length=self.delta_win_length,
            mode="replicate",
        )
        channels = [spec]
        if self.add_delta:
            channels.append(delta)
        if self.add_delta_delta:
            channels.append(
                torchaudio.functional.compute_deltas(
                    delta,
                    win_length=self.delta_win_length,
                    mode="replicate",
                )
            )
        return torch.cat(channels, dim=1)

    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: Mel-frequency features [B, C, F, T].
        """
        if x.ndim == 2:
            x = x.unsqueeze(1)
        if x.ndim != 3:
            raise ValueError(f"Expected waveform [B, C, T], got {tuple(x.shape)}")

        if self.representation == "complex_mel":
            spec = self._complex_mel(x)
        else:
            spec = self.mel_spec(x)
            spec = self.amplitude_to_db(spec)

        return self._append_deltas(spec)