Generalised Wasserstein Dice Score for Imbalanced Multi-class Segmentation Using Holistic Convolutional Networks

Abstract

The Dice score is widely used for binary segmentation due to its robustness to class imbalance. Soft generalisations of the Dice score allow it to be used as a loss function for training convolutional neural networks (CNN). Although CNNs trained using mean-class Dice score achieve state-of-the-art results on multi-class segmentation, this loss function does neither take advantage of inter-class relationships nor multi-scale information. We argue that an improved loss function should balance misclassifications to favour predictions that are semantically meaningful. This paper investigates these issues in the context of multi-class brain tumour segmentation. Our contribution is threefold. 1) We propose a semantically-informed generalisation of the Dice score for multi-class segmentation based on the Wasserstein distance on the probabilistic label space. 2) We propose a holistic CNN that embeds spatial information at multiple scales with deep supervision. 3) We show that the joint use of holistic CNNs and generalised Wasserstein Dice score achieves segmentations that are more semantically meaningful for brain tumour segmentation.

Automatic Brain Tumor Segmentation Using Convolutional Neural Networks with Test-Time Augmentation

“…As a pre-processing, each image was normalized by the mean value and standard deviation. The Dice loss function [25,9] was used for training. At test time, the augmented prediction number was set to N = 20 for all the network structures.…”

Automatic Brain Tumor Segmentation Using Cascaded Anisotropic Convolutional Neural Networks

“…C l equals to 2 in our method. A combination of predictions from multiple scales has also been used in [31,9]. Multi-view Fusion.…”

“…HighRes3DNet [20] proposes a compact end-to-end 3D CNN structure that maintains high-resolution multi-scale features with dilated convolution and residual connection [29,6]. Other works also propose using fully convolutional networks [13,8], incorporating large visual contexts by employing a mixture of convolution and downsampling operations [16,12], and handling imbalanced training data by designing new loss functions [9,27] and sampling strategies [27].…”

Deep convolutional neural networks for brain image analysis on magnetic resonance imaging: a review

“…For instance, in tumour or MS lesion segmentation, obtaining good generalisation is a challenge as most of the lesions are smaller than the entire volume. Two-phase training [24], careful patch selection [22,25,132] and loss functions [21,159,160] are among the proposed strategies to overcome this problem. Although we observed the success of CNNs; their full capacity has not yet been fully leveraged in brain MRI analysis.…”