| """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 = [ |
| ("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)) |
|
|
| |
| del model |
| torch.cuda.empty_cache() |
| print(f" elapsed: {time.time()-t0:.1f}s", flush=True) |
|
|
| |
| 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() |
|
|
