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:

checkpoint	what it is
*_pretrained	trained on a large public dose dataset (the source domain)
*_finetuned	that model fine-tuned on the private target cohort
*_fromscratch	same architecture and schedule, random init — the no-transfer baseline

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.

encoding	channels	used by
source_8ch	8	all three *_pretrained checkpoints
hn_5ch	5	c3d/mednext/swinunetr_finetuned and _fromscratch
pancreas_5ch	5	pancreas_c3d_finetuned and _fromscratch

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:

ch	name	content
0	comb_optptv	optimisation PTVs (OPTName structures) × prescription/10, merged by voxel-wise max
1	comb_ptv	planning PTVs (StructName structures) × prescription/10, merged by max
2	comb_oar	all OARs in one label-encoded channel, each mask scaled by 1 + 4·i/30
3	body	binary external / BODY mask
4	img	CT clipped to [-1000, 1000] HU, divided by 500 → [-2, 2]
5	beam_plate	ray-traced beam paths from source through the PTV surface, inverse-square weighted, summed over gantry angles
6	angle_plate	2D gantry-angle map replicated over the ±5 slices around the isocentre in Z, zero elsewhere
7	prompt_extend	[isVMAT, n_PTVs, site, cohort] broadcast into four equal blocks along Z

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:

ch	hn_5ch	pancreas_5ch
0 comb_optptv	PTV × Rx/10	PTV with every overlapping OAR subtracted (ptv · (1 − any_oar)), × Rx/10
1 comb_ptv	PTV × Rx/10	PTV × min(Rx/10, 3.0) — the value is capped at 30 Gy
2 comb_oar	23 HN OARs, each scaled 1 + 4·i/30 → [1.13, 4.07]	14 PancVMAT OARs in two tiers: dose-limiting OARs (Stomach, Bowel_Small, Bowel_Large + PRVs) → [2.0, 2.5]; all others → [0.5, 0.8]

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.

file	arch	input encoding	params	size
c3d_pretrained.pt	C3D	source_8ch	32.4 M	129 MB
c3d_finetuned.pt	C3D	hn_5ch	32.3 M	129 MB
c3d_fromscratch.pt	C3D	hn_5ch	32.3 M	129 MB
mednext_pretrained.pt	MedNeXt-B	source_8ch	10.5 M	42 MB
mednext_finetuned.pt	MedNeXt-B	hn_5ch	10.5 M	42 MB
mednext_fromscratch.pt	MedNeXt-B	hn_5ch	10.5 M	42 MB
swinunetr_pretrained.pt	SwinUNETR-L	source_8ch	291.4 M	1169 MB
swinunetr_finetuned.pt	SwinUNETR-L	hn_5ch	291.3 M	1169 MB
swinunetr_fromscratch.pt	SwinUNETR-L	hn_5ch	291.3 M	1169 MB
pancreas_c3d_pretrained.pt	C3D	source_8ch	32.4 M	129 MB
pancreas_c3d_finetuned.pt	C3D	pancreas_5ch	32.3 M	129 MB
pancreas_c3d_fromscratch.pt	C3D	pancreas_5ch	32.3 M	129 MB

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.

architecture	transferred	re-initialised
C3D	170 / 172	net_A.encoder.encoder_1.0.single_conv.0.weight (16,8,3,3,3 → 16,5,3,3,3), net_B... (32,24 → 32,21)
MedNeXt-B	228 / 229	stem.weight (32,8,1,1,1 → 32,5,1,1,1)
SwinUNETR-L	164 / 167	swinViT.patch_embed.proj.weight, encoder1.layer.conv1.conv.weight, encoder1.layer.conv3.conv.weight

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.