|
1 | | -import os |
| 1 | +#!/usr/bin/env python |
| 2 | +# coding: utf-8 |
| 3 | + |
| 4 | +# # <center>Pipeline for Deep learning based Dose Prediction and Fluence Optimization </center> |
| 5 | +# |
| 6 | +# |
| 7 | +# This example showcases the core functionalities of the portpy.ai module, which is designed for deep learning-based dose prediction in the context of radiotherapy treatment planning. The main functionalities demonstrated are as follows: |
| 8 | +# |
| 9 | +# 1. **Data Access and Preprocessing**: The portpy.ai module allows easy access to curated portpy data, specifically tailored for training deep learning models for dose prediction. It also includes preprocessing steps to prepare the data for use in the models. |
| 10 | +# |
| 11 | +# 2. **Model Training and Testing**: With the portpy.ai module, users can efficiently train and test their deep learning models for dose prediction. This involves using the prepared data to create and fine-tune the models for optimal performance. |
| 12 | +# |
| 13 | +# 3. **Patient-Specific Dose Prediction**: The module enables the creation of personalized dose prediction for available patients in the PortPy database. This feature allows users to estimate the dose distribution for individual patients, aiding in treatment planning. Also, PortPy uses simple UNET model with MAE loss to create dose prediction. Users can implement their own customized models for better prediction and accuracy |
| 14 | +# |
| 15 | +# |
| 16 | +# 4. **Fluence Optimization with Predicted Dose**: The predicted dose distribution can be used to optimize the fluence of radiation beams during treatment planning. The optimized fluence can be imported back to eclipse for final dose calculation. Please refer to [imrt_tps_import.ipynb](https://github.com/PortPy-Project/PortPy/blob/master/examples/imrt_tps_import.ipynb) for more details. |
| 17 | +# |
| 18 | +# 5. **Plan Visualization**: The module offers tools for visualizing treatment plans based on the predicted dose distribution. |
| 19 | +# |
| 20 | +# Overall, the portpy.ai module provides a comprehensive solution for dose prediction and treatment planning in radiotherapy, leveraging the power of deep learning and data-driven methodologies to enhance precision and patient-specific care. |
| 21 | +# |
| 22 | +# **Note: At present, the PortPy dataset contains only 10 lung patients. However, using the PortPy AI module, treatment planning using dose prediction pipeline has been demonstrated successfully. By the end of August, an additional 100 lung patients will be added to the PortPy dataset, significantly enhancing the training data and thereby improving the performance of the dose prediction model. ** |
| 23 | + |
2 | 24 | import portpy.photon as pp |
3 | 25 | import SimpleITK as sitk |
4 | 26 | import os |
5 | 27 | import cvxpy as cp |
6 | 28 | import numpy as np |
7 | 29 | import matplotlib.pyplot as plt |
8 | 30 | from portpy.ai.preprocess.predict_using_model import predict_using_model |
| 31 | +from portpy.ai.preprocess.data_preprocess import data_preprocess |
| 32 | +from portpy.ai.train import train |
| 33 | +from portpy.ai.test import test |
| 34 | + |
| 35 | + |
| 36 | +# ### 1. Preprocess, train and test using PortPy AI module |
| 37 | + |
| 38 | + |
| 39 | + |
| 40 | +in_dir = r'../data' # directory where portpy raw data is located |
| 41 | +out_dir = r'../ai_data' # directory where processed data to be stored for training and testing model |
| 42 | + |
| 43 | + |
| 44 | +# preprocess portpy data |
| 45 | +data_preprocess(in_dir, out_dir) |
| 46 | + |
| 47 | +# **Note** split the data in train and test folder in the output directory before running further code |
| 48 | +# e.g. out_dir\train\Lung_Patient_2 out_dir\test\Lung_Patient_9 |
| 49 | + |
| 50 | +# ### 2. Training and testing the model |
| 51 | +# |
| 52 | +# Train the model. You can change the parameters for the training as show below |
| 53 | + |
| 54 | +# Provide only the arguments you want to override |
| 55 | +train_options = { |
| 56 | + "dataroot": "../../ai_data", |
| 57 | + "checkpoints_dir": "../../checkpoints", |
| 58 | + "netG": "unet_128", |
| 59 | + "name": "portpy_test_3", |
| 60 | + "model": "doseprediction3d", |
| 61 | + "direction": "AtoB", |
| 62 | + "lambda_L1": 1, |
| 63 | + "dataset_mode": "dosepred3d", |
| 64 | + "norm": "batch", |
| 65 | + "batch_size": 1, |
| 66 | + "pool_size": 0, |
| 67 | + "display_port": 8097, |
| 68 | + "lr": 0.0002, |
| 69 | + "input_nc": 8, |
| 70 | + "output_nc": 1, |
| 71 | + "display_freq": 10, |
| 72 | + "print_freq": 1, |
| 73 | + "gpu_ids": [0] # Converted to a list since multiple GPUs may be supported |
| 74 | +} |
| 75 | + |
| 76 | +train(train_options) # Run training directly in Jupyter Notebook |
| 77 | + |
| 78 | +# You can uncomment and run below in case if you want to run train script from CLI |
| 79 | +#!python ../portpy/ai/train.py --dataroot ../ai_data --netG unet_128 --name portpy_test_3 --model doseprediction3d --direction AtoB --lambda_L1 1 --dataset_mode dosepred3d --norm batch --batch_size 1 --pool_size 0 --display_port 8097 --lr 0.0002 --input_nc 8 --output_nc 1 --display_freq 10 --print_freq 1 --gpu_ids 0 |
| 80 | + |
| 81 | + |
| 82 | +# Test the model |
| 83 | +test_options = { |
| 84 | + "dataroot": "../ai_data", |
| 85 | + "netG": "unet_128", |
| 86 | + "checkpoints_dir": "../../checkpoints", |
| 87 | + "results_dir": "../../results", |
| 88 | + "name": "portpy_test_3", |
| 89 | + "phase": "test", |
| 90 | + "mode": "eval", |
| 91 | + "eval": True, # Boolean flag |
| 92 | + "model": "doseprediction3d", |
| 93 | + "input_nc": 8, |
| 94 | + "output_nc": 1, |
| 95 | + "direction": "AtoB", |
| 96 | + "dataset_mode": "dosepred3d", |
| 97 | + "norm": "batch" |
| 98 | +} |
| 99 | +test(test_options) |
| 100 | +# !python ../portpy/ai/test.py --dataroot ../ai_data --netG unet_128 --checkpoints_dir ../checkpoints --results_dir ../results --name portpy_test_2 --phase test --mode eval --eval --model doseprediction3d --input_nc 8 --output_nc 1 --direction Ato |
| 101 | + |
| 102 | + |
| 103 | +# ### 3. Create dose prediction for the PortPy patient |
| 104 | +# |
| 105 | +# After training the model, users can use their customized model to create dose prediction for any portpy patient. Below script would preprocess and predict for the patient |
| 106 | +# For users who does not want to train and test, they can directly preprocess and predict using AI model |
| 107 | +patient_id = 'Lung_Patient_4' |
| 108 | +model_name = 'portpy_test_3' |
| 109 | +pred_dose = predict_using_model(patient_id=patient_id, in_dir=in_dir, out_dir=out_dir, model_name=model_name, checkpoints_dir='../../checkpoints', results_dir='../../results') |
9 | 110 |
|
10 | | -# change directory to portpy ai module to preprocess the portpy data, train and test the model |
11 | | -os.chdir('../../portpy/ai') |
12 | | - |
13 | | -# # preprocess portpy data |
14 | | -in_dir = r'../../data' |
15 | | -out_dir = r'../../ai_data' |
16 | | - |
17 | | -os.system('python ./preprocess/data_preprocess.py --in_dir ../../data --out_dir ../../ai_data') |
18 | | - |
19 | | -# Train preprocess data |
20 | | -os.system('python train.py --dataroot ../../ai_data --netG unet_128 --name portpy_test_2 --model doseprediction3d --direction AtoB --lambda_L1 1 --dataset_mode dosepred3d --norm batch --batch_size 1 --pool_size 0 --display_port 8097 --lr 0.0002 --input_nc 8 --output_nc 1 --display_freq 10 --print_freq 1 --gpu_ids 0') |
21 | | - |
22 | | -# Test preprocess data |
23 | | -os.system('python test.py --dataroot ../../ai_data --netG unet_128 --name portpy_test_2 --phase test --mode eval --model doseprediction3d --input_nc 8 --output_nc 1 --direction AtoB --dataset_mode dosepred3d --norm batch') |
24 | 111 |
|
25 | | -# predicted dose back to portpy |
26 | | -patient_id = 'Lung_Patient_4' |
27 | | -model_name = 'portpy_test_2' |
28 | | -pred_dose = predict_using_model(patient_id=patient_id, in_dir=in_dir, model_name=model_name) |
| 112 | +# ### 4. Fluence optimization using predicted dose |
29 | 113 |
|
30 | 114 | # load portpy data |
31 | 115 | data = pp.DataExplorer(data_dir=in_dir) |
32 | 116 | data.patient_id = patient_id |
33 | 117 | # Load ct and structure set for the above patient using CT and Structures class |
34 | 118 | ct = pp.CT(data) |
35 | | -ct_arr = ct.ct_dict['ct_hu_3d'][0] |
36 | 119 | structs = pp.Structures(data) |
37 | 120 |
|
38 | 121 | beams = pp.Beams(data) |
39 | 122 |
|
| 123 | +# load clinical criteria from the config files for which plan to be optimized |
| 124 | +protocol_name = 'Lung_2Gy_30Fx' |
| 125 | +clinical_criteria = pp.ClinicalCriteria(data, protocol_name=protocol_name) |
| 126 | + |
40 | 127 | # create rinds based upon rind definition in optimization params |
41 | 128 | protocol_name = 'Lung_2Gy_30Fx' |
42 | 129 | opt_params = data.load_config_opt_params(protocol_name=protocol_name) |
|
45 | 132 | # load influence matrix based upon beams and structure set |
46 | 133 | inf_matrix = pp.InfluenceMatrix(ct=ct, structs=structs, beams=beams) |
47 | 134 |
|
48 | | -# load clinical criteria from the config files for which plan to be optimized |
49 | | -clinical_criteria = pp.ClinicalCriteria(data, protocol_name=protocol_name) |
50 | | - |
51 | 135 | pred_dose_1d = inf_matrix.dose_3d_to_1d(dose_3d=pred_dose) |
52 | 136 |
|
53 | 137 | # create a plan using ct, structures, beams and influence matrix. Clinical criteria is optional |
54 | | -my_plan = pp.Plan(ct, structs, beams, inf_matrix, clinical_criteria) |
| 138 | +my_plan = pp.Plan(ct=ct, structs=structs, beams=beams, inf_matrix=inf_matrix, clinical_criteria=clinical_criteria) |
55 | 139 |
|
56 | 140 | # create cvxpy problem using the clinical criteria and optimization parameters |
57 | 141 | opt = pp.Optimization(my_plan, opt_params=opt_params) |
|
75 | 159 | opt.obj += [(smoothness_X_weight * (1 / num_cols) * cp.sum_squares(Qx @ x) + |
76 | 160 | smoothness_Y_weight * (1 / num_rows) * cp.sum_squares(Qy @ x))] |
77 | 161 |
|
78 | | -sol = opt.solve(solver='MOSEK', verbose=True) |
79 | | -# sol = opt.solve(solver='MOSEK', verbose=True) |
| 162 | +sol = opt.solve(solver='MOSEK', verbose=False) |
80 | 163 | sol = {'optimal_intensity': x.value, 'inf_matrix': inf_matrix} |
81 | 164 |
|
| 165 | + |
| 166 | +# ### 5. Plan visualization |
| 167 | + |
82 | 168 | # plot fluence 3d and 2d for the 1st beam |
83 | 169 | pp.Visualization.plot_fluence_3d(sol=sol, beam_id=my_plan.beams.get_all_beam_ids()[0]) |
84 | 170 |
|
|
89 | 175 | norm_flag=True) |
90 | 176 | ax.set_title('- Optimized .. Predicted') |
91 | 177 | plt.show() |
| 178 | + |
| 179 | + |
| 180 | + |
| 181 | + |
0 commit comments