| """Extract PET/SUVR embeddings and per-subject predictions for paper figures.""" |
| import sys, os, argparse |
| sys.path.insert(0, os.path.dirname(__file__)) |
| import numpy as np |
| import torch |
| from pathlib import Path |
| from torch.utils.data import DataLoader |
| from pet_vlm_dataset import PETSUVRDataset, collate_pet_suvr |
| from train_pet_foundation import PETSUVRFoundationModel, build_encoder |
|
|
| def make_args(backbone): |
| args = argparse.Namespace() |
| args.backbone = backbone |
| args.medicalnet_weights = Path("pretrained/medicalnet/resnet_50_23dataset.pth") |
| args.brainiac_weights = Path("pretrained/brainiac/backbone.safetensors") |
| args.brainfm_weights = Path("pretrained/brainfm/assets/brainfm_pretrained.pth") |
| args.brainfm_code_root = Path("pretrained/brainfm") |
| args.swinunetr_weights = Path("pretrained/swinunetr/model_swinvit.pt") |
| args.sam_med3d_weights = Path("pretrained/sam-med3d/sam_med3d_turbo.pth") |
| args.output_size = [96, 96, 96] |
| return args |
|
|
| def extract(checkpoint_path, backbone, output_prefix): |
| device = torch.device("cuda:0") |
| args = make_args(backbone) |
| output_size = tuple(args.output_size) |
|
|
| encoder = build_encoder(args) |
| n_regions = 120 |
| model = PETSUVRFoundationModel(encoder, n_regions=n_regions).to(device) |
|
|
| ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False) |
| model.load_state_dict(ckpt["model"], strict=False) |
| model.eval() |
|
|
| ds = PETSUVRDataset(Path("metadata/splits/test.csv"), output_size=output_size) |
| loader = DataLoader(ds, batch_size=4, shuffle=False, num_workers=2, collate_fn=collate_pet_suvr) |
|
|
| pet_zs, suvr_zs, preds, targets = [], [], [], [] |
| with torch.no_grad(): |
| for batch in loader: |
| img = batch["image"].to(device) |
| suvr = batch["suvr"].to(device) |
| outputs = model(img, suvr) |
|
|
| pet_feat = model.pet_encoder(img) |
| 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_zs.append(pet_z.cpu().numpy()) |
| suvr_zs.append(suvr_z.cpu().numpy()) |
| preds.append(outputs["pred_suvr"].cpu().numpy()) |
| targets.append(suvr.cpu().numpy()) |
|
|
| pet_zs = np.concatenate(pet_zs) |
| suvr_zs = np.concatenate(suvr_zs) |
| preds = np.concatenate(preds) |
| targets = np.concatenate(targets) |
|
|
| out_dir = Path("runs/paper_figures") |
| out_dir.mkdir(exist_ok=True) |
| np.savez(out_dir / f"{output_prefix}_embeddings.npz", |
| pet_z=pet_zs, suvr_z=suvr_zs, pred_suvr=preds, true_suvr=targets) |
| print(f"Saved {output_prefix}: pet_z={pet_zs.shape}, suvr_z={suvr_zs.shape}, pred={preds.shape}") |
|
|
| if __name__ == "__main__": |
| os.chdir("/data/Albus/Brain") |
| extract("runs/foundation/medicalnet_layer4_regalign_best.pt", "medicalnet", "remap_pet") |
| extract("runs/foundation/medicalnet_frozen_mlp.pt", "medicalnet", "medicalnet_frozen") |
| print("Done!") |
|
|
