ConvWin-UNet: UNet-like hierarchical vision Transformer combined with convolution for medical image segmentation

1.   Introduction

Along with the development of deep learning, image segmentation algorithms, such as UNet ^[1,2,3,4], have achieved great success in medical applications, and the algorithms are significant for medical auxiliary diagnosis.

The classical UNet is a full convolution neural network (FCN) that mainly depends on the U-shaped structure ^[1,5,6]. The early UNet and its variants, such as U-Net++ ^[2], UNet3+ ^[3], U$^2$-Net ^[4], 3D U-Net ^[7], Res UNet ^[8], etc. are constructed based on CNN. At present, there are still much excellent research emerging in the UNet based on CNN, such as KUB-UNet ^[9]. Recently, the self-attention ^[10] structure, represented by Transformer ^[11], has swept the field of computer vision. Meanwhile, a series of UNet models based on Transformer have emerged, such as TransUNet ^[12], SwinUNet ^[13], TransClaw U-Net ^[14], and MT-UNet ^[15]. The performance of UNet has been further improved by the emergence of these networks.

In recent studies, the Transformer structure is feted. And driven by the success of the Transformer structure, the pioneering vision Transformer (ViT) is introduced in ^[11], which achieves an impressive speed-accuracy trade-off on image recognition tasks. To alleviate the difficulty that the training of ViT requires large-scale training datasets, Deit ^[16] introduces several training strategies that allow ViT to be effectively trained on ImageNet. Recently, the Swin Transformer ^[17] architecture achieves the best performance, and it shows a good speed-accuracy trade-off on various vision tasks including image classification. But the Transformer structure is rarely combined with CNN structure.

This paper focuses on the integration of the Transformer and CNN, and combines this structure with UNet to apply it to medical image segmentation. It can not only use the Transformer to obtain global features, but also use CNN to fuse local features, and realize information exchange between patches. The rest of this paper is organized as follows. In Section 2, the related work of CNN and Transformer is introduced briefly. The ConvWin Transformer structure and the integration of this structure into the U-shaped network in UNet are elaborated. Also, the improvements based on the patch expanding layer and various loss functions are described. In Section 3, the experiments and results are discussed. Finally, Section 4 concludes this paper.

2.   Materials and methods

2.1.   Datasets

Hubmap - Hacking the kidney identify glomeruli in human kidney tissue images ^[18]: The Hubmap dataset includes 11 fresh frozen and 9 Formalin-Fixed Paraffin-Embedded (FFPE) PAS kidney images. Among 15 labeled samples publicly, there are Glomeruli FTU annotations for all 20 tissue samples. In this paper, to verify the effectiveness of the algorithm, only the 15 labeled samples are used, where 10 samples are divided into the training set and 5 samples into the testing set. Each picture is cut into small pieces, and the effective area size of each piece is $1024\times 1024.$ See Figure 1A for the specific data pretreatment process. After this, there are 5867 pieces in the training set and 2599 pieces in the test set. The average Dice-Similarity coefficient (DSC) is used to evaluate our method on this dataset.

Synapse multi-organ segmentation dataset ^[19]: This dataset consists of 8 abdominal organs (aorta, gallbladder, spleen, left kidney, right kidney, liver, pancreas, spleen, stomach). In this paper, the use of the training or testing images is the same as that in TransUNet ^[12]. They have used 30 abdominal CT scans in the MICCAI 2015 Multi-Atlas Abdomen Labeling Challenge, with 3779 axial contrast-enhanced abdominal clinical CT images in total. Each CT volume consists of $85 \sim 198$ slices of $512\times 512$ pixels, with a voxel spatial resolution of $([0.54 \sim 0.54]\times[0.98 \sim 0.98]\times [2.5 \sim 5.0])$ mm$^3$. See Figure 1B for the specific data pretreatment process. Similarly, the average DSC and average Hausdorff Distance (HD) are used to evaluate our method on this dataset.

2.2.   Related work

CNN-based UNet: CNN is the first automatic learning algorithm that can successfully train the multi-layer network model in a real sense. The network structure of CNN is mainly designed for image pretreatment. With the development of deep CNN, AlexNet ^[20] is proposed for image classification. Since then, deeper and more effective convolutional neural architectures have been proposed, such as VGG ^[21], GoogleNet ^[22], ResNet ^[23], DenseNet ^[24], and EfficientNet ^[25]. At the same time, UNet, which has been widely used in medical images, comes into being. The original UNet creatively combines the U-shaped structure and convolution network based on FCN to realize the semantic segmentation of medical images at the pixel level. Subsequently, UNet++, UNet3+ add the information interaction between encoder and decoder based on its structure. UNet++ introduces the dense connection into the U-shaped structure, while UNet3+ constructs full scale skip connections to realize the multi-scale information interaction between encoder and decoder. In addition, Res-UNet introduces the resblock structure in ResNet into UNet to make it have stronger feature understanding ability. U$^2$Net creatively replaces each block in UNet with a U-shaped structure. At present, UNet based on CNN still has important applications and research in different fields. For example, the latest KUB-UNet ^[9] applies UNet to the segmentation of the organs of urinary system and achieves the best results.

Transformer-based UNet: Recently, Transformer has been introduced into the field of medical image segmentation in many studies, and the Transformer structure was combined with UNet ^{[12,13,14,15]}. Among them, Trans Unet ^[12] uses CNN to extract shallow features and uses Transformer to extract depth features. MT UNet ^[14] further constructs the mixed Transformer module based on Transformer structure and takes it as the encoder and decoder with deep features. TransClaw U-Net ^[15] also uses the Transformer structure to extract deep-seated features. The difference is that the algorithm additionally performs the same level of upsampling on the features extracted by the Transformer structure and combines the results with the features of the original encoder and decoder to use concatenation. Finally, SwinUNet ^[13] is based on Swin Transformer ^[17], and it uses pure Transformer structure for feature extraction. At the same time, it constructs a patch expanding layer that do not depend on convolution to realize upsampling based on Transformer.

2.3.   Methods

Our research focuses on the integration of Transformer and CNN. First, the ConvWin Transformer structure is constructed in this paper, which refers to the W-MSA structure in Swin and combines with the convolution. The convolution is added after the window attention structure. Then, the ConvWin Transformer structure is integrated as the basic module into the U-shaped structure that is commonly used for medical image segmentation, and the new network is called ConvWin-UNet. Extensive experiments demonstrate that the Transformer structure can be integrated well with CNN, and the use of group convolution is more effective than the use of traditional CNN. Meanwhile, the order of LayerNorm and the convolution has an important impact on the experimental results. Besides, the convergence can be faster if the linear layer in the upsampling of the decoder is replaced with the convolution layer, but the convergence is sometimes not ideal.

2.3.1.   Architecture overview & ConvWin Transformer block

The overall architecture of the proposed ConvWin-UNet is presented in Figure 2A, which mainly includes the encoding part and the decoding part. As shown in Figure 2A, the encoding part and the decoding part are divided into four stages. The first stage of the encoding part is composed of a patch embedding layer ^[11] and two ConvWin Transformer blocks. In the other stages of the encoding part, the patch embedding layer is replaced with the patch merging layer. Every stage of the decoding part consists of two ConvWin Transformer blocks and a patch expanding layer. Between every stage of the encoding part and the decoding part, the low-level features are connected through skip connections. As shown in Figure 2B, the skip connections are composed of 1-layer LayerNorm (LN) and 1-layer Linear, which is different from CNN.

The basic unit of ConvWin-UNet is the ConvWin Transformer block (as shown in Figure 2C), which is mainly constructed by convolution and window based multi-head self-attention module (W-MSA). Here, the W-MSA divides the input feature map into $M\times M$ windows, then uses MSA for each window, and finally reassembles the results obtained from each window into a complete output feature map. Figure 2D shows the specific implementation form of this structure (taking $8\times 8$ feature map and $2\times 2$ windows as examples). Further, the ConvWin Transformer block is composed of an LN layer, a window based multi-head self-attention module (W-MSA), a group convolution, a residual connection, and a MLP block which contains two linear layers and uses GELU as the activation function. Different from Swin Transformer, this paper uses a convolution layer (without activation function and BN) after each multi-head self-attention layer and realizes the interaction between adjacent patches without the shift window in Swin Transformer. And through the experiments, it is found that the group convolution (GC) has a better effect than the ordinary convolution, and adding a LayerNorm module before the convolution further improves the effectiveness of the model. Similarly, ConvWin Transformer blocks can be formulated as:

where $\hat{z}^l$ and $z^l$ represent the outputs of the W-MSA module and the MLP module of the $l^{th}$ block, respectively. The self-attention is calculated as follows, which is the same as the previous works ^[10,11]:

where $Q, K, V\in \mathbb{R}^{M^2\times d}$ are the query, key, and value matrices; $M^2$ is the number of patches in a window, and $d$ is the query or key dimension. The matrix $B$ is taken from the bias matrix $\hat{B}\in \mathbb{R}^{(2M-1)\times(2M-1)}$.

2.3.2.   Encoder

In this paper, the encoder part is divided into four stages. The first stage is composed of a patch embedding layer and 2 ConvWin Transformer blocks, and the other stages consist of a patch merging layer and 2 ConvWin Transformer blocks. Among them, the patch embedding layer is composed of a patch partition layer and linear embedding layer, which is the same as that in Swin Transformer. It concatenates the features of each group of $4 \times 4$ neighboring patches and then passes the 48 dimensional concatenated features through a linear layer to project them to an arbitrary dimension (denoted by $C$). The patch merging layer is similar to the patch embedding layer, and it concatenates the features of each group of $2\times 2$ neighboring patches. In this way, the number of patchs reduces to $1/4$ of the original, and the output dimension is set to $2C$.

For an input image with a dimension of $H \times W \times 3$, the dimension is reduced to $\frac{H}{4}\times \frac{W}{4}\times C$ after the first stage. In this paper, $C = 96$. Each of the remaining stage performs $2\times$ downsampling. Thus, the feature dimensions of the output are respectively $\frac{H}{8}\times \frac{W}{8}\times 2C$, $\frac{H}{16}\times \frac{W}{16}\times 4C$, and $\frac{H}{32}\times \frac{W}{32}\times 8C$ after the other stages. Finally, the feature vector with a dimension of $\frac{H}{32}\times \frac{W}{32}\times 8C$ is obtained through the encoder.

2.3.3.   Decoder

Correspondingly, the decoder part is also divided into four stages in this paper. Each stage is composed of a patch expanding layer and 2 ConvWin Transformer blocks. Similar to that described in the previous paragraph, the features of each stage are fused with its corresponding encoder features through skip connections. For the patch expanding layer, several comparison structures have been constructed in this paper, as shown in Table 1. The main differences between these structures are 1) the order of LayerNorm operations; 2) whether using linear or convolution to expand feature dimensions; 3) whether using realrange or pixelshuffle for downsampling. The specific effects of these structures will be compared in Section 3.3. It should be noted that the patch expanding layer of the last stage uses $4\times$ upsampling, the others use $2\times$ upsampling.

Therefore, the dimension of the feature vector obtained by the encoder changes from $\frac{H}{32}\times \frac{W}{32}\times 8C$ to $\frac{H}{16}\times \frac{W}{16}\times 4C$, $\frac{H}{8}\times \frac{W}{8}\times 2C$, and $\frac{H}{4}\times \frac{W}{4}\times C$ respectively after the first three stages of the decoder. After the last stage, a feature vector with a dimension of $H\times W\times N$ is obtained, where $N$ is the number of classes.

2.3.4.   Loss functions & multi-stage outputs

In this paper, some experiments are conducted on multi-stage outputs. Based on the structure shown in Figure 2, by using the same patch expanding layer, the output of multiple stages in the decoder is directly upsampled as the output of the network structure. The ground truth is denoted as $P$, and the output of each layer is denoted as $Q^i$, where $i$ indicates the i-th stage.

For the loss function, a combination of various loss functions is used for experiments, including CrossEntropy Loss ^[26], Dice Loss ^[27], MSSIM (Mean Structure Similarity Index Measure) ^[28,29].

For image $x$ and image $y$, denote $l (x, y)$ as the mean value to estimate brightness, $c(x, y)$ as the variance to estimate the contrast, and $s (x, y)$ as the covariance to estimate the structural similarity.

Specifically,

where $C_1, C_2, C_3$ are positive constants. $C_1 = (K_1 L)^2$, $C_2 = (K_2 L)^2$, and $C_3 = C_2/2$, where the default values of $K_1$, $K_2$, and $L$ are used, i.e., $K_1 = 0.01$, $K_2 = 0.03$, $L = 2^B-1$ ($B$ is the bit depth, and $B = 256$). Denoting CrossEntropy Loss as $L_{ce}$, Dice Loss as $L_{dice}$ and MSSIM as $L_{mssim}$, we have:

where $H, W, D$ is the height, width and depth (the output channel size) of the image, $\alpha_1 = \beta_1 = 0.0448$, $\alpha_2 = \beta_2 = 0.2856$, $\alpha_3 = \beta_3 = 0.3001$, $\alpha_4 = \beta_4 = 0.2363$, $\alpha_5 = \beta_5 = \gamma_5 = 0.1333$. Thus, the loss function of each stage is:

where $w_{ce}$, $w_{dice}$, $w_{mssim}$ are weights for the loss functions, and they are set to 1 by default. If a loss is not used, the corresponding weight is set to 0. Then, the total loss function is

where $w_i$ is the weight of the i-th stage, and it is a super parameter; $n$ is the number of the output stages. If multi-stage outputs are not used, $n = 1$.

3.   Experiments and results

3.1.   Implementation details

In the experiment on the Hubmap dataset, more than 50% of the data is cut around the effective slice to prevent the black edge caused by rotations in data expansion. During the test, only the $1024\times 1024$ effective area in the middle of the data is taken. Also, a comparative experiment is conducted on the use of data augmentation for each algorithm. The size of the input image is resized to $512\times 512$ when it is input to the network. All experimental model parameters do not use any pretraining model but only use the same initialization method. During the training period, the batch size is set to 16, and the Adam optimizer with a base learning rate of 0.001 and a weight decay of 1e-4 is used to optimize our model for backpropagation. The default number of training epochs is 200. On this dataset, four pieces of Nvidia RTX2080Ti GPU are used to train each model.

In the experiment on the Synapse dataset, to keep the working environment the same, the input image size of $224\times 224$ is used. Meanwhile, all models are trained with the SGD optimizer with a learning rate of 0.05, the momentum of 0.9, and the weight decay of 1e-4. The default batch size is 24, and the default number of training epochs is 150. Besides, only one piece of Nvidia RTX2080Ti GPU is used to train the models.

In the experiments, this paper compares the effects of whether using group convolution on the models in encoder and decoder respectively. At the same time, the performance differences between linear layer, convolution layer and group convolution layer are compared in the selection of upper sampling layer. Here, the structures of the comparison models in the experiments are as follows:

● 1. Encoder:

– unet: Original UNet encoder structure,

– swin: Original Swin Transformer structure,

– cnvwin: Convwin Transformer structure constructed in this paper,

– group_convwin: Group convlution used in the Convwin Transformer structure.

● 2. Decoder:

– unet: Original UNet encoder structure,

– cnvwin: Convwin Transformer structure constructed in this paper,

– group_convwin: Group convlution used in the Convwin Transformer structure.

● 3. Upsampling in decoder

– linear up: Linear layer used in upsampling (Table 1B),

– conv up: convlution layer used in upsampling (Table 1D),

– group conv up: group convlution layer used in upsampling (Table 1E).

3.2.   Experiment results on the hubmap dataset

In this section, experiments are conducted on the combinations of the encoder, decoder, and loss function. It should be noted that there are 2 groups of comparison models construct, namely the base model and the plus model. The network layer used in the encoder stages of the base model is [3, 4, 6, 3], while that of the plus model is [4, 8, 8, 4]. Moreover, in the decoder, the output channel of each stage is set to 64, and the network layer is set to 1.

In Figure 3, several combinations of the encoder and decoder are compared. The ConvWin structure constructed in this paper has certain advantages. Meanwhile, in the comparison of encoders (Table 2), the performance differences brought by different encoders with the same decoder structure are investigated. The results show that the ConvWin structure has advantages over UNet and Swin. From Table 2, it can also be seen that when the decoder is further replaced with ConvWin structure, the performance is more significantly improved. In addition, it can be seen from Figure 4 that ConvWin structure is more sensitive to data augmentation. When using data augmentation, the performance improvement of ConvWin structure is more obvious than that of the original UNet. It can also be seen here that the structure constructed in this paper needs additional data augmentation to improve the spatial understanding ability (such as rotation invariance).

Besides, the use of different decoders leads to different effects. As shown in Table 3, the effects of using different decoders and the same ConvWin structure as the encoder are investigated. The results show that using the (Group) ConvWin structure as the decoder has advantages, especially compared with the original UNet structure, the ConvWin structure constructed in this paper has been improved to different degrees under different encoders. In addition, it is not good to use group ConvWin in all places. Especially in the patch expanding layer, group ConvWin is not a better choice than ConvWin. As can be further seen from Figure 5, there are differences in using and not using the group convolution in the patch expanding structure of the decoder. The expanding layer with the group convolution performs slightly better than that without the group convolution in the case of using the data expansion, but the result is the opposite without the data expansion. However, it can be seen from the experimental results that there is no significant difference between using and not using the group convolution in expanding. In addition, the experimental results of expanding on the Synapse dataset will be further analyzed.

Furthermore, on the Hubmap dataset, an experiment is performed to compare the loss functions and the pre-trainings of ImageNet. As shown in Figure 6A, the use of MSSIM improves the performance slightly but the use of Dice significantly reduces the performance of the model. Meanwhile, the use of multi-stage outputs ($n = 4$) improves the model performance slightly, and the model performance is also more stable. In this paper, the models pre-trained for 20 epochs on the Imagenet-1K dataset are used as pre-training models to train the corresponding segmentation model from scratch. The results are shown in Figure 6B. Pre-training does not significantly improve the performance of the model. However, the training curve in Figure 6C shows that the pre-training model is helpful for convergence.

3.3.   Experiment results on the synapse dataset

To make a fair comparison of similar methods, the experiments are also conducted on the Synapse dataset. The network structure described in Section 2 is used here, and the loss function is Dice + CE (keep up with TransUNet). As shown in Table 4, our method achieves the best performance among the algorithms taken for comparison. The segmentation results of different methods on the Synapse dataset are shown in Figure 7. Our method handles the details more carefully. This is attributed to the added convolution module, which increases the perception of local features.

Furthermore, further experiments are conducted on the expanding structure (see Table 1 specifically). As shown in Table 5, compared with method A, method B can significantly improve the performance by placing Layernorm before Linear in the expanding structure. However, the use of convolution in the expanding structure reduces the performance, which is quite different from the results obtained on the Hubmap dataset. However, the training curve in Figure 8 indicates that the use of convolution in the expanding structure makes the model converge faster, and the loss value is lower (It may be easier to fit in the training set after adding the convolution).

4.   Discussion

This paper focuses on the integration of Transformer and CNN and combines this structrue with UNet to apply it to medical image segmentation. Based on this, this paper constructs a ConvWin block, and a variety of patch expanding blocks to realize upsampling. It can not only use Transformer to obtain global features, but also use CNN to fuse local features, and realize information exchange between patches. Finally, a new loss function is constructed with MSSIM. At the same time, experiments are carried out on several datasets, and good results are obtained.

In addition, this paper also has some shortcomings and can be further studied. 1) For different datasets, the effect of the patch expanding blocks constructed in this paper are different, and more stable and universal patch expanding block can be further explored in future research. 2) It seems that the understanding of the global information needs to be further improved. Although SSIM is added as part of the loss, the local details need to be further refined and further optimization needs to be made for SSIM. 3) The initialization weight can also be further studied. In the subsequent research, the Self-Supervised Learning algorithm can be further tried to improve the robustness of initialization parameters.

Acknowledgments

The work is support by Fundamental Research Funds for the Provincial Universities of Zhejiang: GK219909299001-025; Scientific Research Foundation of Hangzhou Dianzi University: KYS075621028; Natural Science Foundation of China: 12001144; Zhejiang Provincial Natural Science Foundation of China: LQ20A010007. We would like to thank the referees for their careful reading and helpful comments and suggestions which improved the presentation of the paper.

Conflict of interest

All authors declare no conflicts of interest in this paper.