We gratefully acknowledge support from
the Simons Foundation and member institutions.
Full-text links:


Current browse context:


Change to browse by:

References & Citations


(what is this?)
CiteULike logo BibSonomy logo Mendeley logo del.icio.us logo Digg logo Reddit logo ScienceWISE logo

Electrical Engineering and Systems Science > Image and Video Processing

Title: Unified Focal loss: Generalising Dice and cross entropy-based losses to handle class imbalanced medical image segmentation

Abstract: Automatic segmentation methods are an important advancement in medical image analysis. Machine learning techniques, and deep neural networks in particular, are the state-of-the-art for most medical image segmentation tasks. Issues with class imbalance pose a significant challenge in medical datasets, with lesions often occupying a considerably smaller volume relative to the background. Loss functions used in the training of deep learning algorithms differ in their robustness to class imbalance, with direct consequences for model convergence. The most commonly used loss functions for segmentation are based on either the cross entropy loss, Dice loss or a combination of the two. We propose the Unified Focal loss, a new hierarchical framework that generalises Dice and cross entropy-based losses for handling class imbalance. We evaluate our proposed loss function on five publicly available, class imbalanced medical imaging datasets: CVC-ClinicDB, Digital Retinal Images for Vessel Extraction (DRIVE), Breast Ultrasound 2017 (BUS2017), Brain Tumour Segmentation 2020 (BraTS20) and Kidney Tumour Segmentation 2019 (KiTS19). We compare our loss function performance against six Dice or cross entropy-based loss functions, across 2D binary, 3D binary and 3D multiclass segmentation tasks, demonstrating that our proposed loss function is robust to class imbalance and consistently outperforms the other loss functions. Source code is available at: this https URL
Subjects: Image and Video Processing (eess.IV); Computer Vision and Pattern Recognition (cs.CV); Machine Learning (cs.LG)
ACM classes: I.4.6; J.3
Cite as: arXiv:2102.04525 [eess.IV]
  (or arXiv:2102.04525v4 [eess.IV] for this version)

Submission history

From: Michael Yeung [view email]
[v1] Mon, 8 Feb 2021 20:47:38 GMT (4154kb,D)
[v2] Tue, 18 May 2021 08:49:30 GMT (7087kb,D)
[v3] Mon, 24 May 2021 14:06:06 GMT (2828kb,D)
[v4] Wed, 24 Nov 2021 15:12:27 GMT (8137kb,D)

Link back to: arXiv, form interface, contact.