大幅面遥感图像具有大尺度的地物特征和丰富的上下文信息,因此深度网络在大幅面遥感图像上进行密集预测具有较大挑战。遥感图像密集预测任务主要为语义分割和变化检测,主要模型包括CNN和Transformer。CNN模型具有线性复杂度,但由于它无法对上下文信息进行全局建模,难以完成大幅面遥感图像的密集预测。Transformer模型能够对上下文信息进行全局建模,但由于它的二次方复杂度,需要先把大幅面遥感图像裁剪为若干小图像,再把小图像输入到模型中处理。然而,小图像包含的上下文信息非常有限,导致Transformer模型难以有效进行大幅面遥感图像的密集预测。最近,Mamba[1]基于状态空间模型(SSM)同时实现了自然语言处理的全局建模和线性复杂度,Vim[2]和VMamba[3]将其引入计算机视觉领域,展现了强大的预测能力。
然而,遥感图像与自然图像的获取视角迥然不同:以平视姿态获取的自然图像的空间特征分布在横向或纵向,而以鸟瞰姿态获取的遥感图像的空间特征分布在任意方向。Vim和VMamba在图像的横向和纵向上进行全局建模,不适用于空间特征分布在任意方向的遥感场景。
为了解决该问题,南京大学高分辨率遥感实验室联合上海人工智能实验室,首次将SSM引入大幅面遥感图像的密集预测任务,针对遥感场景设计了全向扫描模块(OSSM),实现了SSM在多个方向对遥感图像进行全局建模。在此基础上,针对大幅面遥感图像的语义分割和变化检测任务,分别构建了RSM-SS和RSM-CD模型,如图1所示。RSM-SS使用经典的U-Net架构,RSM-CD使用经典的孪生网络架构,它们都是遥感密集预测任务中常用和简单的架构,不包含复杂的模块。RSM-SS和RSM-CD的encoder都由若干全向状态空间块(OSS)构成,OSS的内部为类Mamba结构,具有全局建模能力和线性复杂度。
图1 RSM-SS和RSM-CD的整体结构
其中,OSSM在八个方向对图像进行选择性扫描,能够在多个方向对遥感图像进行全局建模,如图2所示:大幅面遥感图像转化为tokens之后,输入的tokens在8个方向展平成8个图像token序列,并独立地输入到状态空间模型块中进行选择性扫描,所有方向的扫描结果最后加在一起,得到输出的tokens。
图2 OSSM的整体结构
为了证明OSSM的有效性,在语义分割的Massachusetts Roads数据集和变化检测的WHU-CD数据集上进行了消融实验。结果表明,与计算机视觉领域的选择性扫描模块(Vim的SS1D和VMamba的SS2D)相比,OSSM能够从更多的方向对大幅面遥感图像进行全局建模,更加适应遥感场景的空间特征多方向分布特点,在语义分割和变化检测任务中均能取得更好的效果,如图3所示。
图3 OSSM在Massachusetts Roads数据集和WHU-CD数据集上的消融实验结果
图4 (a) RSM-SS在Massachusetts Roads数据集上的对比实验结果 (b)RSM-SS在WHU数据集的上对比实验结果 (c)RSM-CD在WHU-CD数据集上的对比实验结果 (d)RSM-CD在LEVIR-CD数据集上的的对比实验结果
为了证明RSM-SS和RSM-CD的有效性,在语义分割的Massachusetts Roads数据集和WHU数据集、变化检测的WHU-CD数据集和LEVIR-CD数据集上进行了对比实验,如图4所示。结果表明,RSM-SS和RSM-CD在语义分割和变化检测任务中的表现均超过了CNN和Transformer模型,实现了最好的效果。大幅面遥感图像中的道路和建筑物具有大尺度的空间联系和多方向的空间特征,RSM-SS和RSM-CD在多个方向对遥感图像进行全局建模,从而能够有效捕获道路和建筑物在多个方向的大尺度空间特征,从而实现准确的语义分割和变化检测。
为了进一步证明SSM相较于CNN和Transformer模型在处理大幅面遥感图像方面的有效性,在变化检测的WHU-CD数据集上针对不同大小的遥感图像对比了三种模型的表现和计算量,如图5所示。结果表明,随着遥感图像增大,RSM-CD的计算量呈现与CNN相似的缓慢增长趋势,远低于Transformer的增长趋势。在大幅面遥感图像的变化检测任务上,CNN模型的表现显著弱于RSM-CD,Transformer模型由于显存的限制无法处理大幅面遥感图像,展现出RSM-CD在大幅面遥感图像密集预测任务上的有效性。
图5 WHU-CD数据集中不同图像大小情况下的对比实验结果
综上所述,针对遥感场景地物多方向分布的特点,RSM-SS和RSM-CD在多个方向对遥感图像进行全局建模,能够有效进行大幅面遥感图像的密集预测。RSM-SS和RSM-CD以简单的架构在大幅面遥感图像的密集预测任务上取得了良好的效果,证明了SSM在遥感场景中的应用潜力。
该研究成果近日以“RS-Mamba for Large Remote Sensing Image Dense Prediction”为题发表于IEEE Transactions on Geoscience and Remote Sensing。南京大学地理与海洋科学学院硕士生赵思杰为论文第一作者,南京大学地理与海洋科学学院张学良副教授、上海人工智能实验室陈浩博士为论文共同通讯作者,南京大学地理与海洋科学学院肖鹏峰教授、上海人工智能实验室白磊博士为论文合作作者。该研究得到国家自然科学基金项目(42071297)和南京大学AI & AI for Science专项项目(02091480605203)资助。
S. Zhao, H. Chen, X. Zhang, P. Xiao, L. Bai and W. Ouyang, "RS-Mamba for Large Remote Sensing Image Dense Prediction," in IEEE Transactions on Geoscience and Remote Sensing, vol. 62, pp. 1-14, 2024, Art no. 5633314, doi: 10.1109/TGRS.2024.3425540.
[1] Gu A, Dao T. Mamba: Linear-time sequence modeling with selective state spaces[J]. arXiv preprint arXiv:2312.00752, 2023.
[2] Zhu L, Liao B, Zhang Q, et al. Vision mamba: Efficient visual representation learning with bidirectional state space model[J]. arXiv preprint arXiv:2401.09417, 2024.
[3] Y. Liu, Y. Tian, Y. Zhao, et al. Vmamba: Visual state space model[J], arXiv preprint arXiv:2401.10166, 2024.
点击“阅读原文”即可跳转到论文页面