Lumbar rootlets - model training on `Draw Tube` labels #67

valosekj · 2024-07-19T18:38:02Z

This issue summarizes model training on T2w lumbar data with relabeled rootlets using the 3D Slicer Draw Tube module.

This is a follow up of #48.

0. Data overview

We have 6 subjects with the following labels:

sub_labels = {"sub-CTS04": { "start": "T11","end": "S1"},
             "sub-CTS05": {"start": "T11","end": "S1"},
             "sub-CTS09": {"start": "T10","end": "S2"},
             "sub-CTS10": {"start": "T11","end": "S1"},
             "sub-CTS14": {"start": "T11","end": "S1"},
             "sub-CTS15": {"start": "T10","end": "S2"}}

1. Preparing nnUNet folders

details

# imagesTr
cp sub-CTS04_ses-SPpre_acq-zoomit_T2w.nii.gz Dataset301_LumbarRootlets/imagesTr/sub-CTS04_ses-SPpre_T2w_001_0000.nii.gz
cp sub-CTS05_ses-SPpre_acq-zoomit_T2w.nii.gz Dataset301_LumbarRootlets/imagesTr/sub-CTS05_ses-SPpre_T2w_001_0000.nii.gz
cp sub-CTS09_ses-SPpre_acq-zoomit_T2w.nii.gz Dataset301_LumbarRootlets/imagesTr/sub-CTS09_ses-SPpre_T2w_001_0000.nii.gz
cp sub-CTS10_ses-SPanat_acq-zoomit_T2w.nii.gz Dataset301_LumbarRootlets/imagesTr/sub-CTS10_ses-SPanat_T2w_001_0000.nii.gz
cp sub-CTS14_ses-SPpre_acq-zoomit_T2w.nii.gz Dataset301_LumbarRootlets/imagesTr/sub-CTS14_ses-SPpre_T2w_001_0000.nii.gz
cp sub-CTS15_ses-SPpre_acq-zoomit_T2w.nii.gz Dataset301_LumbarRootlets/imagesTr/sub-CTS15_ses-SPpre_T2w_001_0000.nii.gz

# labelsTr
cp T11-S1_RD_LD_sub-CTS04_relabeled.nii.gz Dataset301_LumbarRootlets/labelsTr/sub-CTS04_ses-SPpre_T2w_001.nii.gz
cp T11-S1_RD_LD_sub-CTS05_relabeled.nii.gz Dataset301_LumbarRootlets/labelsTr/sub-CTS05_ses-SPpre_T2w_001.nii.gz
cp T10-S2_RD_LD_sub-CTS09_relabeled.nii.gz Dataset301_LumbarRootlets/labelsTr/sub-CTS09_ses-SPpre_T2w_001.nii.gz
cp T11-S1_RD_LD_sub-CTS10_relabeled.nii.gz Dataset301_LumbarRootlets/labelsTr/sub-CTS10_ses-SPanat_T2w_001.nii.gz
cp T11-S1_RD_LD_sub-CTS14_relabeled.nii.gz Dataset301_LumbarRootlets/labelsTr/sub-CTS14_ses-SPpre_T2w_001.nii.gz
cp T10-S2_RD_LD_sub-CTS15_relabeled.nii.gz Dataset301_LumbarRootlets/labelsTr/sub-CTS15_ses-SPpre_T2w_001.nii.gz

2. Changing label values to be consecutive (this is required by nnUNet)

details

Original values

cd labelsTr
for file in *nii.gz;do get_unique_values $file;done
[ 0. 19. 20. 21. 22. 23. 24. 25. 26.]
[ 0. 19. 20. 21. 22. 23. 24. 25. 26.]
[ 0. 18. 19. 20. 21. 22. 23. 24. 25. 26. 27.]
[ 0. 19. 20. 21. 22. 23. 24. 25. 26.]
[ 0. 19. 20. 21. 22. 23. 24. 25. 26.]
[ 0. 18. 19. 20. 21. 22. 23. 24. 25. 26. 27.]

Removing label 18 (presented only for two subjects) for now:

sct_maths -i sub-CTS09_ses-SPpre_T2w_001.nii.gz -thr 19 -o sub-CTS09_ses-SPpre_T2w_001.nii.gz
sct_maths -i sub-CTS15_ses-SPpre_T2w_001.nii.gz -thr 19 -o sub-CTS15_ses-SPpre_T2w_001.nii.gz

Recoding using recode_nii.py:

for file in *nii.gz;do python ~/code/model-spinal-rootlets/utilities/recode_nii.py -i $file -o $file;done
Unique values in the data: [ 0. 19. 20. 21. 22. 23. 24. 25. 26.]
Unique values in the recoded data: [0 1 2 3 4 5 6 7 8]
Unique values in the data: [ 0. 19. 20. 21. 22. 23. 24. 25. 26.]
Unique values in the recoded data: [0 1 2 3 4 5 6 7 8]
Unique values in the data: [ 0. 19. 20. 21. 22. 23. 24. 25. 26. 27.]
Unique values in the recoded data: [0 1 2 3 4 5 6 7 8 9]
Unique values in the data: [ 0. 19. 20. 21. 22. 23. 24. 25. 26.]
Unique values in the recoded data: [0 1 2 3 4 5 6 7 8]
Unique values in the data: [ 0. 19. 20. 21. 22. 23. 24. 25. 26.]
Unique values in the recoded data: [0 1 2 3 4 5 6 7 8]
Unique values in the data: [ 0. 19. 20. 21. 22. 23. 24. 25. 26. 27.]
Unique values in the recoded data: [0 1 2 3 4 5 6 7 8 9]

3. Training

fold1, 4 training and 2 validation images.

Semantic (level-specific) model: `Dataset301_LumbarRootlets`

cd ~/code/model-spinal-rootlets/training
bash run_training.sh 1 301 Dataset301_LumbarRootlets

Binary model (all rootlets set to 1): `Dataset302_LumbarRootlets`

Binarize labels and modify dataset.json:

cd $nnUNet_raw
cp -r Dataset301_LumbarRootlets Dataset302_LumbarRootlets
cd Dataset302_LumbarRootlets/labelsTr
for file in *nii.gz;do sct_maths -i $file -bin 0.5 -o $file;done
cd ..
# modify dataset.json

cd ~/code/model-spinal-rootlets/training
bash run_training.sh 1 302 Dataset302_LumbarRootlets

The text was updated successfully, but these errors were encountered:

valosekj · 2024-07-23T11:39:00Z

The training of both models has finished.

TL;DR: predictions of the binary model (Dataset302_LumbarRootlets) on unseen subjects are quite good.
Interestingly, some rootlets predicted by the model (Dataset202_LumbarRootlets; #48) trained on GT created using the FSLeyes were not predicted by the model (Dataset302_LumbarRootlets) trained on GT created using the Slicer Draw Tube tool and vice versa.

Semantic (level-specific) model (Dataset301_LumbarRootlets)

training_log

2024-07-21 01:07:55.046880: Current learning rate: 5e-05
2024-07-21 01:09:39.849596: train_loss -0.8012
2024-07-21 01:09:39.849765: val_loss -0.3613
2024-07-21 01:09:39.849864: Pseudo dice [0.0, 0.2224, 0.4656, 0.2038, 0.0393, 0.1114, 0.1499, 0.1689, 0.0]
2024-07-21 01:09:39.849931: Epoch time: 104.8 s
2024-07-21 01:09:41.051663:
2024-07-21 01:09:41.051810: Epoch 998
2024-07-21 01:09:41.051903: Current learning rate: 4e-05
2024-07-21 01:11:26.669583: train_loss -0.8066
2024-07-21 01:11:26.669759: val_loss -0.3505
2024-07-21 01:11:26.669865: Pseudo dice [0.0, 0.2126, 0.4824, 0.1682, 0.0279, 0.0927, 0.1335, 0.1464, 0.0]
2024-07-21 01:11:26.669945: Epoch time: 105.62 s
2024-07-21 01:11:27.936484:
2024-07-21 01:11:27.936840: Epoch 999
2024-07-21 01:11:27.937007: Current learning rate: 2e-05
2024-07-21 01:13:13.107770: train_loss -0.7985
2024-07-21 01:13:13.107922: val_loss -0.3357
2024-07-21 01:13:13.108025: Pseudo dice [0.0, 0.1974, 0.4752, 0.1902, 0.0334, 0.1057, 0.1582, 0.1612, 0.0]
2024-07-21 01:13:13.108104: Epoch time: 105.17 s
2024-07-21 01:13:15.006795: Training done.
2024-07-21 01:13:15.025028: Using splits from existing split file: /home/GRAMES.POLYMTL.CA/p118175/data/nnunetv2/nnUNet_preprocessed/Dataset301_LumbarRootlets/splits_final.json
2024-07-21 01:13:15.025259: The split file contains 5 splits.
2024-07-21 01:13:15.025300: Desired fold for training: 0
2024-07-21 01:13:15.025336: This split has 4 training and 2 validation cases.
2024-07-21 01:13:15.025455: predicting sub-CTS10_ses-SPanat_T2w_001
2024-07-21 01:13:15.026371: sub-CTS10_ses-SPanat_T2w_001, shape torch.Size([1, 192, 372, 1023]), rank 0
2024-07-21 01:15:05.478443: predicting sub-CTS15_ses-SPpre_T2w_001
2024-07-21 01:15:05.497945: sub-CTS15_ses-SPpre_T2w_001, shape torch.Size([1, 192, 372, 1023]), rank 0
2024-07-21 01:16:55.437292: Validation complete
2024-07-21 01:16:55.437382: Mean Validation Dice:  0.12612002796626479

The predictions of the semantic model on unseen subjects are not good. The rootlets are predicted only the caudal part of the FOV.

example

Binary (all rootlets set to 1) model (Dataset302_LumbarRootlets)

training_log

2024-07-21 01:44:41.538921: Epoch 998
2024-07-21 01:44:41.539023: Current learning rate: 4e-05
2024-07-21 01:45:34.577191: train_loss -0.8987
2024-07-21 01:45:34.577362: val_loss -0.3704
2024-07-21 01:45:34.577416: Pseudo dice [0.3858]
2024-07-21 01:45:34.577475: Epoch time: 53.04 s
2024-07-21 01:45:35.760479:
2024-07-21 01:45:35.760615: Epoch 999
2024-07-21 01:45:35.760712: Current learning rate: 2e-05
2024-07-21 01:46:28.753630: train_loss -0.8903
2024-07-21 01:46:28.753798: val_loss -0.3714
2024-07-21 01:46:28.753851: Pseudo dice [0.3877]
2024-07-21 01:46:28.753914: Epoch time: 52.99 s
2024-07-21 01:46:30.514437: Training done.
2024-07-21 01:46:30.529201: Using splits from existing split file: /home/GRAMES.POLYMTL.CA/p118175/data/nnunetv2/nnUNet_preprocessed/Dataset302_LumbarRootlets/splits_final.json
2024-07-21 01:46:30.529334: The split file contains 5 splits.
2024-07-21 01:46:30.529372: Desired fold for training: 0
2024-07-21 01:46:30.529406: This split has 4 training and 2 validation cases.
2024-07-21 01:46:30.529503: predicting sub-CTS10_ses-SPanat_T2w_001
2024-07-21 01:46:30.530237: sub-CTS10_ses-SPanat_T2w_001, shape torch.Size([1, 192, 372, 1023]), rank 0
2024-07-21 01:47:25.381866: predicting sub-CTS15_ses-SPpre_T2w_001
2024-07-21 01:47:25.399373: sub-CTS15_ses-SPpre_T2w_001, shape torch.Size([1, 192, 372, 1023]), rank 0
2024-07-21 01:48:20.876391: Validation complete
2024-07-21 01:48:20.876483: Mean Validation Dice:  0.40915237629260454

The predictions of the binary model on unseen subjects are reasonable; see the example on 2 testing subjects below. The comparison shows predictions obtained using:

binary model 202 (trained on GT created using the FSLeyes; details about this model in Lumbar rootlets - first model training #48) - shown in yellow
binary model 302 (trained on GT created using the Slicer Draw Tube tool) - shown in light blue

sub-CTS03_ses-SPpre_acq-zoomit_T2w.nii.gz

sub-CTS17_ses-SPpre_acq-zoomit_T2w.nii.gz

Interestingly, some rootlets predicted by the older model (202) were not predicted by the new model (302) and vice versa.

valosekj · 2024-07-23T12:09:21Z

Next steps: tweak model training parameters to improve the model, namely:

turn off mirroring along S-I, as it might confuse the model
experiment with the patch size; a patch size smaller than the image might crop some rootlets, and the model thus loses context about their position. On the other hand, a larger patch size might be tricky to handle by the memory.

valosekj · 2024-07-24T20:28:28Z

Looking at nnUNetPlans.json for the lumbar models:

(the axis order here is: SI, AP, RL)

        "3d_fullres": {
            "batch_size": 2,
            "patch_size": [
                64,
                112,
                320
            ],
            "median_image_size_in_voxels": [
                192.0,
                372.0,
                1023.0
            ],
            "spacing": [
                0.5,
                0.29296875,
                0.29296875
            ],

The patch_size is relatively small compared to the median_image_size_in_voxels.

For example, when comparing with nnUNetPlans.json for the T2w cervical model; here patch_size is closer to median_image_size_in_voxels:

        "3d_fullres": {
            "batch_size": 2,
            "patch_size": [
                224,
                224,
                48
            ],
            "median_image_size_in_voxels": [
                320.0,
                320.0,
                64.0
            ],
            "spacing": [
                0.800000011920929,
                0.800000011920929,
                0.7999992370605469
            ],

This brings me to the idea that maybe I could try to crop the images around the SC before running the training. The contrast-agnostic model seems to work well! So, the cropping could be done easily.

valosekj · 2024-07-24T23:27:14Z

Training on the cropped images started: binary model (Dataset312_LumbarRootlets).

Crop data around the spinal cord

cd ~/code/model-spinal-rootlets/
git fetch
git checkout jv/lumbar_rootlets

cd $nnUNet_raw
cp -r Dataset302_LumbarRootlets Dataset312_LumbarRootlets
cd Dataset312_LumbarRootlets/imagesTr
bash ~/code/model-spinal-rootlets/training/crop_lumbar_data.sh

Note

Note that I had to change "overwrite_image_reader_writer" to NibabelIO in dataset.json. For some reason, SimpleITKIO read the original image dimensions!

nnUNetPlans.json

        "3d_fullres": {
            "batch_size": 2,
            "patch_size": [
                128,
                128,
                128
            ],
            "median_image_size_in_voxels": [
                192.0,
                161.0,
                166.5
            ],
            "spacing": [
                0.5,
                0.29296875,
                0.29296875
            ],

Start training:

bash ~/code/model-spinal-rootlets/training/run_training.sh 1 312 Dataset312_LumbarRootlets

valosekj · 2024-07-25T19:50:10Z

Okay, training on the cropped images is done (binary model Dataset312_LumbarRootlets).

training_log

2024-07-25 07:57:28.071485: Current learning rate: 5e-05
2024-07-25 07:58:12.384417: train_loss -0.9216
2024-07-25 07:58:12.384571: val_loss -0.3303
2024-07-25 07:58:12.384622: Pseudo dice [0.3783]
2024-07-25 07:58:12.384680: Epoch time: 44.31 s
2024-07-25 07:58:13.608948:
2024-07-25 07:58:13.609235: Epoch 998
2024-07-25 07:58:13.609570: Current learning rate: 4e-05
2024-07-25 07:58:57.742579: train_loss -0.9236
2024-07-25 07:58:57.742764: val_loss -0.3314
2024-07-25 07:58:57.742818: Pseudo dice [0.3775]
2024-07-25 07:58:57.742882: Epoch time: 44.13 s
2024-07-25 07:58:58.964358:
2024-07-25 07:58:58.964548: Epoch 999
2024-07-25 07:58:58.964673: Current learning rate: 2e-05
2024-07-25 07:59:43.275873: train_loss -0.9214
2024-07-25 07:59:43.276055: val_loss -0.3291
2024-07-25 07:59:43.276127: Pseudo dice [0.3773]
2024-07-25 07:59:43.276187: Epoch time: 44.31 s
2024-07-25 07:59:45.474522: Training done.
2024-07-25 07:59:45.489014: Using splits from existing split file: /home/GRAMES.POLYMTL.CA/p118175/data/nnunetv2/nnUNet_preprocessed/Dataset312_LumbarRootlets/splits_final.json
2024-07-25 07:59:45.489179: The split file contains 5 splits.
2024-07-25 07:59:45.489222: Desired fold for training: 0
2024-07-25 07:59:45.489260: This split has 4 training and 2 validation cases.
2024-07-25 07:59:45.489359: predicting sub-CTS10_ses-SPanat_T2w_001
2024-07-25 07:59:45.490096: sub-CTS10_ses-SPanat_T2w_001, shape torch.Size([1, 192, 163, 163]), rank 0
2024-07-25 08:00:06.051738: predicting sub-CTS15_ses-SPpre_T2w_001
2024-07-25 08:00:06.053505: sub-CTS15_ses-SPpre_T2w_001, shape torch.Size([1, 192, 160, 168]), rank 0
2024-07-25 08:00:12.400051: Validation complete
2024-07-25 08:00:12.400232: Mean Validation Dice:  0.36048916255847474

The Mean Validation Dice is 0.360, which is lower than 0.409 for the non-cropped model (binary model Dataset302_LumbarRootlets). Also, visually, the prediction on an unseen test subject is comparable (or maybe even worse) than for the non-cropped model (binary model Dataset302_LumbarRootlets):

sub-CTS03_ses-SPpre_acq-zoomit_T2w.nii.gz

light blue - model trained on non-cropped images (Dataset302_LumbarRootlets)
yellow - model trained on cropped images (Dataset312_LumbarRootlets)

Preliminary conclusion: training on images cropped around the SC does not increase the segmentation performance. In contrast, it introduces a dependency on the SC seg used for cropping (which is a disadvantage).

RaphaSchl · 2024-07-26T09:52:49Z

I've tested Dataset302 model (non-cropped images) on other subject images. When we compare 3D renderings of sub-CTS13, sub-CTS17, and sub-CTS20 model rootlet segmentations (from FSLeyes), we see continuous rootlets, as we aspire to, for sub-CTS20, and for most sub-CTS13. For CTS17, there is a lot of discontinuity.

I'm not sure how this could be improved, but thought this variability in performance was worth noting. It does not seem correlated to artefacts in the CSF as I initially thought, as sub-CTS20 is quite artefacted but has good model segmentation.

valosekj · 2024-07-26T12:16:58Z

Thank you for testing the model on additional images, @RaphaSchl! The 3D rendering is useful here!

I had a discussion with @naga-karthik, and he suggested using the 96x160x192 patch size for the cropped model instead of 128, 128, 128. I'll try it!

commands

cd $nnUNet_preprocessed
cp -r Dataset312_LumbarRootlets Dataset322_LumbarRootlets

modify manually patch_size in nnUNetPlans.json to:

        "3d_fullres": {
            "data_identifier": "nnUNetPlans_3d_fullres",
            "preprocessor_name": "DefaultPreprocessor",
            "batch_size": 2,
            "patch_size": [
                192,
                160,
                96
            ],

(192: SI, 160: AP, 96: RL)

change manually the dataset name to Dataset322_LumbarRootlets in dataset.json and nnUNetPlans.json

start training:

bash ~/code/model-spinal-rootlets/training/run_training.sh 1 322 Dataset322_LumbarRootlets

valosekj · 2024-08-02T13:02:22Z

Okay, the training of the model with 96x160x192 patch size (Dataset322_LumbarRootlets) is done.
I trained this model with the default nnUNet trainer and also with nnUNetTrainerDA5 (extensive data augmentation) and nnUNetTrainer_1000epochs_NoMirroring (no axis mirroring during the data augmentation) to see whether there is any impact on the performance.

Comparison with the non-cropped model Dataset302_LumbarRootlets on sub-CTS17_ses-SPpre_acq-zoomit_T2w.nii.gz:

screenshot

valosekj · 2024-08-06T16:21:23Z

Notes about running the models on other testing images (done by @RaphaSchl -- thank you!):

For CTS03 :

Dataset 302_NoMirroring
Dataset 322

For CTS09 (remarkable!):

Dataset 302_NoMirroring
Dataset 322 or _DA (same)
Dataset 312 or _NoMirroring (same)

For CTS13:

Dataset 302_NoMirroring
Dataset 322 or _DA (similar but both have diff imperfections)

For CTS17:

Dataset 302_NoMirroring
Dataset 322_DA

valosekj added the lumbar rootlets label Jul 19, 2024

This comment was marked as resolved.

Sign in to view

valosekj mentioned this issue Jul 26, 2024

nnUNetv2 problem with changing patch_size (hc-leipzig-7t-mp2rage) #70

Closed

valosekj mentioned this issue Aug 6, 2024

Add script for cropping lumbar data #72

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Lumbar rootlets - model training on `Draw Tube` labels #67

Lumbar rootlets - model training on `Draw Tube` labels #67

valosekj commented Jul 19, 2024 •

edited

Loading

This comment was marked as resolved.

This comment was marked as resolved.

This comment was marked as resolved.

valosekj commented Jul 23, 2024 •

edited

Loading

valosekj commented Jul 23, 2024

valosekj commented Jul 24, 2024 •

edited

Loading

valosekj commented Jul 24, 2024 •

edited

Loading

valosekj commented Jul 25, 2024 •

edited

Loading

RaphaSchl commented Jul 26, 2024

valosekj commented Jul 26, 2024 •

edited

Loading

valosekj commented Aug 2, 2024 •

edited

Loading

valosekj commented Aug 6, 2024

Lumbar rootlets - model training on Draw Tube labels #67

Lumbar rootlets - model training on Draw Tube labels #67

Comments

valosekj commented Jul 19, 2024 • edited Loading

0. Data overview

1. Preparing nnUNet folders

2. Changing label values to be consecutive (this is required by nnUNet)

3. Training

Semantic (level-specific) model: Dataset301_LumbarRootlets

Binary model (all rootlets set to 1): Dataset302_LumbarRootlets

This comment was marked as resolved.

This comment was marked as resolved.

This comment was marked as resolved.

valosekj commented Jul 23, 2024 • edited Loading

Semantic (level-specific) model (Dataset301_LumbarRootlets)

Binary (all rootlets set to 1) model (Dataset302_LumbarRootlets)

valosekj commented Jul 23, 2024

valosekj commented Jul 24, 2024 • edited Loading

valosekj commented Jul 24, 2024 • edited Loading

valosekj commented Jul 25, 2024 • edited Loading

RaphaSchl commented Jul 26, 2024

valosekj commented Jul 26, 2024 • edited Loading

valosekj commented Aug 2, 2024 • edited Loading

valosekj commented Aug 6, 2024

Lumbar rootlets - model training on `Draw Tube` labels #67

Lumbar rootlets - model training on `Draw Tube` labels #67

valosekj commented Jul 19, 2024 •

edited

Loading

Semantic (level-specific) model: `Dataset301_LumbarRootlets`

Binary model (all rootlets set to 1): `Dataset302_LumbarRootlets`

valosekj commented Jul 23, 2024 •

edited

Loading

valosekj commented Jul 24, 2024 •

edited

Loading

valosekj commented Jul 24, 2024 •

edited

Loading

valosekj commented Jul 25, 2024 •

edited

Loading

valosekj commented Jul 26, 2024 •

edited

Loading

valosekj commented Aug 2, 2024 •

edited

Loading