AudioMlSpecTools 0.5.0__py3-none-any.whl

AudioMlSpecTools/__init__.py

@@ -0,0 +1,6 @@
+from ._common import MelType, ScalingType, SpecType, AmplitudeToDB, ExportableSTFT # noqa
+from .efficient_features import ChannelConfig, EfficientFeatureSource # noqa
+from .flexible_features import FeatureSource, FeatureChannel # noqa
+from .preproc import AudioPreprocessor, HighPassFilter # noqa
+from .stft_features import FullRangeStftFeatureSource # noqa
+from .wav import load_wav, set_audio_length, list_audio_files, WavReader # noqa

AudioMlSpecTools/_common.py

@@ -0,0 +1,502 @@
+###############################################################################
+# Global Imports
+###############################################################################
+from abc import ABC, abstractmethod
+from enum import Enum
+import json
+import math
+from typing import Optional, Sequence
+import warnings
+
+###############################################################################
+# 3PP Imports
+###############################################################################
+import numpy as np
+import torch
+
+
+###############################################################################
+# Enumerated Types
+###############################################################################
+class SpecType(Enum):
+    '''
+    Possible spectrum types for feature extraction.
+
+    The literature is conflicted on which, if any, provides the best results.
+    '''
+    STFT = "STFT"
+    LOG_STFT = "LOG_STFT"
+    LFCC = "LFCC"
+    MEL = "MEL"
+    LOG_MEL = "LOG_MEL"
+    MFCC = "MFCC"
+
+
+class MelType(Enum):
+    OSHAUGHNESSY = "O'Shaughnessy"
+    FANT = "Fant"
+    LINDSAY_NORMAN = "Lindsay & Norman"
+    SLANEY = "Slaney"
+
+
+class ScalingType(Enum):
+    POWER = "power"
+    MAGNITUDE = "magnitude"
+    LOG = "log"
+
+
+###############################################################################
+# New Classes
+###############################################################################
+class ExportableSTFT(torch.nn.Module):
+    """
+        Exportable to Executorch. Created by Claude with supervision.
+
+        Results are very slightly different than Torch but show no noticeable difference in model accuracy or training.
+    """
+    def __init__(self,
+                 n_fft: int,
+                 hop_length: int,
+                 *,
+                 win_length: Optional[int] = None,
+                 window: Optional[torch.Tensor] = None
+                 ):
+        super().__init__()
+
+        _w = window if window is not None else torch.hann_window(n_fft)
+        _window_len = win_length or n_fft
+        if _window_len < n_fft:
+            pad_left = (n_fft - _window_len) // 2
+            pad_right = n_fft - _window_len - pad_left
+            _w = torch.nn.functional.pad(_w, (pad_left, pad_right))
+
+        # Only compute onesided bins: n_fft//2+1
+        k = torch.arange(n_fft // 2 + 1).unsqueeze(1)  # (freq, 1)
+        n = torch.arange(n_fft).unsqueeze(0)            # (1, n_fft)
+        angles = -2 * torch.pi * k * n / n_fft          # (freq, n_fft)
+
+        self.hop_length = hop_length
+        self.n_fft = n_fft
+        self.pad = n_fft // 2
+
+        # Shape: (n_fft//2+1, n_fft)
+        self.register_buffer("dft_real", torch.cos(angles) * _w)
+        self.register_buffer("dft_imag", torch.sin(angles) * _w)
+
+    def forward(self, x: torch.Tensor):
+        if x.dim() == 1:
+            x = x.unsqueeze(0).unsqueeze(0)  # (1, 1, T)
+        elif x.dim() == 2:
+            x = x.unsqueeze(1)               # (B, 1, T)
+
+        x = torch.nn.functional.pad(x, (self.pad, self.pad), mode="reflect")
+        x = x.squeeze(1)                                          # (B, T_padded)
+
+        frames = x.unfold(-1, self.n_fft, self.hop_length)        # (B, frames, n_fft)
+
+        real = torch.matmul(frames, self.dft_real.T)              # (B, frames, freq)
+        imag = torch.matmul(frames, self.dft_imag.T)
+
+        power = real ** 2 + imag ** 2                             # (B, frames, freq)
+
+        # Match torch.stft output layout: (B, freq, frames)
+        return power.permute(0, 2, 1)
+
+
+class AudioPreprocessor(torch.nn.Module, ABC):
+    @abstractmethod
+    def __call__(self, wav: torch.Tensor | np.ndarray) -> torch.Tensor:
+        ...
+
+
+class BaseFeatureSource(torch.nn.Module, ABC):
+    def __init__(self, preprocessors: Sequence[AudioPreprocessor]):
+        super(BaseFeatureSource, self).__init__()
+        self.preprocessors = preprocessors
+
+
+###############################################################################
+# New Functions
+###############################################################################
+def load_params(params: str | list | dict):
+    if not isinstance(params, str):
+        return params
+    else:
+        with open(params, "r") as infile:
+            return json.loads(infile.read())
+
+
+def write_params(filename: str, params: list | dict):
+    with open(filename, "r") as outfile:
+        return outfile.write(json.dumps(params, indent=2))
+
+
+def power_of_two(n: int):
+    return (n & (n - 1) == 0) and n != 0
+
+
+def generate_filters(n_fft: int, n_filters: int, sample_rate: int, mel_type: Optional[MelType]):
+    n_freqs = n_fft // 2 + 1
+    f_min = 20.0
+    f_max = float(sample_rate // 2)
+
+    if mel_type:
+        return melscale_fbanks(
+            mel_type,
+            f_min,
+            f_max,
+            n_freqs,
+            n_filters,
+            sample_rate,
+        )
+    else:
+        return linear_fbanks(
+            n_freqs,
+            f_min,
+            f_max,
+            n_filters,
+            sample_rate,
+        )
+
+
+def create_scaler(scaling_type: ScalingType):
+    if scaling_type == ScalingType.LOG:
+        return log_scale
+    else:
+        return AmplitudeToDB(stype=scaling_type.value, top_db=80.0)
+
+
+def scale_spec(spec: torch.Tensor) -> torch.Tensor:
+    min_in_val = torch.min(spec)
+    max_in_val = torch.max(spec)
+    in_span = max_in_val - min_in_val
+
+    min_out_val = torch.zeros(1)
+    max_out_val = torch.ones(1)
+    out_span = max_out_val - min_out_val
+
+    scale_factor = out_span / in_span
+    return (spec - min_in_val) * scale_factor
+
+
+def determine_spec_type(calc_mels: bool, calc_logs: bool, calc_cepstrum: bool):
+    if calc_mels:
+        if calc_cepstrum:
+            return SpecType.MFCC
+        elif calc_logs:
+            return SpecType.LOG_MEL
+        else:
+            return SpecType.MEL
+    else:
+        if calc_cepstrum:
+            return SpecType.LFCC
+        elif calc_logs:
+            return SpecType.LOG_STFT
+        else:
+            return SpecType.STFT
+
+
+###############################################################################
+# Imported Legacy Code
+# This code was extracted from v2.8.0 of [torchaudio](https://github.com/pytorch/audio)
+# torchaudio is licensed under the BSD 2-Clause License, reprinted below
+###############################################################################
+#
+# Copyright (c) 2017 Facebook Inc. (Soumith Chintala),
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+###############################################################################
+
+###############################################################################
+# Generate features
+###############################################################################
+
+class AmplitudeToDB(torch.nn.Module):
+    def __init__(self, stype: str = "power", top_db: Optional[float] = None) -> None:
+        super(AmplitudeToDB, self).__init__()
+        self.stype = stype
+        if top_db is not None and top_db < 0:
+            raise ValueError("top_db must be positive value")
+        self.top_db = top_db
+        self.multiplier = 10.0 if stype == "power" else 20.0
+        self.amin = 1e-10
+        self.ref_value = 1.0
+        self.db_multiplier = math.log10(max(self.amin, self.ref_value))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_db = self.multiplier * torch.log10(torch.clamp(x, min=self.amin))
+        x_db -= self.multiplier * self.db_multiplier
+
+        if self.top_db:
+            max_ref = (x_db.max() - self.top_db)
+            x_db = torch.max(x_db, max_ref)
+
+        return x_db
+
+
+def log_scale(waveform: torch.Tensor) -> torch.Tensor:
+    log_offset = 1e-6
+    return torch.log(waveform + log_offset)
+
+
+def hz_to_mel(freq: float, mel_type: MelType) -> float:
+    if mel_type == MelType.OSHAUGHNESSY:
+        return 2595.0 * math.log10(1.0 + (freq / 700.0))
+    elif mel_type == MelType.FANT:
+        return (1000 / math.log10(2)) * math.log10(1.0 + (freq / 1000.0))
+    elif mel_type == MelType.LINDSAY_NORMAN:
+        return 2410.0 * math.log10(1.0 + (freq / 625.0))
+    else:   # MelType.SLANEY
+        min_log_hz = 1000.0
+        f_sp = 200.0 / 3
+
+        if freq < min_log_hz:
+            return freq / f_sp
+        else:
+            min_log_mel = min_log_hz / f_sp
+            logstep = math.log(6.4) / 27.0
+            return min_log_mel + math.log(freq / min_log_hz) / logstep
+
+
+def mel_to_hz(mels: torch.Tensor, mel_type: MelType) -> torch.Tensor:
+    if mel_type == MelType.OSHAUGHNESSY:
+        return 700.0 * (10.0 ** (mels / 2595.0) - 1.0)
+    elif mel_type == MelType.FANT:
+        mul = math.log10(2) / 1000
+        return 1000.0 * (10 ** (mels * mul) - 1.0)
+    elif mel_type == MelType.LINDSAY_NORMAN:
+        return 625.0 * (10.0 ** (mels / 2410.0) - 1.0)
+    else:
+        min_log_hz = 1000.0
+        f_sp = 200.0 / 3
+
+        freqs = f_sp * mels
+        min_log_mel = min_log_hz / f_sp
+
+        logstep = math.log(6.4) / 27.0
+
+        log_t = mels >= min_log_mel
+        freqs[log_t] = min_log_hz * torch.exp(logstep * (mels[log_t] - min_log_mel))
+
+        return freqs
+
+
+def create_triangular_filterbank(
+    all_freqs: torch.Tensor,
+    f_pts: torch.Tensor,
+) -> torch.Tensor:
+    # Adopted from Librosa
+    # calculate the difference between each filter mid point and each stft freq point in hertz
+    f_diff = f_pts[1:] - f_pts[:-1]  # (n_filter + 1)
+    slopes = f_pts.unsqueeze(0) - all_freqs.unsqueeze(1)  # (n_freqs, n_filter + 2)
+    # create overlapping triangles
+    zero = torch.zeros(1)
+    down_slopes = (-1.0 * slopes[:, :-2]) / f_diff[:-1]  # (n_freqs, n_filter)
+    up_slopes = slopes[:, 2:] / f_diff[1:]  # (n_freqs, n_filter)
+    fb = torch.max(zero, torch.min(down_slopes, up_slopes))
+
+    return fb
+
+
+def melscale_fbanks(
+        mel_type: MelType,
+        f_min: float,
+        f_max: float,
+        n_freqs: int,
+        n_mels: int,
+        sample_rate: int,
+):
+    # freq bins
+    all_freqs = torch.linspace(0, sample_rate // 2, n_freqs)
+
+    # calculate mel freq bins
+    m_min = hz_to_mel(f_min, mel_type=mel_type)
+    m_max = hz_to_mel(f_max, mel_type=mel_type)
+
+    m_pts = torch.linspace(m_min, m_max, n_mels + 2)
+    f_pts = mel_to_hz(m_pts, mel_type=mel_type)
+
+    # create filterbank
+    fb = create_triangular_filterbank(all_freqs, f_pts)
+
+    if (fb.max(dim=0).values == 0.0).any():
+        warnings.warn(
+            "At least one mel filterbank has all zero values. "
+            f"The value for `n_mels` ({n_mels}) may be set too high. "
+            f"Or, the value for `n_freqs` ({n_freqs}) may be set too low."
+        )
+
+    return fb
+
+
+def linear_fbanks(
+    n_freqs: int,
+    f_min: float,
+    f_max: float,
+    n_filter: int,
+    sample_rate: int,
+) -> torch.Tensor:
+    # freq bins
+    all_freqs = torch.linspace(0, sample_rate // 2, n_freqs)
+
+    # filter mid-points
+    f_pts = torch.linspace(f_min, f_max, n_filter + 2)
+
+    # create filterbank
+    fb = create_triangular_filterbank(all_freqs, f_pts)
+
+    return fb
+
+
+def create_dct(n_cepstrum: int, n_filters: int) -> torch.Tensor:
+    # http://en.wikipedia.org/wiki/Discrete_cosine_transform#DCT-II
+    n = torch.arange(float(n_filters))
+    k = torch.arange(float(n_cepstrum)).unsqueeze(1)
+    dct = torch.cos(math.pi / float(n_filters) * (n + 0.5) * k)  # size (n_mfcc, n_mels)
+
+    dct[0] *= 1.0 / math.sqrt(2.0)
+    dct *= math.sqrt(2.0 / float(n_filters))
+    return dct.t()
+
+
+###############################################################################
+# Load audio
+###############################################################################
+
+def _get_sinc_resample_kernel(
+    orig_freq: int,
+    new_freq: int,
+    gcd: int,
+    lowpass_filter_width: int = 6,
+    rolloff: float = 0.99,
+    device: torch.device = torch.device("cpu"),
+    dtype: Optional[torch.dtype] = None,
+):
+    orig_freq = int(orig_freq) // gcd
+    new_freq = int(new_freq) // gcd
+
+    if lowpass_filter_width <= 0:
+        raise ValueError("Low pass filter width should be positive.")
+    base_freq = min(orig_freq, new_freq)
+    # This will perform antialiasing filtering by removing the highest frequencies.
+    # At first I thought I only needed this when downsampling, but when upsampling
+    # you will get edge artifacts without this, as the edge is equivalent to zero padding,
+    # which will add high freq artifacts.
+    base_freq *= rolloff
+
+    # The key idea of the algorithm is that x(t) can be exactly reconstructed from x[i] (tensor)
+    # using the sinc interpolation formula:
+    #   x(t) = sum_i x[i] sinc(pi * orig_freq * (i / orig_freq - t))
+    # We can then sample the function x(t) with a different sample rate:
+    #    y[j] = x(j / new_freq)
+    # or,
+    #    y[j] = sum_i x[i] sinc(pi * orig_freq * (i / orig_freq - j / new_freq))
+
+    # We see here that y[j] is the convolution of x[i] with a specific filter, for which
+    # we take an FIR approximation, stopping when we see at least `lowpass_filter_width` zeros crossing.
+    # But y[j+1] is going to have a different set of weights and so on, until y[j + new_freq].
+    # Indeed:
+    # y[j + new_freq] = sum_i x[i] sinc(pi * orig_freq * ((i / orig_freq - (j + new_freq) / new_freq))
+    #                 = sum_i x[i] sinc(pi * orig_freq * ((i - orig_freq) / orig_freq - j / new_freq))
+    #                 = sum_i x[i + orig_freq] sinc(pi * orig_freq * (i / orig_freq - j / new_freq))
+    # so y[j+new_freq] uses the same filter as y[j], but on a shifted version of x by `orig_freq`.
+    # This will explain the F.conv1d after, with a stride of orig_freq.
+    width = math.ceil(lowpass_filter_width * orig_freq / base_freq)
+    # If orig_freq is still big after GCD reduction, most filters will be very unbalanced, i.e.,
+    # they will have a lot of almost zero values to the left or to the right...
+    # There is probably a way to evaluate those filters more efficiently, but this is kept for
+    # future work.
+    idx_dtype = dtype if dtype is not None else torch.float64
+
+    idx = torch.arange(-width, width + orig_freq, dtype=idx_dtype, device=device)[None, None] / orig_freq
+
+    t = torch.arange(0, -new_freq, -1, dtype=dtype, device=device)[:, None, None] / new_freq + idx
+    t *= base_freq
+    t = t.clamp_(-lowpass_filter_width, lowpass_filter_width)
+
+    # we do not use built in torch windows here as we need to evaluate the window
+    # at specific positions, not over a regular grid.
+    window = torch.cos(t * math.pi / lowpass_filter_width / 2) ** 2
+
+    t *= math.pi
+
+    scale = base_freq / orig_freq
+    kernels = torch.where(t == 0, torch.tensor(1.0).to(t), t.sin() / t)
+    kernels *= window * scale
+
+    if dtype is None:
+        kernels = kernels.to(dtype=torch.float32)
+
+    return kernels, width
+
+
+def _apply_sinc_resample_kernel(
+    waveform: torch.Tensor,
+    orig_freq: int,
+    new_freq: int,
+    gcd: int,
+    kernel: torch.Tensor,
+    width: int,
+):
+    orig_freq = int(orig_freq) // gcd
+    new_freq = int(new_freq) // gcd
+
+    # pack batch
+    shape = waveform.size()
+    waveform = waveform.view(-1, shape[-1])
+
+    num_wavs, length = waveform.shape
+    waveform = torch.nn.functional.pad(waveform, (width, width + orig_freq))
+    resampled = torch.nn.functional.conv1d(waveform[:, None], kernel, stride=orig_freq)
+    resampled = resampled.transpose(1, 2).reshape(num_wavs, -1)
+    target_length = torch.ceil(torch.as_tensor(new_freq * length / orig_freq)).long()
+    resampled = resampled[..., :target_length]
+
+    # unpack batch
+    resampled = resampled.view(shape[:-1] + resampled.shape[-1:])
+    return resampled
+
+
+def resample(
+    waveform: torch.Tensor,
+    orig_freq: int,
+    new_freq: int,
+    lowpass_filter_width: int = 6,
+    rolloff: float = 0.99,
+) -> torch.Tensor:
+    if orig_freq == new_freq:
+        return waveform
+
+    gcd = math.gcd(int(orig_freq), int(new_freq))
+
+    kernel, width = _get_sinc_resample_kernel(
+        orig_freq,
+        new_freq,
+        gcd,
+        lowpass_filter_width,
+        rolloff,
+        waveform.device,
+        waveform.dtype,
+    )
+    resampled = _apply_sinc_resample_kernel(waveform, orig_freq, new_freq, gcd, kernel, width)
+    return resampled

AudioMlSpecTools/efficient_features.py

@@ -0,0 +1,176 @@
+###############################################################################
+# Global Imports
+###############################################################################
+from typing import Optional, Sequence
+import warnings
+
+###############################################################################
+# 3PP Imports
+###############################################################################
+import torch
+
+###############################################################################
+# Local Imports
+###############################################################################
+from ._common import MelType, ScalingType, ExportableSTFT, AudioPreprocessor, BaseFeatureSource
+from ._common import power_of_two, create_dct, create_scaler, generate_filters, scale_spec
+
+
+###############################################################################
+# Helper Classes
+###############################################################################
+class ChannelConfig:
+    def __init__(self,
+                 *,
+                 is_mel: bool,
+                 is_log: bool,
+                 is_cepstrum: bool,
+                 ):
+        self.is_mel = is_mel
+        self.is_log = is_log
+        self.is_cepstrum = is_cepstrum
+
+        self.key = self._make_channel_key()
+
+    def get_key(self):
+        return self.key
+
+    def _make_channel_key(self):
+        if self.is_cepstrum:
+            return "mfcc" if self.is_mel else "lfcc"
+        elif self.is_log:
+            return "mel_log_spec" if self.is_mel else "lin_log_spec"
+        else:
+            return "mel_freq_spec" if self.is_mel else "lin_freq_spec"
+
+
+###############################################################################
+# Classes
+###############################################################################
+class EfficientFeatureSource(BaseFeatureSource):
+    def __init__(self,
+                 sample_rate: int,
+                 channels: list[ChannelConfig],
+                 preprocessors: Sequence[AudioPreprocessor] = [],
+                 *,
+                 # For all spectra
+                 n_fft: Optional[int] = None,
+                 hop_length: Optional[int] = None,
+
+                 # Shared
+                 n_filters: Optional[int] = None,
+
+                 # For all cepstra
+                 cepstral_coefficients: Optional[int] = None,
+                 ):
+        super(EfficientFeatureSource, self).__init__(preprocessors)
+
+        self.channels = channels
+
+        self.has_lin_freq = any([not c.is_mel for c in channels])
+        self.has_mel_freq = any([c.is_mel for c in channels])
+
+        self.has_lfcc = self.has_lin_freq and any([c.is_cepstrum for c in channels])
+        self.has_mfcc = self.has_mel_freq and any([c.is_cepstrum for c in channels])
+        self.has_cepstrum = self.has_lfcc or self.has_mfcc
+
+        self.has_lin_log = any([not c.is_mel and c.is_log for c in channels])
+        self.has_mel_log = any([c.is_mel and c.is_log for c in channels])
+        self.has_log_scale = any([c.is_log for c in channels]) or self.has_cepstrum
+
+        # Required configs
+        self.sample_rate = sample_rate
+
+        # Universal configs
+        if n_fft is not None and not power_of_two(n_fft):
+            raise ValueError("n_fft must be a power of 2")
+        self.n_fft = n_fft or 1024
+
+        if hop_length is not None and hop_length > (self.n_fft // 2):
+            warnings.warn(f"hop_length should be set to no more than 1/2 the FFT window size, or {self.n_fft // 2} mels for n_fft = {self.n_fft} (currently {hop_length})")
+        self.hop_length = hop_length or self.n_fft // 4
+
+        # Shared configs (mels, MFCC, LFCC)
+        if n_filters is not None and n_filters > (self.n_fft // 8):
+            warnings.warn(f"n_filters should be set to no more than 1/8 the FFT window size, or {self.n_fft // 8} filters for n_fft = {self.n_fft} (currently {n_filters})")
+        self.n_filters = n_filters or self.n_fft // 8
+
+        # Cepstral configs
+        if cepstral_coefficients is not None and cepstral_coefficients > self.n_filters:
+            raise ValueError(f"cepstral_coefficients must be no greater than n_mels (currently {cepstral_coefficients}/{self.n_filters})")
+        self.cepstral_coefficients = cepstral_coefficients or self.n_filters
+
+        ###################
+        # Spec gen code
+        ###################
+
+        # Basic spectrogram
+        self._stft = ExportableSTFT(self.n_fft, self.hop_length)
+
+        if self.has_lin_freq:
+            self.register_buffer("lin_filt", generate_filters(self.n_fft, self.n_filters, self.sample_rate, None))
+        else:
+            self.lin_filt = None
+
+        if self.has_mel_freq:
+            self.register_buffer("mel_filt", generate_filters(self.n_fft, self.n_filters, self.sample_rate, MelType.OSHAUGHNESSY))
+        else:
+            self.mel_filt = None
+
+        # DB scaling, if necessary
+        self.amplitude_to_DB = create_scaler(ScalingType.POWER) if self.has_log_scale else None
+
+        # Cepstrum, if necessary
+        if self.has_cepstrum:
+            self.register_buffer("dct", create_dct(self.cepstral_coefficients, self.n_filters))
+        else:
+            self.dct = None
+
+    def forward(self, wav: torch.Tensor) -> torch.Tensor:
+        for preproc in self.preprocessors:
+            wav = preproc(wav)
+
+        spec = self._stft(wav)
+
+        if self.lin_filt is not None:
+            lin_freq_spec = torch.matmul(spec.transpose(-1, -2), self.lin_filt).transpose(-1, -2)
+        else:
+            lin_freq_spec = None
+
+        if self.mel_filt is not None:
+            mel_freq_spec = torch.matmul(spec.transpose(-1, -2), self.mel_filt).transpose(-1, -2)
+        else:
+            mel_freq_spec = None
+
+        if self.has_lin_log and (lin_freq_spec is not None) and self.amplitude_to_DB:
+            lin_log_spec = self.amplitude_to_DB(lin_freq_spec)
+        else:
+            lin_log_spec = None
+
+        if self.has_mel_log and (mel_freq_spec is not None) and self.amplitude_to_DB:
+            mel_log_spec = self.amplitude_to_DB(mel_freq_spec)
+        else:
+            mel_log_spec = None
+
+        if self.has_lfcc and (lin_log_spec is not None) and (self.dct is not None):
+            lfcc = torch.matmul(lin_log_spec.transpose(-1, -2), self.dct).transpose(-1, -2)
+        else:
+            lfcc = None
+
+        if self.has_mfcc and (mel_log_spec is not None) and (self.dct is not None):
+            mfcc = torch.matmul(mel_log_spec.transpose(-1, -2), self.dct).transpose(-1, -2)
+        else:
+            mfcc = None
+
+        return self._collate(**{
+            "lin_freq_spec": lin_freq_spec,
+            "mel_freq_spec": lin_freq_spec,
+            "lin_log_spec": lin_log_spec,
+            "mel_log_spec": mel_log_spec,
+            "lfcc": lfcc,
+            "mfcc": mfcc,
+        })
+
+    def _collate(self, **kwargs):
+        specs = [scale_spec(kwargs[c.key]) for c in self.channels]
+        return torch.stack(specs, dim=1)

AudioMlSpecTools/flexible_features.py

@@ -0,0 +1,218 @@
+###############################################################################
+# Global Imports
+###############################################################################
+from typing import Optional, Sequence
+import warnings
+
+###############################################################################
+# 3PP Imports
+###############################################################################
+import torch
+
+###############################################################################
+# Local Imports
+###############################################################################
+from ._common import MelType, ScalingType, ExportableSTFT, AudioPreprocessor, BaseFeatureSource
+from ._common import power_of_two, create_dct, create_scaler, generate_filters, determine_spec_type, scale_spec, load_params, write_params
+
+
+###############################################################################
+# Export Classes
+###############################################################################
+class FeatureChannel(torch.nn.Module):
+    def __init__(self,
+                 sample_rate: int,
+                 *,
+                 # For all spectra
+                 n_fft: Optional[int] = None,
+                 hop_length: Optional[int] = None,
+                 scale_spec: Optional[bool] = None,
+
+                 # Shared
+                 n_filters: Optional[int] = None,
+
+                 # For all mel spectra
+                 is_mel: Optional[bool] = None,
+                 mel_type: Optional[MelType] = None,
+
+                 # For all log spectra
+                 is_logarithmic: Optional[bool] = None,
+                 scaling_type: ScalingType = ScalingType.POWER,
+
+                 # For all cepstra
+                 is_cepstrum: Optional[bool] = None,
+                 cepstral_coefficients: Optional[int] = None,
+                 ):
+        super(FeatureChannel, self).__init__()
+
+        # Required configs
+        self.sample_rate = sample_rate
+
+        # Universal configs
+        if n_fft is not None and not power_of_two(n_fft):
+            raise ValueError("n_fft must be a power of 2")
+        self.n_fft = n_fft or 1024
+
+        if hop_length is not None and hop_length > (self.n_fft // 2):
+            warnings.warn(f"hop_length should be set to no more than 1/2 the FFT window size, or {self.n_fft // 2} mels for n_fft = {self.n_fft} (currently {hop_length})")
+        self.hop_length = hop_length or self.n_fft // 4
+
+        self.scale_spec = scale_spec if scale_spec is not None else True
+
+        # Shared configs (mels, MFCC, LFCC)
+        if n_filters is not None and n_filters > (self.n_fft // 8):
+            warnings.warn(f"n_filters should be set to no more than 1/8 the FFT window size, or {self.n_fft // 8} filters for n_fft = {self.n_fft} (currently {n_filters})")
+        self.n_filters = n_filters or self.n_fft // 8
+
+        # Mel configs
+        calc_mels = is_mel or (mel_type is not None)
+        self.mel_type = mel_type or (MelType.OSHAUGHNESSY if calc_mels else None)
+
+        # Cepstral configs
+        calc_cepstrum = is_cepstrum or (cepstral_coefficients is not None)
+
+        if cepstral_coefficients is not None and cepstral_coefficients > self.n_filters:
+            raise ValueError(f"cepstral_coefficients must be no greater than n_mels (currently {cepstral_coefficients}/{self.n_filters})")
+        self.cepstral_coefficients = cepstral_coefficients or self.n_filters
+
+        # Log configs
+        calc_logs = is_logarithmic or calc_cepstrum
+
+        ###################
+        # Spec gen code
+        ###################
+
+        # Basic spectrogram
+        self._stft = ExportableSTFT(self.n_fft, self.hop_length)
+
+        fb = generate_filters(self.n_fft, self.n_filters, self.sample_rate, self.mel_type)
+        self.register_buffer("fb", fb)
+
+        # DB scaling, if necessary
+        self.scaling_type = scaling_type
+        self.amplitude_to_DB = create_scaler(scaling_type) if calc_logs else None
+
+        # Cepstrum, if necessary
+        if calc_cepstrum:
+            dct_mat = create_dct(self.cepstral_coefficients, self.n_filters)
+            self.register_buffer("dct_mat", dct_mat)
+        else:
+            self.dct_mat = None
+
+        self.spec_type = determine_spec_type(calc_mels, calc_logs, calc_cepstrum)
+
+    def get_spec_type(self):
+        return self.spec_type
+
+    def forward(self, wav: torch.Tensor) -> torch.Tensor:
+        spec = self._stft(wav)
+
+        if self.fb is not None:
+            spec = torch.matmul(spec.transpose(-1, -2), self.fb).transpose(-1, -2)
+
+        if self.amplitude_to_DB is not None:
+            spec = self.amplitude_to_DB(spec)
+
+        if self.dct_mat is not None:
+            spec = torch.matmul(spec.transpose(-1, -2), self.dct_mat).transpose(-1, -2)
+
+        if self.scale_spec:
+            spec = scale_spec(spec)
+
+        return spec
+
+    @staticmethod
+    def from_json(params: str | dict):
+        loaded_params = load_params(params)
+        if not isinstance(loaded_params, dict):
+            raise ValueError(f"Invalid {FeatureChannel.__name__} parameters")
+
+        mel_type_raw = loaded_params.get("mel_type", None)
+        scaling_type_raw = loaded_params.get("scaling_type", None)
+
+        return FeatureChannel(sample_rate=loaded_params["sample_rate"],
+                              # For all spectra
+                              n_fft=loaded_params.get("n_fft", None),
+                              hop_length=loaded_params.get("hop_length", None),
+                              scale_spec=loaded_params.get("scale_spec", None),
+
+                              # Shared
+                              n_filters=loaded_params.get("n_filters", None),
+
+                              # For all mel spectra
+                              is_mel=loaded_params.get("is_mel", None),
+                              mel_type=MelType(mel_type_raw) if mel_type_raw else None,
+
+                              # For all log spectra
+                              is_logarithmic=loaded_params.get("is_logarithmic", None),
+                              scaling_type=ScalingType(scaling_type_raw) if scaling_type_raw else ScalingType.POWER,
+
+                              # For all cepstra
+                              is_cepstrum=loaded_params.get("is_cepstrum", None),
+                              cepstral_coefficients=loaded_params.get("cepstral_coefficients", None),
+                              )
+
+    def get_params(self):
+        return {
+            "sample_rate": self.sample_rate,
+            "n_fft": self.n_fft,
+            "hop_length": self.hop_length,
+            "scale_spec": self.scale_spec,
+
+            # Shared
+            "n_filters": self.n_filters,
+
+            # For all mel spectra
+            "is_mel": self.mel_type is not None,
+            "mel_type": self.mel_type,
+
+            # For all log spectra
+            "is_logarithmic": self.amplitude_to_DB is not None,
+            "scaling_type": self.scaling_type,
+
+            # For all cepstra
+            "is_cepstrum": self.dct_mat is not None,
+            "cepstral_coefficients": self.cepstral_coefficients if self.dct_mat is not None else None,
+        }
+
+    def to_json(self, filename: str):
+        params = self.get_params()
+        write_params(filename, params)
+
+
+class FeatureSource(BaseFeatureSource):
+    def __init__(self,
+                 feature_channels: list[FeatureChannel],
+                 preprocessors: Sequence[AudioPreprocessor] = [],
+                 ):
+        super(FeatureSource, self).__init__(preprocessors)
+
+        if len(feature_channels) == 0:
+            raise ValueError("Must include at least one spec type")
+
+        self.fc = feature_channels
+        self.feature_channels = torch.nn.ModuleList(feature_channels)
+        self.preprocessors = torch.nn.ModuleList(preprocessors)
+
+    def num_channels(self):
+        return len(self.fc)
+
+    def forward(self, wav: torch.Tensor) -> torch.Tensor:
+        for preproc in self.preprocessors:
+            wav = preproc(wav)
+
+        spectra = [chan(wav) for chan in self.feature_channels]
+        return torch.stack(spectra, dim=1)
+
+    @staticmethod
+    def from_json(params: str | list):
+        loaded_params = load_params(params)
+        if not isinstance(loaded_params, list):
+            raise ValueError(f"Invalid {FeatureSource.__name__} parameters")
+
+        feature_channels = [FeatureChannel.from_json(p) for p in loaded_params]
+        return FeatureSource(feature_channels)
+
+    def to_json(self, filename: str):
+        params = [fc.get_params() for fc in self.feature_channels.children() if isinstance(fc, FeatureChannel)]
+        write_params(filename, params)

AudioMlSpecTools/preproc.py

@@ -0,0 +1,50 @@
+###############################################################################
+# 3PP Imports
+###############################################################################
+import numpy as np
+from scipy.signal import butter, sosfilt, sosfilt_zi
+import torch
+
+###############################################################################
+# Local Imports
+###############################################################################
+from ._common import AudioPreprocessor
+
+
+###############################################################################
+# Classes
+###############################################################################
+class HighPassFilter(AudioPreprocessor):
+    def __init__(self,
+                 *,
+                 sample_rate: int,
+                 cutoff_freq: int,
+                 rolloff_db: int,
+                 ):
+        nyquist = sample_rate // 2
+        if cutoff_freq < 0 or cutoff_freq > nyquist:
+            raise ValueError(f"Cutoff freq must be >0 Hz and less than the Nyquist rate ({nyquist})")
+        elif rolloff_db < 6 or rolloff_db % 6 != 0:
+            raise ValueError("Rolloff dB must be >0 and a multiple of 6")
+
+        self.cutoff_freq = cutoff_freq
+        self.rolloff_db = rolloff_db
+
+        self.sos = butter(rolloff_db // 6,
+                          cutoff_freq,
+                          btype="highpass",
+                          analog=False,
+                          fs=sample_rate,
+                          output="sos",
+                          )
+
+        self.zi = sosfilt_zi(self.sos)
+
+    def __call__(self, wav: torch.Tensor | np.ndarray) -> torch.Tensor:
+        _wav = wav.numpy() if isinstance(wav, torch.Tensor) else wav
+        if _wav.ndim != 1:
+            raise ValueError(f"Improper input dim; must be 1 but is {_wav.ndim}")
+
+        processed_samples, _ = sosfilt(self.sos, _wav, zi=self.zi * _wav[0])
+        processed_samples = processed_samples.astype(np.float32)
+        return torch.from_numpy(processed_samples)

AudioMlSpecTools/stft_features.py

@@ -0,0 +1,62 @@
+###############################################################################
+# Global Imports
+###############################################################################
+from typing import Optional, Sequence
+import warnings
+
+###############################################################################
+# 3PP Imports
+###############################################################################
+import torch
+
+###############################################################################
+# Local Imports
+###############################################################################
+from ._common import ScalingType, ExportableSTFT, AudioPreprocessor, BaseFeatureSource
+from ._common import power_of_two, create_scaler, scale_spec
+
+
+###############################################################################
+# Export Classes
+###############################################################################
+class FullRangeStftFeatureSource(BaseFeatureSource):
+    def __init__(self,
+                 sample_rate: int,
+                 preprocessors: Sequence[AudioPreprocessor] = [],
+                 *,
+                 # For all spectra
+                 n_fft: Optional[int] = None,
+                 hop_length: Optional[int] = None,
+
+                 # For all log spectra
+                 is_logarithmic: bool = True,
+                 ):
+        super(FullRangeStftFeatureSource, self).__init__(preprocessors)
+
+        # Internal configs
+        self.sample_rate = sample_rate
+
+        if n_fft is not None and not power_of_two(n_fft):
+            raise ValueError("n_fft must be a power of 2")
+        self.n_fft = n_fft or 1024
+
+        if hop_length is not None and hop_length > (self.n_fft // 2):
+            warnings.warn(f"hop_length should be set to no more than 1/2 the FFT window size, or {self.n_fft // 2} mels for n_fft = {self.n_fft} (currently {hop_length})")
+        self.hop_length = hop_length or self.n_fft // 4
+
+        # Basic spectrogram
+        self._stft = ExportableSTFT(self.n_fft, self.hop_length)
+
+        # DB scaling, if necessary
+        self.amplitude_to_DB = create_scaler(ScalingType.POWER) if is_logarithmic else None
+
+    def forward(self, wav: torch.Tensor) -> torch.Tensor:
+        for preproc in self.preprocessors:
+            wav = preproc(wav)
+
+        spec = self._stft(wav)
+
+        if self.amplitude_to_DB is not None:
+            spec = self.amplitude_to_DB(spec)
+
+        return scale_spec(spec).unsqueeze(0)

AudioMlSpecTools/wav.py

@@ -0,0 +1,108 @@
+###############################################################################
+# Global Imports
+###############################################################################
+import os
+from typing import Optional
+
+###############################################################################
+# 3PP Imports
+###############################################################################
+import torch
+import torchcodec
+
+###############################################################################
+# Local Imports
+###############################################################################
+from ._common import resample
+
+
+###############################################################################
+# Helpers
+###############################################################################
+def _is_multichannel(wave: torch.Tensor) -> bool:
+    return wave.size(0) > 1
+
+
+###############################################################################
+# Functions
+###############################################################################
+def set_audio_length(wave: torch.Tensor, sr: int, duration_secs: Optional[int]):
+    if duration_secs is None:
+        return wave
+    else:
+        # Truncate as needed
+        start_secs = 0
+        end_secs = duration_secs
+        st_idx, end_idx = int(start_secs * sr), int(end_secs * sr)
+        wave = wave[:, st_idx:end_idx]
+
+        # Zero Padding
+        num_samples = int(sr * duration_secs)
+        padding_size = max(num_samples - wave.size(1), 0)
+        if padding_size > 0:
+            wave = torch.cat([wave, torch.zeros(1, padding_size)], dim=1)
+
+        # Trim
+        wave = wave[:, :num_samples]
+
+        return wave
+
+
+def load_wav(path: str,
+             *,
+             target_sr: Optional[int] = None,
+             duration_secs: Optional[int] = None,
+             ) -> torch.Tensor:
+    '''
+    Load a WAV file into memory for processing.
+
+    Other file types may work but have not been tested.
+
+    Positional arguments:
+        path -- Location of the audio file
+
+    Keyword arguments:
+        target_sr -- Sample rate to which audio should be resampled
+        duration_secs -- Consistent duration of output audio, either by truncation or zero-padding as needed
+    '''
+
+    audio = torchcodec.decoders.AudioDecoder(path).get_all_samples()
+    wave = audio.data
+
+    # Resample
+    final_sr = target_sr or audio.sample_rate
+
+    if target_sr is not None:
+        wave = resample(wave, orig_freq=audio.sample_rate, new_freq=target_sr)
+
+    # Stereo -> Mono
+    if _is_multichannel(wave):
+        wave = torch.mean(wave, dim=0, keepdim=True)
+
+    wave = set_audio_length(wave, final_sr, duration_secs)
+
+    return wave
+
+
+def list_audio_files(dir: str) -> list[str]:
+    return sorted([f for f in os.listdir(dir) if f.endswith(".wav")])
+
+
+###############################################################################
+# Classes
+###############################################################################
+class WavReader:
+    def __init__(self, target_sr: int, duration_secs: Optional[int] = None):
+        self.target_sr = target_sr
+        self.duration_secs = duration_secs
+
+    def load(self, path: str) -> torch.Tensor:
+        return load_wav(path, target_sr=self.target_sr, duration_secs=self.duration_secs)
+
+    def __call__(self, path: str) -> torch.Tensor:
+        return self.load(path)
+
+    def clip(self, wav: torch.Tensor, start_sec: int, end_sec: int) -> torch.Tensor:
+        start_frame = start_sec * self.target_sr
+        end_frame = end_sec * self.target_sr
+        return wav[:, start_frame:end_frame]

audiomlspectools-0.5.0.dist-info/METADATA

@@ -0,0 +1,74 @@
+Metadata-Version: 2.4
+Name: AudioMlSpecTools
+Version: 0.5.0
+Summary: Convenience functions for generating ML features from audio data
+Author-email: Ryan Quinn <ryan.quinn@certusinnovations.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/Stonewall-Defense/team-ml-audio-features
+Classifier: Development Status :: 4 - Beta
+Classifier: Environment :: Console
+Classifier: Intended Audience :: Developers
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch==2.10.0
+Requires-Dist: torchcodec==0.10.0
+Dynamic: license-file
+
+# Audio ML Spec Tools
+
+Convenience functions for generating ML features from audio data. Breaks audio ML dependencies on `torchaudio`. Unlike `pytorch` features, these functions can be exported to ExecuTorch and ONNX with no issues.
+
+## Motivation
+
+Except in specific circumstances like `wav2vec`, raw audio has proven to be a much worse input for ML models than spectrogram-based features across a wide variety of problem domains, including environmental sound classificarion ([Guzhov et al. (2021)](https://arxiv.org/pdf/2104.11587)), singing technique classification ([Yamamoto et al. (2021)](https://www.slis.tsukuba.ac.jp/lspc/0000890.pdf)), and ship classification ([Xie, Ren, and Xu (2024)](https://arxiv.org/pdf/2306.01002)).
+
+There is no scientific consensus on the relative benefits of mel-scale spectrograms, linear spectrograms, and MFCCs. Different researchers have shown good results with each type of spectrogram; see respectively [Raponi, Oligeri, and Ali (2021)](https://arxiv.org/pdf/2004.07948), [Jung at al. (2021)](https://www.mdpi.com/2075-4418/11/4/732), and [Razani et al (2017)](https://www.ece.mcgill.ca/~bchamp/Papers/Conference/ISSPIT2017.pdf).
+
+With this library, you can easily try as many feature extraction methods as you want to see what works for your use case.
+
+## Prerequisites
+
+- Python 3.12 runtime
+- `pip` for package installation
+- Note that `torchcodec` depends on a system installation of FFmpeg
+
+## Installation
+
+Install the dependencies into the environment with [pip](https://pypi.org/project/pip/):
+
+```bash
+pip install -r requirements.txt
+```
+
+Then install the package itself locally:
+
+```bash
+pip install .
+```
+
+## Usage
+
+See `examples/features.py`.
+
+## Testing
+
+```bash
+python3 -m coverage run -m unittest discover -s test -p "*_test.py" && python -m coverage report --skip-covered
+python -m coverage html
+```
+
+## Versioning
+
+We use [SemVer](http://semver.org/) for versioning. For the versions available, see the [tags on this repository](https://github.com/Stonewall-Defense/team-ml-audio-features/tags).
+
+## Authors
+
+- **Ryan Quinn** - *Initial work*
+
+## License
+
+MIT.

audiomlspectools-0.5.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+AudioMlSpecTools/__init__.py,sha256=kQxQCIXzMWWCAu8tXNTggOyrFGdvkFBrgNprjQobnG8,438
+AudioMlSpecTools/_common.py,sha256=GzHnubmVsMN9IW-OM-GkzWobxqrY_z2EEu7N8pBZ74s,17367
+AudioMlSpecTools/efficient_features.py,sha256=F_wfXWWMQtdWudKUPc9VBnjslhz6WelQpa_cH0SMH-w,6966
+AudioMlSpecTools/flexible_features.py,sha256=BSOucAPF687lecOzgDoO1fsk5sys9HnOa9IyyxaJg3Q,8662
+AudioMlSpecTools/preproc.py,sha256=HyIzkz2_NoImvY0_brjtLGkuGofmnB-lNKzOJMbw5Cw,2010
+AudioMlSpecTools/stft_features.py,sha256=UZXSZHjDFDUjmPzYBBq4DSmEz1x_e-YWXAITS30xMwo,2474
+AudioMlSpecTools/wav.py,sha256=3kxu336FFYYnSH949VECPOKKguXhFSaia9kkh4__m2o,3605
+audiomlspectools-0.5.0.dist-info/licenses/LICENSE,sha256=4MM0VuiftnP2VCMsqBWR9luSTrr9kX06xVn6VPisv5I,1069
+audiomlspectools-0.5.0.dist-info/METADATA,sha256=ikkDfP1EzDK_198YV2Y1SRUZhyBwrCMiEhF9CZ-N4MQ,2851
+audiomlspectools-0.5.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+audiomlspectools-0.5.0.dist-info/top_level.txt,sha256=DIiCB2FGOgLnCb0oekQqlbrV9mX6k6EoAdZCnK8cncw,17
+audiomlspectools-0.5.0.dist-info/RECORD,,

audiomlspectools-0.5.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

audiomlspectools-0.5.0.dist-info/licenses/LICENSE

@@ -0,0 +1,7 @@
+Copyright © 2025 Certus Innovations
+
+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.

audiomlspectools-0.5.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+AudioMlSpecTools