Source-linked AI summary
D-UNet: a dimension-fusion U shape network for chronic stroke lesion segmentation
Yongjin Zhou, Weijian Huang, Pei Dong, Yong Xia, Shanshan Wang
TL;DR
Chronic stroke lesion segmentation is needed for clinical diagnosis, surgical planning, prognosis, and recovery-related analysis, but 2D CNNs omit volumetric context while 3D CNNs demand substantial resources. D-UNet fuses 2D and 3D convolutions during encoding and adds Enhanced Mixing Loss; on ATLAS, it achieved the best reported quality metrics among the compared methods.
Problem
Chronic stroke lesion segmentation must capture volumetric information accurately, but 2D CNNs ignore 3D context and 3D CNNs have high computational and storage requirements.
Method
D-UNet performs 3D and 2D feature extraction on consecutive MRI slices during encoding, fuses their features, and uses Enhanced Mixing Loss combining Dice and Focal losses.
Results
DSC = 0.5349±0.2763 and Precision = 0.6331±0.2958, the best performance on two quality metrics among three compared state-of-the-art methods.
Takeaways & Limitations
The method utilizes 3D context information with low hardware requirements while providing faster and smoother convergence through its loss function.
Abstract
from arXiv · showhide
Assessing the location and extent of lesions caused by chronic stroke is critical for medical diagnosis, surgical planning, and prognosis. In recent years, with the rapid development of 2D and 3D convolutional neural networks (CNN), the encoder-decoder structure has shown great potential in the field of medical image segmentation. However, the 2D CNN ignores the 3D information of medical images, while the 3D CNN suffers from high computational resource demands. This paper proposes a new architecture called dimension-fusion-UNet (D-UNet), which combines 2D and 3D convolution innovatively in the encoding stage. The proposed architecture achieves a better segmentation performance than 2D networks, while requiring significantly less computation time in comparison to 3D networks. Furthermore, to alleviate the data imbalance issue between positive and negative samples for the network training, we propose a new loss function called Enhance Mixing Loss (EML). This function adds a weighted focal coefficient and combines two traditional loss functions. The proposed method has been tested on the ATLAS dataset and compared to three state-of-the-art methods. The results demonstrate that the proposed method achieves the best quality performance in terms of DSC = 0.5349+0.2763 and precision = 0.6331+0.295).
1 INTRODUCTION
Chronic stroke lesion segmentation is important for outcome prediction but difficult because lesions vary in location, shape, volume, and boundary clarity. D-UNet combines 2D and 3D processing, while Enhanced Mixing Loss addresses training imbalance and convergence.
- Accurate lesion segmentation supports analyses of lesion size, location, and overlap that can predict long-term stroke recovery and rehabilitation.
- Lesion segmentation is challenging because MRI artifacts, stroke-subtype variability, lesion volumes from hundreds to tens of thousands of cubic millimeters, and unclear boundaries complicate labeling.
- 2D CNNs process slices independently and can produce discontinuous predictions by ignoring spatial characteristics, whereas 3D CNNs require substantial computation and storage.
- D-UNet extracts 3D information from consecutive MRI slices during early encoding and fuses the resulting features with a 2D UNet structure.
- Enhanced Mixing Loss combines Dice and Focal loss advantages while enhancing Dice-loss gradient propagation for faster and smoother convergence.
2 RELATED WORKS
Related work spans hand-crafted feature methods, 2D and 3D deep-learning approaches, and traditional segmentation losses. The reviewed methods illustrate trade-offs between feature design, spatial context, continuity, and computational demands.
- Hand-crafted approaches use expert-designed features or classifiers for automatic disease-area and lesion segmentation.
- Deep learning methods reduce reliance on manually engineered features and include specialized 2D network and cascaded MRI segmentation frameworks.
- Slice-wise 2D segmentation can neglect volumetric context, causing predictions to lose continuity when reconstructed across slices.
- 3D CNNs extract volumetric MRI information through architectures such as two-path networks and 3D FC-DenseNet.
- Binary cross-entropy, Focal, and Dice losses represent established options for segmentation and address prediction error, class imbalance, or overlap quality.
3 METHODS
The method combines early 2D and 3D feature extraction within an improved U-Net, then fuses the representations for efficient dimensional integration. It also introduces Enhanced Mixing Loss to address foreground-background imbalance and improve convergence.
- D-UNet framework: D-UNet performs 2D and 3D downsampling in parallel and combines their results through dimension-transform blocks.The fusion is restricted to the early coding stage because deeper layers require substantially more trainable parameters.
- Dimension transformation: The dimension-transform block reduces 3D feature dimensionality with a 1×1×1 convolution, squeezes depth, and applies a 2D 3×3 convolution to match channels.The resulting 3D representation can then be integrated with the 2D branch.
- Dimension fusion: SE weighting is applied separately to the transformed 3D and 2D feature maps before their channel-weighted outputs are added.The SE block weights feature channels to enhance the fusion of 3D and 2D information.
- Loss functions: Focal loss adds a modulating factor to binary cross-entropy, reducing contributions from easy samples during segmentation training.The formulation uses ground-truth labels, prediction probabilities, class-pixel counts, and adjustable modulation factors α and γ.
- Loss functions: Enhanced Mixing Loss combines log-transformed Dice loss with scaled focal loss so both components contribute comparably to gradient propagation.The proposed loss is intended to increase convergence speed and produce a smoother convergence curve while addressing foreground-background voxel imbalance.
- Implementation details: The implementation crops transverse MRI slices, resizes them to 192 × 192, and applies augmentation including translation, scaling, and horizontal flipping.The loss parameters α, γ, and δ are set to 1.1, 0.48, and 1, respectively, with SGD using a 1e-6 learning rate.
4 EXPERIMENTAL RESULTS AND DISCUS-SIONS
Experiments on ATLAS compare D-UNet with 2D and 3D UNet variants, established frameworks, and alternative losses and fusion blocks. D-UNet improves segmentation while using substantially less training time than 3D UNet, and its EML and SE-based fusion show favorable convergence and fusion behavior.
- Experimental setup: ATLAS contains 229 chronic-stroke MRI T1 cases, with 183 cases used for training in the reported experiments.Each case is sized 233×197×189 with physical resolution 0.9×0.9×3.0mm3.
- Comparison with UNet variants: 3.83% higher metrics with only a 2% parameter increase: D-UNet improved over a single 2D structure while adding 3D information.The authors report that the dimension-transform block effectively utilizes 3D information compared with a simple 2D network.
- Comparison with other methods: The proposed method achieved the best upper edge and median in DSC box plots, despite small-lesion cases causing all methods to reach a lower edge of zero.The reported qualitative results also describe correct detection of a very small lesion and closer agreement with ground-truth boundaries.
- Comparison with other methods: 0.7231 DSC: the proposed method ranked first among the methods compared in Table 3, 0.0383 above UNet.DSC, recall, and precision use per-case mean ± standard deviation, while DSC (global) uses voxel-level calculation.
- Loss validity: EML converged faster than the compared traditional losses, while achieving the highest precision despite slightly lower DSC and recall than DL.Relative to DL, EML was 0.0115 lower in DSC and 0.057±0.002 lower in recall; the authors link its highest precision to fewer false positives.
- Dimension-fusion blocks: All tested dimension-fusion structures outperformed UNet without the fusion block, while deeper 3D fusion generally improved results until Add-123, where DSC decreased.SE-based fusion produced higher DSC than direct fusion for matching layer configurations.
5 CONCLUSION
The proposed end-to-end method uses 3D context with low hardware requirements and introduces a loss function designed for faster, smoother convergence. It achieved the best performance among three compared state-of-the-art methods on DSC and Precision.
- DSC = 0.5349±0.2763 and Precision = 0.6331±0.2958 were the best results among three compared state-of-the-art methods.
- The end-to-end training method effectively utilizes 3D context information while maintaining low hardware requirements.
- The proposed loss function is intended to support faster and smoother convergence during training.
- Future work includes increasing punishment for extremely difficult samples and validating generalization on a larger clinical dataset.