|
|
|
|
|
|
| import numpy as np
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| import torchaudio
|
|
|
|
|
|
|
|
|
| SR = 16000
|
| N_MELS = 80
|
| N_FFT = 1024
|
| HOP = 256
|
|
|
| MAX_LEN_CONF = 160
|
| CROP_SEC = 2.5
|
| NUM_CROPS = 5
|
| TH_CONF = 0.5
|
|
|
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
|
|
|
|
| mel_extractor = torchaudio.transforms.MelSpectrogram(
|
| sample_rate=SR,
|
| n_fft=N_FFT,
|
| hop_length=HOP,
|
| n_mels=N_MELS,
|
| power=2.0
|
| ).to(DEVICE)
|
|
|
| db_transform = torchaudio.transforms.AmplitudeToDB(stype="power")
|
|
|
|
|
|
|
|
|
| class ConvSubsampling(nn.Module):
|
| def __init__(self):
|
| super().__init__()
|
| self.conv = nn.Sequential(
|
| nn.Conv2d(1, 64, 3, stride=2, padding=1),
|
| nn.ReLU(),
|
| nn.Conv2d(64, 64, 3, stride=2, padding=1),
|
| nn.ReLU(),
|
| )
|
|
|
| def forward(self, x):
|
| x = self.conv(x.unsqueeze(1))
|
| b, c, f, t = x.shape
|
| return x.view(b, c * f, t).transpose(1, 2)
|
|
|
|
|
| class ConformerBlock(nn.Module):
|
| def __init__(self, dim, heads=4, dropout=0.3):
|
| super().__init__()
|
|
|
| self.ff1 = nn.Sequential(
|
| nn.LayerNorm(dim),
|
| nn.Linear(dim, dim * 4),
|
| nn.GELU(),
|
| nn.Dropout(dropout),
|
| nn.Linear(dim * 4, dim),
|
| )
|
|
|
| self.attn = nn.MultiheadAttention(
|
| dim, heads, dropout=dropout, batch_first=True
|
| )
|
| self.norm_attn = nn.LayerNorm(dim)
|
|
|
| self.conv = nn.Sequential(
|
| nn.Conv1d(dim, dim, 3, padding=1, groups=dim),
|
| nn.BatchNorm1d(dim),
|
| nn.GELU(),
|
| nn.Conv1d(dim, dim, 1),
|
| )
|
| self.norm_conv = nn.LayerNorm(dim)
|
|
|
| self.ff2 = nn.Sequential(
|
| nn.LayerNorm(dim),
|
| nn.Linear(dim, dim * 4),
|
| nn.GELU(),
|
| nn.Dropout(dropout),
|
| nn.Linear(dim * 4, dim),
|
| )
|
|
|
| self.norm_out = nn.LayerNorm(dim)
|
|
|
| def forward(self, x):
|
| x = x + 0.5 * self.ff1(x)
|
| attn, _ = self.attn(x, x, x)
|
| x = self.norm_attn(x + attn)
|
| conv = self.conv(x.transpose(1, 2)).transpose(1, 2)
|
| x = self.norm_conv(x + conv)
|
| x = x + 0.5 * self.ff2(x)
|
| return self.norm_out(x)
|
|
|
|
|
| class AttentivePooling(nn.Module):
|
| def __init__(self, dim):
|
| super().__init__()
|
| self.att = nn.Linear(dim, 1)
|
|
|
| def forward(self, x):
|
| w = torch.softmax(self.att(x), dim=1)
|
| return (x * w).sum(dim=1)
|
|
|
|
|
| ENC_DIM = 64 * (N_MELS // 4)
|
|
|
| class AntiDeepfakeConformer(nn.Module):
|
| def __init__(self):
|
| super().__init__()
|
| self.subsample = ConvSubsampling()
|
| self.encoder = nn.Sequential(
|
| ConformerBlock(ENC_DIM),
|
| ConformerBlock(ENC_DIM),
|
| ConformerBlock(ENC_DIM),
|
| )
|
| self.pool = AttentivePooling(ENC_DIM)
|
| self.fc = nn.Sequential(
|
| nn.Linear(ENC_DIM, 128),
|
| nn.ReLU(),
|
| nn.Dropout(0.3),
|
| nn.Linear(128, 1)
|
| )
|
|
|
| def forward(self, x):
|
| x = self.subsample(x)
|
| x = self.encoder(x)
|
| x = self.pool(x)
|
| return self.fc(x).squeeze(1)
|
|
|
|
|
|
|
|
|
| conformer_model = AntiDeepfakeConformer().to(DEVICE)
|
| conformer_model.load_state_dict(
|
| torch.load("anti_ai_conformer_best.pt", map_location=DEVICE)
|
| )
|
| conformer_model.eval()
|
|
|
|
|
|
|
|
|
| def normalize_soft(audio, target_rms=0.08):
|
| rms = np.sqrt(np.mean(audio ** 2) + 1e-9)
|
| scale = min(target_rms / rms, 3.0)
|
| return audio * scale
|
|
|
|
|
| def split_crops(audio, sr=SR, crop_sec=CROP_SEC, num_crops=NUM_CROPS):
|
| crop_len = int(sr * crop_sec)
|
|
|
| if len(audio) <= crop_len:
|
| return [audio]
|
|
|
| step = max((len(audio) - crop_len) // (num_crops - 1), 1)
|
|
|
| crops = []
|
| for i in range(num_crops):
|
| start = i * step
|
| seg = audio[start:start + crop_len]
|
| if len(seg) == crop_len:
|
| crops.append(seg)
|
|
|
| return crops
|
|
|
|
|
|
|
|
|
| def predict_conformer(audio: np.ndarray):
|
| audio = normalize_soft(audio)
|
|
|
| crops = split_crops(audio)
|
| if len(crops) == 0:
|
| return 0.0, "UNKNOWN"
|
|
|
| probs = []
|
|
|
| with torch.no_grad():
|
| for seg in crops:
|
| seg = torch.tensor(seg, dtype=torch.float32).to(DEVICE)
|
|
|
| mel = db_transform(mel_extractor(seg))
|
|
|
| if mel.shape[1] < MAX_LEN_CONF:
|
| mel = F.pad(mel, (0, MAX_LEN_CONF - mel.shape[1]))
|
| else:
|
| mel = mel[:, :MAX_LEN_CONF]
|
|
|
| mel = mel.unsqueeze(0)
|
| prob = torch.sigmoid(conformer_model(mel)).item()
|
| probs.append(prob)
|
|
|
| final_score = float(np.mean(probs))
|
| label = "AI" if final_score >= TH_CONF else "HUMAN"
|
|
|
| return final_score, label
|
|
|
