  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:
https://github.com/mlyg/unified-focal-loss
