Dual-branch remote sensing image semantic segmentation method and device based on visual transformer and mamba

Abstract

The application belongs to the technical field of optical image imaging, and particularly relates to a double-branch remote sensing image semantic segmentation method and device based on a vision transformer and Mamba. The specific process of the method is as follows: firstly, given an input image with RGB encoding, the image is sent to a ViT branch and a Mamba branch for parallel processing; secondly, the ViT branch and the Mamba branch respectively generate four feature mappings through a hierarchical Block structure and a downsampling mode, and the four feature mappings are respectively represented as a set S={F Si |i=1, 2, 3, 4} and a set M={F Mi |i=1, 2, 3, 4}; thirdly, different species features of the second and fourth stages of the ViT branch and the Mamba branch are fused for consistency by using a cross-domain fusion module, and the second stage fusion feature map is added to the fourth stage fusion feature map through downsampling; finally, the fourth stage cross-domain fused feature representation is used as a segmentation head to realize semantic segmentation of the input image.

Description

Technical Field

[0001] The present invention belongs to the technical field of optical image imaging, and in particular relates to a dual-branch remote sensing image semantic segmentation method and device based on a visual converter and Mamba. Background Art

[0002] Remote sensing image semantic segmentation uses deep learning techniques to classify and segment different objects or targets in remote sensing images at the pixel level, enabling accurate object identification and location. This technology has important applications in environmental monitoring, urban planning, agricultural management, and other fields. It can provide accurate spatial data support for geographic information systems and intelligent transportation, and has significant practical significance and research value.

[0003] Deep learning-based semantic segmentation of remote sensing images has made significant progress in recent years, primarily due to the rapid development of deep learning technology. Convolutional neural networks (CNNs), as powerful feature extraction and classification models, have become the cornerstone of numerous image processing tasks. However, CNNs have a limited receptive field when extracting features based on convolution operations, making them particularly challenging when processing high-resolution images such as remote sensing images.

[0004] To address the receptive field bottleneck of CNNs, the Vision Transformer (ViT) has gained popularity in semantic segmentation tasks with the introduction of the attention mechanism. ViT, with its powerful long-range modeling capabilities, overcomes the limitations imposed by its receptive field and demonstrates remarkable performance across a variety of image tasks. However, due to the need to compute attention between patch sequences, ViT's computational complexity increases exponentially with increasing image resolution. Therefore, the stringent requirements of remote sensing image semantic segmentation on both model parameter count and computational complexity pose a significant challenge to the application of ViT in remote sensing tasks. A representative solution to this problem is the Swin-Transformer, which introduces a groundbreaking sliding window mechanism that significantly reduces ViT's computational complexity. However, by restricting attention computation to a window of limited size, the Swin-Transformer's lightweight architecture inevitably comes at the expense of global long-range modeling capabilities. Therefore, an effective solution is needed to complement the Swin-Transformer's lost global modeling capabilities within an acceptable computational complexity.

[0005] The Mamba architecture, based on state-space models, has recently garnered widespread attention due to its linear complexity advantage. State-space models hold great potential for modeling long-range dependencies, and Mamba's introduction of a selective parameter learning mechanism makes it practical. Various structured Mamba models have incorporated various architectural enhancements, further demonstrating its versatility and scalability. Summary of the Invention

[0006] The purpose of the present invention is to solve the problem that specific features in a single network cannot simultaneously solve the paradox of receptive field and detail preservation. A dual-branch remote sensing image semantic segmentation method and device based on visual transformer and Mamba is provided. The ViT branch based on Swin Transformer is used to extract local semantic information with detail preservation, and the Mamba branch is used to efficiently capture global semantic information. The cross-domain fusion module is used to realize the fusion of dual-space semantic information, thereby achieving high-precision semantic segmentation.

[0007] To achieve the above object, the present invention is achieved through the following technical solutions:

[0008] In the first aspect, the embodiment of the present application provides a dual-branch remote sensing image semantic segmentation method based on visual converter and Mamba, the specific process is as follows:

[0009] First, given an RGB-encoded input image, the image is sent to the ViT branch and the Mamba branch for parallel processing;

[0010] Secondly, the ViT branch and the Mamba branch generate four feature maps respectively through the hierarchical Block structure and downsampling, which are represented as the set S = {F Si |i=1,2,3,4} and set M={F Mi |i=1,2,3,4};

[0011] Thirdly, the cross-domain fusion module is used to consistently fuse the different features of the second and fourth stages of the ViT branch and the Mamba branch, and the second-stage fused feature map is added to the fourth-stage fused feature map through downsampling;

[0012] Finally, the feature representation after cross-domain fusion in the fourth stage is used as the segmentation head to achieve semantic segmentation of the input image.

[0013] Furthermore, the process of performing consistency fusion by the cross-domain fusion module of the present invention is as follows:

[0014] First, the fusion feature maps of the two branches at the same stage are flattened;

[0015] Secondly, a linear mapping of query, key, and value is generated independently for the fusion features expanded in each branch, and the corresponding elements of the linear mapping results are fused;

[0016] Thirdly, the cross attention weight ω is calculated based on the fusion result using the softmax function, and the result of the cross attention calculation is calculated according to the weight;

[0017] Finally, the result F calculated based on the cross attention o Obtain cross-model fusion feature maps.

[0018] Furthermore, the cross-domain fusion module of the present invention sets the cross-model correlation loss L in the training phase. CMR , using L CMR To update the parameters of the cross-domain fusion module;

[0019]

[0020] Among them, CMRM is the calculation of cross-model association graph based on cross-attention weight ω, λ is the set hyperparameter, L CE Represents the cross entropy loss, GT represents the true pixel label value Ground-Truth of (i, j), and W is the weight calculated based on the preliminary prediction results of Ground-Truth and dual-branch features.

[0021] Furthermore, the cross-model relationship graph CMRM is calculated based on the cross attention weight ω of the present invention as follows:

[0022]

[0023] Among them, Conv represents the convolution layer, BN represents the batch normalization operation, and ReLU is the activation layer. Provides matrix addition.

[0024] Furthermore, the present invention for:

[0025]

[0026] Among them, P S and P M The prediction results of the fusion features calculated by the two branches at the same stage.

[0027] In a second aspect, an embodiment of the present application provides a dual-branch remote sensing image semantic segmentation device based on a visual converter and Mamba, comprising: a ViT branch, a Mamba branch, a cross-domain fusion module, and an image segmentation module;

[0028] The ViT branch is used to generate four feature maps for the input RGB encoded image through a hierarchical Block structure and downsampling, which is represented as a set S = {F Si |i=1,2,3,4} and set M={F Mi |i=1,2,3,4};

[0029] The Mamba branch is used to generate four feature maps for the input RGB encoded image through a hierarchical Block structure and downsampling, which is represented as a set M = {F Mi |i=1,2,3,4};

[0030] The second-order cross-domain fusion module is used to consistently fuse the different features of the second stage of the ViT branch and the Mamba branch, downsample the second-stage fused feature map and output it to the fourth-order cross-domain fusion module;

[0031] The 4th-order cross-domain fusion module is used to consistently fuse the different features of the 4th stage of the ViT branch and the Mamba branch, and superimpose them with the fused feature map of the 2nd stage to achieve cross-domain fusion;

[0032] The image segmentation module is used to use the feature representation after superimposing the cross-domain fusion module in the fourth stage as the segmentation head to realize the semantic segmentation of the RGB-encoded remote sensing image input by the network.

[0033] Beneficial effects:

[0034] First, unlike existing single networks that fail to address the receptive field paradox and detail-preserving features, this paper designs a dual-branch network combining a ViT branch and a Mamba branch. The ViT branch extracts detail-preserving local information, while the Mamba branch captures global semantic information. Simultaneously, the network effectively fuses these two semantic features through a cross-domain fusion module, achieving the fusion of detail-preserving local information and global semantic information in dual spaces. Ultimately, this achieves semantic segmentation of high-resolution remote sensing images. The feature fusion effect is further enhanced through a cross-model correlation loss.

[0035] Second, cross-model correlation loss is introduced to enhance feature consistency between different architectures and effectively suppress noise in the cross-domain fusion process. BRIEF DESCRIPTION OF THE DRAWINGS

[0036] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the following briefly introduces the drawings required for use in the embodiments. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0037] Figure 1This is a schematic diagram of the overall structure of a dual-branch remote sensing image semantic segmentation method based on visual converter and Mamba provided by the present invention. DETAILED DESCRIPTION

[0038] The embodiments of the present invention are described in detail below with reference to the accompanying drawings.

[0039] It should be noted that, in the absence of conflict, the following embodiments and features in the embodiments may be combined with each other; and, based on the embodiments in this disclosure, all other embodiments obtained by persons of ordinary skill in the art without creative work are within the scope of protection of this disclosure.

[0040] It should be noted that various aspects of the embodiments within the scope of the appended claims are described below. It should be apparent that the aspects described herein can be embodied in a wide variety of forms, and any specific structure and / or function described herein is merely illustrative. Based on this disclosure, it should be understood by those skilled in the art that an aspect described herein can be implemented independently of any other aspect, and two or more of these aspects can be combined in various ways. For example, any number of aspects described herein can be used to implement an apparatus and / or practice a method. In addition, other structures and / or functionalities other than one or more of the aspects described herein can be used to implement this apparatus and / or practice this method.

[0041] In view of the fact that specific features in a single ViT network cannot simultaneously solve the paradox of receptive field and the problem of detail preservation, inspired by Mamba, the present invention proposes a dual-branch network composed of ViT and Mamba to jointly and efficiently model image features. This invention rethinks the dual-branch architecture and uses a ViT branch based on Swin Transformer to extract local semantic information with detail preservation; at the same time, the Mamba calculation mechanism is used in the Mamba branch to replace the multi-head self-attention in ViT, and combined with multiple multi-layer perceptrons (MLPs) to efficiently capture global semantic information from image data. In addition, although the dual-branch network defined above is already capable of extracting effective features with double spatial size, due to the inherent inconsistency of the feature extraction architecture, the final features obtained are suboptimal. Therefore, the present invention designs a cross-model fusion module based on the attention mechanism, and proposes an attention-guided loss to further guide consistent fusion.

[0042] like Figure 1As shown, the embodiment of the present application provides a dual-branch remote sensing image semantic segmentation method based on visual transformer and Mamba. The local semantic features and global semantic features of the input image are extracted through the ViT and Mamba branches respectively; at the same time, the cross-domain fusion module is used to achieve the effective fusion of the two local features and global features, and the cross-model correlation loss is used to further enhance the feature fusion effect, ultimately achieving high-resolution remote sensing image semantic segmentation. The specific process is as follows:

[0043] Step 0: Get the source data

[0044] Capture at least 500 high-resolution optical images and manually label the target areas pixel by pixel, forming the dataset for the subsequent semantic segmentation model. Manual classification labeling refers to labeling pixels with corresponding labels based on the image content. For example, if a pixel belongs to category 1, it is labeled 1; if a pixel does not belong to any target category, it is usually labeled 0 (background).

[0045] Step 1: ViT branch local feature extraction

[0046] like Figure 1 As shown, given an RGB input image I∈R 3×H×W First, the ViT backbone network module based on Swin-Transformer is used to hierarchically extract four local semantic feature maps S of different sizes from the input image in RGB form through a hierarchical Block structure:

[0047] S={F Si |i=1,2,3,4}

[0048] Among them, H and W represent the height and width of the image respectively, i represents the stage number of the progressive layered downsampling process, and F Si The size is The local semantic feature map S extracted in this step retains the detail information of the RGB image.

[0049] Step 2: Mamba branch global feature extraction

[0050] like Figure 1 As shown in Figure 2, the Mamba branch uses downsampling operations to extract four global semantic feature maps M at different depths:

[0051] M={F Mi |i=1,2,3,4}

[0052] Where i represents the stage number of the progressive layered downsampling process, F Mi The size is The global semantic feature map M extracted in this step retains the global semantic information of the image.

[0053] Step 3: Use the cross-domain fusion module to consistently fuse the different features of the 2nd and 4th stages of the ViT branch and the Mamba branch, and add the 2nd stage fused feature map to the 4th stage fused feature map through downsampling;

[0054] like Figure 1 As shown in the figure, the present invention designs a Cross-Model Fusion Module (CMFM) to consistently fuse the second and fourth stage features captured in steps 1 and 2, and adds the second stage fusion features to the fourth stage fusion result through downsampling. The consistency of different architecture features is enhanced through cross-model correlation loss, and cross-domain fusion noise is suppressed.

[0055] Taking the second stage feature fusion as an example, the feature F obtained by the Swin-Transformer sliding window attention mechanism in step 1 is S2 With rich local detail information, the F obtained in step 2 M2 This step uses the cross attention mechanism to remodel F S2 and F M2 Based on the similarity relationship between the two branches, the key global semantics and spatial detail features of each branch are emphasized, while irrelevant noise and contradictory information are suppressed, thereby improving the performance of semantic segmentation. The specific process is as follows:

[0056] First, F S2 and F M2 Flattened to get and

[0057] Secondly, a linear map of Query, Key, and Value is generated independently for each branch, and the corresponding elements are merged:

[0058]

[0059] Among them, Q, K, and V represent the corresponding linear mapping layers, and LN represents layer normalization.

[0060] Again, the softmax function is used to obtain the cross attention weights:

[0061]

[0062] Where d is the dimension of Q and K, represents matrix multiplication, ω represents cross attention weight, F o is the result of cross attention calculation.

[0063] Finally, the result of the cross attention calculation Fo Input to the MLP layer, enhance the learning ability to produce consistent cross-model fusion output. The cross-domain fusion output of the second stage is downsampled to a size of The feature map is added to the output of the fourth stage cross-domain fusion module to obtain the final cross-domain fusion feature map.

[0064] In the process of cross attention, the present invention calculates the cross-model relationship map CMRM based on the cross attention weight ω:

[0065]

[0066] Among them, Conv represents the convolution layer, BN represents the batch normalization operation, and ReLU is the activation layer. Provides matrix addition.

[0067] In order to better achieve the consistent fusion of the two branch features, the present invention designs a cross-model correlation loss CMRLoss to supervise CMRM:

[0068]

[0069] Among them, λ is a hyperparameter used to adjust The impact of L CE Represents the cross entropy loss, and GT represents the true pixel label value Ground-Truth of (i, j). It is a Ground-Truth and dual-branch feature F S2 、F M2 Weights calculated for preliminary forecast results:

[0070]

[0071] Among them, P S and P M According to F S2 and F M2 Calculated prediction results. L CMR Used to update the parameters of the cross-domain fusion module.

[0072] Step 4: Implement semantic segmentation using cross-domain fusion feature representation

[0073] like Figure 1 As shown, the present invention finally uses the feature representation after cross-domain fusion in the fourth stage as the segmentation head to realize the semantic segmentation of the input high-resolution remote sensing image.

[0074] In a second aspect, an embodiment of the present application provides a dual-branch remote sensing image semantic segmentation device based on a visual converter and Mamba, comprising: a ViT branch, a Mamba branch, a cross-domain fusion module, and an image segmentation module;

[0075] The ViT branch is used to generate four feature maps for the input RGB encoded image through a hierarchical Block structure and downsampling, which is represented as a set S = {F Si |i=1,2,3,4} and set M={F Mi |i=1,2,3,4};

[0076] The Mamba branch is used to generate four feature maps for the input RGB encoded image through a hierarchical Block structure and downsampling, which is represented as a set M = {F Mi |i=1,2,3,4};

[0077] The second-order cross-domain fusion module is used to consistently fuse the different features of the second stage of the ViT branch and the Mamba branch, downsample the second-stage fused feature map and output it to the fourth-order cross-domain fusion module;

[0078] The 4th-order cross-domain fusion module is used to consistently fuse the different features of the 4th stage of the ViT branch and the Mamba branch, and superimpose them with the fused feature map of the 2nd stage to achieve cross-domain fusion;

[0079] The image segmentation module is used to input the feature representation after superimposing the cross-domain fusion module in the fourth stage as the segmentation head into the high-resolution remote sensing image to achieve semantic segmentation of the input image.

[0080] The above description is merely a specific embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any changes or substitutions that can be easily conceived by a person skilled in the art within the technical scope disclosed in the present invention should be included in the scope of protection of the present invention. Therefore, the scope of protection of the present invention should be based on the scope of protection of the claims.

Claims

1. A dual-branch remote sensing image semantic segmentation method based on visual converter and Mamba, characterized by: The specific process is: First, given an RGB-encoded input image, the image is sent to the ViT branch and the Mamba branch for parallel processing; Secondly, the ViT branch and the Mamba branch generate four feature maps respectively through the hierarchical Block structure and downsampling, which are represented as the set S = {F Si |i=1,2,3,4} and set M={F Mi |i=1,2,3,4}; Thirdly, the cross-domain fusion module is used to consistently fuse the different features of the second and fourth stages of the ViT branch and the Mamba branch, and the second-stage fused feature map is added to the fourth-stage fused feature map through downsampling; Finally, the feature representation after cross-domain fusion in the fourth stage is used as the segmentation head to achieve semantic segmentation of the input image; The process of consistency fusion performed by the cross-domain fusion module is as follows: First, the fusion feature maps of the two branches at the same stage are flattened; Secondly, a linear mapping of query, key, and value is generated independently for the fusion features expanded in each branch, and the corresponding elements of the linear mapping results are fused; Thirdly, the cross attention weight ω is calculated based on the fusion result using the softmax function, and the result of the cross attention calculation is calculated according to the weight; Finally, the result F calculated based on the cross attention o Obtain cross-model fusion feature maps.

2. The dual-branch remote sensing image semantic segmentation method based on visual converter and Mamba according to claim 1, characterized in that: The cross-domain fusion module sets the cross-model correlation loss L during the training phase. CMR , using L CMR To update the parameters of the cross-domain fusion module; Among them, CMRM is the calculation of cross-model association graph based on cross-attention weight ω, λ is the set hyperparameter, L CE Represents cross entropy loss, GT represents the true pixel label value Ground-Truth of (i, j), The weights calculated based on the preliminary prediction results of Ground-Truth and dual-branch features.

3. The dual-branch remote sensing image semantic segmentation method based on visual converter and Mamba according to claim 2, characterized in that: The cross-model association graph CMRM is calculated based on the cross attention weight ω: Among them, Conv represents the convolution layer, BN represents the batch normalization operation, and ReLU is the activation layer. Provides matrix addition.

4. The dual-branch remote sensing image semantic segmentation method based on visual converter and Mamba according to claim 2, characterized in that: described for: Among them, P S and P M The prediction results of the fusion features calculated by the two branches at the same stage.

5. A dual-branch remote sensing image semantic segmentation device based on visual converter and Mamba, characterized in that: include: ViT branch, Mamba branch, cross-domain fusion module, and image segmentation module; The ViT branch is used to generate four feature maps for the input RGB encoded image through a hierarchical Block structure and downsampling, which is represented as a set S = {F Si |i=1,2,3,4}; The Mamba branch is used to generate four feature maps for the input RGB encoded image through a hierarchical Block structure and downsampling, which is represented as a set M = {F Mi |i=1,2,3,4}; The second-order cross-domain fusion module is used to consistently fuse the different features of the second stage of the ViT branch and the Mamba branch, downsample the second-stage fusion feature map and output it to the fourth-order cross-domain fusion module; The 4th-order cross-domain fusion module is used to consistently fuse the different features of the 4th stage of the ViT branch and the Mamba branch, and superimpose them with the fused feature map of the 2nd stage to achieve cross-domain fusion; The image segmentation module is used to use the feature representation after superimposing the cross-domain fusion module in the fourth stage as the segmentation head to realize the semantic segmentation of the RGB-encoded remote sensing image input by the network; The process of consistency fusion performed by the cross-domain fusion module is as follows: First, the fusion feature maps of the two branches at the same stage are flattened; Secondly, a linear mapping of query, key, and value is generated independently for the fusion features expanded in each branch, and the corresponding elements of the linear mapping results are fused; Thirdly, the cross attention weight ω is calculated based on the fusion result using the softmax function, and the result of the cross attention calculation is calculated according to the weight; Finally, the result F calculated based on the cross attention o Obtain cross-model fusion feature maps.

Images