RS-Mamba for Large Remote Sensing Image Dense Prediction

arXiv      PDF

English | 简体中文

Introduction

The repository for this project is the code implementation of the paper RS-Mamba for Large Remote Sensing Image Dense Prediction, developed based on the environment of the VMamba.

If you find this project helpful, please give us a star ⭐️.

Main Contribution

• For the first time, the state space model (SSM) is introduced to dense prediction tasks in remote sensing, where our Recurrent State Space Model (RSM) maintains a global effective receptive field with linear complexity. ·
• Reflecting the spatial feature distribution across multiple directions in remote sensing images, RSM employs selective scanning over multiple directions. ·
• Experiments on semantic segmentation and change detection tasks has demonstrated that RSM can achieve state-of-the-art (SOTA) performance even with the simplest model architecture and training methodologies.

Updates

🌟 2024.03.29 Released the RS-Mamba project

🌟 2024.03.29 Released the code of models and training framwork

TODO

• Open source the model code
• Open source the training code

Table of Contents

• Introduction
• Updates
• TODO
• Table of Contents
• Description of Folders and Files
• Installation
  • Environment Setting
• Dataset Preparation
  • Remote Sensing Semantic Segmentation Dataset
    • Massachusetts Roads Dataset
    • WHU Dataset
    • Organization
  • Remote Sensing Change Detection Dataset
    • WHU-CD Dataset
    • LEIVR-CD Dataset
    • Organization
• Model Training and Inference
  • Training and Inference of Semantic Segmentation Model
  • Training and Inference of Change Detection Model
• FAQ
  • 1. Problems when installing selective_scan package in Vmamba
  • 2. ModuleNotFoundError: No module named 'selective_scan_cuda' when running the code
• Citation
• License

Description of Folders and Files

The folders semantci_segmentation_mamba and change_detection_mamba respectively contain the code for remote sensing semantic segmentation and change detection tasks, organized in a similar manner.

Taking the change_detection_mamba folder as an example, train.py and inference.py are the training and inference scripts respectively, rs_mamba_cd.py contains the model code, and the utils folder stores various other code files.

Within the utils folder, data_loading.py is the code file for data loading, dataset_process.py is the dataset processing file containing various functions for preprocessing change detection datasets, losses.py contains the code for loss functions, path_hyperparameter.py stores various model and training hyperparameters, dataset names, and model hyperparameters are also set therein, and utils.py includes code for training and validation.

Installation

Environment Setting

Step 1: Follow the installation instructions of the VMamba project to install the "rs_mamba" environment.

Step 2: Run the following command to install dependencies.

If you only use the model code, this step is not needed.

pip install -r requirements.txt

Dataset Preparation

Remote Sensing Semantic Segmentation Dataset

Massachusetts Roads Dataset

• Dataset Download: Massachusetts Roads Dataset。

WHU Dataset

• Dataset Download: WHU Dataset。

Organization

You need to organize the data set into the following format:

${DATASET_ROOT} # dataset root dir
├── train
    ├── image
        └── 0001.tif
        └── 0002.tif
        └── ...
    ├── label
        └── 0001.tif
        └── 0002.tif
        └── ...
├── val
    ├── image
        └──0001.tif
        └── 0002.tif
        └── ...
    ├── label
        ├── 0001.tif
        └── 0002.tif
        └── ...
├── test
    ├── image
        └── 0001.tif
        └── 0002.tif
        └── ...
    ├── label
        └── 0001.tif
        └── 0002.tif
        └── ...

Remote Sensing Change Detection Dataset

WHU-CD Dataset

• Dataset Download: WHU-CD Dataset。

LEIVR-CD Dataset

• Dataset Download: LEVIR-CD Dataset。

Organization

You need to organize the data set into the following format:

${DATASET_ROOT} # dataset root dir
├── train
    ├── t1
        └── 0001.tif
        └── 0002.tif
        └── ...
    ├── t2
        └── 0001.tif
        └── 0002.tif
        └── ...
    ├── label
        └── 0001.tif
        └── 0002.tif
        └── ...
├── val
    ├── t1
        └── 0001.tif
        └── 0002.tif
        └── ...
    ├── t2
        └── 0001.tif
        └── 0002.tif
        └── ...
    ├── label
        ├── 0001.tif
        └── 0002.tif
        └── ...
├── test
    ├── t1
        └── 0001.tif
        └── 0002.tif
        └── ...
    ├── t2
        └── 0001.tif
        └── 0002.tif
        └── ...
    ├── label
        └── 0001.tif
        └── 0002.tif
        └── ...

Model Training and Inference

All hyperparameters for model training and inference are located in the utils/path_hyperparameter.py file, with corresponding comments for explanation.

Training and Inference of Semantic Segmentation Model

Firstly, navigate to the semantic segmentation folder by running the following command in the command line:

cd semantic_segmentation_mamba

To start training, run the following code in the command line:

python train.py

If you want to debug during training, run the following command in the command line:

python -m ipdb train.py

To perform testing or inference in the command line, run the following code:

python inference.py

Training and Inference of Change Detection Model

Firstly, navigate to the change detection folder by running the following command in the command line:

cd change_detection_mamba

To start training, run the following code in the command line:

python train.py

If you want to debug during training, run the following command in the command line:

python -m ipdb train.py

To perform testing or inference in the command line, run the following code:

python inference.py

FAQ

Here we have listed some common problems in use and their corresponding solutions, and if there are frequently recurring problems in the issue, they will also be listed here.

1. Problems when installing selective_scan package in Vmamba

You can refer to issue102 and issue95 of VMamba, My approach is to ask GPT4, and get a feasible solution, using conda to update GCC, relevant questions and answers are here.

2. ModuleNotFoundError: No module named 'selective_scan_cuda' when running the code

You can refer to issue55, selective_scan_cuda is dispensable.

Citation

If you are using this project's code or performance benchmarks in your research, please refer to bibtex to reference RSMamba below.

@misc{zhao2024rsmamba,
      title={RS-Mamba for Large Remote Sensing Image Dense Prediction}, 
      author={Sijie Zhao and Hao Chen and Xueliang Zhang and Pengfeng Xiao and Lei Bai and Wanli Ouyang},
      year={2024},
      eprint={2404.02668},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

License

This project is licensed under the Apache 2.0 License。