Official_Remote_Sensing_Mamba - 创新的大型遥感图像密集预测模型

RS-Mamba用于大型遥感图像密集预测

arXiv PDF

English | 简体中文

简介

本项目的代码仓库是论文RS-Mamba for Large Remote Sensing Image Dense Prediction的代码实现，基于VMamba的环境开发。

如果您觉得这个项目有帮助，请给我们一个星标 ⭐️。

主要贡献

首次将状态空间模型（SSM）引入遥感密集预测任务，我们的循环状态空间模型（RSM）以线性复杂度保持全局有效感受野。
反映遥感图像多方向的空间特征分布，RSM采用多方向的选择性扫描。
在语义分割和变化检测任务的实验表明，即使采用最简单的模型架构和训练方法，RSM也能达到最先进（SOTA）的性能。

更新

🌟 2024.03.29 发布RS-Mamba项目

🌟 2024.03.29 发布模型和训练框架代码

待办事项

开源模型代码
开源训练代码

简介
更新
待办事项
目录
文件夹和文件说明
安装
- 环境设置
数据集准备
- 遥感语义分割数据集
- 遥感变化检测数据集
模型训练和推理
- 语义分割模型的训练和推理
- 变化检测模型的训练和推理
常见问题
- 1. 在Vmamba中安装selective_scan包时遇到问题
- 2. 运行代码时出现ModuleNotFoundError: No module named 'selective_scan_cuda'
引用
许可证

文件夹和文件说明

semantci_segmentation_mamba和change_detection_mamba文件夹分别包含遥感语义分割和变化检测任务的代码，组织方式类似。

以change_detection_mamba文件夹为例，train.py和inference.py分别是训练和推理脚本，rs_mamba_cd.py包含模型代码，utils文件夹存储各种其他代码文件。

在utils文件夹中，data_loading.py是数据加载的代码文件，dataset_process.py是数据集处理文件，包含各种预处理变化检测数据集的函数，losses.py包含损失函数的代码，path_hyperparameter.py存储各种模型和训练超参数，数据集名称和模型超参数也在其中设置，utils.py包括训练和验证的代码。

安装

环境设置

步骤1: 按照VMamba项目的安装说明安装"rs_mamba"环境。

步骤2: 运行以下命令安装依赖项。

如果您只使用模型代码，则不需要此步骤。

pip install -r requirements.txt

数据集准备

遥感语义分割数据集

马萨诸塞州道路数据集

数据集下载: Massachusetts Roads Dataset。

武汉大学数据集

数据集下载: WHU Dataset。

组织结构

您需要将数据集组织成以下格式：

${DATASET_ROOT} # 数据集根目录
├── 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
        └── ...

遥感变化检测数据集

WHU-CD数据集

数据集下载: WHU-CD Dataset。

LEIVR-CD数据集

数据集下载: LEVIR-CD Dataset。

组织结构

您需要将数据集组织成以下格式：

${DATASET_ROOT} # 数据集根目录
├── 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
        └── ...

模型训练和推理

所有模型训练和推理的超参数都位于utils/path_hyperparameter.py文件中，有相应的注释进行解释。

语义分割模型的训练和推理

首先，通过在命令行中运行以下命令进入语义分割文件夹：

cd semantic_segmentation_mamba

要开始训练，在命令行中运行以下代码：

python train.py

如果您想在训练过程中进行调试，请在命令行中运行以下命令：

python -m ipdb train.py

要在命令行中进行测试或推理，请运行以下代码：

python inference.py

变化检测模型的训练和推理

首先，通过在命令行中运行以下命令进入变化检测文件夹：

cd change_detection_mamba

要开始训练，在命令行中运行以下代码：

python train.py

如果您想在训练过程中进行调试，请在命令行中运行以下命令：

python -m ipdb train.py

要在命令行中进行测试或推理，请运行以下代码：

python inference.py

work_dir：模型训练的输出路径，通常无需修改。
code_root：代码根目录，需修改为本项目根目录的绝对路径。
data_root：数据集根目录，需修改为数据集根目录的绝对路径。
batch_size：单卡的批量大小，需根据显存大小调整。
max_epochs：最大训练轮数，通常无需修改。
vis_backends/WandbVisBackend：网络端可视化工具配置，取消注释后，需在wandb官网注册账号，可在浏览器中查看训练过程的可视化结果。
model/backbone/arch：模型骨干网络类型，需根据所选模型修改，包括b、l、h。
model/backbone/path_type：模型路径类型，需根据所选模型修改。
default_hooks-CheckpointHook：模型训练过程中的检查点保存配置，通常无需修改。
num_classes：数据集类别数，需根据数据集类别数修改。
dataset_type：数据集类型，需根据数据集类型修改。
resume：是否断点续训，通常无需修改。
load_from：模型预训练检查点路径，通常无需修改。
data_preprocessor/mean/std：数据预处理的均值和标准差，需根据数据集的均值和标准差修改，通常无需修改，参考Python脚本。

部分参数继承自_base_，可在基础配置文件文件夹中找到。

单卡训练

python tools/train.py configs/rsmamba/name_to_config.py  # name_to_config.py为你想使用的配置文件

多卡训练

sh ./tools/dist_train.sh configs/rsmamba/name_to_config.py ${GPU_NUM}  # name_to_config.py为你想使用的配置文件，GPU_NUM为使用的GPU数量

其他图像分类模型

如需使用其他图像分类模型，可参考MMPretrain进行模型训练，也可将其Config文件放入本项目的configs文件夹，然后按上述方法进行训练。

模型测试

单卡测试：

python tools/test.py configs/rsmamba/name_to_config.py ${CHECKPOINT_FILE}  # name_to_config.py为你想使用的配置文件，CHECKPOINT_FILE为你想使用的检查点文件

多卡测试：

sh ./tools/dist_test.sh configs/rsmamba/name_to_config.py ${CHECKPOINT_FILE} ${GPU_NUM}  # name_to_config.py为你想使用的配置文件，CHECKPOINT_FILE为你想使用的检查点文件，GPU_NUM为使用的GPU数量

图像预测

单张图像预测：

python demo/image_demo.py ${IMAGE_FILE} configs/rsmamba/name_to_config.py --checkpoint ${CHECKPOINT_FILE} --show-dir ${OUTPUT_DIR}  # IMAGE_FILE为你想预测的图像文件，name_to_config.py为你想使用的配置文件，CHECKPOINT_FILE为你想使用的检查点文件，OUTPUT_DIR为预测结果的输出路径

多张图像预测：

python demo/image_demo.py ${IMAGE_DIR} configs/rsmamba/name_to_config.py --checkpoint ${CHECKPOINT_FILE} --show-dir ${OUTPUT_DIR}  # IMAGE_DIR为你想预测的图像文件夹，name_to_config.py为你想使用的配置文件，CHECKPOINT_FILE为你想使用的检查点文件，OUTPUT_DIR为预测结果的输出路径

常见问题

这里列出了一些使用中常见的问题及其解决方案，如果issue中有频繁出现的问题，也会在此列出。

1. Vmamba中安装selective_scan包时的问题

你可以参考VMamba的issue102和issue95。我的方法是询问GPT4，并得到一个可行的解决方案，使用conda更新GCC，相关问答在这里。

2. 运行代码时出现ModuleNotFoundError: No module named 'selective_scan_cuda'

你可以参考issue55，selective_scan_cuda是可有可无的。

引用

如果你在研究中使用了本项目的代码或性能基准，请参考以下bibtex引用RSMamba。

@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}
}

许可证

本项目遵循Apache 2.0许可证。