audio.py 2 KB
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import librosa
from typing import List, Optional, Tuple


class AudioProcessor:
    def __init__(self, sr: int = 16000, n_mels: int = 80, n_fft: int = 512, hop_length: int = 160):
        self.sr = sr
        self.n_mels = n_mels
        self.n_fft = n_fft
        self.hop_length = hop_length

    def load(self, path: str, sr: Optional[int] = None, duration: Optional[float] = None) -> np.ndarray:
        y, _ = librosa.load(path, sr=sr or self.sr, mono=True, duration=duration)
        return y

    def to_mel(self, y: np.ndarray) -> np.ndarray:
        mel = librosa.feature.melspectrogram(
            y=y, sr=self.sr, n_mels=self.n_mels,
            n_fft=self.n_fft, hop_length=self.hop_length
        )
        return librosa.power_to_db(mel, ref=np.max)

    def to_mel_tensor(self, y: np.ndarray) -> torch.Tensor:
        mel = self.to_mel(y)
        return torch.FloatTensor(mel).unsqueeze(0)

    def sliding_windows(self, y: np.ndarray, window_sec: float = 5.0, stride_sec: float = 2.5) -> List[np.ndarray]:
        win_len = int(window_sec * self.sr)
        stride = int(stride_sec * self.sr)
        if len(y) < win_len:
            pad = win_len - len(y)
            y = np.pad(y, (0, pad))
        windows = []
        for start in range(0, len(y) - win_len + 1, stride):
            windows.append(y[start:start + win_len])
        if not windows:
            windows.append(y[:win_len])
        return windows

    def mel_from_path(self, path: str) -> Tuple[torch.Tensor, float]:
        y = self.load(path)
        duration = len(y) / self.sr
        return self.to_mel_tensor(y), duration

    def extract_chroma(self, y: np.ndarray) -> np.ndarray:
        chroma = librosa.feature.chroma_cqt(y=y, sr=self.sr)
        return chroma

    def extract_f0(self, y: np.ndarray, fmin=65, fmax=2093) -> np.ndarray:
        f0, _, _ = librosa.pyin(y, sr=self.sr, fmin=fmin, fmax=fmax)
        f0 = np.nan_to_num(f0, nan=0.0)
        return f0