U-shaped image segmentation method based on convolution enhanced cross self-attention transformer

Abstract

The application discloses a U-shaped image segmentation method based on a convolution enhanced cross self-attention transformer, which comprises a convolution embedding module, a convolution merging module and a convolution enhanced cross self-attention transformer module as an encoder, and a convolution upsampling module and a convolution decoding module as a decoder. The image block is divided into horizontal and vertical strips by the convolution enhanced cross self-attention transformer module, and feature merging is carried out in the channel dimension, so that the perception range of the self-attention is greatly enhanced, the local coding module is fused into the deformation to enhance the local information processing capacity of the model, the feature representation of the overall style and contour of the object can be provided, the model based on a convolution neural network (CNN) has more stable recognition performance and more accurate segmentation effect for the object with a local change.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of image processing, and particularly relates to a U-shaped image segmentation method based on a convolution-enhanced cross self-attention transformer. BACKGROUND

[0002] The transformer-based model is widely used in image analysis due to its ability to obtain long-term context information. The transformer can provide feature representation of the overall style and contour of an object, and has more stable recognition performance than the convolutional neural network (CNN) based model for objects with severe local changes. However, the existing segmentation network based on the transformer architecture has the following problems: 1) there is a big difference between the feature extraction method of CNN and the transformer model, and the fusion of the feature encoder and the decoder will cause the problem of feature mismatch. 2) When the input of the network is a high-resolution image such as a medical image, it is difficult to establish a model with fewer parameters while maintaining performance due to the high computational complexity of self-attention. 3) The performance of the transformer on small-scale data sets is often poor because it lacks translation invariance, resulting in insensitivity to the position and size of the identified object. SUMMARY

[0003] The present application proposes a U-shaped image segmentation method based on a convolution-enhanced cross self-attention transformer to overcome the defects of the existing transformer architecture feature encoder that cannot obtain detailed local information and the limitation of global perception. The method divides the image into horizontal and vertical strips through the convolution-enhanced cross self-attention transformer, and merges the features in the channel dimension, greatly enhancing the perception range of self-attention while integrating the local encoding module into the transformer to enhance the local information processing capability of the model. The method can provide feature representation of the overall style and contour of an object, and has more stable recognition performance and more accurate segmentation effect than the convolutional neural network (CNN) based model for objects with severe local changes.

[0004] The present application is implemented through the following technical solutions:

[0005] The application relates to a U-shaped image segmentation method based on a convolution-enhanced cross self-attention transformer, which comprises a convolution embedding module, a convolution merging module and a convolution-enhanced cross self-attention transformer as an encoder and a convolution upsampling module and a convolution decoding module as a decoder, wherein the encoder is composed of five sequentially connected sub-encoding blocks, the decoder is composed of four sequentially connected sub-decoding blocks, the first sub-encoding block is directly connected with the fourth sub-decoding block, the second sub-encoding block and the third sub-decoding block, the third sub-encoding block and the second sub-decoding block and the fourth sub-encoding block and the first sub-decoding block are connected through feature conversion units for eliminating semantic gaps, the fifth sub-encoding block is directly connected with the first sub-decoding block, the fourth sub-decoding block is provided with the convolution upsampling module and a dimension mapping module for splicing coding, the dimension mapping module generates a segmentation mask through dimension mapping of the convolution upsampling of the fused feature map, and image segmentation is realized.

[0006] The first sub-encoding block comprises a convolution embedding module and a corresponding convolution merging module, wherein the convolution embedding module embeds an input image into a feature map and performs local coding to extract detailed information, and the convolution merging module performs convolution down-sampling on the feature map.

[0007] The second to fourth sub-encoding blocks each comprise a convolution-enhanced cross self-attention transformer and a corresponding convolution merging module, wherein the convolution-enhanced cross self-attention transformer performs local and global feature learning on feature maps at the same level, and the convolution merging module performs convolution down-sampling on the feature map.

[0008] The fifth sub-encoding block only comprises a convolution-enhanced cross self-attention transformer.

[0009] The first to fourth sub-decoding blocks each comprise a convolution upsampling module and a corresponding convolution decoding module, wherein the convolution upsampling module performs up-sampling on the feature map output by the encoder, the convolution decoding module fuses the feature map through a feature conversion unit in a jump connection mode, and performs dimension reduction and feature extraction to generate a fused feature map. BRIEF DESCRIPTION OF DRAWINGS

[0010] Figure 1 It is a system schematic diagram of the application;

[0011] Figure 2 It is a convolution-enhanced cross self-attention transformer schematic diagram;

[0012] Figure 3 It is a feature conversion unit schematic diagram;

[0013] Figure 4 It is an embodiment flowchart;

[0014] Figure 5 It is an embodiment effect schematic diagram of the application. Detailed Implementation

[0015] like Figure 1 As shown in this embodiment, a U-shaped image segmentation method based on convolutional enhancement cross-self-attention deformers is proposed. It includes: a convolutional embedding module and its convolutional merging module, and four convolutional enhancement cross-self-attention deformers and their convolutional merging modules located at the encoder end; and five convolutional upsampling modules, four convolutional decoding modules, three feature transformation units, and a dimension mapping module located at the decoder end to eliminate the semantic gap between the convolutional enhancement cross-self-attention deformers and the convolutional decoding modules. Specifically, the convolutional embedding module embeds the input image into a feature map of half its size and performs local encoding to extract detailed information; the convolutional merging module performs local encoding on the convolutional embedding module or the convolutional enhancement cross-self-attention deformer. The feature map output by the cross self-attention deformer is downsampled by 2x convolution; the convolution-enhanced cross self-attention deformer learns local and global features sequentially on the feature map at the same level; the convolution upsampling module upsamples the feature map output by the fourth convolution-enhanced cross self-attention deformer and the convolution decoding module; the convolution decoding module fuses the feature map after the convolution upsampling module with the feature map encoded by the three convolution-enhanced cross self-attention deformers and one convolution embedding module through skip connections, and performs dimensionality reduction and feature extraction; the dimension mapping module projects the fused feature map containing context information obtained after upsampling the features after convolution decoding to generate a segmentation mask.

[0016] The convolutional embedding module includes: a first convolutional coding block for high-dimensional embedding and downsampling, and a second convolutional coding block for local representation of the embedded feature map.

[0017] The convolutional coding block consists of a convolutional layer, a batch normalization layer, and an activation function Gelu.

[0018] The convolution merging module performs local information fusion on the input feature map through a built-in 3×3 convolutional layer with a stride of 2 and padding of 1, and then performs downsampling and layer normalization (LN) to accelerate the convergence speed of subsequent encoding.

[0019] like Figure 2As shown in Figure a, the convolutional enhanced cross-self-attention deformer includes: a local encoding module and a cross-self-attention module connected to its residual, a multilayer perceptron (MLP), and two layer normalization units (LNs). Specifically, the local encoding module extracts local features from the input feature map; the cross-self-attention module performs global modeling on the feature map output by the local encoding module, dividing the feature map into multiple windows and calculating self-attention within and between windows to obtain features containing global information; the MLP and LNs integrate semantic information from the feature map containing both local and global information and accelerate convergence. Specifically, the local contraction feature X obtained from the (l-1)th convolutional merging module... l-1 The input local encoding module obtains the encoded output. ,Will +X l-1 Normalizing the input LN yields the regularized feature output. ,Will Inputting the cross-attention module yields features containing global information. Finally +X l-1 The residual connection after computation is MLP(LN( +X l-1 ))+( +X l-1 This yields the local contraction feature X of the l-th convolution-enhanced cross-self-attention deformer. l .

[0020] The local encoding module includes: a depthwise separable convolutional layer DWconv with a stride of 1 and padding of 1 and a kernel size of 3×3, a batch normalization layer BN, and a 1×1 convolutional layer Conv with a stride of 1 and padding of 1, wherein: the local contraction feature X obtained by the (l-1)th convolution merging module is... l-1 The depthwise separable convolution operation DWconv(BN(Conv(X) is used. l-1 This yields the output of the local encoding portion in the l-th layer convolutional enhanced cross-attention deformer. .

[0021] like Figure 2 As shown in b, the cross-self-attention module is calculated as follows: The normalized dimension of the local encoding module in the l-th layer convolutional enhanced cross-self-attention deformer is obtained by normalizing the normalized dimension. Feature output The feature dimension is divided into horizontal and vertical groups, where C is the feature dimension, H and W are the length and width of the feature, respectively, and the vertical group attention... and level group attention Then, cross-self-attention is calculated separately, as follows:

[0022] Vertical group attention in the cross self-attention module For example, Feature map of each dimension Uniformly divided into non-overlapping vertical bands {V} of width d i1 V i2 , ..., V iM}, where M=W / d, and d is set empirically based on the size of the feature map at different stages, and then by applying V to each non-overlapping vertical band. ij Perform a linear projection, i.e., linear(V) i1 ), to obtain the corresponding query (Q) ij ), key (K) ij ) and value (V) ij ), where: linear refers to a fully connected layer with an output dimension three times that of the input, j = 1, 2, ..., M; corresponding to each non-overlapping vertical band V ij The corresponding method for calculating self-attention is as follows: Where: Softmax is the probability mapping function that maps values ​​to a probability distribution between 0 and 1; k is the scaling factor; Dwconv refers to a 3×3 depthwise separable convolution with a stride of 1 and padding of 1 used to learn the positional information of the feature map; and vertical group attention. The feature output after self-attention is the concatenation of matrices after self-attention for each vertical band. ].

[0023] Correspondingly, the level group attention hor- The calculation involves dividing the data into non-overlapping horizontal bands and performing the same process. Finally, the feature outputs of the vertical and horizontal groups after self-attention are merged into a matrix in dimension C, which is the feature containing global information obtained after passing through the cross-self-attention module. .

[0024] The convolutional upsampling module includes: a bilinear interpolation layer of 2x, a convolutional layer with a stride of 1, padding of 1, and a kernel of 3×3, a batch normalization layer (BN), and an activation function Gelu.

[0025] The convolutional decoding module is a feature module for fusing information at different scales. It includes: a convolutional layer with a stride of 1, padding of 1, a kernel of 3×3 and an output dimension that is half the input dimension; a convolutional layer with a stride of 1, padding of 1, a kernel of 3×3 and a default output dimension that is the same as the input dimension; two batch normalization (BN) layers; and the activation function Gelu.

[0026] The feature conversion unit is as follows Figure 3As shown, it includes: a 1x1 convolutional layer with a stride of 1, a BN layer and a LeakyReLU, wherein: the 1x1 convolutional layer is used to adjust the number of channels to maintain the channel consistency of the CNN and the Transformer branch; the BN layer is used to normalize the output of the Transformer branch to make it consistent with the CNN branch. The LeakyReLU provides nonlinear mapping capability and speeds up branch fusion.

[0027] The jump connection splices the output of the encoding stage after the feature conversion unit and the output of the corresponding decoding stage in the feature dimension, and the jump connection includes a sum operation in the channel dimension.

[0028] The dimension mapping module performs 2 times up-sampling on the feature map containing context information through convolution up-sampling, and then uses a projection with a stride of 1, padding of 1, a kernel of 3x3 and an output dimension corresponding to the segmentation category to obtain a segmentation mask.

[0029] As Figure 4 shown, the embodiment relates to an image segmentation method based on the U-shaped network of the convolution-enhanced cross-attention transformer, which includes the following steps:

[0030] Step 1) Dataset making: In this embodiment, prostate MRI transverse T2W imaging data is selected. The T2W images corresponding to the data of two centers are labeled by two professional clinicians, and are packaged and paired. The transition zone (TZ) and peripheral zone (PZ) of the prostate are labeled in the label, and the size of the label is the same as that of the original T2W. Figure 1 A pair of.

[0031] Step 2) Dividing the dataset: For the MRI dataset, the training set and the validation set are divided according to the ratio of 8:2. The image enhancement method is used to expand the data amount for the data of the training set, including: random rotation of 5°, random horizontal flip, random Gaussian blur and random elastic transformation.

[0032] Step 3) Network model building: as Figure 2 shown, the input image is a gray image. The image uses the center cropping method to remove the irrelevant area of the image, and is resampled to the size of 1x224x224. At the same time, the batch size of the network input is set to 24. Therefore, the dimension of the network input is 24x1x224x224. Finally, the output with a dimension of 24x3x224x224 is obtained through the network model shown in Table 1, wherein 3 is the channel dimension, which is the segmentation mask of background, PZ and TZ respectively.

[0033]

[0034] Through specific experiments, on an Ubuntu 18.04 system server with a CPU (Intel Core i7-9700K 32 GB) and two GPUs (NVIDIA GeForce RTX 2080 8 GB), Pytorch1.10 is used to implement the algorithm. In this image segmentation task, the initial learning rate is 0.01, the Adamw optimizer is used, and the Weightdecay is set to 0.0001. The Dice loss function and the cross-entropy loss function are used jointly to control the direction of back propagation, specifically:

[0035] wherein: y i j is a pixel, i is a channel, j is a pixel in an image, H and W are the width and length of the input image respectively.

[0036] Table 2. Experimental results of the embodiment on the MRI dataset.

[0037] Compared with the prior art, the convolution embedding module in the method re-encodes the feature map obtained by the upper layer at an input resolution of 1 / 2, instead of directly sending it to the deformator, solving the problem of local information loss and edge sawtooth when directly processing the input of linear embedding in a pure deformator structure network such as SwinU-Net. The convolution-enhanced cross-attention deformator fuses the unprocessed feature mapping of the previous level with the feature mapping of the proposed local encoding module through skip connection, realizing residual learning of local information, completely retaining the output of the feedforward layer, and accelerating the convergence speed of the local encoding module to prevent gradient disappearance. By introducing the cross-attention module based on the cross-shaped window, the feature map encoded by the local encoding module is combined with the residual connection to extract global features, preventing local information from being weakened in the self-attention calculation process.

[0038] The embodiment uses the MRI dataset to output the prediction result through the network. Further, the Dice coefficient (DC) and Jaccard similarity (JS) are used to evaluate the segmentation result. The results in Table 2 show that, compared with the prior art, the method achieves the best performance on the MRI dataset, with DC and JS on PZ being 79.97% and 74.30% respectively, and on TZ being 87.39% and 83.24% respectively. The DC on PZ is 3.3% higher than the second-ranked SwinU-Net, and the JS is 3.9% higher than the second-ranked SwinU-Net, while the segmentation accuracy on TZ is also better than all other methods in all cases.

[0039] The above specific implementation can be adjusted in different ways by those skilled in the art without departing from the principles and purposes of the embodiments, the protection scope of the embodiments is subject to the claims and is not limited by the above specific implementation, and each implementation within the scope is subject to the embodiments.

Claims

1. A U-shaped image segmentation method based on a convolution-enhanced cross-attention deformer, wherein the U-shaped image segmentation network used includes: The encoder comprises a convolutional embedding module, a convolutional merging module, and a convolutional enhancement cross-self-attention deformation module, while the decoder comprises a convolutional upsampling module and a convolutional decoding module. The encoder consists of five sequentially connected sub-coding blocks, and the decoder consists of four sequentially connected sub-decoding blocks. The first and fourth sub-coding blocks are directly connected. The second and third sub-coding blocks, the third and second sub-coding blocks, and the fourth and first sub-coding blocks are connected via feature transformation units to eliminate semantic gaps. The fifth sub-coding block is directly connected to the first decoding block. The output of the fourth sub-decoding block includes a convolutional upsampling module and a dimension mapping module for concatenating the encoding. The dimension mapping module generates a segmentation mask from the fused feature map through the dimension mapping of the convolutional upsampling, thus achieving image segmentation. The convolutional enhancement cross-self-attention deformation module includes: a local encoding module and a cross-self-attention deformer connected to its residual, a multilayer perceptron (MLP), and two layer normalization units (LNs). The local encoding module extracts local features from the input feature map, the cross-self-attention deformer performs global modeling on the feature map output by the local encoding module, divides the feature map into multiple windows and calculates self-attention within the windows and attention between windows to obtain features containing global information, and the MLP and LNs integrate semantic information on the feature map that contains both local and global information and accelerate convergence.

2. The U-shaped image segmentation method based on convolutional enhanced cross-self-attention deformer according to claim 1, characterized in that, The first sub-encoding block includes a convolutional embedding module and a corresponding convolutional merging module, wherein: the convolutional embedding module embeds the input image into a feature map and performs local encoding to extract detailed information, and the convolutional merging module performs convolutional downsampling on the feature map; The second to fourth sub-encoding blocks each include a convolutional enhancement cross-self-attention deformation module and a corresponding convolutional merging module, wherein: the convolutional enhancement cross-self-attention deformation module performs local and global feature learning on the feature map at the same level, and the convolutional merging module performs convolutional downsampling on the feature map; The fifth sub-encoding block only includes the convolution-enhanced cross-self-attention deformation module.

3. The U-shaped image segmentation method based on convolutional enhanced cross-self-attention deformer according to claim 1, characterized in that, The first to fourth sub-decoding blocks each include a convolutional upsampling module and a corresponding convolutional decoding module, wherein: the convolutional upsampling module upsamples the feature map output by the encoder, and the convolutional decoding module fuses the feature map through skip connections and feature transformation units, and performs dimensionality reduction and feature extraction to generate a fused feature map.

4. The U-shaped image segmentation method based on convolution-enhanced cross-self-attention deformer according to claim 1 or 2, characterized in that, The convolutional embedding module includes: a first convolutional coding block for high-dimensional embedding and downsampling, and a second convolutional coding block for local representation of the embedded feature map; The convolutional coding block consists of a convolutional layer, a batch normalization layer, and an activation function Gelu.

5. The U-shaped image segmentation method based on convolutional enhanced cross-self-attention deformer according to claim 1 or 2, characterized in that, The convolution merging module performs local information fusion on the input feature map through a built-in 3×3 convolutional layer with a stride of 2 and padding of 1, and then performs downsampling and layer normalization (LN) to accelerate the convergence speed of subsequent encoding.

6. The U-shaped image segmentation method based on convolutional enhanced cross-self-attention deformer according to claim 1, characterized in that, The self-attention within a window and the attention between windows are calculated as follows: The cross-self-attention deformer of the l-th layer normalizes the local encoding modules to obtain the normalized dimension. Feature output The feature dimension is divided into horizontal and vertical groups, where C is the feature dimension, H and W are the length and width of the feature, respectively, and the vertical group attention... and level group attention hor- Then, cross-self-attention is calculated separately.

7. The U-shaped image segmentation method based on convolutional enhanced cross-self-attention deformer according to claim 1, characterized in that, The feature transformation unit includes: a 1×1 convolutional layer with a stride of 1, a BN layer, and LeakyReLU, wherein: the 1×1 convolutional layer is used to adjust the number of channels to maintain channel consistency between the CNN and Transformer branches; the BN layer is used to normalize the output of the Transformer branch to make it consistent with the CNN branch.

8. The U-shaped image segmentation method based on convolutional enhanced cross-self-attention deformer according to claim 1, characterized in that, The dimension mapping module upsamples the feature map containing context information by a factor of 2 through convolutional upsampling, and then uses a projection with a stride of 1, padding of 1, kernel size of 3×3, and output dimension as the corresponding segmentation category to obtain the segmentation mask.

Images