Automatic medical image segmentation method of U-shaped network based on channel mixing and self-attention mechanism

Abstract

The invention discloses a medical image automatic segmentation method of a U-shaped network based on channel mixing and a self-attention mechanism. According to the method, channel grouping and a disturbance strategy are introduced into a self-attention mechanism, so that the problems of channel redundancy, high operation cost and the like of a traditional self-attention method are effectively relieved; and the capability of extracting local space information and global channel information by the convolutional layer and the capability of extracting long-distance space dependence and modeling local channel information by the self-attention layer are combined more closely, so that the method has more comprehensive modeling capability for features. According to a network used in the method, firstly, an input medical nuclear magnetic resonance image or a computed tomography image passes through a down-sampling layer to preliminarily extract feature information and reduce the spatial dimension of the image, then a feature map is sent into a multi-level symmetric encoder-decoder network, and each level comprises a designed module; the module structure comprises a patch embedding module formed by depth separable convolution, a self-attention layer based on channel grouping and channel disturbance, and a bottleneck convolution layer formed by convolution layers of different dimensions. Wherein the patch embedding module is used for extracting position information, and the middle self-attention layer is used for extracting a long-distance dependency relationship and extracting local channel information through channel grouping and disturbance. The bottleneck convolution layer is used for exploring inductive bias and aggregating global channel information. Besides, in the jump connection process of the decoder, an efficient fusion module of channel and space cross attention is constructed, and the recovery capability of the feature information is further enhanced. The provided model is evaluated on a plurality of common data sets, and experimental results show that the method makes full use of channel information and spatial information between features, obtains a better segmentation result compared with an existing method, can be used for improving the performance of medical image segmentation, and has important application prospects and practical significance.

Description

Technical Field

[0001] The present invention relates to an automatic medical image segmentation method, in particular to an automatic medical image segmentation method based on a U-shaped network with channel mixing and self-attention mechanism. Background Art

[0002] Convolutional neural networks (CNNs) have been shown to be very effective in medical image segmentation, thanks to their strong ability to comprehensively model channel information and understand local spatial relationships. However, due to the size of the receptive field, traditional convolutions have limitations in capturing long-range dependencies and global context. In recent years, Transformers have shown good prospects in medical image segmentation due to their ability to model long-range dependencies and implicit spatial relationships. In the Transformer model, the self-attention design has attracted much attention, especially the local window self-attention, which is gradually becoming a key component in computer vision through its efficient computation and excellent spatial information capture ability. However, despite the significant progress made by Transformer, current research has mainly focused on the spatial dimension of self-attention, but lacks attention to the channel dimension. Among them, the widely used local window self-attention mechanism shares weights in the channel dimension, which limits its channel modeling ability. In addition, the current multi-head self-attention mechanism is usually implemented by splitting the channel dimension, but due to the lack of interaction between channel information, the channel features are not fully utilized. Some studies have found that some attention heads tend to learn similar linear projections, resulting in redundant attention maps. Therefore, how to effectively alleviate the limitations of self-attention in channel feature modeling still needs further exploration. In recent years, the following methods have emerged:

[0003] (1) Method based on sequential splitting of attention heads: By performing serial calculations on each attention head and inputting the calculation results of the previous head into the calculation of the next head, this serial cascade method can increase the information interaction between different attention heads, thereby reducing channel redundancy and improving computational efficiency.

[0004] (2) Method based on attention head grouping and feature aggregation: By grouping the attention heads and combining the outputs of the last multiple attention heads, the designed module is used to aggregate feature channels. This method can fully utilize the characteristics of multi-head attention parallel computing, and at the same time aggregate the channel feature information from multiple heads through the final feature fusion module.

[0005] Both of the above approaches have their limitations, resulting in limited performance of the model in medical image segmentation tasks. Although the serial cascade method promotes information exchange through sequential interaction, this method destroys the parallelism of multi-head attention calculations and reduces computational efficiency. On the other hand, although the use of a feature aggregation module at the end helps to fuse the feature information of multiple heads. However, considering that in the current Transformer model, the use of multi-head attention layers is usually a series of continuous processes. In this process, there is still a lack of information interaction between multiple heads. Therefore, it is necessary to design an appropriate structure that can make good use of the advantages of parallel computing of multi-head attention layers while promoting information exchange between multiple heads. The present invention aims to propose a new method to overcome this difficulty and bring better reasoning results. Summary of the invention

[0006] The purpose of the present invention is to provide a method for automatic segmentation of medical images based on a U-shaped network with channel mixing and self-attention mechanism. The technical solution to achieve the purpose of the present invention is to extract features of sampled image blocks through a model constructed based on the idea of ​​self-attention layer with channel grouping and channel perturbation, including the following steps:

[0007] Step 1: Preprocess the data based on the existing medical image dataset.

[0008] Step 2: Divide the processed data into training set, validation set and test set.

[0009] Step 3: Train the designed segmentation network model based on the training set:

[0010] Step 3.1: Use a multi-level encoder network to extract graphic features.

[0011] Step 3.2: In each encoder, the position information is first extracted through the front patch embedding module. The middle self-attention layer based on channel mixing and channel shuffling extracts the global spatial information and local channel information. Then, the rear convolutional layer extracts the local spatial information and integrates the global channel feature information.

[0012] Step 3.3: Use a decoder structure that is symmetrical to the encoder to recover features, and use an efficient fusion module to pass detailed information in the encoder level.

[0013] Step 3.4: After calculating the loss based on the loss function, use the gradient descent algorithm to converge the model.

[0014] Step 4: Repeat step 3 to train the segmentation network model.

[0015] Step 5: Use the trained segmentation network to test the test data.

[0016] Compared with the prior art, the present invention has the following significant advantages:

[0017] (1) In order to solve the channel redundancy of multi-head attention, a strategy based on channel grouping and channel perturbation is proposed. This strategy can fully utilize the advantages of parallel computing between multi-head attention and realize information exchange between multiple heads without introducing additional parameters.

[0018] (2) A global channel fusion module that combines convolutions of different dimensions. This module is composed of convolutions of different dimensions and is used at the end of each encoder stage. It can not only more comprehensively fuse feature information from different heads, but also further extract local spatial information of the model. This combination is conducive to combining the advantages of convolutional networks and Transformer networks.

[0019] (3) An efficient feature fusion module with cross-spatial and channel attention is added to further enhance the recovery of feature information in the decoder.

[0020] The present invention is further described in detail below in conjunction with the accompanying drawings. BRIEF DESCRIPTION OF THE DRAWINGS

[0021] Figure 1 It is the overall flow chart of the automatic medical image segmentation method of the U-type network based on channel mixing and attention mechanism of the present invention and the internal details display.

[0022] Figure 2 It is the experimental result of comparing the model of the method of the present invention with several existing models on a multi-organ medical image dataset.

[0023] Figure 3 The model of the method of the present invention is compared with the results of several models on actual test data.

[0024] Figure 4 It is a summary of the specification of the present invention, which describes the overall process logic of the present invention. DETAILED DESCRIPTION

[0025] Combination Figure 1 The present invention provides a method for automatic segmentation of medical images based on a U-shaped network with channel mixing and attention mechanism, comprising the following steps:

[0026] Step 1: Preprocess the data based on the existing medical image dataset.

[0027] Step 2: Divide the processed data into training set, validation set and test set.

[0028] Step 3: Train the designed segmentation network model based on the training set:

[0029] Step 3.1: Use a multi-level encoder network to extract graphic features.

[0030] Step 3.2: In each encoder, the position information is first extracted through the front patch embedding module. The middle self-attention layer based on channel mixing and channel shuffling extracts the global spatial information and local channel information. Then, the rear convolutional layer extracts the local spatial information and integrates the global channel feature information.

[0031] Step 3.3: Use a decoder structure that is symmetrical to the encoder to recover features, and use an efficient fusion module to pass detailed information in the encoder level.

[0032] Step 3.4: After calculating the loss based on the loss function, use the gradient descent algorithm to converge the model.

[0033] Step 4: Repeat step 3 to train the segmentation network model.

[0034] Step 5: Use the trained segmentation network model to test the test data.

[0035] The present invention is further described in detail below with reference to examples:

[0036] The present invention uses medical magnetic resonance images or computer tomography images as input, and generates results after automatic segmentation of lesions through data preprocessing and feature extraction networks.

[0037] The process of the present invention is as follows Figure 1 shown.

[0038] (1) The Synapse dataset used for training and testing of the present invention is a public multi-organ segmentation dataset. There are 30 cases of abdominal clinical CT enhancement in the dataset. According to the official recommended settings, 18 cases are used for training and 12 cases are used for testing. The annotations of each image include 8 abdominal organs, including the aorta, gallbladder, left kidney, right kidney, liver, pancreas, spleen and stomach. At the same time, each case data is processed into multiple two-dimensional slices containing targets of 224 pixels * 224 pixels as input. Since the amount of data involved in the training is small, in order to obtain a better model, we augmented the source data, including randomly selected image flipping, image rotation, adding Gaussian noise, blurring and other operations.

[0039] (2) Figure 1 (a) shows the model architecture diagram of the present invention.

[0040] First, the input image is passed through a convolutional downsampling layer for preliminary feature extraction. The network structure consists of two simple convolutional layers, which reduces the input image size to 1 / 4 of the original image and the feature channel dimension to 128.

[0041] Convolution-self-attention combined feature extractor: The encoder is mainly divided into three parts. The first part is constructed by depth-wise separable convolution, such as Figure 1 As shown in (e), in this module, the position information is extracted through the depthwise separable convolution. Compared with the traditional convolution operation, the depthwise separable convolution decomposes the standard convolution into two separate steps: depthwise convolution and pointwise convolution. Its advantage is that it reduces the number of model parameters and the amount of calculation. It uses fewer parameters to implement the convolution operation, thereby reducing the complexity of the model and improving the calculation efficiency. The second part consists of a multi-axis attention module based on channel grouping and channel perturbation strategies, such as Figure 1 (d) and Figure 1 As shown in (e), the channel grouping and perturbation strategy aims to split the multi-head attention and perform parallel calculations. After calculation, the multiple attention heads are perturbed by channels to disrupt the order of feature channels and then passed to the next multi-head attention layer. This can effectively promote the flow of information between multiple attention heads. The attention calculation method used in this method is multi-axis attention, which is another implementation of the local window attention method. The third part, such as Figure 1 As shown in (c), it consists of multiple convolutions of different dimensions, and adopts a bottleneck design to fully integrate the feature channel information from multiple heads after the self-attention layer calculation, and further supplement the local spatial information. The symmetrical decoder structure uses the same feature extraction module, and the high-order features are restored to low-dimensional feature maps through upsampling of convolutions. The efficient feature fusion module proposed in this paper is used to enhance the recovery of information, and finally the segmentation result is generated after upsampling.

[0042] Efficient feature fusion module: The semantic gap between the encoder and decoder levels limits the segmentation performance. This method proposes an efficient feature fusion module based on cross-spatial attention and channel attention. By better fusing the information of the encoding layer and the decoding layer in space and channels to help the decoding layer recover information, to achieve this, we retain the jump connection part of the original U-Net structure before each downsampling layer and construct an additional efficient feature fusion module. Figure 1 As shown in (b), in the left branch, the encoder-decoder features are concatenated and spatially aggregated by average pooling, and then passed through a 1×1 convolution layer to enhance and select information channels. In the right branch, the encoder and decoder features are separately 1×1 convolved before summing, ReLU activated, and fused through another 1×1 convolution and generating spatial weights through a sigmoid activation function. Finally, channel and spatial information are fused through element-wise multiplication.

[0043] Loss function: The method of the present invention uses the sum of the DICE index and the cross entropy as the loss function to perform model gradient descent convergence.

[0044] (3) Repeat step 3 to train the segmentation network model. The method of the present invention uses average accuracy and segmentation error as indicators to judge the quality of the model.

[0045] (4) The trained segmentation network model is used to test the test data. The method of the present invention and several existing methods (UNet, R50 AttUNet, TransUNet, Swin Unet, MISSFormer, DAEFormer, TransCASCADE, PAG-TransYnet) are compared on the test set. Based on the experimental results, the average accuracy and segmentation error of the present method achieved good experimental results. An average regional measurement accuracy of 84.82% was achieved on the Synapse dataset, which is better than other Transformer-based methods, such as Swin-Unet and MISSFormer with pure Transformer, and TransUNet with multiple Transformer modules stacked in the deep layer of the encoder. It is worth noting that the present method achieved the best results in four organs, namely gallbladder, left kidney, right kidney, and pancreas, and has better effects than other models in distinguishing targets with low contrast and fuzzy boundaries (such as liver, pancreas, and stomach). Benefiting from the self-attention module of channel grouping and channel perturbation, the convolution module constructed in a bottleneck manner, and the efficient feature fusion module based on the cross-spatial channel attention mechanism, this model can achieve better medical image segmentation results.

Claims

1. A method for automatic segmentation of medical images based on channel mixing and self-attention mechanism, which is characterized by the following steps: Step 1: Preprocess the data based on the existing medical image dataset. Step 2: Divide the processed data into training set, test set and validation set. Step 3: Train the designed segmentation network model based on the training set: Step 3.1: Use a multi-level encoder network to extract graphic features. Step 3.2: In each encoder, the position information is first extracted through the front patch embedding module. The middle self-attention layer based on channel mixing and channel shuffling extracts the global spatial information and local channel information. Then, the rear convolutional layer extracts the local spatial information and integrates the global channel feature information. Step 3.3: Use a decoder structure that is symmetrical to the encoder to recover features, and use an efficient fusion module to pass detailed information in the encoder level. Step 3.4: After calculating the loss based on the loss function, use the gradient descent algorithm to converge the model. Step 4: Repeat step 3 to train the segmentation network model. Step 5: Use the trained segmentation network to test the test data.

2. The method for automatic segmentation of medical images based on channel mixing and self-attention mechanism according to claim 1, characterized in that: The specific method of training the network model described in step 3 is as follows: The network is a multi-order U-shaped network based on channel mixing and self-attention mechanism. First, the image is convolutionally downsampled to perform preliminary feature extraction and reduce its computational scale. The front patch embedding module in each encoder extracts position information, the middle self-attention module based on channel grouping and perturbation extracts global spatial information and local channel information, and the rear global channel module extracts local bias and integrates feature channel information. The symmetric decoder structure uses the same feature extraction module, and the high-order features are restored to low-dimensional feature maps through convolution upsampling. The encoder information is passed using an efficient fusion module to help restore image information, and the segmentation result is finally generated after upsampling.

3. The method for automatic segmentation of medical images based on channel mixing and self-attention mechanism according to claim 1, characterized in that: The feature extraction module in each encoder of the segmentation network described in step 3.2 is as follows: The patch embedding module composed of deep separable convolution is used to advance the position information. The deep separable convolution has the inductive characteristics of convolution on local positions, and has less computational complexity than ordinary convolution. The features are then passed through a multi-axis attention layer built based on channel grouping and channel perturbation strategies to extract global spatial information and local channel information and reduce channel redundancy. Finally, a module composed of bottleneck convolution blocks is used to fuse global channel information and extract local spatial information.

Images