An image segmentation method based on an improved U2-Net network

Abstract

An image segmentation method based on an improved U2-Net network includes the following steps: 1) preprocessing the image; 2) performing adaptive histogram smoothing on the image; 3) data augmentation; 4) obtaining output 1; 5) obtaining output 2; 6) obtaining output 3; 7) obtaining output 4; 8) obtaining output 5; 9) obtaining output 6, a significant probability map 6; 10) obtaining output 7, a significant probability map 5; 11) obtaining output 8, a significant probability map 4; 12) obtaining output 9, a significant probability map 3; 13) obtaining output 10, a significant probability map 2; 14) obtaining a significant probability map 1; 15) mapping significant probability maps 1 to 6 to pixel values ​​between 0 and 1 after passing them through a convolutional layer and a sigmoid activation function, representing the probability that each pixel belongs to the foreground (i.e., an object), thus obtaining the segmentation map. This method reduces the manpower and material resources required for manual processing, reduces computational load, and improves segmentation accuracy.

Description

Technical Field

[0001] This invention relates to image segmentation technology, specifically an image segmentation method based on an improved U2-Net network. Background Technology

[0002] In 2015, Ronneberger et al. proposed the U-Net network, which consists of an encoder, a decoder, and skip connections. U-Net achieved excellent results in medical image segmentation tasks, leading to its widespread application. To meet the practical needs of different tasks, various researchers have proposed a series of improvements to U-Net. For example, Oktay et al. proposed Attention U-Net, which designed an attention gate (AG) mechanism and then replaced the skip connections in U-Net with AG modules, effectively suppressing the network's learning of irrelevant regions while focusing on learning task-relevant regions. Zhou et al. proposed U-Net++, which designed a nested structure and dense skip connections based on U-Net, addressing the network depth requirements of different application scenarios and enabling the decoder to better integrate multi-scale information. Huang et al. proposed U-Net3+, which uses full-scale skip connections to combine semantic information from feature maps of different scales, while learning hierarchical representations from the aggregated feature maps of all scales through deep supervision. In 2020, the U2-Net network proposed by Xuebin Q et al. achieved significant results in tasks such as salient object detection, human image segmentation, and natural scene segmentation, and has become one of the important reference models in the field of image segmentation. However, the above models generally suffer from problems such as inaccurate boundary prediction, high memory consumption, and poor information richness.

[0003] Therefore, our model is based on the U2-Net network and introduces a dense connection module to capture features at different levels, obtain receptive fields at different levels, and improve segmentation accuracy. Simultaneously, a channel attention mechanism module is introduced, assigning different weights to each feature, thereby focusing on important features while suppressing unimportant ones, reducing computational cost. Summary of the Invention

[0004] The purpose of this invention is to address the shortcomings of existing technologies by proposing an image segmentation method based on an improved U2-Net network. This method reduces the human and material resources required for manual processing, decreases computational load, and improves segmentation accuracy.

[0005] The technical solution to achieve the objective of this invention is:

[0006] An image segmentation method based on an improved U2-Net network includes the following steps:

[0007] 1) Obtain the dataset and divide it into training and testing sets. Then, preprocess all images in the dataset and use a Gaussian filter to reduce image noise.

[0008] 2) Perform adaptive histogram smoothing on the image from step 1) to increase the contrast between the multiple malformation regions and other regions, so as to reduce the influence of irrelevant regions on the segmentation of the target region.

[0009] 3) Data augmentation: Perform angle rotation, inversion, and rotation operations on the image from step 2). The angle rotation operation is set to rotate clockwise and counterclockwise by 30° and 60°, respectively. The inversion operation is set to horizontal flip, and the rotation operation is set to vertical rotation.

[0010] 4) The images in the training set are fed into the encoder En_1 of the RSU-7 layer as input to the model. The convolutional blocks in the RSU-7 layer encoder are dense connection modules 1, with a total of 7 convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to integrate image features and extract image features. The result of the RSU-7 layer encoder is sent into the channel attention module and is denoted as output 1.

[0011] 5) Input output 1 into the encoder En_2 of the RSU-6 layer. The convolutional block of the RSU-6 layer encoder is the dense connection module 2, which has a total of 6 convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to further extract image features and reduce the spatial scale. Input the result of the RSU-6 layer encoder into the channel attention module, which is denoted as output 2.

[0012] 6) Input output 2 into the encoder En_3 of the RSU-5 layer. The convolutional block of the RSU-5 layer encoder is the dense connection module 3, which has 5 layers of convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to further extract image features and reduce the spatial scale. Input the result of the RSU-5 layer encoder into the channel attention module, which is denoted as output 3.

[0013] 7) Input output 3 into the encoder En_4 of the RSU-4 layer. The convolutional block of the RSU-4 layer encoder is a dense connection module 4. There are 4 convolutional blocks in total. The skip connections in the middle adopt a connection mode that combines long and short connections to further extract image features and reduce the spatial scale. Input the result of the RSU-4 layer encoder into the channel attention module, which is denoted as output 4.

[0014] 8) Input output 4 into the encoder En_5 of the RSU-4F layer. Keep the skip connections of the convolutional blocks of the RSU-4F layer unchanged from the original U2-Net network design. Input the encoder result of the RSU-4F layer into the channel attention module and record it as output 5.

[0015] 9) Input output 5 into the encoder En_6 of the RSU-4F layer. The skip connections of the RSU-4F layer convolutional blocks remain unchanged from the original U2-Net network design. Input the RSU-4F layer encoder result into the channel attention module to obtain output 6, with a significant probability graph 6. 10) Input output 6 into the decoder De_5 of the RSU-4F layer. The skip connections of the RSU-4F layer convolutional blocks remain unchanged from the original U2-Net network design. Input the RSU-4F layer decoder result into the channel attention module to obtain output 7, with a significant probability graph 6. Figure 5 ;

[0016] 11) Output 7 is fed into the decoder De_4 of the RSU-4 layer. The convolutional blocks of the RSU-4 layer decoder are densely connected modules 4, with a total of 4 convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to extract image features and enlarge the spatial scale. The result of the RSU-4 layer decoder is fed into the channel attention module to obtain the output result 8, which has a significant probability. Figure 4 ;

[0017] 12) Output 8 is fed into the decoder De_3 of the RSU-5 layer. The convolutional blocks of the RSU-5 layer decoder are densely connected modules 3, with a total of 5 convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to further extract image features and enlarge the spatial scale. The result of the RSU-5 layer decoder is fed into the channel attention module to obtain the output result 9, which has a significant probability. Figure 3 ;

[0018] 13) Output 9 is fed into the RSU-6 layer decoder De2. The convolutional blocks of the RSU-6 layer decoder are densely connected modules 2, with a total of 6 convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to further extract image features and enlarge the spatial scale. The result of the RSU-6 layer decoder is fed into the channel attention module to obtain the output result 10, which has a significant probability. Figure 2 ;

[0019] 14) The output 10 is fed into the decoder De1 of the RSU-7 layer. The convolutional blocks of the RSU-7 layer decoder are densely connected modules 1, with a total of 7 convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to further extract image features, enlarge the spatial scale, and obtain significant probabilities. Figure 1 ;

[0020] 15) Significant probability Figure 1 The saliency probability map 6 is mapped to pixel values ​​between 0 and 1 through a convolutional layer and a sigmoid activation function, representing the probability that each pixel belongs to the foreground, i.e., the object, thus obtaining the segmentation map.

[0021] This technical solution includes a U2-Net network model, dense connection modules, and a channel attention mechanism. The dense connection module refers to a nested, dense network structure combining long and short skip connections. This solution redesigns the skip connections at the RSU-7, RSU-6, RSU-5, and RSU-4 levels of the U2-Net network using dense connections, fusing features from different levels to give the network receptive fields of varying sizes. Simultaneously, a channel attention module is added to each layer, enabling the network to automatically learn the importance of each channel, thereby enhancing its ability to focus on specific features.

[0022] This technical solution has the following advantages compared to existing technologies:

[0023] Compared to U2-Net's naive skip connections, the dense connection module proposed in this technical solution can more accurately capture features at different levels and exhibits varying sensitivities to target objects of different sizes. Therefore, through feature overlay, the dense connection module can fuse features from different levels, giving the network receptive fields of varying sizes. The addition of a channel attention mechanism enhances the network's ability to represent different channels, enabling it to automatically learn the importance of each channel and thus improve its ability to focus on specific features. It also improves model performance; by adjusting the feature representation of channels, channel attention helps the model better distinguish and capture information between different features, thereby improving the model's performance in computer vision tasks. Channel attention can also reduce network parameters and computational cost to some extent. By adaptively adjusting channel weights, channel attention makes the model more compact, retaining only channel information useful for the task, reducing redundancy and unnecessary computation, and thus lowering the model's storage and computational costs. Attached Figure Description

[0024] Figure 1 Flowchart for an embodiment;

[0025] Figure 2 This is a segmentation model diagram for an embodiment;

[0026] Figure 3 This is a diagram of the densely connected modules in the embodiment;

[0027] Figure 4 This is a diagram of the channel attention module in the embodiment;

[0028] Figure 5 The image shows the segmentation result in the example. Detailed Implementation

[0029] The present invention will be further described below with reference to the accompanying drawings and embodiments, but this is not intended to limit the scope of the invention.

[0030] Example:

[0031] In this example, the batch size is 4, the initial learning rate is 1e-3, and the learning rate decreases by a factor of 10 every 30 epochs, for a total of 120 epochs. The training loss function is BCELoss(), which calculates the binary cross-entropy loss between the target value and the predicted value. Its calculation method is as follows:

[0032] l n =-w n ·[y n ·logx n +(1-y n )·log(1-x n )]

[0033] Where w n Let x represent the weight matrix. n The predicted value matrix (the result of the input matrix being processed by the activation function) represents the matrix of predicted values, y. n This represents the target value matrix. Here, log is base e, which is ln in mathematics.

[0034] Reference Figure 1 An image segmentation method based on an improved U2-Net network is presented here. This example method is a segmentation algorithm for multifidus muscle images, which can output radiological suggestion text about the input multifidus muscle MRI image, including the following steps:

[0035] We collected MRI images of the multifidus muscle from patients with lumbar disc herniation as a dataset, and divided the images into training and testing sets in a 4:1 ratio.

[0036] 1) Preprocess all images in the dataset and use a Gaussian filter to reduce image noise;

[0037] 2) Perform adaptive histogram smoothing on the image from step 1) to increase the contrast between the multiple malformation regions and other regions, so as to reduce the influence of irrelevant regions on the segmentation of the target region.

[0038] 3) Data augmentation: Perform angle rotation, inversion, and rotation operations on the image from step 2). The angle rotation operation is set to rotate clockwise and counterclockwise by 30° and 60°, respectively. The inversion operation is set to horizontal flip, and the rotation operation is set to vertical rotation.

[0039] 4) The images in the training set from step 3) are fed as input to the encoder En_1 of the RSU-7 layer, and the model is as follows: Figure 2 As shown, the convolutional blocks in the RSU-7 layer encoder are densely connected modules 1, with a total of 7 convolutional blocks, as follows: Figure 3 As shown, the intermediate skip connections employ a combination of long and short connections to integrate image features, thereby extracting image features. The results from the RSU-7 layer encoder are then fed into the channel attention module, as follows: Figure 4 As shown, this is denoted as output 1;

[0040] 5) Input output 1 into the RSU-6 layer encoder En_2. The convolutional blocks of the RSU-6 layer encoder are densely connected modules 2, with a total of 6 convolutional blocks, such as... Figure 3 As shown, the middle skip connection adopts a connection mode that combines long and short connections to further extract image features and reduce the spatial scale. The result of the RSU-6 layer encoder is sent to the channel attention module and is denoted as output 2.

[0041] 6) Input output 2 into the encoder En_3 of the RSU-5 layer. The convolutional blocks of the RSU-5 layer encoder are densely connected modules 3, with a total of 5 convolutional blocks, such as... Figure 3 As shown, the middle skip connection adopts a connection mode that combines long and short connections to further extract image features, reduce the spatial scale, and send the RSU-5 layer encoder result into the channel attention module, which is denoted as output 3.

[0042] 7) Output 3 is fed into the encoder En_4 of the RSU-4 layer. The convolutional blocks of the RSU-4 layer encoder are densely connected modules 4, with a total of 4 convolutional blocks, such as... Figure 3 As shown, the middle skip connection adopts a connection mode that combines long and short connections to further extract image features, reduce the spatial scale, and send the RSU-4 layer encoder result into the channel attention module, which is denoted as output 4.

[0043] 8) Input output 4 into the encoder En_5 of the RSU-4F layer. Keep the skip connections of the convolutional blocks of the RSU-4F layer unchanged from the original U2-Net network design. Input the encoder result of the RSU-4F layer into the channel attention module and record it as output 5.

[0044] 9) Input output 5 into the encoder En_6 of the RSU-4F layer. The skip connections of the RSU-4F layer convolutional blocks remain unchanged from the original U2-Net network design. Input the RSU-4F layer encoder result into the channel attention module to obtain output 6, with a significant probability graph 6. 10) Input output 6 into the decoder De_5 of the RSU-4F layer. The skip connections of the RSU-4F layer convolutional blocks remain unchanged from the original U2-Net network design. Input the RSU-4F layer decoder result into the channel attention module to obtain output 7, with a significant probability graph 6. Figure 5 ;

[0045] 11) Output 7 is fed into the decoder De_4 of the RSU-4 layer. The convolutional blocks of the RSU-4 layer decoder are densely connected modules 4, with a total of 4 convolutional blocks, such as... Figure 3 As shown, the intermediate skip connections employ a combination of long and short connections to extract image features, enlarge the spatial scale, and feed the RSU-4 layer decoder results into the channel attention module, yielding output result 8, with a significant probability. Figure 4 ;

[0046] 12) Output 8 is fed into the decoder De_3 of the RSU-5 layer. The convolutional blocks of the RSU-5 layer decoder are densely connected modules 3, with a total of 5 convolutional blocks, such as... Figure 3 As shown, the intermediate skip connections employ a combination of long and short connections to further extract image features and amplify the spatial scale. The RSU-5 layer decoder result is then fed into the channel attention module, yielding output result 9, with a significant probability. Figure 3 ;

[0047] 13) Output 9 is fed into the RSU-6 layer decoder De_2. The convolutional blocks of the RSU-6 layer decoder are densely connected modules 2, with a total of 6 convolutional blocks, such as... Figure 3 As shown, the intermediate skip connections employ a combination of long and short connections to further extract image features and amplify the spatial scale. The RSU-6 layer decoder results are then fed into the channel attention module, yielding an output result of 10, with a significant probability. Figure 2 ;

[0048] 14) The output 10 is fed into the decoder De_1 of the RSU-7 layer. The convolutional blocks of the RSU-7 layer decoder are densely connected modules 1, with a total of 7 convolutional blocks, such as... Figure 3 As shown, the intermediate skip connections employ a combination of long and short connections to further extract image features, magnify the spatial scale, and obtain significant probabilities. Figure 1 ;

[0049] 15) Significant probability Figure 1 The saliency probability map in Figure 6 is mapped to pixel values ​​between 0 and 1 using a convolutional layer and a sigmoid activation function, representing the probability that each pixel belongs to the foreground (i.e., the object), resulting in a segmentation map, as shown below. Figure 5 As shown.

[0050] Compared to traditional segmentation methods Unet and Unet++, the segmentation results achieved by the method in this example are closer to the correct segmentation map (Ground-Truth), thus achieving good segmentation results.

Claims

1. An image segmentation method based on an improved U2-Net network, characterized in that, Includes the following steps: 1) Obtain the dataset and divide it into training and testing sets. Then, preprocess all images in the dataset and use a Gaussian filter to reduce image noise. 2) Perform adaptive histogram smoothing on the image from step 1) to increase the contrast between the multiple malformation regions and other regions, thereby reducing the impact of irrelevant regions on the segmentation of the target region; 3) Data augmentation: Perform angle rotation, inversion, and rotation operations on the image from step 2). The angle rotation operation is set to rotate clockwise and counterclockwise by 30° and 60°, respectively. The inversion operation is set to horizontal flip, and the rotation operation is set to vertical rotation. 4) The images in the training set are fed into the encoder En_1 of the RSU-7 layer as input to the model. The convolutional blocks in the RSU-7 layer encoder are dense connection modules 1, with a total of 7 convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to integrate image features and extract image features. The result of the RSU-7 layer encoder is sent into the channel attention module and is recorded as output 1. 5) Input output 1 into the encoder En_2 of the RSU-6 layer. The convolutional block of the RSU-6 layer encoder is the dense connection module 2, which has a total of 6 convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to further extract image features and reduce the spatial scale. Input the result of the RSU-6 layer encoder into the channel attention module, which is denoted as output 2. 6) Input output 2 into the encoder En_3 of the RSU-5 layer. The convolutional block of the RSU-5 layer encoder is the dense connection module 3, which has 5 layers of convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to further extract image features and reduce the spatial scale. Input the result of the RSU-5 layer encoder into the channel attention module, which is denoted as output 3. 7) Input output 3 into the encoder En_4 of the RSU-4 layer. The convolutional block of the RSU-4 layer encoder is a dense connection module 4. There are 4 convolutional blocks in total. The skip connections in the middle adopt a connection mode that combines long and short connections to further extract image features and reduce the spatial scale. Input the result of the RSU-4 layer encoder into the channel attention module, which is denoted as output 4. 8) Input output 4 into the encoder En_5 of the RSU-4F layer. Keep the skip connections of the convolutional blocks of the RSU-4F layer unchanged from the original U2-Net network design. Input the encoder result of the RSU-4F layer into the channel attention module and record it as output 5. 9) Input output 5 into the encoder En_6 of the RSU-4F layer. The skip connections of the convolutional blocks of the RSU-4F layer remain unchanged from the original design of the U2-Net network. Input the encoder result of the RSU-4F layer into the channel attention module to obtain output 6, and the saliency probability diagram 6. 10) Input output 6 into the decoder De_5 of the RSU-4F layer. The skip connections of the convolutional blocks of the RSU-4F layer remain unchanged from the original design of the U2-Net network. Input the decoder result of the RSU-4F layer into the channel attention module to obtain output 7, and the saliency probability diagram 5. 11) The output 7 is fed into the decoder De_4 of the RSU-4 layer. The convolutional block of the RSU-4 layer decoder is a dense connection module 4, with a total of 4 convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to extract image features and enlarge the spatial scale. The result of the RSU-4 layer decoder is fed into the channel attention module to obtain the output result 8, the saliency probability map 4. 12) The output 8 is fed into the decoder De_3 of the RSU-5 layer. The convolutional block of the RSU-5 layer decoder is a dense connection module 3, with a total of 5 convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to further extract image features and enlarge the spatial scale. The result of the RSU-5 layer decoder is fed into the channel attention module to obtain the output result 9, and the saliency probability map 3. 13) The output 9 is fed into the decoder De_2 of the RSU-6 layer. The convolutional block of the RSU-6 layer decoder is a dense connection module 2, with a total of 6 convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to further extract image features and enlarge the spatial scale. The result of the RSU-6 layer decoder is fed into the channel attention module to obtain the output result 10, and the saliency probability diagram 2. 14) The output 10 is fed into the decoder De_1 of the RSU-7 layer. The convolutional block of the RSU-7 layer decoder is a dense connection module 1, with a total of 7 convolutional blocks. The skip connections in the middle adopt a connection mode that combines long and short connections to further extract image features and enlarge the spatial scale, thus obtaining the saliency probability map 1. 15) After passing the saliency probability maps 1 to 6 through a convolutional layer and a sigmoid activation function, they are mapped to pixel values ​​between 0 and 1, representing the probability that each pixel belongs to the foreground or object, thus obtaining the segmentation map.

Images