Variance‐aware attention U‐Net for multi‐organ segmentation
Haoneng Lin, Zongshang Li, Zefan Yang, Yi Wang
Abstract
PURPOSE: With the continuous development of deep learning based medical image segmentation technology, it is expected to attain more robust and accurate performance for more challenging tasks, such as multi-organs, small/irregular areas, and ambiguous boundary issues. METHODS: We propose a variance-aware attention U-Net to solve the problem of multi-organ segmentation. Specifically, a simple yet effective variance-based uncertainty mechanism is devised to evaluate the discrimination of each voxel via its prediction probability. The proposed variance uncertainty is further embedded into an attention architecture, which not only aggregates multi-level deep features in a global-level but also enforces the network to pay extra attention to voxels with uncertain predictions during training. RESULTS: Extensive experiments on challenging abdominal multi-organ CT dataset show that our proposed method consistently outperforms cutting-edge attention networks with respect to the evaluation metrics of Dice index (DSC), 95% Hausdorff distance (95HD) and average symmetric surface distance (ASSD). CONCLUSIONS: The proposed network provides an accurate and robust solution for multi-organ segmentation and has the potential to be used for improving other segmentation applications.