Pathological image segmentation method of U-Net model based on improved encoder and decoder

Abstract

The invention discloses a pathological image segmentation method based on a U-Net model of an improved encoder and decoder, and the method comprises the following steps: carrying out the preprocessing of data, and improving the generalization capability of a model; a U-Net model encoder and a U-Net model decoder are improved, and an EEDU-Net model is established; training the model by using the training set data; optimal model parameters are obtained, test set data are input, and a prediction result of the model is subjected to visualization processing. According to the invention, each layer of encoder is changed into a multi-branch form, so that the capability of encoding to extract information of different scales is improved. And meanwhile, semantic differences are eliminated through channel attention and axial attention processing on the feature map subjected to jump connection, and the effectiveness of each jump connection is ensured. The robustness and precision of pathological image segmentation are improved by using the EEDU-Net model, and a segmentation effect superior to that of an existing model is shown.

Description

Technical Field

[0001] The present invention relates to the technical field of medical image segmentation, and in particular to a pathological image segmentation method based on a U-Net model with improved encoders and decoders. Background Art

[0002] In recent years, with the rapid development of deep learning in the field of images, deep learning has also made great breakthroughs in pathological image segmentation. Pathological image segmentation can accurately separate the diseased area from the normal area on the image to help pathologists diagnose the disease more accurately. Pathological image segmentation tasks often use U-Net and its derivative models. The reason why U-Net has become the cornerstone of this field is mainly due to its network model design. The U-Net model combines the encoder and decoder of the same level through jump connections to help restore the original information after the image is upsampled. However, the single encoding form in U-Net cannot fully extract features, resulting in different degrees of loss of context information in the feature map of each layer. Jump connections between the same levels will have semantic differences.

[0003] In view of the U-Net framework, many models have improved its encoder, jump connection, and decoder, such as Res-UNet, U-Net++, and U2-Net. These CNN-based models have continuously improved the segmentation accuracy, proving that CNN has strong learning ability. Although CNN-based models have made great progress in the field of pathological image segmentation, traditional CNN lacks the ability to model long-range dependencies in images. In other words, in traditional CNN, the receptive field of the model is limited by the size of the convolution kernel, which causes the model to only focus on a local area in the image rather than the global context. To address this defect, image pyramids and hole convolutions are introduced into the model to expand the receptive field and add the ability to model long-range dependencies to traditional CNN. However, this series of operations cannot completely solve this problem, so it still cannot meet the high-precision requirements required for medical applications.

[0004] Subsequently, influenced by the great success of Transformer in natural language processing, Transformer was also applied to the visual field and achieved satisfactory results in image classification, proving that Transformer is also applicable to the visual field. Through the unique multi-head self-attention mechanism in Transformer, the defect of traditional CNN that cannot effectively model long-range dependencies is effectively made up. Therefore, a large number of subsequent models choose to combine U-Net with Transformer to improve the segmentation ability of the model. Although the segmentation accuracy is improved, the introduction of Transformer will cause a sharp increase in the number of model parameters. The increase in the number of parameters will inevitably lead to an increase in model training time. Not only that, too many parameters may cause overfitting on a limited data set. Summary of the invention

[0005] The purpose of the present invention is to provide a pathological image segmentation method based on a U-Net model with improved encoder and decoder, which improves the robustness and accuracy of pathological image segmentation and shows a segmentation effect that is better than the existing model.

[0006] In order to achieve the above object, the present invention provides the following technical solutions:

[0007] A pathological image segmentation method based on a U-Net model with improved encoder and decoder comprises the following steps:

[0008] Step 1: Preprocess the data;

[0009] Step 2: Improve the encoder and decoder of the U-Net model to establish the EEDU-Net model;

[0010] Step 3: Use the training set data to train the model and obtain the optimal model parameters;

[0011] Step 4: Input the test set data and visualize the model's prediction results.

[0012] Furthermore, the EEDU-Net model described in step 2 is composed of an encoder, a decoder and a skip connection. The encoder hierarchically extracts features of different stages of the image; the decoder is responsible for reconstructing the encoded feature map and gradually restoring it to the original input size through upsampling; the skip connection establishes a direct connection between the encoder and decoder at the same level.

[0013] Furthermore, the encoder includes a dense multi-scale extraction module DMSEM and a downsampling module, and the decoder includes a channel and axial hybrid attention module CAHAM.

[0014] Furthermore, the dense multi-scale extraction module DMSEM is specifically divided into four branches: branch 1 is composed of a 1×1 convolution, branch 2 is composed of maximum pooling and 1×1 convolution, branch 3 is composed of 3×3 LDM and two front and back 1×1 convolutions, branch 4 is composed of 5×5 LDM and two front and back 1×1 convolutions, and the data are spliced ​​in the channel dimension after passing through the four branches.

[0015] Furthermore, the operations of each layer of the LDM are as follows:

[0016] (1) Batch normalization normalizes the features of a batch to a distribution with a mean of 0 and a variance of 1, reducing the change in data distribution and accelerating the convergence of the network during training;

[0017] (2) Rectified Linear Unit introduces nonlinear features, allowing the network to learn more complex models;

[0018] (3) Depthwise Separable Convolution First, each input channel is subjected to a separate convolution operation through depthwise convolution, and then point-by-point convolution is used to realize the linear combination of information from different channels.

[0019] Furthermore, the process of extracting features by the dense multi-scale extraction module DMSEM is as follows: the input feature map first enters the point-by-point convolution of different branches, two of which pass through LDM, and the difference in convolution kernel size in different LDMs represents different learning scales. After different convolutions or pooling of each branch, it is finally spliced ​​in the channel dimension.

[0020] Furthermore, before splicing in the channel dimension, a point-by-point convolution operation is required. When DMSEM is in the shallow layer of the network, the proportion of the branch where the smaller convolution kernel is located is increased, and when DMSEM is in the deep layer of the network, the proportion of the branch where the larger convolution kernel is located is increased. After the feature map passes through DMSEM, the length, width, and number of channels of the feature map will not change, and the number of channels of the output feature map is equal to the sum of the number of channels output by the four branches.

[0021] Furthermore, the specific operation of the downsampling module is as follows: the input data first passes through maximum pooling and average pooling, and then is spliced ​​in the channel dimension, and then the difference between the two is eliminated through 1×1 convolution; at the same time, the input data passes through another convolution, and the data passing through two different convolutions are spliced ​​to obtain the final downsampling result.

[0022] The specific implementation of the channel and axial hybrid attention module CAHAM is:

[0023] (1) The input feature map is subjected to maximum pooling and average pooling to extract feature maps in two different spaces, and then a fully convolutional neural network consisting of point-by-point convolution is used to generate a channel attention map, which is then merged in an element-by-element manner and multiplied element-by-element with the original input;

[0024] (2) After completing the above channel attention operation, the result is input into two branches, one of which first performs the column attention operation and then the row attention operation, while the other branch performs the opposite;

[0025] (3) Finally, the results of the two branches are merged and the number of channels is halved through point-by-point convolution.

[0026] Compared with the prior art, the present invention has the following beneficial effects:

[0027] The present invention first preprocesses the data to lay the foundation for subsequent operations to enhance generalization. The EEDU-Net model is constructed by improving the encoder and decoder of the U-Net model. The encoder adopts a multi-branch form to improve the ability to capture information of different scales, and the decoder combines transposed convolution and attention modules to accurately restore feature maps. The optimal model parameters are obtained through the training set, and finally the test set data is processed with the optimal parameters and visualized. By using this model, the robustness and accuracy of pathological image segmentation are improved, showing a segmentation effect that is better than the existing model.

[0028] Furthermore, the present invention enhances the encoder and decoder of U-Net. In the encoder part, the multi-branch design gives the model the ability of multi-scale learning, ensuring that the model can adapt to objects of different sizes and improving the robustness and generalization ability of the model. In the decoder part, when constructing the Transformer, the feature map is divided into spatial attention and channel attention, and the spatial attention is further refined into row attention and column attention. The mixed use of the two allows the model to maintain efficient global modeling capabilities while reducing the number of parameters. At the same time, the feature map after the jump connection is processed by channel attention and axial attention to eliminate semantic differences and ensure the effectiveness of each jump connection. BRIEF DESCRIPTION OF THE DRAWINGS

[0029] Figure 1 It is a structural diagram of the EEDU-Net model of the present invention;

[0030] Figure 2 Schematic diagram of the dense multi-scale extraction module DMSEM structure of the encoder of the present invention;

[0031] Figure 3 It is a schematic diagram of the structure of a lightweight dense module LDM of the encoder of the present invention;

[0032] Figure 4 Schematic diagram of the downsampling module structure of the encoder of the present invention;

[0033] Figure 5 Schematic diagram of the structure of the channel and axial hybrid attention module CAHAM of the decoder of the present invention;

[0034] Figure 6 To test the segmentation effect of the image of the model of the present invention and other models on the MoNuSeg dataset;

[0035] Figure 7 The image segmentation effect of the visualization part of the model of the present invention and other models is tested on the GlaS dataset. DETAILED DESCRIPTION

[0036] In order to make the above-mentioned objects, features and advantages of the present invention more obvious and easy to understand, the present invention is described in detail below with reference to the accompanying drawings and specific embodiments.

[0037] The pathological image segmentation method based on the U-Net model with improved encoder and decoder described in this embodiment includes the following steps:

[0038] Step 1: Preprocess the data to improve the model’s generalization ability

[0039] Resize the image to 224×224 and perform a horizontal flip with a probability of 0.5. 224×224 is a widely used standard size. This size strikes a good balance between computing resources and model performance. Unifying the image size can facilitate batch processing of the model and improve computing efficiency. At the same time, it also helps to reduce the performance differences of the model on images of different sizes. Horizontal flipping is a common data augmentation technique that can randomly flip images horizontally to increase data diversity and allow the model to learn image features at different angles and directions. At the same time, data augmentation can effectively prevent model overfitting and improve generalization capabilities.

[0040] Step 2: Improve the U-Net model encoder and decoder to build the EEDU-Net model

[0041] The EEDU-Net model is as follows Figure 1 As shown in the figure, it consists of an encoder, a decoder and a skip connection. The encoder hierarchically extracts features of different stages of the image. Specifically, the shallow encoder is responsible for extracting specific features of the image, while the deep encoder is responsible for extracting abstract features of the image. The decoder is responsible for reconstructing the encoded feature map and gradually restoring it to the original input size through upsampling. The skip connection establishes a direct connection between the encoder and decoder at the same level, helping the encoder to better restore spatial details.

[0042] In the encoder, in order to solve the problem of insufficient feature extraction at each layer, it is necessary to enhance the network's ability to extract features of different scales at each layer. To this end, a dense multi-scale extraction module (DMSEM) is proposed based on Google inception. This module can use different branches to extract feature map information of different scales, making the information extracted at each layer as rich as possible. The extraction part of each layer in the encoder uses a dense multi-scale extraction module DMSEM to enhance the network's ability to extract different feature scales at each layer. The DMSEM structure is as follows Figure 2 As shown in the figure, the dense multi-scale extraction module DMSEM is specifically divided into four branches: branch 1 consists of a 1×1 convolution, branch 2 consists of maximum pooling and 1×1 convolution, branch 3 consists of 3×3 LDM and two 1×1 convolutions, and branch 4 consists of 5×5 LDM and two 1×1 convolutions. After passing through the four branches, the data is spliced ​​in the channel dimension.

[0043] LDM (Lightweight Dense Module) Figure 3 As shown, the operations at each layer of LDM are:

[0044] (1) Batch normalization (BN) normalizes the features in batches to a distribution with a mean of 0 and a variance of 1, reducing the variation in data distribution and accelerating the convergence of the network during training;

[0045] (2) Rectified Linear Unit (ReLU) introduces nonlinear features, allowing the network to learn more complex models;

[0046] (3) Depthwise Separable Convolution (DSC) first performs a separate convolution operation on each input channel through depthwise convolution, and then uses point-by-point convolution to achieve linear combination of information from different channels.

[0047] The number of channels output by L1, L2, and L3 in LDM is 32, 64, and 128 respectively. Depthwise separable convolution is used to further reduce the number of parameters to achieve lightweight. This design not only retains the property of residual connection to alleviate the vanishing gradient, but also promotes the flow of information in the network. The input of each layer contains the information of all previous layers, which strengthens feature propagation and reduces the number of model parameters by reusing the feature maps of the previous layers.

[0048] The process of extracting features by the dense multi-scale extraction module DMSEM is as follows: the input feature map first enters the point-by-point convolution of different branches, which not only reduces the dimension and reduces the amount of calculation, but also enables each branch to independently learn the correlation between channels. This design fully decouples the learning of spatial correlation and channel correlation, giving the model stronger learning ability. Two of the branches will pass through LDM, and the difference in the size of the convolution kernel in different LDMs represents different learning scales. After different convolutions or pooling of each branch, the channel dimension is finally spliced. Before splicing, a point-by-point convolution operation is required. The purpose of this step is to change the proportion of different branches when they are finally spliced. When the DMSEM is in the shallow layer of the network, the proportion of the branch where the smaller convolution kernel is located should be increased, and when the DMSEM is in the deep layer of the network, the proportion of the branch where the larger convolution kernel is located should be increased. After the feature map passes through DMSEM, the length, width, and number of channels of the feature map will not change, and the number of channels of the output feature map is equal to the sum of the number of channels output by the four branches.

[0049] The downsampling module is an important part of the encoder. Its structure is as follows: Figure 4 As shown in the figure, this module is responsible for reducing the size of the feature map, retaining important features, and reducing the amount of calculation. The specific operation of the downsampling module in the encoder part is as follows: the input data first passes through the maximum pooling and average pooling, and then is spliced ​​in the channel dimension, and then the difference between the two is eliminated through 1×1 convolution. At the same time, the input data passes through another convolution, and the data that passes through two different convolutions is spliced ​​to obtain the final downsampling result.

[0050] Using only maximum pooling or average pooling will result in different degrees of loss of context information. Therefore, both pooling methods are used in the downsampling module, and an additional branch consisting of convolution is designed. This branch has learning behavior, which makes up for the defect of fixed behavior of pooling operation and enhances the ability of downsampling to extract features. The feature maps after maximum pooling and average pooling are spliced ​​in the channel dimension, and then the difference between the two is eliminated through convolution to merge them. Finally, the fused feature map is added element by element to the other feature map after convolution. The length and width of the feature map after downsampling are reduced by half, and the number of channels is doubled.

[0051] In the decoder, in order to solve the semantic differences caused by jump connections between the same levels, a channel and axial hybrid attention module (CAHAM) is proposed, such as Figure 5 As shown, the decoder uses a channel and axial mixed attention module, which is implemented as follows:

[0052] (1) The input feature map is subjected to maximum pooling and average pooling to extract feature maps in two different spaces, and then a fully convolutional neural network consisting of point-by-point convolution is used to generate a channel attention map, which is then merged in an element-by-element manner and multiplied element-by-element with the original input;

[0053] In particular, since the dimensions of the two tensors are different when multiplying, the broadcast mechanism will be triggered, that is, the low-dimensional tensors will be expanded to the same size for calculation.

[0054] (2) After completing the above channel attention operation, the result is input into two branches, one of which first performs the column attention operation and then the row attention operation, while the other branch performs the opposite;

[0055] (3) Finally, the results of the two branches are merged and the number of channels is halved through point-by-point convolution.

[0056] When building the Transformer, the feature map is divided into spatial attention and channel attention, and the spatial attention is further refined into row attention and column attention. Using only row attention or column attention cannot fuse global information, so each branch uses these two attention mechanisms at the same time. The mixed use of the two allows the model to maintain efficient global modeling capabilities while reducing the number of parameters. The order of using row attention and column attention in different branches is different to enhance the robustness and generalization of the model. Compared with traditional attention, axial attention greatly reduces memory consumption, has lower computational complexity and higher computational efficiency. The length and width of the feature map after CAHAM remain unchanged, but the number of channels is halved.

[0057] Step 3: Use the training set data to train the model and obtain the best model parameters

[0058] The learning rate is set to 0.001. The input image is extracted and downsampled four times by the encoder to obtain a highly abstract feature map. The feature map is restored layer by layer through transposed convolution. The feature map in the recovery process needs to be combined with the feature map of the corresponding layer of the encoder, and through the channel and axial mixed attention module. The feature map of the top layer in the decoder is convolved to obtain the final result. The final result and the true value are passed through the binary cross entropy loss function to obtain the loss. Then the parameters are optimized according to the Adam optimizer. The model parameters are continuously adjusted and the optimal model parameters are saved.

[0059] Step 4: Input the test set data and visualize the model’s prediction results

[0060] First, load the trained model and switch it to evaluation mode. Next, resize the input image to 224×224 to ensure that it is consistent with the training. Then use the model to predict the image and remove the extra dimensions in the prediction result. Then convert the result to a binary image and multiply it by 255 to convert the pixel value range of the binary image from 0 and 1 to an integer range of 0 to 255, which is convenient for storage and visualization. Finally, save the processed binary image.

[0061] Experimental part:

[0062] In order to comprehensively evaluate the performance of the model, we compare it with other models, including U-Net, UNet++, TransUNet, UCTransNet, and Swin-Unet. In the training of EEDU-Net, the batch size is set to 4, the loss function selects the cross entropy loss, the Adam optimizer is used, the learning rate is 0.001, and the epoch is 400. The test results are shown in Table 1.

[0063] Table 1 Comparison with other models on MoNuSeg and GlaS datasets

[0064]

[0065] The experimental results show that the proposed model ranks first in different indicators on both MoNuseg and GlaS datasets. In particular, the IoU indicator in the MoNuSeg dataset is improved by 1.04% compared with TransUNet. Moreover, thanks to the design of CAHAM, EEDU-Net has fewer parameters than other models that introduce Transformer.

[0066] The visualization part tests the segmentation effect of the image on different models, such as Figure 6 , Figure 7 As shown in the figure, the red frame part is the area with poor segmentation effect in the comparison model. It can be seen from the figure that the CNN-based model is too rough on the segmentation boundary, the model with attention mechanism handles it relatively well, and the model of the present invention is more smooth in segmentation.

[0067] To explore the impact of different modules on model performance, ablation experiments were conducted on the MoNuSeg and GlaS datasets. Four versions were compared: the original model, only DMSEM, only CAHAM, and both DMSEM and CAHAM. The results are shown in Table 2, which reflects the improvement of the original model by different modules.

[0068] Table 2 Ablation experiment results of the effectiveness of DMSEM and CAHAM

[0069]

[0070] To prove the design value of multi-branch feature extraction, the original single convolution operation was replaced by DMSEM. Experimental results show that multi-branch feature extraction is an effective way to improve model performance. Compared with a single convolution operation, multi-branch contains convolution and pooling operations with different convolution kernel sizes, which enables the model to capture information of different scales, enrich the representation of features, and help the model better locate and segment.

[0071] In order to avoid the inconsistency between the information transmitted by the encoder and the semantics of the decoder after the jump connection, CAHAM is used to replace the original module. This operation significantly improves the indicators of the model on both datasets. The use of the attention mechanism not only ensures the effectiveness of the jump connection, but also ensures the model's global perception. The separation of channel attention and spatial attention allows the model to focus more on a certain dimension. At the same time, the separated attention mechanism can be optimized independently, making the model more flexible.

Claims

1. A pathological image segmentation method based on a U-Net model with improved encoder and decoder, characterized in that: The following steps are involved: Step 1: Preprocess the data; Step 2: Improve the encoder and decoder of the U-Net model to establish the EEDU-Net model; Step 3: Use the training set data to train the model and obtain the optimal model parameters; Step 4: Input the test set data and visualize the model's prediction results.

2. The pathological image segmentation method based on the U-Net model with improved encoder and decoder according to claim 1, characterized in that: The EEDU-Net model described in step 2 consists of an encoder, a decoder and a skip connection. The encoder hierarchically extracts features at different stages of the image. The decoder is responsible for reconstructing the encoded feature map and gradually restoring it to the original input size through upsampling; the skip connection establishes a direct connection between the encoder and decoder at the same level.

3. The pathological image segmentation method based on the U-Net model of improved encoder and decoder according to claim 2, characterized in that: The encoder includes a dense multi-scale extraction module DMSEM and a downsampling module, and the decoder includes a channel and axial hybrid attention module CAHAM.

4. The pathological image segmentation method based on the U-Net model with improved encoder and decoder according to claim 3, characterized in that: The dense multi-scale extraction module DMSEM is specifically divided into four branches: branch 1 consists of a 1×1 convolution, branch 2 consists of maximum pooling and 1×1 convolution, branch 3 consists of 3×3 LDM and two 1×1 convolutions, and branch 4 consists of 5×5 LDM and two 1×1 convolutions. The data are spliced ​​in the channel dimension after passing through the four branches.

5. The pathological image segmentation method based on the U-Net model with improved encoder and decoder according to claim 4, characterized in that: The operations of each layer of the LDM are as follows: (1) Batch normalization normalizes the features of a batch to a distribution with a mean of 0 and a variance of 1, reducing the change in data distribution and accelerating the convergence of the network during training; (2) Rectified Linear Unit introduces nonlinear features, allowing the network to learn more complex models; (3) Depthwise Separable Convolution First, each input channel is subjected to a separate convolution operation through depthwise convolution, and then point-by-point convolution is used to realize the linear combination of information from different channels.

6. The pathological image segmentation method based on the U-Net model with improved encoder and decoder according to claim 4, characterized in that: The process of extracting features by the dense multi-scale extraction module DMSEM is as follows: the input feature map first enters the point-by-point convolution of different branches, two of which pass through LDM. The difference in the size of the convolution kernel in different LDMs represents the different learning scales. After different convolutions or pooling of each branch, it is finally spliced ​​in the channel dimension.

7. The pathological image segmentation method based on the U-Net model with improved encoder and decoder according to claim 6, characterized in that: Before splicing in the channel dimension, a point-by-point convolution operation is required. When DMSEM is in the shallow layer of the network, the proportion of the branch where the smaller convolution kernel is located is increased, and when DMSEM is in the deep layer of the network, the proportion of the branch where the larger convolution kernel is located is increased. After the feature map passes through DMSEM, the length, width, and number of channels of the feature map will not change, and the number of channels of the output feature map is equal to the sum of the number of channels output by the four branches.

8. The pathological image segmentation method based on the U-Net model with improved encoder and decoder according to claim 2, characterized in that: The specific operation of the downsampling module is as follows: the input data first passes through maximum pooling and average pooling, and then is spliced ​​in the channel dimension, and then the difference between the two is eliminated through 1×1 convolution; at the same time, the input data passes through another convolution, and the data through two different convolutions are spliced ​​to obtain the final downsampling result.

9. The pathological image segmentation method based on the U-Net model with improved encoder and decoder according to claim 2, characterized in that: The specific implementation of the channel and axial hybrid attention module CAHAM is: (1) The input feature map is subjected to maximum pooling and average pooling to extract feature maps in two different spaces, and then a fully convolutional neural network consisting of point-by-point convolution is used to generate a channel attention map, which is then merged in an element-by-element manner and multiplied element-by-element with the original input; (2) After completing the above channel attention operation, the result is input into two branches, one of which first performs the column attention operation and then the row attention operation, while the other branch performs the opposite; (3) Finally, the results of the two branches are merged and the number of channels is halved through point-by-point convolution.

Images