Spaces:
Sleeping
Sleeping
| """Hybrid Feature Extraction Module for ASVspoof 5 Telephony Audio. | |
| Implements: | |
| 1. Branch A (Physics-Based Vocal Tract & Glottal Features, 0-4 kHz): | |
| - Linear Prediction Analysis (LPC) via Levinson-Durbin recursion. | |
| - Linear Prediction Cepstral Coefficients (LPCC) capturing the all-pole vocal tract transfer function. | |
| - Glottal inverse filtering residual energy and excitation constraints. | |
| 2. Branch B (Spectrogram / Filterbank Features, 0-4 kHz): | |
| - Log-Mel Spectrogram and Linear Frequency Cepstral Coefficients (LFCC). | |
| """ | |
| from typing import Dict, Optional, Tuple | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| import torchaudio.transforms as T | |
| def compute_autocorr(frames: torch.Tensor, order: int) -> torch.Tensor: | |
| """Compute autocorrelation for framed signals using FFT. | |
| Args: | |
| frames: Tensor of shape (..., num_frames, frame_len) | |
| order: Maximum lag to compute (P) | |
| Returns: | |
| Autocorrelation tensor of shape (..., num_frames, order + 1) | |
| """ | |
| frame_len = frames.shape[-1] | |
| window = torch.hamming_window(frame_len, device=frames.device, dtype=frames.dtype) | |
| w_frames = frames * window | |
| n_fft = 2 ** ((2 * frame_len - 1).bit_length()) | |
| freq = torch.fft.rfft(w_frames, n=n_fft, dim=-1) | |
| power = freq.abs() ** 2 | |
| r = torch.fft.irfft(power, n=n_fft, dim=-1)[..., :order + 1] | |
| return r | |
| def levinson_durbin(r: torch.Tensor, order: int, eps: float = 1e-8) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """Vectorized Levinson-Durbin algorithm. | |
| Solves R * a = -r for all-pole vocal tract filter coefficients. | |
| Args: | |
| r: Autocorrelation tensor of shape (..., order + 1) | |
| order: LPC order P | |
| Returns: | |
| a: LPC filter coefficients of shape (..., order) | |
| e: Prediction error / residual power of shape (...) | |
| """ | |
| batch_shape = r.shape[:-1] | |
| r_flat = r.reshape(-1, order + 1) | |
| B = r_flat.shape[0] | |
| a = torch.zeros(B, order + 1, device=r.device, dtype=r.dtype) | |
| a[:, 0] = 1.0 | |
| e = torch.clamp(r_flat[:, 0].clone(), min=eps) | |
| for i in range(1, order + 1): | |
| acc = torch.zeros(B, device=r.device, dtype=r.dtype) | |
| for j in range(1, i): | |
| acc += a[:, j] * r_flat[:, i - j] | |
| k = -(r_flat[:, i] + acc) / e | |
| # Reflection coefficient stability clamp | |
| k = torch.clamp(k, -0.999, 0.999) | |
| new_a = a.clone() | |
| for j in range(1, i): | |
| new_a[:, j] = a[:, j] + k * a[:, i - j] | |
| new_a[:, i] = k | |
| a = new_a | |
| e = e * (1.0 - k ** 2) | |
| e = torch.clamp(e, min=eps) | |
| lpc_coeffs = a[:, 1:].reshape(*batch_shape, order) | |
| residual_power = e.reshape(*batch_shape) | |
| return lpc_coeffs, residual_power | |
| def lpc_to_lpcc(lpc: torch.Tensor, num_cepstrum: int) -> torch.Tensor: | |
| """Recursively convert LPC coefficients to Linear Prediction Cepstral Coefficients (LPCC). | |
| Captures the cepstral envelope of the vocal tract transfer function. | |
| """ | |
| *batch_shape, P = lpc.shape | |
| lpc_flat = lpc.reshape(-1, P) | |
| B = lpc_flat.shape[0] | |
| cep = torch.zeros(B, num_cepstrum, device=lpc.device, dtype=lpc.dtype) | |
| cep[:, 0] = -lpc_flat[:, 0] | |
| for n in range(1, num_cepstrum): | |
| idx = n + 1 | |
| sum_val = torch.zeros(B, device=lpc.device, dtype=lpc.dtype) | |
| limit = min(n, P) | |
| for k in range(limit): | |
| k_1idx = k + 1 | |
| sum_val += (1.0 - k_1idx / idx) * lpc_flat[:, k] * cep[:, n - k_1idx] | |
| if n < P: | |
| cep[:, n] = -lpc_flat[:, n] - sum_val | |
| else: | |
| cep[:, n] = -sum_val | |
| return cep.reshape(*batch_shape, num_cepstrum) | |
| class PhysicsVocalTractExtractor(nn.Module): | |
| """Branch A: Physics-based Vocal Tract & Glottal Feature Extractor (0-4 kHz). | |
| Extracts: | |
| 1. Linear Prediction Cepstral Coefficients (LPCC): vocal tract resonance shape. | |
| 2. Glottal residual energy ratio & predictability: glottal excitation constraints. | |
| """ | |
| def __init__( | |
| self, | |
| sample_rate: int = 8000, | |
| frame_len: int = 200, # 25 ms @ 8 kHz | |
| hop_len: int = 80, # 10 ms @ 8 kHz | |
| lpc_order: int = 14, # Appropriate for 8kHz narrowband speech | |
| num_lpcc: int = 16 | |
| ): | |
| super().__init__() | |
| self.sample_rate = sample_rate | |
| self.frame_len = frame_len | |
| self.hop_len = hop_len | |
| self.lpc_order = lpc_order | |
| self.num_lpcc = num_lpcc | |
| self.out_channels = num_lpcc + 2 # LPCC + log residual power + residual ratio | |
| def forward(self, waveform: torch.Tensor) -> torch.Tensor: | |
| """Extract physics-based features. | |
| Args: | |
| waveform: (B, num_samples) at 8 kHz | |
| Returns: | |
| features: (B, out_channels, num_frames) | |
| """ | |
| # 1. Frame signal | |
| # (B, num_frames, frame_len) | |
| frames = waveform.unfold(dimension=-1, size=self.frame_len, step=self.hop_len) | |
| # 2. Compute Autocorrelation | |
| r = compute_autocorr(frames, self.lpc_order) | |
| # 3. Levinson-Durbin LPC & Residual Power | |
| lpc, res_power = levinson_durbin(r, self.lpc_order) | |
| # 4. Convert LPC to LPCC | |
| lpcc = lpc_to_lpcc(lpc, self.num_lpcc) # (B, num_frames, num_lpcc) | |
| # 5. Glottal residual constraints | |
| # Frame signal energy | |
| frame_energy = torch.clamp(r[..., 0], min=1e-8) | |
| log_res_power = torch.log(torch.clamp(res_power, min=1e-8)).unsqueeze(-1) | |
| res_ratio = (res_power / frame_energy).unsqueeze(-1) # Predictability ratio | |
| # Combine: (B, num_frames, num_lpcc + 2) | |
| phys_features = torch.cat([lpcc, log_res_power, res_ratio], dim=-1) | |
| # Permute to (B, channels, time) for 1D convolutions | |
| return phys_features.transpose(1, 2) | |
| class SpectrogramExtractor(nn.Module): | |
| """Branch B: Spectrogram & Filterbank Extractor (0-4 kHz). | |
| Supports Log-Mel Spectrogram and Linear Frequency Cepstral Coefficients (LFCC). | |
| """ | |
| def __init__( | |
| self, | |
| sample_rate: int = 8000, | |
| n_fft: int = 512, | |
| frame_len: int = 200, # 25 ms @ 8 kHz | |
| hop_len: int = 80, # 10 ms @ 8 kHz | |
| n_mels: int = 64, | |
| f_min: float = 0.0, | |
| f_max: float = 4000.0, | |
| feature_type: str = "mel" # "mel" or "lfcc" | |
| ): | |
| super().__init__() | |
| self.sample_rate = sample_rate | |
| self.feature_type = feature_type | |
| self.n_mels = n_mels | |
| self.mel_transform = T.MelSpectrogram( | |
| sample_rate=sample_rate, | |
| n_fft=n_fft, | |
| win_length=frame_len, | |
| hop_length=hop_len, | |
| n_mels=n_mels, | |
| f_min=f_min, | |
| f_max=f_max, | |
| power=2.0 | |
| ) | |
| def forward(self, waveform: torch.Tensor) -> torch.Tensor: | |
| """Extract spectral features. | |
| Args: | |
| waveform: (B, num_samples) at 8 kHz | |
| Returns: | |
| features: (B, 1, freq_bins, time_frames) | |
| """ | |
| # Mel spectrogram: (B, n_mels, T) | |
| mel = self.mel_transform(waveform) | |
| log_mel = torch.log(torch.clamp(mel, min=1e-6)) | |
| # Add channel dimension for 2D CNN: (B, 1, n_mels, T) | |
| return log_mel.unsqueeze(1) | |
| class HybridFeatureExtractor(nn.Module): | |
| """Dual-branch hybrid feature extractor orchestrating physics and spectral branches.""" | |
| def __init__( | |
| self, | |
| sample_rate: int = 8000, | |
| frame_len: int = 200, | |
| hop_len: int = 80, | |
| lpc_order: int = 14, | |
| num_lpcc: int = 16, | |
| n_mels: int = 64 | |
| ): | |
| super().__init__() | |
| self.branch_a = PhysicsVocalTractExtractor( | |
| sample_rate=sample_rate, | |
| frame_len=frame_len, | |
| hop_len=hop_len, | |
| lpc_order=lpc_order, | |
| num_lpcc=num_lpcc | |
| ) | |
| self.branch_b = SpectrogramExtractor( | |
| sample_rate=sample_rate, | |
| frame_len=frame_len, | |
| hop_len=hop_len, | |
| n_mels=n_mels | |
| ) | |
| def forward(self, waveform: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """Extract both Branch A and Branch B features. | |
| Args: | |
| waveform: (B, num_samples) at 8 kHz | |
| Returns: | |
| phys_features: (B, 18, T) | |
| spec_features: (B, 1, 64, T) | |
| """ | |
| phys = self.branch_a(waveform) | |
| spec = self.branch_b(waveform) | |
| return phys, spec | |