Three-dimensional medical image segmentation method and system based on long short-term memory self-attention model
By combining CNN and Transformer's long short-term memory self-attention model, the problem of insufficient global and local features in tumor segmentation in three-dimensional medical images is solved, and tumor segmentation with high precision and low computational complexity is achieved, which is suitable for real-time segmentation needs.
Patent Information
- Application Number
- CN202211440881.X
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2022-11-17
- Publication Date
- 2025-10-03
- Estimated Expiration
- 2042-11-17
AI Technical Summary
Existing 3D medical image tumor segmentation algorithms are insufficient in capturing global and local features, resulting in low segmentation accuracy and high computational complexity, making them difficult to be effectively applied on small sample datasets.
Combining the convolutional neural network (CNN) and the Transformer model, the long short-term memory self-attention (SLM-SA) model is adopted. Through the memory forgetting and updating strategy, a small number of Transformer modules are used to perform self-attention calculations on low-resolution feature maps, combined with CNN for feature extraction, limiting the amount of calculation and improving segmentation accuracy.
High-precision 3D medical image tumor segmentation was achieved on small sample datasets, reducing computational complexity and making it suitable for clinical scenarios requiring real-time segmentation.
Smart Images

Figure CN116129107B_ABST
Abstract
Description
Technical Field
[0001] The present invention belongs to the field of image segmentation in medical image processing. More specifically, it relates to a three-dimensional medical image segmentation method and system based on a long short-term memory self-attention model. The method and system are based on deep learning and can be particularly used to segment tumor areas in three-dimensional medical images. Background Art
[0002] Medical images are an important source of lesion information and are of great value. Early detection and treatment of cancer are crucial. In clinical practice, tumor segmentation in 3D medical images can provide information such as tumor measurement and location, greatly aiding early diagnosis. However, due to differences in texture and shape between individual tumors, the possibility of multiple tumors and blurred tumor boundaries, efficient and accurate tumor segmentation is a challenging task.
[0003] Automatic tumor segmentation in 3D medical images plays a crucial role in the early diagnosis and treatment of cancer. However, traditional segmentation algorithms, such as threshold segmentation, region growing, and active contour models, are mostly based on differences in edges and grayscale distribution. Due to the significant variability in morphology, texture, and grayscale distribution among tumor images, these traditional methods are relatively unsuitable. Consequently, deep learning-based image segmentation algorithms have attracted considerable attention in recent years, and numerous CNN-based models have been employed for image segmentation. The U-Net is currently the most commonly used network architecture, with its unique encoder-decoder structure demonstrating excellent performance in segmentation tasks. Further research has found that while convolution-based segmentation models can effectively extract local features, they cannot capture global features well, which compromises their performance. The Transformer model, originally applied in natural language processing tasks, can improve segmentation accuracy by capturing global image features to predict missing edges and reduce noise-induced errors. Currently, the application of the Transformer to image processing tasks still has several drawbacks. First, the Transformer requires a large number of training samples to achieve good performance, but obtaining sufficient tumor images for training is difficult in clinical applications. At the same time, the computational complexity of the Transformer is huge and proportional to the square of the number of patches, which makes its training and application in 3D image segmentation tasks very time-consuming. In addition, the Transformer lacks the inductive bias of CNN, such as translation invariance and locality, which further limits the model performance. Since CNN and Transformer are good at extracting local features and global features respectively, they can be combined in different ways to improve the model's image segmentation performance on small training samples. One way is to further optimize the attention module by adding convolution operations to the Transformer, and another way is to combine the Transformer and CNN in the network structure.
[0004] Currently, the combined structures of CNN and Transformer for image segmentation can be divided into three types. The first type uses only multi-layer Transformer as an encoder, and its performance is related to the number of training samples. Although the simple superposition of Transformer can improve performance to a certain extent, it will also significantly increase the amount of computation, which is not cost-effective. The second structure inserts the Transformer module between the CNN-based encoder and decoder for further feature extraction. This structure can only extract low-resolution global features, cannot fully play the role of the self-attention mechanism, and has limited improvement in network performance. The third structure uses CNN and Transformer for dual-branch encoding at the same time, fusing encoding features of different resolutions to guide the decoding process. This parallel structure cannot effectively fuse the two together, and the superposition of the two models makes it have a high computational cost. Summary of the Invention
[0005] In response to the above defects or improvement needs of the prior art, the purpose of the present invention is to provide a three-dimensional medical image segmentation method and system based on the Short-term and Long-term Memory Self-attention (SLM-SA) model, wherein by integrating CNN and Transformer, a new model for real-time and accurate segmentation of target areas (such as tumors, etc.) in three-dimensional medical images is obtained, and a learning strategy similar to memory forgetting and updating is adopted. Taking tumors as an example, it is beneficial to improve the segmentation accuracy of complex tumors in medical images. This method uses a CNN module for preliminary extraction of image features, effectively avoiding the disadvantage that Transformer requires a large number of samples for training. A small number of Transformer modules are used to implement the memory structure and apply it to 1 / 4 and 1 / 8 resolution feature maps, effectively exerting the role of the self-attention mechanism while limiting the amount of calculation. This method has fewer parameters and lower computational complexity, and can be applied to clinical scenarios that require real-time and accurate segmentation.
[0006] To achieve the above objectives, according to one aspect of the present invention, a three-dimensional medical image segmentation method based on a long short-term memory self-attention model is provided, which can segment target area images in medical images and is characterized by comprising the following steps:
[0007] (1) Collect three-dimensional medical images containing target areas and known gold standards for target area segmentation, crop the original images according to the region of interest, and perform pixel value normalization to obtain training set samples;
[0008] (2) For the original three-dimensional medical image to be segmented, the original image is cropped according to the region of interest, and the pixel value of the cropped image is normalized to obtain the sample to be segmented;
[0009] (3) Constructing a three-dimensional medical image segmentation network based on long-short-term memory self-attention, the network first uses a feature extraction module composed of multiple convolutional layers, and uses convolution calculation to convert the grayscale information in the network input image into feature information to obtain initial image features; then uses multiple convolution modules with residual structures connected in sequence to further process the initial image features to obtain corresponding convolution features; then uses several short-term memory modules based on multi-head self-attention mechanism to calculate the global correlation between convolution features at different levels, and generates short-term memory features based on the global correlation and the convolution features used in the correlation calculation; then, a long-term memory module also based on the multi-head self-attention mechanism will perform a global correlation calculation on all the obtained short-term memory features, thereby generating features to be reconstructed for reconstructing the segmentation results; wherein the long-term memory module has the same structure as any short-term memory module; finally, the reconstruction module will process the features to be reconstructed through multiple convolution and deconvolution layers, thereby obtaining the predicted probability of the category to which each pixel of the network input image belongs, thereby achieving accurate segmentation of the target area in the three-dimensional medical image;
[0010] The three-dimensional medical image segmentation network is constructed by using the training set samples obtained in step (1) and the corresponding segmentation gold standard as labels to train the trained network so that the trained network can segment the target area contained in the network input image;
[0011] (4) The sample to be segmented obtained in step (2) is input as a network input image into the trained long short-term memory self-attention based three-dimensional medical image segmentation network obtained in step (3), and the segmentation result of the target area in the sample to be segmented can be output.
[0012] As a further preferred embodiment of the present invention, the three-dimensional medical image segmentation network based on long short-term memory self-attention constructed in step (3) specifically includes:
[0013] The feature extraction module is used to extract image features. This feature extraction module reduces the resolution and increases the number of channels through two convolution calculations with a step size of 2, converts the network input image into a multi-channel feature map, and obtains the initial image features.
[0014] A plurality of convolution modules with residual structures connected in sequence are used to process the initial image features to obtain corresponding convolution features; wherein the resolution of the convolution features obtained by any convolution module is the same as the resolution of the initial image features; preferably, the plurality of convolution modules with residual structures connected in sequence are specifically four convolution modules with residual structures connected in sequence, respectively recorded as convolution module 1, convolution module 2, convolution module 3 and convolution module 4, and the convolution features of different levels obtained by these convolution modules are respectively recorded as convolution feature 1, convolution feature 2, convolution feature 3 and convolution feature 4;
[0015] Several short-term memory modules based on the multi-head self-attention mechanism, the number of the short-term memory modules is the number of the convolution modules with residual structure minus 1, wherein the first-level short-term memory module corresponds to the first two levels of convolution modules with residual structure, and each subsequent level of short-term memory module corresponds to the next level of convolution modules with residual structure in sequence; the first-level short-term memory module is used to calculate the correlation between the convolution features obtained by the first two levels of convolution modules with residual structure, and each subsequent level of short-term memory module is used to calculate the correlation between the short-term memory features obtained by the previous level of short-term memory module and the convolution features obtained by the corresponding first-level convolution module with residual structure, thereby calculating the convolution features obtained by all convolution modules with residual structure The global correlation between the product features; preferably, the several short-term memory modules based on the multi-head self-attention mechanism are specifically three short-term memory modules based on the multi-head self-attention mechanism, and the three short-term memory modules based on the multi-head self-attention mechanism are respectively short-term memory module 1, short-term memory module 2 and short-term memory module 3, and each short-term memory module has two inputs; wherein, the convolution feature 1 is downsampled and input into the short-term memory module 1 together with the convolution feature 2 to obtain the short-term memory feature 1; the short-term memory feature 1 is downsampled and input into the short-term memory module 2 together with the convolution feature 3 to obtain the short-term memory feature 2; the short-term memory feature 2 is downsampled and input into the short-term memory module 3 together with the convolution feature 4 to obtain the short-term memory feature 3;
[0016] Furthermore, except for the last short-term memory module, the short-term memory features obtained by each short-term memory module are reduced in number of channels through convolution calculation, and are added to the lowest level convolution feature input to the short-term memory module as the input of the next level convolution module; preferably, the short-term memory feature 1 and the convolution feature 2 are added as the input of the convolution module 3, and the short-term memory feature 2 and the convolution feature 3 are added as the input of the convolution module 4;
[0017] A long-term memory module based on a multi-head self-attention mechanism is used to first concatenate the short-term memory features obtained by each short-term memory module in the dimension where the channel is located to generate a merged feature, and then perform a global correlation calculation to generate a feature to be reconstructed for reconstructing the segmentation result; preferably, a long-term memory module based on a multi-head self-attention mechanism is used to first concatenate short-term memory feature 1, short-term memory feature 2, and short-term memory feature 3 in the dimension where the channel is located to generate a merged feature, and then perform a global correlation calculation on the merged feature to generate a feature to be reconstructed for reconstructing the segmentation result;
[0018] The reconstruction module is used to take the features to be reconstructed as input, use a deconvolution layer with a stride of 2 to increase the image resolution, and use two convolution layers with a stride of 1 to reduce the number of channels, so as to finally obtain the predicted probability of the category to which each pixel of the network input image belongs.
[0019] As a further preferred embodiment of the present invention, the feature extraction module has a total of 12 layers of detail structures, and the output of the previous layer of detail structure serves as the input of the next layer of detail structure, wherein:
[0020] The first layer is a convolutional layer with a stride of 2; the output resolution is 1 / 2 of the network input image, and the number of output channels is a pre-set value;
[0021] The second layer is the batch normalization (BN) layer;
[0022] The third layer is the rectified linear unit (ReLU) activation layer;
[0023] The fourth layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels;
[0024] The 5th layer is the batch normalization (BN) layer;
[0025] The 6th layer is the rectified linear unit (ReLU) activation layer;
[0026] The 7th layer is a convolutional layer with a stride of 2; compared to the input, the output resolution is halved and the number of output channels is doubled;
[0027] The 8th layer is the batch normalization (BN) layer;
[0028] The 9th layer is the rectified linear unit (ReLU) activation layer;
[0029] The 10th layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels;
[0030] The 11th layer is the batch normalization (BN) layer;
[0031] The 12th layer is a rectified linear unit (ReLU) activation layer.
[0032] As a further preferred embodiment of the present invention, for any convolution module, there are 12 layers of detail structures, and the output of the previous layer of detail structure serves as the input of the next layer of detail structure, wherein:
[0033] The first layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels;
[0034] The second layer is the batch normalization (BN) layer;
[0035] The third layer is the rectified linear unit (ReLU) activation layer;
[0036] The fourth layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels;
[0037] The 5th layer is the batch normalization (BN) layer;
[0038] The sixth layer is a rectified linear unit (ReLU) activation layer; its input is the sum of the output of the fifth layer and the input of the first layer;
[0039] The 7th layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels;
[0040] The 8th layer is the batch normalization (BN) layer;
[0041] The 9th layer is the rectified linear unit (ReLU) activation layer;
[0042] The 10th layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels;
[0043] The 11th layer is the batch normalization (BN) layer;
[0044] The 12th layer is a rectified linear unit (ReLU) activation layer; its input is the sum of the output of the 11th layer and the input of the 7th layer.
[0045] As a further preferred embodiment of the present invention, for the short-term memory module:
[0046] In the input of the short-term memory module, for the first-level short-term memory module, the convolution feature obtained by the first-level convolution module with a residual structure is recorded as the memory feature, and the convolution feature obtained by the second-level convolution module with a residual structure is recorded as the current feature; for each subsequent level of short-term memory module, the short-term memory feature obtained by the previous level of short-term memory module is recorded as the memory feature, and the convolution feature obtained by the first-level convolution module with a residual structure corresponding to the short-term memory module is recorded as the current feature;
[0047] For any short-term memory module, for its current feature, a convolution calculation is performed before input to double the number of channels; for its memory feature, a convolution calculation with a step size of 2 is performed before input, while reducing the resolution and increasing the number of channels; preferably, reducing the resolution is specifically halving the resolution, and increasing the number of channels is specifically doubling the number of channels;
[0048] For any short-term memory module, the current feature is regarded as the query in the self-attention calculation process, and the memory feature is regarded as the key and value to perform the calculation of the self-attention module in the short-term memory module; specifically:
[0049] First, the query, key, and value are each encoded using a convolutional layer. The resulting feature encodings of the query, key, and value are then divided into a preset number of patches in the spatial dimension. Next, the attention weight is calculated based on the sequence generated for the query and key. The formula is:
[0050]
[0051] Where d is the length of the feature vector, Q and K represent the sequences corresponding to Query and Key respectively, the Softmax function is applied to the dimension corresponding to the Key sequence, and Attn represents the attention weight;
[0052] Then, the number of self-attention groups in the attention weight is doubled through a linear layer, and then input into a convolution layer for a convolution calculation with unchanged resolution and number of channels. Then, another opposite linear layer is used to restore the number of self-attention groups to obtain the final attention weight; finally, the final attention weight is matrix multiplied by Value, and the result is restored to a three-dimensional feature map through an inverse serialization operation; for the three-dimensional feature map, a learnable weight is pre-assigned according to its number of channels, so that the channels correspond to the learnable weights one-to-one, and then each channel of the three-dimensional feature map is multiplied by the corresponding learnable weight to obtain the overall output of the self-attention module; the learnable weights can be automatically updated through training;
[0053] The overall output of the self-attention module is processed by the first LN (Layer Normalization) layer and added to the current feature; the output is then passed through a feedforward neural network and a second LN layer, and the result is added to the output of the first LN layer to obtain the short-term memory feature.
[0054] As a further preferred embodiment of the present invention, for the long-term memory module:
[0055] First, the short-term memory features output by each short-term memory module are spliced in the channel dimension to obtain the merged features. The merged features are regarded as the query, key and value in the self-attention calculation process, and the calculation of the self-attention module in the long-term memory module is performed. Specifically:
[0056] First, the query, key, and value are each encoded using a convolutional layer. The resulting feature encodings of the query, key, and value are then divided into a preset number of patches in the spatial dimension. Next, the attention weight is calculated based on the sequence generated for the query and key. The formula is:
[0057]
[0058] Where d is the length of the feature vector, Q and K represent the sequences corresponding to Query and Key respectively, the Softmax function is applied to the dimension corresponding to the Key sequence, and Attn represents the attention weight;
[0059] Then, the number of self-attention groups in the attention weight is doubled through a linear layer, and then input into a convolution layer for a convolution calculation with unchanged resolution and number of channels. Then, another opposite linear layer is used to restore the number of self-attention groups to obtain the final attention weight; finally, the final attention weight is matrix multiplied by Value, and the result is restored to a three-dimensional feature map through an inverse serialization operation; for the three-dimensional feature map, a learnable weight is pre-assigned according to its number of channels, so that the channels correspond to the learnable weights one-to-one, and then each channel of the three-dimensional feature map is multiplied by the corresponding learnable weight to obtain the overall output of the self-attention module; the learnable weights can be automatically updated through training;
[0060] The overall output of the self-attention module is processed by the first LN (Layer Normalization) layer and added to the merged features; the output is then passed through a feedforward neural network and a second LN layer, and the result is added to the output of the first LN layer to obtain a long-term memory feature, which is used as the input of the reconstruction module.
[0061] As a further preferred embodiment of the present invention, the reconstruction module has a total of 18 layers of detail structures, and the output of the previous layer of detail structure serves as the input of the next layer of detail structure, wherein:
[0062] The first layer is the batch normalization (BN) layer;
[0063] The second layer is the rectified linear unit (ReLU) activation layer;
[0064] The third layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels;
[0065] The fourth layer is the batch normalization (BN) layer;
[0066] The fifth layer is the rectified linear unit (ReLU) activation layer;
[0067] The sixth layer is a deconvolution layer with a stride of 2. Compared with the input, the output resolution is doubled, and the number of output channels remains unchanged.
[0068] The 7th layer is the batch normalization (BN) layer;
[0069] The 8th layer is the rectified linear unit (ReLU) activation layer;
[0070] The 9th layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels;
[0071] The 10th layer is the batch normalization (BN) layer;
[0072] The 11th layer is the rectified linear unit (ReLU) activation layer;
[0073] The 12th layer is a deconvolution layer with a stride of 2. Compared with the input, the output resolution is doubled, and the number of output channels remains unchanged.
[0074] The 13th layer is the batch normalization (BN) layer;
[0075] The 14th layer is the rectified linear unit (ReLU) activation layer;
[0076] The 15th layer is a convolutional layer with a stride of 1. Compared to the input, the output resolution remains unchanged, but the number of output channels is reduced by half.
[0077] The 16th layer is the batch normalization (BN) layer;
[0078] The 17th layer is the rectified linear unit (ReLU) activation layer;
[0079] The 18th layer is a convolutional layer with a stride of 1. Compared with the input, the output resolution remains unchanged and the number of output channels is 2.
[0080] As a further preferred embodiment of the present invention, in step (1), the training set samples further include rotating, translating and / or flipping the cropped images, and performing pixel value normalization processing to obtain samples;
[0081] The target region is a tumor region, an organ region or a blood vessel region.
[0082] According to another aspect of the present invention, a three-dimensional medical image segmentation system based on a long short-term memory self-attention model is provided, characterized in that it includes the following functional modules:
[0083] Image preprocessing function module: used to crop the original 3D medical image to be segmented according to the region of interest, and perform pixel value normalization on the cropped image to obtain the sample to be segmented;
[0084] A three-dimensional medical image segmentation network based on long short-term memory self-attention: first, a feature extraction module composed of multiple convolutional layers is used to convert the grayscale information in the network input image into feature information using convolution calculation to obtain initial image features; then, multiple convolutional modules with residual structures connected in sequence are used to further process the initial image features to obtain corresponding convolutional features; then, several short-term memory modules based on a multi-head self-attention mechanism are used to calculate the global correlation between convolutional features at different levels, and based on the global correlation, short-term memory features are generated in combination with the convolutional features used in the correlation calculation; then, a long-term memory module also based on a multi-head self-attention mechanism will perform a global correlation calculation on all the obtained short-term memory features, thereby generating features to be reconstructed for reconstructing the segmentation result; wherein, the long-term memory module has the same structure as any short-term memory module; finally, the reconstruction module will process the features to be reconstructed through multiple convolutional and deconvolution layers, thereby obtaining a predicted probability of the category to which each pixel of the network input image belongs, thereby achieving accurate segmentation of the target area in the three-dimensional medical image;
[0085] The three-dimensional medical image segmentation network based on long short-term memory self-attention has been trained so as to use the sample to be segmented obtained by the image preprocessing function module as the network input image input and output the segmentation result of the target area in the sample to be segmented; wherein, the training is performed using training set samples and the corresponding segmentation gold standard as labels; the training set samples are obtained by collecting three-dimensional medical images containing the target area and the known target area segmentation gold standard, cropping the original image according to the area of interest, and performing pixel value normalization processing.
[0086] As a further preferred embodiment of the present invention, the three-dimensional medical image segmentation network based on long short-term memory self-attention specifically includes:
[0087] The feature extraction module is used to extract image features. This feature extraction module reduces the resolution and increases the number of channels through two convolution calculations with a step size of 2, converts the network input image into a multi-channel feature map, and obtains the initial image features.
[0088] A plurality of convolution modules with residual structures connected in sequence are used to process the initial image features to obtain corresponding convolution features; wherein the resolution of the convolution features obtained by any convolution module is the same as the resolution of the initial image features; preferably, the plurality of convolution modules with residual structures connected in sequence are specifically four convolution modules with residual structures connected in sequence, respectively recorded as convolution module 1, convolution module 2, convolution module 3 and convolution module 4, and the convolution features of different levels obtained by these convolution modules are respectively recorded as convolution feature 1, convolution feature 2, convolution feature 3 and convolution feature 4;
[0089] Several short-term memory modules based on the multi-head self-attention mechanism, the number of the short-term memory modules is the number of the convolution modules with residual structure minus 1, wherein the first-level short-term memory module corresponds to the first two levels of convolution modules with residual structure, and each subsequent level of short-term memory module corresponds to the next level of convolution modules with residual structure in sequence; the first-level short-term memory module is used to calculate the correlation between the convolution features obtained by the first two levels of convolution modules with residual structure, and each subsequent level of short-term memory module is used to calculate the correlation between the short-term memory features obtained by the previous level of short-term memory module and the convolution features obtained by the corresponding first-level convolution module with residual structure, thereby calculating the convolution features obtained by all convolution modules with residual structure The global correlation between the product features; preferably, the several short-term memory modules based on the multi-head self-attention mechanism are specifically three short-term memory modules based on the multi-head self-attention mechanism, and the three short-term memory modules based on the multi-head self-attention mechanism are respectively short-term memory module 1, short-term memory module 2 and short-term memory module 3, and each short-term memory module has two inputs; wherein, the convolution feature 1 is downsampled and input into the short-term memory module 1 together with the convolution feature 2 to obtain the short-term memory feature 1; the short-term memory feature 1 is downsampled and input into the short-term memory module 2 together with the convolution feature 3 to obtain the short-term memory feature 2; the short-term memory feature 2 is downsampled and input into the short-term memory module 3 together with the convolution feature 4 to obtain the short-term memory feature 3;
[0090] Furthermore, except for the last short-term memory module, the short-term memory features obtained by each short-term memory module are reduced in number of channels through convolution calculation, and are added to the lowest level convolution feature input to the short-term memory module as the input of the next level convolution module; preferably, the short-term memory feature 1 and the convolution feature 2 are added as the input of the convolution module 3, and the short-term memory feature 2 and the convolution feature 3 are added as the input of the convolution module 4;
[0091] A long-term memory module based on a multi-head self-attention mechanism is used to first concatenate the short-term memory features obtained by each short-term memory module in the dimension where the channel is located to generate a merged feature, and then perform a global correlation calculation to generate a feature to be reconstructed for reconstructing the segmentation result; preferably, a long-term memory module based on a multi-head self-attention mechanism is used to first concatenate short-term memory feature 1, short-term memory feature 2, and short-term memory feature 3 in the dimension where the channel is located to generate a merged feature, and then perform a global correlation calculation on the merged feature to generate a feature to be reconstructed for reconstructing the segmentation result;
[0092] The reconstruction module is used to take the features to be reconstructed as input, use a deconvolution layer with a stride of 2 to increase the image resolution, and use two convolution layers with a stride of 1 to reduce the number of channels, so as to finally obtain the predicted probability of the category to which each pixel of the network input image belongs.
[0093] Compared with the prior art, the above technical solution conceived by the present invention can achieve the following beneficial effects, taking the tumor area as the target area as an example:
[0094] 1) This paper combines the advantages of both CNN and Transformer models, proposing a long-short-term memory self-attention structure. This improves the segmentation accuracy of tumor regions in 3D medical images, effectively eliminating interference from similar regions and improving the completeness and accuracy of multiple tumor detection. Using a CNN module for initial image feature extraction, this method overcomes the inherent shortcomings of the Transformer model and achieves better tumor segmentation results even with smaller sample sizes.
[0095] 2) The model in the present invention minimizes module stacking and model splicing that do not help improve accuracy, has fewer model parameters and lower computational complexity, improves the speed of automatic tumor segmentation in three-dimensional medical images, and is suitable for application scenarios with higher real-time requirements. BRIEF DESCRIPTION OF THE DRAWINGS
[0096] Figure 1 Schematic diagram of the overall structure of the deep learning model in this invention.
[0097] Figure 2 Schematic diagram of the self-attention module structure in the present invention.
[0098] Figure 3These are the visual effects of the segmented images corresponding to Example 1 and Comparative Examples 1-5; wherein, from left to right, the small images in each column correspond to: original image (Original image), processed by Example 1 (Our method), processed by Comparative Example 1 (UNet), processed by Comparative Example 2 (Attention UNet, i.e. AttUNet), processed by Comparative Example 3 (UNet++, i.e. NestedUNet), processed by Comparative Example 4 (TransBTS), and processed by Comparative Example 5 (UNETR). DETAILED DESCRIPTION
[0099] In order to make the objectives, technical solutions and advantages of the present invention more clearly understood, the present invention is further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely for the purpose of explaining the present invention and are not intended to limit the present invention. In addition, the technical features involved in the various embodiments of the present invention described below may be combined with each other as long as they do not conflict with each other.
[0100] This paper combines deep learning with the Short-term and Long-term Memory Self-attention (SLM-SA) model to develop a method for tumor region segmentation in three-dimensional medical images. In SLM-SA, an improved self-attention module is constructed based on the Transformer to exploit the correlation between different levels and regional features extracted by the CNN. At the same time, a short-term and long-term memory structure is designed to filter and combine the correlations between multi-level features to achieve the exclusion of similar regions and accurate detection of multiple tumors. Combined with a feature extraction and image reconstruction module based on a multi-layer convolutional neural network, accurate segmentation of tumor regions in three-dimensional medical images is achieved.
[0101] The aforementioned 3D medical image segmentation network based on long-short-term memory self-attention corresponds to the following steps: The 3D medical image segmentation network first uses a feature extraction module consisting of multiple convolutional layers to convert the grayscale information in the image into feature information using convolution calculations to obtain initial image features. Multiple convolutional modules with residual structures are then used to further process the initial image features to obtain convolutional features containing higher-level semantic information. A short-term memory module based on a multi-head self-attention mechanism then mines the global correlations between convolutional features at different levels. Based on these correlations, short-term memory features are generated in combination with the convolutional features used in the correlation calculation. Each calculation is equivalent to a forgetting and updating of memory by combining the relationship between current knowledge and deeper information. A long-term memory module, also based on a multi-head self-attention mechanism, then performs a global self-attention calculation on all input short-term memory features. Furthermore, the semantic information contained in all the obtained short-term memory features is combined to generate features to be reconstructed for the segmentation result. The short-term and long-term memory modules have the same structure. Finally, the reconstruction module will gradually restore the image resolution and reduce the number of channels through multiple convolution and deconvolution layers to obtain the predicted probability of the category to which each pixel belongs, thereby achieving accurate segmentation of the target area in the three-dimensional medical image.
[0102] Specifically, such as Figure 1 、 Figure 2 As shown ( Figure 1 The figure shows the overall structure of the deep learning model implemented in the present invention. Figure 2 Schematic diagram of the self-attention module structure implemented in the present invention):
[0103] First, the feature extraction module is used to extract image features. Two convolution operations with a stride of 2 are performed to reduce the resolution and increase the number of channels, converting the medical image into a multi-channel feature map to obtain the initial image features.
[0104] Then, the initial image features are processed through repeated convolution modules, with the output of the previous convolution module serving as the input of the next convolution module, while the resolution of the feature map remains unchanged. For example, the model includes 4 convolution modules, which respectively obtain convolution features 1-4;
[0105] After obtaining convolutional features at multiple levels, a short-term memory module based on a multi-head self-attention mechanism is used to mine the global correlation between convolutional features at different levels. Based on this correlation, short-term memory features are generated in combination with the convolutional features used in the correlation calculation. For example, the model includes three short-term memory modules, each of which has two inputs. Convolutional feature 1 is downsampled and input into short-term memory module 1 along with convolutional feature 2 to obtain short-term memory feature 1. Short-term memory feature 1 is downsampled and input into short-term memory module 2 along with convolutional feature 3 to obtain short-term memory feature 2. Short-term memory feature 2 is downsampled and input into short-term memory module 3 along with convolutional feature 4 to obtain short-term memory feature 3. Except for the last level of short-term memory modules, the short-term memory features obtained from each short-term memory module are convolutionally calculated to reduce the number of channels and then added to the convolutional features input to the short-term memory module as the input to the next level of convolutional module. In the above example, short-term memory feature 1 is added to convolutional feature 2 as the input to convolution module 3, and short-term memory feature 2 is added to convolutional feature 3 as the input to convolution module 4.
[0106] After the calculation of the short-term memory module is completed, the obtained short-term memory features 1-3 are spliced in the dimension where the channel is located to generate a merged feature, and the merged feature is input into the long-term memory module based on the multi-head self-attention mechanism. The long-term memory module will perform a global self-attention calculation to combine the semantic information in all the obtained short-term memory features to generate the features to be reconstructed for reconstructing the segmentation result.
[0107] Finally, the features to be reconstructed are input into the reconstruction module, where a deconvolution layer with a stride of 2 is used to increase the image resolution, and the last two convolution layers with a stride of 1 are used to reduce the number of channels. Finally, the predicted probability of the category to which each pixel belongs is obtained, that is, the segmentation result of the target area in the medical image.
[0108] The feature extraction module has 12 layers of detail structures. The output of the previous layer of detail structure serves as the input of the next layer of detail structure.
[0109] The first layer is a convolutional layer with a stride of 2 and an output resolution of the original Figure 1 / 2, the number of output channels is a preset value. The second layer is a batch normalization (BN) layer. The third layer is a rectified linear unit (ReLU) activation layer. The fourth layer is a convolutional layer with a stride of 1, which does not change the resolution and the number of channels. The fifth layer is a BN layer. The sixth layer is a ReLU activation layer. The seventh layer is a convolutional layer with a stride of 2, which halves the output resolution and doubles the number of output channels. The eighth layer is a BN layer. The ninth layer is a ReLU activation layer. The tenth layer is a convolutional layer with a stride of 1, which does not change the resolution and the number of channels. The eleventh layer is a BN layer. The twelfth layer is a ReLU activation layer.
[0110] Each convolutional module has 12 layers of detail structures, with the output of the previous layer serving as the input to the next layer. The first layer is a convolutional layer with a stride of 1, which does not change the resolution or number of channels. The second layer is a batch normalization layer. The third layer is a ReLU activation layer. The fourth layer is a convolutional layer with a stride of 1, which does not change the resolution or number of channels. The fifth layer is a batch normalization layer. The sixth layer is a ReLU activation layer. The output of the fifth layer is first added to the input of the first layer before serving as the input to the sixth layer. The seventh layer is a convolutional layer with a stride of 1, which does not change the resolution or number of channels. The eighth layer is a batch normalization layer. The ninth layer is a ReLU activation layer. The tenth layer is a convolutional layer with a stride of 1, which does not change the resolution or number of channels. The eleventh layer is a batch normalization layer. The twelfth layer is a ReLU activation layer. The output of the eleventh layer is first added to the input of the seventh layer before serving as the input to the twelfth layer.
[0111] For the short-term memory module, specifically: in the input of the short-term memory module, convolution features 2-4 are recorded as current features, and convolution feature 1 and short-term memory features 1-2 are recorded as memory features. For the current features, a convolution calculation will be performed before input to double the number of channels. For the memory features, a convolution calculation with a stride of 2 will be performed before input to reduce the feature map resolution to 1 / 8 of the original image to reduce the amount of calculation and speed up the inference process. The number of channels is doubled to retain sufficient feature information.
[0112] In the short-term memory module, the current feature is regarded as the query in the self-attention calculation process, and the memory feature is regarded as the key and value to perform the calculation of the self-attention module in the short-term memory module; specifically:
[0113] First, the query, key, and value are each encoded using a convolutional layer. The resulting feature encodings of the query, key, and value are then divided into a preset number of patches in the spatial dimension. Next, the attention weight is calculated based on the sequence generated for the query and key. The formula is:
[0114]
[0115] Where d is the length of the feature vector, Q and K represent the sequences corresponding to Query and Key respectively, the Softmax function is applied to the dimension corresponding to the Key sequence, and Attn represents the attention weight;
[0116] Then, the number of self-attention groups in the attention weight is doubled through a linear layer, and then input into a convolution layer for a convolution calculation with unchanged resolution and number of channels. Then, another opposite linear layer is used to restore the number of self-attention groups to obtain the final attention weight; finally, the final attention weight is matrix multiplied by Value, and the result is restored to a three-dimensional feature map through an inverse serialization operation; for the three-dimensional feature map, a learnable weight is pre-assigned according to its number of channels, so that the channels correspond to the learnable weights one-to-one, and then each channel of the three-dimensional feature map is multiplied by the corresponding learnable weight to obtain the overall output of the self-attention module; the learnable weights can be automatically updated through training;
[0117] The overall output of the self-attention module is processed by the first LN (Layer Normalization) layer and added to the current feature; the output is then passed through a feedforward neural network and a second LN layer, and the result is added to the output of the first LN layer to obtain the short-term memory feature.
[0118] For the long-term memory module, specifically: the short-term memory features output by each short-term memory module are spliced in the channel dimension to obtain the merged features, and the merged features are regarded as the query, key and value in the self-attention calculation process, and the calculation of the self-attention module in the long-term memory module is performed. Specifically:
[0119] First, the query, key, and value are each encoded using a convolutional layer. The resulting feature encodings of the query, key, and value are then divided into a preset number of patches in the spatial dimension. Next, the attention weight is calculated based on the sequence generated for the query and key. The formula is:
[0120]
[0121] Where d is the length of the feature vector, Q and K represent the sequences corresponding to Query and Key respectively, the Softmax function is applied to the dimension corresponding to the Key sequence, and Attn represents the attention weight;
[0122] Then, the number of self-attention groups in the attention weight is doubled through a linear layer, and then input into a convolution layer for a convolution calculation with unchanged resolution and number of channels. Then, another opposite linear layer is used to restore the number of self-attention groups to obtain the final attention weight; finally, the final attention weight is matrix multiplied by Value, and the result is restored to a three-dimensional feature map through an inverse serialization operation; for the three-dimensional feature map, a learnable weight is pre-assigned according to its number of channels, so that the channels correspond to the learnable weights one-to-one, and then each channel of the three-dimensional feature map is multiplied by the corresponding learnable weight to obtain the overall output of the self-attention module; the learnable weights can be automatically updated through training;
[0123] The overall output of the self-attention module is processed by the first LN (Layer Normalization) layer and added to the input of this self-attention module; the output is then passed through a feedforward neural network and a second LN layer, and the result is added to the output of the first LN layer to obtain long-term memory features, which are used as input to the reconstruction module.
[0124] The reconstruction module has 18 layers of detail structures, with the output of the previous layer serving as the input to the next layer. The first layer is a batch normalization (BN) layer. The second layer is a ReLU activation layer. The third layer is a convolutional layer with a stride of 1, which does not change the resolution or the number of channels. The fourth layer is a batch normalization (BN) layer. The fifth layer is a ReLU activation layer. The sixth layer is a deconvolutional layer with a stride of 2, which doubles the resolution but does not change the number of channels. The seventh layer is a batch normalization (BN) layer. The eighth layer is a ReLU activation layer. The ninth layer is a convolutional layer with a stride of 1, which does not change the resolution or the number of channels. The tenth layer is a batch normalization (BN) layer. The eleventh layer is a ReLU activation layer. The twelfth layer is a deconvolutional layer with a stride of 2, which doubles the resolution but does not change the number of channels. The thirteenth layer is a batch normalization (BN) layer. The fourteenth layer is a ReLU activation layer. The fifteenth layer is a convolutional layer with a stride of 1, which maintains the resolution but halves the number of channels. The sixteenth layer is a batch normalization (BN) layer. The seventeenth layer is a ReLU activation layer. The 18th layer is a convolutional layer with a step size of 1, unchanged resolution, and 2 output channels, which obtains the segmentation results of the target area in the medical image.
[0125] Of course, the feature extraction module, the convolution module with a residual structure, and the reconstruction module can also be constructed with reference to relevant existing technologies.
[0126] In addition, the above network can use the sum of cross entropy and Dice as the loss function to train the model. The formula is as follows:
[0127]
[0128]
[0129]
[0130] Where TP, FN, and FP are the number of true positive, false negative, and false positive pixels, respectively. N represents the number of segmented categories, and M represents the total number of samples. If the true category of sample i is c, y ic Takes 1, otherwise takes 0. p ic is the predicted probability that sample i belongs to category c, and a is 0.5.
[0131] Taking a medical image with a size of 96×96×96 (the unit of each dimension value is pixel; the same below) as an example, Figure 1 The figure shows the overall structure of the deep learning model implemented in the present invention. The calculation process is applied to 3D medical images to obtain accurate tumor region segmentation results. The specific steps may include:
[0132] (1) The input image with a size of 1×96×96×96 is calculated by the feature extraction module to obtain a feature map of 32×24×24×24. The input and output feature maps of each convolution module have the same size.
[0133] (2) For the calculation of the short-term memory module at each stage, the memory feature is convolutionally scaled from 32×24×24×24 to 64×12×12×12. The current feature is convolutionally scaled from 32×24×24×24 to 64×24×24×24. The output size of the short-term memory module is 64×24×24×24. The convolution is used to reduce the channel size to 32×24×24×24 to obtain the final output. The output is then added to the current feature as an update.
[0134] (3) The output of each short-term memory module is concatenated in the channel dimension to obtain the input of the long-term memory module, which has a size of 96×24×24×24, and the output size is the same.
[0135] (4) The reconstruction module gradually obtains a 2×96×96×96 category prediction probability matrix, i.e., the segmentation result, from the feature map of size 96×24×24×24 through convolution calculation.
[0136] like Figure 2 The figure shows a schematic diagram of the self-attention module structure implemented in the present invention. Taking the input feature maps of size 64×24×24×24 (input 1) and 64×12×12×12 (input 2) as examples, the specific steps include:
[0137] (1) Apply the corresponding convolution calculation and encoding operation with a patch size of 2×2×2 and the number of attention heads of 4, respectively. The query is obtained from input 1 with a size of 4×1728×128, and the key and value are obtained from input 2 with a size of 4×216×128.
[0138] (2) The query and key matrices are multiplied to obtain a weight map with a size of 4×1728×216. After the Reattention calculation, the size remains unchanged.
[0139] (3) The weight map is multiplied by the Value matrix, and then the features are restored to a three-dimensional feature map. Each channel is multiplied by a learnable parameter, and the output feature map size is 64×24×24×24.
[0140] Example 1
[0141] Based on the above method, actual 3D medical image segmentation processing is performed, which specifically includes the following steps:
[0142] (1) Using the KiTS19 dataset, the left and right kidneys of each sample were separated, retaining only the data containing the tumor. A total of 215 96×96×96 kidney images were obtained. Based on the provided kidney segmentation labels, binary labels containing only tumor and non-tumor were generated. Through rotation, translation, and flipping data augmentation, 5375 samples were obtained, of which 4350, 525, and 500 samples were used for training, validation, and testing, respectively.
[0143] (2) The model proposed in the present invention is trained using the above training set and validation set, the optimal model parameters are loaded, and the model is applied to the test set to output the prediction results of the tumor area in the image. Finally, the segmentation results are quantitatively evaluated.
[0144] Furthermore, in order to verify the method of the present invention, the following comparative examples were designed (each comparative example also uses the KiTS19 dataset):
[0145] Comparative Example 1
[0146] The tumor region segmentation task in 3D medical images was achieved using U-Net (convolutional networks for biomedical images segmentation. In: Proceedings of the 18th International Conference on Medical Image Computing and Computer-Assisted Intervention, 234-241, 2015), trained using the same dataset, learning rate, number of iterations, and optimizer parameters as the method of the present invention.
[0147] Comparative Example 2
[0148] The Attention U-Net (Batch target recognition count under complex conditions based on att-unet. In: Proceedings of the 5th International Conference on Computer and Communication Systems, 261-265, 2020) was used to implement the tumor area segmentation task in 3D medical images, and the same dataset, learning rate, number of iterations, and optimizer parameters were used for training as the method of the present invention.
[0149] Comparative Example 3
[0150] The tumor region segmentation task in 3D medical images was implemented using UNet++ (Unet plus plus: a nested u-net architecture for medical images segmentation. Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support 11045, 3-11, 2018), using the same dataset, learning rate, number of iterations, and optimizer parameters as the method of the present invention.
[0151] Comparative Example 4
[0152] TransBTS (multimodal brain tumor segmentation using transformer. In: Proceedings of the International Conference on Medical Image Computing and Computer Assisted Intervention, 109-119, 2021) was used to implement the segmentation task of tumor areas in three-dimensional medical images, using the same dataset, learning rate, number of iterations, and optimizer parameter training as the method of the present invention.
[0153] Comparative Example 5
[0154] Use UNETR (UNETR: transformers for 3D medical image segmentation. In: Proceedings of the IEEE / CVF Winter Conference on Applications of Computer Vision, 1748-1758, 2022) to implement the task of segmenting tumor areas in 3D medical images, using the same dataset, learning rate, number of iterations, and optimizer parameters as the method of the present invention for training.
[0155] Result Analysis
[0156] To demonstrate the advantages of the present invention, the segmentation effects of Example 1 were compared with those of Comparative Examples 1-5. The quantitative comparison was evaluated using the Dice coefficient, intersection over union (IOU), 95% Hausdorff distance (HD95), and relative volume error (RVE). IOU, HD95, and RVE are defined as follows:
[0157]
[0158] HD95=max k95% [d(X,Y),d(Y,X)]
[0159]
[0160]
[0161] Where min is the minimum value, max is the minimum value. k95%Indicates taking the maximum value of the 95th percentile, abs represents the calculation of the absolute value, |R a | and |R b |Represents region R a and R b volume.
[0162] Table 1 lists the four quantitative evaluation metrics for the segmentation results of Example 1 and Comparative Examples 1-5, as well as the running speed, computational complexity, and parameter count. As can be seen from the table, Example 1 significantly improves on the four quantitative evaluation metrics of Dice, IoU, HD95, and RVE compared to Comparative Examples 1-5. Furthermore, Example 1 has only 0.871M parameters and a computational complexity of 52G. At 55 frames per second (FPS), it achieves a Dice of 82.4%, enabling more accurate real-time segmentation of tumors in 3D medical images.
[0163] Table 1 Segmentation accuracy of each segmentation method on KiTS19
[0164]
[0165] In order to more intuitively demonstrate the superiority of the present invention, we provide the visual effect diagrams of the segmented images corresponding to Example 1 and Comparative Examples 1-5. Figure 3 As shown in the figure, the first column shows the original medical image, and columns 2-7 show the segmentation results obtained by different models. The white, light gray, and dark gray areas represent true positive, false positive, and false negative areas in the tumor prediction results, respectively. Clearly, the tumor boundaries predicted by our method are highly consistent with their benchmarks, and our method can accurately exclude highly similar targets in the second row of samples.
[0166] The above embodiments are merely examples. In addition to tumors, the method and system of the present invention are also applicable to segmenting three-dimensional medical images of other target areas such as organs and blood vessels.
[0167] It will be easily understood by those skilled in the art that the above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention should be included in the scope of protection of the present invention.
Claims
1. A 3D medical image segmentation method based on a long short-term memory self-attention model, capable of segmenting target area images in medical images, characterized by: The following steps are involved: (1) Collect 3D medical images containing target regions and known gold standards for target region segmentation, crop the original images according to the region of interest, and perform pixel value normalization to obtain training set samples; (2) For the original three-dimensional medical image to be segmented, the original image is cropped according to the region of interest, and the pixel values of the cropped image are normalized to obtain the sample to be segmented; (3) Constructing a three-dimensional medical image segmentation network based on long-short-term memory self-attention, the network first uses a feature extraction module composed of multiple convolutional layers, and uses convolution calculation to convert the grayscale information in the network input image into feature information to obtain initial image features; then uses multiple convolution modules with residual structures connected in sequence to further process the initial image features to obtain corresponding convolution features; then uses several short-term memory modules based on multi-head self-attention mechanism to calculate the global correlation between convolution features at different levels, and generates short-term memory features based on the global correlation and the convolution features used in the correlation calculation; then, a long-term memory module also based on the multi-head self-attention mechanism will perform a global correlation calculation on all the obtained short-term memory features, thereby generating features to be reconstructed for reconstructing the segmentation results; wherein the long-term memory module has the same structure as any short-term memory module; finally, the reconstruction module will process the features to be reconstructed through multiple convolution and deconvolution layers, thereby obtaining the predicted probability of the category to which each pixel of the network input image belongs, thereby achieving accurate segmentation of the target area in the three-dimensional medical image; Utilizing the training set samples obtained in step (1) and the corresponding segmentation gold standard as labels for training and constructing the three-dimensional medical image segmentation network, the trained network is capable of segmenting the target area contained in the network input image; (4) The sample to be segmented obtained in step (2) is input as a network input image into the trained long short-term memory self-attention based three-dimensional medical image segmentation network obtained in step (3), and the segmentation result of the target area in the sample to be segmented can be output.
2. The three-dimensional medical image segmentation method according to claim 1, wherein: The three-dimensional medical image segmentation network based on long short-term memory self-attention constructed in step (3) specifically includes: The feature extraction module is used to extract image features. This feature extraction module reduces the resolution and increases the number of channels through two convolution calculations with a step size of 2, converts the network input image into a multi-channel feature map, and obtains the initial image features. A plurality of convolution modules with residual structures connected in sequence are used to process the initial image features to obtain corresponding convolution features; wherein the resolution of the convolution features obtained by any convolution module is the same as the resolution of the initial image features; Several short-term memory modules based on the multi-head self-attention mechanism, the number of the short-term memory modules is the number of the convolution modules with residual structures minus 1, wherein the first-level short-term memory module corresponds to the first two levels of convolution modules with residual structures, and each subsequent level of short-term memory module corresponds sequentially to the next level of convolution modules with residual structures; the first-level short-term memory module is used to calculate the correlation between the convolution features obtained by the first two levels of convolution modules with residual structures, and each subsequent level of short-term memory module is used to calculate the correlation between the short-term memory features obtained by the previous level of short-term memory module and the convolution features obtained by the corresponding first-level convolution module with residual structure, thereby calculating the global correlation between the convolution features obtained by all convolution modules with residual structures; In addition, except for the last short-term memory module, the short-term memory features obtained by each short-term memory module are reduced by convolution calculation to reduce the number of channels, and added to the lowest level convolution feature input to the short-term memory module as the input of the next level convolution module; The long-term memory module based on the multi-head self-attention mechanism is used to first concatenate the short-term memory features obtained by each short-term memory module in the channel dimension to generate a merged feature, and then perform a global correlation calculation to generate the features to be reconstructed for reconstructing the segmentation result; The reconstruction module is used to take the features to be reconstructed as input, use a deconvolution layer with a stride of 2 to increase the image resolution, and use two convolution layers with a stride of 1 to reduce the number of channels, so as to finally obtain the predicted probability of the category to which each pixel of the network input image belongs.
3. The three-dimensional medical image segmentation method according to claim 2, wherein: The multiple convolution modules with residual structures connected in sequence are specifically four convolution modules with residual structures connected in sequence, which are respectively recorded as convolution module 1, convolution module 2, convolution module 3 and convolution module 4. The convolution features of different levels obtained by these convolution modules are respectively recorded as convolution feature 1, convolution feature 2, convolution feature 3 and convolution feature 4; The multiple short-term memory modules based on the multi-head self-attention mechanism are specifically three short-term memory modules based on the multi-head self-attention mechanism. These three short-term memory modules based on the multi-head self-attention mechanism are respectively short-term memory module 1, short-term memory module 2 and short-term memory module 3. Each short-term memory module has two inputs; among them, the convolution feature 1 is input into the short-term memory module 1 after downsampling and the convolution feature 2 to obtain the short-term memory feature 1; the short-term memory feature 1 is input into the short-term memory module 2 after downsampling and the convolution feature 3 to obtain the short-term memory feature 2; the short-term memory feature 2 is input into the short-term memory module 3 after downsampling and the convolution feature 4 to obtain the short-term memory feature 3; The short-term memory feature 1 and the convolution feature 2 are added as the input of the convolution module 3, and the short-term memory feature 2 and the convolution feature 3 are added as the input of the convolution module 4; The long-term memory module based on the multi-head self-attention mechanism is used to first concatenate short-term memory feature 1, short-term memory feature 2, and short-term memory feature 3 in the channel dimension to generate a merged feature, and then perform a global correlation calculation on the merged feature to generate the features to be reconstructed for reconstructing the segmentation result.
4. The three-dimensional medical image segmentation method according to claim 2, wherein: The feature extraction module has a total of 12 layers of detail structures, and the output of the previous layer of detail structure serves as the input of the next layer of detail structure, where: The first layer is a convolutional layer with a stride of 2; the output resolution is 1 / 2 of the network input image, and the number of output channels is a pre-set value; The second layer is the batch normalization (BN) layer; The third layer is the rectified linear unit (ReLU) activation layer; The fourth layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels; The 5th layer is the batch normalization (BN) layer; The 6th layer is the rectified linear unit (ReLU) activation layer; The 7th layer is a convolutional layer with a stride of 2; compared to the input, the output resolution is halved and the number of output channels is doubled; The 8th layer is the batch normalization (BN) layer; The 9th layer is the rectified linear unit (ReLU) activation layer; The 10th layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels; The 11th layer is the batch normalization (BN) layer; The 12th layer is a rectified linear unit (ReLU) activation layer.
5. The three-dimensional medical image segmentation method according to claim 2, wherein: For any convolution module, there are 12 layers of detail structures, and the output of the previous layer of detail structure serves as the input of the next layer of detail structure, where: The first layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels; The second layer is the batch normalization (BN) layer; The third layer is the rectified linear unit (ReLU) activation layer; The fourth layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels; The 5th layer is the batch normalization (BN) layer; The sixth layer is a rectified linear unit (ReLU) activation layer; its input is the sum of the output of the fifth layer and the input of the first layer; The 7th layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels; The 8th layer is the batch normalization (BN) layer; The 9th layer is the rectified linear unit (ReLU) activation layer; The 10th layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels; The 11th layer is the batch normalization (BN) layer; The 12th layer is a rectified linear unit (ReLU) activation layer; its input is the sum of the output of the 11th layer and the input of the 7th layer.
6. The three-dimensional medical image segmentation method according to claim 2, wherein: For the short-term memory module: In the input of the short-term memory module, for the first-level short-term memory module, the convolution feature obtained by the first-level convolution module with a residual structure is recorded as the memory feature, and the convolution feature obtained by the second-level convolution module with a residual structure is recorded as the current feature; for each subsequent level of short-term memory module, the short-term memory feature obtained by the previous level of short-term memory module is recorded as the memory feature, and the convolution feature obtained by the first-level convolution module with a residual structure corresponding to the short-term memory module is recorded as the current feature; For any short-term memory module, for its current feature, a convolution calculation is performed before input to double the number of channels; for its memory feature, a convolution calculation with a step size of 2 is performed before input, while reducing the resolution and increasing the number of channels; the resolution reduction is specifically halved, and the channel increase is specifically doubled; For any short-term memory module, the current feature is regarded as the query in the self-attention calculation process, and the memory feature is regarded as the key and value to perform the calculation of the self-attention module in the short-term memory module; specifically: First, the query, key, and value are each encoded using a convolutional layer. The resulting feature encodings of the query, key, and value are then divided into a preset number of patches in the spatial dimension. Next, the attention weight is calculated based on the sequence generated for the query and key. The formula is: Where d is the length of the feature vector, Q and K represent the sequences corresponding to Query and Key respectively, and the Softmax function is applied to the dimension corresponding to the Key sequence. Attn represents the attention weight; Then, the number of self-attention groups in the attention weight is doubled through a linear layer, and then input into a convolution layer for a convolution calculation with unchanged resolution and number of channels. Then, another opposite linear layer is used to restore the number of self-attention groups to obtain the final attention weight; finally, the final attention weight is matrix multiplied by Value, and the result is restored to a three-dimensional feature map through an inverse serialization operation; for the three-dimensional feature map, a learnable weight is pre-assigned according to its number of channels, so that the channels correspond to the learnable weights one-to-one, and then each channel of the three-dimensional feature map is multiplied by the corresponding learnable weight to obtain the overall output of the self-attention module; the learnable weights can be automatically updated through training; The output of the self-attention module as a whole is processed by the first LN layer and added to the current feature; The output is then passed through a feedforward neural network and a second LN layer, and the result is added to the output of the first LN layer to obtain the short-term memory feature.
7. The three-dimensional medical image segmentation method according to claim 2, wherein: For the long-term memory module: First, the short-term memory features output by each short-term memory module are spliced in the channel dimension to obtain the merged features. The merged features are regarded as the query, key and value in the self-attention calculation process, and the calculation of the self-attention module in the long-term memory module is performed. Specifically: First, the query, key, and value are each encoded using a convolutional layer. The resulting feature encodings of the query, key, and value are then divided into a preset number of patches in the spatial dimension. Next, the attention weight is calculated based on the sequence generated for the query and key. The formula is: Where d is the length of the feature vector, Q and K represent the sequences corresponding to Query and Key respectively, and the Softmax function is applied to the dimension corresponding to the Key sequence. Attn represents the attention weight; Then, the number of self-attention groups in the attention weight is doubled through a linear layer, and then input into a convolution layer for a convolution calculation with unchanged resolution and number of channels. Then, another opposite linear layer is used to restore the number of self-attention groups to obtain the final attention weight; finally, the final attention weight is matrix multiplied by Value, and the result is restored to a three-dimensional feature map through an inverse serialization operation; for the three-dimensional feature map, a learnable weight is pre-assigned according to its number of channels, so that the channels correspond to the learnable weights one-to-one, and then each channel of the three-dimensional feature map is multiplied by the corresponding learnable weight to obtain the overall output of the self-attention module; the learnable weights can be automatically updated through training; The output of the self-attention module as a whole is processed by the first LN layer and then added to the merged feature; The output is then passed through a feedforward neural network and a second LN layer, and the result is added to the output of the first LN layer to obtain long-term memory features, which are used as input to the reconstruction module.
8. The three-dimensional medical image segmentation method according to claim 2, wherein: The reconstruction module has a total of 18 layers of detail structures, and the output of the previous layer of detail structure serves as the input of the next layer of detail structure, where: The first layer is the batch normalization (BN) layer; The second layer is the rectified linear unit (ReLU) activation layer; The third layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels; The fourth layer is the batch normalization (BN) layer; The fifth layer is the rectified linear unit (ReLU) activation layer; The sixth layer is a deconvolution layer with a stride of 2. Compared with the input, the output resolution is doubled, and the number of output channels remains unchanged. The 7th layer is the batch normalization (BN) layer; The 8th layer is the rectified linear unit (ReLU) activation layer; The 9th layer is a convolutional layer with a stride of 1, which does not change the input resolution and number of channels; The 10th layer is the batch normalization (BN) layer; The 11th layer is the rectified linear unit (ReLU) activation layer; The 12th layer is a deconvolution layer with a stride of 2. Compared with the input, the output resolution is doubled, and the number of output channels remains unchanged. The 13th layer is the batch normalization (BN) layer; The 14th layer is the rectified linear unit (ReLU) activation layer; The 15th layer is a convolutional layer with a stride of 1. Compared to the input, the output resolution remains unchanged, but the number of output channels is reduced by half. The 16th layer is the batch normalization (BN) layer; The 17th layer is the rectified linear unit (ReLU) activation layer; The 18th layer is a convolutional layer with a stride of 1. Compared with the input, the output resolution remains unchanged and the number of output channels is 2.
9. The three-dimensional medical image segmentation method according to claim 1, wherein: In the step (1), the training set samples further include rotating, translating and / or flipping the cropped image, and performing pixel value normalization processing to obtain samples; The target region is a tumor region, an organ region or a blood vessel region.
10. A 3D medical image segmentation system based on long short-term memory self-attention model, characterized in that: Includes the following functional modules: Image preprocessing function module: used to crop the original 3D medical image to be segmented according to the region of interest, and perform pixel value normalization on the cropped image to obtain the sample to be segmented; A three-dimensional medical image segmentation network based on long short-term memory self-attention: first, a feature extraction module composed of multiple convolutional layers is used to convert the grayscale information in the network input image into feature information using convolution calculation to obtain initial image features; then, multiple convolutional modules with residual structures connected in sequence are used to further process the initial image features to obtain corresponding convolutional features; then, several short-term memory modules based on a multi-head self-attention mechanism are used to calculate the global correlation between convolutional features at different levels, and based on the global correlation, short-term memory features are generated in combination with the convolutional features used in the correlation calculation; then, a long-term memory module also based on a multi-head self-attention mechanism will perform a global correlation calculation on all the obtained short-term memory features, thereby generating features to be reconstructed for reconstructing the segmentation result; wherein, the long-term memory module has the same structure as any short-term memory module; finally, the reconstruction module will process the features to be reconstructed through multiple convolutional and deconvolution layers, thereby obtaining a predicted probability of the category to which each pixel of the network input image belongs, thereby achieving accurate segmentation of the target area in the three-dimensional medical image; The three-dimensional medical image segmentation network based on long short-term memory self-attention has been trained so as to use the sample to be segmented obtained by the image preprocessing function module as the network input image input and output the segmentation result of the target area in the sample to be segmented; wherein, the training is performed using training set samples and the corresponding segmentation gold standard as labels; the training set samples are obtained by collecting three-dimensional medical images containing the target area and the known target area segmentation gold standard, cropping the original image according to the area of interest, and performing pixel value normalization processing.
11. The three-dimensional medical image segmentation system according to claim 10, wherein: The three-dimensional medical image segmentation network based on long short-term memory self-attention specifically includes: The feature extraction module is used to extract image features. This feature extraction module reduces the resolution and increases the number of channels through two convolution calculations with a step size of 2, converts the network input image into a multi-channel feature map, and obtains the initial image features. A plurality of convolution modules with residual structures connected in sequence are used to process the initial image features to obtain corresponding convolution features; wherein the resolution of the convolution features obtained by any convolution module is the same as the resolution of the initial image features; Several short-term memory modules based on the multi-head self-attention mechanism, the number of the short-term memory modules is the number of the convolution modules with residual structures minus 1, wherein the first-level short-term memory module corresponds to the first two levels of convolution modules with residual structures, and each subsequent level of short-term memory module corresponds sequentially to the next level of convolution modules with residual structures; the first-level short-term memory module is used to calculate the correlation between the convolution features obtained by the first two levels of convolution modules with residual structures, and each subsequent level of short-term memory module is used to calculate the correlation between the short-term memory features obtained by the previous level of short-term memory module and the convolution features obtained by the corresponding first-level convolution module with residual structure, thereby calculating the global correlation between the convolution features obtained by all convolution modules with residual structures; In addition, except for the last short-term memory module, the short-term memory features obtained by each short-term memory module are reduced by convolution calculation to reduce the number of channels, and added to the lowest level convolution feature input to the short-term memory module as the input of the next level convolution module; The long-term memory module based on the multi-head self-attention mechanism is used to first concatenate the short-term memory features obtained by each short-term memory module in the channel dimension to generate a merged feature, and then perform a global correlation calculation to generate the features to be reconstructed for reconstructing the segmentation result; The reconstruction module is used to take the features to be reconstructed as input, use a deconvolution layer with a stride of 2 to increase the image resolution, and use two convolution layers with a stride of 1 to reduce the number of channels, so as to finally obtain the predicted probability of the category to which each pixel of the network input image belongs.
12. The three-dimensional medical image segmentation system according to claim 11, wherein: The multiple convolution modules with residual structures connected in sequence are specifically four convolution modules with residual structures connected in sequence, which are respectively recorded as convolution module 1, convolution module 2, convolution module 3 and convolution module 4. The convolution features of different levels obtained by these convolution modules are respectively recorded as convolution feature 1, convolution feature 2, convolution feature 3 and convolution feature 4; The multiple short-term memory modules based on the multi-head self-attention mechanism are specifically three short-term memory modules based on the multi-head self-attention mechanism. These three short-term memory modules based on the multi-head self-attention mechanism are respectively short-term memory module 1, short-term memory module 2 and short-term memory module 3. Each short-term memory module has two inputs; among them, the convolution feature 1 is input into the short-term memory module 1 after downsampling and the convolution feature 2 to obtain the short-term memory feature 1; the short-term memory feature 1 is input into the short-term memory module 2 after downsampling and the convolution feature 3 to obtain the short-term memory feature 2; the short-term memory feature 2 is input into the short-term memory module 3 after downsampling and the convolution feature 4 to obtain the short-term memory feature 3; The short-term memory feature 1 and the convolution feature 2 are added as the input of the convolution module 3, and the short-term memory feature 2 and the convolution feature 3 are added as the input of the convolution module 4; The long-term memory module based on the multi-head self-attention mechanism is used to first concatenate short-term memory feature 1, short-term memory feature 2, and short-term memory feature 3 in the channel dimension to generate a merged feature, and then perform a global correlation calculation on the merged feature to generate the features to be reconstructed for reconstructing the segmentation result.
Citation Information
Patent Citations
Remote sensing image content description method based on variational self-attention reinforcement learning
CN111126282A
Multi-channel neural network method for long-term memory learning
CN113065650A