U-net medical image contour automatic extraction network with attention mechanism

Abstract

The application discloses a U-Net medical image profile automatic extraction network fusing an attention mechanism, comprising an RGB image input module, an output end of the RGB image input module being connected to an input end of a feature extraction module, the feature extraction module comprising a feature encoding module, a feature decoding module and an attention module; an output end of the feature extraction module being connected to an input end of an MLP, the attention module comprising spatial attention and channel attention, and being used for inhibiting neurons of non-attention regions; the MLP output end being provided with 2 neurons, which represent probabilities of foreground and background respectively, and a Softmax and a Marching Square being sequentially connected after the MLP output end. The application fuses the attention module, improves edge profile extraction precision, preliminarily solves the problem of fuzzy edges generated by a traditional framework, and reduces interference of background noise, so that the precision requirement of the medical field on medical image profile extraction is basically met; the process of the traditional framework is simplified, and the time and cost of obtaining a target model are greatly saved.

Description

Technical Field

[0001] This invention relates to the field of medical image processing technology, specifically to a U-Net medical image contour automatic extraction network that incorporates an attention mechanism. Background Technology

[0002] Medical images can reflect the anatomical structures or functional tissues within the human body. Dividing a medical image into several non-overlapping regions based on certain similarity features—a process known as medical image segmentation—is the most fundamental aspect of medical image analysis. Accurate, robust, and rapid image segmentation is the most crucial step before subsequent steps such as quantitative analysis and 3D visualization, and it also lays the fundamental foundation for important clinical applications such as image-guided surgery, radiotherapy planning, and treatment evaluation.

[0003] In recent years, with the development of deep neural networks in the field of medical image processing, deep learning has become the mainstream method for medical image segmentation tasks. Numerous researchers have demonstrated that deep learning-based segmentation methods have strong application potential in medical image segmentation. Deep learning segmentation methods achieve medical image segmentation by classifying 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, thus achieving end-to-end segmentation. Among them, U-Net is currently the most widely used framework in the field of medical image segmentation.

[0004] Existing technology: The U-Net network structure can acquire detailed and contour information of the image in the encoder part; then, through a skip connection stage, the extracted features are passed to the decoder part; finally, the decoder part combines features from multiple scales to perform feature recovery. Due to its U-shaped structure, U-Net can train a model with good performance using fewer images. The U-Net network can be divided into a feature extraction network and a feature fusion network. The feature extraction network uses convolutional layers and pooling layers to perform downsampling operations, while the feature fusion network performs upsampling operations, which can restore the image resolution while the network gradually converges to the target region. In the feature fusion stage, features extracted at the same level are fused again to avoid loss of details.

[0005] While U-Net-based medical image segmentation methods have achieved remarkable results, obtaining accurate segmentation results remains challenging due to noise issues. Most methods still suffer from problems such as blurred edges, neglected details, and the need for manual parameter tuning. To address these issues, we introduce a U-Net-based automatic medical image contour extraction network that incorporates an attention mechanism. Summary of the Invention

[0006] The purpose of this invention is to provide a U-Net medical image contour automatic extraction network that incorporates an attention mechanism to solve the problems mentioned in the background art.

[0007] To achieve the above objectives, the present invention provides the following technical solution: a U-Net medical image contour automatic extraction network integrating an attention mechanism, comprising an RGB image input module, wherein the output end of the RGB image input module is connected to the input end of a feature extraction module, and the feature extraction module comprises a feature encoding module, a feature decoding module, and an attention module;

[0008] The output of the feature extraction module is connected to the input of the MLP, and the attention module includes spatial attention and channel attention, which are used to suppress neurons in non-interested regions.

[0009] The MLP is used to classify the extracted features. Its output element is set to two neurons, which represent the probabilities of the foreground and background, respectively, and then connected to Softmax and Marching Square in sequence.

[0010] The RGB image input module is used to input an RGB image.

[0011] The feature extraction module is used to extract features from the RGB image. After obtaining the features of the RGB image, the feature extraction module has a C-dimensional feature representation for each pixel. By fusing local and global information, only one MLP inference is needed for the C-dimensional features of each pixel. Since the task at this stage is binary classification, two-dimensional information is finally obtained, representing the probabilities of the target and the non-target respectively. By comparing the target and the non-target, a binary image can be obtained. Finally, a binary image contour extraction algorithm is used.

[0012] The attention module is used to remove interference information from RGB images.

[0013] The feature encoding module uses ResNet18.

[0014] The channel attention is expressed as follows:

[0015] M c (F)=F*Sigmoid(MLP(AvgPool(F))+MLP(MaxPool(F)));

[0016] The spatial attention is calculated as shown in the following formula, where This indicates that A and B will be joined together by channel;

[0017]

[0018] Combining channel attention and spatial attention, the calculation formula for CBAM is obtained:

[0019] M(F)=M s (M c (G(F)))+F; where: G(F)=Conv2(Conv1(F)).

[0020] The MLP employs a multilayer perceptron with three hidden layers.

[0021] Compared with the prior art, the beneficial effects of the present invention are: by fusing the attention module in a certain way, the present invention improves the edge contour extraction accuracy, initially solves the problem of blurred edges generated by traditional frames, and reduces the interference of background noise, thereby basically meeting the accuracy requirements of medical image contour extraction in the medical field.

[0022] This invention simplifies the process of the traditional framework, making the time required for training and inference relatively short, and greatly saving the time and cost of obtaining the target model.

[0023] This invention uses the Marching Square algorithm to extract the final contour. The algorithm is simple and fast to implement and can be processed in parallel. Attached Figure Description

[0024] Figure 1 This is a schematic diagram of the framework structure of the U-Net medical image contour automatic extraction network that integrates the attention mechanism of the present invention;

[0025] Figure 2 This is a schematic diagram of the Backbone framework structure;

[0026] Figure 3 This is a schematic diagram of the MLP framework structure;

[0027] Figure 4 This is a diagram illustrating the MLP invocation process;

[0028] Figure 5 This is a schematic diagram of the first basic case of Marching Squares;

[0029] Figure 6 A schematic diagram of the second basic case of Marching Squares;

[0030] Figure 7 A diagram illustrating the third basic case of Marching Squares;

[0031] Figure 8 This is a diagram illustrating the fourth basic case of Marching Squares;

[0032] Figure 9 This is a diagram illustrating the fifth basic case of Marching Squares;

[0033] Figure 10 This is a diagram illustrating the sixth basic case of Marching Squares. Detailed Implementation

[0034] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0035] Please see Figure 1-10 The present invention provides a technical solution: a U-Net medical image contour automatic extraction network that integrates an attention mechanism, including an RGB image input module, the output end of which is connected to the input end of a feature extraction module, the feature extraction module including a feature encoding module, a feature decoding module and an attention module;

[0036] The output of the feature extraction module is connected to the input of the MLP, and the attention module includes spatial attention and channel attention, which are used to suppress neurons in non-interested regions.

[0037] The MLP is used to classify the extracted features. Its output element is set to two neurons, which represent the probabilities of the foreground and background, respectively, and then connected to Softmax and Marching Square in sequence.

[0038] Given an RGB image, we first extract its features. The feature extraction module is similar to U-Net, consisting of an encoding network (feature encoding module) and a decoding network (feature decoding module), and also incorporates an attention module. Subsequent ablation experiments showed that the added attention module effectively removes interfering information from the image.

[0039] After obtaining the features, each pixel has a C-dimensional feature representation, fusing local and global information. Therefore, only one MLP inference is needed for the C-dimensional features of each pixel. Since the task at this stage is binary classification, we finally obtain 2-dimensional information, representing the probabilities of the target and non-target. By comparing the target and non-target, a binary image is obtained. Finally, a contour extraction algorithm for the binarized image is used. In the experiment, the Marching Squares algorithm was used to extract the contour.

[0040] The feature extraction module is divided into two parts: a feature encoding module and a feature decoding module. The feature encoding module uses ResNet18, which is easy to train, easy to implement, and has a relatively small number of parameters. More importantly, it has a downsampling structure, making it very suitable for fast feature extraction.

[0041] 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 an attention mechanism, including spatial attention and channel attention, used to suppress neurons in non-interested regions.

[0042] Channel attention is expressed as follows:

[0043] M c (F)=F*Sigmoid(MLP(AvgPool(F))+MLP(MaxPool(F)));

[0044] Spatial attention is calculated as shown in the following formula, where This indicates that A and B will be joined together by channel;

[0045]

[0046] By combining channel attention and spatial attention, the calculation formula for CBAM is obtained.

[0047] M(F)=M s (M c (G(F)))+F; where: G(F)=Conv2(Conv1(F));

[0048] Feature encoding often yields image details, including both target and non-target details. If attention processing is performed at this point to suppress the features of non-target parts in advance, the predicted target region is more likely to be considered as 1 by subsequent MLP modules, which is beneficial for later MLP inference.

[0049] The feature extraction module is similar to U-Net, but unlike U-Net, its output is not a probability map, but rather a description of the features. Assuming the input image is W*H*3, the extracted dimension is W*H*C, where C is the number of feature descriptions. Experiments showed that setting C=512 makes the network more robust. Each pixel has a C-dimensional feature description. These features not only describe the information around the pixel but also incorporate global information; therefore, subsequent judgments only require inference from the vector of this pixel.

[0050] The multilayer perceptron (MLP) designed in this invention is used to classify extracted features. Unlike PIFU, which samples and reclassifies the entire feature set, this invention does not sample features because it deals with two-dimensional image information, requiring minimal computation. Furthermore, current devices have significantly improved computing power and memory capacity, allowing training on all features without sampling. For each pixel, a C-dimensional vector describes its global feature information.

[0051] The framework of this invention has three hidden layers, namely [512, 256, 128]. Since the purpose of image segmentation is binary classification, the output element is set to two neurons, representing the probabilities of the foreground and background respectively, followed by Softmax and Marching Squares.

[0052] Finally, by comparing the probabilities of the foreground and background, if the probability of the foreground is greater than that of the background, the pixel is set to 1. Through the above process, a clean binary image can be obtained.

[0053] The benefits of doing this:

[0054] (1) Improve accuracy. Previous image segmentation methods yielded a probability map, which calculated the probabilities of both the foreground and background. Ablation experiments showed that introducing MLP can improve accuracy by about 0.3%.

[0055] (2) No need for manually setting a threshold. Traditional image contour extraction often requires a manually set threshold on the probability map to obtain a clean binary image. This framework eliminates the need to set this hyperparameter, avoiding the influence of manually set thresholds on the results.

[0056] After obtaining a clean binary image, theoretically any binary image contour extraction algorithm can extract the contour. This 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, it uses linear interpolation based on a threshold to obtain the curve containing the threshold.

[0057] Since the output of this invention is a binary graph, the Marching Squares algorithm yields the same result as long as the threshold is greater than 0. The Marching Squares algorithm is used because it is simple to implement and computationally inexpensive. For four points located within a cell, there are six basic cases, which can be further refined to 16 cases through rotation and mirroring. Only the case needs to be determined to construct the edge. This algorithm can be run in parallel, meaning it can be further optimized and accelerated.

[0058] Key points:

[0059] 1. Based on the UNet framework, this invention introduces an attention module and uses a fully convolutional network to obtain a binary mask of objects in medical images, segmenting the objects from the background image, improving the accuracy of contour extraction, and providing better features for the contour extraction stage.

[0060] 2. This invention uses a multilayer perceptron with 3 hidden layers as a classifier to perform a second correction based on the mask information from the previous step, making the result more accurate and the contour boundaries clearer, which is beneficial for extracting contours using the MarchingSquare algorithm.

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

[0062] Protection point:

[0063] 1. The invention is based on the traditional U-Net, and introduces an attention module in the upsampling stage of U-Net. It uses a fully convolutional network to obtain a binary mask of the object in the medical image and segment the object from the background image. This should be within the protection scope of the present invention.

[0064] 2. This invention uses a multilayer perceptron with three hidden layers as a classifier to perform a second correction on the binary mask, making the result more accurate and the contour boundaries clearer, thereby better meeting the requirements of the medical field for image contour extraction. This should be within the scope of protection of this invention.

[0065] 3. This invention uses Marching Square as the final contour extraction algorithm. Compared to neural network methods, this method improves the performance of contour extraction and can accelerate the process in parallel. The application of the Marching Square algorithm in medical contour extraction should be within the scope of protection of this invention.

[0066] Existing technology: The U-Net network structure can acquire detailed and contour information of the image in the encoder part; then, through a skip connection stage, the extracted features are passed to the decoder part; finally, the decoder part combines features from multiple scales to perform feature recovery. Due to its U-shaped structure, U-Net can train a model with good performance using fewer images. The U-Net network can be divided into a feature extraction network and a feature fusion network. The feature extraction network uses convolutional layers and pooling layers to perform downsampling operations, while the feature fusion network performs upsampling operations, which can restore the image resolution while the network gradually converges to the target region. In the feature fusion stage, features extracted at the same level are fused again to avoid loss of details.

[0067] Although U-Net-based medical image segmentation methods have achieved remarkable results, obtaining accurate segmentation results remains very difficult due to noise issues. Most methods still suffer from problems such as blurred edges, ignored details, and the need for manual parameter tuning.

[0068] Disadvantages of existing technology:

[0069] The limitations of existing medical image contour extraction networks:

[0070] 1. The detection accuracy is relatively low, and it is easy to have false detections or missed detections.

[0071] 2. It requires a large dataset for training, which consumes a lot of cost and time.

[0072] 3. The output results have blurry edges and details are ignored.

[0073] 4. Traditional networks have one or more hyperparameters that need to be manually adjusted, and the adjustment of hyperparameters has a direct impact on the results.

[0074] The reasons for the above-mentioned shortcomings are as follows:

[0075] 1. The model does not focus on certain important features and is easily affected by noise.

[0076] 2. The model itself is quite complex, which results in a lot of time being spent on training.

[0077] 3. Convolution operations on tensors in CNNs can easily cause details to be ignored, making the results tend to be in a stable state between certainty and uncertainty.

[0078] 4. Determining the contour edge requires manually setting a threshold, usually 0.5, but this value is not optimal. There are one or more different optimal values ​​for different data.

[0079] To address issues such as blurred edges and neglected details, and to improve the accuracy and training efficiency of contour extraction while reducing the influence of human intervention on hyperparameters, this invention focuses on a pixel-based network implemented through an attention mechanism that eliminates the need for analog-to-digital (ACMs). This attention mechanism enables faster convergence and achieves good training results even with limited data samples. Using this network, the probability of a pixel being located within an interior region is inferred, and through certain steps, a binary image of the object can be directly output without manually setting thresholds. The main improvements of this invention are as follows:

[0080] (1) To address the issue of interference, this paper proposes a U-Net module with an attention mechanism, which can effectively remove interference and improve the recall rate by 6.3% in the Drosophila embryo dataset compared to U-Net.

[0081] (2) The probability that a pixel is located within the target is represented using an MLP, and the features extracted by U-Net are recalculated instead of being obtained directly from U-Net. This improves accuracy by 0.3% compared to U-Net on the Drosophila embryo dataset.

[0082] (3) The value of manually set threshold has a great impact on the results. To avoid manually setting threshold, one-hot encoding is used to solve this problem, which is simpler and more effective than adaptive threshold setting.

[0083] This invention improves the accuracy of edge contour extraction by fusing attention modules in a certain way, initially solves the problem of blurred edges generated by traditional frames, and reduces the interference of background noise, thereby basically meeting the accuracy requirements of medical image contour extraction in the medical field.

[0084] This invention simplifies the process of the traditional framework, making the time required for training and inference relatively short, and greatly saving the time and cost of obtaining the target model.

[0085] This invention uses the Marching Square algorithm to extract the final contour. The algorithm is simple and fast to implement and can be processed in parallel.

[0086] The technical problem solved by this invention:

[0087] 1. Add an attention module and use the CBAM module to focus the network on the area of ​​interest and remove interference.

[0088] 2. Use one-hot encoding to solve the problem of manually setting thresholds.

[0089] 3. Improve and optimize existing models to shorten training time and make the models as lightweight as possible.

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

Claims

1. A U-Net medical image contour automatic extraction network incorporating an attention mechanism, comprising an RGB image input module, characterized in that: The output of the RGB image input module is connected to the input of the feature extraction module, which includes a feature encoding module, a feature decoding module, and an attention module. The output of the feature extraction module is connected to the input of the multilayer perceptron (MLP), and the attention module includes spatial attention and channel attention, which are used to suppress neurons in non-interested regions. The MLP is used to classify the extracted features. Its output element is set to two neurons, which represent the probabilities of the foreground and background respectively, and then connected to Softmax and Marching Square in sequence. The feature extraction module is used to extract features from the RGB image. After obtaining the features of the RGB image, the feature extraction module has a C-dimensional feature representation for each pixel. This C-dimensional feature representation integrates local and global information. An MLP inference is performed on the C-dimensional features of each pixel to obtain 2-dimensional information, which represents the probabilities of the foreground and background respectively. By comparing the foreground and background, a binary image is obtained. The contour extraction of the binary image is performed using a binary image contour extraction algorithm.

2. The U-Net medical image contour extraction network with an attention-integrated mechanism according to claim 1, characterized in that: The RGB image input module is used to input an RGB image.

3. The U-Net medical image contour extraction network with an attention mechanism as described in claim 1, characterized in that: The attention module is used to remove interference information from RGB images.

4. The U-Net medical image contour extraction network with an attention mechanism as described in claim 1, characterized in that: The feature encoding module uses ResNet18.

5. The U-Net medical image contour extraction network with an integrated attention mechanism as described in claim 1, characterized in that: The channel attention is expressed as follows: ; The spatial attention is calculated as shown in the following formula, where A⊕B means concatenating A and B along the channel; ; Combining channel attention and spatial attention, the calculation formula for CBAM is obtained: ;in: , Features of RGB images.

6. The U-Net medical image contour automatic extraction network with an attention mechanism as described in claim 1, characterized in that: The MLP employs a multilayer perceptron with three hidden layers.

Images