File size: 2,996 Bytes
e51b27b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """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!")
|