NifTK / NiftyNet

[unmaintained] An open-source convolutional neural networks platform for research in medical image analysis and image-guided therapy
http://niftynet.io
Apache License 2.0
1.36k stars 403 forks source link

Inference produces dotted segmentation output [deepmedic] #417

Open andreasstefani opened 5 years ago

andreasstefani commented 5 years ago

Hello,

I'm trying to train the 'deepmedic' model with the BraTS 2018 dataset.

Data: The BraTS 2018 dataset is very similar to the BraTS 2017 dataset. The only difference is that there are 3 tumour sub-regions (BraTS 2018) instead of 4 tumour sub-regions (BraTS2017). Therefore, including the background, the number of classes is 4 for BraTS 2017. The BraTS 2018 images are from size 240x240x155. No other pre-processing was done despite the image normalization to zero-mean and unit variance with the setting "whitening = True" in the config.

Training: First, I trained the model using all cases from the BraTS 2018 dataset (Figure 1). Then, I decided to use just 10 BraTS 2018 cases: 5 HGG, and 5 LGG (Figure 2). As one can see in the config file according to the fraction parameters, the 10 cases are split in: 7 cases for training, 2 cases for validation, and 1 case for inference.

After training two 'deepmedic' models with the settings in the attached config file and the two different data sets, the following segmentations will be produced. The coloured area is the ground truth segmentation. In Figure 1, the similarity between both segmentations seems not too bad besides the fact of all the other grey pixels.

Figure 1: all cases, and settings exactly as in the attached config file. snapshot0010 image

Figure 2: just 10 cases, and settings exactly as in the attached config file. snapshot0012 image

I read a lot through all the issues related to this problem, tried different settings etc. but wasn't able to train the model properly. For example, I tried 30000 iterations and then the segmentation output was entirely white (a plain white output image). The Dice loss was nearly the same as in the training/validation curves above (Figure 3). I am not sure but I guess this has to do with the number of classes I am trying to segment and is normal?

Figure 3: Training and validation curve for 30k iterations with Dice loss. image

At the moment I have no idea how to fix this problem and train the model properly. I hope you can help me out, because I am using NiftyNet for my masters dissertation and also try to implement a novel solution in parallel.

Thank you in advance.

Best, Andreas

dm_config.txt

JunMa11 commented 5 years ago

So weird. I do not have experience with deepmedic, but I try U-Net with brats2018, it works well. Do you have similar problem with U-Net in niftynet.

Zach-ER commented 5 years ago

Hi there, I had a look at the config file and I think the most likely culprit is the 'window sampling' parameter. You have chosen uniform, which means that the 9x9x9 label window is equally likely to be anywhere in the volume. This means that you will see relatively few foreground patches.

In the DeepMedic MIA paper, they say

Specifically, we build the training batches by extracting segments from the training images with 50% probability being centred on a foreground or background voxel, alleviating class-imbalance.

What I'd try is making a foreground mask (binary mask that is >0 for any non-background label) and using that as an input to use balanced window sampling.

I would also use the Dice + Cross Entropy loss instead of just Dice: the cross-entropy bit does well at pushing down the probabilities of foreground in empty regions.

Please update when you've tried these things!

josefmalmstrom commented 4 years ago

So weird. I do not have experience with deepmedic, but I try U-Net with brats2018, it works well. Do you have similar problem with U-Net in niftynet.

@JunMa11 Any chance you could share your config for training U-Net on Brats2018?

Thanks in advance!