Uncertainty-Aware Hierarchical Aggregation Network for Medical Image Segmentation

• Published in: IEEE transactions on circuits and systems for video technology Vol. 34; no. 8; pp. 7440 - 7453
• Main Authors: Zhou, Tao, Zhou, Yi, Li, Guangyu, Chen, Geng, Shen, Jianbing
• Format: Journal Article
• Language: English
• Published: New York IEEE 01.08.2024; The Institute of Electrical and Electronics Engineers, Inc. (IEEE)
• Subjects: Artificial neural networks; Biomedical imaging; Feature extraction; Fuses; hierarchical feature fusion; Image segmentation; Lesions; Medical diagnostic imaging; Medical image segmentation; Medical imaging; Modules; scale aggregation module; Semantics; Uncertainty; uncertainty-induced cross-level fusion
• ISSN: 1051-8215; 1558-2205
• DOI: 10.1109/TCSVT.2024.3370685

Medical image segmentation is an essential process to assist clinics with computer-aided diagnosis and treatment. Recently, a large amount of convolutional neural network (CNN)-based methods have been rapidly developed and achieved remarkable performances in several different medical image segmentation tasks. However, the same type of infected region or lesions often has a diversity of scales, making it a challenging task to achieve accurate medical image segmentation. In this paper, we present a novel Uncertainty-aware Hierarchical Aggregation Network, namely UHA-Net, for medical image segmentation, which can fully make utilization of cross-level and multi-scale features to handle scale variations. Specifically, we propose a hierarchical feature fusion (HFF) module to aggregate high-level features, which is used to produce a global map for the coarse localization of the segmented target. Then, we propose an uncertainty-induced cross-level fusion (UCF) module to fully fuse features from the adjacent levels, which can learn knowledge guidance to capture the contextual information from adjacent resolutions. Further, a scale aggregation module (SAM) is presented to learn multi-scale features by using different convolution kernels, to effectively deal with scale variations. At last, we formulate a unified framework to simultaneously fuse inter-layer convolutional features and learn the discriminability of multi-scale representations from the intra-layer features, leading to accurate segmentation results. We carry out experiments on three different medical image segmentation tasks, and the results demonstrate that our UHA-Net outperforms state-of-the-art segmentation methods. Our implementation code and segmentation maps will be publicly at https://github.com/taozh2017/UHANet .