test.py

import os
import torch
import torch.nn as nn
import numpy as np
import librosa
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter

# ── CONFIG ────────────────────────────────────────────────────
AUDIO_FILE  = "file2.wav"                    # your audio file path
MODEL_PATH  = "crnn_transformer_model (2).pth"  # your saved model path
THRESHOLD   = 0.2
SAVE_PLOT   = None                          # e.g. "result.png" or None to show inline
# ─────────────────────────────────────────────────────────────


# ── Model ─────────────────────────────────────────────────────

class TransformerEncoder(nn.Module):
    def __init__(self, d_model=128, nhead=4, num_layers=2, dropout=0.4):
        super().__init__()
        self.input_proj    = nn.Linear(3200, d_model)
        self.pos_embedding = nn.Parameter(torch.randn(1, 200, d_model) * 0.01)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 2,
            dropout=dropout, batch_first=True, norm_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.norm         = nn.LayerNorm(d_model)

    def forward(self, x):
        x = self.input_proj(x)
        x = x + self.pos_embedding[:, :x.size(1), :]
        x = self.transformer(x)
        return self.norm(x)


class CRNN(nn.Module):
    def __init__(self, num_classes=2, d_model=128, nhead=4, num_transformer_layers=2):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(2, 32,  kernel_size=3, padding=1), nn.BatchNorm2d(32),  nn.ReLU(), nn.Dropout2d(0.1), nn.MaxPool2d(2, 2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64),  nn.ReLU(), nn.Dropout2d(0.1), nn.MaxPool2d(2, 2),
            nn.Conv2d(64, 128,kernel_size=3, padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.Dropout2d(0.2), nn.MaxPool2d(2, 2),
        )
        self.transformer_enc = TransformerEncoder(d_model, nhead, num_transformer_layers, dropout=0.4)
        self.fc = nn.Sequential(
            nn.Linear(d_model, 64), nn.ReLU(), nn.Dropout(0.6), nn.Linear(64, num_classes)
        )

    def forward(self, x):
        x = self.conv(x)
        b, c, h, w = x.size()
        x = x.permute(0, 2, 1, 3).reshape(b, h, c * w)
        x = self.transformer_enc(x)
        return self.fc(x.mean(dim=1))


# ── Feature extraction ────────────────────────────────────────

def extract_features(audio, sr=16000, n_mels=128, n_fft=1024, hop_length=512, max_len=200):
    # Mel
    mel = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=n_mels, n_fft=n_fft, hop_length=hop_length)
    mel = librosa.power_to_db(mel, ref=np.max)
    mel = (mel - mel.mean()) / (mel.std() + 1e-8)

    # LFCC
    stft    = np.abs(librosa.stft(audio, n_fft=n_fft, hop_length=hop_length))
    n_bins  = stft.shape[0]
    freqs   = np.linspace(0, sr // 2, n_bins)
    centers = np.linspace(0, sr // 2, n_mels + 2)
    lf      = np.zeros((n_mels, n_bins))
    for i in range(n_mels):
        lo, mid, hi = centers[i], centers[i+1], centers[i+2]
        up   = (freqs >= lo) & (freqs <= mid)
        down = (freqs >  mid) & (freqs <= hi)
        lf[i, up]   = (freqs[up]   - lo) / (mid - lo + 1e-8)
        lf[i, down] = (hi - freqs[down]) / (hi - mid + 1e-8)
    lfcc = np.dot(lf, stft)
    lfcc = librosa.power_to_db(lfcc ** 2 + 1e-8, ref=np.max)
    lfcc = (lfcc - lfcc.mean()) / (lfcc.std() + 1e-8)

    def pad(spec):
        l = spec.shape[1]
        if l < max_len:
            return np.pad(spec, ((0, 0), (0, max_len - l)), mode='edge' if l >= 2 else 'constant')
        return spec[:, :max_len]

    return pad(mel), pad(lfcc)


# ── Main predict + explain ────────────────────────────────────

def predict_and_explain(audio_file, model_path, threshold, save_plot=None):
    if not os.path.exists(model_path):
        print(f"❌  Model not found: {model_path}"); return
    if not os.path.exists(audio_file):
        print(f"❌  Audio file not found: {audio_file}"); return

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Load model
    model = CRNN().to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()

    # Load audio
    try:
        audio, _ = librosa.load(audio_file, sr=16000)
        if len(audio) == 0:
            audio = np.zeros(16000)
    except Exception as e:
        print(f"Error loading audio: {e}")
        audio = np.zeros(16000)

    mel, lfcc = extract_features(audio)
    features  = np.stack([mel, lfcc], axis=0)          # (2, 128, 200)
    feat_t    = torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(device)

    # ── Prediction ────────────────────────────────────────────
    with torch.no_grad():
        prob      = torch.softmax(model(feat_t), dim=1)
        fake_prob = prob[0][1].item()
        real_prob = prob[0][0].item()

    label = "FAKE 🔴" if fake_prob > threshold else "REAL ✅"

    print(f"\n  File      : {os.path.basename(audio_file)}")
    print(f"  Result    : {label}")
    print(f"  FAKE prob : {fake_prob*100:.1f}%")
    print(f"  REAL prob : {real_prob*100:.1f}%")
    bar = 40
    f   = int(fake_prob * bar)
    print(f"\n  REAL {'█' * (bar-f)}{'░' * f} FAKE")
    print(f"       {real_prob*100:>5.1f}%{' '*(bar-10)}{fake_prob*100:>5.1f}%")

    # ── Saliency map (gradient-based XAI) ────────────────────
    # How it works: we do a forward pass with gradients enabled,
    # then backpropagate from the predicted class score back to
    # the input. Large gradients = input regions the model relied
    # on most heavily for its decision.
    model.eval()
    feat_grad = feat_t.clone().requires_grad_(True)
    output     = model(feat_grad)
    pred_class = output.argmax(dim=1).item()
    model.zero_grad()
    output[0, pred_class].backward()

    # Average absolute gradient across Mel + LFCC channels → (128, 200)
    saliency       = feat_grad.grad.abs().mean(dim=1).squeeze().cpu().numpy()
    saliency_smooth = gaussian_filter(saliency, sigma=1.5)
    saliency_norm   = (saliency_smooth - saliency_smooth.min()) / \
                      (saliency_smooth.max() - saliency_smooth.min() + 1e-8)

    # ── Text summary ──────────────────────────────────────────
    hop_sec    = 512 / 16000
    time_imp   = saliency_norm.mean(axis=0)
    freq_imp   = saliency_norm.mean(axis=1)
    top_frames = np.argsort(time_imp)[-5:][::-1]
    top_bands  = np.argsort(freq_imp)[-5:][::-1]
    mel_freqs  = librosa.mel_frequencies(n_mels=128, fmin=0, fmax=8000)

    print("\n" + "═" * 50)
    print("  EXPLAINABILITY SUMMARY")
    print("═" * 50)
    print(f"  Prediction : {'FAKE' if fake_prob > threshold else 'REAL'}  "
          f"({max(fake_prob, real_prob)*100:.1f}% confidence)")
    print("\n  Most suspicious TIME regions:")
    for fr in top_frames:
        print(f"    {fr*hop_sec:.3f}s – {(fr+1)*hop_sec:.3f}s  "
              f"(importance {time_imp[fr]:.3f})")
    print("\n  Most suspicious FREQUENCY bands:")
    for b in top_bands:
        print(f"    ~{mel_freqs[b]:.0f} Hz  "
              f"(mel band {b}, importance {freq_imp[b]:.3f})")
    print("═" * 50)

    # ── Plot ──────────────────────────────────────────────────
    extent = [0, 200 * 512 / 16000, 0, 8]   # time (s) vs freq (kHz)

    fig, axes = plt.subplots(3, 1, figsize=(12, 11))
    fig.suptitle(
        f"Explainability Report  |  {os.path.basename(audio_file)}\n"
        f"Prediction: {'FAKE' if fake_prob > threshold else 'REAL'}  "
        f"|  FAKE prob: {fake_prob*100:.1f}%  |  Threshold: {threshold}",
        fontsize=13, fontweight='bold'
    )

    # Plot 1 — Mel
    im0 = axes[0].imshow(mel, aspect='auto', origin='lower', cmap='magma', extent=extent)
    axes[0].set_title("Channel 1: Mel Spectrogram")
    axes[0].set_xlabel("Time (s)"); axes[0].set_ylabel("Frequency (kHz)")
    plt.colorbar(im0, ax=axes[0], label="dB (normalised)")

    # Plot 2 — LFCC
    im1 = axes[1].imshow(lfcc, aspect='auto', origin='lower', cmap='viridis', extent=extent)
    axes[1].set_title("Channel 2: LFCC  (linear filterbank — captures synthesis artifacts)")
    axes[1].set_xlabel("Time (s)"); axes[1].set_ylabel("Frequency (kHz)")
    plt.colorbar(im1, ax=axes[1], label="dB (normalised)")

    # Plot 3 — Saliency
    im2 = axes[2].imshow(saliency_norm, aspect='auto', origin='lower', cmap='hot', extent=extent)
    axes[2].set_title("Saliency Map  —  bright regions = where model focused for FAKE/REAL decision")
    axes[2].set_xlabel("Time (s)"); axes[2].set_ylabel("Frequency (kHz)")
    plt.colorbar(im2, ax=axes[2], label="Importance (normalised)")

    plt.tight_layout()

    if save_plot:
        plt.savefig(save_plot, dpi=150, bbox_inches='tight')
        print(f"\n  Plot saved → {save_plot}")
    else:
        plt.show()

    plt.close(fig)


if __name__ == "__main__":
    predict_and_explain(AUDIO_FILE, MODEL_PATH, THRESHOLD, SAVE_PLOT)