Skip to content

Latest commit

 

History

History
416 lines (300 loc) · 13.9 KB

README_Chinese.md

File metadata and controls

416 lines (300 loc) · 13.9 KB
Raw

RS-Mamba for Large Remote Sensing Image Dense Prediction



arXiv      PDF


GitHub stars license arXiv

English | 简体中文

简介

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

如果你觉得本项目对你有帮助，请给我们一个 star ⭐️。

主要贡献
  • 首次将状态空间模型（SSM）引入遥感密集预测任务，RSM在具有全局有效感受野的同时具有线性复杂度，从而能够处理保留了丰富上下文信息的大遥感图像 ·
  • 针对遥感图像的空间特征分布在多个方向的特点，RSM通过在多个方向对图像进行选择性扫描，从而能够在多个方向上进行全局建模，并提取出多个方向的大尺度空间特征 ·
  • 语义分割和变化检测任务上的实验证明了，RSM在使用最简单的模型架构和训练方法的情况下，仍然能够达到SOTA效果，具有很大的潜力

更新日志

🌟 2024.03.29 发布了 RS-Mamba 项目

🌟 2024.04.06 模型代码和训练框架代码已开源

TODO

  • 开源模型代码
  • 开源训练框架

目录

文件夹与文件说明

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 数据集

WHU 数据集

组织方式

你需要将数据集组织成如下的格式：

${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 数据集

LEIVR-CD 数据集

组织方式

你需要将数据集组织成如下的格式：

${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

常见问题

我们在这里列出了使用时的一些常见问题及其相应的解决方案，后续如果issue中有经常出现的问题，也会在这里列出来。

1. 安装VMamba环境中的selective_scan库出现问题

可以参考VMamba的issue102, issue95, 我的做法是询问GPT4之后，得到了可行的解决方法，使用conda更新GCC即可，相关询问和回答在这里.

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

可以参考VMamba的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 开源许可证。

论文解读

关于这篇论文的中文解读，你可以参照这篇知乎文章

关于SSM和Mamba的学习，可以参照这两篇知乎回答：如何理解 Mamba 模型 Selective State Spaces?, 如何理解语言模型的训练和推理复杂度?