Partial Decoder Attention Network with Contour-weighted Loss Function for Data-Imbalance Medical Image Segmentation

TL;DR

This work tackles data imbalance in medical image segmentation by introducing contour-weighted learning and a contour-weighted compound loss (CWCD) that emphasizes boundary regions while balancing contour vs non-contour areas. A lightweight encoder-decoder network, PDANet, integrates a partial decoder, receptive field blocks, and channel-wise attention to preserve fine boundary details with limited parameters. Contour extraction via erosion creates contour maps used to weight cross-entropy, while a separable Dice loss balances interior and boundary regions, together forming a model-independent loss framework. Across AMOS, BraTS, and PENGWIN, CWCD and PDANet yield consistent accuracy and robustness gains over nine baselines, demonstrating practical potential for handling data imbalance in diverse medical segmentation tasks.

Abstract

Image segmentation is pivotal in medical image analysis, facilitating clinical diagnosis, treatment planning, and disease evaluation. Deep learning has significantly advanced automatic segmentation methodologies by providing superior modeling capability for complex structures and fine-grained anatomical regions. However, medical images often suffer from data imbalance issues, such as large volume disparities among organs or tissues, and uneven sample distributions across different anatomical structures. This imbalance tends to bias the model toward larger organs or more frequently represented structures, while overlooking smaller or less represented structures, thereby affecting the segmentation accuracy and robustness. To address these challenges, we proposed a novel contour-weighted segmentation approach, which improves the model's capability to represent small and underrepresented structures. We developed PDANet, a lightweight and efficient segmentation network based on a partial decoder mechanism. We evaluated our method using three prominent public datasets. The experimental results show that our methodology excelled in three distinct tasks: segmenting multiple abdominal organs, brain tumors, and pelvic bone fragments with injuries. It consistently outperformed nine state-of-the-art methods. Moreover, the proposed contour-weighted strategy improved segmentation for other comparison methods across the three datasets, yielding average enhancements in Dice scores of 2.32%, 1.67%, and 3.60%, respectively. These results demonstrate that our contour-weighted segmentation method surpassed current leading approaches in both accuracy and robustness. As a model-independent strategy, it can seamlessly fit various segmentation frameworks, enhancing their performance. This flexibility highlighted its practical importance and potential for broad use in medical image analysis.

Figures (9)

• 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 $E$ ($C = G-E$). The second row visualizes the contour calculation process.
• Figure 2: Overview of the proposed contour-weighted segmentation approach. (a) The main framework is an encoder-decoder structure with a partial decoder module (PDM) inserted in the middle. The loss function consists of two components: a contour-weighted cross-entropy loss ($\mathcal{L}_{SDL}$) and a separable dice loss ($\mathcal{L}_{CWCE}$). Where both the segmented contour and the erosion mask are obtained through Eq. (\ref{['equ: contour extraction']}). (b) The channel-wise attention module (CWCA) is used to fuse the feature information from the encoder and decoder adaptively. (c) The partial decoder module (PDM) integrates multi-scale features through a U-Net-like structure and a holistic attention mechanism. (d) The receptive field block (RFB) employs different kernel sizes to capture multi-grained feature information.
• Figure 3: Data distribution bias. (a) The mean volume of different organs in the AMOS dataset. (b) The mean volume of different tissues in the BraTS dataset. (c) The mean volume of different fragments in the PENGWIN dataset. (d) The number of different segmentation fragments in the PENGWIN dataset. Orange lines indicate the sacrum, green indicates the left hipbone, and red indicates the right hipbone.
• Figure 4: Influence of different parameters on the segmentation performance of 3D U-Net on the AMOS dataset. In the left column, $\alpha$ is set to 0. In the middle column, $\alpha$ is set to 0.5, and $\lambda$ is set to 2. In the right column, $\beta$ is set to 0.5, and $\lambda$ is set to 2.
• Figure 5: Qualitative comparison of different loss functions based on 3D 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 (ET) denotes the yellow region.