--- license: cc-by-nc-4.0 tags: - radiotherapy - dose-prediction - medical-imaging - transfer-learning - 3d-segmentation - pytorch library_name: pytorch pipeline_tag: image-to-image --- # 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](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 | | `han_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: ```bash 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 paper's models use 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](#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 identical between HaN and pancreas. **Channels 0, 1 and 2 are not** — this is the trap: | ch | `han_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 HaN 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 consequences worth knowing: * **In `han_5ch`, channel 0 and channel 1 are byte-identical.** The HaN loader builds the optimisation and planning PTV sets from the same `StructName` masks, so the channel is duplicated. It is kept as a separate channel only to preserve the pretrained model's layout. (Verified on real cases, not inferred from the code.) * **In `pancreas_5ch`, channel 1 is capped at 3.0.** A 40 Gy PTV appears as `3.0`, not `4.0`, while channel 0 for the same PTV is uncapped at `4.0`. If you reimplement the pipeline and skip the cap, the model sees inputs it was never trained on. The dose *target* differs too: HaN normalises the planned dose to the PTV_High D97 before dividing by `dose_div_factor`; pancreas does not renormalise (`norm_scale = 1.0`). --- ## Available weights All files are plain `state_dict` + metadata, safe under `torch.load(..., weights_only=True)`. SHA-256 sums are in [`weights/MANIFEST.json`](weights/MANIFEST.json). | file | arch | input encoding | params | size | |---|---|---|---|---| | `c3d_pretrained.pt` | C3D | `source_8ch` | 32.4 M | 129 MB | | `c3d_finetuned.pt` | C3D | `han_5ch` | 32.3 M | 129 MB | | `c3d_fromscratch.pt` | C3D | `han_5ch` | 32.3 M | 129 MB | | `mednext_pretrained.pt` | MedNeXt-B | `source_8ch` | 10.5 M | 42 MB | | `mednext_finetuned.pt` | MedNeXt-B | `han_5ch` | 10.5 M | 42 MB | | `mednext_fromscratch.pt` | MedNeXt-B | `han_5ch` | 10.5 M | 42 MB | | `swinunetr_pretrained.pt` | SwinUNETR-L | `source_8ch` | 291.4 M | 1169 MB | | `swinunetr_finetuned.pt` | SwinUNETR-L | `han_5ch` | 291.3 M | 1169 MB | | `swinunetr_fromscratch.pt` | SwinUNETR-L | `han_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 ```bash pip install -r requirements.txt # MedNeXt checkpoints only: pip install git+https://github.com/MIC-DKFZ/MedNeXt.git ``` Minimal load-and-predict: ```python 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"]) # han_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: ```bash cp -r meta_files_template/han meta_files # then fill in your own cases # edit configs/c3d_han.yaml -> loader_params.data_root python src/inference.py \ --ckpt weights/c3d_finetuned.pt \ --config configs/c3d_han.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 `_pred.nii.gz`. Inspect a checkpoint's expected channels without running anything: ```bash 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. Because only the input channel count changed, that means **only the input stem was re-initialised**: | 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` | **No channel slicing, averaging or repetition was applied to the input stem** — it was simply retrained from scratch alongside the rest of the network. `src/transfer.py` reproduces this exactly: ```python 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 ``` Training hyperparameters used for the released head-and-neck C3D run: Adam, `lr = 3e-5`, `weight_decay = 1e-4`, cosine schedule to `eta_min = 1e-7`, 275 epochs. The `_fromscratch` baseline used the same schedule at `lr = 1e-4`. --- ## Repository layout ``` ├── README.md ├── LICENSE CC BY-NC 4.0 + no-clinical-use notice ├── requirements.txt ├── configs/ one YAML per model/anatomy │ ├── c3d_han.yaml │ ├── mednext_han.yaml │ ├── swinunetr_han.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_han.py produces han_5ch │ ├── toolkit_han.py │ ├── data_loader_pancreas.py produces pancreas_5ch │ └── toolkit_pancreas.py ├── meta_files_template/ de-identified metadata schema + data layout │ ├── han/ 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: ```bibtex @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](mailto:ke.sheng@ucsf.edu) **Paper:** *Medical Physics* (in press)  ·  **Model repository:** `https://huggingface.co//radiotherapy-dose-prediction-transfer-learning` ## Please also cite the work this builds on ```bibtex @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](LICENSE) — attribution required, non-commercial use only. **Research use only; not for clinical use.**