Wavelet Multi-scale Region-Enhanced Network for Medical Image Segmentation

1 Introduction

Medical image segmentation aims to segment anatomical or pathological structures within medical images. It plays a crucial role in computer-aided diagnosis and intelligent healthcare and has been widely studied [Shaker et al., 2024; Ates et al., 2023; Messaoudi et al., 2023; Yin et al., 2023; Chen et al., 2024].

In recent years, deep learning methods, especially the convolutional neural network (CNN) based methods, such as fully convolutional network (FCN) [Shelhamer et al., 2017] and UNet architecture [Ronneberger et al., 2015] have achieved significant advancements in medical image segmentation. For example, [Tureˇckov´a et al., 2020] proposed attention gates to improve segmentation accuracy. [Fang et al., 2019] combined UNet architecture with residual channel attention blocks for segmentation. [Xiang et al., 2020] proposed Bio-Net, which achieved performance improvement without adding extra parameters by recurrent bi-directional skip connections.

∗ Peng Zhou is the corresponding author.

Although UNet has been widely adopted, there still exist some issues of detailed information missing during the segmentation. More specifically, first, in the encoding stage, the downsampling approach often leads to the loss of important information [Xu et al., 2023a]. Second, in the decoding stage, the fusion that directly concatenates shallow semantic information from the skip connections and deep semantic information from the previous layer of the decoder may be inappropriate in medical image segmentation. In medical image segmentation, we should focus on various regions of interest (ROI) with different shapes and sizes. However, directly concatenating shallow and deep features may cause information interference in the decoder due to background noises in shallow features, hindering the model's ability to effectively capture ROIs and detailed information [Xu et al., 2023b; Oktay et al., 2018].

To address these issues, in this paper, we propose a novel Wavelet Multi-scale Region-Enhanced Network (WMREN), which can effectively preserve the detail information and enhance the ROIs in the encoding and decoding process during the segmentation. To tackle the first problem, we apply the Haar Wavelet Downsampling (HWD) [Xu et al., 2023a] to extract the features. HWD achieves lossless information transformation via the Haar wavelet transform, preserving detailed image information by increasing the number of feature map channels while reducing resolution. However, we observe that, when directly applying HWD to handle medical images, it cannot achieve satisfactory segmentation on multiscale organs due to its fixed receptive field. Figure 1 shows an example. Figure 1a is the ground truth segmentation of an image in the multi-organ segmentation dataset Synapse. Figure 1b is the segmentation result with HWD as its encoder. Considering the small organ pancreas, HWD cannot segment it accurately.

To further address this issue, we propose a novel bi-branch encoder that simultaneously applies the Haar wavelet transform and the residual blocks to extract features. As shown in Figure 1b, HWD may fail to extract high-frequency features such as details and small objects like pancreas. On the contrary, the convolution based network can capture the local detailed features with a small convolution kernel. Figure 1c shows the result of the network with residual blocks as its encoder. The small organ pancreas can be segmented correctly. Therefore, we try to tackle the problem in Haar wavelet transform by parallelly extracting features with a residual network. To effectively harness the complementarity of the Haar wavelet transform and residual network, we design a novel Wavelet Residual Multi-scale Module (WRMM), which has a bi-branch structure to handle multi-scale objects. Figure 1d shows the result of our bi-branch encoder, which performs better than both the HWD based and Residual based encoders.

To tackle the second problem about the decoder, we develop a novel module named Spatial Adaptive Fusion Module (SAFM) for effectively fusing the features in the skip connections and previous layer, instead of directly concatenating them, which can effectively enhance the ROIs. SAFM module can focus on and enhance key regions and suppress background noise during the feature fusion. However, we observe that this fusion process has a side effect that may blur the boundary of objects. To address this issue, we further propose a Contrastive Refinement Enhancement Module (CREM) and plug it into the SAFM to preserve the edge information. CREM employs a bi-branch architecture, where the main branch enhances multi-scale features and fuses them with decoder features, and the other branch extracts boundary information by comparing the differences before and after boundary enhancement, which can effectively preserve the boundary information. By integrating these two modules into the decoder, the segmentation results can maintain both global semantic information and local detail information.

Our main contributions are summarized as follows:

• We propose a novel WMREN framework for medical image segmentation, in which we introduce a collaborative downsampling approach combining the wavelet transform and CNN. This framework can effectively extract multi-scale features frommedical images by reducing information loss.
• We design an innovative SAFM module and a CREM module in the decoder. These modules can highlight and enhance the ROIs and preserve the boundary informa-

tion, while suppressing the noises and artifacts.

• We conduct comprehensive evaluations on medical image datasets. The results demonstrate the superiority of our method compared to state-of-the-art medical image segmentation methods.

2 RelatedWorks

2.1 Medical Image Segmentation

Medical image segmentation is an important task in medical image processing. In medical segmentation, U-shaped networks, which are in an encoder-decoder framework, are widely used in this field [Ronneberger et al., 2015; Isensee et al., 2018; Rahman et al., 2024]. For example, UNet++ introduced a nested encoder-decoder structure with dense feature connections [Zhou et al., 2020]. To better utilize multi-scale information, UNet 3+ was designed to comprehensive feature transmission pathways, achieving more thorough feature fusion [Huang et al., 2020]. Some works improved UNet by incorporating attention mechanisms [Oktay et al., 2018; Xiang et al., 2020; Jin et al., 2020] and residual structures [He et al., 2016; Ibtehaz and Rahman, 2020; Sharma and Mishra, 2023; Rahman et al., 2021]. For example, the DeepLab series, including DeepLabv3 [Chen et al., 2018a] and DeepLabv3+ [Chen et al., 2018b], incorporated atrous convolution and spatial pyramid pooling to handle multi-scale information for medical image segmentation.

Compared to CNNs, Transformer can capture long-range dependencies and enhance global semantic understanding of images more effectively [Zhang et al., 2024; Xie et al., 2021; Chen et al., 2023; Rahman and Marculescu, 2023]. For example, TransUNet effectively captured global and local features through CNN-Transformer fusion [Chen et al., 2021]. MISSFormer optimized multi-scale feature processing through innovative feed-forward networks and context bridging modules [Huang et al., 2022]. DAE-Former enhanced feature dimension capture and spatial localization through improved self-attention and cross attention [Azad et al., 2023]. However, the Transformer architecture still has unresolved issues in computational complexity and detail preservation during patch embedding. Therefore, in this paper, we still focus on full convolutional network based methods.

2.2 Wavelet-based Feature Learning in CNNs

Recent studies have explored integrating the wavelet transform into CNN architectures to avoid the information loss caused by downsampling. For example, Multi-level Wavelet CNN (MWCNN) was designed to achieve a good balance between receptive field size and computational efficiency in image restoration tasks [Liu et al., 2018]. SFFNet was proposed to maximize spatial information utilization for segmentation tasks while effectively handling regions with diverse grayscale intensities in remote sensing imagery" [Yang et al., 2024]. WTConv tackled CNN's limited receptive field problem while avoiding excessive parameter growth [Finder et al., 2025]. [Xu et al., 2023a] proposed a wavelet-based downsampling method (HWD) to tackle the spatial information loss during CNN downsampling.

Despite the promising performance, we observe that the wavelet downsampling may fail to handle the organs with complicated shapes and different scales. To tackle this problem, we propose a bi-branch encoder by combining the wavelet downsampling and residual blocks to extract features.

3 Method

In this section, we introduce our WMREN in more detail. Figure 2 illustrates the overall framework of WMREN. It is a UNet structure that contains an encoder and decoder. The encoder contains four Wavelet Resdiual Multi-scale Modules (WRMM) and the decoder contains three Region Enhancement Modules (REM). WRMM will be introduced in Section 3.1 and REM will be introduced in Section 3.2. In the skip connection between encoder and decoder, we apply 1×1 convolutions to adjust the number of channels, ensuring consistent channel dimensions when fusing features with the decoder. The details of the encoder and decoder will be introduced in the following subsections.

3.1 Wavelet Residual Multi-scale Module

The conventional encoder in UNet consists of convolution and downsampling operations. In our WRMM encoder, we design a Wavelet ResNet Downsampling Module (WRDM) for downsampling and apply a Multi-scale Large Kernel Module (MLKM) [Wang et al., 2024] for convolution. The structure of WRMM is shown in Figure 3a. In the following, we will introduce our designed WRDM in more detail.

Figure 3b shows the structure of WRDM. It has a bi-branch structure. The left branch applies the Haar wavelet transform for downsampling and the right branch applies a residual block to extract features. The Haar wavelet branch first uses

(a) WRMM

HWD [Xu et al., 2023a] for downsampling. In more detail, given an input feature map X, it decomposes it into a lowfrequency component XL and three high-frequency components: horizontal detail coefficients XHL, vertical detail coefficients XLH, and diagonal detail coefficients XHH:

where DWT denotes a wavelet transformation function.

Then, it combines the low-frequency component and three high-frequency components and feeds them into a 1 × 1 convolutional layer with a batch normalization (BN) and Sigmoid activation function (σ(·)) to obtain the feature maps of HWD (XHWD). More formally, this downsampling can be formulated as follows:

where Concat represents the feature concatenation along the channel dimension. Conv1×1 denotes the 1 × 1 convolution.

To preserve space information more effectively, in the right branch, we also apply several residual blocks to extract features in parallel. In more detail, we first perform downsampling of feature maps via convolutional layers and then use residual blocks to extract features. Each residual block consists of three consecutive convolutional layers with kernel sizes of 1 × 1, 3 × 3, and 1 × 1, and then adds the input features to the output feature. We denote the result of the residual blocks as XRes.

After the Haar wavelet downsampling and residual downsampling, we fuse them to obtain a new representation. To this end, we first obtain a fusion feature F1 by concatenating two features (XHWD, XRes) followed by a 3×3 convolution, which is shown as follows:

Then, for each branch (i.e., the HWD branch and residual branch), we apply 1 × 1 convolutional layer and Softmax activation function to calculate an adaptive weight (i.e., W1 for the HWD branch and W2 for the residual branch). More formally, we calculate W1 and W2 as:

where Split represents the equal slice along the channel dimension. Softmax represents the Softmax activation function. The weighted features are integrated through a 3 × 3 convolution to obtain the fusion feature F2, which is:

where ⊙ denotes the element-wise multiplication. Finally, F1 and F2 are summed to obtain the final fusion representation FDown = F1 + F2.

3.2 Region Enhancement Module

In the conventional UNet decoder, the fusion between deep features from the decoder and shallow features from the encoder is just a concatenation. However, in medical image segmentation, we should focus on some ROIs, especially the boundaries of the organs. Therefore, when generating the segmentation images in the decoder, we should enhance the ROIs and boundaries. To this end, we design a Region Enhancement Module in the decoder, which can enhance the important regions and boundaries when fusing the features from the encoder and decoder. The structure of REM is shown in Figure 4a. It consists of three modules: Spatial Adaptive Fusion Module (SAFM), Contrast Refinement Enhancement Module (CREM), and Convolutional Block Attention Module (CBAM) [Woo et al., 2018]. SAFM enhances the main region of segmentation targets by fusing deep-layer features from the decoder with shallow-layer features from the encoder at an early stage. CREM enhances the shallow-layer features of the encoder through multi-scale boundary feature extraction.

CBAM [Woo et al., 2018] effectively highlights crucial information through the lightweight channel and spatial attention mechanisms.

In more detail, the feature from the previous decoder layer is first fed into CBAM to obtain an enhanced feature map, denoted as FCBAM, which serves as an input of the subsequent SAFMmodule. Meanwhile, the shallow feature Xskip from the encoder is fed into CREM, which outputs an edgeenhanced feature map XCREM and an edge weight information map WCREM, where WCREM guides boundary recovery during subsequent feature fusion. Next, SAFM receives the edge-enhanced shallow features XCREM from CREM and the features FCBAM from CBAM as inputs. Injecting deep feature information into shallow features effectively prevents the loss of main detail information during the subsequent fusion process. Its output is denoted as XSAFM. During the process of skip connection and decoder feature fusion, XSAFM and FCBAM are first concatenated along the channel dimension, followed by 1 × 1 convolution to obtain the fused feature map Ffusion. Then, using the previously obtained edge weight WCREM to guide the detail recovery process, we can obtain the reconstructed feature map for decoder

With this complete process, the decoder module ensures effective fusion of deep and shallow features and accurate preservation of details, thereby achieving high-quality feature segmentation. In the following subsections, we will introduce our designed CREM and SAFM in more detail.

Contrast Refinement Enhancement Module

In medical image segmentation, the boundary information is important and we should prevent the boundary blurring during the segmentation image reconstruction. To this end, we propose a Contrast Refinement Enhancement Module, as shown in Figure 4b. We tackle the boundary-blurring issue in two ways. On one hand, we enhance boundaries in the shallow feature maps of the encoder. On the other hand, we implement boundary guidance on the feature maps following feature fusion in the decoder.

Specifically, at first, we need to obtain multi-scale information from the shallow feature Xskip. To this end, the input feature Xskip first is fed into a 1 × 1 convolution to obtain P0 = Conv1×1 (Xskip). Then, P0 is fed into three cascaded AveragePool layers (AvgPool) to obtain feature representations at different scales, which are shown as follows:

P1, P2, and P3 are the three scales features, respectively.

To obtain boundary information through comparison, we first perform the boundary blurring method on the feature maps Pis (i = 1, 2, 3) to obtain boundary blurring feature maps, and then subtract the blurred feature map from the original feature map Pi to obtain edge information. The boundary information is then added to the original feature map Pi to achieve edge enhancement. The process is shown as follows:

where Ei is the boundary enhanced features for the i-th scale. To obtain multi-scale boundary feature maps, we concatenate P0 with the obtained Ei and then integrate them using 1 × 1 convolution, as shown in Eq.(9). The result is denoted as XCREM, which will serve as an input to SAFM to provide shallow-layer information.

Besides, CREM also generates weights like the space attention for the feature maps to highlight the boundary. The weights, denoted as WCREM, are generated via a 1 × 1 convolutional layer on the difference between the enhanced representation XCREM and the original feature map Xskip as follows:

WCREM will be used as the weights in Eq.(6) to obtain the reconstructed feature maps in decoder Fdecoder.

Spatial Adaptive Fusion Module

To better highlight the important regions in the encoder's shallow representations and effectively fuse them with the decoder's representation, we propose the SAFM. As illustrated in Figure 4c, let XCREM and FCBAM denote the shallow features from CREM and deep features from CBAM, respectively, both with dimensions H ×W ×C. The feature fusion process can be described as follows.

To achieve adaptive feature selection and preserve important features, attention weights are generated by applying Sigmoid functions to both feature maps, enabling the emphasis of crucial information while suppressing less relevant features. Finally, to utilize complementary information from both features, the weighted features are combined through element-wise addition. We apply FMUL as the result of fusion, which is shown as follows:

Notice that only the regions that are highlighted in both XCREM and FCBAM, which means these may be the ROIs, can be highlighted in Eq.(11). Therefore, Eq.(11) can effectively enhance the main body of ROIs.

Subsequently, the fused features FMUL are fed into a 1 × 1 convolution, followed by a BatchNorm normalization and ReLU activation function. The integrated features are then processed through another Sigmoid function to generate the final attention weights WSAFM, which are shown as follows:

where ReLU denotes the ReLU activate function.

To ensure that useful low-level information is not lost while highlighting important features, we apply residual connections to preserve the original information. Specially, these weights in Eq.(12) are multiplied with XCREM and combined with the input features through a residual connection as: We denote the result of the SAFM as FSAFM

3.3 Loss Function

We employ a composite loss function to optimize the network parameter, which combines Cross Entropy Loss LCE and Dice Loss LDice . Specifically, the total loss function Ltotal

The Cross Entropy Loss (LCE) measures pixel-level discrepancies between predicted results and ground truth labels, contributing to improving classification accuracy. The Dice Loss (LDice) focuses on evaluating the overlap between predicted segmentation masks and ground truth masks to improve the segmentation accuracy.

4 Experiment

4.1 Experimental Setup and Implementation Details

We evaluate our method on three benchmark datasets: Synapse 1 ), ACDC 2 , and ISIC17 [Codella et al., 2018]. On all data sets, we follow the default protocol for splitting the dataset into training, validation, and testing sets. The Synapse dataset is a multi-organ segmentation dataset, which consists of 3779 axial abdominal CT images collected from 30 clinical cases. Each image is annotated with labels for eight abdominal organs: aorta, gallbladder, spleen, left kidney, right kidney, liver, pancreas, and stomach. The performance evaluation is conducted using two metrics: the average Dice score (DSC) in percentage and the average Hausdorff Distance (HD95) in millimeters across all eight abdominal organs. The ACDC dataset is an automated cardiac diagnostic challenge dataset that consists 100 cardiac cine-MRI scans from clinical cases, and corresponding labels for three cardiac regions: the left ventricle (LV), right ventricle (RV), and myocardium (MYO). All methods are evaluated with DSC over all three cardiac regions. The ISIC17 dataset is a skin lesion segmentation dataset containing high-resolution medical images of melanoma skin lesions. Each image includes pixel-level annotations marking the lesion area. Performance evaluation is conducted using the DSC to measure the segmentation accuracy of the lesion regions.

1 https://www.synapse.org#!Synapse:syn3193805/wiki/217789

2 https://www.creatis.insa-lyon.fr/Challenge/acdc/

All experiments are implemented using PyTorch on Windows 10 system and Nvidia GeForce RTX 4090 GPU. Following [Huang et al., 2022], various data augmentation techniques are employed, including random cropping, random rotation, random Gaussian noise, random blurring, as well as transformations in luminance, contrast, and resolution. Images from Synapse and ACDC datasets are reshaped to 224 × 224 pixels. In the ISIC17 dataset, images are resized to 256 × 256. The model is trained via SGD optimizer with a momentum of 0.90, weight decay of 0.0001, and an initial learning rate of 0.05 and following a polynomial decay policy. It is trained for 400 epochs with a batch size of 24.

4.2 Experimental Results

Table 1 shows the results of WMREN and other methods on the Synapse dataset. It shows that WMREN achieves state-ofthe-art performance with an average DSC of 84.40%, surpassing both CNN-based and Transformer-based methods. Our method also has the lowest HD95 at 15.65 mm, indicating its effectiveness in organ boundary detection. Besides, compared to other state-of-the-art methods, our method achieves better segmentation accuracy for the small organs, such as the pancreas, gallbladder, and kidneys. It demonstrates superior performance in segmenting small organs, while maintaining high accuracy for large organs such as the stomach and liver.

Tables 2 and 3 show the results on ACDC and ISIC17 datasets, respectively. WMREN also performs the best on these datasets. Our method achieves a leading average DSC of 92.12% on ACDC dataset and 87.67% on ISIC17 dataset, demonstrating the superiority of the proposed method.

Figures 5, 6, and 7 show the qualitative results on the Synapse, ACDC, and ISIC17 datasets, respectively. As demonstrated in Figure 5, our method exhibits superior performance in capturing fine details such as small objects and boundaries compared to other approaches. Figure 6 shows that our method demonstrates superior segmentation performance in areas where other approaches lose details, particularly in the right ventricle and myocardial boundaries. As shown in Figure 7, our method effectively mitigates common issues found in other approaches, including blurred segmentation boundaries and spurious segmented objects.

4.3 Ablation Study

In this section, we conduct comprehensive ablation studies to evaluate the effects of our designed modules, including the WRDM in the encoder, and SAFM and CREM in the decoder.

In the encoder, we design a WRDM module that simultaneously applies the Haar wavelet branch and residual branch to extract features. To show the effectiveness of the bi-branch architecture, we compare it with two single-branch encoders, i.e., H-Net that only uses the Haar wavelet branch and R-Net that only uses the residual branch. The results on Synapse are shown in Table 4. It can be seen that bi-branch network performs better than both the single-branch network, demonstrating the effectiveness of bi-branch encoder WRDM.

In the decoder, we design a CREM module to enhance the boundaries and SAFM to enhance the ROIs. To show the effectiveness of these two modules, we conduct ablation studies by comparing with the following degenerated models: baseline denotes the UNet backbone without the CREM and SAFM. Baseline+SAFM denotes the network with SAFM but without CREM. Baseline+SAFM+CREM denotes our whole network. We can see that, compared to the baseline model, baseline+SAFM improved DSC by 0.8% and the combination of CREM and SAFM improved DSC by 1.68%. This demonstrates that SAFM is effective in enhancing the model's performance. Furthermore, we can observe that baseline+SAFM increased baseline by 4.58 w.r.t. HD95. The main reason may be that although SAFM can enhance the ROIs, it may blur the boundaries which causes the increase of the HD95. To tackle this problem, we design the CREM module to enhance the boundaries. We can see that, with CREM, our model achieves 15.65 w.r.t. HD95, which is much lower than both baseline and baseline SAFM, demonstrating the effectiveness of the CREM on boundaries enhance.

Figure 8 shows a qualitative result. Considering the pancreas, the baseline decoder cannot detect the whole region of the pancreas. Together with SAFM, baseline+SAFM can capture the regions, showing the effectiveness of the region enhancement. However, the boundary of baseline+SAFM is incorrect. After plugging CREM into the model, we can segment the pancreas more accurately. It demonstrates the effectiveness of CREM in boundary enhancement.

5 Conclusion

In this paper, to address the detail loss issues of traditional medical image segmentation methods, we proposed an innovative Wavelet Multi-scale Region-Enhanced Network for medical image segmentation. In the encoder, we integrated wavelet downsampling with residual blocks to minimize detail information loss during the downsampling. In the decoder, we carefully designed a Contrastive Refinement Enhancement Module to highlight the boundaries of organs and a Spatial Adaptive Fusion Module to enhance the ROIs. Extensive experiments on benchmark datasets show that our method outperforms other state-of-the-art medical images segmentation methods. The ablation studies also demonstrates the effectiveness of all our designed modules and their contributions to the overall performance.

Acknowledgments

This paper is supported by the National Natural Science Foundation of China grants 62176001, and the Natural Science Project of Anhui Provincial Education Department under grant 2023AH030004.