AFFUNet: adaptive feature fusion Transformer U-Net with joint loss function for medical image segmentation.

Authors

Abstract

High-precision medical image segmentation is crucial for enabling computer-aided accurate diagnosis and personalized treatment planning. Although Transformer-based methods have shown promising results, they still face significant challenges. These include a fixed feature fusion strategy that cannot adapt to dynamic changes in feature importance, the standard Dice loss that treats all pixels equally and struggles to focus on difficult samples, and inadequate boundary precision optimization, which limits the clinical applicability of segmentation results. To address these challenges, this study proposes Adaptive Feature Fusion U-Net (AFFUNet), an adaptive feature fusion (AFF) Transformer U-Net with a joint loss function, aiming to improve segmentation accuracy through dynamic feature fusion, hard sample reweighting, and explicit boundary optimization. Based on the Cross-Shaped Window Transformer UNet (CSWin-UNet) architecture, we propose three modifications to address the challenges discussed above: (I) an AFF module that replaces standard skip connection concatenation to dynamically learn fusion weights for encoder and decoder features; (II) focal Dice loss, which adjusts the Dice loss based on prediction confidence to better prioritize difficult samples; (III) boundary-aware loss, which explicitly enhances boundary precision through a gradient-based boundary extraction technique. We evaluated our method on two benchmark datasets: the Synapse multi-organ computed tomography (CT) dataset and the Automated Cardiac Diagnosis Challenge (ACDC) cardiac magnetic resonance imaging (MRI) dataset. On the Synapse dataset, our method achieved an average Dice coefficient of 81.21% and an average 95th percentile Hausdorff distance (HD95) of 18.41 mm, representing a 0.09% improvement in Dice and a 0.45 mm reduction in HD95 compared to the state-of-the-art CSWin-UNet. On the ACDC dataset, our method achieved a mean Dice of 89.77%, consistently outperforming CSWin-UNet (88.44%). Ablation studies confirmed that each component contributes to the overall performance, particularly in small organs and boundary regions. Notably, the added AFF module introduces negligible computational overhead: compared to CSWin-UNet, the parameters and floating point operations per second (FLOPs) remained virtually unchanged (-0.02 M, -0.036 G). Thus, our method achieves improved segmentation accuracy while maintaining practical efficiency. This work integrates AFF, focal Dice loss, and boundary-aware loss into a collaborative framework that works across the feature, objective, and boundary levels. The resulting end-to-end mechanism achieves improved performance on standard benchmarks while maintaining practical efficiency, warranting further evaluation in clinical settings.

Tags

Topics