This repository implements an end-to-end deep learning framework for stochastic dose prediction in radiotherapy planning. Our method uses a VQ-VAE architecture augmented with Adaptive Instance Normalization (ADAIN) and a KL divergence–based loss to capture the inherent variability in clinical dose distributions. By perturbing latent style features with Gaussian noise, the model generates multiple clinically plausible dose maps from CT images and ROI masks.
Radiotherapy treatment planning is complex and inherently variable. Traditional deterministic models produce only a single dose distribution per patient, while our framework generates multiple, clinically acceptable dose distributions. Key features include:
- VQ-VAE Backbone: Encodes high-dimensional CT images and ROI masks into a discrete latent space via a learned codebook.
- Adaptive Instance Normalization (ADAIN): Enables the decoder to adapt to varying input distributions.
- Stochastic Prediction: Injects Gaussian noise into style features to produce diverse dose predictions.
- KL Divergence Loss: Aligns the latent distributions of style features (extracted from CT images and ROI masks) and dose-related features (extracted from multi-case dose inputs) to capture clinical variability.
The overall workflow of the model (illustrated in Fig. 1 of the paper) comprises two main components:
-
Dose Generation Network (VQ-VAE):
- Encoder: Processes the input CT images and ROI masks. It partitions the input into patches and generates latent feature maps through down-sampling convolution blocks.
- Codebook Quantization: The latent features are discretized by mapping them to the nearest vector in a learned codebook (with L2 normalization).
- Decoder: Receives the stacked quantized features and perturbed style features; it reconstructs the dose distribution using adaptive instance normalization.
-
Style Feature Extraction & Random Perturbation:
- Dual CNN Encoders:
- One extracts style features from the concatenated CT image and ROI masks.
- A second extracts dose-related features from the CT image, ROI masks, and dose distributions from multiple clinical replans.
- KL Divergence Alignment:
The extracted feature distributions ( P(i) ) and ( Q(i) ) are aligned using KL divergence: [ D_{KL}(P \parallel Q) = \sum_i P(i) \log \frac{P(i)}{Q(i)} ] - Stochastic Perturbation:
A Gaussian noise tensor ( \epsilon ) is applied to the computed style feature mean and variance, following the reparameterization: [ q_{\text{st_re}} = q_{\text{st_mean}} + \epsilon \times q_{\text{st_var}} ] This perturbation encourages diversity in the predicted dose maps.
- Dual CNN Encoders:
The overall loss function is a weighted sum of several components: [ \mathcal{L} = \gamma_d,\mathcal{L}{\text{dose}} + \gamma_m,\mathcal{L}{\text{mask}} + \gamma_p,\mathcal{L}{\text{perceptual}} + \gamma_c,\mathcal{L}{\text{codebook}} + \gamma_k, D_{KL} ] where:
- (\mathcal{L}_{\text{dose}}): An ( L_1 ) loss measuring the voxel-wise difference between predicted and ground truth doses.
- (\mathcal{L}_{\text{mask}}): A loss term that emphasizes dose differences within critical ROI masks.
- (\mathcal{L}_{\text{perceptual}}): Operates on high-level feature maps to preserve structural details.
- (\mathcal{L}_{\text{codebook}}): Ensures that the encoder outputs match the corresponding codebook entries.
- (D_{KL}): Enforces consistency between the latent distributions of the style and dose features.
Two datasets are used for training and evaluation:
-
In-House Dataset:
- Contains CT images, structure contours, and radiotherapy (RT) dose distributions for 76 patients.
- Each patient has five slightly varied replanning cases, providing a range of dose distributions.
- The images are preprocessed (e.g., resampled to 512×512) and normalized (CT values from –1024 to 2048 HU).
-
OpenKBP Dataset:
- A public head-and-neck dataset from the 2020 AAPM Grand Challenge with 340 cases.
- Partitioned into training (n=200), validation (n=40), and test (n=100) sets.
- Consistent image spacing and resolution allow for reproducible dose prediction experiments.
-
Clone the Repository:
git clone https://github.com/yourusername/yourrepository.git cd yourrepository -
Create and Activate a Virtual Environment:
python -m venv venv source venv/bin/activate # For Windows: venv\Scripts\activate
-
Install Dependencies:
pip install -r requirements.txt
(Make sure to include dependencies such as PyTorch, torchvision, and other required libraries.)
- Input Resolution: 128×128 pixels (modifiable based on dataset characteristics).
- Batch Size: 16.
- Learning Rate: Initially set to 1e-5 with decay after 100 epochs.
- Epochs: 200.
- Data Augmentation: Applied during training to improve model robustness.
To start training:
python main.pyThe model’s performance is assessed using:
- Mean Absolute Error (MAE): [ MAE(y, \hat{y}) = \frac{1}{N}\sum_{i=1}^{N} \left|y_i - \hat{y}_i\right| ]
- Dosimetric Metrics: (D_{2}), (D_{95}), and (D_{\text{mean}}) for targets and organs-at-risk (OARs).
- Dose-Volume Histograms (DVH): To compare the clinical relevance of the predicted dose distributions.
Run evaluation with:
python main.py
# command line about trainingThis will output:
- Predicted dose distributions.
- DVH curves for visual and quantitative analysis.
Detailed performance metrics, error maps, and DVH curves are available in the supplementary documentation and accompanying figures.
We thank our collaborators at Fudan University Shanghai Cancer Center and the broader research community in radiotherapy dose prediction. Special thanks to the developers of the OpenKBP dataset.
This project is licensed under the MIT License. See the LICENSE file for further details.
If you find this work useful, please cite:
@article{yourcitation2025,
title={Stochastic Dose Prediction in Radiotherapy via VQ-VAE with KL Divergence},
author={Your Name and Collaborators},
journal={IEEE Transactions on Medical Imaging},
year={2025},
publisher={IEEE}
}Feel free to open issues or submit pull requests if you have suggestions for improvements.