""" Run a released checkpoint over a test set and write predicted dose as NIfTI. python src/inference.py \ --ckpt weights/c3d_finetuned.pt \ --config configs/c3d_han.yaml \ --out predictions/c3d_finetuned The architecture, input-channel count and dose rescaling all come from the checkpoint itself, so the same command works for any file under weights/. Post-processing applied to every case (same as during evaluation): 1. raw output -> Gy (sigmoid * scale_out * dose_div_factor) 2. clip to [0, 1.2 x prescription of PTV_High] (1.1x for C3D-HaN, see --clip_factor) 3. zero everything outside the BODY mask 4. copy spacing/origin/direction from the case's CT and write _pred.nii.gz """ import argparse import json import os import sys import time import numpy as np import SimpleITK as sitk import torch import yaml from tqdm.auto import tqdm sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from channels import describe, pad_to_source_8ch # noqa: E402 from checkpoint import load_model, to_dose # noqa: E402 def copy_sitk_imageinfo(reference, image): image.SetSpacing(reference.GetSpacing()) image.SetDirection(reference.GetDirection()) image.SetOrigin(reference.GetOrigin()) return image def get_loader(channel_spec, loader_cfig): """ Pick the data loader that produces the channel encoding this checkpoint wants. The two 5-channel encodings are NOT interchangeable -- channels 0, 1 and 2 are computed differently. Feeding one model the other's inputs runs cleanly and silently produces wrong dose. See channels.py. """ if channel_spec == "pancreas_5ch": import data_loader_pancreas as dl else: import data_loader_han as dl return dl.GetLoader(cfig=loader_cfig) def main(): p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument("--ckpt", required=True, help="a file under weights/") p.add_argument("--config", required=True, help="a file under configs/") p.add_argument("--out", required=True, help="directory for the predicted NIfTIs") p.add_argument("--phase", default="test", choices=["train", "valid", "test"], help="which dev_split of meta_data.csv to run on") p.add_argument("--data_root", default=None, help="overrides loader_params.data_root in the config") p.add_argument("--clip_factor", type=float, default=None, help="clip dose at this multiple of the PTV_High prescription; " "defaults to the value the checkpoint was evaluated with") p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") p.add_argument("--pad_to_8ch", action="store_true", help="zero-pad 5-channel inputs to the 8-channel source layout so a " "*_pretrained checkpoint will run. Shape-debugging only -- the " "padded channels carry real signal during pretraining, so the " "output is not a valid dose prediction.") p.add_argument("--describe", action="store_true", help="print the checkpoint's channel spec and exit") args = p.parse_args() cfig = yaml.safe_load(open(args.config)) loader_cfig = cfig["loader_params"] if args.data_root: loader_cfig["data_root"] = args.data_root device = torch.device(args.device) model, ckpt = load_model(args.ckpt, device=device) print(f"Loaded {ckpt['arch']} / {ckpt['regime']} / {ckpt['anatomy']} " f"({ckpt['in_channels']}-channel input) from {args.ckpt}") print(describe(ckpt["channel_spec"])) if args.describe: return if ckpt["in_channels"] == 8 and not args.pad_to_8ch: p.error( "This is a source-domain *pretrained* checkpoint: it takes 8 input channels, " "but the released data loaders emit 5 (no beam_plate / angle_plate / " "prompt_extend for this data). It is published as a fine-tuning starting " "point, not as a runnable predictor. Use a *_finetuned or *_fromscratch " "checkpoint, or pass --pad_to_8ch if you really want a zero-padded forward " "pass. See README.md.") clip_factor = args.clip_factor if clip_factor is None: clip_factor = ckpt.get("eval_clip_factor", 1.2) loaders = get_loader(ckpt["channel_spec"], loader_cfig) test_loader = {"train": loaders.train_dataloader, "valid": loaders.val_dataloader, "test": loaders.test_dataloader}[args.phase]() ptv_dict = json.load(open(loader_cfig["scale_dose_dict"])) os.makedirs(args.out, exist_ok=True) split_dir = "test" if args.phase == "test" else "train_valid" t0 = time.time() n_written = 0 with torch.no_grad(): for data_dict in tqdm(test_loader, desc="Running inference"): inputs = data_dict["data"].to(device) if args.pad_to_8ch: inputs = pad_to_source_8ch(inputs) dose = to_dose(model(inputs), ckpt) for i in range(len(dose)): pred = dose[i, 0].cpu().numpy() case_id = data_dict["id"][i] key = case_id.split("+")[0] entry = ptv_dict.get(key) or ptv_dict.get(key.zfill(3)) if entry and "PTV_High" in entry: pred = np.clip(pred, 0, entry["PTV_High"]["PDose"] * clip_factor) else: print(f"[Warning] no PTV_High prescription for {case_id}; clipping at 80 Gy") pred = np.clip(pred, 0, 80.0) case_dir = os.path.join(loader_cfig["data_root"], split_dir, case_id) ct_path = os.path.join(case_dir, "CT.nii.gz") body_path = os.path.join(case_dir, "Masks", "BODY.nii.gz") if not (os.path.exists(ct_path) and os.path.exists(body_path)): print(f"[Skip] missing CT or BODY for {case_id}") continue pred = pred * (sitk.GetArrayFromImage(sitk.ReadImage(body_path)) > 0) pred_img = sitk.GetImageFromArray(pred.astype(np.float32)) pred_img = copy_sitk_imageinfo(sitk.ReadImage(ct_path), pred_img) sitk.WriteImage(pred_img, os.path.join(args.out, f"{case_id}_pred.nii.gz")) n_written += 1 elapsed = time.time() - t0 print(f"Wrote {n_written} predictions to {args.out}") if n_written: print(f"Avg time/scan: {elapsed / n_written:.2f} s") if device.type == "cuda": print(f"Peak GPU memory: {torch.cuda.max_memory_allocated(device) / 1024 ** 3:.2f} GB") if __name__ == "__main__": main()