File size: 3,850 Bytes
212e9d7 | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """Export one test subject's data for Figure 4 case study."""
from __future__ import annotations
import argparse
from pathlib import Path
import nibabel as nib
import numpy as np
import pandas as pd
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 main():
parser = argparse.ArgumentParser()
parser.add_argument("--checkpoint", type=Path, required=True)
parser.add_argument("--split", type=Path, default=Path("data/metadata/splits/test.csv"))
parser.add_argument("--subject-index", type=int, default=0, help="which test subject to use (0=first)")
parser.add_argument("--backbone", default=None)
parser.add_argument("--medicalnet-weights", type=Path, default=Path("pretrained/medicalnet/resnet_50_23dataset.pth"))
parser.add_argument("--batch-size", type=int, default=1)
parser.add_argument("--num-workers", type=int, default=0)
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", type=bool, default=None)
parser.add_argument("--out", type=Path, required=True)
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, 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")
manifest = pd.read_csv(args.split)
row = manifest.iloc[args.subject_index]
dataset = PETSUVRDataset(args.split, output_size=tuple(args.output_size))
sample = dataset[args.subject_index]
n_regions = int(sample["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)
model.eval()
image = sample["image"].unsqueeze(0).to(device)
suvr = sample["suvr"].unsqueeze(0).to(device)
with torch.no_grad():
outputs = model(image, suvr)
true_suvr_arr = suvr.cpu().numpy().flatten()
pred_suvr_arr = outputs["pred_suvr"].cpu().numpy().flatten()
# Load PET volume for slice extraction
pet_path = str(row["pet_path"])
pet_vol = nib.load(pet_path).get_fdata(dtype=np.float32)
# Extract middle slices in each orientation
d, h, w = pet_vol.shape
axial_slice = pet_vol[d // 2, :, :]
coronal_slice = pet_vol[:, h // 2, :]
sagittal_slice = pet_vol[:, :, w // 2]
# Load region labels from SUVR CSV
suvr_csv_path = str(row["suvr_csv_path"])
suvr_df = pd.read_csv(suvr_csv_path)
region_labels = [str(lbl) for lbl in suvr_df["label_name"].tolist() if str(lbl) != "Background"]
args.out.parent.mkdir(parents=True, exist_ok=True)
np.savez(args.out,
subject_id=str(row["sample_id"]),
true_suvr=true_suvr_arr,
pred_suvr=pred_suvr_arr,
region_labels=np.array(region_labels, dtype=object),
axial_slice=axial_slice,
coronal_slice=coronal_slice,
sagittal_slice=sagittal_slice)
print(f"wrote case study for subject {row['sample_id']} to {args.out}")
print(f" true SUVR range: [{true_suvr_arr.min():.3f}, {true_suvr_arr.max():.3f}]")
print(f" pred SUVR range: [{pred_suvr_arr.min():.3f}, {pred_suvr_arr.max():.3f}]")
print(f" PET shape: {pet_vol.shape}, slices extracted")
if __name__ == "__main__":
main()
|