A Registration- and Uncertainty-based Framework for White Matter Tract Segmentation with Only One Annotated Subject
Hao Xu, Tengfei Xue, Dongnan Liu, Fan Zhang, Carl-Fredrik Westin, Ron Kikinis, Lauren J. O’Donnell, and Weidong Cai
*The IEEE International Symposium on Biomedical Imaging (ISBI) 2023 [paper,project,code]
White matter (WM) tract segmentation based on diffusion magnetic resonance imaging (dMRI) plays an important role in the analysis of human health and brain diseases. However, the annotation of WM tracts is time-consuming and needs experienced neuroanatomists. In this study, to explore tract segmentation in the challenging setting of minimal annotations, we propose a novel framework utilizing only one annotated subject (subject-level one-shot) for tract segmentation. Our method is constructed by proposed registration-based peak augmentation (RPA) and uncertainty-based refining (URe) modules. Registration-based data augmentation is employed for synthesizing pseudo subjects and their corresponding labels to improve the tract segmentation performance. The proposed uncertainty-based refining module alleviates the negative influence of the low-confidence predictions on pseudo labels. Comparison results indicate the effectiveness of our proposed modules, by achieving accurate whole-brain tract segmentation with only one annotated subject.
-
PyTorch >= 3.6
-
Mrtrix 3 >= 3.0
conda install -c mrtrix3 mrtrix3 python=3.6
-
boto3
pip install boto3
-
tractseg
pip install tractseg
-
nibabel
pip install nibabel
You can prepare datasets by yourself or follow the following steps.
- Download Human Connectome Project (HCP) datasets.
- Register a HCP account: HCP
- Enable Amazon S3 Access: AWS
- Download HCP datasets by running download_HCP_1200_diffusion_mri.py:
python /download_HCP_1200_diffusion_mri.py --id your_aws_id --key your_aws_key --out_dit your_hcp_dir
- Download Corresponding WM tract labels from Zenodo.
- Transform a trk streamline file to a binary map. Transform a trk streamline file to a binary map by running trk2bin.py:
python /trk2bin.py --tract_dir your_tract_dir --ref_dir your_hcp_dir
and finally, the tract dataset directory should look like:
$your_tract_dir
├─992774
│ ├─tracts
│ │ ├─AF_left.nii.gz
│ │ ├─AF_rgiht.nii.gz
| | .
| | .
| | .
│ │ ├─UF_rgiht.nii.gz
├─991267
│ ├─tracts
│ │ ├─AF_left.nii.gz
│ │ ├─AF_rgiht.nii.gz
| | .
| | .
| | .
│ │ ├─UF_rgiht.nii.gz
.
.
.
├─599469
│ ├─tracts
│ │ ├─AF_left.nii.gz
│ │ ├─AF_rgiht.nii.gz
| | .
| | .
| | .
│ │ ├─UF_rgiht.nii.gz
- Transform dMRI datasets to peak data. Transform dMRI datasets to peak data using multi-shell multi-tissue constrained spherical deconvolution (MSMT-CSD) by running HCP2MSMT_CSD.py:
python /HCP2MSMT_CSD.py --hcp_dir your_hcp_dir --out_dir your_msmt_csd_dir
and finally, the peak data directory should look like:
$your_msmt_csd_dir
├─992774
│ ├─peaks.nii.gz
├─991267
│ ├─peaks.nii.gz
.
.
.
├─599469
│ ├─peaks.nii.gz
- Train
python \code\engine.py --action step_1_train --step_1_ckpt_dir your_step_1_ckpt_dir --batch_size 1 --lr 0.001 --epoch 100 --model vm2 --sim_loss mse --alpha 0.02
- Generate pseudo data
python \code\engine.py --action step_1_generation --step_1_ckpt_dir your_step_1_ckpt_dir --step_1_pseudo_data_dir your_step_1_pseudo_data_dir
python \code\engine.py --action step_2_train --step_2_ckpt_dir your_step_2_ckpt_dir --batch_size 48 --lr 0.002 --epoch 200 --LOSS_WEIGHT 10 --LOSS_WEIGHT_LEN 200
python \code\engine.py --action step_3_train --step_1_pseudo_data_dir your_step_1_pseudo_data_dir --step_2_ckpt_dir your_step_2_ckpt_dir --step_3_ckpt_dir your_step_3_ckpt_dir --batch_size 48 --lr 0.001 --epoch 200 --LOSS_WEIGHT 10 --LOSS_WEIGHT_LEN 200
python \code\engine.py --action test --step_3_ckpt_dir your_step_3_ckpt_dir
Methods | Dice score | |
Comparison | U-Net | 43.19±15.20% |
TractSeg | 48.85±17.58% | |
Ablation study | Ours (RPA) | 69.45±9.53% |
Ours (RPA + UDA) | 73.01±8.14% |
Our code is based on TractSeg.
@inproceedings{Xu_2023_ISBI,
title={A Registration- and Uncertainty-based Framework for White Matter Tract Segmentation With Only One Annotated Subject },
author={Xu, Hao and Xue, Tengfei and Liu, Dongnan and Zhang, Fan and Westin, Carl-Fredrik and Kikinis, Ron and J. O'Donnell, Lauren and Cai, Weidong},
booktitle={2023 IEEE 20th International Symposium on Biomedical Imaging (ISBI)},
year={2023},
organization={IEEE}
}