import json import math from collections import defaultdict from pathlib import Path import numpy as np import torch from torch.utils.data import ConcatDataset, Subset from tqdm import tqdm from sklearn.metrics import average_precision_score, accuracy_score, roc_auc_score from data import create_dataloader from util import ConceptSparsityRecorder def _compute_binary_metrics(y_true: np.ndarray, y_pred: np.ndarray, y_prob: np.ndarray, reverse: bool) -> dict: if reverse: real_mask = y_true == 1 fake_mask = y_true == 0 else: real_mask = y_true == 0 fake_mask = y_true == 1 try: auc = roc_auc_score(y_true, y_prob) if y_prob.size > 0 else 0.5 except ValueError: # roc_auc_score is undefined when only one class is present in y_true. auc = 0.5 metrics = { "acc": accuracy_score(y_true, y_pred), "auc": float(auc), "ap": average_precision_score(y_true, y_prob) if y_prob.size > 0 else 0.0, "r_acc": accuracy_score(y_true[real_mask], y_pred[real_mask]) if real_mask.any() else 0.0, "f_acc": accuracy_score(y_true[fake_mask], y_pred[fake_mask]) if fake_mask.any() else 0.0, "y_pred": y_pred, "y_prob": y_prob, } return metrics def _format_ratio_key(ratio: float) -> str: ratio_str = f"{ratio:.6f}".rstrip("0").rstrip(".") return ratio_str if ratio_str else "0" def _extract_dataset_paths(dataset): """Return a flat list of file paths for ImageFolder-style datasets, or None.""" if isinstance(dataset, Subset): base_paths = _extract_dataset_paths(dataset.dataset) if base_paths is None: return None return [base_paths[i] for i in dataset.indices] if isinstance(dataset, ConcatDataset): paths = [] for sub in dataset.datasets: sub_paths = _extract_dataset_paths(sub) if sub_paths is None: return None paths.extend(sub_paths) return paths if hasattr(dataset, "samples"): return [path for path, _ in dataset.samples] if hasattr(dataset, "imgs"): return [path for path, _ in dataset.imgs] return None def _get_patch_size_from_model(model) -> int: """Best-effort extraction of patch size from CLIP/ViT style models.""" visual = getattr(model, "visual", None) if visual is None and hasattr(model, "model"): visual = getattr(model.model, "visual", None) if visual is None: return 0 patch_size = getattr(visual, "patch_size", None) if isinstance(patch_size, (tuple, list)): patch_size = patch_size[0] if patch_size: return int(patch_size) conv1 = getattr(visual, "conv1", None) if conv1 is not None and hasattr(conv1, "kernel_size"): ks = conv1.kernel_size return int(ks[0]) if isinstance(ks, (tuple, list)) else int(ks) return 0 def _infer_grid_size(num_patches: int, height: int, width: int, patch_size: int): """Infer (grid_h, grid_w) for patch tokens.""" if patch_size and height % patch_size == 0 and width % patch_size == 0: grid_h = height // patch_size grid_w = width // patch_size if grid_h * grid_w == num_patches: return grid_h, grid_w if num_patches: side = int(math.sqrt(num_patches)) if side * side == num_patches: return side, side if patch_size and num_patches and width: grid_w = max(1, int(round(width / patch_size))) grid_h = max(1, int(math.ceil(num_patches / grid_w))) return grid_h, grid_w return 0, 0 def validate(model, opt): print( f"[validate] dataroot={getattr(opt, 'dataroot', None)} " f"data_aug={getattr(opt, 'data_aug', None)} " f"isTrain={getattr(opt, 'isTrain', None)}" ) data_loader = create_dataloader(opt) # Keep evaluation on the same device as the model to avoid cross-GPU tensor mixups. try: model_device = next(model.parameters()).device except StopIteration: model_device = torch.device("cpu") num_classes = getattr(opt, "num_classes", 1) reverse = getattr(opt, "reverse", False) try: max_eval_samples = int(getattr(opt, "max_eval_samples", 0) or 0) except (TypeError, ValueError): max_eval_samples = 0 eval_ratios = ( getattr(opt, "clip_concept_eval_sparsity", None) or getattr(opt, "dinov3_concept_eval_sparsity", None) or [] ) sparsity_recorder = ConceptSparsityRecorder( getattr(opt, "clip_concept_eval_log_path", None) or getattr(opt, "dinov3_concept_eval_log_path", None), dataset=getattr(opt, "dataroot", None), extra_meta={ "mode": "validate", "experiment": getattr(opt, "name", None), }, ) coords_path_raw = getattr(opt, "codebook_coords_path", None) dataset_name = getattr(opt, "dataset_name", None) or Path(getattr(opt, "dataroot", "dataset")).name or "dataset" if coords_path_raw and "{dataset}" in str(coords_path_raw): coords_path_raw = str(coords_path_raw).replace("{dataset}", dataset_name) coords_limit = getattr(opt, "codebook_coords_limit", None) try: coords_limit = int(coords_limit) if coords_limit is not None else None except (TypeError, ValueError): coords_limit = None coords_path = Path(coords_path_raw).expanduser() if coords_path_raw else None coords_mode = None coords_base_dir = None if coords_path: suffix = coords_path.suffix.lower() if suffix in {".dat", ".jsonl"}: coords_mode = "dat" if suffix == ".dat" else "jsonl" else: coords_base_dir = coords_path coords_mode = "dat" coords_fp = None coords_bin = None indices_bin = None scores_bin = None concepts_bin = None paths_fp = None sample_paths = None dumped_labels = [] dumped_samples = 0 meta_template = {} if coords_path: if coords_base_dir is not None: coords_path = coords_base_dir / dataset_name / ( "top_patch_coords.dat" if coords_mode == "dat" else "coords.jsonl" ) coords_path.parent.mkdir(parents=True, exist_ok=True) sample_paths = _extract_dataset_paths(data_loader.dataset) if coords_mode == "jsonl": coords_fp = coords_path.open("w", encoding="utf-8") print(f"[validate] dumping codebook coords to {coords_path}") else: coords_bin = coords_path.open("wb") indices_bin = coords_path.with_name("top_patch_indices.dat").open("wb") scores_bin = coords_path.with_name("top_patch_scores.dat").open("wb") concepts_bin = coords_path.with_name("top_patch_concept_ids.dat").open("wb") paths_fp = coords_path.with_name("paths.txt").open("w", encoding="utf-8") print(f"[validate] dumping codebook coords to {coords_path} (dat format)") sample_offset = 0 with torch.no_grad(): y_true: list = [] head_probs: defaultdict = defaultdict(list) head_preds: defaultdict = defaultdict(list) for img, label in tqdm(data_loader, desc="Testing", leave=False): if max_eval_samples > 0 and sample_offset >= max_eval_samples: break if max_eval_samples > 0: remain = max_eval_samples - sample_offset if remain <= 0: break if img.shape[0] > remain: img = img[:remain] label = label[:remain] in_tens = img.to(model_device, non_blocking=True) if coords_path is not None: try: raw_output = model(in_tens, return_debug=True) except TypeError: raw_output = model(in_tens) else: raw_output = model(in_tens) head_outputs = {} concept_raw = None if isinstance(raw_output, dict): if raw_output.get("logits") is not None: head_outputs["main"] = raw_output["logits"] if raw_output.get("cb_only") is not None: head_outputs["cb_only"] = raw_output["cb_only"] if raw_output.get("concept_logits") is not None: head_outputs["concept"] = raw_output["concept_logits"] if raw_output.get("token_logits") is not None: head_outputs["token"] = raw_output["token_logits"] concept_raw = raw_output.get("concept_features_raw") elif isinstance(raw_output, tuple): head_outputs["main"] = raw_output[0] else: head_outputs["main"] = raw_output if concept_raw is not None and eval_ratios: for ratio in eval_ratios: try: ratio_val = float(ratio) except (TypeError, ValueError): continue key = f"concept_s_{_format_ratio_key(ratio_val)}" if sparsity_recorder.enabled(): logits_variant, keep_mask = model.compute_concept_logits_with_ratio( concept_raw, ratio_val, return_mask=True ) sparsity_recorder.update(key, ratio_val, keep_mask) else: logits_variant = model.compute_concept_logits_with_ratio(concept_raw, ratio_val) head_outputs[key] = logits_variant labels_cpu = label.flatten().cpu().tolist() y_true.extend(labels_cpu) for head_name, logits in head_outputs.items(): logits_cpu = logits.detach().cpu() if logits_cpu.ndim == 1: logits_cpu = logits_cpu.unsqueeze(1) if num_classes == 1 or logits_cpu.shape[1] == 1: probs_fake = torch.sigmoid(logits_cpu).flatten() preds = (probs_fake >= 0.5).long() head_probs[head_name].extend(probs_fake.tolist()) head_preds[head_name].extend(preds.tolist()) else: probs = torch.softmax(logits_cpu, dim=1) probs_fake = probs[:, 1] preds = torch.argmax(probs, dim=1) head_probs[head_name].extend(probs_fake.tolist()) head_preds[head_name].extend(preds.tolist()) if coords_path is not None and isinstance(raw_output, dict): topk_idx = raw_output.get("codebook_topk_idx") if topk_idx is not None: num_patches = None tokens = raw_output.get("tokens") if tokens is not None and hasattr(tokens, "shape"): num_patches = int(tokens.shape[1]) if "embedding_dim" not in meta_template: meta_template["embedding_dim"] = int(tokens.shape[-1]) if num_patches is None: num_patches = int(getattr(model, "num_patches", 0) or 0) if not num_patches: num_patches = int(getattr(opt, "clip_num_patches", 0) or 0) if not num_patches: num_patches = int(topk_idx.max().item() + 1) height, width = int(in_tens.shape[-2]), int(in_tens.shape[-1]) patch_size = _get_patch_size_from_model(model) grid_h, grid_w = _infer_grid_size(num_patches, height, width, patch_size) if coords_mode == "jsonl": topk_idx_cpu = topk_idx.detach().cpu().tolist() for b, idx_list in enumerate(topk_idx_cpu): global_idx = sample_offset + b if coords_limit is not None and global_idx >= coords_limit: continue path = None if sample_paths is not None and global_idx < len(sample_paths): path = sample_paths[global_idx] coords_rc = None coords_xyxy = None if grid_w and grid_h: coords_rc = [] coords_xyxy = [] for idx in idx_list: row = int(idx) // grid_w col = int(idx) % grid_w coords_rc.append([row, col]) if patch_size: x0 = col * patch_size y0 = row * patch_size coords_xyxy.append([x0, y0, x0 + patch_size, y0 + patch_size]) entry = { "index": int(global_idx), "path": path, "label": int(labels_cpu[b]) if b < len(labels_cpu) else None, "topk_idx": [int(i) for i in idx_list], "coords_rc": coords_rc, "coords_xyxy": coords_xyxy if coords_xyxy else None, "grid_size": [int(grid_h), int(grid_w)] if grid_h and grid_w else None, "patch_size": int(patch_size) if patch_size else None, "image_size": [height, width], } coords_fp.write(json.dumps(entry, ensure_ascii=True) + "\n") else: idx_np = topk_idx.detach().cpu().to(torch.int32).numpy() if idx_np.ndim == 1: idx_np = idx_np.reshape(-1, 1) top_k = idx_np.shape[1] scores = raw_output.get("codebook_score") topk_concept = raw_output.get("codebook_topk_concept") if scores is not None: score_np = scores.detach().cpu().numpy() score_sel = np.take_along_axis(score_np, idx_np.astype(np.int64), axis=1).astype(np.float32) else: score_sel = np.zeros_like(idx_np, dtype=np.float32) if topk_concept is not None: concept_sel = topk_concept.detach().cpu().to(torch.int32).numpy() else: concept_sel = np.zeros_like(idx_np, dtype=np.int32) if grid_w and grid_h: row = idx_np // grid_w col = idx_np % grid_w coords_packed = (row.astype(np.int32) << 16) | col.astype(np.int32) else: coords_packed = idx_np.astype(np.int32) for b in range(idx_np.shape[0]): global_idx = sample_offset + b if coords_limit is not None and global_idx >= coords_limit: continue coords_packed[b].tofile(coords_bin) idx_np[b].tofile(indices_bin) score_sel[b].tofile(scores_bin) if concepts_bin is not None: concept_sel[b].reshape(-1).tofile(concepts_bin) if paths_fp is not None and sample_paths is not None and global_idx < len(sample_paths): paths_fp.write(f"{sample_paths[global_idx]}\n") dumped_labels.append(int(labels_cpu[b]) if b < len(labels_cpu) else 0) dumped_samples += 1 if "num_patches" not in meta_template: meta_template["num_patches"] = int(num_patches) if "top_k" not in meta_template: meta_template["top_k"] = int(top_k) if "top_r" not in meta_template and concept_sel is not None and concept_sel.ndim == 3: meta_template["top_r"] = int(concept_sel.shape[-1]) if "grid_h" not in meta_template and grid_h and grid_w: meta_template["grid_h"] = int(grid_h) meta_template["grid_w"] = int(grid_w) if "patch_size" not in meta_template and patch_size: meta_template["patch_size"] = int(patch_size) sample_offset += len(labels_cpu) y_true_arr = np.array(y_true) metrics = {} for head_name, preds in head_preds.items(): y_pred_arr = np.array(preds) y_prob_arr = np.array(head_probs[head_name]) metrics[head_name] = _compute_binary_metrics(y_true_arr, y_pred_arr, y_prob_arr, reverse) sparsity_recorder.finalize(total_samples=len(y_true)) if coords_fp is not None: coords_fp.close() if coords_mode == "dat": if coords_bin is not None: coords_bin.close() if indices_bin is not None: indices_bin.close() if scores_bin is not None: scores_bin.close() if concepts_bin is not None: concepts_bin.close() paths_path = coords_path.with_name("paths.txt") labels_path = coords_path.with_name("labels.npy") meta_path = coords_path.with_name("meta.json") if paths_fp is not None: paths_fp.close() if dumped_labels: np.save(labels_path, np.array(dumped_labels, dtype=np.int64)) patch_grid = None if meta_template.get("grid_h") and meta_template.get("grid_w"): patch_grid = [meta_template["grid_h"], meta_template["grid_w"]] patch_size = meta_template.get("patch_size") meta = { "top_patch_tokens.dat": None, "top_patch_scores.dat": "top_patch_scores.dat", "cls_tokens.dat": None, "top_patch_indices.dat": "top_patch_indices.dat", "top_patch_coords.dat": coords_path.name, "top_patch_concept_ids.dat": "top_patch_concept_ids.dat", "num_samples": dumped_samples, "num_patches": meta_template.get("num_patches"), "top_k": meta_template.get("top_k"), "top_r": meta_template.get("top_r"), "embedding_dim": meta_template.get("embedding_dim"), "projection": "post", "dtype": "float32", "labels.npy": labels_path.name, "paths.txt": paths_path.name, "patch_grid": patch_grid, "patch_size": [patch_size, patch_size] if patch_size else None, } with meta_path.open("w", encoding="utf-8") as fp: json.dump(meta, fp, indent=2) return metrics, y_true_arr