A Lightweight Semantic Segmentation Network Based on Self-Attention Mechanism and State Space Model for Efficient Urban Scene Segmentation

Abstract

In the semantic segmentation of remote sensing images, methods based on convolutional neural networks (CNNs) and Transformers have been extensively studied. Nevertheless, CNN struggles to capture the global context due to its local feature extraction, while Transformer is constrained by the complexity of quadratic calculations. Recently, there has been a great deal of interest in Mamba-based state space models. However, the existing Mamba-based methods do not adequately consider the significance of local information in remote sensing image segmentation tasks. In this article, a codec style network UMFormer is constructed for the semantic segmentation of remote sensing images. Specifically, UMFormer employs the ResNet18 as the encoder, with the objective of performing a preliminary image feature extraction. Subsequently, a self-attention mechanism is optimized to extract the global information pertaining to the objects of disparate sizes within the context of a multiscale condition. For fusing the codec feature map information, another attention structure is built to reconstruct the space information and to capture the relative position relationship. Finally, a decoder based on Mamba is designed to effectively model both global and local information. Concurrently, a feature fusion mechanism utilizing feature similarity is devised with the objective of embedding local information into global ones. Numerous experiments on UAV Imagery Dataset (UAVid), Vaihingen, and Potsdam datasets have demonstrated that the proposed UMFormer exhibits enhanced accuracy while maintaining an efficient running speed. The code will be freely available at: https://github.com/takeyoutime/UMFormer