Remote sensing image semantic segmentation method based on U-Net

Abstract

The invention provides a remote sensing image semantic segmentation method based on U-Net, and belongs to the technical field of image segmentation, and the method comprises the steps: enabling an input image to sequentially pass through an encoder part composed of four processing modules after the input image is obtained, thereby obtaining feature maps of four different dimensions; and inputting the feature map with the highest channel number and the smallest size output by the last module into the bottleneck part of the network, performing feature enhancement, entering the decoder part of the network, and gradually recovering the spatial resolution of the image through up-sampling. In the recovery process, the output feature map of the encoder is connected to the decoder by utilizing jump connection, so that the decoder is helped to better recover the details of the image. And finally, the output of the decoder passes through a 1 * 1 convolutional layer to generate an output image with the same size as the input image. According to the method, the U-shaped network is improved by using the three new modules, the segmentation precision is improved, and the high efficiency and accuracy of the method in remote sensing image semantic segmentation are ensured.

Description

Technical Field

[0001] The present invention relates to the field of remote sensing image processing, and in particular to a remote sensing image semantic segmentation method based on U-Net. Background Art

[0002] With the continuous advancement of remote sensing technology, remote sensing images have been widely used in environmental monitoring, urban planning, agricultural detection, disaster assessment and other fields. Remote sensing images have the characteristics of wide coverage and flexible acquisition methods, and can provide important spatial information for various applications. However, remote sensing images often contain a large amount of background noise, complex landform types and irregular object shapes, which makes it difficult for traditional image processing methods to effectively extract detailed features of the target area. Especially in complex environments, the diversity and scale changes of objects in remote sensing images make image segmentation tasks very challenging.

[0003] As an important task in the field of image processing, semantic segmentation aims to assign each pixel in an image to a specific category. Traditional semantic segmentation methods mostly rely on manual feature extraction and classifiers, but these methods are often limited in effect when processing high-resolution remote sensing images, especially in complex backgrounds and fine-grained classification. In recent years, semantic segmentation methods based on deep learning have gradually become mainstream, among which the U-Net network has become a commonly used architecture in semantic segmentation of remote sensing images due to its excellent performance and low computing resource requirements. U-Net can effectively improve segmentation accuracy by extracting multi-scale features from the input image and fusing feature information at different levels through jump connections, especially in detail preservation and small target detection.

[0004] However, although U-Net has achieved good results in many remote sensing image segmentation tasks, in some specific scenarios, such as small objects or complex background areas in high-resolution images, the segmentation accuracy of U-Net still has room for improvement. Therefore, researchers are exploring ways to improve the accuracy and efficiency of remote sensing image semantic segmentation by improving the U-Net architecture, introducing new feature fusion mechanisms, and optimizing training strategies to better meet the needs of practical applications. Summary of the invention

[0005] In order to solve the above technical problems, the present invention provides a remote sensing image semantic segmentation method based on U-Net.

[0006] The present invention provides a remote sensing image semantic segmentation method based on U-Net, comprising the following steps:

[0007] S1. Improved U-Net model trained by Vaihingen and Potsdam public datasets

[0008] S2, input the three-channel RGB remote sensing image into the encoder part of the network;

[0009] S3, through the Shunted Transformer Block in the encoder, using its diversion self-attention mechanism and convolution operation, extract the low-level and global feature information of the image;

[0010] S4, using skip connections to fuse the low-level features output by the encoder with the high-level features restored by the decoder;

[0011] S5, in the decoder part, the spatial resolution of the image is gradually restored by upsampling, and finally restored to the resolution consistent with the input image;

[0012] S6. Use the segmentation head to process the output of the decoder part and classify each pixel to obtain the category label of each pixel.

[0013] Furthermore, the loss function of the training part in S1 is a composite function that combines the cross entropy loss function, the Dice loss function, and the edge binary cross entropy loss function. It pays special attention to the edge area and promotes the consistency of the model between two different prediction results through KL divergence. The specific calculation method is as follows:

[0014]

[0015] in, and is the main loss calculated for two different predictions, and and is the marginal loss obtained from the two predictions, L KL is the KL divergence loss.

[0016] It should be noted that the input image in S2 needs to undergo a series of data enhancement preprocessing during the training process to improve the generalization ability of the model, including: randomly flipping the image horizontally, randomly adjusting the brightness, contrast, saturation and hue of the image, normalizing the image, randomly cropping the image, and converting the image into a PyTorch tensor.

[0017] Furthermore, the specific process of passing through the Shunted Transformer Block in the encoder in S3 is as follows:

[0018] S31, a three-channel RGB remote sensing image (assuming the size is 3×256×256) is received by the encoder;

[0019] S32, the input image is processed by the stage1 module of the encoder to obtain a feature map of size 64×64×64;

[0020] S33, the feature map of size 64×64×64 is further input into the stage2 module for processing to obtain a feature map of size 128×32×32;

[0021] S34, the feature map of size 128×32×32 is further input into the stage3 module for processing to obtain a feature map of size 256×16×16;

[0022] S35, input the feature map of size 256×16×16 into the last module stage4 of the encoder for processing, and obtain a feature map of size 512×8×8 as the output of the encoder.

[0023] It should be noted that traditional convolutional neural networks directly process raw pixel-level information, while Transformer needs to process sequence data. Therefore, before passing through the Shunted Transformer Block, it needs to go through PatchEmbedding processing. Its purpose is to convert the input image into a set of image blocks of fixed size and use the image blocks as input sequences to meet the requirement that Transformer needs to process sequence data.

[0024] Furthermore, the jump connection in S4 is not limited to the encoder and decoder in the present invention, but also exists inside the decoder. This connection passes deeper features to shallower layers through upsampling, so that shallower features can obtain contextual information from deeper layers, making up for the defect that shallow features are not able to capture global information. This fusion of multi-level features enhances the model's ability to understand input images at various scales.

[0025] It should be noted that this new type of jump connection refers to the residual connection structure, which can alleviate the gradient vanishing problem in the training process without increasing the number of parameters.

[0026] Furthermore, in the upsampling process in S5, the present invention proposes a new fusion method, namely, using learnable weights and PReLU for nonlinear combination, which allows more flexible adjustment of the contribution weights between features. This method can capture more complex feature relationships than simple addition or multiplication fusion methods. In addition, after the feature map is fused, a channel attention mechanism is introduced to enable the network to automatically focus on important features, improve the expressiveness of the fused features, and enhance the model's understanding of multi-scale features.

[0027] Furthermore, the segmentation head output in S6 is a category prediction for each pixel, and the category is a predefined category of different types of objects in the remote sensing image. The predefined object categories in the present invention may be: opaque surface, building, low shrub, tree, vehicle.

[0028] Compared with the related art, the remote sensing image semantic segmentation method based on U-Net provided by the present invention has the following beneficial effects:

[0029] 1. The multi-scale feature fusion module proposed in the present invention can effectively fuse image features from different levels, and the segmentation accuracy is significantly improved, especially when processing complex backgrounds and small objects.

[0030] 2. The Edge Loss loss function proposed in the present invention enables the model to focus on the edge information of the image during training, thereby better processing the boundary part of the object in the image.

[0031] 3. The new skip connection method introduced in the present invention can not only effectively avoid information loss, but also help the model overcome the gradient vanishing and overfitting problems in deep networks when processing complex remote sensing images, thereby improving the anti-interference ability of the model. BRIEF DESCRIPTION OF THE DRAWINGS

[0032] Figure 1 This is a structural diagram of the improved U-Net model of the present invention;

[0033] Figure 2 This is the structure diagram of Shunted Transformer Block;

[0034] Figure 3 This is a comparison diagram between the skip connection structure proposed in the present invention and other model structures, where A is TransUNet, B is ST-UNet, C is HST-UNet, and D is the structure proposed in the present invention.

[0035] Figure 4 This is the structure diagram of the multi-scale feature fusion module. DETAILED DESCRIPTION

[0036] In order to make the purpose, technical solutions and advantages of the embodiments of the present invention clearer, the technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are part of the embodiments of the present invention, not all of the embodiments. The components of the embodiments of the present invention described and shown in the drawings here can be arranged and designed in various different configurations.

[0037] Therefore, the following detailed description of the embodiments of the present invention provided in the accompanying drawings is not intended to limit the scope of the invention claimed for protection, but merely represents selected embodiments of the present invention. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0038] The specific implementation of the remote sensing image semantic segmentation method based on U-Net is described in detail below in conjunction with specific embodiments.

[0039] The present invention provides a remote sensing image semantic segmentation method based on U-Net, comprising the following steps:

[0040] S1. The improved U-Net model is trained by the Vaihingen and Potsdam public data sets. The structure of the improved U-Net model of the present invention is as follows: Figure 1 As shown;

[0041] S2, input the three-channel RGB remote sensing image into the encoder part of the network;

[0042] S3, through the Shunted Transformer Block in the encoder, using its diversion self-attention mechanism and convolution operation, extract the low-level and global feature information of the image;

[0043] S4, using skip connections to fuse the low-level features output by the encoder with the high-level features restored by the decoder;

[0044] S5, in the decoder part, the spatial resolution of the image is gradually restored by upsampling, and finally restored to the resolution consistent with the input image;

[0045] S6. Use the segmentation head to process the output of the decoder part and classify each pixel to obtain the category label of each pixel.

[0046] Furthermore, the loss function of the training part in S1 is a composite function that combines the cross entropy loss function, the Dice loss function, and the edge binary cross entropy loss function. It pays special attention to the edge area and promotes the consistency of the model between two different prediction results through KL divergence. The specific calculation method is as follows:

[0047] Calculate L first CE =-∑ i y i log(p i ), where y i represents the true label, p i is the model prediction result. At the same time, y i and pi Calculate Using weight coefficient combination L CE and L Dice You can get L main =α·L CE +β·L Dice On this basis, the Laplacian kernel is introduced and applied to the ground truth labels and model predictions of edge extraction: Edge(x) = ReLU(Conv(x, Laplacian Kernel)), where Conv represents the convolution operation, ReLU represents the linear rectification function, x represents the input image, and the Laplacian kernel is defined as The calculated edge features are then thresholded (i.e., pixels greater than or equal to the threshold are set to 1 and other pixels are set to 0) to obtain a binary edge map. Subsequently, the binary cross entropy loss is used to calculate the difference between the estimated edge and the actual edge: L εdge =BCE(Edge(Prediction), Edge(Target)), where BCE is the binary cross entropy loss, calculated as Where y represents the true label, Represents the model prediction result. In addition, we introduce KL divergence to ensure the consistency between the two outputs of the model. The final expression is Here, t is the main loss calculated for two different predictions, and and is the marginal loss obtained from the two predictions, L KL is the KL divergence loss.

[0048] Furthermore, the specific process of passing through the Shunted Transformer Block in the encoder in S3 is as follows:

[0049] S31, a three-channel RGB remote sensing image (assuming the size is 3×256×256) is received by the encoder;

[0050] S32, the input image is processed by the stage1 module of the encoder to obtain a feature map of size 64×64×64;

[0051] S33, the feature map of size 64×64×64 is further input into the stage2 module for processing to obtain a feature map of size 128×32×32;

[0052] S34, the feature map of size 128×32×32 is further input into the stage3 module for processing to obtain a feature map of size 256×16×16;

[0053] S35, input the feature map of size 256×16×16 into the last module stage4 of the encoder for processing, and obtain a feature map of size 512×8×8 as the output of the encoder.

[0054] It should be noted that Shunted Transformer is a network architecture that combines Transformer and convolution features. Its structure diagram is as follows Figure 2 As shown in the figure, its core advantage is that it optimizes the computational efficiency and enhances the modeling ability of image features based on the traditional Transformer structure through Shunted Attention.

[0055] Furthermore, the skip connection structure in S4 is compared with other model structures as shown in the figure below. Figure 3 As shown in the figure, in the present invention, the jump connection is not limited to the encoder and decoder, but also exists inside the decoder. This connection passes deeper features to shallower layers through upsampling, so that shallower features can obtain contextual information from deeper layers, making up for the defect that shallow features are not able to capture global information. This fusion of multi-level features enhances the model's ability to understand input images at various scales.

[0056] Furthermore, in the upsampling process in S5, the present invention proposes a new fusion method, namely, using learnable weights and PReLU for nonlinear combination, which allows more flexible adjustment of the contribution weights between features. This method can capture more complex feature relationships than simple addition or multiplication fusion methods. In addition, after the feature map is fused, a channel attention mechanism is introduced to enable the network to automatically focus on important features, improve the expressive power of the fused features, and enhance the model's understanding of multi-scale features. The specific upsampling process is:

[0057] S51, input the output feature map (512×8×8) of the bottleneck part of the model into the first fusion module in the decoder, which first upsamples the feature map to 256×16×16 through an upsampling operation, and then fuses the upsampled feature map with the output feature map of the encoder stage3;

[0058] S52, input the fused feature map of size 256×16×16 into the second fusion module in the decoder, and repeat similar operations to obtain a feature map of size 64×64×64;

[0059] S53, upsample the obtained feature map of size 64×64×64 and restore it to a size of 256×256.

[0060] Furthermore, the segmentation head in S6 adjusts the number of channels of the obtained feature map of size 64×256×256 to the number of predicted categories, so as to obtain the category prediction of each pixel. The categories are different types of ground object categories in the predefined remote sensing image. The predefined ground object categories in the present invention are: opaque surface, building, low shrub, tree, and vehicle.

[0061] The convolutional network and Transformer architecture involved in the present invention are all existing technologies and will not be described in detail here.

[0062] The above descriptions are merely embodiments of the present invention and are not intended to limit the patent scope of the present invention. Any equivalent structure or equivalent process transformation made using the contents of the present invention specification and drawings, or directly or indirectly applied in other related technical fields, are also included in the patent protection scope of the present invention.

Claims

1. A remote sensing image semantic segmentation method based on U-Net, characterized in that: The steps include: S1, the improved U-Net model is trained by the public datasets of Vaihingen and Potsdam; S2, input the three-channel RGB remote sensing image into the encoder part of the network; S3, through the Shunted Transformer Block in the encoder, using its diversion self-attention mechanism and convolution operation, extract the low-level and global feature information of the image; S4, using skip connections to fuse the low-level features output by the encoder with the high-level features restored by the decoder; S5, in the decoder part, the spatial resolution of the image is gradually restored by upsampling, and finally restored to the resolution consistent with the input image; S6. Use the segmentation head to process the output of the decoder part and classify each pixel to obtain the category label of each pixel.

2. The remote sensing image semantic segmentation method based on U-Net according to claim 1, characterized in that: The loss function of the training part in S1 is a composite function that combines the cross entropy loss function, the Dice loss function, and the edge binary cross entropy loss function. At the same time, the KL divergence is used to promote the consistency of the model between two different prediction results. The specific calculation method is as follows: in, and is the main loss calculated for two different predictions, and and is the marginal loss obtained from the two predictions, L KL is the KL divergence loss.

3. The remote sensing image semantic segmentation method based on U-Net according to claim 1, characterized in that: The input image in S2 will undergo data augmentation preprocessing during the training process, including randomly flipping the image horizontally, randomly adjusting one or more of the brightness, contrast, saturation and hue of the image, normalizing the image, randomly cropping the image, and converting the image into a PyTorch tensor.

4. The remote sensing image semantic segmentation method based on U-Net according to claim 1, characterized in that: The specific process of passing through the ShuntedTransformerBlock in the encoder in S3 is: S31, receiving a three-channel RGB remote sensing image by an encoder; S32, the input image is processed by the stage1 module of the encoder to obtain a feature map of size 64×64×64; S33, the feature map of size 64×64×64 is further input into the stage2 module for processing to obtain a feature map of size 128×32×32; S34, the feature map of size 128×32×32 is further input into the stage3 module for processing to obtain a feature map of size 256×16×16; S35, input the feature map of size 256×16×16 into the last module stage4 of the encoder for processing, and obtain a feature map of size 512×8×8 as the output of the encoder.

5. The remote sensing image semantic segmentation method based on U-Net according to claim 1, characterized in that: The jump connection in S4 is not limited to the encoder and decoder, but also exists inside the decoder; by upsampling, deeper features are passed to shallower layers, so that shallower features can obtain contextual information from deeper layers.

6. The remote sensing image semantic segmentation method based on U-Net according to claim 1, characterized in that: In the upsampling process in S5, learnable weights and PReLU are used for nonlinear combination to adjust the contribution weights between features and capture more complex feature relationships than simple addition or multiplication fusion methods. After the feature maps are fused, a channel attention mechanism is introduced to enable the network to automatically focus on important features.

7. The remote sensing image semantic segmentation method based on U-Net according to claim 1, characterized in that: The segmentation head output in S6 is a category prediction for each pixel, and the category is a predefined category of different types of ground objects in the remote sensing image.

Images