nnUNet training for MS lesion seg

[plbenveniste]

In this issue, I detail the steps to training several nnUNet to automatically segment MS lesions in the SC. More specifically I trained:

• single class nnUNet: at inference : 1 input (image), 1 output (1 class = lesions)
• region-based nnUNet: at inference : 1 input (image), 1 output (2 classes = sc and lesions).
• multichannel nnUNet: at inference : 2 inputs (image and sc seg), 1 output (1 class = lesions)

The comments below describe the investigation and the results.

[plbenveniste]

When running nnUNetv2_plan_and_preprocess -d 101 --verify_dataset_integrity -c 3d_fullres, I saw that some label files from sct-testing-large contain lesion segmentation with values 1 and 2.

Example of subject sub-rennesMS002_acq-sagcerv_T2w.nii.gz

I couldn't find an explanation for this in neither the .json file or the dataset README.txt

To solve this, the conversion script must include a binarization of all label file which include more than 1 label and label file which contain values greater than 1.

[plbenveniste]

To Do :

• [x] Add the canproco exclude file when building the dataset json file

[plbenveniste]

Here we report the performances of the different model trained on data from Canproco, basel-mp2rage, sct-testing-largeand bavaria-quebec-spine-ms. The dataset used for training, testing and inference are stored in ~/duke/temp/plben/ms-lesion-agnostic/nnunet-json-datasets. The nnUNet models which were trained are stored in ~/duke/temp/plben/ms-lesion-agnostic/nnUNet-trainings.

Here are the performances achieved by various nnUNet configurations:

	Training Dice	Mean Validation Dice	Comment
3D singleclass (fold 1)	0.5823	0.43379	Training could have been longer
2D singleclass 2000 epochs (fold 1)	0.6476	0.473	Converged before 1000
3D singleclass 2000 epochs (fold 1)	0.600	0.45436	Started overfiting
3D singleclass no smooth (fold 1)	0.5835	0.4369	Could try longer
3D regionbased (fold 1)	[0.9204,0.601]*	mean = 0.69360**	Could be trained longer
3D multichannel (fold 1)	0.3088	0.4421	? training failure
3D multichannel no smooth (fold 2)	0.5569	0.4179	Should try training longer

*For region-based training, the Dice score correspond respectively to SC seg and lesion seg **We didn't have access to the validation's two different dice score

Next steps: -> Moving away from nnUNet models to use monai trained models (Unet, AttentionUnet, SwinUnetR, Diffusion models) -> Why ? : Monai allows for more tailored data augmentation, more control over the loss functions and more fun !