Upload folder using huggingface_hub

e23cf4c verified 2 months ago

9.56 kB

"""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()

	"""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()