A Semantic Segmentation Method for Remote Sensing Images Based on U-Net

Abstract

The present invention provides a semantic segmentation method for remote sensing images based on U-Net, belonging to the technical field of image segmentation. After obtaining the input image, the method sequentially passes the image through an encoder part composed of four processing modules, thereby obtaining four feature maps with different dimensions. The feature map with the highest number of channels and the smallest size output by the last module is input into the bottleneck part of the network. After feature enhancement, it enters the decoder part of the network, and the spatial resolution of the image is gradually restored through upsampling. During the restoration process, skip connections are used to connect the output feature maps of the encoder to the decoder to help the decoder better restore the details of the image. Finally, the output of the decoder will pass through a 1x1 convolutional layer to generate an output image with the same size as the input image. The present invention improves the U-shaped network by using three newly proposed modules, improves the segmentation accuracy, and ensures the efficiency and accuracy of the method in the semantic segmentation of remote sensing images.

Description

Technical Field

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

[0002] With the continuous progress of remote sensing technology, remote sensing images have been widely used in various fields such as environmental monitoring, urban planning, agricultural detection, and disaster assessment. 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 land cover types, and irregular object shapes, making it difficult for traditional image processing methods to effectively extract the detailed features of the target area. Especially in complex environments, the object diversity and scale variations in remote sensing images pose great challenges to the image segmentation task.

[0003] Semantic segmentation, as an important task in the field of image processing, aims to assign each pixel in an image to a specific category. Most traditional semantic segmentation methods rely on handcrafted feature extraction and classifiers, but these methods often have limited effectiveness when dealing with high-resolution remote sensing images, especially facing great difficulties in complex backgrounds and fine-grained classification. In recent years, semantic segmentation methods based on deep learning have gradually become the mainstream. Among them, the U-Net network, with its excellent performance and relatively low computational resource requirements, has become a commonly used architecture in remote sensing image semantic segmentation. U-Net can effectively improve the segmentation accuracy by performing multi-scale feature extraction on the input image and fusing feature information at different levels through skip connections, and it performs particularly well in detail preservation and small object 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 regions in high-resolution images, there is still room for improvement in the segmentation accuracy of U-Net. 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] To solve the above technical problems, the present invention provides a method for semantic segmentation of remote sensing images based on U-Net.

[0006] The method for semantic segmentation of remote sensing images based on U-Net provided by the present invention includes the following steps:

[0007] S1. Train an improved U-Net model using the 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, utilize its shunted self-attention mechanism and convolutional operations to extract the low-level and global feature information of the image;

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

[0011] S5. In the decoder part, gradually restore the spatial resolution of the image through upsampling and finally restore it to the same resolution as the input image;

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

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

[0014]

[0015] Among them, and are the main losses calculated for two different predictions, while and are the edge losses obtained from the two predictions, and 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 augmentation preprocessings during the training process to improve the generalization ability of the model, including: randomly horizontally flipping the image, 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 sequentially passing through the Shunted Transformer Block in the encoder in S3 is as follows:

[0018] S31. The encoder receives the three-channel RGB remote sensing image (assumed size is 3×256×256);

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

[0020] S33. Continuously input the feature map with a size of 64×64×64 into the stage2 module for processing, and obtain a feature map with a size of 128×32×32;

[0021] S34. Continuously input the feature map with a size of 128×32×32 into the stage3 module for processing, and obtain a feature map with a size of 256×16×16;

[0022] S35. Input the feature map with a size of 256×16×16 into the last module stage4 of the encoder for processing, and obtain a feature map with a size of 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 is necessary to first go through PatchEmbedding processing, the purpose of which is to convert the input image into a set of image patches with a fixed size, and use the image patches as the input sequence to meet the requirement that Transformer needs to process sequence data.

[0024] Furthermore, the skip connection in S4 is not limited to between the encoder and the decoder in the present invention, but also exists inside the decoder. This connection passes the deeper features to the shallower layers through upsampling, enabling the shallower features to obtain context information from the deeper layers, and making up for the deficiency of the shallow features in capturing global information. The fusion of such multi-level features enhances the model's ability to understand the input image at various scales.

[0025] It should be noted that: This new type of skip connection refers to the residual connection structure, which can, without increasing the number of parameters, also play a role in alleviating the problem of gradient disappearance during the training process.

[0026] Furthermore, during the upsampling process in S5, the present invention proposes a new fusion method, that is, using learnable weights and PReLU for non-linear combination, which allows for 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 fusion, a channel attention mechanism is introduced, enabling the network to automatically focus on important features, improving the expression ability of the fused features, and enhancing the model's understanding of multi-scale features.

[0027] Further, in S6, the segmentation head outputs the class prediction for each pixel, and the class is different types of ground object classes in the predefined remote sensing image. The predefined ground object classes in the present invention may be: opaque surface, building, low shrub, tree, vehicle.

[0028] Compared with the related technology, a 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 by the present invention can effectively fuse image features from different levels. Especially when dealing with complex backgrounds and small objects, the segmentation accuracy is significantly improved.

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

[0031] 3. The novel skip connection method introduced by the present invention can not only effectively avoid information loss, but also help the model overcome the problems of gradient disappearance and overfitting in the deep network when dealing with complex remote sensing images. Thus, the anti-interference ability of the model is improved. BRIEF DESCRIPTION OF THE DRAWINGS

[0032] Figure 1 It is the structure diagram of the improved U-Net model of the present invention;

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

[0034] Figure 3 It is the comparison diagram of the skip connection structure proposed by the present invention and other model structures. Among them, A is TransUNet., B is ST-UNet., C is HST-UNet., and D is the structure proposed by the present invention.

[0035] Figure 4 It is the structure diagram of the multi-scale feature fusion module. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0036] In order to make the objectives, 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 with reference to the accompanying drawings in the embodiments of the present invention. Obviously, the described embodiments are some, but not all, of the embodiments of the present invention. Usually, 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 claimed invention, but merely represents selected embodiments of the present invention. All other embodiments obtained by those of ordinary skill in the art based on the embodiments of the present invention without creative efforts shall fall within the scope of protection of the present invention.

[0038] The following specifically describes the specific implementation of the remote sensing image semantic segmentation method based on U-Net in combination with specific embodiments.

[0039] A remote sensing image semantic segmentation method based on U-Net provided by the present invention includes the following steps:

[0040] S1. Train an improved U-Net model through the Vaihingen and Potsdam public datasets. The improved U-Net model structure of the present invention is as Figure 1 shown;

[0041] S2. Input a three-channel RGB remote sensing image into the encoder part of the network;

[0042] S3. Through the Shunted Transformer Block in the encoder, utilize its shunted self-attention mechanism and convolution operation to extract low-level and global feature information of the image;

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

[0044] S5. In the decoder part, gradually restore the spatial resolution of the image through upsampling and finally restore it to the same resolution as the input image;

[0045] S6. Process the output of the decoder part using a segmentation head and classify each pixel to obtain the class label of each pixel.

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

[0047] First calculate L CE =-∑ i y i log(p i ), where y i represents the true label and p i is the model prediction result. At the same time, use y i and pi Calculate Combine L using the weight coefficient CE and L Dice to obtain L main = α·L CE + β·L Dice . On this basis, introduce the Laplacian kernel and apply it to the ground truth labels and model predictions for edge extraction: Edge(x) = ReLU(Conv(x, Laplacian Kernel)), where Conv represents the convolution operation, ReLU represents the rectified linear unit function, x represents the input image, and the Laplacian kernel is defined as Then, threshold the calculated edge features (i.e., set pixels greater than or equal to the threshold to 1 and other pixels to 0) to obtain a binary edge map. Subsequently, use binary cross-entropy loss 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, and the calculation method is where y represents the true label, represents the model prediction result. In addition, we introduce the KL divergence to ensure the consistency between the two outputs of the model, and the final expression is Here, t is the main loss calculated for the two different predictions, while and[[ID=LOREM]] are the edge losses obtained from the two predictions, and L KL is the KL divergence loss.

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

[0049] S31: The encoder receives a three-channel RGB remote sensing image (assumed to be 3×256×256);

[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 continuously 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 continuously 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 the 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 the features of Transformer and convolution. Its structural diagram is as Figure 2 shown. Its core advantage lies in optimizing the computational efficiency based on the traditional Transformer structure through Shunted Attention, and enhancing the ability to model image features.

[0055] Furthermore, the comparison diagram of the skip connection structure in S4 with other model structures is as Figure 3 shown. In the present invention, the skip connection is not only limited to between the encoder and the decoder, but also exists inside the decoder. This connection passes the deeper features to the shallower layers through upsampling, enabling the shallower features to obtain context information from the deeper layers, making up for the deficiency of the shallow features in capturing global information. The fusion of such multi-level features enhances the model's ability to understand the input image at various scales.

[0056] Furthermore, in the upsampling process of S5, the present invention proposes a new fusion method, that is, using learnable weights and PReLU for non-linear combination, which allows for 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 fusion, a channel attention mechanism is introduced, enabling the network to automatically focus on important features, improving the expressive ability of the fused features, and enhancing the model's understanding of multi-scale features. The specific upsampling process is as follows:

[0057] S51. Input the output feature map (512×8×8) of the model bottleneck part into the first fusion module in the decoder. This fusion module will first upsample the feature map to 256×16×16 through an upsampling operation, and then fuse the upsampled feature map with the output feature map of stage3 in the encoder part;

[0058] S52. Input the fused feature map of size 256×16×16 into the second fusion module in the decoder, and repeat similar operations until a feature map of size 64×64×64 is obtained;

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

[0060] Further, in S6, the splitting head will adjust the number of channels of the obtained feature map with a size of 64×256×256 to the number of predicted categories, and thus the category prediction of each pixel can be obtained. The categories are predefined land cover categories in remote sensing images. The predefined land cover categories in the present invention are respectively: opaque surface, building, low shrub, tree, vehicle.

[0061] The convolutional network and Transformer architecture involved in the present invention are all prior arts and will not be elaborated here.

[0062] The above are only the embodiments of the present invention, and thus do not limit the patent scope of the present invention. Any equivalent structure or equivalent process transformation made by using the content of the specification and drawings of the present invention, or directly or indirectly applied in other related technical fields, shall be equally included in the patent protection scope of the present invention.

Claims

1. A semantic segmentation method for remote sensing images based on U-Net, characterized in that, It includes the following steps: S1. Train an improved U-Net model with 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, utilize its shunted self-attention mechanism and convolutional operations to extract the low-level and global feature information of the image; S4. Use skip connections to fuse the low-level features output by the encoder part with the high-level features restored by the decoder part; S5. In the decoder part, gradually restore the spatial resolution of the image through upsampling and finally restore it to the same resolution as the input image; S6. Process the output of the decoder part using the segmentation head and classify each pixel to obtain the class label of each pixel; The specific process of successively passing through the Shunted Transformer Block in the encoder in S3 is as follows: S31. The encoder receives the three-channel RGB remote sensing image; S32. The input image is processed by the stage1 module of the encoder to obtain a feature map with a size of 64×64×64; S33. The feature map with a size of 64×64×64 is continuously input into the stage2 module for processing to obtain a feature map with a size of 128×32×32; S34. The feature map with a size of 128×32×32 is continuously input into the stage3 module for processing to obtain a feature map with a size of 256×16×16; S35. The feature map with a size of 256×16×16 is input into the last module stage4 of the encoder for processing to obtain a feature map with a size of 512×8×8 as the output of the encoder.

2. The method for remote sensing image semantic segmentation based on U-Net according to claim 1, wherein The input image in S2 will undergo data augmentation preprocessing during training, including randomly horizontally flipping the image, 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.

3. A method for semantic segmentation of remote sensing images based on U-Net according to claim 1, characterized in that, The skip connections in S4 are not limited to between the encoder and the decoder, but also exist inside the decoder; the deeper features are passed to the shallower layers through upsampling so that the shallower features can obtain context information from the deeper layers.

4. A method for semantic segmentation of remote sensing images based on U-Net according to claim 1, characterized in that, During the upsampling process in S5, a learnable weight and PReLU are used for non-linear combination to adjust the contribution weights between features, capture more complex feature relationships than simple addition or multiplication fusion methods, and after the feature map fusion, a channel attention mechanism is introduced to enable the network to automatically focus on important features.

5. A semantic segmentation method for remote sensing images based on U-Net according to claim 1, characterized in that The output of the segmentation head in S6 is the class prediction for each pixel, and the class is different types of ground object classes in the predefined remote sensing image.

Images