Spaces:
Sleeping
Sleeping
| """Dual-Stream Spectrogram and Physics Neural Network Architecture for Audio Deepfake Detection.""" | |
| from typing import Dict, Optional, Tuple | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from src.features import HybridFeatureExtractor, SpectrogramExtractor | |
| from src.physics_features import StackedPhysicsFeatureExtractor | |
| class SpectrogramBranchEncoder(nn.Module): | |
| """Stream A: 2D Convolutional encoder for Log-Mel spectrogram features.""" | |
| def __init__(self, in_channels: int = 1, embed_dim: int = 128): | |
| super().__init__() | |
| self.conv = nn.Sequential( | |
| nn.Conv2d(in_channels, 32, kernel_size=3, stride=1, padding=1, bias=False), | |
| nn.BatchNorm2d(32), | |
| nn.GELU(), | |
| nn.MaxPool2d((2, 2)), | |
| nn.Dropout(0.1), | |
| nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1, bias=False), | |
| nn.BatchNorm2d(64), | |
| nn.GELU(), | |
| nn.MaxPool2d((2, 2)), | |
| nn.Dropout(0.1), | |
| nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=False), | |
| nn.BatchNorm2d(128), | |
| nn.GELU(), | |
| ) | |
| # Pool over frequency bins (height), retain temporal variation | |
| self.freq_pool = nn.AdaptiveAvgPool2d((1, None)) | |
| self.proj = nn.Sequential( | |
| nn.Linear(256, embed_dim), | |
| nn.LayerNorm(embed_dim), | |
| nn.GELU(), | |
| nn.Dropout(0.2), | |
| ) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| """Args: | |
| x: (B, 1, n_mels, T) | |
| Returns: | |
| embedding: (B, embed_dim) | |
| """ | |
| feats = self.conv(x) # (B, 128, F', T') | |
| feats = self.freq_pool(feats).squeeze(2) # (B, 128, T') | |
| mean_pool = torch.mean(feats, dim=-1) | |
| max_pool = torch.max(feats, dim=-1)[0] | |
| pooled = torch.cat([mean_pool, max_pool], dim=-1) # (B, 256) | |
| return self.proj(pooled) | |
| class PhysicsStreamEncoder(nn.Module): | |
| """Stream B: Multi-modal physics encoder unifying 1D dynamics, 2D phase, and 2D bicoherence.""" | |
| def __init__(self, embed_dim: int = 128): | |
| super().__init__() | |
| # 1. 1D Temporal Dynamics (TKEO stats + Glottal residual stats: 6 channels) | |
| self.enc_1d = nn.Sequential( | |
| nn.Conv1d(6, 32, kernel_size=5, stride=1, padding=2, bias=False), | |
| nn.BatchNorm1d(32), | |
| nn.GELU(), | |
| nn.MaxPool1d(2), | |
| nn.Dropout(0.1), | |
| nn.Conv1d(32, 64, kernel_size=5, stride=1, padding=2, bias=False), | |
| nn.BatchNorm1d(64), | |
| nn.GELU(), | |
| ) | |
| self.proj_1d = nn.Sequential( | |
| nn.Linear(128, 64), | |
| nn.LayerNorm(64), | |
| nn.GELU() | |
| ) | |
| # 2. 2D Modified Group Delay (Phase Continuity) | |
| self.enc_mgd = nn.Sequential( | |
| nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1, bias=False), | |
| nn.BatchNorm2d(32), | |
| nn.GELU(), | |
| nn.MaxPool2d((2, 2)), | |
| nn.Dropout(0.1), | |
| nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1, bias=False), | |
| nn.BatchNorm2d(64), | |
| nn.GELU(), | |
| nn.AdaptiveAvgPool2d((1, None)), | |
| ) | |
| self.proj_mgd = nn.Sequential( | |
| nn.Linear(128, 64), | |
| nn.LayerNorm(64), | |
| nn.GELU() | |
| ) | |
| # 3. 2D Bicoherence (Quadratic Phase Coupling) | |
| self.enc_bico = nn.Sequential( | |
| nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1, bias=False), | |
| nn.BatchNorm2d(32), | |
| nn.GELU(), | |
| nn.MaxPool2d(2), | |
| nn.Dropout(0.1), | |
| nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1, bias=False), | |
| nn.BatchNorm2d(64), | |
| nn.GELU(), | |
| nn.AdaptiveAvgPool2d((2, 2)), | |
| ) | |
| self.proj_bico = nn.Sequential( | |
| nn.Linear(256, 64), | |
| nn.LayerNorm(64), | |
| nn.GELU() | |
| ) | |
| # Unified Physics Latent Projection (64 * 3 = 192 -> embed_dim) | |
| self.out_proj = nn.Sequential( | |
| nn.Linear(192, embed_dim), | |
| nn.LayerNorm(embed_dim), | |
| nn.GELU(), | |
| nn.Dropout(0.2) | |
| ) | |
| def forward(self, feats_1d: torch.Tensor, mgd: torch.Tensor, bico: torch.Tensor) -> torch.Tensor: | |
| """Args: | |
| feats_1d: (B, 6, T) | |
| mgd: (B, 1, 64, T) | |
| bico: (B, 1, 32, 32) | |
| Returns: | |
| h_phys: (B, embed_dim) | |
| """ | |
| # 1D Temporal Pooling | |
| f1d = self.enc_1d(feats_1d) | |
| p1d = torch.cat([torch.mean(f1d, dim=-1), torch.max(f1d, dim=-1)[0]], dim=-1) | |
| h_1d = self.proj_1d(p1d) | |
| # 2D MGD Phase Pooling | |
| fmgd = self.enc_mgd(mgd).squeeze(2) | |
| pmgd = torch.cat([torch.mean(fmgd, dim=-1), torch.max(fmgd, dim=-1)[0]], dim=-1) | |
| h_mgd = self.proj_mgd(pmgd) | |
| # 2D Bicoherence Pooling | |
| fbico = self.enc_bico(bico).flatten(1) | |
| h_bico = self.proj_bico(fbico) | |
| # Assemble and project | |
| h_stacked = torch.cat([h_1d, h_mgd, h_bico], dim=-1) # (B, 192) | |
| return self.out_proj(h_stacked) | |
| class CrossModalGatedFusion(nn.Module): | |
| """Attention-based Gating Fusion layer uniting Spectrogram and Physics streams.""" | |
| def __init__(self, embed_dim: int = 128): | |
| super().__init__() | |
| self.gate_fc = nn.Sequential( | |
| nn.Linear(embed_dim * 2, embed_dim), | |
| nn.Sigmoid() | |
| ) | |
| self.fuse_proj = nn.Sequential( | |
| nn.Linear(embed_dim * 4, embed_dim * 2), | |
| nn.LayerNorm(embed_dim * 2), | |
| nn.GELU(), | |
| nn.Dropout(0.3) | |
| ) | |
| def forward(self, h_spec: torch.Tensor, h_phys: torch.Tensor) -> torch.Tensor: | |
| """Args: | |
| h_spec: (B, embed_dim) | |
| h_phys: (B, embed_dim) | |
| Returns: | |
| fused: (B, embed_dim * 2) | |
| """ | |
| concat = torch.cat([h_spec, h_phys], dim=-1) | |
| gate = self.gate_fc(concat) # (B, embed_dim) | |
| gated_blend = gate * h_spec + (1.0 - gate) * h_phys | |
| # Multi-scale fusion: gated blend + individual streams | |
| full_repr = torch.cat([gated_blend, h_spec, h_phys, gate], dim=-1) | |
| return self.fuse_proj(full_repr) | |
| class TandemPhysicsDualStreamDetector(nn.Module): | |
| """End-to-end Dual-Stream Deepfake Detector under Tandem Telephony Compression. | |
| Combines: | |
| - Stream A: Log-Mel Spectrogram representation (2D CNN). | |
| - Stream B: Physics & Math representation (TKEO, GIF jitter/shimmer, MGD phase, Bicoherence QPC). | |
| - Fusion: Cross-modal gated attention. | |
| - Classification: Multi-layer perceptron. | |
| """ | |
| def __init__(self, embed_dim: int = 128, n_mels: int = 64, **kwargs): | |
| super().__init__() | |
| # Feature Extractors | |
| self.spec_extractor = SpectrogramExtractor(sample_rate=8000, n_mels=n_mels) | |
| self.physics_extractor = StackedPhysicsFeatureExtractor(sample_rate=8000) | |
| # Encoders | |
| self.spec_encoder = SpectrogramBranchEncoder(in_channels=1, embed_dim=embed_dim) | |
| self.physics_encoder = PhysicsStreamEncoder(embed_dim=embed_dim) | |
| # Fusion & Classifier | |
| self.fusion = CrossModalGatedFusion(embed_dim=embed_dim) | |
| self.classifier = nn.Sequential( | |
| nn.Linear(embed_dim * 2, 64), | |
| nn.GELU(), | |
| nn.Dropout(0.3), | |
| nn.Linear(64, 1) | |
| ) | |
| def forward(self, waveform: torch.Tensor) -> torch.Tensor: | |
| """Forward pass directly from 8 kHz lossy compressed audio (B, N). | |
| Returns: | |
| logits: (B,) log-odds for binary deepfake classification. | |
| """ | |
| # 1. Feature Extraction | |
| spec_feats = self.spec_extractor(waveform) # (B, 1, 64, T) | |
| phys_dict = self.physics_extractor(waveform) # feats_1d, mgd, bico | |
| # 2. Dual-Stream Encodings | |
| h_spec = self.spec_encoder(spec_feats) | |
| h_phys = self.physics_encoder( | |
| feats_1d=phys_dict["feats_1d"], | |
| mgd=phys_dict["mgd"], | |
| bico=phys_dict["bico"] | |
| ) | |
| # 3. Cross-Modal Fusion | |
| fused = self.fusion(h_spec, h_phys) | |
| # 4. Classification Logit | |
| logits = self.classifier(fused).squeeze(-1) | |
| return logits | |
| # Backward compatibility alias | |
| DualBranchTelephonyDetector = TandemPhysicsDualStreamDetector | |