By Team 22/7, Chaitra V. Hegde | Advisor : Narges Razavian
Background and Problem Definition
- Accurate brain structural segmentation is central to nearly all neuroimaging analyses.
- Freesurfer and other traditional tools take 2-4 hrs to segment a volume.
- Freesurfer also has systematic biases.
- Deep Learning based models can increase the efficiency and quality of segmentation.
Contribution
- Novel training methodology of scheduling dice and cross entropy loss to optimally train segmentation models.
- Revealed systematic biases in Freesurfer tool and built model free of those biases.
- Inference time < 20 secs / MRI volume
Model Architecture and Data
- Vanilla U-Net architecture
- Pre-trained the model using Human Connectome Project (HCP) Data: Healthy Individuals (# = 1113), Ground Truth = Freesurfer Segmentation.
- Fine-tuned using manually annotated MICCAI challenge train dataset (train = 15 subjects and test = 20 subjects).
Training Methodology
- Xavier initialization for model parameters.
- Augmented data using Gaussian blurring, contrast adjustment, rotation and translation.
- Experimented with different loss function schedule:
- Only dice loss for the entire training period.
- Only cross entropy loss (CEL) for the entire training period.
- Both CEL and dice loss for the entire training period.
- CEL and dice loss for 30 epochs then switched to pure dice loss (Loss scheduling).
- Pre-trained on HCP with Freesurfer labels as the ground truth.
- Fine-tuned on MICCAI train data.
Results
- Evaluated on MICCAI Test Data:
| Name | Loss Function | Pre-Trained | Fine-Tuned |
|---|---|---|---|
| U-Net | Fixed (w-Dice Loss) | Did not converge | |
| U-Net | Fixed (w-Cel) | 0.7602 ± 0.085 | |
| U-Net | Fixed (w-Cel + w-dice loss) | 0.7819 ± 0.072 | |
| U-Net | Loss scheduling | 0.8049 ± 0.067 | 0.885± 0.042 |
| QuickNAT | Fixed (w-Cel + w-dice loss + Boundary Loss) | 0.798 ± 0.097 | 0.0901 ± 0.045 |
| U-Net | Fixed (w-Cel + w-dice loss + Boundary Loss) | 0.681 ± 0.193 | 0.857 ± 0.079 |
- Vanilla U-Net trained with loss scheduling performed better than the state-of-the-art. QuickNAT for the pre-trained model.
- A sample segment is visualized below for all the different pre-trained models.
- Visualization shows the loss scheduling helps the model to get close to Ground Truth (i.e. Freesurfer)
- Coordinated U-net’s performance was same as vanilla U-net
Systematic Bias in Freesurfer
- Difference in performance of pre-trained and fine-tuned model can be seen below:
- Reveals systematic bias of Freesurfer.
- Large bias for Lat-Ventricle, Pallidum, Hippocampus and Amygdala.
Model Free of Systematic Bias
- Visualization of fine-tuned and pre-trained model on the MICCAI test data.
Model)
- Pre-trained model is very liberal for Hippocampus. This implies Freesurfer is also liberal while segmenting it.
- Freesurfer being the tool based on template matching mechanism, doesn’t emphasize on the edges- hence, rough edges and over predictions. Where as our model, segments the tissue based on edge detection
Conclusion and Future Work
- Loss function and training methodology are as important as the model architecture.
- Deep learning models can overcome biases and inference time is low.
- Future work: Experiment with other data sets and segmentation models to conclusively prove that loss scheduling is a helpful training methodology.
- Future work: Build an open source tool and
- release models for ~190 segments
References
- Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. CoRR, abs/1505.04597, 2015.
- Abhijit Guha Roy, Sailesh Conjeti, Nassir Navab, and Christian Wachinger. Quicknat: Segmenting MRI neuroanatomy in 20 seconds. CoRR, abs/1801.04161,
2018 - Bennett A. Landman and Simon K. Warfield. Miccai 2012 workshop on multi-atlas labeling (volume 2). CreateSpace Independent Publishing Platform(EDS), 2012