Code for Medical Image Analysis paper Residual Aligner-based Network (RAN): Motion-Separable Structure for Coarse-to-fine Discontinuous Deformable Image Registration
This repo provides an implementation of the training and inference pipeline of RAN based on TensorFlow and Keras.
The research in this paper focuses on solving the problem of multi-organ discontinuous deformation alignment. An innovative quantitative metric, Motion Separability, is proposed in the paper. This metric is designed to measure the ability of deep learning networks to predict organ discontinuous deformations. Based on this metric, a novel network structure skeleton, the Motion-Separable structure, is designed. In addition, we introduce a Motion disentanglement module to help the network distinguish and process complex motion patterns among different organs.
To verify the validity of this quantitative metric as well as the accuracy and efficiency of our method, a series of unsupervised alignment experiments are conducted in the paper. These experiments cover nine major organs of the abdomen and lung images. The experimental results show that the method in the paper is not only able to effectively identify and process the complex motions among the organs, but also improves the accuracy and efficiency of the alignment.
The main contributions include:
- Discontinuous alignment network: this is the first quantitative study targeting discontinuous deformation alignment based on a deep learning network.
- Theoretical analysis: this paper quantifies and defines the maximum range of capturable motion and the upper bound of motion separability in neural networks, providing a theoretical analysis of the upper bound of motion separability. This helps us to understand the range of motion that can be recognised by the network and guides the optimisation of the network structure and parameter settings.
- Motion separable backbone structure: based on the theoretical analysis in this paper, a novel multi-scale skeleton structure is designed in the paper. This structure enables the network to efficiently predict motion patterns with larger separable upper bounds by using optimized dilated convolution on high-resolution feature maps, while maintaining a capturable motion range with low computational complexity.
- Motion decoupling and refinement module: in addition, we propose a Residual Aligner module (RAM) that utilizes confidence levels and mechanisms based on semantic and contextual information to differentiate predicted displacements in different organs or regions. This means that our method can more accurately deal with specific movements in each region.
- Accurate and Efficient Registration Results: The above-proposed components constitute a novel residual alignment network (RAN) that performs efficient, coarse-to-fine, unsupervised alignment of separable motions on publicly available lung and abdominal CT data, achieving higher accuracy and lower computational cost.
Clone code from Github repo: https://github.com/jianqingzheng/res_aligner_net.git
git clone https://github.com/jianqingzheng/res_aligner_net.git
cd res_aligner_net/
install packages
pip install tensorflow==2.3.1
pip install numpy==1.19.5
pip install pyquaternion==0.9.9
Other versions of the packages could also be applicable
[$DOWNLOAD_DIR]/res_aligner_net/
├── data/[$data_name]/dataset
| | # experimental dataset for training and testing (.nii|.nii.gz files)
| ├── train/
| | ├── images/
| | | ├──0001.nii.gz
| | | └── ...
| | ├── labels/
| | | ├──0001.nii.gz
| | | └── ...
| ├── test/
| | ├── images/
| | | ├──0001.nii.gz
| | | └── ...
| | └── labels/
| | ├──0001.nii.gz
| | └── ...
├── models/[$data_name]/
| └── [$data_name]-[$model_name]/
| | # the files of model parameters (.tf.index and .tf.data-000000-of-00001 files)
| ├── model_1_[$model_num].tf.index
| ├── model_1_[$model_num].tf.data-000000-of-00001
| └── ...
└── ...
- Run
python external/deepreg/abd_data.py
to download and setup abdominal CT,
or Runpython external/deepreg/lung_data.py
to download and setup lung CT - Run
python main_preprocess.py --proc_type train --data_name $data_name
- Run
python main_preprocess.py --proc_type test --data_name $data_name
Argument | Description |
---|---|
--data_name |
The data folder name |
* Example for the setup (unpaired_ct_abdomen):
- Run
python external/deepreg/abd_data.py
- Run
python main_preprocess.py --proc_type train --data_name unpaired_ct_abdomen
python main_preprocess.py --proc_type test --data_name unpaired_ct_abdomen
The data used for experiments in this paper are publicly available from abdomen CT and lung CT.
- Run
python main_train.py --model_name $model_name --data_name $data_name --max_epochs $max_epochs
- Check the saved model in
res_aligner_net/models/$data_name/$data_name-$model_name/
Argument | Description |
---|---|
--data_name |
The data folder name |
--model_name |
The used model |
--max_epochs |
The max epoch number for training |
max_epochs==0
for training from scratch
* Example for training (default):
- Run
python main_train.py --model_name RAN4 --data_name unpaired_ct_abdomen --max_epochs 0
- Check the saved model in
res_aligner_net/models/unpaired_ct_abdomen/unpaired_ct_abdomen-RAN4/
- Run
python main_infer.py --model_name $model_name --data_name $data_name
- Check the results in
res_aligner_net/data/$data_name/dataset/test_proc/warped_img
Argument | Description |
---|---|
--data_name |
The data folder name |
--model_name |
The used network structure |
--model_id |
The index of the model |
model_id==1
for a model after synthetic training,model_id==2
for a model after real training,model_id==3
for the model trained according to the paper's settings.
* Example for inference (default):
- Run
python main_infer.py --model_name RAN4 --data_name unpaired_ct_abdomen
- Check the results in
res_aligner_net/data/unpaired_ct_abdomen/dataset/test_proc/warped_img
A demo can be found in the provided notebook.
Alternatively, it can be easily run via .
Any publication that discloses findings arising from using this source code or the network model should cite:
- Zheng, J. Q., Wang, Z., Huang, B., Lim, N. H., & Papież, B. W. "Residual Aligner-based Network (RAN): Motion-separable structure for coarse-to-fine discontinuous deformable registration." Medical Image Analysis, 2024, 91: 103038.
@article{ZHENG2024103038,
title = {Residual Aligner-based Network (RAN): Motion-separable structure for coarse-to-fine discontinuous deformable registration},
journal = {Medical Image Analysis},
volume = {91},
pages = {103038},
year = {2024},
issn = {1361-8415},
doi = {https://doi.org/10.1016/j.media.2023.103038},
url = {https://www.sciencedirect.com/science/article/pii/S1361841523002980},
author = {Jian-Qing Zheng and Ziyang Wang and Baoru Huang and Ngee Han Lim and Bartłomiej W. Papież},
keywords = {Discontinuous deformable registration, Motion-separable structure, Motion disentanglement, Coarse-to-fine registration},
}