TotalSpineSeg

Tool for automatic segmentation and labelling of all vertebrae and intervertebral discs (IVDs), spinal cord, and spinal canal. We follow TotalSegmentator classes with an additional class for IVDs, spinal cord and spinal canal (See list of class here). We used nnUNet as our backbone for model training and inference.

If you use this model, please cite our work:

Warszawer Y, Molinier N, Valošek J, Shirbint E, Benveniste PL, Achiron A, Eshaghi A and Cohen-Adad J. Fully Automatic Vertebrae and Spinal Cord Segmentation Using a Hybrid Approach Combining nnU-Net and Iterative Algorithm. Proceedings of the 32th Annual Meeting of ISMRM. 2024

• Dependencies
• Installation
• First Model
  • First Model - Train
  • First Model - Inference
• List of class

Dependencies

• Spinal Cord Toolbox (SCT)

Installation

1. Open Terminal in a directory you want to work on.

2. Create and activate Virtual Environment (Highly recommanded):

   python -m venv venv
   source venv/bin/activate

3. Install PyTorch as described on their website.

4. Clone and install this repository:

   git clone https://github.com/neuropoly/totalsegmentator-mri.git
   python -m pip install -e totalsegmentator-mri

5. Install requirements:

   python -m pip install -r totalsegmentator-mri/requirements.txt

First Model

A hybrid approach integrating nnU-Net with an iterative algorithm for segmenting vertebrae, IVDs, spinal cord, and spinal canal. To tackle the challenge of having many classes and class imbalance, we developed a two-step training process. A first model (model 1 - 206) was trained (single input channel: image) to identify 4 classes (IVDs, vertebrae, spinal cord and spinal canal) as well as specific IVDs (C2-C3, C7-T1 and L5-S1) representing key anatomical landmarks along the spine, so 7 classes in total (Figure 1A). The output segmentation was processed using an algorithm that distinguished odd and even IVDs based on the C2-C3, C7-T1 and L5-S1 IVD labels output by the model (Figure 1B). Then, a second nnU-Net model (model 2 - 210) was trained (two input channels: 1=image, 2=odd IVDs), to output 12 classes (Figure 1C). Finally, the output of model 2 was processed in order to assign an individual label value to each vertebrae and IVD in the final segmentation mask (Figure 1D).

Figure 1: Illustration of the hybrid method for automatic segmentation of the spine and spinal cord structures. T1w image (A) is used to train model 1, which outputs 7 classes (B). These output labels are processed to extract odd IVDs (C). The T1w and odd IVDs are used as two input channels to train model 2, which outputs 12 classes (D). These output labels are processed to extract individual IVDs and vertebrae (E).

First Model - Train

1. Download the corresponding content from SPIDER dataset into 'data/raw/spider/images' and 'data/raw/spider/masks' (you can use mkdir -p data/raw/spider to create the folder first).

2. Make sure git and git-annex are installed (You can install with sudo apt-get install git-annex -y).

3. Extract data-multi-subject_PAM50_seg.zip into 'data/bids/data-multi-subject'.

4. Extract data-single-subject_PAM50_seg.zip into 'data/bids/data-single-subject'.

5. Extract whole-spine.zip (private dataset) into 'data/bids/whole-spine'.

6. Get the required datasets from Spine Generic Project:

   source totalsegmentator-mri/run/get_spine_generic_datasets.sh

7. Prepares SPIDER datasets in BIDS structure:

   source totalsegmentator-mri/run/prepare_spider_bids_datasets.sh

8. Prepares datasets in nnUNetv2 structure:

   source totalsegmentator-mri/run/prepare_nnunet_datasets.sh

9. Train the model:

   source totalsegmentator-mri/run/train_nnunet.sh

First Model - Inference

Run the model on a folder containing the images in .nii.gz format (Make sure to train the model or extract the trained nnUNet_results into data/nnUNet/nnUNet_results befor running):

source totalsegmentator-mri/run/inference_nnunet.sh INPUT_FOLDER OUTPUT_FOLDER

List of class