Rethinking-the-readout / src /eval_opensource.py

Mirror code and weights from GitHub

63f8a47 verified 5 days ago

8.41 kB

"""Evaluate a trained V-PVP head on the AIGVDBench Open-Source split.
Computes per-generator AUC, mean-of-generators AUC (mga AUC), and overall
AUC on the test set. Each backbone is loaded once (frozen, fp16 autocast),
patch features ``z`` of shape ``(B, P, T', D)`` are extracted, the velocity
field ``v = z[:, :, 1:] - z[:, :, :-1]`` is computed, and the head produces
a scalar logit per video.
Usage:
    python -m src.eval_opensource \
        --backbone videomae_k400 \
        --backbone_weights <path-or-hub-id> \
        --head_ckpt    checkpoints/pvp_heads/videomae_k400_vpvp.pt \
        --data_root    <path-to-aigvdbench> \
        --num_frames   16 \
        --batch_size   32 \
        --out_json     results/videomae_open_source.json
Expected mga AUC on the Open-Source split with the released checkpoint
(20 generators, 3000 real + 60000 fake videos):
    videomae_k400     ->  0.9528   (V-PVP row of Table 1)
"""
import argparse
import json
import os
import sys
import time

import numpy as np
import torch
import torch.nn as nn
from sklearn.metrics import roc_auc_score
from torch.utils.data import DataLoader

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from backbones import BACKBONE_REGISTRY
from dataset import OpenSourceConfig, OpenSourceTest, collate
from pvp_head import PVPHead


# Frame count the released checkpoint was trained at.
DEFAULT_NUM_FRAMES = {
    "videomae_k400": 16,
}

# Dataset-side normalization applied by ``OpenSourceTest`` (mean 0.45 /
# std 0.225). The eval loop undoes it and re-applies the backbone's own
# statistics below; this is the exact preprocessing path the released
# checkpoint was trained and evaluated with.
DS_MEAN = (0.45, 0.45, 0.45)
DS_STD  = (0.225, 0.225, 0.225)


def renormalize(clips, ds_mean, ds_std, bb_mean, bb_std):
    """Undo dataset normalization, re-apply backbone's. clips: (B, C, T, H, W)."""
    x = clips * ds_std + ds_mean
    return (x - bb_mean) / bb_std


def compute_metrics(scores, labels, cats, threshold=0.5):
    """mga AUC = mean-of-generators AUC over the per-(real, generator) pools."""
    cats   = np.asarray(cats)
    labels = np.asarray(labels)
    scores = np.asarray(scores)
    real   = (labels == 0)
    preds  = (scores > threshold).astype(int)

    auc_per_gen, acc_per_gen = {}, {}
    for c in sorted(set(cats[~real].tolist())):
        gen_mask = (cats == c) & (labels == 1)
        if gen_mask.sum() < 2 or real.sum() < 2:
            continue
        y = np.r_[np.zeros(real.sum()), np.ones(gen_mask.sum())]
        s = np.r_[scores[real],          scores[gen_mask]]
        p = np.r_[preds[real],           preds[gen_mask]]
        auc_per_gen[c] = float(roc_auc_score(y, s))
        acc_per_gen[c] = float((p == y).mean())

    return {
        "auc_per_gen": auc_per_gen,
        "acc_per_gen": acc_per_gen,
        "mga_auc":     float(np.mean(list(auc_per_gen.values()))),
        "mga_acc":     float(np.mean(list(acc_per_gen.values()))),
        "overall_auc": float(roc_auc_score(labels, scores)),
        "overall_acc": float((preds == labels).mean()),
        "n_real":      int(real.sum()),
        "n_fake":      int((~real).sum()),
    }


@torch.no_grad()
def evaluate(backbone, head, loader, device, ds_norm, bb_norm):
    head.eval()
    backbone.eval()
    scores_all, labels_all, cats_all = [], [], []
    n_done = 0
    t0 = time.time()
    for clips, labels, cats in loader:
        clips = clips.to(device, non_blocking=True)
        x = renormalize(clips, *ds_norm, *bb_norm)
        with torch.autocast(device_type="cuda", dtype=torch.float16):
            z = backbone.extract_patch_features(x)
        z = z.to(torch.float32)
        if z.size(2) < 2:
            raise RuntimeError(f"backbone returned T'={z.size(2)} < 2; need >=2 frames")
        v = z[:, :, 1:, :] - z[:, :, :-1, :]
        logit = head(z, v)
        scores_all.append(torch.sigmoid(logit).cpu().numpy())
        labels_all.append(labels.numpy())
        cats_all.extend(cats)
        n_done += clips.size(0)
        if n_done % 1024 == 0:
            print(f"  [{n_done}] {n_done/(time.time()-t0):.1f} v/s", flush=True)
    return (np.concatenate(scores_all),
            np.concatenate(labels_all),
            np.array(cats_all))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--backbone", required=True,
                    choices=list(BACKBONE_REGISTRY.keys()))
    ap.add_argument("--backbone_weights", required=True,
                    help="path (or HuggingFace repo id) of the K400 backbone weights")
    ap.add_argument("--head_ckpt", required=True,
                    help="path to the V-PVP head state_dict (.pt)")
    ap.add_argument("--data_root", required=True,
                    help="AIGVDBench root containing Real/ and OpenSource/ folders")
    ap.add_argument("--num_frames", type=int, default=0,
                    help="0 -> backbone default (16 for VideoMAE)")
    ap.add_argument("--batch_size", type=int, default=32)
    ap.add_argument("--num_workers", type=int, default=8)
    ap.add_argument("--threshold", type=float, default=0.5)
    ap.add_argument("--out_json", default="",
                    help="if set, write per-generator AUC + summary to this JSON")
    args = ap.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    nf = args.num_frames or DEFAULT_NUM_FRAMES[args.backbone]
    print(f"backbone={args.backbone}  num_frames={nf}  bs={args.batch_size}", flush=True)
    print(f"backbone_weights={args.backbone_weights}", flush=True)
    print(f"head_ckpt={args.head_ckpt}", flush=True)

    # ---- backbone (frozen) ----
    bb_cls   = BACKBONE_REGISTRY[args.backbone]
    backbone = bb_cls(pretrained_model=args.backbone_weights,
                      num_frames=nf, img_size=224).to(device).eval()
    for p in backbone.parameters():
        p.requires_grad = False
    D = backbone.embed_dim
    print(f"  P={backbone.num_patches_spatial}  D={D}  "
          f"input_mean={backbone.input_mean}  input_std={backbone.input_std}", flush=True)

    # ---- head ----
    head = PVPHead("velocity_gated", D=D).to(device)
    sd = torch.load(args.head_ckpt, map_location=device)
    if isinstance(sd, dict) and "state_dict" in sd:
        sd = sd["state_dict"]
    head.load_state_dict(sd)
    head.eval()
    n_head = sum(p.numel() for p in head.parameters())
    print(f"  head params: {n_head:,}", flush=True)

    # ---- data ----
    cfg = OpenSourceConfig(data_root=args.data_root, num_frames=nf,
                           crop_size=224, max_per_category=0)
    ds  = OpenSourceTest(cfg)
    print(f"  test N={len(ds)}", flush=True)
    loader = DataLoader(ds, batch_size=args.batch_size, shuffle=False,
                        num_workers=args.num_workers, pin_memory=True,
                        collate_fn=collate)

    # ---- normalization tensors ----
    ds_mean = torch.tensor(DS_MEAN, device=device).view(1, 3, 1, 1, 1)
    ds_std  = torch.tensor(DS_STD,  device=device).view(1, 3, 1, 1, 1)
    bb_mean = torch.tensor(backbone.input_mean, device=device).view(1, 3, 1, 1, 1)
    bb_std  = torch.tensor(backbone.input_std,  device=device).view(1, 3, 1, 1, 1)

    # ---- run ----
    t0 = time.time()
    scores, labels, cats = evaluate(backbone, head, loader, device,
                                    (ds_mean, ds_std), (bb_mean, bb_std))
    elapsed = (time.time() - t0) / 60
    print(f"  inference done in {elapsed:.1f} min  ({len(scores)} videos)", flush=True)

    m = compute_metrics(scores, labels, cats, threshold=args.threshold)
    print(f"\nOpen-Source split results ({m['n_real']} real / {m['n_fake']} fake):")
    print(f"  mga AUC   = {m['mga_auc']:.4f}")
    print(f"  mga ACC   = {m['mga_acc']:.4f}")
    print(f"  overall AUC = {m['overall_auc']:.4f}")
    print(f"  overall ACC = {m['overall_acc']:.4f}")
    print("\nPer-generator AUC:")
    for c, a in m["auc_per_gen"].items():
        print(f"  {c:<32s} {a*100:6.2f}")

    if args.out_json:
        os.makedirs(os.path.dirname(args.out_json) or ".", exist_ok=True)
        out = {
            "backbone":   args.backbone,
            "num_frames": nf,
            "head_ckpt":  args.head_ckpt,
            **m,
        }
        with open(args.out_json, "w") as f:
            json.dump(out, f, indent=2)
        print(f"\nsaved {args.out_json}")


if __name__ == "__main__":
    main()

	"""Evaluate a trained V-PVP head on the AIGVDBench Open-Source split.

	Computes per-generator AUC, mean-of-generators AUC (mga AUC), and overall
	AUC on the test set. Each backbone is loaded once (frozen, fp16 autocast),
	patch features ``z`` of shape ``(B, P, T', D)`` are extracted, the velocity
	field ``v = z[:, :, 1:] - z[:, :, :-1]`` is computed, and the head produces
	a scalar logit per video.

	Usage:
	python -m src.eval_opensource \
	--backbone videomae_k400 \
	--backbone_weights <path-or-hub-id> \
	--head_ckpt checkpoints/pvp_heads/videomae_k400_vpvp.pt \
	--data_root <path-to-aigvdbench> \
	--num_frames 16 \
	--batch_size 32 \
	--out_json results/videomae_open_source.json

	Expected mga AUC on the Open-Source split with the released checkpoint
	(20 generators, 3000 real + 60000 fake videos):

	videomae_k400 -> 0.9528 (V-PVP row of Table 1)
	"""
	import argparse
	import json
	import os
	import sys
	import time

	import numpy as np
	import torch
	import torch.nn as nn
	from sklearn.metrics import roc_auc_score
	from torch.utils.data import DataLoader

	sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

	from backbones import BACKBONE_REGISTRY
	from dataset import OpenSourceConfig, OpenSourceTest, collate
	from pvp_head import PVPHead


	# Frame count the released checkpoint was trained at.
	DEFAULT_NUM_FRAMES = {
	"videomae_k400": 16,
	}

	# Dataset-side normalization applied by ``OpenSourceTest`` (mean 0.45 /
	# std 0.225). The eval loop undoes it and re-applies the backbone's own
	# statistics below; this is the exact preprocessing path the released
	# checkpoint was trained and evaluated with.
	DS_MEAN = (0.45, 0.45, 0.45)
	DS_STD = (0.225, 0.225, 0.225)


	def renormalize(clips, ds_mean, ds_std, bb_mean, bb_std):
	"""Undo dataset normalization, re-apply backbone's. clips: (B, C, T, H, W)."""
	x = clips * ds_std + ds_mean
	return (x - bb_mean) / bb_std


	def compute_metrics(scores, labels, cats, threshold=0.5):
	"""mga AUC = mean-of-generators AUC over the per-(real, generator) pools."""
	cats = np.asarray(cats)
	labels = np.asarray(labels)
	scores = np.asarray(scores)
	real = (labels == 0)
	preds = (scores > threshold).astype(int)

	auc_per_gen, acc_per_gen = {}, {}
	for c in sorted(set(cats[~real].tolist())):
	gen_mask = (cats == c) & (labels == 1)
	if gen_mask.sum() < 2 or real.sum() < 2:
	continue
	y = np.r_[np.zeros(real.sum()), np.ones(gen_mask.sum())]
	s = np.r_[scores[real], scores[gen_mask]]
	p = np.r_[preds[real], preds[gen_mask]]
	auc_per_gen[c] = float(roc_auc_score(y, s))
	acc_per_gen[c] = float((p == y).mean())

	return {
	"auc_per_gen": auc_per_gen,
	"acc_per_gen": acc_per_gen,
	"mga_auc": float(np.mean(list(auc_per_gen.values()))),
	"mga_acc": float(np.mean(list(acc_per_gen.values()))),
	"overall_auc": float(roc_auc_score(labels, scores)),
	"overall_acc": float((preds == labels).mean()),
	"n_real": int(real.sum()),
	"n_fake": int((~real).sum()),
	}


	@torch.no_grad()
	def evaluate(backbone, head, loader, device, ds_norm, bb_norm):
	head.eval()
	backbone.eval()
	scores_all, labels_all, cats_all = [], [], []
	n_done = 0
	t0 = time.time()
	for clips, labels, cats in loader:
	clips = clips.to(device, non_blocking=True)
	x = renormalize(clips, ds_norm, bb_norm)
	with torch.autocast(device_type="cuda", dtype=torch.float16):
	z = backbone.extract_patch_features(x)
	z = z.to(torch.float32)
	if z.size(2) < 2:
	raise RuntimeError(f"backbone returned T'={z.size(2)} < 2; need >=2 frames")
	v = z[:, :, 1:, :] - z[:, :, :-1, :]
	logit = head(z, v)
	scores_all.append(torch.sigmoid(logit).cpu().numpy())
	labels_all.append(labels.numpy())
	cats_all.extend(cats)
	n_done += clips.size(0)
	if n_done % 1024 == 0:
	print(f" [{n_done}] {n_done/(time.time()-t0):.1f} v/s", flush=True)
	return (np.concatenate(scores_all),
	np.concatenate(labels_all),
	np.array(cats_all))


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--backbone", required=True,
	choices=list(BACKBONE_REGISTRY.keys()))
	ap.add_argument("--backbone_weights", required=True,
	help="path (or HuggingFace repo id) of the K400 backbone weights")
	ap.add_argument("--head_ckpt", required=True,
	help="path to the V-PVP head state_dict (.pt)")
	ap.add_argument("--data_root", required=True,
	help="AIGVDBench root containing Real/ and OpenSource/ folders")
	ap.add_argument("--num_frames", type=int, default=0,
	help="0 -> backbone default (16 for VideoMAE)")
	ap.add_argument("--batch_size", type=int, default=32)
	ap.add_argument("--num_workers", type=int, default=8)
	ap.add_argument("--threshold", type=float, default=0.5)
	ap.add_argument("--out_json", default="",
	help="if set, write per-generator AUC + summary to this JSON")
	args = ap.parse_args()

	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	nf = args.num_frames or DEFAULT_NUM_FRAMES[args.backbone]
	print(f"backbone={args.backbone} num_frames={nf} bs={args.batch_size}", flush=True)
	print(f"backbone_weights={args.backbone_weights}", flush=True)
	print(f"head_ckpt={args.head_ckpt}", flush=True)

	# ---- backbone (frozen) ----
	bb_cls = BACKBONE_REGISTRY[args.backbone]
	backbone = bb_cls(pretrained_model=args.backbone_weights,
	num_frames=nf, img_size=224).to(device).eval()
	for p in backbone.parameters():
	p.requires_grad = False
	D = backbone.embed_dim
	print(f" P={backbone.num_patches_spatial} D={D} "
	f"input_mean={backbone.input_mean} input_std={backbone.input_std}", flush=True)

	# ---- head ----
	head = PVPHead("velocity_gated", D=D).to(device)
	sd = torch.load(args.head_ckpt, map_location=device)
	if isinstance(sd, dict) and "state_dict" in sd:
	sd = sd["state_dict"]
	head.load_state_dict(sd)
	head.eval()
	n_head = sum(p.numel() for p in head.parameters())
	print(f" head params: {n_head:,}", flush=True)

	# ---- data ----
	cfg = OpenSourceConfig(data_root=args.data_root, num_frames=nf,
	crop_size=224, max_per_category=0)
	ds = OpenSourceTest(cfg)
	print(f" test N={len(ds)}", flush=True)
	loader = DataLoader(ds, batch_size=args.batch_size, shuffle=False,
	num_workers=args.num_workers, pin_memory=True,
	collate_fn=collate)

	# ---- normalization tensors ----
	ds_mean = torch.tensor(DS_MEAN, device=device).view(1, 3, 1, 1, 1)
	ds_std = torch.tensor(DS_STD, device=device).view(1, 3, 1, 1, 1)
	bb_mean = torch.tensor(backbone.input_mean, device=device).view(1, 3, 1, 1, 1)
	bb_std = torch.tensor(backbone.input_std, device=device).view(1, 3, 1, 1, 1)

	# ---- run ----
	t0 = time.time()
	scores, labels, cats = evaluate(backbone, head, loader, device,
	(ds_mean, ds_std), (bb_mean, bb_std))
	elapsed = (time.time() - t0) / 60
	print(f" inference done in {elapsed:.1f} min ({len(scores)} videos)", flush=True)

	m = compute_metrics(scores, labels, cats, threshold=args.threshold)
	print(f"\nOpen-Source split results ({m['n_real']} real / {m['n_fake']} fake):")
	print(f" mga AUC = {m['mga_auc']:.4f}")
	print(f" mga ACC = {m['mga_acc']:.4f}")
	print(f" overall AUC = {m['overall_auc']:.4f}")
	print(f" overall ACC = {m['overall_acc']:.4f}")
	print("\nPer-generator AUC:")
	for c, a in m["auc_per_gen"].items():
	print(f" {c:<32s} {a*100:6.2f}")

	if args.out_json:
	os.makedirs(os.path.dirname(args.out_json) or ".", exist_ok=True)
	out = {
	"backbone": args.backbone,
	"num_frames": nf,
	"head_ckpt": args.head_ckpt,
	**m,
	}
	with open(args.out_json, "w") as f:
	json.dump(out, f, indent=2)
	print(f"\nsaved {args.out_json}")


	if __name__ == "__main__":
	main()