Uncertainty- and hardness-weighted loss functions for medical image segmentation

Accurate segmentation of medical images is essential for various image processing tasks and is now predominantly achieved using deep learning techniques. However, existing approaches often employ loss functions that fail to account for pixel-level differences in prediction uncertainty or hardness. This limitation frequently results in relatively large segmentation errors, particularly in object boundary regions. To address the limitation, we developed a novel class of uncertainty- / hardness-weighted loss functions by introducing two distinct pixel-wise weighting schemes: probability-guided uncertainty (PGU) and region-enhanced hardness (REH) weights. These weights, derived from the differences between network predictions and their corresponding ground truths, were designed to emphasize challenging pixels while reducing segmentation uncertainties. We validated these loss functions by integrating them with two classical neural networks, i.e., Swin Transformer based U-shape network (Swin-Unet) and V-shape network (V-Net) to segment two- and three-dimensional target objects across four different images datasets, including Retinal Fundus Glaucoma Challenge (REFUGE) dataset, Retinal Vascular Tree Analysis (RETA) dataset, optical coherence tomography (OCT) dataset, and Atria Segmentation Challenge (ASC) dataset. Extensive experiments demonstrated that our developed loss functions outperformed classical losses, such as cross-entropy (CE) and Dice losses, along with their variants, highlighting the effectiveness and generalization of the introduced weighting schemes. The source code is available at https://github.com/wmuLei/uhLoss.