Uncertainty-Aware Hierarchical Aggregation Network for Medical Image Segmentation

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.

I. Introduction

Medical image segmentation plays a vital role in the medical computing domain [1], such as tumor segmentation in magnetic resonance (MR) images [2], polyp segmentation in colonoscopy images [3], breast lesion segmentation in ultrasound images [4], lung infection segmentation in computed tomography (CT) Images [5], nuclei segmentation in microscope cell images [6], and so on. Medical image segmentation provides a particular region of the lesion and extracts useful information for clinicians to achieve a more accurate diagnosis. Manual annotation is typically adopted in clinical practice, however, it is time-consuming and subjective. Therefore, there is a high demand for developing automatic and accurate segmentation methods, which can be used to derive a quantitative assessment for subsequent diagnosis, treatment planning, and lesion progression monitoring [7].