scripts/compare_models.py

"""So sánh DINOv3 ViT-Small vs DINOv3 ConvNeXt-Tiny trên cùng điều kiện.

Điều kiện công bằng:
  - Cùng dataset (DeepFakeFace splits via CSV)
  - Cùng preprocessing (resize 256×256, ImageNet normalization)
  - Cùng protocol đánh giá (linear probe: StandardScaler + LogisticRegression balanced)
  - Cùng seed, cùng batch size
  - Cả 2 đều là frozen feature extractor (không fine-tune)

Output:
  - Bảng so sánh metrics (accuracy, precision, recall, F1, ROC-AUC)
  - Báo cáo JSON chi tiết
"""
import argparse
import json
import os
import sys
import time

import numpy as np
import torch
from PIL import Image
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
)
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from torchvision import transforms
from tqdm import tqdm

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
from src.models.dinov3_vit import load_dinov3
from src.models.dinov3_convnext import load_dinov3_convnext

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
IMG_SIZE = 256


def build_transform():
    return transforms.Compose([
        transforms.Resize((IMG_SIZE, IMG_SIZE)),
        transforms.ToTensor(),
        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
    ])


def load_csv(csv_path: str):
    """Đọc CSV (path,label) → list[(path, label)]."""
    import csv

    items = []
    with open(csv_path, newline="") as f:
        reader = csv.reader(f)
        next(reader, None)  # skip header
        for row in reader:
            if len(row) >= 2:
                p, lbl = row[0], int(row[1])
                if os.path.exists(p):
                    items.append((p, lbl))
    return items


@torch.no_grad()
def extract_features(model, items, device, batch_size: int = 32, desc: str = "Extracting"):
    """Trích xuất feature vectors từ backbone (frozen)."""
    model.to(device).eval()
    tf = build_transform()
    feats, labels = [], []

    for i in tqdm(range(0, len(items), batch_size), desc=desc, unit="batch"):
        batch = items[i : i + batch_size]
        imgs = [tf(Image.open(p).convert("RGB")) for p, _ in batch]
        x = torch.stack(imgs).to(device)
        f = model(x)
        feats.append(f.cpu().numpy())
        labels.extend(lbl for _, lbl in batch)

    return np.vstack(feats), np.array(labels)


def evaluate_linear_probe(X_train, y_train, X_test, y_test, seed: int = 42):
    """Train linear probe trên train, đánh giá trên test."""
    clf = make_pipeline(
        StandardScaler(),
        LogisticRegression(class_weight="balanced", max_iter=2000, random_state=seed),
    )
    clf.fit(X_train, y_train)

    y_pred = clf.predict(X_test)
    y_prob = clf.predict_proba(X_test)[:, 1]

    cm = confusion_matrix(y_test, y_pred, labels=[0, 1])
    return {
        "accuracy": float(accuracy_score(y_test, y_pred)),
        "precision": float(precision_score(y_test, y_pred)),
        "recall": float(recall_score(y_test, y_pred)),
        "f1": float(f1_score(y_test, y_pred)),
        "roc_auc": float(roc_auc_score(y_test, y_prob)),
        "confusion_matrix": cm.tolist(),
        "tn": int(cm[0][0]),
        "fp": int(cm[0][1]),
        "fn": int(cm[1][0]),
        "tp": int(cm[1][1]),
    }


def print_model_result(name: str, n_params: float, feat_dim: int, elapsed: float, m: dict):
    """In kết quả của 1 model."""
    print(f"\n{'='*65}")
    print(f"  {name}")
    print(f"  Params: {n_params:.1f}M | Feature dim: {feat_dim} | Time: {elapsed:.0f}s")
    print(f"{'='*65}")
    print(f"  Accuracy : {m['accuracy']:.4f}")
    print(f"  Precision: {m['precision']:.4f}   (fake = positive)")
    print(f"  Recall   : {m['recall']:.4f}")
    print(f"  F1       : {m['f1']:.4f}")
    print(f"  ROC-AUC  : {m['roc_auc']:.4f}")
    print(f"  CM       : TN={m['tn']}  FP={m['fp']}  FN={m['fn']}  TP={m['tp']}")


def print_comparison_table(vit_result: dict, cnn_result: dict):
    """In bảng so sánh side-by-side."""
    print(f"\n{'='*70}")
    print(f"{'METRIC':<15} {'ViT-S/16':>12} {'ConvNeXt-T':>12} {'Δ (CNN-ViT)':>14}")
    print(f"{'-'*15} {'-'*12} {'-'*12} {'-'*14}")

    for key, label in [
        ("accuracy", "Accuracy"),
        ("precision", "Precision"),
        ("recall", "Recall"),
        ("f1", "F1 Score"),
        ("roc_auc", "ROC-AUC"),
    ]:
        vit_val = vit_result[key]
        cnn_val = cnn_result[key]
        delta = cnn_val - vit_val
        direction = "↑" if delta > 0 else "↓" if delta < 0 else "—"
        print(f"  {label:<13} {vit_val:>10.4f}  {cnn_val:>10.4f}  {delta:+10.4f} {direction}")

    print(f"{'='*70}")


def main():
    parser = argparse.ArgumentParser(
        description="So sánh DINOv3 ViT vs ConvNeXt (cùng điều kiện)"
    )
    parser.add_argument("--train-csv", default="data/splits/train.csv",
                        help="CSV file cho tập train")
    parser.add_argument("--test-csv", default="data/splits/test.csv",
                        help="CSV file cho tập test")
    parser.add_argument("--vit-model", default="models/dinov3_small/model.safetensors",
                        help="Path đến DINOv3 ViT safetensors")
    parser.add_argument("--cnn-model", default="models/dinov3_next_cnn/model-2.safetensors",
                        help="Path đến DINOv3 ConvNeXt safetensors")
    parser.add_argument("--output", default="outputs/results/comparison_report.json",
                        help="Path output JSON report")
    parser.add_argument("--device", default="auto",
                        choices=["auto", "cuda", "mps", "cpu"])
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    # ---------- Device ----------
    if args.device == "auto":
        device = "cuda" if torch.cuda.is_available() else \
                 "mps" if torch.backends.mps.is_available() else "cpu"
    else:
        device = args.device
    print(f"Device: {device}")

    # ---------- Data ----------
    print(f"\nLoad data:")
    train_items = load_csv(args.train_csv)
    test_items = load_csv(args.test_csv)
    print(f"  Train: {len(train_items)} ảnh (real={sum(1 for _,l in train_items if l==0)}, "
          f"fake={sum(1 for _,l in train_items if l==1)})")
    print(f"  Test : {len(test_items)} ảnh (real={sum(1 for _,l in test_items if l==0)}, "
          f"fake={sum(1 for _,l in test_items if l==1)})")

    # ---------- Models ----------
    models_config = [
        {
            "name": "DINOv3 ViT-S/16",
            "key": "vit",
            "path": args.vit_model,
            "loader": load_dinov3,
            "loader_kwargs": {"img_size": IMG_SIZE},
        },
        {
            "name": "DINOv3 ConvNeXt-Tiny",
            "key": "cnn",
            "path": args.cnn_model,
            "loader": load_dinov3_convnext,
            "loader_kwargs": {},
        },
    ]

    report = {
        "config": {
            "img_size": IMG_SIZE,
            "batch_size": args.batch_size,
            "seed": args.seed,
            "device": device,
        },
        "data": {
            "train_csv": args.train_csv,
            "test_csv": args.test_csv,
            "train_samples": len(train_items),
            "test_samples": len(test_items),
        },
        "models": {},
    }

    metrics = {}

    for cfg in models_config:
        print(f"\n{'─'*65}")
        print(f"→ {cfg['name']}")

        # Load
        print(f"  Load: {cfg['path']}")
        model = cfg["loader"](cfg["path"], **cfg["loader_kwargs"])
        n_params = sum(p.numel() for p in model.parameters()) / 1e6
        feat_dim = model.embed_dim
        print(f"  Params: {n_params:.1f}M | Feature dim: {feat_dim}")

        # Extract features
        t0 = time.time()
        X_train, y_train = extract_features(model, train_items, device, args.batch_size,
                                            desc=f"  Train features ({cfg['key']})")
        X_test, y_test = extract_features(model, test_items, device, args.batch_size,
                                          desc=f"  Test features  ({cfg['key']})")
        elapsed = time.time() - t0

        # Evaluate
        m = evaluate_linear_probe(X_train, y_train, X_test, y_test, args.seed)
        metrics[cfg["key"]] = m

        print_model_result(cfg["name"], n_params, feat_dim, elapsed, m)

        report["models"][cfg["key"]] = {
            "name": cfg["name"],
            "params_M": round(n_params, 1),
            "feature_dim": feat_dim,
            "extract_time_s": round(elapsed, 1),
            "metrics": m,
        }

        # Giải phóng model khỏi GPU
        del model
        if device in ("cuda", "mps"):
            import gc
            gc.collect()
            if device == "cuda":
                torch.cuda.empty_cache()
            elif device == "mps":
                torch.mps.empty_cache()

    # ---------- Comparison Table ----------
    print_comparison_table(metrics["vit"], metrics["cnn"])

    # ---------- Save Report ----------
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    with open(args.output, "w") as f:
        json.dump(report, f, indent=2, ensure_ascii=False)
    print(f"\nĐã lưu báo cáo: {args.output}")


if __name__ == "__main__":
    main()