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