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 scores achieves segmentations that are more semantically meaningful for brain tumour segmentation.

Accepted as an oral presentation at the MICCAI 2017 Brain Lesion (BrainLes) Workshop

References in corpus (4)

• Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations
• NiftyNet: a deep-learning platform for medical imaging
• On the Compactness, Efficiency, and Representation of 3D Convolutional Networks: Brain Parcellation as a Pretext Task
• Scalable multimodal convolutional networks for brain tumour segmentation

Cited by in corpus (11)

• Automatic Brain Tumor Segmentation using Convolutional Neural Networks with Test-Time Augmentation
• Generalized Wasserstein Dice Score, Distributionally Robust Deep Learning, and Ranger for brain tumor segmentation: BraTS 2020 challenge
• Label-set Loss Functions for Partial Supervision: Application to Fetal Brain 3D MRI Parcellation
• Distributionally Robust Segmentation of Abnormal Fetal Brain 3D MRI
• blob loss: instance imbalance aware loss functions for semantic segmentation
• Generalized Wasserstein Dice Loss, Test-time Augmentation, and Transformers for the BraTS 2021 challenge
• AMD-HookNet for Glacier Front Segmentation
• Distributionally Robust Deep Learning using Hardness Weighted Sampling
• Comparing Machine Learning based Segmentation Models on Jet Fire Radiation Zones
• A Mosquito Pick-and-Place System for PfSPZ-based Malaria Vaccine Production
• Tree-based Semantic Losses: Application to Sparsely-supervised Large Multi-class Hyperspectral Segmentation