An image segmentation method based on dual attention fusion

Abstract

This invention discloses an image segmentation method based on dual attention fusion. First, a dataset is constructed according to the target task. Then, a segmentation network model is built, and image samples from the dataset are input into the constructed segmentation network model for training. The segmentation network model adopts a U-shaped structure, including an encoding module, a dual attention gating fusion module, and a decoding module. The encoding module encodes the input image to obtain an initial feature map. The dual attention gating fusion module includes a multi-scale weighted channel attention branch and a global spatial self-attention branch, and the features extracted from the two branches are fused through a gating mechanism module to obtain the final feature map. The decoding module decodes the final feature map to obtain the segmented image. Finally, the image data to be segmented is input into the trained segmentation network model to obtain a high-precision segmentation result of the target image.

Description

Technical Field

[0001] This invention belongs to the technical field of image processing, specifically relating to an image segmentation method based on dual attention fusion. Background Technology

[0002] Medical imaging plays a crucial role in disease diagnosis. Accurately segmenting organs from medical images manually is extremely difficult and time-consuming, and relies heavily on personal experience. The rapid development of modern image processing and artificial intelligence (AI) provides doctors with more critical information about lesions. AI technology is improving the accuracy of disease diagnosis and reducing diagnostic time. Accurately segmenting specific organs from medical images remains a challenging task, but it is vital for clinical diagnosis.

[0003] In recent years, research on medical image segmentation based on deep learning has made some breakthroughs. However, segmenting multiple organs from medical images remains very difficult, due to the following challenges: First, multiple organs in the human body are compressed into a small space, leading to significant deformation of the same organ in different individuals. For example, the colon and pancreas in the abdomen have very different shapes and appearances in different people. Second, the scanning quality of existing scanners is limited, resulting in blurred boundaries, low contrast, and high noise. In abdominal CT slices, there is often a blurred boundary between the head of the pancreas and the duodenum, leading to limited segmentation accuracy. Third, human organs have different sizes and shapes in medical images, so we need more abstract features and high-level features extracted from different dimensions to extract semantic information from the images.

[0004] Most existing methods are image segmentation algorithms based on convolutional neural networks (CNNs). By adding new modules to the basic segmentation framework, the feature extraction capability of the network is improved, thereby obtaining higher-precision segmentation results. However, CNN-based methods are good at capturing detailed information in local features, but they are insufficient for modeling long-distance relationships in the global image. The Transformer self-attention algorithm has been proposed in the field of natural language processing and has been widely used in the field of computer vision, achieving good results. However, how to make good use of the respective advantages of convolution and Transformer is still a topic worthy of exploration. Summary of the Invention

[0005] This invention provides an image segmentation method based on dual attention fusion. The proposed cut network model with dual attention gating fusion adopts an encoder-decoder structure and adds a dual attention gating fusion module at the bottom of the model, which greatly improves the feature extraction capability of the model. The network directly outputs high-precision image segmentation results at the end, which greatly improves the accuracy of automatic segmentation and solves the technical problems of low segmentation accuracy and poor segmentation effect of existing methods.

[0006] This invention can be achieved through the following technical solutions:

[0007] An image segmentation method based on dual attention fusion includes the following steps:

[0008] S1. Construct a dataset according to the target task, wherein the dataset includes multiple image samples with pixel-level annotations for a specific target;

[0009] S2. Construct a segmentation network model and input the image samples from the dataset into the constructed segmentation network model for training.

[0010] The segmentation network model adopts a U-shaped structure, including an encoding module, a dual-attention gating fusion module, and a decoding module. The encoding module encodes the input image to obtain an initial feature map. The dual-attention gating fusion module includes a CNN-based multi-scale weighted channel attention branch and a Transformer-based global spatial self-attention branch, and fuses the features extracted from the two branches through a gating mechanism module to obtain the final feature map. The decoding module decodes the final feature map to obtain the segmented image.

[0011] The multi-scale weighted channel attention branch is used to extract inter-class response features to improve classification accuracy and obtain multi-scale feature maps. The global spatial self-attention branch is used to extract long-distance dependency features to improve localization accuracy and obtain global feature maps.

[0012] S3. Input the image data to be segmented into the trained segmentation network model to obtain the high-precision segmentation result of the target image.

[0013] Furthermore, the multi-scale weighted channel attention branch includes a cascaded multi-scale convolution operation and a weighted channel attention operation.

[0014] For the initial input feature map, on the one hand, a 1×1 convolution is first performed on the initial feature map to reduce the number of channels to 1 / 8 of the original, generating feature map t1; on the other hand, after halving the size and number of channels of the initial feature map, 3×3 convolution, 5×5 convolution and 7×7 convolution are performed in parallel, compressing the number of channels to 1 / 8 of the initial feature map, thus obtaining three more feature maps t2, t3 and t4 with 1 / 8 of the channels, for a total of four feature maps t1, t2, t3 and t4.

[0015] Then, SE channel attention operations are performed on the four feature maps t1, t2, t3, and t4 respectively, resulting in four sets of channel attention coefficients. These four sets of channel attention coefficients are then assigned four trainable weight values ​​to obtain four sets of weighted channel attention coefficients. These coefficients are then concatenated and normalized using the Softmax function. Next, they are multiplied by the corresponding channels onto the combined feature map formed by concatenating the four feature maps t1, t2, t3, and t4. After another convolution, a multi-scale feature map with half the original number of channels is finally obtained, thus realizing the extraction of multi-scale information from the initial feature map.

[0016] Furthermore, the global spatial self-attention branch includes multiple Transformer attention blocks connected in series, and each Transformer attention block includes a layer normalization module, a multi-head self-attention MSA, a multilayer perceptron MLP, and a residual module.

[0017] Furthermore, the Transformer attention block is configured with 12 blocks, the 5×5 convolution is configured with four groups, and the 7×7 convolution is configured with eight groups.

[0018] Furthermore, the gating mechanism module adopts a GRU gated recursive unit structure, including a "reset gate" and an "update gate". It uses the sigmoid function as the control activation function to enhance the beneficial features and suppress the detrimental factors in the input multi-scale feature map and global feature map, thereby fully fusing feature information and removing redundant information to obtain the final feature map.

[0019] Furthermore, the encoding module uses the residual convolutional blocks of the ResNet-50 network as the main structure, and inserts several convolution and pooling operations as connections within the main structure;

[0020] The decoding module includes three identical upsampling blocks and a segmentation map output block. The upsampling block includes one upsampling with doubled resolution, one concatenation of encoded and decoded features at the same scale, and two convolution operations, wherein the latter convolution halves the number of channels in the feature map. The segmentation map output block includes one convolution operation and one double upsampling operation, wherein the number of channels in the convolution output is the number of categories to be segmented.

[0021] The beneficial technical effects of this invention are as follows:

[0022] (1) The proposed dual attention gating fusion segmentation network model adopts a U-shaped encoding-decoding structure. A dual attention gating fusion module is added between the encoding module and the decoding module at the bottom of the model, which greatly improves the feature extraction capability of the model and directly outputs the high-precision segmentation result of the image at the end of the network, which greatly improves the accuracy of automatic segmentation.

[0023] (2) Different branches in the dual attention gating fusion module play different feature extraction functions. The multi-scale weighted channel attention branch uses learnable weight parameters to adjust the importance of multi-scale features, effectively extracting inter-class response features in the feature map to improve classification accuracy. The Transformer global space self-attention branch, which is formed by the continuous stacking of multiple Transformer self-attention blocks, effectively extracts long-distance dependencies in the feature map to improve localization accuracy.

[0024] (3) For the feature maps obtained from different branches in the dual attention gating fusion module, the gated recurrent unit (GRU) can effectively fuse high-level features from different branches. The update gate of the GRU determines the number of low-level and high-level features entering the next stage to emphasize important information, while the reset gate is used to forget information that is detrimental to segmentation, effectively avoiding redundancy of feature information.

[0025] (4) The powerful encoding module and dual attention gating fusion module obtain high-level features with both global and multi-scale information. Based on these high-level features, decoding is performed. At the same time, the features are concatenated with the same-scale features from the encoding module during the decoding stage, which further enriches the decoding features. This is ultimately reflected in the segmentation results and effectively improves the segmentation accuracy of the network. Attached Figure Description

[0026] Figure 1 This is a flowchart illustrating the image segmentation method according to an embodiment of the present invention;

[0027] Figure 2 This is a schematic diagram of the overall structure of the segmentation network model with dual attention gating fusion according to an embodiment of the present invention;

[0028] Figure 3 This is a schematic diagram of the structure of the dual attention gating fusion module according to an embodiment of the present invention;

[0029] Figure 4 This is a schematic diagram of the structure of the Transformer global self-attention module in an embodiment of the present invention;

[0030] Figure 5 This is a schematic diagram of the structure of the gated loop unit (GRU) according to an embodiment of the present invention;

[0031] Figure 6 This is a schematic diagram comparing the segmentation results of an embodiment of the present invention with those of other algorithms. Detailed Implementation

[0032] The specific embodiments of the present invention will now be described in detail with reference to the accompanying drawings and preferred embodiments.

[0033] The following example, using the segmentation of a three-dimensional CT image of multiple abdominal organs, illustrates the segmentation process of this invention. This invention provides a dual attention-gated fusion image segmentation method. In a preferred embodiment, the overall schematic diagram of the method is shown below. Figure 1 As shown, the method includes the following steps:

[0034] S1. Construct a dataset based on the target task. This dataset includes multiple image samples (which may be 3D images) with pixel-level annotations for a specific target. Preprocess these sample data according to specific requirements and split them into training and validation datasets according to a certain proportion.

[0035] In this embodiment, an abdominal multi-organ dataset is selected. The dataset contains multiple pre-annotated 3D organ scan images, manually segmented images, and original multi-organ scan images used for automatic segmentation. The manually segmented images are divided into 9 regions based on the aorta, gallbladder, left kidney, right kidney, liver, pancreas, spleen, stomach, and background.

[0036] When used as training and validation sets, the original 3D scanned images are sliced ​​at equal intervals along the horizontal axis, resulting in multiple 2D slices for each 3D image. These 2D slices are then cropped at a fixed scale, randomly rotated and flipped with a certain probability, and finally normalized to augment the data. In this embodiment, the total dataset is randomly split into training and validation sets in an 8:2 sample ratio. Only the training dataset is used during training iterations, while the validation dataset is used only for testing accuracy during the model selection phase.

[0037] S2. Construct a segmentation network model based on dual attention gating fusion. The segmentation network model includes an encoding module, a dual attention gating fusion module and a decoding module. The dual attention gating fusion module includes a multi-scale weighted channel attention branch based on CNN and a global spatial self-attention branch based on Transformer. The features of the two branches are fused through a gating mechanism module.

[0038] like Figure 2As shown, the segmentation network model adopts a U-shaped encoder-decoder structure and adds a dual attention gating fusion module at the bottom of the model. For the encoding module, the residual convolutional blocks of the ResNet-50 network are mainly used as the core structure, with several convolutional and pooling operations inserted as connections. The dual attention gating fusion module includes a CNN-based multi-scale weighted channel attention branch and a Transformer-based global spatial self-attention branch, which fuses the features of the two branches through a gating mechanism. For the decoding module, it consists of three identical upsampling blocks and a segmentation map output block. The multi-scale weighted channel attention branch is used to extract inter-class response features to improve classification accuracy and obtain multi-scale feature maps, while the global spatial self-attention branch is used to extract long-range dependency features to improve localization accuracy and obtain global feature maps.

[0039] In this embodiment, after the input 224×224×3 two-dimensional image slice enters the encoding module, the following four stages of operations are performed, as follows:

[0040] The first stage consists of two downsampling operations. The first downsampling is achieved through a downsampling convolutional block, which is composed of a convolutional operation with a stride of 2, a group normalization operation, and a ReLU activation function concatenated together. The output at this time is 112×112×64. The second downsampling is a max pooling operation, which yields an output of 56×56×64.

[0041] The second stage consists of two residual block operations at the same scale. Each residual block has two branches: one branch is an identity connection, and the other branch consists of three concatenated convolutional blocks. The first and second convolutional blocks are composed of a concatenated convolutional operation with a stride of 1, a group normalization operation, and a ReLU activation function. The third convolutional block consists of only a concatenated convolutional operation with a stride of 1 and a group normalization operation. The first convolutional block compresses the number of channels in the feature map, and the third convolutional block restores the original number of channels. Finally, the two branches are added together and activated by the ReLU function before output.

[0042] The first residual block operation increases the number of channels by 4 times without changing the feature map size, resulting in a 56×56×256 feature map. The second residual block does not change the shape of the feature map.

[0043] The third and fourth stages both involve a downsampling residual block operation followed by a residual block operation at the same scale. The downsampling residual block consists of two branches: one branch is a downsampling convolutional block, which is composed of a convolution operation with a stride of 2 and a group normalization operation. The other branch consists of three convolutional blocks. The first convolutional block is composed of a convolution operation with a stride of 1, a group normalization operation, and a ReLU activation function. The second convolutional block is composed of a convolution operation with a stride of 2, a group normalization operation, and a ReLU activation function. The third convolutional block is composed of only a convolution operation with a stride of 1 and a group normalization operation. The first convolutional block compresses the number of channels in the feature map, and the third convolutional block restores the original number of channels. Finally, the two branches are added together and activated by the ReLU function before output.

[0044] After the third stage of operation, the feature map shape is 28×28×512. After the fourth stage of operation, it becomes 14×14×1024, thus obtaining an initial feature map that is downsampled by 4 times compared to the original input and has 1024 channels.

[0045] For the multi-scale weighted channel attention branch in the dual attention gating fusion module, such as Figure 3 As shown. In this embodiment, after the initial feature map of size 14×14×1024 enters the multi-scale weighted channel attention branch, on the one hand, a 14×14×128 feature map t1 is obtained through 1×1 convolution; on the other hand, a downsampling convolution is first performed to obtain a 7×7×512 feature map. Then, a 3×3 convolution, a 5×5 grouped convolution with four groups, and a 7×7 grouped convolution with eight groups are performed on this feature map to obtain feature maps t2, t3, and t4. Then, SE channel attention operation is performed on feature maps t1, t2, t3, and t4 respectively to obtain four groups of channel attention coefficients c1, c2, c3, and c4. These four groups of coefficients are then assigned learnable weight coefficients w1, w2, w3, and w4 to obtain new channel attention coefficients w1c1, w2c2, w3c3, and w4c4. These four groups of coefficients are concatenated and then normalized by the Softmax function to obtain a 1×1×512 channel attention coefficient c.

[0046] On the other side, the four feature maps t1, t2, t3, and t4 are also concatenated to obtain a 14×14×512 combined feature map. The channel attention coefficient c is multiplied by the combined feature map for the corresponding channel to finally obtain a 14×14×512 multi-scale feature map output by the multi-scale weighted channel attention branch.

[0047] For the global self-attention branch in the dual-attention gating fusion module, after the initial feature map of size 14×14×1024 enters the global self-attention branch, it is first reshaped into 196 feature vectors of dimension 728 through convolution and reshape operations, and then positional encoding is performed before being input into a sequence of multiple stacked Transformer attention blocks. The aforementioned feature vectors, after entering each Transformer attention block, undergo the following process... Figure 4 The feature processing shown performs global long-distance dependency modeling. In this embodiment, the feature processing of the Transformer attention block is performed 12 times, and the resulting feature vector is finally reconstructed into a global feature map of 14×14×512.

[0048] For the gating mechanism module in the dual-attention gating fusion module, a GRU gated recursive unit structure is adopted. It has two main gates: a "reset gate" and an "update gate," using the sigmoid function as the control activation function. The "reset gate" controls the state information of one input, and the "update gate" controls the information of the other input. Through these two gates, the GRU module can flexibly control the flow of feature information, achieving the enhancement of beneficial features and the suppression of detrimental factors in the multi-scale feature maps and global feature maps, thereby fully fusing feature information and removing redundant information to obtain the final feature map. Specifically, the 14×14×512 feature maps generated by the multi-scale weighted channel attention branch and the global self-attention branch are input into a... Figure 5 In the gated recurrent unit (GRU) shown, the update gate in the GRU structure determines the information flow to enter the next stage to highlight important feature information, while the reset gate is used to discard information that is detrimental to segmentation, effectively avoiding redundancy of feature information, and outputting a final feature map of 14×14×512 with the same shape as the input.

[0049] For the decoding module of the network model, decoding is performed based on the 14×14×512 high-level semantic feature map output by the dual attention gating fusion module. The decoding module consists of three identical upsampling blocks and a segmentation map output block. The upsampling block includes one upsampling that doubles the resolution, one concatenation of encoded and decoded features at the same scale, and two convolution operations, where the last convolution halves the number of channels in the feature map. In this embodiment, the 14×14×512 feature map is upsampled four times to finally recover to a size of 224×224×64, and then convolved once to reduce the number of channels to 9, which is the number of categories to be segmented, finally resulting in a 224×224×9 segmentation map.

[0050] S3. After configuring the loss function, optimizer, and training parameters for the model training phase, input the training dataset into the segmentation network for model training, and select the optimal image segmentation model by verifying the performance of the validation dataset.

[0051] In this embodiment, the input image size is a 224×224×3 slice. The initial learning rate of the model is set to 0.01. According to the Ploy learning rate decay strategy, the learning rate gradually decreases to 0.001 as the number of iterations increases. The batch size of the input data for each iteration is set to 12, and the number of training epochs is 150. The optimizer selected is the Stochastic Gradient Descent (SGD) optimizer with a momentum of 0.9 and a weight decay coefficient of 1e-4. The loss function is a weighted sum of the Dice loss function and the cross-entropy loss function, with weights set to 0.5. The settings of other network parameters can be understood in the conventional way and will not be elaborated here. In this embodiment, after the model has been trained for 100 epochs, a validation evaluation is performed on the current breakpoint model after each training epoch, and the evaluation accuracy is recorded. After the model has been trained for all 150 epochs, the breakpoint model with the highest validation index is taken as the optimal model weight. The final performance of the model can be judged by commonly used evaluation parameters such as prediction accuracy, or by human judgment of the segmentation effect.

[0052] S4. Input the 3D image data to be segmented into the optimal dual attention gating fusion segmentation network to obtain the high-precision segmentation result of the target image.

[0053] The original multi-organ scan images used for automatic segmentation are preprocessed as described above and then input into the trained optimal image segmentation network to obtain the segmentation results. A comparison diagram of the image segmentation results obtained using the segmentation method of this invention with other algorithms is shown below. Figure 6 As shown.

[0054] In this application, a combination of CNN and Transformer is used to improve segmentation performance. A U-shaped framework is used as the network backbone, and multi-scale channel attention is obtained through weighting, so that the features contain more spatial context information. By stacking multiple Transformer blocks, a Transformer global self-attention branch is designed to extract long-range dependencies. A gated dual attention module is proposed to effectively combine the advantages of multiple attention branch features and Transformer global self-attention branch features, thereby significantly improving the beneficial features of medical image segmentation and suppressing unfavorable information. Finally, decoding is performed based on these high-level features, and at the same scale, the features are concatenated with those from the encoding module during the decoding stage, which further enriches the decoding features and effectively improves the segmentation accuracy of the network.

[0055] While specific embodiments of the present invention have been described above, those skilled in the art should understand that these are merely illustrative examples. Various changes or modifications can be made to these embodiments without departing from the principles and essence of the present invention. Therefore, the scope of protection of the present invention is defined by the appended claims.

Claims

1. An image segmentation method based on dual attention fusion, characterized in that... Includes the following steps: S1. Construct a dataset according to the target task, wherein the dataset includes multiple image samples with pixel-level annotations for a specific target; S2. Construct a segmentation network model and input the image samples from the dataset into the constructed segmentation network model for training. The segmentation network model adopts a U-shaped structure, including an encoding module, a dual-attention gating fusion module, and a decoding module. The encoding module encodes the input image to obtain an initial feature map. The dual-attention gating fusion module includes a CNN-based multi-scale weighted channel attention branch and a Transformer-based global spatial self-attention branch, and fuses the features extracted from the two branches through a gating mechanism module to obtain the final feature map. The decoding module decodes the final feature map to obtain the segmented image. The multi-scale weighted channel attention branch is used to extract inter-class response features to improve classification accuracy and obtain multi-scale feature maps. The global spatial self-attention branch is used to extract long-distance dependency features to improve localization accuracy and obtain global feature maps. S3. Input the image data to be segmented into the trained segmentation network model to obtain a high-precision segmentation result of the target image; The multi-scale weighted channel attention branch includes a cascaded multi-scale convolution operation and a weighted channel attention operation. The global spatial self-attention branch includes multiple Transformer attention blocks connected in series. Each Transformer attention block includes a layer normalization module, a multi-head self-attention MSA, a multilayer perceptron MLP, and a residual module.

2. The image segmentation method based on dual attention fusion according to claim 1, characterized in that: For the initial input feature map, on the one hand, a 1×1 convolution is first performed on the initial feature map to reduce the number of channels to 1 / 8 of the original, generating feature map t1; on the other hand, after halving the size and number of channels of the initial feature map, 3×3 convolution, 5×5 convolution and 7×7 convolution are performed in parallel, compressing the number of channels to 1 / 8 of the initial feature map, thus obtaining three more feature maps t2, t3 and t4 with 1 / 8 of the channels, for a total of four feature maps t1, t2, t3 and t4. Then, SE channel attention operations are performed on the four feature maps t1, t2, t3, and t4 respectively, resulting in four sets of channel attention coefficients. These four sets of channel attention coefficients are then assigned four trainable weight values ​​to obtain four sets of weighted channel attention coefficients. These coefficients are then concatenated and normalized using the Softmax function. Next, they are multiplied by the corresponding channels onto the combined feature map formed by concatenating the four feature maps t1, t2, t3, and t4. After another convolution, a multi-scale feature map with half the original number of channels is finally obtained, thus realizing the extraction of multi-scale information from the initial feature map.

3. The image segmentation method based on dual attention fusion according to claim 2, characterized in that: The Transformer has 12 attention blocks, four sets of 5×5 convolutions, and eight sets of 7×7 convolutions.

4. The image segmentation method based on dual attention fusion according to claim 1, characterized in that: The gating mechanism module adopts a GRU gated recursive unit structure, including a "reset gate" and an "update gate". It uses the sigmoid function as the control activation function to enhance the beneficial features and suppress the detrimental factors in the input multi-scale feature map and global feature map, thereby fully fusing feature information and removing redundant information to obtain the final feature map.

5. The image segmentation method based on dual attention fusion according to claim 1, characterized in that: The encoding module uses the residual convolutional blocks of the ResNet-50 network as the main structure, and inserts several convolution and pooling operations as connections in the main structure; The decoding module includes three identical upsampling blocks and a segmentation map output block. The upsampling block includes one upsampling with doubled resolution, one concatenation of encoded and decoded features at the same scale, and two convolution operations, wherein the latter convolution halves the number of channels in the feature map. The segmentation map output block includes one convolution operation and one double upsampling operation, wherein the number of channels in the convolution output is the number of categories to be segmented.

Images