scripts/evaluate_pet_foundation.py

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