firstmover/cr-seg

Source code for "Consistency Regularization Improves Placenta Segmentation in Fetal EPI MRI Time Series" paper.

Consistency Regularization Improves Placenta Segmentation in Fetal EPI MRI Time Series

This repo includes the code for the paper Liu et al. 2023.

---

Environment and dependency

Create a conda environment with:

conda create --name cr_seg --file requirements.txt

Specify your paths to data, cache, and results directories in:

1. ./envs/default
2. ./configs/segm_release/spatial_temporal_cr.py

Results

Pre-process of data

srun --partition=gpu \
--job-name=segm \
--gres=gpu:1 \
--ntasks=1 \
--ntasks-per-node=1 \
--cpus-per-task=16 \
--time=24:00:00 \
python scripts/pre_compute_data.py

Training

pretrain registration models for all cross-validation folds

python ./scripts/submit_job_registraion.py \
--exp-name regi_release \
--config-name voxelmorph \
--job-name=regi \
--num-gpus-per-node=1 \
--cpus-per-task=20 \
--num-nodes=1 \
--array-parallelism=5

Train UNet with spatial and temporal consistency regularization for all cross-validation folds

python ./scripts/submit_job_segmentation.py \
--task-mode train \
--exp-name segm_release \
--config-name spatial_temporal_cr \
--lambda-list '0.001' \
--lambda-t-list '0.001' \
--job-name segm_regi \
--num-gpus-per-node 4 \
--cpus-per-task 8 \
--array-parallelism 5

Inference and evaluation

Run inference for labeled data for all cross-validation folds

python ./scripts/submit_job_segmentation.py \
--task-mode inference_labeled \
--exp-name segm_release \
--config-name spatial_temporal_cr \
--lambda-list '0.001' \
--lambda-t-list '0.001' \
--tta \
--tta-cfg-path ./configs/segm_release/_base_/tta_all_crop.py \
--save-data-name-list 'img,pred_seg_map,gt_seg_map' \
--job-name inference \
--partition gpu \
--num-gpus-per-node 1 \
--cpus-per-task 16 \
--array-parallelism 5

Run inference for time series data (unlabeled and labeled data) for all cross-validation folds

python ./scripts/submit_job_segmentation.py \
--task-mode inference_time_series \
--exp-name segm_release \
--config-name spatial_temporal_cr \
--lambda-list '0.001' \
--lambda-t-list '0.001' \
--tta \
--tta-cfg-path ./configs/segm_release/_base_/tta_all_crop.py \
--save-data-name-list 'pred_seg_map' \
--job-name inference \
--partition gpu \
--num-gpus-per-node 1 \
--cpus-per-task 16 \
--array-parallelism 5

Visualization

Pre-compute and evulate time series data

python scripts/visualization/pre_compute_time_series.py \
--result_root ./results/segm_release \
--job-name eval \
--partition gpu \
--num-gpus-per-node 1 \
--cpus-per-task 24 \
--array-parallelism 8

Visualize labeled results

streamlit run scripts/visualization/labeled.py -- --result_root ./results/segm_release --model_name epoch_100_all

Visualize time series results

streamlit run scripts/visualization/time_series.py -- --result_root ./results/segm_release --model_name epoch_100_all

todos

• Add commands for non-slurm users
• Add more details for data set structures
• Update citation

Acknowledgement

• mabulnaga/automatic-placenta-segmentation
• voxelmorph/voxelmorph

License

This repo is licensed under the MIT License and the copyright belongs to all authors - see the LICENSE file for details.

Citation

@inproceedings{liu2023consistency,
  title={Consistency Regularization Improves Placenta Segmentation in Fetal EPI MRI Time Series},
  author={Liu, Yingcheng and Karani, Neerav and Abulnaga, S Mazdak and Xu, Junshen and Grant, P Ellen and Abaci Turk, Esra and Golland, Polina},
  booktitle={International Workshop on Preterm, Perinatal and Paediatric Image Analysis},
  pages={77--87},
  year={2023},
  organization={Springer}
}

Contact

Email: [email protected]