FFUNet: A novel feature fusion makes strong decoder for medical image segmentation

IET Signal ProcessingVolume 16, Issue 5 p. 501-514 ORIGINAL RESEARCHOpen Access FFUNet: A novel feature fusion makes strong decoder for medical image segmentation Junsong Xie, Corresponding Author Junsong Xie [email protected] orcid.org/0000-0002-6845-0441 Department of Computer Science, Anhui Medical University, Hefei, China Correspondence Junsong Xie, Department of Computer Science, Anhui Medical University, No. 81 Meishan Rd, Hefei 230032, China. Email: [email protected]Search for more papers by this authorRenju Zhu, Renju Zhu Department of Critical Care Medicine, Hefei BOE Hospital, Hefei, ChinaSearch for more papers by this authorZezhi Wu, Zezhi Wu Department of Computer Science, Anhui Medical University, Hefei, ChinaSearch for more papers by this authorJinling Ouyang, Jinling Ouyang First Clinical Medical College, Anhui Medical University, Hefei, ChinaSearch for more papers by this author Junsong Xie, Corresponding Author Junsong Xie [email protected] orcid.org/0000-0002-6845-0441 Department of Computer Science, Anhui Medical University, Hefei, China Correspondence Junsong Xie, Department of Computer Science, Anhui Medical University, No. 81 Meishan Rd, Hefei 230032, China. Email: [email protected]Search for more papers by this authorRenju Zhu, Renju Zhu Department of Critical Care Medicine, Hefei BOE Hospital, Hefei, ChinaSearch for more papers by this authorZezhi Wu, Zezhi Wu Department of Computer Science, Anhui Medical University, Hefei, ChinaSearch for more papers by this authorJinling Ouyang, Jinling Ouyang First Clinical Medical College, Anhui Medical University, Hefei, ChinaSearch for more papers by this author First published: 14 May 2022 https://doi.org/10.1049/sil2.12114AboutSectionsPDF ToolsRequest permissionExport citationAdd to favoritesTrack citation ShareShare Give accessShare full text accessShare full-text accessPlease review our Terms and Conditions of Use and check box below to share full-text version of article.I have read and accept the Wiley Online Library Terms and Conditions of UseShareable LinkUse the link below to share a full-text version of this article with your friends and colleagues. Learn more.Copy URL Share a linkShare onFacebookTwitterLinkedInRedditWechat Abstract Convolutional neural networks (CNNs) have strong ability to extract local features, but it is slightly lacking in extracting global contexts. In contrast, transformers are good at long-distance modelling due to the global self-attention mechanisms while its performance in localization is limited. On the other hand, the feature gap between an encoder and decoder is also challenging for a U-shaped network, which adopts a plain skip connection. Inherited from convolutional networks and transformers, FFUNet, a hybrid network structure with a novel module named the Feature Fusion Module (FFM) is proposed for medical image segmentation. The proposed FFM consisting of Feature Attention Selection, Cross Offset Generation and Deformable Convolution Layer, aims to replace the original plain skip connection to alleviate the ambiguous semantic information between the encoder and decoder for a more powerful medical image segmentation network. Extensive experiments demonstrate that the proposed FFUNet has amazing performance in segmentation gains on the Synapse dataset. In addition, consistent improvements are also achieved across other four popular datasets and CNN-based or transformer-based segmentation networks, which illustrate that the proposed method has advantages in generalization and compactness. 1 INTRODUCTION Medical image segmentation has a positive effect on the diagnosis of patients and the formulation of subsequent treatment plans. The segmentation usually requires experienced doctors to perform, which is very time-consuming and may cause some subjective errors. In recent years, benefiting from the great accomplishments of convolutional neural networks (CNNs; e.g. fully convolutional networks [1]), automatic and high-precision segmentation of medical images has become possible. In the field of medical image segmentation, a U-shaped structure including a symmetric encoder–decoder with skip connections [2, 3], has become the mainstream segmentation network architecture. The main part of the subsequent networks is to finely design these three parts (namely Encoder, Decoder and Skip Connections) to obtain better segmentation performance. Before transformers [4] are applied in computer vision (CV) tasks, encoders and decoders mainly rely on convolution operations. The encoder is used to extract semantic information from the original feature maps to generate low-resolution and high-dimensional image features, and the decoder upsamples the final output of the encoder to gradually restore the resolution to the original feature map size for dense prediction. In the encoder process, continuous downsampling operations are usually used to obtain high-level semantic feature maps with a large receptive field. Then, in the decoding process, corresponding skip connections are used to make up for the loss of spatial information caused by downsampling in the encoder. Using such a concise structure, U-Net and its variants (Res-UNet [5], MultiResUNet [6], U-Net++ [7] and U-Net3+ [8]) have achieved encouraging results on various medical image datasets. Although the CNN-based U-Net structure has achieved great segmentation achievements on multiple medical image datasets, the inherent flaws of convolution operations focus on localized feature map results in lacking the capability to model long-distance. On the other hand, benefiting from the global self-attention mechanisms, transformers are the expert in long-distance modelling. Thus, more and more researchers are exploring to adopt the hybrid combination of the CNN and transformer (e.g., Conformer [9], BotNet [10] and TransFuse [11]) or pure transformer (e.g., Swin Transformer [12] and Shuffle Transformer [13]) to extract high-level semantic information for precise classification while eliminating the intrinsic locality of convolution. The transformer-based recognition networks have matched or surpassed the state-of-the-art network based on CNN on some large datasets (e.g., ImageNet-22K) while have similar or even less parameters and computational cost. To go with the flow, motivated by the excellent performance, the transformer is graduallyemployed by the encoder and decoder of the U-Net structure in segmentation tasks. For example, SwinUNet [14], a U-shaped pure transformer for medical image segmentation tasks outperforms the state-of-the-art methods on the Synapse multi-organ CT dataset (Synapse) [15] and Automated cardiac diagnosis challenge dataset (ACDC) [16]. In addition to attentively design the powerful encoder and decoder, an appropriate skip connection is also essential to fuse the feature correctly. Since feature gaps between the stages of the encoder and decoder will degrade segmentation performance, a series of approaches are proposed to explore how to ease the gaps between the encoder and decoder. A typical network named U-Net++ [7] that uses dense connectivity in skip connection between the encoder and decoder has achieved better segmentation results than the plain skip connections. UCTransNet [17] designs the skip connections from a channel-wise perspective, which consists of Channel Cross fusion with Transformer (CCT) and Channel-wise Cross-Attention (CCA) and produces superior segmentation performance on the datasets mentioned in the paper. Inspired by the Conformer's [9] stunning success, we propose FFUNet to leverage the power of CNNs and transformers for 2-D medical image segmentation in this work. FFUNet is a hybrid U-shaped architecture that consists of the encoder, decoder, and skip connections. The encoder is built based on the Conformer module, which has dual branches (i.e., CNN and transformer branch). It is worth mentioning that we did not use the transformer branch in the decoder. Besides, to strengthen the ability of the encoder to extract high-level semantic information, skip connection also plays an important role in the decoder's ability to restore spatial details. Different channels and spatial dimensions usually focus on different semantic patterns; adaptively and efficiently fusing features is favourable for the complicated dense prediction (e.g., medical image segmentation) to improve the precision. We rethink the feature gaps of skip connections and redesign an alternative architecture instead of a plain connection operation (e.g., element-wise addition or channel-wise concatenation) for better fusing the features between the encoder and decoder stages. More specifically, a novel module named Feature Fusion Module (FFM), which learns to align the encoder feature maps to the decoder by selecting each computational location in a convolutional kernel with a learnable offset, is proposed to efficiently fuse the incompatible feature sets of the encoder and decoder stage. FFM can be easily embedded in and applied for the popular UNets in medical image segmentation tasks. A series of experiments show that FFUNet can greatly improve segmentation pipelines by the following absolute gains of 1.50%, 0.70% and 1.24% DSC over SwinUNet on Synapse, ACDC and Gland segmentation (GlaS) datasets, respectively. Our contributions can be summarized as follows: Motivated by the excellent performance of the Conformer, we build the encoder based on the Conformer module, which is composed of the CNN block, the transformer block, and the Feature Coupling Unit (FCU) [9]. The decoder retains Conformer's CNN block only. ConvUpBlock is developed to upsample the extracted deep features with interpolation operation for pixel-level semantic prediction. In order to alleviate the feature gap between encoder and decoder stages, FFM, which consists of Feature Attention Selection (FAS), Cross Offset Generation (COG) and Deformable Convolution Layer (DCL), is elaborately designed to boost the skip connection effectiveness. We introduce a new perspective to bridge the feature gap between the encoder and decoder by a more effective feature fusion with the attention mechanism [18] and deformable convolution operation [19] for capturing more sophisticated dependencies. We present a thorough experimental evaluation indicating the excellent performance of FFUNet on five public datasets, including Synapse [15], ACDC [16], GlaS [20], MoNuSeg [21, 22] and ISIC-2017 [23]. The proposed FFUNet leads to a significant segmentation precision, and DSC/HD reaches 80.63/18.15 on the Synapse dataset, yielding an overall improvement of 1.50/3.40 scores compared with the recently state-of-the-art SwinUNet [14]. The consistent improvements have achieved across the remaining four datasets compared with the CNN-based or transformer-based segmentation network. In general, extensive experiments on multiple datasets reveal that our network has strong robustness and generalization capabilities. The rest of this paper is organized as follows. In Section 2, we review related work about the medical image segmentation tasks. Section 3 introduces the network structure and detailed method. Then, the comparison and ablation experiments are presented in Section 4, and Section 5 concludes our proposed method. 2 RELATED WORK CNN-based encoder and decoder methods In the past period, CNNs [24-29] that rely on convolution operations and residual connections occupied a mainstream position in image recognition. Medical image segmentation is a dense prediction task to identify every pixel to a classification. Due to the outstanding performance and elegance of the U-shaped structure, UNet-based methods are widely used for segmentation tasks. Res-UNet [5] is a residual network building on the original UNet model and adding a weighted attention mechanism. Such design allows the model to learn more to distinguish the characteristics of blood vessels and non-vascular pixels, and to better maintain the retinal vessel tree structure. The liver and tumour segmentation from CT Volumes faces the two challenges that 2-D CNN cannot learn spatial information in the third dimension and 3-D CNN requires high computational cost and GPU memory consumption. To overcome the aforementioned issues, Li et al. [30] proposed a novel hybrid densely connected UNet (H-DenseUNet), which consists of a simple ResNet, 2D DenseUNet, 3D DenseUNet and hybrid feature fusion layer. At present, CNN-based methods have got impressive success not only in natural image classification but also in medical image segmentation fields attributed to its powerful feature extraction capability. Transformers for medical image segmentation Recently, a pioneering Vision Transformer [4] depending on self-attention mechanisms that can eliminate the limited long-distance dependency abilities of convolution operation has achieved state-of-the-art performance on natural image classification. Motivated by the impressive accuracy of transformers in various CV fields, new frameworks for medical image segmentation have recently been developed in order to model global contexts. For instance, Chen et al. [15] firstly attempt to introduce the transformer into the encoder named TransUNet, and achieved good performance in medical image segmentation on Synapse and ACDC datasets. Unlike TransUNet, TransBTS [31] is based on 3-D CNN and processes all image slices at once, allowing for better representation of continuous information between slices. The Medical Transformer [32] is designed to meet the fact that data in medical imaging is relatively low compared with the datasets for CV using a gated axial-attention model and a Local-Global training strategy. Driven by the great performance of Swin Transformer [12], Cao et al. [14] proposed SwinUNet, which is composed of a pure transformer encoder and a pure transformer decoder, to reveal the feasibility and potential of pure transformer in U-shaped structure's encoder and decoder. As a summary, we present the application of the transformer in the medical image segmentation encoder and decoder in Table 1. TABLE 1. Transformers for medical image segmentation Networks Encoder Decoder TransUNet CNN + transformer CNN SwinUNet Transformer Transformer TransBTS 3-D CNN + transformer 3-D CNN + transformer MedT CNN + transformer CNN FFUNet CNN + transformer CNN + FFM Abbreviations: CNN, convolutional neural network; FFM, Feature Fusion Module. It is worth mentioning that a dual network framework (CNN branch and transformer branch), called Conformer [9] that utilizes CNN to extract local features and transformer to capture global contexts, is proposed as a general backbone network. The conformer which exploits the FCU to fuse the CNN's local features and transformer's global representations achieves striking performance on various CV tasks including image classification, object detection and instance segmentation. In this work, we explore to use the Conformer module as basic unit to construct a qualified encoder–decoder network with a well-designed encoder for medical image segmentation. Skip connections between the encoder and decoder When UNet [2] was proposed, the skip connection mechanism was applied to bridge the feature between the encoder and decoder and played an important role in restoring finer details [33]. The skip connections of UNet impose unnecessary restrictive fusion schemes, only forcing fusion on the same scale feature maps of the encoder and decoder subnets. To overcome the unnecessary restriction, UNet++ [7] combines not only deep features but also shallow features for collecting multi-scale features to alleviate the sematic gap, and more conducive to feature fusion. Wang et al. [17] analysed the impact of the skip connection to medical image segmentation on GlaS and MoNuSeg datasets and proposed the UCTransNet, which consists of CCT and CCA from channel attention perspective and produced superior segmentation performance. Huang et al. [34] believed that adopting pixel addition or channel-wise concatenation between the encoder and decoder results in misaligned feature maps and degrades the segmentation accuracy, especially on object boundaries. They present the Feature-aligned Pyramid Network (FaPN), which consists of a feature alignment module and a feature selection module that outperforms FPN [35] when paired with Faster/Mask R-CNN in Average Precision/mean Intersection-over-Union. Feature fusion network Features fusion involving different stages is an important means to improve segmentation performance The low-level features have higher resolution and contain more detailed information; however, they lack semantics and have more noise due to fewer convolutions. High-level features have more semantic information, while their resolution is very low, and the perception capability of details is poor. Therefore, how to efficiently integrate these two kind of features is the key to the segmentation model. As shown in Table 2, we present the recently popular fusion methods. TABLE 2. Feature fusion methods of different networks Networks Methods U-Net Channel-wise concatenation ResNet Element-wise addition UCTransNet CTrans module SwinUNet Channel-wise concatenation FFUNet Feature Fusion Module In this work, we designed FFM which leverages not only the feature map of the current stage but also the feature map of the previous stage to fuse the features effectively as shown in Figure 1. Detailed introduction of FFM will be illustrated in section 3. FIGURE 1Open in figure viewerPowerPoint Overview comparison among our FFUNet using FFM (d) and other three networks. FFM, Feature Fusion Module 3 METHODS Figure 2 illustrates an overview of our elaborately designed FFUNet architecture. It is believed that excellent segmentation performance requires strong cooperation of among the three parts of Encoder, Decoder and Skip Connections. Then, a question arises as to how the decoder effectively uses the encoder and skip connections to perform accurate pixel-wise classification, which is the key to dense prediction. Driven by the question, we elaborated design of the decoder, which employs a novel module named FFM to bridge the feature gap between the encoder and decoder by a more effective feature fusion with an attention mechanism [18] and deformable convolution operation [19] for capturing more sophisticated dependencies. FIGURE 2Open in figure viewerPowerPoint The architecture of FFUNet, which is composed of encoder, decoder and skip connection. The encoder is built on the Conformer module. To better fuse features, we elaborate the design decoder, which exploits skip connections, FAS, COG and DCL. COG, Cross Offset Generation; DCL, Deformable Convolution Layer; FAS, Feature Attention Selection Architecture overview As shown in Figure 2, the proposed FFUNet includes four vital modules, namely, the Conformer module, which is composed of TransBlock and CNNBlock, ConvUpBlock, and the FFM, which consists of FAS, COG and DCL, and skip connections. Specifically, the basic unit of FFUNet's encoder is the Conformer module [9], which is employed to extract multi-scale features and model long-distance dependencies. The ConvUpBlock is specially designed to perform 2× upsampling with interpolation operation for pixel-level semantic prediction. The FAS employs the attention mechanism to select the importance of feature maps from the encoder and as an input of COG to generate offsets for DCL. The COG is adopted to fuse features produced by ConvUpBlock, FAS and the last stage low-level feature, effectively generating beneficial feature locations for DCL. The DCL is designed to choose helpful features according to learned offsets by COG. The effective cooperation of the three parts makes the FFUNet perform well on multiple datasets. Given an input image X ∈ ℝ H × W × C , H × W and C denote the spatial resolution and number of channels, respectively. Based on the principle of employing both local features and global contexts, the Conformer is designed with a dual network structure, including the CNN branch and transformer branch. We first use the stem and Conformer module to generate five hierarchical features F i ∈ ℝ C i × H 2 i × W 2 i , where i ∈ { 1 , 2 , 3 , 4 , 5 } and C i ∈ { 64 , 64 , 128 , 256 , 256 } . The last decoder block is employed to restore the full spatial resolution with a size of H × W. Encoder with two branches Following [9], our encoder consists of TransBlock, CNNBlock, FCU [9] and ConvDownBlock besides the stem block. The input image X first passes through stem block and a maxpooling layer, yielding feature maps F 1 ∈ C × W 2 × H 2 and F 2 ∈ C × W 4 × H 4 (C = 64), respectively. Then, F2 is fed into the CNN branch and transformer branch to further perform representation learning. Specifically, the CNN branch, which is composed of seven CNNBlocks and ConvDownBlocks, produces a feature pyramid structure. ConvDownBlock performs downsampling to reduce the resolution of the feature map and expand the receptive field. Transformer branch is composed of six TransBlocks, aiming to provide long-distance contexts for the CNN branch, while the CNN branch offers fine-detailed local features to the transformer branch. In order to mitigate the misalignment between the two branches, the FCU [9] is proposed to transform local features with global contexts in an interactive fashion. Similar to the previous literature [36-38], self-attention in TransBlock is as follows: Attention ( Q , K , V ) = Soft Max Q K T / d + B V (1)where Q , K , V ∈ ℝ M 2 × d denote the query, key and value matrices; M2 represents the number of patches in a window; d is query or key's dimension; B is taken from a smaller-sized bias matrix B ˆ ∈ ℝ ( 2 M − 1 ) × ( 2 M − 1 ) [36]. SoftMax maps the input to a real number ranging from 0 to 1, and the normalization guarantees that the sum is 1. Self-attention, which calculates the attention among elements inside the TransBlock, allows the network to integrate information across the entire image even in the lowest layers. Attention can be explained as a means of biasing the allocation of available computational resources towards the most helpful components of a signal. The focussing process is reflected in the calculation of the weight coefficients. The larger the weight, the more focussed on its corresponding value, that is, the weight represents the importance of the information, and the value is the corresponding information. ConvUpBlock As shown in Figure 3, ConvUpBlock is designed to 2× upsample the feature maps, which increases the resolution of maps while the channel number decreases. Firstly, we upsample the input feature map to increase the spatial resolution using bilinear interpolation. Following the layout in ResNet [25], ConvUpBlock is composed of an upsampling layer and two bottlenecks, which include a 1 × 1 convolution to reduce the channel dimension (by a factor 4), a 3 × 3 spatial convolution, a 1 × 1 convolution to recover the channel number and a residual connection making the network going deeper. FIGURE 3Open in figure viewerPowerPoint Schematic of the ConvUpBlock Feature fusion module There is a foreseeable sematic gap between the skip connection of the encoder and decoder due to the downsampling operations by ConvDownBlock. In order to alleviate the gap, FFM is presented (see Figure 2). Specifically, we define the output of the ith skip connection as S i ∈ ℝ H i × W i , where i ∈ { 1 , 2 , 3 , 4 } , Hi = H/2i Wi = W/2i, and H × W is the spatial resolution of the original input image. Take the first skip connection as an example; FFM has three inputs, that is, X ∈ ℝ 14 × 14 from ConvUpBlock, S 4 ∈ ℝ 14 × 14 and S 3 ∈ ℝ 28 × 28 from skip connections. The S 4 passes through FAS (denote as f F A S ( ⋅ ) ), leading to A 4 ∈ ℝ 14 × 14 . The COG takes X, S 4 and S 3 as inputs to generate feature offsets. Then, X and the output of COG (denote as f C O G ( ⋅ ) ) are fed into the DCL (denote as f d ( ⋅ ) ) to sample important points for accurate prediction. In summary, the process of FFM can be formulated as follows: F 4 ′ = f d X , f C O G S 3 , X , f F A S S 4 + f F A S S 4 (2)where [ ] is the concatenation of two parameters. 3.4.1 Feature Attention Selection According to Refs. [39-42], the attention mechanism plays an important role in segmentation tasks. The decoder attempts to reduce the channel number while increasing the spatial resolution for dense prediction using detailed features provided by skip connections. Thus, it is vital and challenging to reinforce the useful feature maps for accurate allocations while suppressing useless features from skip connections. As shown in Figure 4, FAS is proposed to overcome the challenge. Like CBAM in Ref. [18], FAS consists of a channel attention module and a spatial attention module [43], which can be formulated as follows: A i = Conv A t t s A t t c S i (3)where Conv ( ⋅ ) is a 1 × 1 convolutional layer, and At t s ( ⋅ ) and At t c ( ⋅ ) denote the spatial attention module and channel attention module, respectively. According to Figure 4, At t c ( ⋅ ) and At t s ( ⋅ ) can be formulated as follows: At t c S i = f c S i ⊗ S i + S i (4) At t s S i = f s At t c S i ⊗ At t c S i + At t c S i (5)where f c ( ⋅ ) and f s ( ⋅ ) implement the same as CBAM in Ref. [18]. FIGURE 4Open in figure viewerPowerPoint Feature Attention Selection. S i = s 1 , s 2 , … , s D and A i = a 1 , a 2 , … , a D ′ refer to the input and output feature maps respectively, where s d and a d ′ ∈ ℝ H i × W i , D/D′ denotes input/output channel It is worth noting that the design of our FAS is inspired by the CBAM and there are mainly two differences between them. Firstly, we adopt addition operation after the element-wise multiplication to increase the value of important features. Secondly, at the last of the module, a convolution operation with 1 × 1 kernel size followed by group normalization [44] is employed to adjust the channel number. 3.4.2 Cross Offset Generation The dataflow of the proposed COG is shown in Figure 5. X refers to the feature map from the bottom-up path ConvUpBlock, S i and S i − 1 come from skip connections, and COG can be mathematically formulated as follows: C i = Conv U i , A i , X (6) U i = Up S i − 1 ⊗ A i (7)where Conv ( ⋅ ) is a 3×3 convolutional layer with a padding set to 1, Up ( ⋅ ) denotes 2× upsampling operation, and U i , A i , X is the concatenation of U i , A i and X , which provides spatial difference between the encoder and the corresponding decoder feature maps X. In the FaPN [34] module, only A i and X are used to produce offsets, while our COG module not only utilizes A i and X but also the previous skip connection feature denoted as S i − 1 , is exploited to generate offsets for the following deformable convolution operation. FIGURE 5Open in figure viewerPowerPoint The dataflow of the proposed Cross Offset Generation. FAS, Feature Attention Selection 4 EXPERIMENTS In this section, a series of experiments are conducted to evaluate the performance of the proposed FFUNet. Datasets Synapse multi-organ segmentation dataset (Synapse), which has eight abdominal organs (aorta, gallbladder, spleen, left kidney, right kidney, liver, pancreas, spleen, stomach) contains 3779 axial abdominal clinical CT images. Following Ref. [14, 15], 18 cases from Synapse are used for training and other 12 cases are employed for testing in our experiment. Similar to Refs. [14, 15], ACDC we used is collected from 100 patients using MRI scanners. Left ventricle, right ventricle and myocardium are labelled in the dataset, which includes 70 training samples, 10 validation samples and 20 for testing. The ISIC-2017 challenge dataset, provided by the International Skin Image Collaboration (ISIC) archive, is usually divided into 2000 training images, 150 validation images and 600 testing images. Moreover, GlaS has 85 images for training and 80 for testing, while the MoNuSeg dataset has 30 images for training and 14 for testing. Implementation details The hardware equipment used in experiments are Intel Intel(R) Xeon(R) CPU E5-2680 v4 2.40GHz and NVIDIA GeForce RTX 3060 (12G). PyTorch 1.9 and Python 3.6 are used to implement the proposed network. The input image is adjusted to 224 ×224 for all experiments, and batch size is set as 24 except on GlaS and MoNuSeg is set as 4 following Ref. [17, 32]. We use the Conformer pre-training weights on ImageNet to initialize the network, unless otherwise specified. During the training stage, the stochastic gradient descent method with a momentum parameter of 0.9 and weight decay of 1e−4 is employed to optimize the model. For all experiments, only simple data augmentations (e.g., random flipping and rotation) are applied for increasing data diversity to enhance the generalization ability of the model. Both dice loss and binary cross entropy loss is computed to optimize the model for back propagation. It is worth to point that Synapse and ACDC datasets, which are 3-D volumes, are trained in a slice-