from __future__ import annotations import argparse import csv from pathlib import Path import numpy as np import torch from torch.utils.data import DataLoader from pet_vlm_dataset import PETSUVRDataset, collate_pet_suvr from train_pet_foundation import PETSUVRFoundationModel, build_encoder def _pearson(pred: np.ndarray, target: np.ndarray) -> float: pred_flat = pred.reshape(-1) target_flat = target.reshape(-1) if pred_flat.std() < 1e-8 or target_flat.std() < 1e-8: return float("nan") return float(np.corrcoef(pred_flat, target_flat)[0, 1]) def _rankdata(values: np.ndarray) -> np.ndarray: order = np.argsort(values) ranks = np.empty_like(order, dtype=np.float64) ranks[order] = np.arange(len(values), dtype=np.float64) return ranks def _spearman(pred: np.ndarray, target: np.ndarray) -> float: pred_ranks = np.apply_along_axis(_rankdata, 1, pred) target_ranks = np.apply_along_axis(_rankdata, 1, target) return _pearson(pred_ranks, target_ranks) def _topk_overlap(pred: np.ndarray, target: np.ndarray, k: int, largest: bool) -> float: pred_idx = np.argsort(pred, axis=1) target_idx = np.argsort(target, axis=1) if largest: pred_idx = pred_idx[:, -k:] target_idx = target_idx[:, -k:] else: pred_idx = pred_idx[:, :k] target_idx = target_idx[:, :k] overlaps = [] for p, t in zip(pred_idx, target_idx): overlaps.append(len(set(p.tolist()) & set(t.tolist())) / k) return float(np.mean(overlaps)) def _retrieval_metrics(logits: np.ndarray) -> dict[str, float]: ranks = [] for i in range(logits.shape[0]): order = np.argsort(-logits[i]) rank = int(np.where(order == i)[0][0]) + 1 ranks.append(rank) ranks_np = np.asarray(ranks) return { "recall@1": float(np.mean(ranks_np <= 1)), "recall@5": float(np.mean(ranks_np <= 5)), "recall@10": float(np.mean(ranks_np <= 10)), "mrr": float(np.mean(1.0 / ranks_np)), "median_rank": float(np.median(ranks_np)), } def evaluate(model: PETSUVRFoundationModel, loader: DataLoader, device: torch.device) -> dict[str, float]: model.eval() pred_chunks: list[np.ndarray] = [] target_chunks: list[np.ndarray] = [] pet_z_chunks: list[torch.Tensor] = [] suvr_z_chunks: list[torch.Tensor] = [] with torch.no_grad(): for batch in loader: image = batch["image"].to(device, non_blocking=True) suvr = batch["suvr"].to(device, non_blocking=True) outputs = model(image, suvr) pred_chunks.append(outputs["pred_suvr"].detach().cpu().numpy()) target_chunks.append(suvr.detach().cpu().numpy()) pet_feat = model.pet_encoder(image) pet_z = torch.nn.functional.normalize(model.pet_projector(pet_feat), dim=-1) suvr_z = torch.nn.functional.normalize(model.suvr_encoder(suvr), dim=-1) pet_z_chunks.append(pet_z.cpu()) suvr_z_chunks.append(suvr_z.cpu()) pred = np.concatenate(pred_chunks, axis=0) target = np.concatenate(target_chunks, axis=0) pet_z = torch.cat(pet_z_chunks, dim=0) suvr_z = torch.cat(suvr_z_chunks, dim=0) logits = (pet_z @ suvr_z.T).numpy() diff = pred - target metrics = { "samples": float(target.shape[0]), "mae": float(np.mean(np.abs(diff))), "rmse": float(np.sqrt(np.mean(diff**2))), "pearson": _pearson(pred, target), "spearman": _spearman(pred, target), "top5_high_overlap": _topk_overlap(pred, target, 5, largest=True), "top5_low_overlap": _topk_overlap(pred, target, 5, largest=False), } metrics.update({f"pet_to_suvr_{k}": v for k, v in _retrieval_metrics(logits).items()}) metrics.update({f"suvr_to_pet_{k}": v for k, v in _retrieval_metrics(logits.T).items()}) return metrics def main() -> None: parser = argparse.ArgumentParser(description="Evaluate PET-SUVR foundation checkpoints.") parser.add_argument("--checkpoint", type=Path, required=True) parser.add_argument("--manifest", type=Path, default=Path("metadata/splits/test.csv")) parser.add_argument("--backbone", choices=["small_cnn", "medicalnet", "brainiac", "brainfm", "swinunetr", "sam_med3d"], default=None) parser.add_argument("--medicalnet-weights", type=Path, default=Path("pretrained/medicalnet/resnet_50_23dataset.pth")) parser.add_argument("--brainiac-weights", type=Path, default=Path("pretrained/brainiac/backbone.safetensors")) parser.add_argument("--brainfm-weights", type=Path, default=Path("pretrained/brainfm/assets/brainfm_pretrained.pth")) parser.add_argument("--brainfm-code-root", type=Path, default=Path("pretrained/brainfm")) parser.add_argument("--swinunetr-weights", type=Path, default=Path("pretrained/swinunetr/model_swinvit.pt")) parser.add_argument("--sam-med3d-weights", type=Path, default=Path("pretrained/sam-med3d/sam_med3d_turbo.pth")) parser.add_argument("--batch-size", type=int, default=4) parser.add_argument("--num-workers", type=int, default=2) parser.add_argument("--output-size", type=int, nargs=3, default=None) parser.add_argument("--embed-dim", type=int, default=None) parser.add_argument("--freeze-encoder", action=argparse.BooleanOptionalAction, default=None) parser.add_argument("--csv-out", type=Path, default=None) args = parser.parse_args() ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False) saved_args = ckpt.get("args", {}) for name in ("backbone", "embed_dim", "freeze_encoder"): if getattr(args, name.replace("-", "_"), None) is None and name in saved_args: setattr(args, name, saved_args[name]) if args.output_size is None: args.output_size = tuple(saved_args.get("output_size", (96, 96, 96))) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dataset = PETSUVRDataset(args.manifest, output_size=tuple(args.output_size)) loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, collate_fn=collate_pet_suvr) n_regions = int(dataset[0]["suvr"].numel()) encoder = build_encoder(args) model = PETSUVRFoundationModel(encoder, n_regions, args.embed_dim or 256, bool(args.freeze_encoder)).to(device) model.load_state_dict(ckpt["model"], strict=True) metrics = evaluate(model, loader, device) print(f"checkpoint={args.checkpoint}") print(f"manifest={args.manifest}") for key, value in metrics.items(): print(f"{key}={value:.6f}") if args.csv_out: args.csv_out.parent.mkdir(parents=True, exist_ok=True) write_header = not args.csv_out.exists() with args.csv_out.open("a", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=["checkpoint", "manifest", *metrics.keys()]) if write_header: writer.writeheader() writer.writerow({"checkpoint": str(args.checkpoint), "manifest": str(args.manifest), **metrics}) if __name__ == "__main__": main()