"""Bootstrap 95% CI for ALL frozen baseline encoders (Stage-1 only: MAE + R@1). Runs 1000-resample bootstrap on the 153-subject test set for each baseline. Usage (from /data/Albus/Brain): CUDA_VISIBLE_DEVICES=2 python scripts/bootstrap_all_baselines.py """ from __future__ import annotations import sys import time from pathlib import Path import numpy as np import torch import torch.nn.functional as F from torch.utils.data import DataLoader sys.path.insert(0, str(Path(__file__).resolve().parent)) from pet_vlm_dataset import PETSUVRDataset, collate_pet_suvr from train_pet_foundation import PETSUVRFoundationModel, build_encoder # -- baselines --------------------------------------------------------------- BASELINES = [ ("MedicalNet frozen", "runs/foundation/medicalnet_frozen_mlp.pt"), ("BrainIAC frozen", "runs/foundation/brainiac_frozen_mlp.pt"), ("BrainFM frozen", "runs/foundation/brainfm_frozen_mlp_b4_best.pt"), ("SAM-Med3D frozen", "runs/foundation/sam_med3d_frozen_mlp_best.pt"), ("SwinUNETR frozen", "runs/foundation/swinunetr_frozen_mlp_best.pt"), ] TEST_MANIFEST = Path("metadata/splits/test.csv") B = 1000 SEED = 42 BATCH_SIZE = 4 def _retrieval_recall_at_1(logits: np.ndarray) -> 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) return float(np.mean(np.asarray(ranks) <= 1)) @torch.no_grad() def collect_stage1(model, loader, device): model.eval() pred_c, tgt_c, pz_c, sz_c = [], [], [], [] 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_c.append(outputs["pred_suvr"].cpu().numpy()) tgt_c.append(suvr.cpu().numpy()) pet_feat = model.pet_encoder(image) pet_z = F.normalize(model.pet_projector(pet_feat), dim=-1) suvr_z = F.normalize(model.suvr_encoder(suvr), dim=-1) pz_c.append(pet_z.cpu().numpy()) sz_c.append(suvr_z.cpu().numpy()) return { "pred": np.concatenate(pred_c), "target": np.concatenate(tgt_c), "pet_z": np.concatenate(pz_c), "suvr_z": np.concatenate(sz_c), } def stage1_metrics(d, idx): pred = d["pred"][idx] target = d["target"][idx] uid = np.unique(idx) logits = d["pet_z"][uid] @ d["suvr_z"][uid].T return { "mae": float(np.mean(np.abs(pred - target))), "pet_suvr_r1": _retrieval_recall_at_1(logits), } def bootstrap_ci(metric_fn, n, B=1000, seed=42): rng = np.random.RandomState(seed) all_idx = np.arange(n) point = metric_fn(all_idx) boots = np.empty(B) for b in range(B): idx = rng.choice(n, size=n, replace=True) boots[b] = metric_fn(idx) lo = float(np.percentile(boots, 2.5)) hi = float(np.percentile(boots, 97.5)) return point, lo, hi def load_model(ckpt_path, device): ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) saved = ckpt.get("args", {}) class _A: pass a = _A() a.backbone = saved.get("backbone", "medicalnet") a.medicalnet_weights = Path(saved.get("medicalnet_weights", "pretrained/medicalnet/resnet_50_23dataset.pth")) a.brainiac_weights = Path(saved.get("brainiac_weights", "pretrained/brainiac/backbone.safetensors")) a.brainfm_weights = Path(saved.get("brainfm_weights", "pretrained/brainfm/assets/brainfm_pretrained.pth")) a.brainfm_code_root = Path(saved.get("brainfm_code_root", "pretrained/brainfm")) a.swinunetr_weights = Path(saved.get("swinunetr_weights", "pretrained/swinunetr/model_swinvit.pt")) a.sam_med3d_weights = Path(saved.get("sam_med3d_weights", "pretrained/sam-med3d/sam_med3d_turbo.pth")) a.output_size = tuple(saved.get("output_size", (96, 96, 96))) embed_dim = saved.get("embed_dim", 256) freeze = bool(saved.get("freeze_encoder", False)) ds_tmp = PETSUVRDataset(TEST_MANIFEST, output_size=a.output_size) n_regions = int(ds_tmp[0]["suvr"].numel()) encoder = build_encoder(a) model = PETSUVRFoundationModel(encoder, n_regions, embed_dim, freeze).to(device) model.load_state_dict(ckpt["model"], strict=True) model.eval() return model, a.output_size def main(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}", flush=True) results = [] for name, ckpt_path in BASELINES: t0 = time.time() print(f"\n{'='*60}", flush=True) print(f" {name} ({ckpt_path})", flush=True) print(f"{'='*60}", flush=True) model, output_size = load_model(ckpt_path, device) bs = 2 if "sam_med3d" in ckpt_path else BATCH_SIZE ds = PETSUVRDataset(TEST_MANIFEST, output_size=output_size) loader = DataLoader(ds, batch_size=bs, shuffle=False, num_workers=2, collate_fn=collate_pet_suvr) d = collect_stage1(model, loader, device) N = d["pred"].shape[0] print(f" N = {N}", flush=True) for metric_name in ("mae", "pet_suvr_r1"): fn = lambda idx, _m=metric_name: stage1_metrics(d, idx)[_m] pt, lo, hi = bootstrap_ci(fn, N, B=B, seed=SEED) print(f" {metric_name:20s} {pt:.4f} 95% CI [{lo:.4f}, {hi:.4f}]", flush=True) results.append((name, metric_name, pt, lo, hi)) # free GPU memory del model torch.cuda.empty_cache() print(f" elapsed: {time.time()-t0:.1f}s", flush=True) # ---- summary table ---- print(f"\n\n{'='*70}", flush=True) print(f"SUMMARY: Bootstrap 95% CI (B={B}, seed={SEED})", flush=True) print(f"{'='*70}", flush=True) print(f"{'Model':<22s} {'MAE':>8s} {'MAE 95% CI':>18s} {'R@1':>8s} {'R@1 95% CI':>18s}", flush=True) print("-"*70, flush=True) for i in range(0, len(results), 2): nm = results[i][0] mae_pt, mae_lo, mae_hi = results[i][2], results[i][3], results[i][4] r1_pt, r1_lo, r1_hi = results[i+1][2], results[i+1][3], results[i+1][4] print(f"{nm:<22s} {mae_pt:8.4f} [{mae_lo:.4f}, {mae_hi:.4f}] {r1_pt:8.4f} [{r1_lo:.4f}, {r1_hi:.4f}]", flush=True) print(f"{'='*70}", flush=True) if __name__ == "__main__": main()