Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Contour-weighted loss for class-imbalanced image segmentation

Zhhengyong Huang, Yao Sui

TL;DR

This work tackles severe data imbalance in medical image segmentation by introducing a contour-weighted compound loss (CWCD) that blends contour-weighted cross-entropy with separable dice loss. The contour maps derived from morphological erosion focus learning on object boundaries, while separable dice loss treats contour and non-contour regions separately to address intra- and inter-class imbalance. Empirical evaluation on BraTS and AMOS demonstrates that CWCD yields higher DSC than several state-of-the-art losses across multiple networks, with pronounced gains under strong imbalance, indicating improved boundary accuracy and robustness for multi-organ and tumor segmentation. The approach is practical for clinical-grade segmentation and is made available with code to facilitate adoption and further development.

Abstract

Image segmentation is critically important in almost all medical image analysis for automatic interpretations and processing. However, it is often challenging to perform image segmentation due to data imbalance between intra- and inter-class, resulting in over- or under-segmentation. Consequently, we proposed a new methodology to address the above issue, with a compact yet effective contour-weighted loss function. Our new loss function incorporates a contour-weighted cross-entropy loss and separable dice loss. The former loss extracts the contour of target regions via morphological erosion and generates a weight map for the cross-entropy criterion, whereas the latter divides the target regions into contour and non-contour components through the extracted contour map, calculates dice loss separately, and combines them to update the network. We carried out abdominal organ segmentation and brain tumor segmentation on two public datasets to assess our approach. Experimental results demonstrated that our approach offered superior segmentation, as compared to several state-of-the-art methods, while in parallel improving the robustness of those popular state-of-the-art deep models through our new loss function. The code is available at https://github.com/huangzyong/Contour-weighted-Loss-Seg.

Contour-weighted loss for class-imbalanced image segmentation

TL;DR

This work tackles severe data imbalance in medical image segmentation by introducing a contour-weighted compound loss (CWCD) that blends contour-weighted cross-entropy with separable dice loss. The contour maps derived from morphological erosion focus learning on object boundaries, while separable dice loss treats contour and non-contour regions separately to address intra- and inter-class imbalance. Empirical evaluation on BraTS and AMOS demonstrates that CWCD yields higher DSC than several state-of-the-art losses across multiple networks, with pronounced gains under strong imbalance, indicating improved boundary accuracy and robustness for multi-organ and tumor segmentation. The approach is practical for clinical-grade segmentation and is made available with code to facilitate adoption and further development.

Abstract

Image segmentation is critically important in almost all medical image analysis for automatic interpretations and processing. However, it is often challenging to perform image segmentation due to data imbalance between intra- and inter-class, resulting in over- or under-segmentation. Consequently, we proposed a new methodology to address the above issue, with a compact yet effective contour-weighted loss function. Our new loss function incorporates a contour-weighted cross-entropy loss and separable dice loss. The former loss extracts the contour of target regions via morphological erosion and generates a weight map for the cross-entropy criterion, whereas the latter divides the target regions into contour and non-contour components through the extracted contour map, calculates dice loss separately, and combines them to update the network. We carried out abdominal organ segmentation and brain tumor segmentation on two public datasets to assess our approach. Experimental results demonstrated that our approach offered superior segmentation, as compared to several state-of-the-art methods, while in parallel improving the robustness of those popular state-of-the-art deep models through our new loss function. The code is available at https://github.com/huangzyong/Contour-weighted-Loss-Seg.
Paper Structure (11 sections, 12 equations, 3 figures, 3 tables)

This paper contains 11 sections, 12 equations, 3 figures, 3 tables.

Table of Contents

  1. INTRODUCTION
  2. METHODS
  3. Contour Extraction
  4. Separable Dice Loss
  5. Contour-Weighted Cross Entropy Loss
  6. Implementation Details and Datasets
  7. RESULTS
  8. Assessments of Different Loss Functions
  9. Assessments with Different Networks
  10. Ablation Study
  11. DISCUSSION

Figures (3)

  • Figure 1: Illustration of our proposed contour-weighted map. We calculate the contours $C$ from the difference between the mask $G$ and its morphological erosion correspondence. The obtained contours $C = G - E$ are used to weight the cross-entropy loss and compute the separable dice loss. The second and third rows show the contour calculation process on two representative images, respectively.
  • Figure 2: Qualitative comparison of different loss functions based on U-Net on the BraTS dataset. The whole tumor (WT) encompasses a union of red, yellow, and green regions. The tumor core (TC) includes the union of red and yellow regions. The enhancing tumor core (ET) denotes the green region.
  • Figure 3: Qualitative comparison of different loss functions based on U-Net on representative slices from the AMOS dataset.