Medical image segmentation method and system based on multi-axis attention
By introducing a multi-axis attention module in U-Net, combining the main feature retention and compact separation attention module, the information loss and overfitting problems of U-Net in complex structures and small-objective segmentation are solved, and medical images are improved segmentation accuracy and model performance.
Patent Information
- Application Number
- CN202311724150.2
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2023-12-15
- Publication Date
- 2024-12-17
AI Technical Summary
Existing U-Net medical image segmentation method based on CBAM has oversegmentation or information loss in complex structures and small-objective segmentation. Jump connections lead to information redundancy and overfitting. The expression of CBAM attention mechanism is limited, especially in small samples. And the model is overfitted on medical image datasets with uneven data distribution.
Using a multi-axis attention medical image segmentation method, the U-Net structure is improved through the main feature retention module, compact separation attention module and multi-axis attention module, and the channel axis, height axis and width axis attention module are introduced to optimize feature extraction and expressive ability, reduce computational overhead, and enhance model performance.
It effectively solves the problems of information loss and information redundancy in complex structures and small-object segmentation, improves the expression ability and performance of the model, especially when the data set is constrained, and improves the accuracy and reliability of the model. sex.
Smart Images

Figure CN119152216A_ABST
Abstract
Description
Technical Field
[0001] The present invention belongs to the fields of computer science and medical imaging, and more specifically, to a medical image segmentation method and system based on multi-axis attention. Background Art
[0002] Medical image segmentation plays a vital role in the field of medical imaging, helping doctors accurately locate and segment abnormal areas in images, such as tumors, lesions, blood vessels, etc. By accurately depicting the lesion area, doctors can more accurately diagnose diseases and plan treatments. During surgery, medical image segmentation can provide doctors with real-time, high-resolution image information, assisting doctors to better navigate and locate surgical targets, and helping to improve the accuracy and safety of surgery.
[0003] There are two commonly used medical image segmentation models. The first is based on the Convolutional Neural Network (CNN) model, such as U-Net and its variants, DeepLab series, and Fully Convolutional Network (FCN). The second is based on the Transformer structure model, such as TransUNet, SwinUNet, SegFormer, etc. Since the Transformer-based segmentation model has higher complexity and more complex structure, the existing medical image segmentation framework is still concentrated on the Convolutional Neural Network model, and among these models, U-Net and its variants are still the first choice for segmentation due to their better performance and intuitive structure.
[0004] When considering improving the U-Net structure, how to introduce smaller and more efficient modules is the key to improvement. This problem can be partially solved by introducing the Convolutional Block Attention Module (CBAM).
[0005] However, the current U-Net medical image segmentation method based on CBAM still has some defects that cannot be ignored:
[0006] First, the U-Net model may suffer from over-segmentation or information loss for some complex structures and small objects due to a lack of sufficient attention to the region of interest. In addition, the jump connection in the U-Net model directly concatenates low-level features with high-level features, which can lead to problems such as information redundancy and overfitting.
[0007] Second, the introduced CBAM attention mechanism improves the model’s expressiveness to a limited extent;
[0008] Third, in datasets such as medical images with few samples and uneven data distribution, excessive superposition of computing units by introducing CBAM will lead to model overfitting. Summary of the invention
[0009] In response to the above defects or improvement needs of the prior art, the present invention provides a medical image segmentation method and system based on multi-axis attention, which aims to solve the technical problem that the U-Net model in the existing CBAM-based U-Net medical image segmentation method may suffer from over-segmentation or information loss for certain complex structures and small targets due to lack of sufficient attention to the region of interest, and the technical problem that the jump connection in the U-Net model directly splices low-level features with high-level features, thereby leading to technical problems such as information redundancy and overfitting, and the technical problem that the introduced CBAM attention mechanism improves the limited expressiveness of the model, and the technical problem that in data sets such as medical images with few samples and uneven data distribution, excessive superposition of computing units by introducing CBAM will lead to model overfitting.
[0010] To achieve the above object, according to one aspect of the present invention, a medical image segmentation method based on multi-axis attention is provided, comprising the following steps:
[0011] (1) acquiring a medical image, and preprocessing the medical image to obtain a preprocessed medical image;
[0012] (2) The medical image preprocessed in step (1) is input into a pre-trained multi-axis attention-based medical image segmentation model to obtain the final lesion segmentation prediction result.
[0013] Preferably, the medical image is a two-dimensional static image, which may be a tumor image acquired by an ultrasound device or a dermoscopic image;
[0014] The preprocessing process in step (1) first extracts a region of interest from the medical image according to the specific location of the lesion in the preset segmentation label, and then converts the extracted region of interest into a medical image of 256×256 size and jpg format as the preprocessed medical image;
[0015] The lesion segmentation result in step (2) is a binary image corresponding to the medical image obtained in step (1), in which black is the background and white is the foreground (i.e., the lesion).
[0016] The medical image segmentation model based on multi-axis attention includes a main feature retention module, an encoder, a decoder, and a multi-axis attention module connected between the main feature retention module and the decoder.
[0017] The input of the main feature retention module is the preprocessed medical image, and the output is a feature map with a size of 256×256 and 64 channels;
[0018] The encoder contains four compact separation attention modules. The input of the encoder is the feature map of size 256×256 and 64 channels output by the main feature retention module, and its output is the intermediate feature map;
[0019] The multi-axis attention module includes a channel-axis attention module, a height-axis attention module, and a width-axis attention module. The input of the multi-axis attention module is a feature map with a size of 256×256 and 64 channels output by the main feature retention module, and the output is a multi-axis attention feature map with the same size as the input feature map.
[0020] The decoder contains four compact separation attention modules. The input of the decoder is the feature map with a size of 32×32 and 512 channels output by the third compact separation attention module in the encoder, the intermediate feature map with a size of 16×16 and 512 channels output by the fourth compact separation attention module, and the multi-axis attention feature map output by the multi-axis attention module. The output is the final lesion segmentation prediction result.
[0021] Preferably, the main feature retention module first uses a 7×7 convolution kernel to convolve the preprocessed medical image, and the convolution kernel has in_ch input channels and out_ch output channels. Then, the convolved medical image is corrected using the ReLU function to obtain the corrected features, and then the corrected features are batch normalized to obtain the normalized features. Then, another 7×7 convolution kernel is used to convolve the normalized features, and then the convolution features are corrected using the ReLU function to obtain the corrected features. Then, the corrected features are batch normalized to obtain a normalized intermediate feature map. Then, a 1×1 convolution kernel with no padding is used to convolve the normalized intermediate feature map, and then the convolution features are corrected using the ReLU function to obtain the corrected features. Then, the corrected features are batch normalized to obtain a normalized feature map with a size of 256×256 and 64 channels.
[0022] The encoder first inputs the feature map output by the main feature retention module into the maximum pooling layer to obtain a feature map with a size of 128×128 and a number of channels of 64. Next, the encoder inputs the feature map into the first compact separation attention module to obtain a feature map with a size of 128×128 and a number of channels of 128; then, the feature map with a size of 128×128 and a number of channels of 128 obtained by the first compact separation attention module is input into the maximum pooling layer, so that the height and width of the feature map are halved, and the number of channels remains unchanged, so as to achieve spatial downsampling, thereby obtaining a feature map with a size of 64×64 and a number of channels of 128 after downsampling; then, the feature map with a size of 64×64 and a number of channels of 128 after downsampling is input into the second compact separation attention module, and the feature map with a size of 64×64 and a number of channels of 128 is obtained. To an output feature map with a size of 64×64 and 256 channels, the feature map is then input into the maximum pooling layer to obtain a downsampled feature map with a size of 32×32 and 256 channels; then, the downsampled 256-channel feature map is input into the third compact separation attention module to obtain an output feature map with a size of 32×32 and 512 channels, to obtain a downsampled feature map with a size of 16×16 and 512 channels; finally, the 512-channel feature map is input into the fourth compact separation attention module to obtain an intermediate feature map with a size of 16×16 and 512 channels.
[0023] Preferably, the channel axis attention module first uses average pooling and maximum pooling on the channel dimension to obtain two channel dimension information, and then inputs the two channel dimension information into the double convolution module respectively to obtain two channel attention maps with a size of 1×1 and a channel number of 64. Thereafter, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F c Then, F c Perform point-by-point multiplication with the feature map of size 256×256 and 64 channels output by the main feature retention module to obtain the feature map F′ c ; Then, the obtained feature map F′ c Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 256×256 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. The sigmoid activation function is then used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map. Then, the spatial attention map With the feature map F″ c Perform multiplication operation to obtain the final output feature map F c "as follows:
[0024]
[0025] Then, the feature map F″ is output c The height axis attention module and the width axis attention module are respectively input to extract the height axis feature map and the width axis feature map to obtain the feature map F″ h and F″ w ;
[0026] Finally, the multi-axis attention module converts the final height axis feature map F″ h And the final width axis feature map F″ w Superimpose them to obtain the multi-axis attention feature map F:
[0027] F=F″ h +F″ w .
[0028] Preferably, the decoder first performs an upsampling operation with a step size of 2*2 and a stride of 2 on the intermediate feature map of 16×16 and 512 channels output by the fourth compact separation attention module to obtain a feature map of size 32×32 and 512 channels, and then splices the feature map with the feature map of size 32×32 and 512 channels output by the third compact separation attention module in the channel dimension to obtain a spliced feature map of size 32×32 and 1024 channels; then, the spliced feature map is input into the fifth compact separation attention module to obtain a feature map of size 32×32 and 256 channels; then, the feature map is upsampled with a step size of 2*2 and a stride of 2 A sampling operation is performed to obtain a feature map with a size of 64×64 and a number of channels of 256. Subsequently, the feature map is spliced in the channel dimension with the feature map with a size of 64×64 and a number of channels of 256 output by the second compact separation attention module to obtain a feature map with a size of 64×64 and a number of channels of 512; then, the feature map is input into the sixth compact separation attention module to obtain a feature map with a size of 64×64 and a number of channels of 128. Subsequently, an upsampling operation with a step size of 2*2 and a stride of 2 is performed on the feature map, and the obtained feature map is spliced in the channel dimension with the feature map with a size of 128×128 and a number of channels of 128 output by the first compact separation attention module to obtain To a spliced feature map with a size of 128×128 and a number of channels of 256, and then, the spliced feature map is input into the seventh compact separation attention module to obtain a feature map with a size of 128×128 and a number of channels of 64; then, the feature map is upsampled with a step size of 2*2 and a stride of 2 to obtain a feature map with a size of 256×256 and a dimension of 64; then, the feature map is spliced with the multi-axis attention feature map output by the multi-axis attention module into a feature map with a size of 256×256 and a number of channels of 128; then, the feature map is input into the eighth compact separation attention module to obtain a feature map with a size of 256×256 and a number of channels of 64; thereafter, through a 1 The ×1 convolutional layer reduces the number of channels of the feature map to two channels, and then obtains the probability distribution matrix of the tumor location; then, the probability distribution matrix is mapped to the [0,1] interval through the Softmax function; finally, in the probability distribution matrix of the first channel, the element value with a probability value greater than 0.5 in the probability distribution matrix is set to 1, that is, the background (that is, the position is considered to be the background, that is, there is no lesion), and the other element values are set to 0; in the second channel, the element value with a probability value greater than 0.5 in the probability distribution matrix is set to 1, that is, the foreground (that is, the position is considered to be the foreground, that is, the lesion), and the other element values are set to 0; finally, the binarized probability distribution matrix is saved as a binarized image, which is the final lesion segmentation prediction result.
[0029] Preferably, the medical image segmentation model based on multi-axis attention is obtained by the following steps:
[0030] (2-1) Obtain a medical image dataset, preprocess the medical image dataset to obtain a preprocessed medical image dataset, and divide the preprocessed medical image dataset into a training set and a test set in a ratio of 8:2.
[0031] (2-2) Set the total number of iterations epoch of the training process and initialize the learning rate η of the medical image segmentation model based on multi-axis attention of the present invention; the total number of iterations epoch ranges from 100 to 300, and the learning rate η ranges from 0 to 1.
[0032] (2-3) For each medical image in the training set obtained in step (2-1), the medical image is input into the main feature retention module to obtain a feature map corresponding to the medical image with a size of 256×256 and a channel number of 64.
[0033] (2-4) For each medical image in the training set obtained in step (2-1), the feature map corresponding to the medical image obtained in step (2-3) is input into the encoder to obtain an intermediate feature map corresponding to the medical image.
[0034] (2-5) For each medical image in the training set obtained in step (2-1), the feature map corresponding to the medical image obtained in step (2-3) is input into the multi-axis attention module to obtain the multi-axis attention feature map corresponding to the medical image.
[0035] (2-6) For each medical image in the training set obtained in step (2-1), the intermediate feature map corresponding to the medical image obtained in step (2-4) and the multi-axis attention feature map corresponding to the medical image obtained in step (2-5) are input into the decoder to obtain the final lesion segmentation prediction result F corresponding to the medical image. pred ;
[0036] (2-7) For each medical image in the training set obtained in step (2-1), the lesion segmentation prediction result F corresponding to the medical image obtained in step (2-6) is pred Calculate the loss function and use the loss function to iteratively train the multi-axis attention medical image segmentation model to obtain the preliminarily trained multi-axis attention medical image segmentation model W k The preliminarily trained medical image segmentation model is verified using the test set obtained in step (2-1) to obtain the segmentation accuracy of the test set.
[0037] (2-8) Repeat the above steps (2-3) to (2-7) until the total number of iterations set in step (2-2) is reached, or the segmentation accuracy on the test set reaches the optimal value, thereby obtaining a trained medical image segmentation model based on multi-axis attention.
[0038] Preferably, step (2-3) is specifically as follows: first, the main feature retention module uses a 7×7 convolution kernel to perform convolution processing on the preprocessed medical image, and the convolution kernel has in_ch input channels and out_ch output channels. Then, the convolved medical image is corrected using the ReLU function to obtain corrected features, and then the corrected features are batch normalized to obtain normalized features. Then, another convolution kernel of size 7×7 is used to perform convolution processing on the normalized features, and then, The convolutional features are corrected using the ReLU function to obtain the corrected features, and then the corrected features are batch normalized to obtain the normalized intermediate feature map. Thereafter, a convolution kernel of size 1×1 and no padding is used to convolve the normalized intermediate feature map, and then the convolutional features are corrected using the ReLU function to obtain the corrected features. Thereafter, the corrected features are batch normalized to obtain a normalized feature map of size 256×256 and number of channels 64.
[0039] Specifically, step (2-4) is as follows: the encoder first inputs the feature map output by the main feature retention module into the maximum pooling layer to obtain a feature map with a size of 128×128 and a number of channels of 64. Next, the encoder inputs the feature map into the first compact separation attention module to obtain a feature map with a size of 128×128 and a number of channels of 128; then, the feature map with a size of 128×128 and a number of channels of 128 obtained by the first compact separation attention module is input into the maximum pooling layer, so that the height and width of the feature map are halved, and the number of channels remains unchanged, thereby achieving spatial downsampling, thereby obtaining a feature map with a size of 64×64 and a number of channels of 128 after downsampling; then, the feature map with a size of 64×64 and a number of channels of 128 after downsampling is input into the second compact separation attention module to obtain an output size of 64× 64, a feature map with 256 channels, and then input the feature map into the maximum pooling layer to obtain a downsampled feature map with a size of 32×32 and a channel number of 256; then, the downsampled 256-channel feature map is input into the third compact separation attention module to obtain an output feature map with an output size of 32×32 and a channel number of 512; subsequently, the feature map is input into the maximum pooling layer to obtain a downsampled feature map with a size of 16×16 and a channel number of 512; finally, the 512-channel feature map is input into the fourth compact separation attention module to obtain an intermediate feature map with a size of 16×16 and a channel number of 512.
[0040] Preferably, step (2-5) is specifically as follows: first, the channel axis attention module uses average pooling and maximum pooling on the channel dimension to obtain two channel dimension information, and then, the two channel dimension information are respectively input into the double convolution module to obtain two channel attention maps with a size of 1×1 and a channel number of 64. Thereafter, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F c Then, F c Perform point-by-point multiplication with the feature map of size 256×256 and 64 channels output by the main feature retention module to obtain the feature map F′ c ; Then, the obtained feature map F′ c Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 256×256 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. The sigmoid activation function is then used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map. Then, the spatial attention map With the feature map F″ c Perform multiplication operation to obtain the final output feature map F c "as follows:
[0041]
[0042] Then, the feature map F″ is output c The height axis attention module and the width axis attention module are respectively input to extract the height axis feature map and the width axis feature map to obtain the feature map F″ h and F″ w ;
[0043] Finally, the multi-axis attention module converts the final height axis feature map F″ h And the final width axis feature map F″ w Superimpose them to obtain the multi-axis attention feature map F:
[0044] F=F″ h +F″ w
[0045] The output feature map F″ c The height axis attention module and the width axis attention module are respectively input to extract the height axis feature map and the width axis feature map to obtain the feature map F″ h and F″ w The specific process is as follows:
[0046] First, the height axis attention module takes the feature map F″ output by the channel axis attention module c The height axis and the channel axis are swapped so that the height axis becomes the new channel axis, and the original channel axis is converted to the height axis to obtain a height axis feature map with a size of 64×256 and a channel number of 256. Then, the height axis attention module uses the average pooling and maximum pooling on the channel dimension to obtain two channel dimension information. Then, the two channel dimension information is input into the double convolution module respectively to obtain two channel attention maps with a size of 1×1 and a channel number of 256. After that, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F. hThen, F h Perform point-by-point multiplication with the height axis feature map of size 64×256 and number of channels 256 to obtain the feature map F′ h ; Then, the obtained feature map F′ h Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 64×256 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. The sigmoid activation function is then used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map. Then, the spatial attention map With the feature map F′ h Perform multiplication operation to obtain the feature map F″ h as follows:
[0047]
[0048] Then, the height axis attention module swaps the height axis and channel axis of the obtained height axis feature map again to obtain the final height axis feature map F″ with a size of 256×256 and 64 channels h .
[0049] At the same time, the width axis attention module first converts the feature map F″ output by the channel axis attention module c The width axis and the channel axis are swapped so that the width axis becomes the new channel axis, and the original channel axis is converted to the width axis to obtain a width axis feature map with a size of 256×64 and a channel number of 256; then the width axis attention module uses the average pooling and maximum pooling on the channel dimension to obtain two channel dimension information, and then the two channel dimension information is respectively input into the double convolution module to obtain two channel attention maps with a size of 1×1 and a channel number of 256. Afterwards, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F w Then, F w Perform point-by-point multiplication with the width axis feature map of size 256×64 and number of channels 256 to obtain the feature map F′ w ; Then, the obtained feature map F′ w Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 256×64 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. The sigmoid activation function is then used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map. Then, the spatial attention map With the feature map F′ w Perform multiplication operation to obtain the feature map F″ w as follows:
[0050]
[0051] Then, the width axis attention module swaps the width axis and channel axis of the obtained width axis feature map again to obtain the final width axis feature map F″ with a size of 256×256 and 64 channels w .
[0052] Preferably, step (2-6) is specifically as follows: the decoder first performs an upsampling operation with a step size of 2*2 and a stride of 2 on the intermediate feature map of 16×16 and 512 channels output by the fourth compact separation attention module to obtain a feature map of size 32×32 and 512 channels, and then splices the feature map with the feature map of size 32×32 and 512 channels output by the third compact separation attention module in the channel dimension to obtain a spliced feature map of size 32×32 and 1024 channels; then, the spliced feature map is input into the fifth compact separation attention module to obtain a feature map of size 32×32 and 256 channels; then, the feature map is step-wise multiplied. An upsampling operation with a length of 2*2 and a stride of 2 is performed to obtain a feature map with a size of 64×64 and a number of channels of 256. Subsequently, the feature map is concatenated in the channel dimension with the feature map with a size of 64×64 and a number of channels of 256 output by the second compact separation attention module to obtain a feature map with a size of 64×64 and a number of channels of 512; then, the feature map is input into the sixth compact separation attention module to obtain a feature map with a size of 64×64 and a number of channels of 128. Subsequently, an upsampling operation with a step size of 2*2 and a stride of 2 is performed on the feature map, and the obtained feature map is concatenated with the feature map with a size of 128×128 and a number of channels of 128 output by the first compact separation attention module. The feature map is spliced in the channel dimension to obtain a spliced feature map with a size of 128×128 and a number of channels of 256. Then, the spliced feature map is input into the seventh compact separation attention module to obtain a feature map with a size of 128×128 and a number of channels of 64. Subsequently, the feature map is upsampled with a step size of 2*2 and a stride of 2 to obtain a feature map with a size of 256×256 and a dimension of 64. Subsequently, the feature map is spliced with the multi-axis attention feature map output by the multi-axis attention module into a feature map with a size of 256×256 and a number of channels of 128. Subsequently, the feature map is input into the eighth compact separation attention module to obtain a feature map with a size of 256×256 and a number of channels of 6. 4 feature map; then, a 1×1 convolution layer is used to reduce the number of channels of the feature map to two channels, and then the probability distribution matrix of the tumor location is obtained; then, the probability distribution matrix is mapped to the [0,1] interval through the Softmax function; finally, on the probability distribution matrix of the first channel, the element value with a probability value greater than 0.5 in the probability distribution matrix is set to 1, that is, the background (that is, the position is considered to be the background, that is, there is no lesion), and the other element values are set to 0; on the second channel, the element value with a probability value greater than 0.5 in the probability distribution matrix is set to 1, that is, the foreground, and the other element values are set to 0; finally, the binarized probability distribution matrix is saved as a binarized image, which is the final lesion segmentation prediction result;
[0053] The loss function in step (2-7) is a comprehensive function, which is composed of the sum of the Focal loss function and the Dice loss function. The specific calculation formula is:
[0054] L=λ 1 ·L focal +λ 2 ·L dice
[0055]
[0056]
[0057] Among them, λ 1 and λ 2 are the weighted values of the two losses, both set to 0.5, α and γ are the built-in weights of Focal loss, set to 0.25 and 2 respectively, ∈ is a constant, whose value is 1*10 -8 , used to prevent the denominator from being 0, represents the true result of lesion segmentation of the i-th medical image in the training set obtained in step (2-1), represents the lesion segmentation prediction result of the i-th medical image in the training set obtained in step (2-1), i∈[1,N], and N represents the total number of medical images in the training set obtained in step (2-1).
[0058] According to another aspect of the present invention, a medical image segmentation system based on multi-axis attention is provided, comprising:
[0059] The first module is used to acquire a medical image and preprocess the medical image to acquire a preprocessed medical image;
[0060] The second module is used to input the medical image preprocessed by the first module into a pre-trained multi-axis attention-based medical image segmentation model to obtain the final lesion segmentation prediction result.
[0061] In general, the above technical solutions conceived by the present invention can achieve the following beneficial effects compared with the prior art:
[0062] 1. Since the present invention adopts steps (2-1) to (2-6), it can effectively improve the utilization of image information by retaining the main features of medical images. Therefore, it can solve the technical problem of over-segmentation or information loss of certain complex structures and small targets in the existing U-Net model.
[0063] 2. Since the present invention adopts steps (2-1) to (2-6), it can effectively make up for the information difference between low-level features and high-level features by optimizing the low-level features in the jump connection. Therefore, it can solve the technical problems such as information redundancy and overfitting caused by the existing U-Net model directly splicing low-level features with high-level features due to the jump connection.
[0064] 3. Since the present invention adopts steps (2-1) to (2-6), it improves the original CBAM attention module through a multi-axis attention module, which only considers features from the channel axis and ignores the potential feature information in the height axis and width axis directions. The multi-axis attention module proposed in the present invention increases the expression ability of the attention module by adding attention modules in the height axis and width axis directions. Therefore, it can solve the technical problem that the originally introduced CBAM attention mechanism has limited expression ability of the model.
[0065] 4. Since the present invention adopts steps (2-1) to (2-6), it proposes to use compact separation attention to replace the feature extraction module in the traditional U-Net encoder to enhance the feature extraction capability of U-Net, and uses compact separation attention in the decoder to enhance the feature representation capability of U-Net to comprehensively improve the model performance. Therefore, it can solve the problem of insufficient performance of the original U-Net model.
[0066] 5. The present invention adopts steps (2-1) to (2-6), which enables the model to extract more useful spatial information by paying more attention to important channels and ignoring channels that are not very beneficial to the segmentation task, and more spatial information helps the model to extract more semantic information through subsequent feature extraction modules. Ultimately, the multi-axis attention mechanism can improve the overall performance of the model on other tasks with less computational overhead.
[0067] 6. The medical image segmentation method based on multi-axis attention, which is suitable for small-scale medical images, proposed in the present invention can effectively improve the performance and accuracy of the model when the data set is limited, and provide reliable and efficient tools and methods for data analysis and diagnosis in the medical field. BRIEF DESCRIPTION OF THE DRAWINGS
[0068] Figure 1 is an overall flow chart of the medical image segmentation method based on multi-axis attention of the present invention;
[0069] Figure 2 It is a structural schematic diagram of the medical image segmentation model based on multi-axis attention of the present invention. DETAILED DESCRIPTION
[0070] In order to make the purpose, technical solutions and advantages of the present invention more clearly understood, the present invention is further described in detail below in conjunction with the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain 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 can be combined with each other as long as they do not conflict with each other.
[0071] The basic idea of the present invention is to propose a new deep learning U-Net medical image segmentation method by introducing a main feature retention module, a compact separation attention module and a multi-axis attention module. The main process includes: using the main feature retention module and the compact separation attention module to improve the model model computing efficiency, integrating the proposed multi-axis attention mechanism into the jump connection structure of U-Net to reduce the information redundancy existing in the direct fusion of low-level features and high-level features, and improving information utilization; compared with the traditional feature extraction module, the main feature retention module and the compact separation attention module require fewer learning parameters; compared with the traditional attention mechanism, the improved attention module can effectively improve the model expression ability. In short, the improved model is smaller, has higher performance, and is easier to deploy.
[0072] The purpose of this invention is to propose a new type of attention mechanism, called the multi-axis attention mechanism, and then propose a higher performance deep learning segmentation method for small-scale medical image datasets based on the proposed attention mechanism.
[0073] like Figure 1 As shown, the present invention provides a medical image segmentation method based on multi-axis attention, comprising the following steps:
[0074] (1) acquiring a medical image, and preprocessing the medical image to obtain a preprocessed medical image;
[0075] Specifically, the medical image in this step is a two-dimensional static image, which can be a tumor image obtained by an ultrasound device or a dermatoscopic image.
[0076] Since the acquired medical images are large in size and contain a lot of irrelevant information related to imaging equipment, patient privacy, etc., this step first extracts the region of interest from the medical image according to the specific location of the lesion in the pre-set segmentation label, and then converts the extracted region of interest into a medical image of 256×256 size and jpg format as the preprocessed medical image.
[0077] (2) The medical image preprocessed in step (1) is input into a pre-trained multi-axis attention-based medical image segmentation model to obtain the final lesion segmentation prediction result.
[0078] Specifically, the final lesion segmentation result is a binary image (only black and white) corresponding to the medical image obtained in step (1), where black is the background and white is the foreground (i.e., the lesion).
[0079] The medical image segmentation model based on multi-axis attention of the present invention is a U-Net network model optimized in many aspects, such as Figure 2 As shown, it includes a main feature retention module, an encoder, a decoder, and a multi-axis attention module connected between the main feature retention module and the decoder.
[0080] The input of the main feature retention module is the preprocessed medical image, and the output is a feature map with a size of 256×256 and 64 channels.
[0081] Specifically, the main feature retention module first uses a 7×7 convolution kernel to convolve the preprocessed medical image, which has in_ch (usually 3) input channels and out_ch (here 64) output channels. Then, the convolved medical image is corrected using the ReLU function to obtain the corrected features. Then, the corrected features are batch normalized to obtain the normalized features. Then, another 7×7 convolution kernel is used to convolve the normalized features. Then , the convolutional features are corrected using the ReLU function to obtain the corrected features, and then the corrected features are batch normalized to obtain the normalized intermediate feature map, and then a convolution kernel with a size of 1×1 and no padding is used to convolve the normalized intermediate feature map, and then the convolutional features are corrected using the ReLU function to obtain the corrected features, and then the corrected features are batch normalized to obtain a normalized feature map with a size of 256×256 and 64 channels.
[0082] The encoder contains four compact separation attention modules. The input of the encoder is the feature map of size 256×256 and 64 channels output by the main feature retention module, and its output is the intermediate feature map.
[0083] Specifically, the encoder first inputs the feature map output by the main feature retention module into the maximum pooling layer (pooling window size is 2×2, stride is 2) to obtain a feature map with a size of 128×128 and 64 channels. Next, the encoder inputs the feature map into the first compact separation attention module to obtain a feature map with a size of 128×128 and 128 channels (specifically, the feature map is first convolved with a convolution kernel with a size of 1×1 and 32 channels, and the convolution result is calculated using two sets of convolution kernels with a size of 3×3 and 16 channels to obtain the first feature map and the second feature map. Then, the first feature map and the second feature map are added to obtain the superimposed feature map, and then the superimposed feature map is processed using a convolution kernel with a size of 3×3 and 16 channels to obtain the optimized third feature map, and then the window size and the third feature map are used. Figure 1The third feature map is processed by a pooling window of the same size to obtain a channel feature map with a size of 1×1 and 16 channels. Then, a convolution kernel with a size of 1×1 and 32 channels is used to convolve the channel feature map to obtain a feature map with a size of 1×1 and 32 channels. Subsequently, the feature map is batch normalized, and then the normalized feature map is corrected using the ReLU function to obtain a corrected feature map. Then, the corrected feature map is convolved again with a convolution kernel of size 1×1 and number of channels 32, and then, it is corrected using the Softmax activation function to obtain the fourth feature map. Then, the first and third feature maps are respectively combined with the fourth feature map by dot multiplication to obtain two intermediate feature maps. Then, the two intermediate feature maps are obtained by channel addition. The intermediate feature map is then processed with a convolution kernel of size 1×1 and number of channels 128 to obtain the fifth feature map. Finally, the fifth feature map is added to the feature map output by the main feature retention module to obtain a feature map of size 128×128 and number of channels 128); then, the feature map of size 128×128 and number of channels 128 obtained by the first compact separation attention module is input into the maximum pooling layer (pooling window size is 2×2, stride is 2) . ), so that the height and width of the feature map are halved, the number of channels remains unchanged, and spatial downsampling is achieved, so as to obtain a feature map with a size of 64×64 and a number of channels of 128 after downsampling; then, the feature map with a size of 64×64 and a number of channels of 128 after downsampling is input into the second compact separation attention module to obtain a feature map with an output size of 64×64 and a number of channels of 256 (specifically, a convolution kernel with a size of 1×1 and a number of channels of 64 is first used to convolve the feature map, and two sets of convolution kernels with a size of 3×3 and a number of channels of 32 are used to calculate the convolution processing results to obtain the first feature map and the second feature map. Then, the first feature map and the second feature map are added to obtain a superimposed feature map, and then, the superimposed feature map is processed using a convolution kernel with a size of 3×3 and a number of channels of 32 to obtain an optimized third feature map, and then, the window size and the third feature map are used to optimize the feature map. Figure 1The third feature map is processed by a pooling window of the same size to obtain a channel feature map with a size of 1×1 and 32 channels. Then, a convolution kernel with a size of 1×1 and 64 channels is used to convolve the channel feature map to obtain a feature map with a size of 1×1 and 64 channels. Subsequently, the feature map is batch normalized, and then the normalized feature map is corrected using the ReLU function to obtain a corrected feature map. Then, the corrected feature map is convolved again using a convolution kernel with a size of 1×1 and a number of channels of 64, and then modified using the Softmax activation function to obtain a fourth feature map. Then, the first and third feature maps are respectively combined with the fourth feature map by dot multiplication to obtain two intermediate feature maps. Then, the two intermediate feature maps are obtained by channel addition to obtain an intermediate feature map. Then, the intermediate feature map is processed using a convolution kernel with a size of 1×1 and a number of channels of 256 to obtain a fifth feature map. Finally, the fifth feature map is added to the feature map output by the main feature retention module to obtain a feature map with a size of 64×64 and a number of channels of 256), and then the feature map is input into the maximum pooling layer (to further reduce the size of the feature map) to proceed. The downsampled feature map is 32×32 in size and 256 in number of channels. Then, the downsampled 256-channel feature map is input into the third compact separation attention module to obtain a feature map with an output size of 32×32 and a number of channels of 512 (specifically, a convolution kernel with a size of 1×1 and a number of channels of 128 is first used to convolve the feature map, and two sets of convolution kernels with a size of 3×3 and a number of channels of 64 are used to calculate the convolution processing results to obtain the first feature map and the second feature map. Then, the first feature map and the second feature map are added to obtain the superimposed feature map, and then the superimposed feature map is processed using a convolution kernel with a size of 3×3 and a number of channels of 64 to obtain the optimized third feature map, and then the window size and the third feature map are used to optimize the feature map. Figure 1The third feature map is processed with a pooling window of the same size to obtain a channel feature map with a size of 1×1 and 64 channels. Then, the channel feature map is convolved with a convolution kernel with a size of 1×1 and a number of channels of 128 to obtain a feature map with a size of 1×1 and a number of channels of 128. Then, the feature map is batch normalized. After that, the normalized feature map is corrected using the ReLU function to obtain a corrected feature map. Then, the corrected feature map is convolved again with a convolution kernel with a size of 1×1 and a number of channels of 128. Then, the Softmax activation function is used to correct it to obtain a fourth feature map. Then, the first and third feature maps are respectively combined with the fourth feature map by dot multiplication to obtain two intermediate feature maps. Then, the two intermediate feature maps are obtained by channel addition. Then, the intermediate feature map is processed with a convolution kernel with a size of 1×1 and a number of channels of 512 to obtain a fifth feature map. Finally, the fifth feature map and the main feature map are combined. The feature maps output by the retention module are added to obtain a feature map with unchanged size and 512 channels); then, the feature map is input into the maximum pooling layer (the purpose is to halve the size of the feature map) to obtain a downsampled feature map with a size of 16×16 and 512 channels; finally, the feature map of 512 channels is input into the fourth compact separation attention module to obtain an intermediate feature map with a size of 16×16 and 512 channels (the specific processing process of the fourth compact separation attention module is exactly the same as that of the third compact separation attention module, which will not be repeated here).
[0084] The multi-axis attention module includes a channel-axis attention module, a height-axis attention module, and a width-axis attention module. The input of the multi-axis attention module is a feature map with a size of 256×256 and 64 channels output by the main feature retention module, and the output is a multi-axis attention feature map with the same size as the input feature map.
[0085] Specifically, the channel axis attention module first uses the average pooling and maximum pooling on the channel dimension to obtain two channel dimension information, and then inputs the two channel dimension information into the double convolution module respectively (first reduce the number of input channels through a 1×1 convolution layer, then, through another 1×1 convolution layer, restore the number of channels to the number of input feature channels, and finally map the value output by the convolution layer to between 0 and 1 through the Sigmoid activation function) to obtain two channel attention maps with a size of 1×1 and a channel number of 64. Afterwards, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F c Then, F c Perform point-by-point multiplication with the feature map of size 256×256 and 64 channels output by the main feature retention module to obtain the feature map F′c ; Then, the obtained feature map F′ c Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 256×256 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. The sigmoid activation function is then used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map. Then, the spatial attention map With the feature map F″ c Perform multiplication operation to obtain the final output feature map F c ″ is as follows (where Indicates that the matrix elements correspond to the product operation):
[0086]
[0087] Then, the feature map F″ is output c The height axis attention module and the width axis attention module are respectively input to extract the height axis feature map and the width axis feature map to obtain the feature map F″ h and F″ w ;
[0088] Specifically, the height axis attention module first converts the feature map F″ output by the channel axis attention module c The height axis and the channel axis are swapped so that the height axis becomes the new channel axis, and the original channel axis is transformed into the height axis to obtain a height axis feature map with a size of 64×256 and a channel number of 256; then the height axis attention module uses the average pooling and maximum pooling on the channel dimension to obtain two channel dimension information, and then the two channel dimension information are respectively input into the double convolution module (first through a 1×1 convolution layer to reduce the number of input channels, then through another 1×1 convolution layer, the number of channels is restored to the number of input feature channels, and finally the Sigmoid activation function is used to map the value output by the convolution layer to between 0 and 1) to obtain two channel attention maps with a size of 1×1 and a channel number of 256. Afterwards, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F h Then, F h Perform point-by-point multiplication with the height axis feature map of size 64×256 and number of channels 256 to obtain the feature map F′ h; Then, the obtained feature map F′ h Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 64×256 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. The sigmoid activation function is then used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map. Then, the spatial attention map With the feature map F′ h Perform multiplication operation to obtain the feature map F″ h As follows (where Indicates that the matrix elements correspond to the product operation):
[0089]
[0090] Then, the height axis attention module swaps the height axis and channel axis of the obtained height axis feature map again to obtain the final height axis feature map F″ with a size of 256×256 and 64 channels h .
[0091] At the same time, the width axis attention module first converts the feature map F″ output by the channel axis attention module c The width axis and the channel axis are swapped so that the width axis becomes the new channel axis, and the original channel axis is transformed into the width axis to obtain a width axis feature map with a size of 256×64 and a channel number of 256; then the width axis attention module uses the average pooling and maximum pooling on the channel dimension to obtain two channel dimension information, and then the two channel dimension information are respectively input into the double convolution module (first through a 1×1 convolution layer to reduce the number of input channels, then through another 1×1 convolution layer, the number of channels is restored to the number of input feature channels, and finally the Sigmoid activation function is used to map the value output by the convolution layer to between 0 and 1) to obtain two channel attention maps with a size of 1×1 and a channel number of 256. Afterwards, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F w Then, F w Perform point-by-point multiplication with the width axis feature map of size 256×64 and number of channels 256 to obtain the feature map F′ w ; Then, the obtained feature map F′ wPerform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 256×64 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. The sigmoid activation function is then used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map. Then, the spatial attention map With the feature map F′ w Perform multiplication operation to obtain the feature map F″ w As follows (where Indicates that the matrix elements correspond to the product operation):
[0092]
[0093] Then, the width axis attention module swaps the width axis and channel axis of the obtained width axis feature map again to obtain the final width axis feature map F″ with a size of 256×256 and 64 channels w (The purpose is to make the final height axis feature map and width axis feature map consistent with the feature map F" output by the channel axis attention module c of the same size).
[0094] Finally, the multi-axis attention module converts the final height axis feature map F″ h And the final width axis feature map F″ w Superimpose them to obtain the multi-axis attention feature map F:
[0095] F=F″ h +F″ w
[0096] The decoder contains four compact separation attention modules. The input of the decoder is the feature map with a size of 32×32 and 512 channels output by the third compact separation attention module in the encoder, the intermediate feature map with a size of 16×16 and 512 channels output by the fourth compact separation attention module, and the multi-axis attention feature map output by the multi-axis attention module. The output is the final lesion segmentation prediction result.
[0097] Specifically, the decoder first performs an upsampling operation on the 16×16 intermediate feature map with a channel number of 512 output by the fourth compact separation attention module with a step size of 2*2 and a stride of 2 to obtain a feature map with a size of 32×32 and a channel number of 512. Then, the feature map is spliced in the channel dimension with the feature map with a size of 32×32 and a channel number of 512 output by the third compact separation attention module to obtain a spliced feature map with a size of 32×32 and a channel number of 1024; then, the spliced feature map is input into the fifth compact separation attention module to obtain a size of A feature map of 32×32 with 256 channels (specifically, a convolution kernel of 1×1 and 64 channels is first used to convolve the feature map, and two sets of convolution kernels of 3×3 and 32 channels are used to calculate the convolution results to obtain the first feature map and the second feature map. Then, the first feature map and the second feature map are added to obtain the superimposed feature map, and then the superimposed feature map is processed using a convolution kernel of 3×3 and 32 channels to obtain the optimized third feature map, and then the window size and the third feature map are used to optimize the feature map. Figure 1The third feature map is processed by a pooling window of the same size to obtain a channel feature map with a size of 1×1 and 32 channels. Then, a convolution kernel with a size of 1×1 and 64 channels is used to convolve the channel feature map to obtain a feature map with a size of 1×1 and 64 channels. Subsequently, the feature map is batch normalized, and then the normalized feature map is corrected using the ReLU function to obtain a corrected feature map. Then, the corrected feature map is convolved again with a convolution kernel of size 1×1 and channel number 64, and then, it is corrected using the Softmax activation function to obtain the fourth feature map. Then, the first and third feature maps are respectively combined with the fourth feature map by dot multiplication to obtain two intermediate feature maps. Then, the two intermediate feature maps are obtained by channel addition. Then, the intermediate feature map is processed with a convolution kernel of size 1×1 and channel number 256 to obtain the fifth feature map. Then, the input feature map is downsampled using a convolution kernel of size 1×1 and channel number 256 to obtain a downsampled feature with a size of 32×32 and channel number 256. Finally, the fifth feature map is added to the downsampled feature to obtain the final output); then, the feature map is upsampled with a step size of 2*2 and a stride of 2 to obtain a feature map with a size of 64×64 and a channel number of 64×64. A feature map with a channel number of 256 is obtained, and then the feature map is spliced with the feature map with a size of 64×64 and a channel number of 256 output by the second compact separation attention module in the channel dimension to obtain a feature map with a size of 64×64 and a channel number of 512; then, the feature map is input into the sixth compact separation attention module to obtain a feature map with a size of 64×64 and a channel number of 128 (specifically, a convolution kernel with a size of 1×1 and a channel number of 32 is first used to convolve the feature map, and two sets of convolution kernels with a size of 3×3 and a channel number of 16 are used to calculate the convolution processing result to obtain the first feature map and the second feature map. Then, the first feature map and the second feature map are added to obtain a superimposed feature map, and then the superimposed feature map is processed using a convolution kernel with a size of 3×3 and a channel number of 16 to obtain an optimized third feature map, and then, the window size and the third feature map are used. Figure 1The third feature map is processed by a pooling window of the same size to obtain a channel feature map with a size of 1×1 and 16 channels. Then, a convolution kernel with a size of 1×1 and 32 channels is used to convolve the channel feature map to obtain a feature map with a size of 1×1 and 32 channels. Subsequently, the feature map is batch normalized, and then the normalized feature map is corrected using the ReLU function to obtain a corrected feature map. Then, the corrected feature map is convolved again with a convolution kernel of size 1×1 and channel number 32, and then, it is corrected using the Softmax activation function to obtain the fourth feature map. Then, the first and third feature maps are respectively combined with the fourth feature map by dot multiplication to obtain two intermediate feature maps. Then, the two intermediate feature maps are obtained by channel addition. Then, the intermediate feature map is processed using a convolution kernel of size 1×1 and channel number 128 to obtain the fifth feature map. Then, the input feature map is downsampled using a convolution kernel of 1×1 and channel number 128 to obtain a downsampled feature of size 64×64 and channel number 128. Finally, the fifth feature map is added to the downsampled feature to obtain the final output), then, the feature map is upsampled with a step size of 2*2 and a stride of 2, and the obtained The feature map is concatenated with the feature map of size 128×128 and number of channels 128 output by the first compact separation attention module in the channel dimension to obtain a concatenated feature map of size 128×128 and number of channels 256, and then the concatenated feature map is input into the seventh compact separation attention module to obtain a feature map of size 128×128 and number of channels 64; (specifically, a convolution kernel of size 1×1 and number of channels 16 is first used to convolve the feature map, and two sets of convolution kernels of size 3×3 and number of channels 8 are used to calculate the convolution processing result to obtain the first feature map and the second feature map. Then, the first feature map and the second feature map are added to obtain the superimposed feature map, and then the superimposed feature map is processed using a convolution kernel of size 3×3 and number of channels 8 to obtain the optimized third feature map, and then the window size and the third feature map are used to optimize the feature map. Figure 1The third feature map is processed with a pooling window of the same size to obtain a channel feature map with a size of 1×1 and 8 channels. Then, the channel feature map is convolved with a convolution kernel with a size of 1×1 and 16 channels to obtain a feature map with a size of 1×1 and 16 channels. Then, the feature map is batch normalized. After that, the normalized feature map is corrected using the ReLU function to obtain a corrected feature map. Then, the corrected feature map is convolved again with a convolution kernel with a size of 1×1 and 16 channels. Then, the Softmax activation function is used to correct it to obtain a fourth feature map. Then, the first and third feature maps are respectively combined with the fourth feature map by dot multiplication to obtain two intermediate feature maps. Then, the two intermediate feature maps are obtained by channel addition. Then, the intermediate feature map is processed with a convolution kernel with a size of 1×1 and 64 channels to obtain the fifth feature map. Then, the convolution kernel with a size of 1×1 and a channel number of 64 is used to downsample the input feature map to obtain a downsampled feature with a size of 128×128 and a channel number of 64. Finally, the fifth feature map is added to the downsampled feature to obtain the final output). Subsequently, the feature map is upsampled with a step size of 2*2 and a stride of 2 to obtain a feature map with a size of 256×256 and a dimension of 64. Subsequently, the feature map is concatenated with the multi-axis attention feature map output by the multi-axis attention module to obtain a feature map with a size of 2 56×256 feature map with 128 channels; then, the feature map is input into the eighth compact separation attention module to obtain a feature map with a size of 256×256 and a number of channels of 64; (the specific processing process of the eighth compact separation attention module is exactly the same as that of the seventh compact separation attention module, which will not be repeated here) after that, the number of channels of the feature map is reduced to two channels through a 1×1 convolution layer, and then the probability distribution matrix of the tumor position is obtained; then, the probability distribution matrix is mapped to the [0,1] interval through the Softmax function; finally, on the probability distribution matrix of the first channel, the element value with a probability value greater than 0.5 in the probability distribution matrix is set to 1, that is, the background (that is, the position is considered to be the background, that is, there is no lesion), and the other element values are set to 0; on the second channel, the element value with a probability value greater than 0.5 in the probability distribution matrix is set to 1, that is, the foreground (that is, the position is considered to be the foreground, that is, the lesion), and the other element values are set to 0; finally, the binarized probability distribution matrix is saved as a binarized image, which is the final lesion segmentation prediction result.
[0098] Specifically, the multi-axis attention-based medical image segmentation model of the present invention is obtained by the following steps:
[0099] (2-1) Obtain a medical image dataset, preprocess the medical image dataset to obtain a preprocessed medical image dataset, and divide the preprocessed medical image dataset into a training set and a test set in a ratio of 8:2.
[0100] The process of dividing the training set and the test set in this step is as follows: first, all the preprocessed medical images are randomly shuffled, and then all the shuffled medical images are evenly divided into 5 parts, and any one of the 5 parts of data is taken as the test set, and the remaining 4 parts are taken as the training sets.
[0101] (2-2) Setting the total number of iterations epoch of the training process and initializing the learning rate η of the medical image segmentation model based on multi-axis attention of the present invention;
[0102] Specifically, the total number of iterations epoch in this step ranges from 100 to 300, preferably 200; the learning rate η ranges from 0 to 1, preferably 0.001.
[0103] (2-3) For each medical image in the training set obtained in step (2-1), the medical image is input into the main feature retention module to obtain a feature map corresponding to the medical image with a size of 256×256 and a channel number of 64.
[0104] Specifically, the main feature retention module first uses a 7×7 convolution kernel to convolve the preprocessed medical image. The convolution kernel has in_ch (usually 3) input channels and out_ch (here
[0105] 64) output channels, then, the convolved medical image is corrected using the ReLU function to obtain the corrected features, then, the corrected features are batch normalized to obtain the normalized features, then, another 7×7 convolution kernel is used to convolve the normalized features, then, the convolved features are corrected using the ReLU function to obtain the corrected features, then, the corrected features are batch normalized to obtain the normalized intermediate feature map, then, a 1×1 convolution kernel with no padding is used to convolve the normalized intermediate feature map, then, the convolved features are corrected using the ReLU function to obtain the corrected features, then, the corrected features are batch normalized to obtain a normalized feature map with a size of 256×256 and 64 channels.
[0106] The advantage of the above steps (2-3) is that the main feature retention module can retain the main features of each medical image and can effectively improve the utilization of image information. Therefore, it can solve the technical problem of over-segmentation or information loss of the existing U-Net model for some complex structures and small targets.
[0107] (2-4) For each medical image in the training set obtained in step (2-1), the feature map corresponding to the medical image obtained in step (2-3) is input into the encoder to obtain an intermediate feature map corresponding to the medical image.
[0108] Specifically, the encoder first inputs the feature map output by the main feature retention module into the maximum pooling layer (pooling window size is 2×2, stride is 2) to obtain a feature map with a size of 128×128 and 64 channels. Next, the encoder inputs the feature map into the first compact separation attention module to obtain a feature map with a size of 128×128 and 128 channels (specifically, the feature map is first convolved with a convolution kernel with a size of 1×1 and 32 channels, and the convolution result is calculated using two sets of convolution kernels with a size of 3×3 and 16 channels to obtain the first feature map and the second feature map. Then, the first feature map and the second feature map are added to obtain the superimposed feature map, and then the superimposed feature map is processed using a convolution kernel with a size of 3×3 and 16 channels to obtain the optimized third feature map, and then the window size and the third feature map are used. Figure 1The third feature map is processed by a pooling window of the same size to obtain a channel feature map with a size of 1×1 and 16 channels. Then, a convolution kernel with a size of 1×1 and 32 channels is used to convolve the channel feature map to obtain a feature map with a size of 1×1 and 32 channels. Subsequently, the feature map is batch normalized, and then the normalized feature map is corrected using the ReLU function to obtain a corrected feature map. Then, the corrected feature map is convolved again with a convolution kernel of size 1×1 and number of channels 32, and then, it is corrected using the Softmax activation function to obtain the fourth feature map. Then, the first and third feature maps are respectively combined with the fourth feature map by dot multiplication to obtain two intermediate feature maps. Then, the two intermediate feature maps are obtained by channel addition. The intermediate feature map is then processed with a convolution kernel of size 1×1 and number of channels 128 to obtain the fifth feature map. Finally, the fifth feature map is added to the feature map output by the main feature retention module to obtain a feature map of size 128×128 and number of channels 128); then, the feature map of size 128×128 and number of channels 128 obtained by the first compact separation attention module is input into the maximum pooling layer (pooling window size is 2×2, stride is 2) . ), so that the height and width of the feature map are halved, the number of channels remains unchanged, and spatial downsampling is achieved, so as to obtain a feature map with a size of 64×64 and a number of channels of 128 after downsampling; then, the feature map with a size of 64×64 and a number of channels of 128 after downsampling is input into the second compact separation attention module to obtain a feature map with an output size of 64×64 and a number of channels of 256 (specifically, a convolution kernel with a size of 1×1 and a number of channels of 64 is first used to convolve the feature map, and two sets of convolution kernels with a size of 3×3 and a number of channels of 32 are used to calculate the convolution processing results to obtain the first feature map and the second feature map. Then, the first feature map and the second feature map are added to obtain a superimposed feature map, and then, the superimposed feature map is processed using a convolution kernel with a size of 3×3 and a number of channels of 32 to obtain an optimized third feature map, and then, the window size and the third feature map are used to optimize the feature map. Figure 1The third feature map is processed by a pooling window of the same size to obtain a channel feature map with a size of 1×1 and 32 channels. Then, a convolution kernel with a size of 1×1 and 64 channels is used to convolve the channel feature map to obtain a feature map with a size of 1×1 and 64 channels. Subsequently, the feature map is batch normalized, and then the normalized feature map is corrected using the ReLU function to obtain a corrected feature map. Then, the corrected feature map is convolved again using a convolution kernel with a size of 1×1 and a number of channels of 64, and then modified using the Softmax activation function to obtain a fourth feature map. Then, the first and third feature maps are respectively combined with the fourth feature map by dot multiplication to obtain two intermediate feature maps. Then, the two intermediate feature maps are obtained by channel addition to obtain an intermediate feature map. Then, the intermediate feature map is processed using a convolution kernel with a size of 1×1 and a number of channels of 256 to obtain a fifth feature map. Finally, the fifth feature map is added to the feature map output by the main feature retention module to obtain a feature map with a size of 64×64 and a number of channels of 256), and then the feature map is input into the maximum pooling layer (to further reduce the size of the feature map) to proceed. The downsampled feature map is 32×32 in size and 256 in number of channels. Then, the downsampled 256-channel feature map is input into the third compact separation attention module to obtain a feature map with an output size of 32×32 and a number of channels of 512 (specifically, a convolution kernel with a size of 1×1 and a number of channels of 128 is first used to convolve the feature map, and two sets of convolution kernels with a size of 3×3 and a number of channels of 64 are used to calculate the convolution processing results to obtain the first feature map and the second feature map. Then, the first feature map and the second feature map are added to obtain the superimposed feature map, and then the superimposed feature map is processed using a convolution kernel with a size of 3×3 and a number of channels of 64 to obtain the optimized third feature map, and then the window size and the third feature map are used to optimize the feature map. Figure 1The third feature map is processed with a pooling window of the same size to obtain a channel feature map with a size of 1×1 and 64 channels. Then, the channel feature map is convolved with a convolution kernel with a size of 1×1 and a number of channels of 128 to obtain a feature map with a size of 1×1 and a number of channels of 128. Then, the feature map is batch normalized. After that, the normalized feature map is corrected using the ReLU function to obtain a corrected feature map. Then, the corrected feature map is convolved again with a convolution kernel with a size of 1×1 and a number of channels of 128. Then, the Softmax activation function is used to correct it to obtain a fourth feature map. Then, the first and third feature maps are respectively combined with the fourth feature map by dot multiplication to obtain two intermediate feature maps. Then, the two intermediate feature maps are obtained by channel addition. Then, the intermediate feature map is processed with a convolution kernel with a size of 1×1 and a number of channels of 512 to obtain a fifth feature map. Finally, the fifth feature map and the main feature map are combined. The feature maps output by the retention module are added to obtain a feature map with unchanged size and 512 channels); then, the feature map is input into the maximum pooling layer (the purpose is to halve the size of the feature map) to obtain a downsampled feature map with a size of 16×16 and 512 channels; finally, the feature map of 512 channels is input into the fourth compact separation attention module to obtain an intermediate feature map with a size of 16×16 and 512 channels (the specific processing process of the fourth compact separation attention module is exactly the same as that of the third compact separation attention module, which will not be repeated here).
[0109] The advantage of the above steps (2-4) is that the compact separation attention in the encoder can effectively enhance the feature extraction capability of U-Net to extract more effective features to provide effective support for the subsequent decoder work.
[0110] (2-5) For each medical image in the training set obtained in step (2-1), the feature map corresponding to the medical image obtained in step (2-3) is input into the multi-axis attention module to obtain the multi-axis attention feature map corresponding to the medical image.
[0111] Specifically, this step is as follows: first, the channel axis attention module uses the average pooling and maximum pooling on the channel dimension to obtain two channel dimension information, and then, the two channel dimension information is respectively input into the double convolution module (first through the 1×1 convolution layer to reduce the number of input channels, then through another 1×1 convolution layer, the number of channels is restored to the number of input feature channels, and finally the value output by the convolution layer is mapped to between 0 and 1 through the Sigmoid activation function) to obtain two channel attention maps with a size of 1×1 and a channel number of 64. Subsequently, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map Fc Then, F c Perform point-by-point multiplication with the feature map of size 256×256 and 64 channels output by the main feature retention module to obtain the feature map F′ c ; Then, the obtained feature map F′ c Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 256×256 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. The sigmoid activation function is then used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map. Then, the spatial attention map With the feature map F″ c Perform multiplication operation to obtain the final output feature map F c ″ is as follows (where Indicates that the matrix elements correspond to the product operation):
[0112]
[0113] Then, the feature map F″ is output c The height axis attention module and the width axis attention module are respectively input to extract the height axis feature map and the width axis feature map to obtain the feature map F″ h and F″ w ;
[0114] Finally, the multi-axis attention module converts the final height axis feature map F″ h And the final width axis feature map F″ w Superimpose them to obtain the multi-axis attention feature map F:
[0115] F=F″ h +F″ w
[0116] The above will output the feature map F″ c The height axis attention module and the width axis attention module are respectively input to extract the height axis feature map and the width axis feature map to obtain the feature map F″ h and F″ w The specific process is as follows:
[0117] First, the height axis attention module takes the feature map F″ output by the channel axis attention module cThe height axis and the channel axis are swapped so that the height axis becomes the new channel axis, and the original channel axis is transformed into the height axis to obtain a height axis feature map with a size of 64×256 and a channel number of 256; then the height axis attention module uses the average pooling and maximum pooling on the channel dimension to obtain two channel dimension information, and then the two channel dimension information are respectively input into the double convolution module (first through a 1×1 convolution layer to reduce the number of input channels, then through another 1×1 convolution layer, the number of channels is restored to the number of input feature channels, and finally the Sigmoid activation function is used to map the value output by the convolution layer to between 0 and 1) to obtain two channel attention maps with a size of 1×1 and a channel number of 256. Afterwards, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F h Then, F h Perform point-by-point multiplication with the height axis feature map of size 64×256 and number of channels 256 to obtain the feature map F′ h ; Then, the obtained feature map F′ h Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 64×256 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. The sigmoid activation function is then used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map. Then, the spatial attention map With the feature map F′ h Perform multiplication operation to obtain the feature map F′ h ' is as follows (where Indicates that the matrix elements correspond to the product operation):
[0118]
[0119] Then, the height axis attention module swaps the height axis and channel axis of the obtained height axis feature map again to obtain the final height axis feature map F′ with a size of 256×256 and 64 channels h ′.
[0120] At the same time, the width axis attention module first converts the feature map F′ output by the channel axis attention module c′ Swap the width axis and the channel axis so that the width axis becomes the new channel axis, and the original channel axis is transformed into the width axis to obtain a width axis feature map with a size of 256×64 and a channel number of 256; then the width axis attention module uses the average pooling and maximum pooling on the channel dimension to obtain two channel dimension information, and then, the two channel dimension information are respectively input into the double convolution module (first through a 1×1 convolution layer to reduce the number of input channels, then, through another 1×1 convolution layer, the number of channels is restored to the number of input feature channels, and finally through the Sigmoid activation function The value output by the convolution layer is mapped to between 0 and 1) to obtain two channel attention maps with a size of 1×1 and a channel number of 256. Afterwards, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F w Then, F w Perform point-by-point multiplication with the width axis feature map of size 256×64 and number of channels 256 to obtain the feature map F′ w ; Then, the obtained feature map F′ w Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 256×64 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. The sigmoid activation function is then used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map. Then, the spatial attention map With the feature map F′ w Perform multiplication operation to obtain the feature map F′ w ' is as follows (where Indicates that the matrix elements correspond to the product operation):
[0121]
[0122] Then, the width axis attention module swaps the width axis and channel axis of the obtained width axis feature map again to obtain the final width axis feature map F″ with a size of 256×256 and 64 channels w (The purpose is to make the final height axis feature map and width axis feature map consistent with the feature map F" output by the channel axis attention module c of the same size).
[0123] The advantage of the above steps (2-5) is that the multi-axis attention module improves the problem that the original CBAM attention module only considers features from the channel axis and ignores the potential feature information in the height axis and width axis directions, and increases the expression ability of the attention module by adding attention modules in the height axis and width axis directions.
[0124] (2-6) For each medical image in the training set obtained in step (2-1), the intermediate feature map corresponding to the medical image obtained in step (2-4) and the multi-axis attention feature map corresponding to the medical image obtained in step (2-5) are input into the decoder to obtain the final lesion segmentation prediction result F corresponding to the medical image. pred ;
[0125] Specifically, the decoder first performs an upsampling operation on the intermediate feature map of 16×16 and 512 channels output by the fourth compact separation attention module with a step size of 2*2 and a stride of 2 to obtain a feature map of size 32×32 and 512 channels, and then splices the feature map with the feature map of size 32×32 and 512 channels output by the third compact separation attention module in the channel dimension to obtain a spliced feature map of size 32×32 and 1024 channels; then, the spliced feature map is input into the fifth compact separation attention module to obtain a size A feature map with a size of 32×32 and a number of channels of 256 is first convolved with a convolution kernel with a size of 1×1 and a number of channels of 64, and the convolution result is calculated using two sets of convolution kernels with a size of 3×3 and a number of channels of 32 to obtain a first feature map and a second feature map. Then, the first feature map and the second feature map are added to obtain a superimposed feature map, and then, the superimposed feature map is processed using a convolution kernel with a size of 3×3 and a number of channels of 32 to obtain an optimized third feature map, and then, the window size and the third feature map are used. Figure 1The third feature map is processed by a pooling window of the same size to obtain a channel feature map with a size of 1×1 and 32 channels. Then, a convolution kernel with a size of 1×1 and 64 channels is used to convolve the channel feature map to obtain a feature map with a size of 1×1 and 64 channels. Subsequently, the feature map is batch normalized, and then the normalized feature map is corrected using the ReLU function to obtain a corrected feature map. Then, the corrected feature map is convolved again with a convolution kernel of size 1×1 and channel number 64, and then, it is corrected using the Softmax activation function to obtain the fourth feature map. Then, the first and third feature maps are respectively combined with the fourth feature map by dot multiplication to obtain two intermediate feature maps. Then, the two intermediate feature maps are obtained by channel addition. Then, the intermediate feature map is processed with a convolution kernel of size 1×1 and channel number 256 to obtain the fifth feature map. Then, the input feature map is downsampled using a convolution kernel of size 1×1 and channel number 256 to obtain a downsampled feature with a size of 32×32 and channel number 256. Finally, the fifth feature map is added to the downsampled feature to obtain the final output); then, the feature map is upsampled with a step size of 2*2 and a stride of 2 to obtain a feature map with a size of 64×64 and a channel number of 64×64. A feature map with a channel number of 256 is obtained, and then the feature map is spliced with the feature map with a size of 64×64 and a channel number of 256 output by the second compact separation attention module in the channel dimension to obtain a feature map with a size of 64×64 and a channel number of 512; then, the feature map is input into the sixth compact separation attention module to obtain a feature map with a size of 64×64 and a channel number of 128 (specifically, a convolution kernel with a size of 1×1 and a channel number of 32 is first used to convolve the feature map, and two sets of convolution kernels with a size of 3×3 and a channel number of 16 are used to calculate the convolution processing result to obtain the first feature map and the second feature map. Then, the first feature map and the second feature map are added to obtain a superimposed feature map, and then the superimposed feature map is processed using a convolution kernel with a size of 3×3 and a channel number of 16 to obtain an optimized third feature map, and then, the window size and the third feature map are used. Figure 1The third feature map is processed by a pooling window of the same size to obtain a channel feature map with a size of 1×1 and 16 channels. Then, a convolution kernel with a size of 1×1 and 32 channels is used to convolve the channel feature map to obtain a feature map with a size of 1×1 and 32 channels. Subsequently, the feature map is batch normalized, and then the normalized feature map is corrected using the ReLU function to obtain a corrected feature map. Then, the corrected feature map is convolved again with a convolution kernel of size 1×1 and channel number 32, and then, it is corrected using the Softmax activation function to obtain the fourth feature map. Then, the first and third feature maps are respectively combined with the fourth feature map by dot multiplication to obtain two intermediate feature maps. Then, the two intermediate feature maps are obtained by channel addition. Then, the intermediate feature map is processed using a convolution kernel of size 1×1 and channel number 128 to obtain the fifth feature map. Then, the input feature map is downsampled using a convolution kernel of 1×1 and channel number 128 to obtain a downsampled feature of size 64×64 and channel number 128. Finally, the fifth feature map is added to the downsampled feature to obtain the final output), then, the feature map is upsampled with a step size of 2*2 and a stride of 2, and the obtained The feature map of the first compact separation attention module is spliced with the feature map of size 128×128 and number of channels 128 output by the first compact separation attention module to obtain a spliced feature map of size 128×128 and number of channels 256. Then, the spliced feature map is input into the seventh compact separation attention module to obtain a feature map of size 128×128 and number of channels 64 (specifically, a convolution kernel of size 1×1 and number of channels 16 is first used to convolve the feature map, and two sets of convolution kernels of size 3×3 and number of channels 8 are used to calculate the convolution result to obtain the first feature map and the second feature map. Then, the first feature map and the second feature map are added to obtain the superimposed feature map, and then the superimposed feature map is processed using the convolution kernel of size 3×3 and number of channels 8 to obtain the optimized third feature map, and then the window size and the third feature map are used. Figure 1The third feature map is processed with a pooling window of the same size to obtain a channel feature map with a size of 1×1 and 8 channels. Then, the channel feature map is convolved with a convolution kernel with a size of 1×1 and 16 channels to obtain a feature map with a size of 1×1 and 16 channels. Then, the feature map is batch normalized. After that, the normalized feature map is corrected using the ReLU function to obtain a corrected feature map. Then, the corrected feature map is convolved again with a convolution kernel with a size of 1×1 and 16 channels. Then, the Softmax activation function is used to correct it to obtain a fourth feature map. Then, the first and third feature maps are respectively combined with the fourth feature map by dot multiplication to obtain two intermediate feature maps. Then, the two intermediate feature maps are obtained by channel addition. Then, the intermediate feature map is processed with a convolution kernel with a size of 1×1 and 64 channels to obtain the fifth feature map. Then, the convolution kernel with a size of 1×1 and a channel number of 64 is used to downsample the input feature map to obtain a downsampled feature with a size of 128×128 and a channel number of 64. Finally, the fifth feature map is added to the downsampled feature to obtain the final output); then, the feature map is upsampled with a step size of 2*2 and a stride of 2 to obtain a feature map with a size of 256×256 and a dimension of 64; then, the feature map is concatenated with the multi-axis attention feature map output by the multi-axis attention module to obtain a feature map with a size of 256×256 feature map with 128 channels; then, the feature map is input into the eighth compact separation attention module to obtain a feature map with a size of 256×256 and a number of channels of 64 (the specific processing process of the eighth compact separation attention module is exactly the same as that of the seventh compact separation attention module, which will not be repeated here); then, the number of channels of the feature map is reduced to two channels through a 1×1 convolution layer, and then the probability distribution matrix of the tumor position is obtained; then, the probability distribution matrix is mapped to the [0,1] interval through the Softmax function; finally, on the probability distribution matrix of the first channel, the element value with a probability value greater than 0.5 in the probability distribution matrix is set to 1, that is, the background (that is, the position is considered to be the background, that is, there is no lesion), and the other element values are set to 0; on the second channel, the element value with a probability value greater than 0.5 in the probability distribution matrix is set to 1, that is, the foreground (that is, the position is considered to be the foreground, that is, the lesion), and the other element values are set to 0; finally, the binarized probability distribution matrix is saved as a binarized image, which is the final lesion segmentation prediction result.
[0126] The advantage of the above steps (2-6) is that the compact separation attention in the decoder can fully integrate the features from the encoder part with the features of the decoder itself, thereby effectively enhancing the feature representation ability of U-Net to improve model performance.
[0127] (2-7) For each medical image in the training set obtained in step (2-1), the lesion segmentation prediction result F corresponding to the medical image obtained in step (2-6) is pred Calculate the loss function and use the loss function to iteratively train the multi-axis attention medical image segmentation model to obtain the preliminarily trained multi-axis attention medical image segmentation model W k The preliminarily trained medical image segmentation model is verified using the test set obtained in step (2-1) to obtain the segmentation accuracy of the test set.
[0128] Specifically, the loss function in this step is a comprehensive function, which is composed of the sum of the Focal loss function and the Dice loss function. The specific calculation formula is:
[0129] L=λ 1 ·L focal +λ 2 ·L dice
[0130]
[0131]
[0132] Among them, λ 1 and λ 2 are the weighted values of the two losses, both set to 0.5, α and γ are the built-in weights of Focal loss, set to 0.25 and 2 respectively, ∈ is a constant, whose value is 1*10 -8 , used to prevent the denominator from being 0, represents the true result of lesion segmentation of the i-th medical image in the training set obtained in step (2-1), represents the lesion segmentation prediction result of the i-th medical image in the training set obtained in step (2-1), i∈[1,N], and N represents the total number of medical images in the training set obtained in step (2-1).
[0133] (2-8) Repeat the above steps (2-3) to (2-7) until the total number of iterations set in step (2-2) is reached, or the segmentation accuracy on the test set reaches the optimal value, thereby obtaining a trained medical image segmentation model based on multi-axis attention.
[0134] It will be easily understood by those skilled in the art that the above description is only 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 protection scope of the present invention.
Claims
1. A medical image segmentation method based on multi-axis attention, characterized in that: The following steps are involved: (1) acquiring a medical image, and preprocessing the medical image to obtain a preprocessed medical image; (2) The medical image preprocessed in step (1) is input into a pre-trained multi-axis attention-based medical image segmentation model to obtain the final lesion segmentation prediction result.
2. The medical image segmentation method based on multi-axis attention according to claim 1, characterized in that: Medical images are two-dimensional static images, which can be tumor images obtained by ultrasound equipment or dermoscopic images; The preprocessing process in step (1) first extracts a region of interest from the medical image according to the specific location of the lesion in the preset segmentation label, and then converts the extracted region of interest into a medical image of 256×256 size and jpg format as the preprocessed medical image; The lesion segmentation result in step (2) is a binary image corresponding to the medical image obtained in step (1), in which black is the background and white is the foreground (i.e., the lesion). The medical image segmentation model based on multi-axis attention includes a main feature retention module, an encoder, a decoder, and a multi-axis attention module connected between the main feature retention module and the decoder. The input of the main feature retention module is the preprocessed medical image, and the output is a feature map with a size of 256×256 and 64 channels; The encoder contains four compact separation attention modules. The input of the encoder is the feature map of size 256×256 and 64 channels output by the main feature retention module, and its output is the intermediate feature map; The multi-axis attention module includes a channel-axis attention module, a height-axis attention module, and a width-axis attention module. The input of the multi-axis attention module is a feature map with a size of 256×256 and 64 channels output by the main feature retention module, and the output is a multi-axis attention feature map with the same size as the input feature map. The decoder contains four compact separation attention modules. The input of the decoder is the feature map with a size of 32×32 and 512 channels output by the third compact separation attention module in the encoder, the intermediate feature map with a size of 16×16 and 512 channels output by the fourth compact separation attention module, and the multi-axis attention feature map output by the multi-axis attention module. The output is the final lesion segmentation prediction result.
3. The medical image segmentation method based on multi-axis attention according to claim 1 or 2, characterized in that: The main feature retention module first uses a 7×7 convolution kernel to convolve the preprocessed medical image, and the convolution kernel has in_ch input channels and out_ch output channels. Then, the convolution medical image is corrected using the ReLU function to obtain the corrected features, and then the corrected features are batch normalized to obtain the normalized features. Then, another 7×7 convolution kernel is used to convolve the normalized features, and then the convolution features are corrected using the ReLU function to obtain the corrected features. Then, the corrected features are batch normalized to obtain the normalized intermediate feature map. Then, a 1×1 convolution kernel with no padding is used to convolve the normalized intermediate feature map, and then the convolution features are corrected using the ReLU function to obtain the corrected features. Then, the corrected features are batch normalized to obtain a normalized feature map with a size of 256×256 and 64 channels. The encoder first inputs the feature map output by the main feature retention module into the maximum pooling layer to obtain a feature map with a size of 128×128 and a number of channels of 64. Next, the encoder inputs the feature map into the first compact separation attention module to obtain a feature map with a size of 128×128 and a number of channels of 128; then, the feature map with a size of 128×128 and a number of channels of 128 obtained by the first compact separation attention module is input into the maximum pooling layer, so that the height and width of the feature map are halved, and the number of channels remains unchanged, so as to achieve spatial downsampling, thereby obtaining a feature map with a size of 64×64 and a number of channels of 128 after downsampling; then, the feature map with a size of 64×64 and a number of channels of 128 after downsampling is input into the second compact separation attention module, and the feature map with a size of 64×64 and a number of channels of 128 is obtained. To an output feature map with a size of 64×64 and 256 channels, the feature map is then input into the maximum pooling layer to obtain a downsampled feature map with a size of 32×32 and 256 channels; then, the downsampled 256-channel feature map is input into the third compact separation attention module to obtain an output feature map with a size of 32×32 and 512 channels, to obtain a downsampled feature map with a size of 16×16 and 512 channels; finally, the 512-channel feature map is input into the fourth compact separation attention module to obtain an intermediate feature map with a size of 16×16 and 512 channels.
4. The medical image segmentation method based on multi-axis attention according to any one of claims 1 to 3, characterized in that: The channel axis attention module first uses the average pooling and maximum pooling on the channel dimension to obtain two channel dimension information, and then inputs the two channel dimension information into the double convolution module to obtain two channel attention maps with a size of 1×1 and a channel number of 64. Afterwards, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F c Then, F c Perform point-by-point multiplication with the feature map of size 256×256 and 64 channels output by the main feature retention module to obtain the feature map F′ c ; Then, the obtained feature map F′ c Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 256×256 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. Subsequently, the Sigmoid activation function is used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map Then, the spatial attention map With the feature map F″ c Perform multiplication operation to obtain the final output feature map F c "as follows: Then, the output feature map F″ c The height axis attention module and the width axis attention module are respectively input to extract the height axis feature map and the width axis feature map to obtain the feature map F″ h and F″ w ; Finally, the multi-axis attention module converts the final height axis feature map F″ h And the final width axis feature map F″ w Superimpose them to obtain the multi-axis attention feature map F: F=F″ h +F″ w 。 5. The medical image segmentation method based on multi-axis attention according to claim 4, characterized in that: The decoder first performs an upsampling operation with a step size of 2*2 and a stride of 2 on the intermediate feature map of 16×16 and 512 channels output by the fourth compact separation attention module to obtain a feature map of size 32×32 and 512 channels, and then splices the feature map with the feature map of size 32×32 and 512 channels output by the third compact separation attention module in the channel dimension to obtain a spliced feature map of size 32×32 and 1024 channels; then, the spliced feature map is input into the fifth compact separation attention module to obtain a feature map of size 32×32 and 256 channels; then, the feature map is upsampled with a step size of 2*2 and a stride of 2 to obtain a feature map of size 64×64 and 256 channels, and then, the feature map is spliced with the feature map of size 64×64 and 256 channels output by the second compact separation attention module in the channel dimension to obtain a 64×64 feature map.
4. A feature map with a channel number of 512; then, the feature map is input into the sixth compact separation attention module to obtain a feature map with a size of 64×64 and a channel number of 128, and then, the feature map is upsampled with a step size of 2*2 and a stride of 2, and the obtained feature map is spliced with the feature map with a size of 128×128 and a channel number of 128 output by the first compact separation attention module to obtain a spliced feature map with a size of 128×128 and a channel number of 256, and then, the spliced feature map is input into the seventh compact separation attention module to obtain a feature map with a size of 128×128 and a channel number of 64; then, the feature map is upsampled with a step size of 2*2 and a stride of 2 to obtain a feature map with a size of 256×256 and a dimension of 64; then, the feature map is spliced with the multi-axis attention feature map output by the multi-axis attention module into a feature map with a size of 256×256 and a channel number of 128; Subsequently, the feature map is input into the eighth compact separation attention module to obtain a feature map with a size of 256×256 and 64 channels; Afterwards, a 1×1 convolutional layer is used to reduce the number of channels of the feature map to two channels, and then the probability distribution matrix of the tumor location is obtained; then, the probability distribution matrix is mapped to the [0,1] interval through the Softmax function; finally, in the probability distribution matrix of the first channel, the element value with a probability value greater than 0.5 in the probability distribution matrix is set to 1, that is, the background (that is, the position is considered to be the background, that is, there is no lesion), and the other element values are set to 0; in the second channel, the element value with a probability value greater than 0.5 in the probability distribution matrix is set to 1, that is, the foreground (that is, the position is considered to be the foreground, that is, the lesion), and the other element values are set to 0; finally, the binarized probability distribution matrix is saved as a binarized image, which is the final lesion segmentation prediction result.
6. The medical image segmentation method based on multi-axis attention according to claim 5, characterized in that: The multi-axis attention based medical image segmentation model is obtained through the following steps: (2-1) Obtain a medical image dataset, preprocess the medical image dataset to obtain a preprocessed medical image dataset, and divide the preprocessed medical image dataset into a training set and a test set in a ratio of 8:
2. (2-2) Set the total number of iterations epoch of the training process and initialize the learning rate η of the medical image segmentation model based on multi-axis attention of the present invention; the total number of iterations epoch ranges from 100 to 300, and the learning rate η ranges from 0 to 1. (2-3) For each medical image in the training set obtained in step (2-1), the medical image is input into the main feature retention module to obtain a feature map corresponding to the medical image with a size of 256×256 and a channel number of 64. (2-4) For each medical image in the training set obtained in step (2-1), the feature map corresponding to the medical image obtained in step (2-3) is input into the encoder to obtain an intermediate feature map corresponding to the medical image. (2-5) For each medical image in the training set obtained in step (2-1), the feature map corresponding to the medical image obtained in step (2-3) is input into the multi-axis attention module to obtain the multi-axis attention feature map corresponding to the medical image. (2-6) For each medical image in the training set obtained in step (2-1), the intermediate feature map corresponding to the medical image obtained in step (2-4) and the multi-axis attention feature map corresponding to the medical image obtained in step (2-5) are input into the decoder to obtain the final lesion segmentation prediction result F corresponding to the medical image. pred ; (2-7) For each medical image in the training set obtained in step (2-1), the lesion segmentation prediction result F corresponding to the medical image obtained in step (2-6) is pred Calculate the loss function and use the loss function to iteratively train the multi-axis attention medical image segmentation model to obtain the preliminarily trained multi-axis attention medical image segmentation model W k The preliminarily trained medical image segmentation model is verified using the test set obtained in step (2-1) to obtain the segmentation accuracy of the test set. (2-8) Repeat the above steps (2-3) to (2-7) until the total number of iterations set in step (2-2) is reached, or the segmentation accuracy on the test set reaches the optimal value, thereby obtaining a trained medical image segmentation model based on multi-axis attention.
7. The medical image segmentation method based on multi-axis attention according to claim 6, characterized in that: Step (2-3) is specifically as follows: first, the main feature retention module uses a 7×7 convolution kernel to convolve the preprocessed medical image, and the convolution kernel has in_ch input channels and out_ch output channels. Then, the convolved medical image is corrected using the ReLU function to obtain the corrected features. Then, the corrected features are batch normalized to obtain the normalized features. Then, another 7×7 convolution kernel is used to convolve the normalized features. Then, the convolution kernel is used to perform the batch normalization. The ReLU function corrects the convolved features to obtain the corrected features, and then the corrected features are batch normalized to obtain the normalized intermediate feature map. Thereafter, a convolution kernel of size 1×1 and no padding is used to convolve the normalized intermediate feature map, and then the ReLU function is used to correct the convolved features to obtain the corrected features. Thereafter, the corrected features are batch normalized to obtain a normalized feature map of size 256×256 and number of channels 64. Specifically, step (2-4) is that the encoder first inputs the feature map output by the main feature retention module into the maximum pooling layer to obtain a feature map with a size of 128×128 and 64 channels. Next, the encoder inputs the feature map into the first compact separation attention module to obtain a feature map with a size of 128×128 and a number of channels of 128; then, the feature map with a size of 128×128 and a number of channels of 128 obtained by the first compact separation attention module is input into the maximum pooling layer, so that the height and width of the feature map are halved, and the number of channels remains unchanged, so as to achieve spatial downsampling, thereby obtaining a downsampled feature map with a size of 64×64 and a number of channels of 128; then, the downsampled feature map with a size of 64×64 and a number of channels of 128 is input into the second compact separation attention module to obtain an output feature map with a size of 64×64 and a number of channels of 256, and then the feature map is input into the maximum pooling layer to obtain a downsampled feature map with a size of 32×32 and a number of channels of 256; then, the downsampled 256-channel feature map is input into the third compact separation attention module to obtain an output feature map with a size of 32×32 and a number of channels of 512; Subsequently, the feature map is input into the maximum pooling layer to obtain a downsampled feature map with a size of 16×16 and a number of channels of 512. Finally, the feature map with 512 channels is input into the fourth compact separation attention module to obtain an intermediate feature map with a size of 16×16 and a number of channels of 512.
8. The medical image segmentation method based on multi-axis attention according to claim 7, characterized in that: Step (2-5) is as follows: first, the channel axis attention module uses the average pooling and maximum pooling on the channel dimension to obtain two channel dimension information, and then, the two channel dimension information is respectively input into the double convolution module to obtain two channel attention maps with a size of 1×1 and a channel number of 64. Thereafter, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F c Then, F c Perform point-by-point multiplication with the feature map of size 256×256 and 64 channels output by the main feature retention module to obtain the feature map F′ c ; Then, the obtained feature map F′ c Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 256×256 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. Subsequently, the Sigmoid activation function is used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map Then, the spatial attention map With the feature map F″ c Perform multiplication operation to obtain the final output feature map F c "as follows: Then, the output feature map F″ c The height axis attention module and the width axis attention module are respectively input to extract the height axis feature map and the width axis feature map to obtain the feature map F″ h and F″ w ; Finally, the multi-axis attention module converts the final height axis feature map F″ h And the final width axis feature map F″ w Superimpose them to obtain the multi-axis attention feature map F: F=F″ h +F″ w The output feature map F″ c The height axis attention module and the width axis attention module are respectively input to extract the height axis feature map and the width axis feature map to obtain the feature map F″ h and F″ w The specific process is as follows: First, the height axis attention module takes the feature map F″ output by the channel axis attention module c The height axis and the channel axis are swapped so that the height axis becomes the new channel axis, and the original channel axis is converted to the height axis to obtain a height axis feature map with a size of 64×256 and a channel number of 256. Then, the height axis attention module uses the average pooling and maximum pooling on the channel dimension to obtain two channel dimension information. Then, the two channel dimension information is input into the double convolution module respectively to obtain two channel attention maps with a size of 1×1 and a channel number of 256. After that, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F. h Then, F h Perform point-by-point multiplication with the height axis feature map of size 64×256 and number of channels 256 to obtain the feature map F′ h ; Then, the obtained feature map F′ h Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 64×256 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. Subsequently, the Sigmoid activation function is used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map Then, the spatial attention map With the feature map F′ h Perform multiplication operation to obtain the feature map F″ h as follows: Then, the height axis attention module swaps the height axis and channel axis of the obtained height axis feature map again to obtain the final height axis feature map F″ with a size of 256×256 and 64 channels h . At the same time, the width axis attention module first converts the feature map F″ output by the channel axis attention module c The width axis and the channel axis are swapped so that the width axis becomes the new channel axis, and the original channel axis is converted to the width axis to obtain a width axis feature map with a size of 256×64 and a channel number of 256; then the width axis attention module uses the average pooling and maximum pooling on the channel dimension to obtain two channel dimension information, and then the two channel dimension information is respectively input into the double convolution module to obtain two channel attention maps with a size of 1×1 and a channel number of 256. Afterwards, the two channel attention maps are superimposed and input into the Sigmoid activation function to obtain the final channel attention map F w Then, F w Perform point-by-point multiplication with the width axis feature map of size 256×64 and number of channels 256 to obtain the feature map F′ w ; Then, the obtained feature map F′ w Perform maximum pooling and average pooling operations on the channel dimension to obtain the maximum pooling feature map And average pooling feature map The two are then channel-concatenated to obtain a feature map of size 256×64 and dimension 2. The concatenated feature map is then extracted using a convolutional layer of size 7×7 and output channel number 1 to generate a spatial attention map of the same size as the concatenated feature map and channel number 1. Subsequently, the Sigmoid activation function is used to normalize the spatial attention map between 0 and 1 to obtain the final spatial attention map Then, the spatial attention map With the feature map F′ w Perform multiplication operation to obtain the feature map F″ w as follows: Then, the width axis attention module swaps the width axis and channel axis of the obtained width axis feature map again to obtain the final width axis feature map F″ with a size of 256×256 and 64 channels w .
9. The medical image segmentation method based on multi-axis attention according to claim 8, characterized in that: Step (2-6) is specifically as follows: the decoder first performs an upsampling operation with a step size of 2*2 and a stride of 2 on the intermediate feature map of 16×16 and 512 channels output by the fourth compact separation attention module to obtain a feature map of size 32×32 and 512 channels, and then splices the feature map with the feature map of size 32×32 and 512 channels output by the third compact separation attention module in the channel dimension to obtain a spliced feature map of size 32×32 and 1024 channels; then, the spliced feature map is input into the fifth compact separation attention module to obtain a feature map of size 32×32 and 256 channels; then, the feature map is upsampled with a step size of 2*2 and a stride of 2 to obtain a feature map of size 64×64 and 256 channels, and then, the feature map is spliced with the feature map of size 64×64 and 256 channels output by the second compact separation attention module in the channel dimension to obtain a feature map of size A feature map of size 64×64 and number of channels 512 is obtained; then, the feature map is input into the sixth compact separation attention module to obtain a feature map of size 64×64 and number of channels 128, then, the feature map is upsampled with a step size of 2*2 and a stride of 2, and the obtained feature map is spliced with the feature map of size 128×128 and number of channels 128 output by the first compact separation attention module in the channel dimension to obtain a spliced feature map of size 128×128 and number of channels 256, then, the spliced feature map is input into the seventh compact separation attention module to obtain a feature map of size 128×128 and number of channels 64; then, the feature map is upsampled with a step size of 2*2 and a stride of 2 to obtain a feature map of size 256×256 and dimension 64; then, the feature map is spliced with the multi-axis attention feature map output by the multi-axis attention module into a feature map of size 256×256 and number of channels 128; Subsequently, the feature map is input into the eighth compact separation attention module to obtain a feature map with a size of 256×256 and 64 channels; After that, the number of channels of the feature map is reduced to two channels through a 1×1 convolution layer, and then the probability distribution matrix of the tumor location is obtained; then, the probability distribution matrix is mapped to the [0,1] interval through the Softmax function; finally, in the probability distribution matrix of the first channel, the element value with a probability value greater than 0.5 in the probability distribution matrix is set to 1, that is, the background (that is, the position is considered to be the background, that is, there is no lesion), and the other element values are set to 0; in the second channel, the element value with a probability value greater than 0.5 in the probability distribution matrix is set to 1, that is, the foreground, and the other element values are set to 0; finally, the binarized probability distribution matrix is saved as a binarized image, which is the final lesion segmentation prediction result; The loss function in step (2-7) is a comprehensive function, which is composed of the sum of the Focal loss function and the Dice loss function. The specific calculation formula is: L=λ1·L focal +λ2·L dice Among them, λ1 and λ2 are the weighted values of the two losses, both set to 0.5, α and γ are the built-in weights of Focal loss, set to 0.25 and 2 respectively, and ∈ is a constant whose value is 1*10 -8 , used to prevent the denominator from being 0, represents the true result of lesion segmentation of the i-th medical image in the training set obtained in step (2-1), represents the lesion segmentation prediction result of the i-th medical image in the training set obtained in step (2-1), i∈[1,N], and N represents the total number of medical images in the training set obtained in step (2-1).
10. A medical image segmentation system based on multi-axis attention, characterized in that: include: The first module is used to acquire a medical image and preprocess the medical image to acquire a preprocessed medical image; The second module is used to input the medical image preprocessed by the first module into a pre-trained multi-axis attention-based medical image segmentation model to obtain the final lesion segmentation prediction result.
Citation Information
Cited By
Deep learning-based chromosome image segmentation identification method and system
CN120783337A
Medical image segmentation method for multi-modal brain tumor data
CN122244455A
