Wavelet U-Net for Medical Image Segmentation

Abstract

Biomedical image segmentation plays an increasingly important role in medical diagnosis. However, it remains a challenging task to segment the medical images due to their diversity of structures. Convolutional networks (CNNs) commonly uses pooling to enlarge the receptive field, which usually results in irreversible information loss. In order to solve this problem, we rethink the alternative method of pooling operation. In this paper, we embed the wavelet transform into the U-Net architecture to achieve the purpose of down-sampling and up-sampling which we called wavelet U-Net (WU-Net). Specifically, in the encoder module, we use discrete wavelet transform (DWT) to replace the pooling operation to reduce the resolution of the image, and use inverse wavelet transform (IWT) to gradually restore the resolution in the decoder module. Besides, we use Densely Cross-level Connection strategy to encourage feature re-use and to enhance the complementarity between cross-level information. Furthermore, in Attention Feature Fusion module (AFF), we introduce the channel attention mechanism to select useful feature maps, which can effectively improve the segmentation performance of the network. We evaluated this model on the digital retinal images for vessel extraction (DRIVE) dataset and the child heart and health study (CHASEDB1) dataset. The results show that the proposed method outperforms the classic U-Net method and other state-of-the-art methods on both datasets.