CN115187621A - Automatic U-Net medical image contour extraction network integrating attention mechanism - Google Patents

Abstract

The invention discloses a U-Net medical image contour automatic extraction network fused with an attention mechanism, comprising an RGB image input module, the output end of the RGB image input module is connected to the input end of the feature extraction module, and the feature extraction module includes a feature encoding module , feature decoding module and attention module; the output end of the feature extraction module is connected to the input end of the MLP, and the attention module includes spatial attention and channel attention, which are used to suppress the neurons in the non-attention area; the MLP output unit is set as 2 neurons, respectively representing the probability of foreground and background, and followed by Softmax and Marching Square in turn. The invention integrates the attention module, improves the edge contour extraction accuracy, preliminarily solves the problem of blurred edges generated by the traditional frame, and reduces the interference of background noise, thereby basically meeting the accuracy requirements for medical image contour extraction in the medical field; The process of the framework greatly saves the time and cost of obtaining the target model.

Description

U-Net Medical Image Contour Automatic Extraction Network with Attention Mechanism

technical field

The invention relates to the technical field of medical image processing, in particular to a U-Net medical image contour automatic extraction network fused with an attention mechanism.

Background technique

Medical images can reflect anatomical structures or functional tissues in the human body. Dividing a medical image into several disjoint regions according to a certain similarity feature in the medical image, that is, medical image segmentation, is the most important basis in medical image analysis. Accurate, robust and fast image segmentation is the most important step before subsequent steps such as quantitative analysis and 3D visualization, and also lays the most fundamental foundation for important clinical applications such as image-guided surgery, radiotherapy planning, and treatment evaluation.

In recent years, with the development of deep neural network in the field of medical image processing, deep learning has become the mainstream method in medical image segmentation tasks. The practice of many researchers has proved that the segmentation method based on deep learning has a strong role in the field of medical image segmentation. application potential. The deep learning segmentation method is to achieve the segmentation of medical images by classifying the pixels. Unlike traditional pixel or superpixel classification methods that use hand-crafted features, deep learning methods can automatically learn task-relevant features from medical images and classify pixels based on these features, thereby achieving end-to-end performance. segmentation. Among them, U-Net is currently the most widely used framework in the field of medical image segmentation.

Prior art: The U-Net network structure can obtain the detailed information and contour information of the image in the encoder part; then, the extracted features are transferred to the decoder part through the skip connection stage; finally, the decoder part combines multiple scales. feature for feature recovery. Due to its U-shaped structure, U-Net can be trained with fewer pictures to get a good model. The U-Net network can be divided into a feature extraction network and a feature fusion network. The feature extraction network uses a convolution layer and a pooling layer to implement downsampling operations, while the feature fusion network is an upsampling operation, which can restore the image resolution. At the same time, the network gradually converge to the target area. In the feature fusion stage, the features extracted from the same level are fused again to avoid loss of details.

Although the medical image segmentation method based on U-Net has achieved remarkable results, it is still very difficult to obtain accurate segmentation results due to the influence of noise. Most methods still have blurred edges, neglected details, and need to manually adjust parameters. question. To this end, we introduce an automatic U-Net medical image contour extraction network fused with attention mechanism.

SUMMARY OF THE INVENTION

The purpose of the present invention is to provide a U-Net medical image contour automatic extraction network fused with an attention mechanism, so as to solve the problems raised in the above background art.

In order to achieve the above object, the present invention provides the following technical solutions: a U-Net medical image contour automatic extraction network fused with attention mechanism, including an RGB image input module, the output end of the RGB image input module is connected to the feature extraction module. The input end, the feature extraction module includes a feature encoding module, a feature decoding module and an attention module;

The output end of the feature extraction module is connected to the input end of the MLP, and the attention module includes spatial attention and channel attention, and is used to suppress neurons in non-attention areas;

The MLP is used for classifying the extracted features, and its output unit is set to 2 neurons, which represent the probability of foreground and background respectively, and is followed by Softmax and Marching Square in turn.

The RGB image input module is used for inputting an RGB image.

The feature extraction module is used to extract the features of the RGB image. After the feature extraction module obtains the features of the RGB image, there is a C-dimensional feature representation for each pixel, and the local and global information is fused. Therefore, only the The C-dimensional feature performs an MLP inference. Since the task at this stage is a binary classification, 2-dimensional information is finally obtained, representing the probability of the target and the non-target respectively. By comparing the target and non-target, a binary image can be obtained. Finally, use Algorithm for binarizing image contour extraction.

The attention module is used to remove interference information in RGB images.

The feature encoding module adopts ResNet18.

The channel attention is as follows:

Mc(F)=F*Sigmoid(MLP( _AvgPool (F))+MLP(MaxPool(F)));

表示将A和B按通道拼接；The calculation of the spatial attention is as follows, where

Indicates that A and B are spliced by channel;

Combining channel attention and spatial attention, the calculation formula of CBAM is obtained:

M(F)=Ms(Mc(G(F)))+F; where: G(F) _{= Conv2 ( Conv1} (F)).

The MLP uses a multilayer perceptron with 3 hidden layers.

Compared with the prior art, the beneficial effects of the present invention are: the present invention improves the edge contour extraction accuracy by integrating the attention module in a certain way, preliminarily solves the problem of blurred edges generated by the traditional framework, and reduces the interference of background noise. , which basically meets the precision requirements of medical image contour extraction in the medical field.

The invention simplifies the flow of the traditional framework, so that the time used for training and inference is relatively small, and the time and cost for obtaining the target model are greatly saved.

The present invention performs the final extraction of the contour through the Marching Square algorithm, the algorithm is simple and fast to implement, and can be processed in parallel.

Description of drawings

1 is a schematic diagram of the frame structure of a U-Net medical image contour automatic extraction network incorporating an attention mechanism according to the present invention;

Fig. 2 is Backbone frame structure schematic diagram;

Figure 3 is a schematic diagram of the MLP frame structure;

4 is a schematic diagram of an MLP calling process;

Figure 5 is a schematic diagram of the first basic situation of Marching Squares;

Figure 6 is a schematic diagram of the second basic situation of Marching Squares;

Figure 7 is a schematic diagram of the third basic situation of Marching Squares;

Figure 8 is a schematic diagram of the fourth basic situation of Marching Squares;

Figure 9 is a schematic diagram of the fifth basic situation of Marching Squares;

Figure 10 is a schematic diagram of the sixth basic case of Marching Squares.

Detailed ways

The technical solutions in the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings in the embodiments of the present invention. Obviously, the described embodiments are only a part of the embodiments of the present invention, but not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those of ordinary skill in the art without creative efforts shall fall within the protection scope of the present invention.

1-10, the present invention provides a technical solution: a U-Net medical image contour automatic extraction network fused with attention mechanism, including an RGB image input module, the output of the RGB image input module is connected to the feature The input end of the extraction module, the feature extraction module includes a feature encoding module, a feature decoding module and an attention module;

The output end of the feature extraction module is connected to the input end of the MLP, and the attention module includes spatial attention and channel attention, and is used to suppress neurons in non-attention areas;

The MLP is used for classifying the extracted features, and its output unit is set to 2 neurons, which represent the probability of foreground and background respectively, and is followed by Softmax and Marching Square in turn.

Input an RGB image, first perform feature extraction on it. The feature extraction module is similar to U-Net. One part is the encoding network (ie, the feature encoding module), and the other part is the decoding network (feature decoding module). At the same time, the attention module is integrated. Through subsequent ablation experiments, it is found that the added attention module It can well remove the interference information in the picture.

After the feature is obtained, there is a C-dimensional feature representation for each pixel, which combines local and global information, so only one MLP inference is needed for the C-dimensional feature of each pixel. Since the task at this stage is two-classification, the final result is obtained. 2-dimensional information, representing the probability of target and non-target, respectively. By comparing the target and the non-target, a binary image can be obtained. Finally, the algorithm of binary image contour extraction can be used. In the experiment, the Marching Squares algorithm is used to extract the contour.

The feature extraction module is divided into two parts, one is the feature encoding module and the other is the feature decoding module. The feature encoding module uses ResNet18, which is easy to train, easy to implement, and has relatively few parameters. More importantly, it has a down-sampling structure, which is very suitable for fast feature extraction.

The feature decoding module is similar to the feature decoding part of U-Net, but an attention module is added before each upsampling. The CBAM module (i.e., the attention module) is a network with attention mechanisms, including spatial attention and channel attention, for suppressing neurons in non-attention areas.

The channel attention is given by:

Mc(F)=F*Sigmoid(MLP( _AvgPool (F))+MLP(MaxPool(F)));

表示将A和B按通道拼接；The calculation of spatial attention is as follows, where

Indicates that A and B are spliced by channel;

Combining channel attention and spatial attention, the calculation formula of CBAM is obtained.

M(F)=Ms(Mc(G(F)))+F; where: G(F) _{= Conv2 ( Conv1} (F));

After feature encoding, the details of some pictures are often obtained, and these details include the details of the target and non-target. If attention is processed at this time and the features of the non-target part are suppressed in advance, the predicted target area is easier to follow. The MLP module considers it to be 1, which is beneficial to the subsequent MLP inference.

The feature extraction module is similar to U-Net, but unlike U-Net, the output of this module is not a probability map, but a description of features. Assuming that the input image is W*H*3, then the dimension obtained after feature extraction is W*H*C, where C is the number of feature descriptions. In the experiment, it is found that setting C=512 makes the network more robust. For each pixel, there is a feature description of a C-dimensional vector. These features not only describe the information around the pixel, but also integrate the global information, so the subsequent judgment only needs to infer the vector of this pixel.

The multi-layer perceptron (ie MLP) designed in the present invention is used for classifying the extracted features. PIFU will reclassify the entire feature sampling. Different from PIFU, the feature will not be sampled in the framework of the present invention, because this is the information of a two-dimensional picture, and the calculation amount is not large. Furthermore, the current device computing power and video memory capacity There is a significant improvement, and all features can be trained without sampling. For each pixel, there is a C-dimensional vector describing the global feature information of the pixel.

The framework of the present invention sets 3 hidden layers, which are [512, 256, 128] respectively. Since the purpose of image segmentation is binary classification, the output unit is set to 2 neurons, representing the probability of foreground and background respectively, and followed by Softmax and Marching Squares.

Finally, compare the probabilities of the foreground and the background. As long as the foreground is greater than the background, the pixel is set to 1, and a pure binary image can be obtained through the above process.

The benefits of doing this:

(1) Improve the accuracy. In the past image segmentation, a probability map was obtained, and probability calculations were performed on both the foreground and the background. It can be seen from the ablation experiment that the introduction of MLP can improve the accuracy by about 0.3%.

(2) There is no need to manually set the threshold. In the past image contour extraction, it is often necessary to perform an artificially set threshold on the probability map to obtain a pure binary image. The design of this framework will avoid setting this hyperparameter and avoid the impact of artificially set thresholds on the results.

After the pure binary image is obtained, theoretically any binary image contour extraction algorithm can extract the contour, and the present invention uses the Marching Squares algorithm. Similar to Marching Cubes, Marching Squares is an algorithm for extracting contour lines. Given a two-dimensional probability map, according to a threshold, linear interpolation is used to obtain the curve where the threshold is located.

Since the output of the present invention is a binary image, when the Marching Squares algorithm is used, as long as the threshold value is greater than 0, the obtained result is the same. The reason for using the Marching Squares algorithm is that it is simple to implement and requires little computation. For the four points located in the cell, there are the following 6 basic cases, and 16 cases can be obtained by rotating the mirror. It is only necessary to determine which case is the edge, and the algorithm can be run in parallel, which means that it can be further optimized and accelerated.

key point:

1. On the basis of the UNet framework, the present invention introduces an attention module, uses a fully convolutional network to obtain a binary mask of medical image objects, separates the objects from the background image, improves the accuracy of contour extraction, and can Better feed features to the contour extraction stage.

2. The present invention uses a multi-layer perceptron with 3 hidden layers as a classifier, and performs a second correction from the mask information of the previous step, so that the result is more accurate, and the contour boundary is clear, which is conducive to the use of the MarchingSquare algorithm Extract contours.

3. The present invention uses Marching Square as the final contour extraction algorithm. Compared with the neural network method, this method improves the performance of contour extraction and can accelerate the execution of the process in parallel.

Protection point:

1. The invention is based on the traditional U-Net, and introduces an attention module in the up-sampling stage of U-Net. Using a fully convolutional network, the binary mask of the medical image object is obtained, and the object is segmented from the background image. should be within the scope of protection of the present invention.

2. The present invention uses a multi-layer perceptron with 3 hidden layers as a classifier, and performs a second correction on the binary mask, so that the results are more accurate, and the contour boundaries are clear, so as to better achieve the accuracy of the medical field. The requirements for the effect of image contour extraction should be within the protection scope of the present invention.

3. The present invention uses Marching Square as the final contour extraction algorithm. Compared with the neural network method, this method improves the performance of contour extraction and can accelerate the execution of the process in parallel. The application of the Marching Square algorithm in medical contour extraction should fall within the protection scope of the present invention.

Prior art: The U-Net network structure can obtain the detailed information and contour information of the image in the encoder part; then, the extracted features are transferred to the decoder part through the skip connection stage; finally, the decoder part combines multiple scales. feature for feature recovery. Due to its U-shaped structure, U-Net can be trained with fewer pictures to get a good model. The U-Net network can be divided into a feature extraction network and a feature fusion network. The feature extraction network uses a convolution layer and a pooling layer to implement downsampling operations, while the feature fusion network is an upsampling operation, which can restore the image resolution. At the same time, the network gradually converge to the target area. In the feature fusion stage, the features extracted from the same level are fused again to avoid loss of details.

Although the medical image segmentation method based on U-Net has achieved remarkable results, it is still very difficult to obtain accurate segmentation results due to the influence of noise. Most methods still have blurred edges, neglected details, and need to manually adjust parameters. question.

Disadvantages of the prior art:

Disadvantages of existing medical image contour extraction networks:

1. The detection accuracy is relatively low, and it is prone to false detection or missed detection.

2. A large number of datasets are required for training, which consumes a lot of cost and time.

3. The edges of the output result are blurred and details are ignored.

4. The traditional network has one or more hyperparameters that need to be adjusted manually, and the adjustment of the hyperparameters has a direct impact on the results.

The reasons for the above shortcomings are:

1. The model does not pay attention to some important features and is easily disturbed by noise.

2. The design of the model itself is more complicated, resulting in a lot of time spent on training.

3. The convolution operation on tensors in CNN tends to make details ignored, making the results tend to be in a stable state between certainty and uncertainty.

4. It is necessary to manually set a threshold to determine the contour edge, usually set to 0.5, but this value is not optimal, and there are one or more different optimal values for different data.

In order to solve the problems of blurred edges, ignoring details, etc., improve the accuracy and training efficiency of contour extraction, and reduce the artificial influence on hyperparameters, specifically, in view of the above problems, the present invention focuses on a pixel-based pixel-based algorithm realized by an attention mechanism. The network that judges and does not need to use ACMs, combined with the attention mechanism, can converge faster, and can train well without a large number of data samples. Using this network to infer the probability that the pixel is located in the inner area, the binary image of the object can be directly output through certain steps, without the need to manually set the threshold. The main improvement of the present invention is as follows:

(1) In order to solve the problem of interference items, this paper proposes a U-Net module with an attention mechanism, which can effectively remove interference, and the recall rate in the Drosophila embryo dataset is 6.3% higher than that of U-Net.

(2) MLP is used to represent the probability that the pixel is located in the target, and the features extracted by U-Net are recalculated instead of directly obtained by U-Net. In the Drosophila embryo dataset, the accuracy improved by 0.3% over U-Net.

(3) The size of the manual setting of the threshold has a great influence on the results. In order to avoid the manual setting of the threshold, the one-hot encoding method is used to solve this problem, which is simpler and more effective than the method of adaptively setting the threshold.

By integrating the attention module in a certain way, the invention improves the edge contour extraction accuracy, preliminarily solves the problem of blurred edges generated by the traditional frame, and reduces the interference of background noise, thereby basically satisfying the medical image contour extraction accuracy in the medical field Require.

The invention simplifies the flow of the traditional framework, so that the time used for training and inference is relatively small, and the time and cost for obtaining the target model are greatly saved.

The present invention performs the final extraction of the contour through the Marching Square algorithm, the algorithm is simple and fast to implement, and can be processed in parallel.

The technical problem solved by the present invention:

1. Add the attention module, use the CBAM module to make the network focus on the area of interest and remove interference.

2. Use the one-hot encoding method to solve the problem of manually setting the threshold.

3. Improve and optimize the existing model so that the training time is relatively short and the model is as lightweight as possible.

Although embodiments of the present invention have been shown and described, it will be understood by those skilled in the art that various changes, modifications, and substitutions can be made in these embodiments without departing from the principle and spirit of the invention and modifications, the scope of the present invention is defined by the appended claims and their equivalents.

Claims (7)

1. A U-Net medical image contour automatic extraction network integrating attention mechanism comprises an RGB image input module and is characterized in that: the output end of the RGB image input module is connected to the input end of the feature extraction module, and the feature extraction module comprises a feature coding module, a feature decoding module and an attention module;

the output end of the feature extraction module is connected to the input end of the MLP, and the attention module comprises space attention and channel attention and is used for inhibiting neurons in a non-attention area;

the MLP is used for classifying the extracted features, the output elements of the MLP are set to be 2 neurons which respectively represent the probability of the foreground and the probability of the background, and the probability of the foreground and the probability of the background are sequentially connected with Softmax and Marching Square.

2. The attention mechanism-fused U-Net medical image contour automatic extraction network according to claim 1, wherein: the RGB image input module is used for inputting an RGB image.

3. The attention mechanism-fused U-Net medical image contour automatic extraction network according to claim 1, wherein: the feature extraction module is used for extracting features of the RGB image, after the features of the RGB image are obtained, C-dimensional feature representation is carried out on each pixel, local and global information is fused, therefore, only one MLP reasoning needs to be carried out on the C-dimensional features of each pixel, as tasks in the stage are classified in two ways, 2-dimensional information is obtained finally, the probability of a target and the probability of a non-target are represented respectively, a binary image can be obtained by comparing the target with the non-target, and finally an algorithm for extracting binary image contours is adopted.

4. The attention mechanism-fused U-Net medical image contour automatic extraction network according to claim 1, wherein: the attention module is used for removing interference information in the RGB image.

5. The attention mechanism-fused U-Net medical image contour automatic extraction network according to claim 1, wherein: the feature encoding module uses ResNet18.

6. The attention mechanism-fused U-Net medical image contour automatic extraction network according to claim 1, wherein: the channel attention is given by:

M _c (F)＝F*Sigmoid(MLP(AvgPool(F))+MLP(MaxPool(F)))；

the spatial attention is calculated as follows, wherein

The method comprises the following steps of (1) splicing A and B according to channels;

combining the channel attention and the spatial attention, a calculation formula of the CBAM is obtained:

M(F)＝M _s (M _c (G (F))) + F; wherein: g (F) = Conv ₂ (Conv ₁ (F))。

7. The attention mechanism-fused U-Net medical image contour automatic extraction network according to claim 1, wherein: the MLP employs a multi-layer perceptron with 3 hidden layers.

Images