Spaces:
Sleeping
Sleeping
| """Advanced Physics and Mathematical Feature Extractors for Telephony Audio (0-4 kHz). | |
| Implements: | |
| 1. Teager-Kaiser Energy Operator (TKEO): Non-linear mechanical energy tracking. | |
| 2. Glottal Inverse Filtering (GIF): Vocal tract removal, glottal residual isolation, | |
| and cycle-to-cycle Jitter and Shimmer micro-perturbation analysis. | |
| 3. Modified Group Delay (MGD): High-resolution phase continuity representations. | |
| 4. Bicoherence (Higher-Order Spectra): Vectorized Quadratic Phase Coupling (QPC). | |
| """ | |
| from typing import Dict, Optional, Tuple | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from src.features import compute_autocorr, levinson_durbin | |
| class TeagerKaiserEnergyExtractor(nn.Module): | |
| """Teager-Kaiser Energy Operator (TKEO) module. | |
| Discrete non-linear energy tracking: | |
| ψ[x[n]] = x[n]^2 - x[n+1] * x[n-1] | |
| Extracts frame-level energy statistics: | |
| 1. Mean Teager Energy | |
| 2. Teager Energy Variance | |
| 3. Crest Factor (Peak-to-Average Ratio) | |
| """ | |
| def __init__(self, frame_len: int = 200, hop_len: int = 80): | |
| super().__init__() | |
| self.frame_len = frame_len | |
| self.hop_len = hop_len | |
| def forward(self, waveform: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """Args: | |
| waveform: Tensor of shape (B, N) at 8 kHz. | |
| Returns: | |
| raw_tkeo: (B, N) instantaneous TKEO sequence. | |
| frame_stats: (B, 3, num_frames) [mean, var, crest_factor]. | |
| """ | |
| # 1. Compute sample-wise TKEO: ψ[n] = x[n]^2 - x[n+1]*x[n-1] | |
| x = waveform | |
| tkeo_inner = x[..., 1:-1] ** 2 - x[..., 2:] * x[..., :-2] | |
| # Pad edges to maintain exact length N | |
| raw_tkeo = F.pad(tkeo_inner, (1, 1), mode="replicate") | |
| # 2. Frame the TKEO sequence: (B, num_frames, frame_len) | |
| frames = raw_tkeo.unfold(dimension=-1, size=self.frame_len, step=self.hop_len) | |
| # 3. Compute frame statistics | |
| mean_energy = torch.mean(frames, dim=-1) | |
| var_energy = torch.var(frames, dim=-1, unbiased=False) | |
| max_energy = torch.max(frames, dim=-1)[0] | |
| crest_factor = max_energy / torch.clamp(torch.abs(mean_energy), min=1e-6) | |
| # Stack: (B, 3, num_frames) | |
| frame_stats = torch.stack([mean_energy, var_energy, crest_factor], dim=1) | |
| return raw_tkeo, frame_stats | |
| class GlottalInverseFilterExtractor(nn.Module): | |
| """Glottal Inverse Filtering (GIF) with Jitter and Shimmer extraction. | |
| Removes vocal tract formants via LPC inverse filtering to isolate the glottal source residual e[n]. | |
| Measures pitch perturbation (Jitter) and amplitude perturbation (Shimmer) across glottal cycles. | |
| """ | |
| def __init__( | |
| self, | |
| sample_rate: int = 8000, | |
| frame_len: int = 200, | |
| hop_len: int = 80, | |
| lpc_order: int = 14, | |
| min_pitch_hz: float = 60.0, | |
| max_pitch_hz: float = 400.0, | |
| ): | |
| super().__init__() | |
| self.sample_rate = sample_rate | |
| self.frame_len = frame_len | |
| self.hop_len = hop_len | |
| self.lpc_order = lpc_order | |
| # Pitch search range in samples at 8 kHz (e.g. 60 Hz -> 133 samples, 400 Hz -> 20 samples) | |
| self.min_lag = max(2, int(sample_rate / max_pitch_hz)) | |
| self.max_lag = min(frame_len - 1, int(sample_rate / min_pitch_hz)) | |
| def forward(self, waveform: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """Args: | |
| waveform: Tensor of shape (B, N) at 8 kHz. | |
| Returns: | |
| residual: (B, N) glottal excitation waveform e[n]. | |
| glottal_stats: (B, 3, num_frames) [residual_rms, local_jitter, local_shimmer]. | |
| """ | |
| B, N = waveform.shape | |
| # 1. Pre-emphasis to cancel lip radiation (1 - 0.97 z^-1) | |
| pre = torch.cat([waveform[..., :1], waveform[..., 1:] - 0.97 * waveform[..., :-1]], dim=-1) | |
| # 2. Frame signal and estimate vocal tract all-pole filter via Levinson-Durbin | |
| frames = pre.unfold(dimension=-1, size=self.frame_len, step=self.hop_len) | |
| r = compute_autocorr(frames, self.lpc_order) | |
| lpc, _ = levinson_durbin(r, self.lpc_order) # (B, num_frames, lpc_order) | |
| # 3. Time-domain inverse filtering e[n] = x[n] - sum(a_k * x[n-k]) | |
| # Compute global/mean LPC filter per sample or frame-wise | |
| # Mean LPC filter across frames for stable global inverse filtering | |
| mean_lpc = torch.mean(lpc, dim=1) # (B, lpc_order) | |
| # Form FIR inverse filter: [1, -a_1, -a_2, ..., -a_P] | |
| # In conv1d, kernel is (out_channels, in_channels, kernel_size) | |
| residuals = [] | |
| for b in range(B): | |
| inv_filter = torch.cat([torch.tensor([1.0], device=waveform.device), -mean_lpc[b]]) | |
| inv_kernel = inv_filter.flip(dims=[0]).view(1, 1, -1) | |
| sig_b = pre[b:b+1].unsqueeze(1) | |
| pad_left = self.lpc_order | |
| sig_padded = F.pad(sig_b, (pad_left, 0), mode="replicate") | |
| res_b = F.conv1d(sig_padded, inv_kernel).squeeze(1) # (1, N) | |
| residuals.append(res_b) | |
| residual = torch.cat(residuals, dim=0) # (B, N) | |
| # 4. Measure Jitter and Shimmer on framed glottal residual | |
| res_frames = residual.unfold(dimension=-1, size=self.frame_len, step=self.hop_len) | |
| # Autocorrelation of residual frames for pitch cycle tracking | |
| n_fft = 2 ** ((2 * self.frame_len - 1).bit_length()) | |
| res_fft = torch.fft.rfft(res_frames, n=n_fft, dim=-1) | |
| res_r = torch.fft.irfft(res_fft.abs() ** 2, n=n_fft, dim=-1)[..., :self.frame_len] | |
| # Search pitch lag in [min_lag, max_lag] | |
| pitch_lags = res_r[..., self.min_lag:self.max_lag + 1] | |
| best_lag_rel = torch.argmax(pitch_lags, dim=-1) | |
| t0 = best_lag_rel.float() + self.min_lag # Pitch period in samples (B, num_frames) | |
| # Amplitude peak per frame | |
| a0 = torch.max(torch.abs(res_frames), dim=-1)[0] # (B, num_frames) | |
| # Frame residual RMS | |
| res_rms = torch.sqrt(torch.mean(res_frames ** 2, dim=-1) + 1e-8) | |
| # Local Jitter: cycle-to-cycle period perturbation |T0[t] - T0[t-1]| / mean(T0) | |
| diff_t0 = torch.abs(t0[:, 1:] - t0[:, :-1]) | |
| diff_t0 = torch.cat([diff_t0[:, :1], diff_t0], dim=-1) | |
| mean_t0 = torch.clamp(torch.mean(t0, dim=-1, keepdim=True), min=1.0) | |
| jitter = diff_t0 / mean_t0 | |
| # Local Shimmer: cycle-to-cycle amplitude perturbation |A0[t] - A0[t-1]| / mean(A0) | |
| diff_a0 = torch.abs(a0[:, 1:] - a0[:, :-1]) | |
| diff_a0 = torch.cat([diff_a0[:, :1], diff_a0], dim=-1) | |
| mean_a0 = torch.clamp(torch.mean(a0, dim=-1, keepdim=True), min=1e-5) | |
| shimmer = diff_a0 / mean_a0 | |
| # Stack: (B, 3, num_frames) | |
| glottal_stats = torch.stack([res_rms, jitter, shimmer], dim=1) | |
| return residual, glottal_stats | |
| class ModifiedGroupDelayExtractor(nn.Module): | |
| """Modified Group Delay (MGD) feature extractor. | |
| Computes negative derivative of unwrapped phase spectrum using Fourier transforms: | |
| τ(ω) = (X_R * Y_R + X_I * Y_I) / S(ω)^(2γ) | |
| MGD(ω) = sign(τ) * |τ|^α | |
| Exposes unnatural vocoder phase continuity and phase unwrapping irregularities. | |
| """ | |
| def __init__( | |
| self, | |
| n_fft: int = 256, | |
| frame_len: int = 200, | |
| hop_len: int = 80, | |
| gamma: float = 0.9, | |
| alpha: float = 0.4, | |
| num_ceps: int = 64 | |
| ): | |
| super().__init__() | |
| self.n_fft = n_fft | |
| self.frame_len = frame_len | |
| self.hop_len = hop_len | |
| self.gamma = gamma | |
| self.alpha = alpha | |
| self.num_ceps = num_ceps | |
| # Ramp vector n: [0, 1, 2, ..., frame_len - 1] | |
| self.register_buffer("ramp", torch.arange(frame_len, dtype=torch.float32).view(1, 1, -1).contiguous().clone(), persistent=False) | |
| def forward(self, waveform: torch.Tensor) -> torch.Tensor: | |
| """Args: | |
| waveform: Tensor of shape (B, N) at 8 kHz. | |
| Returns: | |
| mgd: Tensor of shape (B, 1, num_ceps, num_frames) | |
| """ | |
| # 1. Frame signal | |
| x_frames = waveform.unfold(dimension=-1, size=self.frame_len, step=self.hop_len) # (B, T, L) | |
| y_frames = x_frames * self.ramp # n * x[n] | |
| # 2. FFT of x[n] and n * x[n] | |
| X = torch.fft.rfft(x_frames, n=self.n_fft, dim=-1) # (B, T, F) | |
| Y = torch.fft.rfft(y_frames, n=self.n_fft, dim=-1) | |
| # 3. Numerator of group delay: X_R * Y_R + X_I * Y_I | |
| num = X.real * Y.real + X.imag * Y.imag | |
| # 4. Smoothed spectral magnitude envelope S(ω) via moving average across frequency | |
| mag = X.abs() | |
| B_dim, T_dim, F_dim = mag.shape | |
| mag_flat = mag.reshape(-1, 1, F_dim) | |
| smooth_mag = F.avg_pool1d(mag_flat, kernel_size=5, stride=1, padding=2).reshape(B_dim, T_dim, F_dim) | |
| denom = torch.clamp(smooth_mag, min=1e-5) ** (2.0 * self.gamma) | |
| # 5. Modified Group Delay | |
| tau = num / denom | |
| mgd = torch.sign(tau) * (torch.clamp(torch.abs(tau), min=1e-8) ** self.alpha) | |
| # Select first num_ceps frequency bins (or interpolate) | |
| if mgd.shape[-1] >= self.num_ceps: | |
| mgd = mgd[..., :self.num_ceps] | |
| else: | |
| mgd = F.pad(mgd, (0, self.num_ceps - mgd.shape[-1])) | |
| # Permute to (B, 1, num_ceps, num_frames) for 2D CNN | |
| mgd_2d = mgd.transpose(1, 2).unsqueeze(1) | |
| return mgd_2d | |
| class VectorizedBicoherenceExtractor(nn.Module): | |
| """Vectorized Bicoherence (Higher-Order Spectra / Bispectrum) module. | |
| Computes Quadratic Phase Coupling (QPC) in the 0-4 kHz band: | |
| B(f1, f2) = E[X(f1) * X(f2) * X*(f1 + f2)] | |
| b^2(f1, f2) = |B(f1, f2)|^2 / (E[|X(f1)X(f2)|^2] * E[|X(f1+f2)|^2]) | |
| Reveals lack of non-linear physical harmonic locking in synthetic neural vocoders. | |
| """ | |
| def __init__(self, n_fft: int = 128, frame_len: int = 200, hop_len: int = 80, num_bins: int = 32): | |
| super().__init__() | |
| self.n_fft = n_fft | |
| self.frame_len = frame_len | |
| self.hop_len = hop_len | |
| self.num_bins = num_bins | |
| # Precompute non-redundant frequency indices grid (i, j) where i + j < num_bins and j >= i | |
| i_idx, j_idx = torch.meshgrid(torch.arange(num_bins), torch.arange(num_bins), indexing="ij") | |
| mask = (i_idx + j_idx < num_bins) & (j_idx >= i_idx) | |
| sum_idx = torch.clamp(i_idx + j_idx, max=num_bins - 1) | |
| self.register_buffer("i_idx", i_idx.contiguous().clone(), persistent=False) | |
| self.register_buffer("j_idx", j_idx.contiguous().clone(), persistent=False) | |
| self.register_buffer("sum_idx", sum_idx.contiguous().clone(), persistent=False) | |
| self.register_buffer("mask", mask.contiguous().clone(), persistent=False) | |
| def forward(self, waveform: torch.Tensor) -> torch.Tensor: | |
| """Args: | |
| waveform: Tensor of shape (B, N) at 8 kHz. | |
| Returns: | |
| bicoherence: Tensor of shape (B, 1, num_bins, num_bins) | |
| """ | |
| # 1. Frame signal | |
| frames = waveform.unfold(dimension=-1, size=self.frame_len, step=self.hop_len) # (B, T, L) | |
| # 2. Window and FFT | |
| window = torch.hamming_window(self.frame_len, device=waveform.device) | |
| w_frames = frames * window | |
| X = torch.fft.rfft(w_frames, n=self.n_fft, dim=-1) # (B, T, F_fft) | |
| # Select first num_bins frequency channels (0 - 4 kHz band) | |
| X_sub = X[..., :self.num_bins] # (B, T, num_bins) complex64 | |
| # 3. Vectorized tensor index expansion | |
| X_i = X_sub[:, :, self.i_idx] # (B, T, num_bins, num_bins) | |
| X_j = X_sub[:, :, self.j_idx] # (B, T, num_bins, num_bins) | |
| X_sum = X_sub[:, :, self.sum_idx] # (B, T, num_bins, num_bins) | |
| # Bispectrum expectation across frames dimension (T) | |
| term = X_i * X_j * torch.conj(X_sum) | |
| bispectrum = torch.mean(term, dim=1) # (B, num_bins, num_bins) | |
| # Normalization denominator | |
| denom1 = torch.mean((X_i * X_j).abs() ** 2, dim=1) | |
| denom2 = torch.mean(X_sum.abs() ** 2, dim=1) | |
| denom = torch.clamp(denom1 * denom2, min=1e-8) | |
| # Normalized bicoherence | |
| bico = (bispectrum.abs() ** 2) / denom | |
| bico = bico * self.mask.unsqueeze(0).float() | |
| # Add channel dimension: (B, 1, num_bins, num_bins) | |
| return bico.unsqueeze(1) | |
| class StackedPhysicsFeatureExtractor(nn.Module): | |
| """Orchestrates all physics and mathematical feature extractors. | |
| Combines: | |
| - 1D Temporal Features (TKEO stats + Glottal residual stats): (B, 6, num_frames) | |
| - 2D Phase Spectrum (Modified Group Delay MGD): (B, 1, 64, num_frames) | |
| - 2D Harmonic Coupling (Vectorized Bicoherence): (B, 1, 32, 32) | |
| """ | |
| def __init__( | |
| self, | |
| sample_rate: int = 8000, | |
| frame_len: int = 200, | |
| hop_len: int = 80, | |
| lpc_order: int = 14, | |
| mgd_ceps: int = 64, | |
| bico_bins: int = 32 | |
| ): | |
| super().__init__() | |
| self.tkeo = TeagerKaiserEnergyExtractor(frame_len=frame_len, hop_len=hop_len) | |
| self.gif = GlottalInverseFilterExtractor( | |
| sample_rate=sample_rate, | |
| frame_len=frame_len, | |
| hop_len=hop_len, | |
| lpc_order=lpc_order | |
| ) | |
| self.mgd = ModifiedGroupDelayExtractor( | |
| frame_len=frame_len, | |
| hop_len=hop_len, | |
| num_ceps=mgd_ceps | |
| ) | |
| self.bico = VectorizedBicoherenceExtractor( | |
| frame_len=frame_len, | |
| hop_len=hop_len, | |
| num_bins=bico_bins | |
| ) | |
| def forward(self, waveform: torch.Tensor) -> Dict[str, torch.Tensor]: | |
| """Args: | |
| waveform: Tensor of shape (B, N) at 8 kHz. | |
| Returns: | |
| Dict containing: | |
| - 'feats_1d': (B, 6, num_frames) [TKEO stats (3) + Glottal stats (3)] | |
| - 'mgd': (B, 1, 64, num_frames) [Modified Group Delay 2D] | |
| - 'bico': (B, 1, 32, 32) [Bicoherence Matrix 2D] | |
| - 'raw_tkeo': (B, N) | |
| - 'residual': (B, N) | |
| """ | |
| raw_tkeo, tkeo_stats = self.tkeo(waveform) | |
| residual, gif_stats = self.gif(waveform) | |
| mgd_2d = self.mgd(waveform) | |
| bico_2d = self.bico(waveform) | |
| # Align 1D temporal frame counts if minor boundary discrepancy | |
| min_frames = min(tkeo_stats.shape[-1], gif_stats.shape[-1], mgd_2d.shape[-1]) | |
| tkeo_stats = tkeo_stats[..., :min_frames] | |
| gif_stats = gif_stats[..., :min_frames] | |
| mgd_2d = mgd_2d[..., :min_frames] | |
| # Stack 1D temporal physics features: (B, 6, min_frames) | |
| feats_1d = torch.cat([tkeo_stats, gif_stats], dim=1) | |
| return { | |
| "feats_1d": feats_1d, | |
| "mgd": mgd_2d, | |
| "bico": bico_2d, | |
| "raw_tkeo": raw_tkeo, | |
| "residual": residual | |
| } | |