| """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() |
|
|
| |
| pet_path = str(row["pet_path"]) |
| pet_vol = nib.load(pet_path).get_fdata(dtype=np.float32) |
| |
| d, h, w = pet_vol.shape |
| axial_slice = pet_vol[d // 2, :, :] |
| coronal_slice = pet_vol[:, h // 2, :] |
| sagittal_slice = pet_vol[:, :, w // 2] |
|
|
| |
| 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() |
|
|
