Abdominal CT image segmentation method based on mixed attention mechanism

Abstract

The invention discloses an abdominal CT image segmentation method based on a local-global mixed attention mechanism. The method comprises the following steps: inputting an abdominal CT image to be segmented into an LG-UNet model, and obtaining a visualization result of image segmentation; wherein the LG-UNet model is obtained based on training of a training set, and the training set comprises the CT images of the multiple organs of the abdomen and the corresponding true values of the CT images. In a model structure, global attention and local attention are calculated in parallel through a mixed attention module, and weighted fusion is carried out on the global attention and the local attention by utilizing a learnable dynamic proportionality coefficient, so that the perception capability of the model on global information and local details of an image is enhanced. And the mixed attention module is integrated in front of and behind the last down-sampling layer of the encoder of the U-Net network to form a dynamic attention adjustment module. The module uses local attention windows of different sizes in different hierarchies to accommodate image features of different hierarchies. Through the design, the model can effectively fuse local and global features, and the accuracy of abdominal CT image segmentation is remarkably improved.

Description

Technical Field

[0001] The present invention belongs to the field of image segmentation and is an abdominal CT image segmentation method based on a hybrid attention mechanism. Background Art

[0002] In the field of medical image analysis, image segmentation as a key step plays a crucial role in disease diagnosis, treatment planning, and surgical navigation. Convolutional Neural Network (CNN) has been widely applied in medical image segmentation tasks due to its powerful feature extraction ability and has made significant progress. Among them, the U-Net model with encoder-decoder structure is particularly prominent. U-Net can effectively handle pixel-level classification problems in medical image segmentation tasks through its unique structural design and has shown good performance in various medical image segmentation scenarios. However, CNN models perform well in capturing local features, but their receptive field is relatively limited, making it difficult to fully capture global information in the image, which affects the segmentation accuracy when dealing with some segmentation tasks that require consideration of global context relationships. To overcome the shortcomings of CNN in capturing global features, researchers have made a lot of improvements and explorations. On the one hand, by introducing dense connections, image pyramids, and other modules, the ability of CNN models to extract multi-scale features can be enhanced to some extent, indirectly improving their ability to perceive global information. On the other hand, the self-attention mechanism as a new feature extraction technology can dynamically focus on areas with important semantic information in the image and enhance the global nature of feature representation by calculating the dependency between different positions. After the Transformer structure based on the self-attention mechanism achieved great success in natural language processing, it has also been gradually introduced into the field of medical image segmentation. For example, the Vision Transformer (ViT) uses the self-attention mechanism to capture the dependency between positions in the input sequence, successfully solving the long-distance dependency problem and showing good performance in medical image segmentation tasks. Chen et al. first combined the global context modeling ability of Transformer with the high-resolution positioning ability of U-Net and applied it to multi-organ segmentation tasks, achieving significant performance improvement. In addition, Zhang et al. proposed a model containing two parallel encoders of Transformer branch and CNN branch, which also achieved good segmentation results in colon polyp, skin, and prostate segmentation tasks, further proving the effectiveness of the self-attention mechanism in medical image segmentation. Chang et al. combined Transformer with the Claw U-Net architecture and achieved better performance than TransUNet on Synapse multi-organ segmentation. Hatamizadeh et al. proposed UNETR, which uses the outputs of the 3rd, 6th, 9th, and 12th Transformer blocks as inputs to four down-sampling convolutional neural network modules in the encoder.Swin UNETR replaces the Transformer module in UNETR with the Swin Transformer module, and better segmentation results are obtained. However, the current improvement methods mostly focus on the application of a single attention mechanism, either focusing on the enhancement of local features or focusing on the extraction of global features, and it is difficult to simultaneously consider the dual needs of local details and global context information in the abdominal CT image segmentation task.

[0003] In order to improve the local and global feature extraction capability of the coding and decoding architecture neural network, the application provides an abdominal CT image segmentation method based on a local-global hybrid attention mechanism SUMMARY

[0004] The purpose of the application is to provide an abdominal CT image segmentation method based on a local-global hybrid attention mechanism, to improve the local and global feature extraction capability of the coding and decoding architecture neural network and to improve the segmentation accuracy of the retinal blood vessel image.

[0005] To achieve the above purpose, the application provides the following scheme:

[0006] The abdominal CT image segmentation method based on the local-global hybrid attention mechanism comprises:

[0007] Obtaining an abdominal CT image to be segmented;

[0008] Inputting the abdominal multi-organ CT image to be segmented into an LG-UNet model to obtain an image segmentation visualization result, wherein the LG-UNet model is obtained based on training of a training set, and the training set comprises abdominal multi-organ CT images and corresponding true values;

[0009] The local-global hybrid attention module is integrated into the last down-sampling layer of the encoder of the U-Net network, and a hybrid attention module is added before and after the last down-sampling layer, respectively;

[0010] The local-global hybrid attention module uses local attention windows of different sizes in different layers to adapt to image features at different levels;

[0011] The model comprises an encoder, a decoder and a dynamic attention adjustment module, wherein:

[0012] The encoder uses the powerful positioning capability of the convolutional network to extract local information;

[0013] The decoder is used for fusing local-global encoding features and skip connection features, and outputs the new prediction result;

[0014] The dynamic attention adjustment module: in the Transformer multi-head attention layer, the global attention weight matrix and the local convolution attention weight matrix are calculated in parallel, and are weighted and fused through a learnable dynamic proportion coefficient;

[0015] The encoder includes 4 down-sampling layers, each of which is composed of two 3x3 convolutions, a batch normalization layer and a ReLU activation layer, uses a pooling layer with a step of 2 for down-sampling, and two local-global hybrid attention layers;

[0016] The decoder includes two 3x3 convolutions, a batch normalization layer, a ReLU activation layer, a bilinear interpolation up-sampling and a parallel jump connection path, and each time the encoder performs a down-sampling, the decoder part correspondingly performs an up-sampling step to gradually restore the spatial resolution;

[0017] The dynamic attention adjustment module: in the Transformer multi-head attention layer, the global attention weight matrix and the local convolution attention weight matrix are calculated in parallel, and are weighted and fused through a learnable dynamic proportion coefficient, including:

[0018] A weighting factor is designed for global attention and local attention, and an input token sequence is processed to obtain a feature vector representing the input, which is input into a fully connected layer to output two values representing the global and local weight factors, respectively;

[0019] The output is converted to a value between 0 and 1 through a fully connected layer and a Sigmoid function, and learnable weighting factors a and b are used to dynamically adjust the proportion of global and local attention, ensuring that the model can automatically adjust the global and local attention according to different task requirements;

[0020] The weighting factors a and b satisfy the constraint condition a+β=1, ensuring that the sum of the two is 1, avoiding the model completely ignoring the attention of one aspect;

[0021] The Transformer maps the input features to queries (Q), keys (K) and values (V);

[0022] Optionally, the global attention is modeled through a standard multi-head self-attention mechanism to generate a global attention matrix :

[0023]

[0024] wherein, represents the dimension of each attention head, and h is the number of attention heads;

[0025] Optionally, the local attention is extracted by the neighborhood feature correlation in the local window of the adjacent n token sequence, and a local attention matrix is generated :

[0026]

[0027] wherein, denotes the local window of position i.

[0028] Optionally, the size of the local window is n, and the index range of the local window for position i is:

[0029]

[0030] The final attention representation The final attention representation can be obtained by weighted summation.

[0031] .

[0032] The beneficial effects of the present application are: the method of the present application introduces a hybrid local-global attention mechanism, which significantly improves the performance of abdominal CT image segmentation. First, the global attention and the local attention are calculated in parallel in the multi-head attention layer of the Transformer, the global attention module can effectively capture the long-distance dependence and the global context information in the image, and the local attention module limited by the window focuses on extracting the local detail features of the image, so that the global information and the local details of the image are considered at the same time in the feature extraction stage. Secondly, the global attention features and the local attention features are weighted and fused by the learnable dynamic proportion coefficient, the dynamic proportion coefficient can adaptively adjust the weight distribution according to the feature distribution of the input image and the task demand, so that the model can flexibly balance the importance of global and local features in different scenarios, and the segmentation of the model is improved. BRIEF DESCRIPTION OF DRAWINGS

[0033] In order to more clearly illustrate the technical solutions of the embodiments of the present application or the prior art, the drawings needed in the embodiments will be briefly introduced below. Obviously, the drawings in the following description are only some embodiments of the present application, and other drawings can be obtained by those skilled in the art without creative labor.

[0034] Figure 1 The flow chart of the abdominal CT image segmentation method based on the hybrid attention mechanism is provided.

[0035] Figure 2 The flow chart of the abdominal CT image data set preprocessing is provided.

[0036] Figure 3 A local-global hybrid attention mechanism diagram of the embodiment of the present application.

[0037] Figure 4 A local-global hybrid attention mechanism diagram of the embodiment of the present application. DETAILED DESCRIPTION

[0038] The technical solutions in the embodiments of the present application will be described clearly and completely below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, rather than all the embodiments. Based on the embodiments in the present application, all other embodiments obtained by those skilled in the art without creative work fall within the protection scope of the present application.

[0039] In order to make the above objectives, characteristics and advantages of the present application more obvious and easy to understand, the present application will be further described in detail below with reference to the drawings and specific embodiments.

[0040] Inspired by Tansformer, a medical image segmentation method based on hybrid attention mechanism is proposed on the basis of UNet, including local, global attention modules and dynamic attention adjustment modules composed of them and convolutional layers, named LG-UNet. As shown in Figure 1 from step 101 to step 104.

[0041] The embodiment provides a hybrid attention mechanism based abdominal CT image segmentation method, which comprises the following steps:

[0042] As shown in Figure 2 , an abdominal multi-organ image to be segmented is obtained, from step 201 to step 204.

[0043] The retinal blood vessel image to be segmented is input into the LG-UNet model, and the image segmentation visualization result is obtained, wherein the LG-UNet model is obtained based on training of a training set, and the training set comprises abdominal multi-organ images and corresponding true values.

[0044] The LG-UNet model comprises an encoder, a decoder and a dynamic attention adjustment module, wherein: the encoder uses the powerful positioning ability of the convolutional network to extract local information; the decoder is used for fusing local-global encoding features and skip connection features, and outputs the new prediction result. The dynamic attention adjustment module: in the Transformer multi-head attention layer, the global attention weight matrix and the local convolution attention weight matrix are calculated in parallel, and are fused by a learnable dynamic proportion coefficient.

[0045] As shown in Figure 3As shown, in this embodiment, based on the proposed dynamic attention adjustment module, an LG-UNet model is constructed on the basis of UNet, which is composed of an encoder and a decoder. The encoder is mainly composed of a CNN hybrid module and an attention module, wherein the CNN down-sampling layer is composed of two 3x3 convolution, one batch normalization layer and one ReLU activation layer, and the down-sampling is performed using a pooling layer with a step of 2, and the attention module includes two local-global attention sub-modules and one convolution layer. The decoder has the same structure as the decoder of UNet, which is mainly composed of a 1x1 convolution, an up-sampling layer and a convolution layer.

[0046] The overall flow of the method of this embodiment is shown in Figure 3 First, the encoder operation is performed using four down-sampling layers, and the dynamic attention adjustment module calculates the global attention weight matrix and the local convolution attention weight matrix in parallel, so that the model can adaptively adjust the proportion of global and local attention according to the task requirements, design a weighting mechanism, and combine the weighted results of global and local attention, so that the model can automatically learn how to adjust the contribution of global and local attention during the training process. Next, the decoder fuses the features obtained by the dynamic attention adjustment module and the skip connection features from the encoder, and finally outputs the prediction result.

[0047] Specifically, in the LG-UNet model of this embodiment, the features extracted by the encoder are processed through multiple down-sampling layers to obtain feature maps of different resolutions. These feature maps contain rich local information and can provide important detail support for subsequent segmentation tasks; in order to further enhance the model's perception ability of global information, a hybrid attention module is integrated into the encoder. This module uses different sizes of local attention windows in different layers to adapt to image features at different levels, so as to realize the comprehensive extraction of local and global features. In the decoder part, different resolution skip connection features are taken. These skip connection features come from different levels of the encoder and contain rich local details. Through the skip connection, these features are transmitted to the decoder and fused with the features of the decoder itself.

[0048] Further, the dynamic attention adjustment method in model training is as follows: in the multi-head attention layer of the Transformer, the global attention weight matrix and the local convolution attention weight matrix are calculated in parallel, and the weighted fusion is performed through a learnable dynamic proportion coefficient.

[0049] Specifically, in the present embodiment, in order to realize dynamic proportional coefficient weighted fusion, a weighting factor is designed for global attention and local attention, the input token sequence is processed to obtain a feature vector representing the input, and the feature vector is input into a fully connected layer to output two values representing the global and local weight factors respectively.

[0050] The output is converted to a value between 0 and 1 by the fully connected layer and the Sigmoid function, and the learnable weighting factors a and b are used to dynamically adjust the proportion of global and local attention, ensuring that the model can automatically adjust the global and local attention according to different task requirements.

[0051] The weighting factors a and b satisfy the constraint condition a+ b = 1, ensuring that the sum of the two is 1, avoiding the model completely ignoring the attention of one aspect.

[0052] The Transformer maps the input features to queries (Q), keys (K) and values (V).

[0053] The global attention is modeled by a standard multi-head self-attention mechanism to generate a global attention matrix :

[0054]

[0055] wherein, represents the dimension of each attention head, and h is the number of attention heads.

[0056] The local attention extracts the correlation of neighborhood features in the local window of the adjacent n token sequence to generate a local attention matrix :

[0057]

[0058] wherein, represents the local window of position i.

[0059] Let the size of the local window be n, and for position i, the index range of the local window is:

[0060]

[0061] The final attention representation can be obtained by weighted summation:

[0062]

[0063] Compared with traditional segmentation models, the method in this paper optimizes global feature extraction while retaining the traditional convolutional skip connection to enhance the ability of local information and global context information feature extraction. The main effects are as follows:

[0064] The hybrid attention module is in the Transformer multi-head attention layer, and the global attention weight matrix and the local convolution attention weight matrix are calculated in parallel. By learning the dynamic proportion coefficient to weight the fusion of the two attention mechanisms, the model can dynamically adjust the attention degree of global information and local details according to different levels of image features.

[0065] A dynamic attention adjustment module is designed: in the encoder of the U-Net network, a hybrid attention module is added before and after the last downsampling layer, forming a dynamic attention adjustment module. This design enables the model to simultaneously utilize the advantages of local and global attention mechanisms during key feature extraction stages, enhancing feature expression capabilities.

[0066] A specific experimental example of this embodiment is shown as follows:

[0067] The experiment is implemented using PyTorch 1.13, running on a computer with an Intel i5-13400F processor, 32GB of 3600 DDR4 memory, and a 12GB NVIDIA RTX3060, with a Windows 11 system environment. The average similarity coefficient (Dice) and the average Hausdorff distance (HD) are used to compare with SOTA methods such as Att-UNet, TransUnet, and SwinUnet to verify the effectiveness of the method. The network structure of U-Net is the same as the initial encoder in the proposed method; TransUnet is a hybrid model that combines CNN and Transformer, thus having the advantages of both CNN and Transformers; SwinUnet is a pure Transformer model. Therefore, these methods are selected for comparison with the proposed method. The experiment uses the Synapse multi-organ dataset for training and testing. The Synapse dataset includes 30 cases of 3779 abdominal axial clinical CT images, with 18 samples divided into the training set and 12 samples divided into the test set. All 3D bodies are segmented in a slice-by-slice manner to evaluate 9 abdominal organs (aorta, gallbladder, spleen, left kidney, right kidney, liver, pancreas, spleen, and stomach).

[0068] The experiment uses a 224x224 picture as input, a batch size of 4, uses the SGD optimization algorithm during training, an initial learning rate of 0.01, a momentum of 0.9, a weight decay of 0.0001, and a cosine annealing method to update the learning rate. The training epoch = 200.

[0069] The experiment uses the average Dice coefficient and the average HD to evaluate the segmentation accuracy of the model. The Dice coefficient is a set similarity index used to calculate the similarity of two samples, and its calculation formula is as follows:

[0070]

[0071] wherein and respectively represent the true label and the prediction result, represents the intersection thereof.

[0072] HD is also a set similarity index, which defines the similarity between two point sets by calculating the distance between them. The smaller the value, the higher the similarity between the two sets. The calculation formula of HD is as follows:

[0073]

[0074]

[0075]

[0076] wherein is the distance norm between point set A and point set B, and are the one-way Hausdorff distance from set A to set B and from set B to set A, respectively.

[0077] First, Table 1 compares the test results of LG-UNet and traditional CNN segmentation models such as V-Net, DARR, R50 U-Net, and Att-UNet on the Synapse dataset. LG-UNet performs best, with an average Dice value of 80.95% and an HD value of 21.54mm on the test set, while Att-UNet, which ranks second among these methods, adds attention mechanisms to the skip connection, with an average Dice value of 77.77% and an HD value of 36.02mm.

[0078] In addition, Table 1 also compares the test results of the CNN+Transformer hybrid model on the Synapse dataset, including R50 ViT, TransUNet, Swin-Unet, and TransClaw U-Net, etc. Again, LG-UNET performs best. Swin-Unet, ranked 2nd, adopts a pure Transformer encoding-decoding architecture, with average Dice and HD values of 79.13% and 21.55mm, respectively. TransClaw U-Net, ranked 3rd, combines Transformer with Claw U-Net architecture, with average Dice and HD values of 78.09% and 26.38mm, respectively. TransUNet, similar to LG-UNET architecture, has average Dice and HD values of 77.48% and 31.69mm, respectively.

[0079] Table 1

[0080]

[0081] The method of the embodiment adopts parallel computing of the global attention weight matrix and the local convolution attention weight matrix in the Transformer multi-head attention layer. The global attention models long-range dependency relationships through a standard multi-head self-attention mechanism to generate a global attention matrix; the local attention extracts neighborhood feature correlations in a local window of adjacent n token sequences to generate a local attention matrix. By weighting and fusing the two attention mechanisms through a learnable dynamic proportion coefficient, the model can dynamically adjust the attention degree to global information and local details according to different levels of image features. Experimental comparison on the abdominal multi-organ CT image dataset proves its effectiveness.

[0082] The above-described embodiments are merely descriptions of the preferred modes of the present application and do not limit the scope of the present application. Without departing from the design spirit of the present application, various modifications and improvements to the technical solutions of the present application made by those of ordinary skill in the art shall fall within the protection scope determined by the claims of the present application.

Claims

1. A method for abdominal CT image segmentation based on a local-global hybrid attention mechanism, characterized in that: include: Data acquisition and preprocessing: Obtain the Synapse abdominal multi-organ CT image to be segmented from the MICCAI challenge official website; The abdominal multi-organ CT image to be segmented is input into the LG-UNet model to obtain an image segmentation visualization result, wherein the LG-UNet model is trained based on a training set, and the training set includes abdominal multi-organ CT images and corresponding true values.

2. The abdominal CT image segmentation method based on the local-global hybrid attention mechanism according to claim 1, characterized in that: The method comprises: Integrate the hybrid attention module into the last downsampling layer of the encoder of the U-Net network and add a hybrid attention module before and after it to form a dynamic attention adjustment module; The hybrid attention module uses local attention windows of different sizes in different layers to adapt to image features at different levels.

3. The abdominal CT image segmentation method based on the local-global hybrid attention mechanism according to claim 1, characterized in that: The model includes an encoder, a decoder, and a dynamic attention adjustment module, where: The encoder uses the powerful positioning capability of the convolutional network to extract local information; The decoder is used to fuse local-global encoding features and skip connection features and output the prediction result; The dynamic attention adjustment module: in the Transformer multi-head attention layer, the global attention weight matrix and the local convolution attention weight matrix are calculated in parallel, and weighted fusion is performed through a learnable dynamic scaling coefficient.

4. The abdominal CT image segmentation method based on the local-global hybrid attention mechanism according to claim 3, characterized in that: Implementation of the encoder: It includes 4 downsampling layers, each of which consists of two 3×3 convolutions, a batch normalization layer, and a ReLU activation layer, using a pooling layer with a stride of 2 for downsampling and two of the aforementioned hybrid attention layers.

5. The abdominal CT image segmentation method based on local-global hybrid attention mechanism according to claim 3, characterized in that: Implementation of the decoder: It consists of two 3×3 convolutions, a batch normalization layer, a ReLU activation layer, bilinear interpolation upsampling, and a parallel skip connection path. Whenever the encoder performs a downsampling, the decoder part performs an upsampling step accordingly to gradually restore the spatial resolution.

6. The abdominal CT image segmentation method based on hybrid attention mechanism according to claim 3, characterized in that: The hybrid attention mechanism includes: Design a weighting factor for global attention and local attention, process the input token sequence, obtain a feature vector representing the input, input the feature vector into a fully connected layer, and output two values, representing the global and local weight factors respectively; The output is converted to a value between 0 and 1 through a fully connected layer and a Sigmoid function. Learnable weighting factors α and β are used to dynamically adjust the proportion of global and local attention, ensuring that the model can automatically adjust the global and local attention according to different task requirements. The weighting factors α and β satisfy the constraint α + β = 1, ensuring that the sum of the two is 1, thus preventing the model from completely ignoring one aspect of attention; Transformer maps input features into query (Q), key (K) and value (V); The global attention is modeled through the standard multi-head self-attention mechanism to model long-range dependencies and generate a global attention matrix : in, represents the dimension of each attention head, h is the number of attention heads; Local attention extracts the correlation of neighborhood features in the local window of adjacent n token sequences to generate a local attention matrix : in, represents the local window at position i; Assume that the size of the local window is n. For position i, the index range of the local window is: Final attention representation It can be obtained by weighted summation: 。

Images