zhennongchen / 2DUNet_CT_Seg_Final_v_ZC

a 2D U-Net pipeline to segment heart chambers in CT 3D volumes.

U-Net-2D

Author: Zhennong Chen, PhD

Description:

This repo enables the deep learning (DL) heart chamber segmentation of CT volumes in a 2D slice-by-slice fashion.

It includes scripts to pre-process the image for DL training, DL model training and DL model prediction.

Citation: Zhennong Chen, Marzia Rigolli, Davis Marc Vigneault, Seth Kligerman, Lewis Hahn, Anna Narezkina, Amanda Craine, Katherine Lowe, Francisco Contijoch, Automated cardiac volume assessment and cardiac long- and short-axis imaging plane prediction from electrocardiogram-gated computed tomography volumes enabled by deep learning, European Heart Journal - Digital Health, Volume 2, Issue 2, June 2021, Pages 311–322, https://doi.org/10.1093/ehjdh/ztab033

Install:

The entire code is containerized. This makes setting up environment swift and easy. Make sure you have nvidia-docker and Docker CE installed on your machine before going further.

• The docker file is included in this repo.
  
• The code relies on dvpy python package, make sure you have installed the latest version of it. If not, do: pip uninstall dvpy; pip install git+https://github.com/zhennongchen/dvpy.git#egg=dvpy
  

Data Preparation:

To train the model, make sure you have prepared your images and manual/ground truth segmentation. Here is a list of things you need to do for data preparation.

1. tool_resample_by_c3d.sh: resample the CT volumes/manual segmentations to a uniform pixel dimension (required for U-Net input). default = 0.625mm^3
2. pre_adapt_image.py: pre-process the image for model training
3. pre_partition.py: randomly split the patient list if to do n-fold cross-validation.

Main Script:

• set_defaults.sh: define the parameters&folders for DL experiments.
• main_train.py: to train the model, using n-1 subsamples for training and the rest 1 subsample for validation.
• main_validate.py: to validate the model for n-fold cross-validation
• main_predict.py: to predict segmentation on new cases by trained DL model
  • in folder post_processing: post_process the predicted segmentation (optional), mainly to exclude the disconnected parts. first run tool_exclude_disconnectivity.m, and then run tool_mat_to_nii.ipynb to turn mat file to nii file.

Additional Guidelines

see comments in the script

Please contact [e-mail] or [e-mail] for any further questions.