namijiang98/radiotherapy-dose-prediction-transfer-learning

Transfer Learning for 3D Radiotherapy Dose Prediction

Model weights and inference code for "End-to-end Automated Radiotherapy Planning Using Transfer Learning to Overcome Data Scarcity" (Jiang et al., Department of Radiation Oncology, UCSF).

Pretrained and transfer-learned 3D dose-prediction models for head-and-neck and pancreas radiotherapy, across three architectures (C3D, MedNeXt-B, SwinUNETR-L).

Each architecture ships three checkpoints so the effect of transfer learning is directly measurable:

Research use only. Not a medical device, not cleared by any regulatory body, and not to be used for treating patients. See LICENSE.

⚠️ Read this first: the models do not all take the same input

There are three different input encodings in this release. All of them are 5-channel or 8-channel volumes of shape (Z, Y, X) = (128, 192, 192), so feeding a model the wrong one runs without error and silently produces wrong dose. Every checkpoint records which encoding it expects in ckpt["channel_spec"], and src/inference.py picks the matching data loader automatically.

src/channels.py is the authoritative spec. Print any of them with:

python src/channels.py

The 8 source-domain channels (*_pretrained)

The pretrained models come from a public challenge dataset that carried beam geometry and plan metadata:

Which channels are unused, and why they were removed rather than zeroed

The models in the paper use the first 5 input channels. The private target cohorts have no beam-geometry or plan-prompt information, so channels 5, 6 and 7 carry no signal at all for this data.

They were dropped, not zero-filled. The released *_finetuned and *_fromscratch checkpoints physically have a 5-channel input stem — you cannot feed them 8 channels, and you should not pad your 5 channels to 8 for them. Channels 0–4 keep the same positions and meanings, which is what makes transfer possible; only the input stem changed shape. See Transfer mechanics below.

The consequence for the *_pretrained checkpoints: they still expect 8 channels, and they were trained with real beam and prompt channels, so they are published as a fine-tuning starting point, not as a runnable predictor on data that lacks those channels. src/inference.py refuses to run them by default. channels.pad_to_source_8ch exists for shape debugging only.

How the two 5-channel encodings differ

Channels 3 (body) and 4 (img) are built identically for HN and pancreas. Channels 0, 1 and 2 are not:

Two design choices worth understanding before you adapt this to your own data:

• —In `hn_5ch`, channels 0 and 1 are byte-identical. Our HN cohort does not distinguish optimisation PTVs from planning PTVs, so the loader builds both from the same StructName masks. The slot is kept separate rather than collapsed so that if your dataset does have distinct optimisation PTV masks, you can populate channel 0 with them and the model will use that information — the pretrained weights were trained with the two channels genuinely differing.

• —In `pancreas_5ch`, channel 1 is capped at `3.0` (30 Gy) while channel 0 is uncapped, so a 40 Gy PTV appears as 3.0 in channel 1 and 4.0 in channel 0. The cap normalises the wide spread of pancreas prescriptions into a range the network sees consistently. If you retrain on your own cohort, set the cap from your own prescription distribution — the highest prescription in your dataset is a reasonable choice. Note that changing it shifts the input distribution away from what the released weights were fine-tuned on, so re-tune rather than expecting the published checkpoint to transfer unchanged.

Available weights

All files are plain state_dict + metadata, safe under torch.load(..., weights_only=True). SHA-256 sums are in `weights/MANIFEST.json`.

Notes:

• —pancreas_c3d_pretrained.pt has byte-identical weights to c3d_pretrained.pt — the pancreas model was fine-tuned directly from the source-domain model, not from the head-and-neck model. It is duplicated so the pancreas set stands alone; only its scale_out metadata differs.
• —The *_pretrained files have marginally more parameters than their 5-channel descendants — that difference is entirely the 8- vs 5-channel input stem.

Output convention

Every model emits a raw tensor that must be rescaled to physical dose:

dose_Gy = sigmoid(output) × scale_out × dose_div_factor

with scale_out = 7.5 for head-and-neck, 5.5 for pancreas, and dose_div_factor = 10 everywhere. checkpoint.to_dose() does this for you using each checkpoint's own constants.

C3D returns a list `[output_A, output_B]` — output_A is the coarse first U-Net, `output_B` is the refined prediction you want. MedNeXt and SwinUNETR return a single tensor.

Quick start

pip install -r requirements.txt
# MedNeXt checkpoints only:
pip install git+https://github.com/MIC-DKFZ/MedNeXt.git

Minimal load-and-predict:

import sys, torch
sys.path.insert(0, "src")
from checkpoint import load_model, to_dose

model, ckpt = load_model("weights/c3d_finetuned.pt", device="cuda")
print(ckpt["channel_spec"], ckpt["channel_names"])
# hn_5ch ['comb_optptv', 'comb_ptv', 'comb_oar', 'body', 'img']

x = torch.randn(1, 5, 128, 192, 192, device="cuda")   # your prepared input
with torch.no_grad():
    dose_gy = to_dose(model(x), ckpt)                 # [1, 1, 128, 192, 192] in Gy

Full inference over a dataset, writing NIfTI volumes:

cp -r meta_files_template/hn meta_files      # then fill in your own cases
# edit configs/c3d_hn.yaml -> loader_params.data_root

python src/inference.py \
    --ckpt   weights/c3d_finetuned.pt \
    --config configs/c3d_hn.yaml \
    --out    predictions/c3d_finetuned

The script derives the architecture, channel encoding, data loader and dose rescaling from the checkpoint, then per case: converts to Gy, clips to clip_factor × PTV_High prescription, zeroes everything outside the BODY mask, copies the CT's spacing/origin/direction, and writes <case_id>_pred.nii.gz.

Inspect a checkpoint's expected channels without running anything:

python src/inference.py --ckpt weights/pancreas_c3d_finetuned.pt \
    --config configs/c3d_pancreas.yaml --out /tmp/x --describe

Transfer mechanics

How the *_finetuned checkpoints were actually produced, so you can reproduce it on your own data.

Every tensor whose shape matched was copied from the pretrained model; the rest kept their fresh initialisation. The input channel count changed.

src/transfer.py shows a demo:

import sys; sys.path.insert(0, "src")
from build_model import build_model
from transfer import load_pretrained_into

model  = build_model("c3d", in_channels=5)      # your own channel count
report = load_pretrained_into(model, "weights/c3d_pretrained.pt")
print(report)      # names every tensor that was NOT transferred -- read it

Repository layout

├── README.md
├── LICENSE                       CC BY-NC 4.0 + no-clinical-use notice
├── requirements.txt
├── configs/                      one YAML per model/anatomy
│   ├── c3d_hn.yaml
│   ├── mednext_hn.yaml
│   ├── swinunetr_hn.yaml
│   └── c3d_pancreas.yaml
├── src/
│   ├── channels.py               THE channel spec -- read this first
│   ├── model_c3d.py              C3D architecture
│   ├── build_model.py            arch name + channel count -> model
│   ├── checkpoint.py             loading weights, rescaling output to Gy
│   ├── transfer.py               warm-starting your own model
│   ├── inference.py              dataset -> NIfTI dose volumes
│   ├── data_loader_hn.py        produces hn_5ch
│   ├── toolkit_hn.py
│   ├── data_loader_pancreas.py   produces pancreas_5ch
│   └── toolkit_pancreas.py
├── meta_files_template/          de-identified metadata schema + data layout
│   ├── hn/  pancreas/  README.md
├── weights/                      12 checkpoints + MANIFEST.json
└── tools/
    └── convert_checkpoints.py    provenance: raw training ckpts -> weights/*.pt

Training code (loss functions, trainer loop, W&B logging, evaluation) is not included — this release is scoped to using the weights. Everything needed for that is here: architectures, both data loaders, inference, and transfer.

What is not published

• —The training data. Both target cohorts are private patient data.
• —The real metadata files. meta_files_template/ gives the schema instead.
• —Few-shot and low-LR checkpoints. The study also produced 5/10/20/50 % few-shot runs (two seeds each) and a low-learning-rate variant. They are not in this release; open an issue if you need them.

Citation

If you use these weights or code, please cite:

@article{jiang_endtoend_rtplanning,
  title   = {End-to-end Automated Radiotherapy Planning Using Transfer Learning
             to Overcome Data Scarcity},
  author  = {Jiang, Lu and Hirata, Emily and Porter, Evan and Xu, Di and Du, Jiayi
             and Yang, Wensha and Lyu, Qihui and Cao, Minsong and Sheng, Ke},
  journal = {Medical Physics},
  year    = {2026},
  note    = {In press; DOI to be assigned}
}

This paper has been provisionally accepted at Medical Physics. The BibTeX entry above will be updated with the final volume, pages and DOI once they are assigned.

Authors: Lu Jiang, Emily Hirata, Evan Porter, Di Xu, Jiayi Du, Wensha Yang, Qihui Lyu, Minsong Cao, Ke Sheng — Department of Radiation Oncology, University of California, San Francisco.

Corresponding author: Ke Sheng, PhD — Professor and Vice Chair of Medical Physics, Department of Radiation Oncology, UCSF — ke.sheng@ucsf.edu

Paper: Medical Physics (in press)  ·  Model repository: https://huggingface.co/<your-username>/radiotherapy-dose-prediction-transfer-learning

Please also cite the work this builds on

@inproceedings{gao2023flexible,
  title     = {Flexible-CM GAN: Towards Precise 3D Dose Prediction in Radiotherapy},
  author    = {Gao, Riqiang and Lou, Bin and Xu, Zhoubing and Comaniciu, Dorin and Kamen, Ali},
  booktitle = {CVPR},
  year      = {2023}
}

@article{liu2021cascade,
  title   = {A cascade 3D U-Net for dose prediction in radiotherapy},
  author  = {Liu, Shuolin and Zhang, Jingjing and Li, Teng and Yan, Hui and Liu, Jianfei},
  journal = {Medical Physics},
  year    = {2021}
}

@inproceedings{roy2023mednext,
  title     = {MedNeXt: Transformer-driven Scaling of ConvNets for Medical Image Segmentation},
  author    = {Roy, Saikat and Koehler, Gregor and Ulrich, Constantin and Baumgartner, Michael
               and Petersen, Jens and Isensee, Fabian and Jaeger, Paul F and Maier-Hein, Klaus},
  booktitle = {MICCAI},
  year      = {2023}
}

@inproceedings{hatamizadeh2022swinunetr,
  title     = {Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images},
  author    = {Hatamizadeh, Ali and Nath, Vishwesh and Tang, Yucheng and Yang, Dong
               and Roth, Holger R and Xu, Daguang},
  booktitle = {MICCAI Brainlesion Workshop},
  year      = {2022}
}

Ethics and funding

This retrospective study was approved by the Institutional Review Board of the University of California, San Francisco (IRB #24-42071); the requirement for informed consent was waived. The private UCSF head-and-neck and pancreas cohorts used for fine-tuning and evaluation are not redistributed here; only model weights and code are.

Supported by NIH R01CA255432, NIH R44CA183390, and NIH R01CA259008.

Disclosures. Dr. Ke Sheng reports grant funding from the National Institutes of Health (NIH). Dr. Minsong Cao reports consulting fees and honoraria from Varian Medical Systems, Siemens Healthineers, and the Medical Dosimetrist Certification Board. The other authors declare no conflicts of interest.

License

CC BY-NC 4.0 — attribution required, non-commercial use only. Research use only; not for clinical use.