app.py

import os
import math
import tempfile
import subprocess
import sys
import numpy as np
import scipy.io.wavfile as wavfile
import imageio_ffmpeg

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio

from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, JSONResponse

SAMPLE_RATE = 16000
TARGET_SAMPLES = 64000
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

app = FastAPI(title="AASIST Audio Deepfake Detector API", version="2.5.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class SincConvFast(nn.Module):
    def __init__(self, out_channels=128, kernel_size=251, sample_rate=16000, min_low_hz=50, min_band_hz=50):
        super().__init__()
        self.out_channels = out_channels
        self.kernel_size = kernel_size if kernel_size % 2 != 0 else kernel_size + 1
        self.sample_rate = sample_rate
        self.min_low_hz = min_low_hz
        self.min_band_hz = min_band_hz
        low_hz = 30
        high_hz = sample_rate / 2 - (min_low_hz + min_band_hz)
        band_hz = (high_hz - low_hz) / out_channels
        hz = torch.linspace(low_hz, high_hz, out_channels + 1)
        self.low_hz_ = nn.Parameter(hz[:-1].unsqueeze(1))
        self.band_hz_ = nn.Parameter(torch.full((out_channels, 1), band_hz))
        half = kernel_size // 2
        n_lin = torch.linspace(1, half, half)
        window = 0.54 - 0.46 * torch.cos(2 * math.pi * n_lin / self.kernel_size)
        n = 2 * math.pi * n_lin / self.sample_rate
        self.register_buffer('window_', window)
        self.register_buffer('n_', n)

    def forward(self, x):
        if x.ndim == 2:
            x = x.unsqueeze(1)
        low = self.min_low_hz + torch.abs(self.low_hz_)
        high = torch.clamp(low + self.min_band_hz + torch.abs(self.band_hz_), self.min_low_hz, self.sample_rate / 2)
        band = (high - low)[:, 0]
        f_times_t_low = torch.matmul(low, self.n_.unsqueeze(0))
        f_times_t_high = torch.matmul(high, self.n_.unsqueeze(0))
        band_pass_left = ((torch.sin(f_times_t_high) - torch.sin(f_times_t_low)) / (self.n_ / 2 + 1e-8)) * self.window_
        band_pass_center = 2 * band.unsqueeze(1)
        band_pass_right = torch.flip(band_pass_left, dims=[1])
        band_pass = torch.cat([band_pass_left, band_pass_center, band_pass_right], dim=1)
        band_pass = band_pass / (2 * band.unsqueeze(1) + 1e-8)
        filters = band_pass.view(self.out_channels, 1, self.kernel_size)
        return F.conv1d(x, filters, stride=10, padding=self.kernel_size // 2)

class AASISTResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm1d(out_channels)
        self.selu = nn.SELU()
        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm1d(out_channels)
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=stride),
                nn.BatchNorm1d(out_channels)
            )

    def forward(self, x):
        res = self.shortcut(x)
        out = self.selu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        return self.selu(out + res)

class AASISTClassifier(nn.Module):
    def __init__(self, num_classes=2, in_channels=128):
        super().__init__()
        self.sinc_conv = SincConvFast(out_channels=in_channels, kernel_size=251)
        self.bn0 = nn.BatchNorm1d(in_channels)
        self.selu = nn.SELU()
        self.b1 = AASISTResidualBlock(in_channels, 32)
        self.b2 = AASISTResidualBlock(32, 32)
        self.p1 = nn.MaxPool1d(3)
        self.b3 = AASISTResidualBlock(32, 64)
        self.b4 = AASISTResidualBlock(64, 64)
        self.p2 = nn.MaxPool1d(3)
        self.b5 = AASISTResidualBlock(64, 128)
        self.b6 = AASISTResidualBlock(128, 128)
        self.p3 = nn.MaxPool1d(3)
        self.att_conv = nn.Conv1d(128, 1, kernel_size=1)
        self.fc1 = nn.Linear(256, 64)
        self.dropout = nn.Dropout(0.25)
        self.fc2 = nn.Linear(64, num_classes)

    def forward(self, input_values, labels=None):
        x = input_values if input_values.ndim == 2 else input_values.squeeze(1)
        x = self.selu(self.bn0(torch.abs(self.sinc_conv(x))))
        x = self.p1(self.b2(self.b1(x)))
        x = self.p2(self.b4(self.b3(x)))
        x = self.p3(self.b6(self.b5(x)))
        w = torch.softmax(self.att_conv(x), dim=-1)
        mu = torch.sum(x * w, dim=-1)
        std = torch.sqrt(torch.clamp(torch.sum((x - mu.unsqueeze(-1))**2 * w, dim=-1), min=1e-5))
        pooled = torch.cat([mu, std], dim=-1)
        logits = self.fc2(self.dropout(self.selu(self.fc1(pooled))))
        loss = nn.CrossEntropyLoss(label_smoothing=0.10)(logits, labels.to(x.device)) if labels is not None else None
        probs = torch.softmax(logits, dim=-1)
        return {'loss': loss, 'logits': logits, 'probs': probs, 'real_score': probs[:, 1]}

GLOBAL_MODEL = AASISTClassifier().to(DEVICE)
CHECKPOINT_PATH = './aasist_speech_only_best.pt'
LOADED = False

if os.path.exists(CHECKPOINT_PATH):
    try:
        GLOBAL_MODEL.load_state_dict(torch.load(CHECKPOINT_PATH, map_location=DEVICE))
        LOADED = True
        print(f"Loaded checkpoint: {CHECKPOINT_PATH}")
    except Exception as e:
        print(f"Failed to load checkpoint: {e}")
else:
    print(f"Checkpoint not found at {CHECKPOINT_PATH}, using baseline model")

GLOBAL_MODEL.eval()

def preprocess_waveform(waveform, orig_sr):
    if waveform.ndim == 2 and waveform.shape[0] > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)
    elif waveform.ndim == 1:
        waveform = waveform.unsqueeze(0)
    waveform = waveform.squeeze(0)
    if orig_sr != SAMPLE_RATE and orig_sr > 0:
        waveform = torchaudio.transforms.Resample(orig_sr, SAMPLE_RATE)(waveform.unsqueeze(0)).squeeze(0)
    peak = waveform.abs().max()
    if peak > 1e-6:
        waveform = waveform / peak * 0.95
    return waveform

def decode_audio_bytes(file_bytes, filename="audio.mp3"):
    ext = os.path.splitext(filename)[1].lower() or ".mp3"
    tmp_in_path = None
    tmp_wav_path = None
    try:
        tmp_in = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
        tmp_in.write(file_bytes)
        tmp_in_path = tmp_in.name
        tmp_in.close()

        tmp_wav_path = tmp_in_path + "_converted.wav"
        ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
        cmd = [
            ffmpeg_exe, "-y", "-i", tmp_in_path,
            "-ar", "16000", "-ac", "1", "-f", "wav", tmp_wav_path
        ]
        res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if not os.path.exists(tmp_wav_path) or os.path.getsize(tmp_wav_path) == 0:
            raise ValueError(f"FFmpeg failed: {res.stderr.decode('utf-8', errors='ignore')}")
        sr, arr = wavfile.read(tmp_wav_path)
        if arr.dtype == np.int16:
            arr = arr.astype(np.float32) / 32768.0
        elif arr.dtype == np.int32:
            arr = arr.astype(np.float32) / 2147483648.0
        elif arr.dtype != np.float32:
            arr = arr.astype(np.float32)
        waveform = torch.from_numpy(arr).float()
        return preprocess_waveform(waveform, sr)
    finally:
        if tmp_in_path and os.path.exists(tmp_in_path):
            try: os.remove(tmp_in_path)
            except: pass
        if tmp_wav_path and os.path.exists(tmp_wav_path):
            try: os.remove(tmp_wav_path)
            except: pass

def analyze_audio_tensor(waveform, threshold=0.5, fake_chunk_threshold=0.15):
    crop_len = TARGET_SAMPLES
    step_len = TARGET_SAMPLES // 2
    crops = []

    if waveform.shape[0] < crop_len:
        crops.append(F.pad(waveform, (0, crop_len - waveform.shape[0])))
    else:
        start = 0
        while start + crop_len <= waveform.shape[0]:
            crops.append(waveform[start:start + crop_len])
            start += step_len
        if len(crops) == 0:
            crops.append(waveform[:crop_len])

    real_scores = []
    with torch.no_grad():
        for crop in crops:
            inp = crop.unsqueeze(0).to(DEVICE)
            out = GLOBAL_MODEL(inp)
            real_scores.append(out['real_score'].item())

    real_scores = np.array(real_scores)
    avg_real = float(np.mean(real_scores))
    min_real = float(np.min(real_scores))
    p25_real = float(np.percentile(real_scores, 25))

    fake_chunks_count = int(np.sum(real_scores < threshold))
    fake_crop_ratio = float(fake_chunks_count / len(crops))

    is_fake = (fake_crop_ratio >= fake_chunk_threshold) or (p25_real < 0.40)
    is_real = not is_fake

    label = "FAKE" if is_fake else "REAL"
    verdict_ar = "🚨 DEEPFAKE (صوت مزيف ذكاء اصطناعي)" if is_fake else "✅ REAL (صوت بشري أصلي)"
    confidence = (1.0 - p25_real) if is_fake else avg_real

    timeline = []
    for i, s in enumerate(real_scores):
        is_chunk_fake = bool(s < threshold)
        timeline.append({
            "segment": i + 1,
            "segment_index": i + 1,
            "start_sec": round(i * 2.0, 1),
            "end_sec": round(i * 2.0 + 4.0, 1),
            "start_time_sec": round(i * 2.0, 1),
            "end_time_sec": round(i * 2.0 + 4.0, 1),
            "real_prob": round(float(s) * 100.0, 2),
            "fake_prob": round((1.0 - float(s)) * 100.0, 2),
            "prediction": "FAKE" if is_chunk_fake else "REAL",
            "is_fake": is_chunk_fake
        })

    return {
        "prediction": label,
        "verdict_code": label,
        "verdict_ar": verdict_ar,
        "is_real": is_real,
        "confidence": round(confidence, 4),
        "overall_confidence_percent": round(confidence * 100.0, 2),
        "real_probability": round(avg_real, 4),
        "fake_probability": round(1.0 - avg_real, 4),
        "fake_segment_ratio_percent": round(fake_crop_ratio * 100.0, 1),
        "segments_count": len(crops),
        "total_chunks": len(crops),
        "fake_chunks_count": fake_chunks_count,
        "p25_real_percent": round(p25_real * 100.0, 2),
        "min_real_percent": round(min_real * 100.0, 2),
        "mean_real_percent": round(avg_real * 100.0, 2),
        "threshold": threshold,
        "model_loaded": LOADED,
        "checkpoint": CHECKPOINT_PATH if LOADED else "Baseline",
        "timeline": timeline
    }

@app.get("/api/health")
def health_check():
    return {
        "status": "healthy",
        "device": str(DEVICE),
        "model_loaded": LOADED,
        "checkpoint": CHECKPOINT_PATH if LOADED else None
    }

@app.post("/api/analyze-file")
async def analyze_file(file: UploadFile = File(...)):
    filename = file.filename or "uploaded_audio.mp3"
    contents = await file.read()
    if len(contents) == 0:
        raise HTTPException(status_code=400, detail="Empty file")
    try:
        waveform = decode_audio_bytes(contents, filename)
        res = analyze_audio_tensor(waveform)
        res["filename"] = filename
        return JSONResponse(content=res)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/analyze-youtube")
async def analyze_youtube(youtube_url: str = Form(...)):
    if not youtube_url or ("youtube.com" not in youtube_url and "youtu.be" not in youtube_url):
        raise HTTPException(status_code=400, detail="Invalid YouTube URL")
    tmp_wav = None
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
            tmp_wav = tmp.name
        ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()

        try:
            import yt_dlp
        except ImportError:
            subprocess.run([sys.executable, "-m", "pip", "install", "-U", "yt-dlp"])
            import yt_dlp

        ydl_opts = {
            'format': 'bestaudio/best',
            'outtmpl': tmp_wav.replace('.wav', '') + '.%(ext)s',
            'ffmpeg_location': ffmpeg_exe,
            'quiet': True,
            'no_warnings': True,
            'nocheckcertificate': True,
            'user_agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
            'extractor_args': {
                'youtube': {
                    'player_client': ['android', 'web', 'tv']
                }
            }
        }

        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(youtube_url, download=True)
            dl_filename = ydl.prepare_filename(info)

        wav_converted = tmp_wav if tmp_wav.endswith('.wav') else tmp_wav + '.wav'
        cmd = [ffmpeg_exe, "-y", "-i", dl_filename, "-ar", "16000", "-ac", "1", "-f", "wav", wav_converted]
        subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)

        if dl_filename != wav_converted and os.path.exists(dl_filename):
            try: os.remove(dl_filename)
            except: pass

        if not os.path.exists(wav_converted) or os.path.getsize(wav_converted) == 0:
            raise HTTPException(status_code=400, detail="Failed to extract audio from YouTube video.")

        sr, arr = wavfile.read(wav_converted)
        if arr.dtype == np.int16:
            arr = arr.astype(np.float32) / 32768.0
        elif arr.dtype == np.int32:
            arr = arr.astype(np.float32) / 2147483648.0
        elif arr.dtype != np.float32:
            arr = arr.astype(np.float32)

        waveform = torch.from_numpy(arr).float()
        waveform = preprocess_waveform(waveform, sr)
        res = analyze_audio_tensor(waveform)
        res["source"] = youtube_url
        
        if os.path.exists(wav_converted):
            try: os.remove(wav_converted)
            except: pass

        return JSONResponse(content=res)
    except HTTPException as he:
        raise he
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

static_dir = os.path.join(os.path.dirname(__file__), "static")
if not os.path.exists(static_dir):
    os.makedirs(static_dir)
app.mount("/static", StaticFiles(directory=static_dir), name="static")

@app.get("/")
def read_root():
    index_path = os.path.join(static_dir, "index.html")
    if os.path.exists(index_path):
        with open(index_path, "r", encoding="utf-8") as f:
            return HTMLResponse(content=f.read())
    return HTMLResponse(content="<h1>AASIST API Running</h1>")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="127.0.0.1", port=8000, reload=True)