scripts/eval_test_json.py

#!/usr/bin/env python3
"""Shardable frame-level evaluator for image-data/test.json."""

import argparse
import datetime as _dt
import json
import os
import random
import sys
import time

import cv2
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.utils.data
import yaml
from PIL import Image, ImageFile

ImageFile.LOAD_TRUNCATED_IMAGES = True


def log(fp, message):
    line = "[{}] {}".format(_dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), message)
    print(line, flush=True)
    fp.write(line + "\n")
    fp.flush()


def load_yaml(path):
    with open(path, "r") as f:
        return yaml.safe_load(f)


def merge_config(detector_path, test_config_path):
    detector_config = load_yaml(detector_path)
    test_config = load_yaml(test_config_path)
    config = dict(test_config)
    config.update({k: v for k, v in detector_config.items() if v is not None})
    return config


class TestJsonDataset(torch.utils.data.Dataset):
    def __init__(self, json_path, root, config, shard_index, shard_count, max_samples=None):
        json_path = os.path.abspath(json_path)
        with open(json_path, "r") as f:
            payload = json.load(f)
        items = payload.get("items")
        if not isinstance(items, list):
            raise ValueError("{} must contain an items list".format(json_path))

        if root is None:
            root = payload.get("metadata", {}).get("root")
        if not root:
            raise ValueError("Dataset root is required.")
        if not os.path.isabs(root):
            root = os.path.normpath(os.path.join(os.path.dirname(json_path), root))

        selected = items[shard_index::shard_count]
        if max_samples is not None:
            selected = selected[:max_samples]
        self.items = selected
        self.root = root
        self.resolution = int(config.get("resolution", 224))
        self.mean = np.asarray(config.get("mean", [0.485, 0.456, 0.406]), dtype=np.float32)
        self.std = np.asarray(config.get("std", [0.229, 0.224, 0.225]), dtype=np.float32)
        label_map = config.get("image_data_json", {}).get("label_map", {"real": 0, "fake": 1})
        self.label_map = dict(label_map)
        self.paths = []
        self.labels = []
        self.datasets = []
        self.fine_labels = []
        self.methods = []
        self.invalid_read_count = 0

        for idx, item in enumerate(self.items):
            rel_path = item.get("path")
            if not rel_path:
                raise ValueError("Missing path at selected item {}".format(idx))
            raw_label = item.get("label")
            if raw_label not in self.label_map:
                raise ValueError("Unsupported label {} for {}".format(raw_label, rel_path))
            self.paths.append(rel_path)
            self.labels.append(int(self.label_map[raw_label]))
            self.datasets.append(str(item.get("dataset", "")))
            self.fine_labels.append(str(item.get("fine_label", "")))
            self.methods.append(str(item.get("method", "")))

    def __len__(self):
        return len(self.paths)

    def _abs_path(self, rel_path):
        if os.path.isabs(rel_path):
            return rel_path
        return os.path.join(self.root, rel_path)

    def __getitem__(self, index):
        rel_path = self.paths[index]
        image_path = self._abs_path(rel_path)
        image = cv2.imread(image_path, cv2.IMREAD_COLOR)
        valid_image = True
        if image is None:
            try:
                pil_image = Image.open(image_path).convert("RGB")
                image = np.asarray(pil_image)
            except Exception:
                image = np.zeros((self.resolution, self.resolution, 3), dtype=np.uint8)
                valid_image = False
        else:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = cv2.resize(image, (self.resolution, self.resolution), interpolation=cv2.INTER_CUBIC)
        image = image.astype(np.float32) / 255.0
        image = (image - self.mean.reshape(1, 1, 3)) / self.std.reshape(1, 1, 3)
        image = torch.from_numpy(image.transpose(2, 0, 1))
        return {
            "image": image,
            "label": int(self.labels[index]),
            "path": rel_path,
            "dataset": self.datasets[index],
            "fine_label": self.fine_labels[index],
            "method": self.methods[index],
            "valid_image": valid_image,
        }


def collate_fn(batch):
    return {
        "image": torch.stack([x["image"] for x in batch], dim=0),
        "label": torch.LongTensor([x["label"] for x in batch]),
        "path": [x["path"] for x in batch],
        "dataset": [x["dataset"] for x in batch],
        "fine_label": [x["fine_label"] for x in batch],
        "method": [x["method"] for x in batch],
        "valid_image": [bool(x["valid_image"]) for x in batch],
        "mask": None,
        "landmark": None,
    }


def init_seed(config):
    seed = config.get("manualSeed")
    if seed is None:
        seed = 1024
        config["manualSeed"] = seed
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if config.get("cuda", True) and torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def count_params(model):
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return int(trainable), int(total)


def adapt_moe_gate_modules_for_checkpoint(model, weights):
    """Match MoE gate scorer modules to the checkpoint key format before strict load."""
    try:
        from networks.fusion import _SharedGateMLP  # noqa: E402
    except Exception:
        return []

    device = next(model.parameters()).device
    adapted = []

    def set_shared_gate(branch_name):
        if not hasattr(model, branch_name):
            return
        branch = getattr(model, branch_name)
        mlp_w0 = "{}.shared_gate.net.0.weight".format(branch_name)
        mlp_w2 = "{}.shared_gate.net.2.weight".format(branch_name)
        linear_w = "{}.shared_gate.weight".format(branch_name)
        if mlp_w0 in weights and mlp_w2 in weights:
            hidden_dim, embed_dim = weights[mlp_w0].shape
            out_dim = weights[mlp_w2].shape[0]
            branch.shared_gate = _SharedGateMLP(
                embed_dim=int(embed_dim),
                gate_hidden_dim=int(hidden_dim),
                out_dim=int(out_dim),
            ).to(device)
            adapted.append("{}:shared_gate=MLP({},{},{})".format(
                branch_name, int(embed_dim), int(hidden_dim), int(out_dim)
            ))
        elif linear_w in weights and not isinstance(getattr(branch, "shared_gate", None), torch.nn.Linear):
            out_dim, embed_dim = weights[linear_w].shape
            branch.shared_gate = torch.nn.Linear(int(embed_dim), int(out_dim)).to(device)
            adapted.append("{}:shared_gate=Linear({},{})".format(
                branch_name, int(embed_dim), int(out_dim)
            ))

    def set_layer_attn_scorer(branch_name):
        if not hasattr(model, branch_name):
            return
        branch = getattr(model, branch_name)
        if not hasattr(branch, "layer_attn_scorer"):
            return
        mlp_w0 = "{}.layer_attn_scorer.net.0.weight".format(branch_name)
        mlp_w2 = "{}.layer_attn_scorer.net.2.weight".format(branch_name)
        linear_w = "{}.layer_attn_scorer.weight".format(branch_name)
        if mlp_w0 in weights and mlp_w2 in weights:
            hidden_dim, embed_dim = weights[mlp_w0].shape
            out_dim = weights[mlp_w2].shape[0]
            branch.layer_attn_scorer = _SharedGateMLP(
                embed_dim=int(embed_dim),
                gate_hidden_dim=int(hidden_dim),
                out_dim=int(out_dim),
            ).to(device)
            adapted.append("{}:layer_attn_scorer=MLP({},{},{})".format(
                branch_name, int(embed_dim), int(hidden_dim), int(out_dim)
            ))
        elif linear_w in weights and not isinstance(branch.layer_attn_scorer, torch.nn.Linear):
            out_dim, embed_dim = weights[linear_w].shape
            branch.layer_attn_scorer = torch.nn.Linear(int(embed_dim), int(out_dim)).to(device)
            adapted.append("{}:layer_attn_scorer=Linear({},{})".format(
                branch_name, int(embed_dim), int(out_dim)
            ))

    for branch_name in ("clip_layer_moe", "dino_layer_moe"):
        set_shared_gate(branch_name)
        set_layer_attn_scorer(branch_name)

    return adapted


def load_detector(config, weights_path, device, training_dir, repo_root):
    if training_dir not in sys.path:
        sys.path.insert(0, training_dir)
    if repo_root not in sys.path:
        sys.path.insert(0, repo_root)
    from detectors import DETECTOR  # noqa: E402

    if config.get("model_name") == "effort":
        patch_effort_svd_builder()

    model = DETECTOR[config["model_name"]](config).to(device)
    ckpt = torch.load(weights_path, map_location=device)
    if isinstance(ckpt, dict) and "state_dict" in ckpt:
        ckpt = ckpt["state_dict"]
    weights = {k.replace("module.", ""): v for k, v in ckpt.items()}
    model_keys = set(model.state_dict())
    for deprecated_key in ("clip_vision.embeddings.position_ids",):
        if deprecated_key in weights and deprecated_key not in model_keys:
            weights.pop(deprecated_key)
    adapted = adapt_moe_gate_modules_for_checkpoint(model, weights)
    if adapted:
        print("Adapted MoE gate modules for checkpoint: {}".format(", ".join(adapted)), flush=True)
    model.load_state_dict(weights, strict=True)
    if config.get("model_name") == "effort":
        model._cached_svd_modules = cache_effort_svd_weights(model)
    model.eval()
    return model


def patch_effort_svd_builder():
    """Skip expensive eval-time SVD; checkpoint loading fills these tensors."""
    try:
        from detectors import effort_detector  # noqa: E402
    except Exception:
        return

    if getattr(effort_detector, "_FAST_EVAL_SVD_PATCHED", False):
        return

    def fast_replace_with_svd_residual(module, r):
        if not isinstance(module, torch.nn.Linear):
            return module
        in_features = int(module.in_features)
        out_features = int(module.out_features)
        bias = module.bias is not None
        new_module = effort_detector.SVDResidualLinear(
            in_features,
            out_features,
            r,
            bias=bias,
            init_weight=module.weight.data.clone(),
        )
        if bias and module.bias is not None:
            new_module.bias.data.copy_(module.bias.data)

        rank = min(int(r), min(in_features, out_features))
        residual_dim = max(min(in_features, out_features) - rank, 0)
        if residual_dim > 0:
            new_module.S_residual = torch.nn.Parameter(torch.empty(residual_dim))
            new_module.U_residual = torch.nn.Parameter(torch.empty(out_features, residual_dim))
            new_module.V_residual = torch.nn.Parameter(torch.empty(residual_dim, in_features))
            new_module.S_r = torch.nn.Parameter(torch.empty(rank), requires_grad=False)
            new_module.U_r = torch.nn.Parameter(torch.empty(out_features, rank), requires_grad=False)
            new_module.V_r = torch.nn.Parameter(torch.empty(rank, in_features), requires_grad=False)
        else:
            new_module.S_residual = None
            new_module.U_residual = None
            new_module.V_residual = None
            new_module.S_r = None
            new_module.U_r = None
            new_module.V_r = None
        new_module.weight_original_fnorm = torch.tensor(0.0)
        new_module.weight_main_fnorm = torch.tensor(0.0)
        return new_module

    effort_detector.replace_with_svd_residual = fast_replace_with_svd_residual
    effort_detector._FAST_EVAL_SVD_PATCHED = True


def cache_effort_svd_weights(model):
    """Cache SVDResidualLinear effective weights for eval-only forward speed."""
    try:
        from detectors import effort_detector  # noqa: E402
    except Exception:
        return 0

    if not getattr(effort_detector.SVDResidualLinear, "_FAST_EVAL_FORWARD_PATCHED", False):
        def fast_forward(self, x):
            cached_weight = getattr(self, "_eval_weight", None)
            if cached_weight is not None:
                return torch.nn.functional.linear(x, cached_weight, self.bias)
            if hasattr(self, 'U_residual') and hasattr(self, 'V_residual') and self.S_residual is not None:
                residual_weight = self.U_residual @ torch.diag(self.S_residual) @ self.V_residual
                weight = self.weight_main + residual_weight
            else:
                weight = self.weight_main
            return torch.nn.functional.linear(x, weight, self.bias)

        effort_detector.SVDResidualLinear.forward = fast_forward
        effort_detector.SVDResidualLinear._FAST_EVAL_FORWARD_PATCHED = True

    cached = 0
    with torch.no_grad():
        for module in model.modules():
            if isinstance(module, effort_detector.SVDResidualLinear):
                if hasattr(module, 'U_residual') and hasattr(module, 'V_residual') and module.S_residual is not None:
                    residual_weight = module.U_residual @ torch.diag(module.S_residual) @ module.V_residual
                    weight = module.weight_main + residual_weight
                else:
                    weight = module.weight_main
                module.register_buffer("_eval_weight", weight.detach().clone())
                cached += 1
    return cached


def build_loader(dataset, batch_size, workers):
    pin_memory = os.environ.get("PIN_MEMORY", "1") not in ("0", "false", "False", "no", "No")
    kwargs = {
        "dataset": dataset,
        "batch_size": batch_size,
        "shuffle": False,
        "num_workers": workers,
        "drop_last": False,
        "collate_fn": collate_fn,
    }
    if workers > 0:
        kwargs["pin_memory"] = pin_memory
        kwargs["persistent_workers"] = True
        kwargs["prefetch_factor"] = 4
    return torch.utils.data.DataLoader(**kwargs)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo-root", required=True)
    parser.add_argument("--workdir", required=True)
    parser.add_argument("--detector-path", required=True)
    parser.add_argument("--test-config-path", required=True)
    parser.add_argument("--weights-path", required=True)
    parser.add_argument("--test-json", required=True)
    parser.add_argument("--dataset-root", default=None)
    parser.add_argument("--batch-size", type=int, required=True)
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--shard-index", type=int, required=True)
    parser.add_argument("--shard-count", type=int, required=True)
    parser.add_argument("--raw-output", required=True)
    parser.add_argument("--log-file", required=True)
    parser.add_argument("--max-samples", type=int, default=None)
    args = parser.parse_args()

    if args.dataset_root and not os.path.isabs(args.dataset_root):
        args.dataset_root = os.path.abspath(os.path.join(args.repo_root, args.dataset_root))

    torch.set_num_threads(int(os.environ.get("TORCH_NUM_THREADS", "1")))
    os.makedirs(os.path.dirname(args.raw_output), exist_ok=True)
    os.makedirs(os.path.dirname(args.log_file), exist_ok=True)
    os.chdir(args.workdir)

    with open(args.log_file, "w") as fp:
        log(fp, "repo_root={} workdir={}".format(args.repo_root, args.workdir))
        log(fp, "visible_cuda={} shard={}/{} batch={} workers={}".format(
            os.environ.get("CUDA_VISIBLE_DEVICES"), args.shard_index, args.shard_count,
            args.batch_size, args.workers,
        ))
        config = merge_config(args.detector_path, args.test_config_path)
        config["test_batchSize"] = args.batch_size
        config["workers"] = args.workers
        config["cuda"] = bool(config.get("cuda", True))
        config["lmdb"] = False
        config["image_data_json"] = {
            "enabled": True,
            "root": args.dataset_root,
            "val": args.test_json,
            "label_map": {"real": 0, "fake": 1},
        }
        init_seed(config)
        if config.get("cudnn", True):
            cudnn.benchmark = True

        dataset = TestJsonDataset(
            json_path=args.test_json,
            root=args.dataset_root,
            config=config,
            shard_index=args.shard_index,
            shard_count=args.shard_count,
            max_samples=args.max_samples,
        )
        loader = build_loader(dataset, args.batch_size, args.workers)
        log(fp, "samples={} batches={}".format(len(dataset), len(loader)))

        device = torch.device("cuda" if config["cuda"] and torch.cuda.is_available() else "cpu")
        model = load_detector(
            config=config,
            weights_path=args.weights_path,
            device=device,
            training_dir=os.path.join(args.workdir, "training"),
            repo_root=args.repo_root,
        )
        trainable_params, total_params = count_params(model)
        log(fp, "model_loaded trainable_params={} total_params={}".format(trainable_params, total_params))
        if hasattr(model, "_cached_svd_modules"):
            log(fp, "cached_svd_modules={}".format(model._cached_svd_modules))

        preds = []
        labels = []
        paths = []
        datasets = []
        fine_labels = []
        methods = []
        valid_images = []
        started = time.time()
        with torch.no_grad():
            for batch_index, batch in enumerate(loader, start=1):
                label = torch.where(batch["label"] != 0, 1, 0)
                model_batch = {
                    "image": batch["image"].to(device, non_blocking=True),
                    "label": label.to(device, non_blocking=True),
                    "mask": None,
                    "landmark": None,
                }
                output = model(model_batch, inference=True)
                preds.extend(output["prob"].detach().cpu().numpy().astype(np.float32).tolist())
                labels.extend(label.cpu().numpy().astype(np.int8).tolist())
                paths.extend(batch["path"])
                datasets.extend(batch["dataset"])
                fine_labels.extend(batch["fine_label"])
                methods.extend(batch["method"])
                valid_images.extend(batch["valid_image"])
                if batch_index == 1 or batch_index % 50 == 0 or batch_index == len(loader):
                    elapsed = time.time() - started
                    if device.type == "cuda":
                        used_mb = torch.cuda.max_memory_allocated(0) / (1024 * 1024)
                        log(fp, "progress={}/{} elapsed_sec={:.1f} max_allocated_mb={:.1f}".format(
                            batch_index, len(loader), elapsed, used_mb,
                        ))
                    else:
                        log(fp, "progress={}/{} elapsed_sec={:.1f}".format(batch_index, len(loader), elapsed))

        np.savez_compressed(
            args.raw_output,
            pred=np.asarray(preds, dtype=np.float32),
            label=np.asarray(labels, dtype=np.int8),
            path=np.asarray(paths, dtype=object),
            dataset=np.asarray(datasets, dtype=object),
            fine_label=np.asarray(fine_labels, dtype=object),
            method=np.asarray(methods, dtype=object),
            valid_image=np.asarray(valid_images, dtype=np.bool_),
            shard_index=np.asarray([args.shard_index], dtype=np.int16),
            shard_count=np.asarray([args.shard_count], dtype=np.int16),
            batch_size=np.asarray([args.batch_size], dtype=np.int32),
            trainable_params=np.asarray([trainable_params], dtype=np.int64),
            total_params=np.asarray([total_params], dtype=np.int64),
        )
        invalid_count = int((~np.asarray(valid_images, dtype=np.bool_)).sum())
        if invalid_count:
            log(fp, "invalid_image_placeholders={}".format(invalid_count))
        log(fp, "wrote_raw={} samples={}".format(args.raw_output, len(preds)))
        log(fp, "done elapsed_sec={:.1f}".format(time.time() - started))


if __name__ == "__main__":
    main()