Remote Sensing Image Segmentation Using Vision Mamba and Multi-Scale Multi-Frequency Feature Fusion

Recently, SSM become a research hotspot. As the current leading SSM, Mamba not only excels in long-range modeling capabilities but also demonstrates linear complexity in handling input sizes, addressing the computational efficiency issues of Transformers in modeling long sequences of state spaces. In the field of visual research, Vim [21] proposes a pure vision backbone model based on SSM, introducing Mamba into the visual domain for the first time. VMamba [22] enhances visual processing capabilities by introducing the Cross-Scan module, enabling the model to selectively scan images in two dimensions and demonstrating superiority in image classification tasks. LocalMamba [30] focuses on a windowed scanning strategy for visual spatial models, optimizing visual information to capture local dependencies and introducing a dynamic scanning method to search for optimal choices at different layers. In downstream visual tasks, Mamba also sees extensive use in the domain of remote sensing image segmentation research. Rs3Mamba [23] proposes using the VSSBlock to provide additional global information, aiding the convolution-centered main stream in the extraction of features, thereby enhancing global awareness of the model while preserving local details. In CM-UNet [24], a CNN encoder is tasked with acquiring detailed image features, and a Mamba decoder is in charge of combining global information. The inclusion of an MSAA module allows for the merging of features at various scales, competently capturing long-distance correlations and multiscale global context information in high-resolution remote sensing images. MambaHSI [31] introduces a new Mamba-based model for HSI classification, MambaHSI, which simultaneously models long-distance interactions across the entire image and adaptively integrates spatial and spectral data in a flexible manner. Although these strategies utilize the strong ability of VMamba to model dependencies over long distances while maintaining linear computational complexity, they overlook the limitations in local information extraction that may arise due to specific scanning directions.

SENet [32] introduces the channel attention mechanism, which performs global average pooling (GAP) along the channel dimension and uses fully connected layers to compute the weight of each channel. Due to its significant performance improvements, attention modules gain widespread attention. Subsequent studies, such as CBAM [33] and scSE [34], use 2D convolution with a kernel size of $k\times k$ to calculate spatial attention and integrate it with channel attention. SRM [35] suggests using GAP with global standard deviation pooling. Although these attention mechanisms perform well in practice, the dimensionality reduction in the fully connected layers of GAP negatively impacts the attention mechanism. Therefore, ECANet [27] proposes using a Adaptive 1D Conv to replace the fully connected layer, avoiding the negative effects of dimensionality reduction and capturing local cross-channel interactions. Additionally, a pivotal aspect of channel attention mechanisms is the calculation of scalar values for individual channels. Due to its uncomplicated design, GAP struggles to effectively extract complex input features. FCANet [28] demonstrates that GAP represents a unique instance of the Discrete Cosine Transform (DCT) and explores different frequency component combinations within a multi-spectral framework, which enhances the ability of GAP to capture complex information. These improved attention mechanisms, by more precisely handling inter-channel and intra-channel interactions, help models better understand and utilize input data features. Therefore, effectively applying these methods in feature fusion strategies becomes key to achieving refined feature fusion.

In U-shaped networks, long skip connections are an important component. They help the network obtain detailed semantic features by bridging finer details from lower layers with coarse resolution and advanced semantic features. Although skip connections are widely used to combine features from various paths, the fusion of these connected features is typically achieved through addition or concatenation. This approach ignores the differences between the features, assigning the same weight to all of them, which prevents the complementary nature of features at different levels from being fully utilized. As a result, the capability of the network to recognize both minor details and advanced semantic features is limited. With the widespread application of attention mechanisms, SKNet [36] introduces a method for specialized feature fusion based on attention mechanisms that uses a non-linear approach. However, its constraints are apparent in its focus only on soft feature selection within the same layer, overlooking the challenge of integrating features across layers in skip connections. Moreover, the scale variation of objects represents a significant challenge within the realm of computer vision. To mitigate the issues arising from scale variation and small objects, as well as to enhance the effectiveness of skip connections, the AFF module [25] proposes a unified and generalized method for feature fusion using attention mechanisms. Similarly, the Multiattention Network (MANet) [26] extends attention mechanisms across multiple scales, aggregating contextual information at different scales to obtain a more complete and nuanced merging of features. Although these feature fusion methods can effectively enhance the feature fusion effects of long skip connections, they often overlook the issues of information transmission loss and insufficient information utilization in the attention mechanisms, as mentioned in ECANet [27] and FCANet [28].

The complete network structure of the proposed CVMH-UNet is depicted in Fig. 1, which is designed and improved based on VMamba. The CVMH-UNet architecture primarily comprises an embedding block, an encoder, a decoder, a multi-frequency multi-scale feature fusion block (MFMSBlock), and a final projection layer. The original remote sensing image $x\in\mathbb{R}^{3\times H\times W}$ is segmented into non-intersecting $4\times 4$ patches in the embedding block and mapped to a C-dimensional feature space. Subsequently, it undergoes processing by a normalization layer to produce the embedded feature map $x^{\prime}\in\mathbb{R}^{C\times\frac{H}{4}\times\frac{W}{4}}$, where the mapping dimension C is set to 96. Subsequently, the embedded features are fed into the encoder, which consists of CVSSBlocks, to extract features. In the first three stages, patch merging is employed for downscaling to progressively decrease spatial extents and acquire hierarchical features. Afterward, the decoder, also composed of CVSSBlocks, reconstructs the features, and in the last three stages, patch expanding is used for upsampling to increase spatial dimensions and further extract higher-level feature representations. Finally, the final result is obtained through the final projection layer, restoring the dimensions of the output image to match those of the input image. It is important to note that both the encoder and decoder are composed of CVSSBlocks, but they follow an asymmetric configuration strategy. In the four stages of the encoding part, the number of CVSSBlocks is [2,2,2,2], while in the four stages of the decoding part, the number of CVSSBlocks is [2,2,2,1]. This design aims to balance model complexity and computational efficiency. Additionally, to better restore details and capture more target information, the decoder interfaces with the encoder through the MFMSBlock for feature fusion. This ensures strong information integrity and precision in information transmission, strengthening the discriminative capacity of the features and refining detailed fusion of small objects, leading to segmentation results that are more accurate and exact.

The basic unit of the CVMH-UNet encoder-decoder architecture is the proposed CVSSBlock, which is an improved version of the VSSBlock from VMamba. The advantage of this module lies in its higher feature representation capability, achieved through the integration of convolutional branches and optimized scanning methods. The specific configuration is illustrated in Fig. 2(a). Inspired by VMamba [22], and offset the deficiency in local information extraction resulting from the SS2D scanning method in the VSSBlock, we designed a structure with two branches in the CVSSBlock by adding a convolutional branch to capture local information. Additionally, to address the inefficiency in information extraction caused by the SS2D method in VMamba, which scans in different directions along the same path, we designed the CS2D method with intersecting scanning paths. This allows the extraction of spatial features from multiple directions without increasing computational complexity, making it more efficient in capturing global information from high-resolution remote sensing images.

As can be seen in Fig. 2(a), two branches are applied on the input feature $F_{in}$ to obtain the global and local information. More specifically, within the global pathway, the input features first go through the Cross Scan module (CS) to extract the global information, and then the global information is input to the channel attention (CA) mechanism to obtain the global feature $F_{g}$. This system leverages the interdependencies among channel mappings to enhance the representation of specific semantic features. In the CS module, after layer normalization, input $F_{in}$ is channeled into two distinct branches. The initial branch includes a linear layer and an activation step. The other branch encompasses a linear layer, a depthwise separable convolution, and an activation step, before entering the CS2D module for extracting global features across various directions. Post this, the features undergo normalization and are subsequently merged element-wise with the output from the global pathway. In conclusion, a linear layer blends the features, which are subsequently integrated with the residual connection to produce the final output of the CS module. In the local branch, the input features first pass through a $3\times 3$ convolution to extract local information, which is then enhanced by the Spatial Attention (SA) mechanism to emphasize local details and suppress irrelevant regions, resulting in the local feature $F_{l}$. To sum up, the process can be represented as follows:

where $Conv_{3\times 3}(\cdot)$ represents a $3\times 3$ convolution, $CS(\cdot)$ represents the Cross Scan module, and $CA(\cdot)$ and $SA(\cdot)$ represent the channel attention mechanism and spatial attention mechanism, respectively. Afterward, the global and local features are added together, followed by depthwise separable convolution, layer normalization, and a $1\times 1$ convolution to refine the combined features. These refined features are then are merged with the starting input $F_{in}$ through a residual connection, generating the fused feature $F_{u}$. Finally, to enhance feature representation, the E-FFN module is introduced to further improve information interaction between different channels, achieving better segmentation performance. This process can be represented as follows:

where $F_{out}$ denotes the output features of the CVSSBlock, $dwConv_{3\times 3}(\cdot)$ denotes the $3\times 3$ depth-separable convolution, $LN(\cdot)$ denotes layer normalization, $Conv_{\mathrm{l\times l}}(\cdot)$ denotes the $1\times 1$ convolution, and $EN(\cdot)$ denotes the E-FNN module.

Additionally, to optimize the feature extraction process and reduce information loss, enabling the network to better extract features at each layer, two CVSSBlocks are connected using a residual connection. As shown in Fig 2(d), this approach reduces information loss, helping the network better identify and segment targets in high-resolution remote sensing images. Moreover, it assists the network in learning more refined feature representations, thus enhancing the precision of the segmentation of the model.

The encoder and decoder of CVMH-UNet are capable of extracting low-level features with rich texture details and high-level features with semantic information, respectively. To more effectively utilize this information, we designed the MFMSBlock. The core idea is to introduce frequency information into global pooling and combine it with additional local scale analysis to construct a multi-frequency, multi-scale feature fusion mechanism. This module adopts a structure with two branches to acquire information across different scales different scales. Within the global branch, 2D DCT is used to introduce frequency information into global pooling, and various global pooling techniques are utilized to capture fine features across different frequency bands, enhancing the completeness of the information. Additionally, Adaptive 1D Conv [27] is employed to directly compute channel attention, preventing data loss due to the dimensional reduction by fully connected layers, thus improving the precision of information transmission and producing more accurate global feature channel attention. In the local branch, point-wise convolution is used to provide local feature channel attention at another scale, capturing complex details and broader structural information. Finally, the two branches aggregate information from different scales along the channel dimension. The fusion approach not only strengthens the capacity of the model to detect local texture details but also enhances its comprehension of global semantic frameworks, substantially elevating the overall efficacy of the model.

The overall structure of the proposed MFMSBlock is depicted in Fig. 4(a). The feature map $X_{i}$, obtained by adding the low-level feature map $F_{i}$ from the ${i}$-th layer of the encoder and the high-level semantic feature map $\tilde{F}_{i}$ from the ${i}$-th layer of the decoder, passes through the multi-frequency multi-scale attention mechanism (MFMS-AM) to generate fusion weights $X_{i}^{\prime}$ and $1-X_{i}^{{}^{\prime}}$. This allows the MFMSBlock to conduct a soft selection process or a weighted mean between $F_{i}$ and $\tilde{F}_{i}$. This module can be expressed as:

where $Z_{i}\in\mathbb{R}^{C\times H\times W}$ represents the fused features, and $MA(\cdot)$ represents the attention weight values generated by MFMS-AM.

The overall structure of MFMS-AM is shown in Fig. 4(b). It uses 2D DCT to express image features as a weighted sum of cosine functions at different frequencies. The representation at the ${k}$-th frequency is denoted as $X_{i}^{k}$, and is calculated using the following formula:

where $(u_{k},v_{k})$ is the index of the frequency corresponding to $X_{i}^{k}$. In addition, the 2D DCT basis image using the top-K selection strategy [28] is defined as:

Next, $X_{i}^{k}$ is compressed by applying global average pooling, global max pooling, and global min pooling, respectively. These are then passed through a Adaptive 1D Conv [27] to avoid dimensionality reduction and directly compute the weights. With the channel dimension $C$ as a reference, the kernel size ${\phi}$ can be dynamically calculated by applying the following equation:

where $\begin{vmatrix}{\lambda}\end{vmatrix}_{odd}$ denotes the nearest odd number to ${\lambda}$. In the experiments conducted for this paper, ${\alpha}$ and ${\beta}$ are uniformly set to 2 and 1, respectively. Finally, the branches are summed to obtain the global channel attention map $G(X_{i})\in\mathbb{R}^{C}$. Additionally, the proposed MFMS-AM follows the concept of the Multi-Scale Channel Attention Module (MS-CAM) [25], where an additional point-wise convolution branch provides local feature channel attention at another scale. It only leverages point-wise interactions at each spatial location to enable channel interaction. The local feature channel attention map $L(X_{i})\in\mathbb{R}^{C\times H\times W}$ is acquired via the subsequent bottleneck architecture:

where $pwConv_{{}_{1}}(\cdot)$ and $pwConv_{{}_{2}}(\cdot)$ represent point-wise convolutions, $\theta(\cdot)$ denotes the rectified linear unit (ReLU), and $\psi(\cdot)$ denotes batch normalization (BN). It is worth noting that $L(X_{i})$, which aligns with the shape of the input features, allows it to preserve and highlight subtle details in low-level features. Given $G(X_{i})$ and $L(X_{i})$, MFMS-AM can produce the fusion weight $X_{i}^{\prime}$ through the activation function (Sigmoid):

The introduction of multi-frequency information improves the integrity of the information, and the Adaptive 1D Conv [27] enhances the precision of information transmission. This allows the MFMSBlock to more effectively utilize the feature information produced by the encoder and decoder. In addition, the point-wise convolution branch provides extra scale information, enabling the MFMSBlock to better handle targets of different scales in high-resolution remote sensing images. Applying these strategies ensures that both low-level and high-level features are fully utilized and complement each other during the fusion process, resulting in refined feature fusion.

Preparing the training set required for network training, including the original remote sensing images and their corresponding labels. Next, preparing the data by adjusting the image dimensions to fulfill the network’s input criteria, applying data augmentation and normalization, and generating pixel-level labels indicating the class to which each pixel belongs. Once the data processing is complete, input the processed data into CVMH-UNet. First, feature extraction is performed through the encoder composed of CVSSBlocks, followed by feature reconstruction through the decoder, which is also composed of CVSSBlocks. The MFMSBlock is then used to fuse the features generated by each layer of the encoder and decoder. Finally, in the final projection layer, the output channel count is modified to correspond with the class count, predicting the probability of each pixel belonging to a given class.

Throughout the training process, the performance of the predicted feature maps is judged against the ground truth labels using multi-class cross-entropy and Dice loss as metrics. Backpropagation is then performed to adjust the model parameters, ensuring the model conforms more precisely to the training data. After setting a reasonable number of iterations, the weights of the model are preserved once the loss function stabilizes, and these weights are used for predictions.

The experiments are carried out on a standalone NVIDIA GeForce RTX 4080 GPU with 16GB of RAM, and using the PyTorch platform. AdamW is used as the optimizer for training all models. The training parameters include a learning rate of 0.001, a weight decay coefficient of 0.05, a batch size of 5, and a total of 300 epochs. Additionally, to ensure rigor and fairness in the experiments, all networks are tested under the same framework, using identical data processing methods, the same optimizer, and the same loss function.

The performance of the proposed CHVM-UNet is juxtaposed with that of several state-of-the-art semantic segmentation methods. The comparison involves models that employ CNN and Transformer methodologies, such as BANet [37], ABCNet [38], UNetFormer [19], A2-FPN [39], MAResU-Net [40], MANet [26], CMTFNet [41], as well as Rs3Mamba [23] and CM-UNet [24], which are based on Vision Mamba. Unless otherwise specified, the optimal results are presented in bold within the tables, and the near-optimal results are signified with underlining.

Assessing the performance of our model involves three metrics—overall accuracy (OA), mean intersection over union (mIoU), and mean F1 score (mF1)—to compare it with the most recent state-of-the-art methods. The following explanation pertains to the calculation of OA, mIoU, and mF1 using the collective confusion matrix:

where $\mathrm{TP}_{k}$, $\mathrm{FP}_{k}$, $\mathrm{TN}_{k}$, $\mathrm{FN}_{k}$ represent the true positives, false positives, true negatives, and false negatives for class k, respectively. OA represents the proportion of accurately predicted pixels out of the overall pixel count. In addition, we evaluate the model’s complexity by considering its floating-point operations per second (FLOPs) and the total number of parameters. Higher values of OA, mIoU, and mF1 indicate better model performance, while lower values of Flops and Params indicate a more lightweight model.

The comparative analysis of different methods on the ISPRS Vaihingen dataset is summarized in Table I. As the background class (clutter) represents a minor fraction of the entire dataset, consistent with other studies, we exclude it from our accuracy reporting. As shown in Table I, our proposed CVMH-UNet achieves optimal scores in the metrics of mF1, mIoU, and OA. In contrast to the next best method, mF1 improves by 0.6%, mIoU by 0.92%, and OA by 0.45%. Our approach garners the top scores in every category. Notably, it achieved a 90.25% mIoU score for the large-scale building class, surpassing the next best method by 0.92%. The small-scale car category saw an IoU score of 68.44%, which is 0.53% higher than the method in second place. These results clearly demonstrate that CVMH-UNet effectively segments targets of different scales, meeting the current demands for high-resolution remote sensing image segmentation. To visually illustrate the differences between our method and others, Fig. 5 illustrates the segmentation outcomes from different methods on the ISPRS Vaihingen dataset. Fig. 5 presents the information that CVMH-UNet better preserves the complete contour of large-scale targets (building category) and more accurately identifies and segments each individual vehicle in small-scale targets (car category). In summary, CVMH-UNet demonstrates excellent segmentation capabilities for targets of varying scales, outperforming other compared methods.

To substantiate the effectiveness of our model, we conducted comparative experiments with other advanced methods on the larger ISPRS Potsdam dataset, with the with results tabulated in Table II. Consistent with the ISPRS Vaihingen dataset, the background class accuracy remains undisclosed. The results in Table II indicate that our method achieves the highest IoU scores in all categories except for the building category, while also obtaining the highest mF1 and mIoU scores. In contrast to the next best method, mF1 improves by 0.49%, and mIoU increases by 0.76%. Furthermore, to visually illustrate the differences between our method and others, a depiction of the segmentation results from various methods on the ISPRS Potsdam dataset can be seen in Fig. 6. By comparing the segmentation results in Fig. 6, it is evident that the proposed CVMH-UNet captures more detailed information, with segmentation results that are closer to the actual ground truth. The segmentation boundaries are more distinct, with a reduced number of misclassified pixels, thus outperforming alternative methods.

We tested the computational complexity of all models using random tensor data with a size of $3\times 256\times 256$. The computational complexity and parameter scale of each method are compared in Table VII, all assessed under equivalent environmental settings. As shown in Table VII, While CVMH-UNet requires more computational resources than streamlined networks(such as ABCNet), its segmentation accuracy is significantly higher than that of these lightweight models. In contrast to other networks founded on CNN or Transformer principles (such as CMTFNet) and networks employing complex fusion strategies (such as MANet), CVMH-UNet exhibits a clear advantage in computational complexity. In particular, when compared to other VMamba-based networks such as Rs3Mamba and CM-UNet, CVMH-UNet achieves the highest segmentation accuracy. In summary, the results of the experiments show that the proposed CVMH-UNet effectively reduces computational costs while maintaining high segmentation accuracy, successfully balancing segmentation accuracy and computational efficiency.