Final nnU-Net model training for zurich-mouse

[plbenveniste]

Here is the followed strategy for the final model training for the zurich-mouse dataset for white and grey matter segmentation.

• Extract annotated 2D slices from the dataset
• Train a 2D nnU-Net model on the annotated slices
• Predict segmentation on full 3D volumes
• Train a 3D nnU-Net model on these segmentations and run inference on entire dataset
• Select some 3D predictions, improve segmentation (by removing sections which are not annotated) and add to initial training dataset
• Train a 3D nnU-Net on the training dataset (slices and selected 3D nnU-Net predictions)
• Test on entire zurich-mouse dataset.
• Save model and write documentation on how to use it

Related to #32 #37

[plbenveniste]

Data pre-processed !
Training dataset :

• 192 files ( 31 3D volume from 3D nnU-Net prediction and 161 annotated slices).
• The 3D volumes are of different sizes because I removed the slices from both the image and the mask when segmentation was of poor quality.

Conversion to nnU-Net format : ok
nnU-Net plan and pre-process : ok
Currently training on fold 0,1,2,3 and 4.

[plbenveniste]

Model trained on the final training dataset.
Average Dice over all folds: 0.91 🎉
Inference currently running on the entire dataset.

To do:

• convert back to the BIDS format (and add python script to the github)
• update github README.md
• upload segmentations in the git-annex : instead of manual-masks use nnUNet_masks
• store the trained model

[plbenveniste]

The results were converted back to BIDS format.
However, observing the results (such as sub_mouse1_chunk-4) showed that some regions remained unlabeled.
I think, it's because those regions are a bit different and have never been labeled before.

[plbenveniste]

Comparison with U-Net results only on problematic segmentations from the nnUNet:

• For sub-mouse1_chunk-4
  
• For sub-mouse10_chunk-2
  
• For sub-mouse11_chunk-4
  There is zero segmentation in both case
• For sub-mouse12_chunk-3
  
• For sub-mouse136_chunk-2
  

Conclusion:
On images on which nnUNet performs poorly, it still performs better than the UNet used in ivadomed. It is more exhaustive in terms of labelling and also performs better in seeing the 3D aspect of the volume.

On a non-problematic example:
For sub-mouse11_chunk-2

Better labelling overall from the nnUNet. This is particularly visible on top and bottom slices of the volume.

[jcohenadad]

Better labelling overall from the nnUNet.

Agreed! I think we can go ahead and publish a release, after making sure there is a clear procedure for running the inference on a single image (ie: update test.py and update README)

[plbenveniste]

Selection of the model :
We want to store the model and at the same time we don't want to take up too much space on Github and we don't want users to have to download a 2GB file in order to use our model. Therefore, we went through the following reasoning process.
First every fold has 2 model stored: thecheckpoint_final.pth model and the checkpoint_best.pth. However, we see that in every case the model follows a continuously growing dice score. Therefore, the ckeckpoint_best.pth is the same as the
checkpoint_final.pth. However, we decided to use the checkpoint_best.pth as it is the one to use in case during training the dice scores diminishes at some point (rule of thumb here).
Also, since every fold perform relatively similar, which means that our dataset is homogeneous, we only retain one fold which has the highest dice score: fold 4.

Fold selection :

• fold 0 : 0.9135
• fold 1 : 0.9083
• fold 2 : 0.9109
• fold 3 : 0.9132
• fold 4 : 0.9173

We also remove files from the model folder which are not useful for inference (validation images, data description ...), except files training_log.txtand progress.png which can be interesting for users to see how the model was trained and how the training evolved throughout time.

Furthermore, we also decided to keep the post-processing files in the folder in case the person want to perform post-processing on the results of the inference.

By doing this, we went from 5.6 GB to 280 MB.

[jcohenadad]

Also, since every fold perform relatively similar, which means that our dataset is homogeneous, we only retain one fold which has the highest dice score: fold 4.

We know there is generally an advantage in doing ensembling in terms of segmentation performance, but of course if inference is too slow, then we should revisit.