基于SwinTransformer+UNet的遥感图像语义分割

您所在的位置:网站首页 图像分割cnn 基于SwinTransformer+UNet的遥感图像语义分割

基于SwinTransformer+UNet的遥感图像语义分割

2023-11-29 20:54| 来源: 网络整理| 查看: 265

fig1

目录 摘要1.Introduction2.Related Work2.1.基于CNN的RS语义分割2.2.自注意力机制2.3.ViT 3.方法3.1.网络结构3.2.SwinTransformer block3.3.Spatial Interaction Module3.4.Feature Compression Module3.5.Relational Aggregation Module 4.Experiments4.1.Datasets4.1.1.Vaihingen Dataset4.1.2.Potsdam Dataset 4.2.实现细节4.2.1.训练设置4.2.2.损失函数4.2.3.评价指标 4.3.实验结果 个人总结

摘要

全局上下文信息(context information)对于遥感图像的语义分割至关重要。然而,现有的大多数方法都依赖于卷积神经网络(CNN),由于卷积运算的局部性(locality),直接获取全局上下文信息充满了挑战性。受具有强大全局建模能力的Swin transformer的启发,我们提出了一种新的 Remote Sensing image 语义分割框架,称为ST-U网络,该框架将Swin transformer嵌入到基于CNN的经典UNet中。ST UNet构成了一种新型的SwinTransformer和CNN并行的双编码器(dual encoder)。首先,我们提出了一种空间交互模块(SIM,spatial interaction module),该模块通过建立像素级的相关性(pixel-level correlation)来对 SwinTransformer block 中的空间信息进行编码,以增强被遮挡对象的特征表示能力。其次,我们构造了一个特征压缩模块(FCM,feature compression module),以减少细节信息的丢失,并在SwinTransformer的patch token的下采样过程中浓缩更多的小尺度特征,从而提高了小尺度地物(ground object)的分割精度。最后,作为双编码器之间的桥梁,设计了一个关系聚合模块(RAM,relational aggregation module),将来自SwinTransformer的全局依赖项(global dependencies)层次化地集成到CNN的特征中。我们的ST-UNet分别在ISPRS Vaihingen和Potsdam数据集上实现了明显提升。

1.Introduction

随着航天技术和传感器技术的快速发展,研究者可以很容易地收集大量的高质量遥感图像,它反映了生态环境的状态和人类活动的痕迹。学习这些图像中包含的知识并有效地筛选感兴趣的信息已成为遥感图像智能解译(intelligent interpretation)的重点。语义分割作为一种可行的解决方案受到了广泛的关注。其核心目标是识别图像中每个像素的语义类别。目前,遥感图像语义分割被应用于许多现实场景,如城市规划、灾害评估和农业生产。

近年来,卷积神经网络(CNN)的快速发展为语义分割提供了技术支持。特别是,全卷积网络(FCN)起到了至关重要的作用。随后,研究者们提出了许多创新性的工作。在此过程中,编码器-解码器结构表现出优异的分割性能,并逐渐成为语义分割网络中的一种流行结构配置。编码器用于提取特征,解码器在融合高级语义和低级空间信息的同时,尽可能精细地恢复图像分辨率。例如,UNet利用解码器通过skip connections来学习相应编码过程中的空间相关性(spatial correlation)。Deeplab V3+引入了基于Deeplab V3的解码器,以集成空间特征,其显著提高了网络性能。

然而,地物的特殊性(小尺度small scale、高相似性high similarity和相互遮挡mutual occlusion)对遥感图像的语义分割提出了新的挑战,如图1所示。基于CNN的模型在特征提取过程中执行特征下采样,以减少计算量,这很容易导致小尺度特征被丢弃。不同语义类别的地物可能具有相似的大小、材质和光谱特征,难以区分。此外,遮挡问题通常会导致语义歧义(semantic ambiguity)。因此,需要更多的全局上下文信息和精细的空间特征作为语义推理的线索。 fig2

图1:RS图像示例,其中第一行图像取自ISPRS Vaihingen数据集,第二行取自ISPRS Potsdam数据集。 a:屋顶上的天窗外观类似于“汽车Car”和“不透水表面Impervious surface”; b:“建筑Building”和“不透水表面”的材质(material)相同,“树木Tree”几乎在“低矮植被Low Vegetation”中看不见; c:“汽车”的大部分被“树Tree”遮住了。

CNN在空间位置表示方面具有优势,但由于卷积运算的局部性,很难直接建模全局语义的相互作用信息和上下文信息。现有的方法采用注意力机制来解决这个问题。DANet用并行的通道注意力和位置注意力构造长范围依赖。此外,多尺度特征融合策略也是一种常用的方法。Zhao等人使用不同级别的特征来综合描述特征。PSPNet及其改进版SuperNet通过金字塔pooling模块利用上下文信息。这些方法从CNN获得的局部特征中聚合全局信息,而不是直接编码全局上下文。因此,很难从背景复杂的遥感图像中获得清晰的全局场景信息。

最近,transformer的成功为全局关系建模开辟了新的研究思路。Transformer是自然语言处理领域中一种流行的序列预测模型。Carion等人提出了DETR,它利用Transformer的编码器-解码器结构来模拟序列中元素之间的相互作用。与基于CNN的模型类似,Chen等人设计了一种双分支Transformer结构,以学习不同尺度的特征。事实证明,多尺度特征表示对ViT也是有效的。根据这种结构,构建了SwinTransformer,并在多个密集预测任务中显示出巨大的潜力。目前,基于Swin transformer的模型在医学图像分割方面取得了很大的进展,但其在RS图像上的分割潜力尚未得到证实。

在本文中,为了缓解CNN在全局建模方面的不足,我们提出了一种新的RS图像语义分割网络框架ST-UNet,该框架利用Swin transformer来辅助UNet。如前所述,UNet是基于CNN的U形解码器-编码器网络,通过skip-connection实现编码器和解码器之间的特征融合。我们将UNet中的编码器作为主编码器,Swin变压器作为辅助编码器,形成并行双编码器结构。具体地,我们通过关系聚合模块(RAM)构建了从辅助编码器到主编码器的单向信息流,RAM是ST-UNet的关键组件。此外,将SIM(空间交互模块)附加到SwinTransformer上,探索全局特征的空间相关性,并使用FCM(特征压缩模块)提高小尺度对象的分割精度。

本文贡献:

我们构建了空间交互模块(SIM)来关注空间维度上像素级的特征相关性,从而缓解地物遮挡引起的语义歧义。此外,SIM还补偿了Swin transformer受其窗口机制限制的全局建模能力。我们在辅助编码器中提出了一个特征压缩模块(FCM),以缓解patch token下采样过程中小规模特征的遗漏。FCM可以收集小规模对象的更多特征,减少细致信息的丢失。为了提取遥感图像中的判别特征,我们设计了RAM,它从辅助编码器(auxiliary encoder)中提取通道的关系信息,作为全局线索来引导主编码器。RAM能有效地识别相似性较高的地物。 2.Related Work 2.1.基于CNN的RS语义分割

由于一些数据集和竞赛的发布,例如IGARSS数据融合竞赛,SpaceNet竞赛、DeepGlobe竞赛、ISPRS benchmarks。基于CNN的遥感图像语义分割受到了广泛关注。Zhang等人采用HRNet中的多分支并行卷积结构来生成多尺度特征图,并设计了一个自适应空间pooling模块来聚合更多的局部上下文。Maggiori等人将MLP引入分割网络,以产生更好的分割结果。Qi等人提出了一种基于递归神经网络和三维卷积的空间信息推理结构,用于学习全局空间上下文和局部视觉特征,有效解决了道路检测中的遮挡问题。

此外,一些研究人员还关注小尺度特征的特征提取。Kammpfmeyer等人组合了基于patch的像素分类和像素到像素分割,这引入了不确定映射,以在小规模对象上实现高性能。受UNet的启发,Dong等人提出了DenseU网络,该网络通过密集融合策略实现小规模特征的聚合。FactSeg提出了一种由前景激活分支和语义修复分支组成的对称双分支解码器。这两个分支通过skip-connection进行多尺度特征融合,从而提高了小尺度对象分割的精度。Chen等人提出了一种自适应感受野的卷积网络,以实现大规模对象和小规模对象特征提取之间的折衷。

物体边界的预测也是一个值得注意的方面。Marmanis等人在语义分割网络中明确引入了边缘检测模块,以监督边界特征学习。ERN提出了两个简单的边缘损失强化模块,以加强物体边界的保护。

上述基于CNN的方法促进了遥感图像语义分割的发展。我们没有放弃CNN,而是致力于提出一个新的框架(例如最近基于Transformer的方法),以继承两者的优点。

2.2.自注意力机制

近年来,自注意机制在计算机视觉任务中得到了广泛的应用。Zhao等人和Li等人提出了视频字幕的区域级注意和帧级注意。Zhao等人探讨了成对自注意和局部自注意在图像识别中的有效性。SENet通过全局平均pooling层表示通道之间的关系,以自动了解不同通道的重要性。CBAM将通道级注意和空间级注意应用于自适应特征细化。Ding等人提出了patch注意模块,以突出特征图的焦点区域。Panboonyuen等人在GCN框架的每个阶段引入了通道注意块,以分层优化特征图。Su等人将注意力集中在一小批图像中的相似对象上,并通过自注意机制对它们之间的交互信息进行编码。与上述方法不同,我们考虑了垂直和水平像素级的自注意力,并结合多个pooling层来提取通道依赖。

2.3.ViT

Transformer首先被提出用于机器翻译任务,并超越了之前基于复杂递归或CNN的序列模型。随后,基于Transformer的模型在自然语言处理的各个领域得到了广泛的应用。标准transformer块由一个多头自注意(MSA)、一个多层感知器(MLP)和一个层规范化(LN)组成,如图2a所示。MSA在建立输入和输出序列之间的全局依赖关系方面发挥了关键作用。 fig3

图2:a为标准的Transformer block,b为SwinTransformer block,c为本文提出的SIM。

最近的研究表明,Transformer也适用于计算机视觉任务。Dosovitskiy等人是第一个将transformer应用于图像分类的人,通过拆分和展平(splitting and flattening)将图像数据转换为一系列token。这样的ViT在大规模数据集的预训练下实现了最先进的性能。之后,Chen等人探索了一种基于通用Transformer的图像处理任务预训练方法。T2T-ViT递归地将相邻token聚合为一个token,以建模由周围token表示的局部结构特征,并减少token的数量。Touvron等人提出了一种新的基于token的蒸馏策略,以提高原始ViT在较小Imagenet-1k数据集上的训练效率。

尽管如此,ViT仍然需要为密集的预测任务付出巨大的训练成本。它只输出与预测目标不匹配的低分辨率特征(与输入图像的分辨率相同)。一些工作修改了ViT体系结构,以适应密集的预测任务,如语义分割和对象检测。SETR将Transformer视为编码器,在每一层中对全局上下文进行建模,并结合一个简单的解码器形成语义分割网络。PVT模仿CNN backbone的特性,将金字塔结构引入ViT,以获得多尺度特征图。具体来说,它通过patch embedding层灵活地控制Transformer的长度。虽然PVT减少了计算资源的消耗,但复杂度依然是图像大小的二次方。因此,Liu等人提出了基于移位窗口策略的SwinTransformer,该策略将MSA的计算限制在非重叠窗口,同时允许跨窗口信息交互。Swin transformer仅具有线性计算复杂度,在各种视觉任务中都取得了优异的性能,包括图像分类、目标检测和语义分割。放弃传统Transformer中的MSA,它利用基于窗口的MSA(W-MSA)和移位的W-MSA(SW-MSA),如图2b所示。

Cao等人和Lin等人以Swin transformer为主干,开发了用于医学图像语义分割的U形编解码器框架。特别是,Lin等人利用基于Swin transformer的双编码器和两个输入图像尺度来提取不同语义尺度的特征表示。

然而,TransUNet和TransFuse指出,纯Transformer分割网络产生的结果并不令人满意,因为Transformer只关注全局建模,缺乏定位能力。因此,他们创建了CNN和transformer的混合结构。TransUnet将CNN和transformer顺序堆叠,形成一个新的编码器结构,而TransFuse同时执行这两个功能,并试图融合这两个功能。此外,TransFuse在Transformer分支的解码器中使用了简单的渐进式上采样,以恢复空间分辨率。

受这些优秀作品的启发,我们采用了由Swin transformer块组成的辅助编码器,为基于CNN的主编码器提供全局上下文信息。据我们所知,本文提出的ST-UNet是第一个将SwinTransformer应用到遥感图像分割任务中的,它弥补了纯CNN的不足,提高了分割精度。

3.方法 3.1.网络结构

fig4

图3:ST-UNet的架构。ST-UNet包含三个重要的模块:RAM,SIM,FCM。

ST-UNet的整体架构如图3所示。作为SwinTransformer和UNet的混合体,我们的ST-UNet遵循UNet的优秀结构,其中kip-connection层连接编码器和解码器。特别地,ST-UNet构造了由基于CNN的残差网络和SwinTransformer组成的双编码器,该编码器通过RAM传输信息,以充分获得RS图像的判别特征。此外,我们还设计了SIM和FCM,以进一步提高SwinTransformer的性能。

对于给定的RS图像 X ∈ R H × W × 3 X\in R^{H\times W\times 3} X∈RH×W×3,ViT将图像划分为非重叠的patch,并用序列数据的"token"来比喻。图像的token之间是存在关联的。patch中同一对象的像素通常聚集在一起,具有很强的语义相关性。因此,为了避免在初始输入阶段失去语义信息的连续性,我们通过卷积从每个图像中获得存在重叠的patch token。在实验中,patch大小被设为 8 × 8 8\times 8 8×8,重叠率为 50 % 50\% 50%。然后,线性embedding层将这些patch展平并投影到 C 1 C_{1} C1​维。这些patch token被放入由SwinTransformer block堆叠的辅助编码器中。辅助编码器有四个特征提取阶段,每个阶段的输出定义为 S n S_{n} Sn​。标准SwinTransformer block包括两种类型,即window-based的transformer(W-Trans)和shifted的W-Trans(SW-Trans)。特别是,我们提出用SIM建立像素级信息交换(pixel-level information exchange),该信息连接到SwinTransformer block的输出。SIM可以有效地弥补基于窗口的自注意的局限性,缓解遮挡(occlusion)引起的语义歧义问题(semantic ambiguity)。此外,为了在与主编码器的特征分辨率匹配的同时获得多尺度特征,我们通过缩短patch token的长度来构建FCM,形成一个四阶段的分层特征编码结构。在此过程中,FCM可以减少对小尺度对象特征的遗漏。stage n n n 输出的分辨率为 ( H 2 n + 1 , W 2 n + 1 ) (\frac{H}{2^{n+1}},\frac{W}{2^{n+1}}) (2n+1H​,2n+1W​),维数是 2 n − 1 C 1 2^{n-1}C_{1} 2n−1C1​。

原始RS图像 X X X首先以大小上一半的压缩量馈送至ResNet50,以获得主编码器中的深度特征。第 n n n个残差模块的输出特征图被表示为 A n ∈ R H 2 n + 1 × W 2 n + 1 × 2 n − 1 C 2 A_{n}\in R^{\frac{H}{2^{n+1}}\times\frac{W}{2^{n+1}}\times 2^{n-1}C_{2}} An​∈R2n+1H​×2n+1W​×2n−1C2​,这里 C 2 = 128 C_{2}=128 C2​=128。然后, A n A_{n} An​和 S n S_{n} Sn​被输入到RAM,融合结果返回主编码器。RAM模块作为主编码器和辅助编码器之间的桥梁,通过可变形卷积和通道注意力机制建立连接。

经过以上编码阶段,我们得到特征 F ∈ R H 32 × W 32 × 1024 F\in R^{\frac{H}{32}\times\frac{W}{32}\times 1024} F∈R32H​×32W​×1024,该特征在卷积层之后被送入解码器。我们将其输入到 2 × 2 2\times 2 2×2反卷积层扩大分辨率。继UNet之后,ST-UNet利用skip-connection层来连接编码器和解码器,同时减少通过 3 × 3 3×3 3×3卷积层的通道数量。此处,每个卷积层都伴随一个batch normalization层和ReLU层。以上过程被执行4次,特征 F F F逐渐恢复到 F ′ ∈ R H 2 × W 2 × 64 F'\in R^{\frac{H}{2}\times\frac{W}{2}\times 64} F′∈R2H​×2W​×64。最后我们用 3 × 3 3\times 3 3×3卷积层和线性插值上采样获得分割结果。

3.2.SwinTransformer block

如前面所述,标准transformer block包含了MSA,MLP和LN(见图2a,LN为layer norm)。因此,我们将第 l l l层的输出 s l s^{l} sl表示为: s ^ l = M S A ( L N ( s l − 1 ) ) + s l − 1 \widehat{s}^{l}=MSA(LN(s^{l-1}))+s^{l-1} s l=MSA(LN(sl−1))+sl−1 s l = M L P ( L N ( s ^ l ) ) + s ^ l s^{l}=MLP(LN(\widehat{s}^{l}))+\widehat{s}^{l} sl=MLP(LN(s l))+s l标准transformer block使用MSA计算所有token之间的全局自注意力,这导致计算复杂度为token数量的二次方,这对于高分辨率遥感图像的特征提取是致命的。为了高效建模,Swin transformer提出了具有两种分区配置的W-MSA,以取代普通MSA,即常规窗口配置(W-MSA)和移位窗口配置(SW-MSA)。他们在窗口内进行自注意,但忽略了窗口外的token,每个窗口只覆盖 D × D D×D D×D patches。在实验中, D = 8 D=8 D=8。如图2b所示,在连续的Swin Transformer block中交替执行W-MSA和SW-MSA,以增强跨窗口的信息连接。为了便于区分,我们将两个Swin transformer block重命名为W-Trans block 和SW-Trans block。其表示如下: s ^ l = W − M S A ( L N ( s l − 1 ) ) + s l − 1 \widehat{s}^{l}=W-MSA(LN(s^{l-1}))+s^{l-1} s l=W−MSA(LN(sl−1))+sl−1 s l = M L P ( L N ( s ^ l ) ) + s ^ l s^{l}=MLP(LN(\widehat{s}^{l}))+\widehat{s}^{l} sl=MLP(LN(s l))+s l s ^ l + 1 = S W − M S A ( L N ( s l ) ) + s l \widehat{s}^{l+1}=SW-MSA(LN(s^{l}))+s^{l} s l+1=SW−MSA(LN(sl))+sl s l + 1 = M L P ( L N ( s ^ l + 1 ) ) + s ^ l + 1 s^{l+1}=MLP(LN(\widehat{s}^{l+1}))+\widehat{s}^{l+1} sl+1=MLP(LN(s l+1))+s l+1其中, s l s^{l} sl代表W-Trans block的输出, s l + 1 s^{l+1} sl+1代表SW-Trans block的输出。

3.3.Spatial Interaction Module

Swin transformer块在有限的窗口内建立patch token的关系,有效地减少了内存开销。然而,这种方法在一定程度上削弱了transformer的全局建模能力(即使它采用了规则窗口和移位窗口的交替执行策略)。此外,遥感图像中地物的遮挡会导致边界模糊,这需要一些空间信息来消除。因此,我们提出跨W-Trans block和SW-Trans block的SIM,以进一步增强信息交换,同时编码更精确的空间信息。SIM在两个空间维度引入了注意力,以考虑像素之间的关系,而不仅仅是patch token之间,使Transformer更适合图像分割任务。SIM的组件如图4所示。 fig5

图4:SIM的结构。

给定stage n n n,我们首先reshape W-Trans block的输入特征 s l − 1 ∈ R ( h × w ) × c 1 s^{l-1}\in R^{(h\times w)\times c_{1}} sl−1∈R(h×w)×c1​为 z ∈ R h × w × c 1 z\in R^{h\times w\times c_{1}} z∈Rh×w×c1​,此处, c 1 = 2 n − 1 C 1 , h = H 2 n + 1 , w = W 2 n + 1 c_{1}=2^{n-1}C_{1},h=\frac{H}{2^{n+1}},w=\frac{W}{2^{n+1}} c1​=2n−1C1​,h=2n+1H​,w=2n+1W​。特征 z z z被输入到 3 × 3 3\times 3 3×3的空洞卷积(dilation rate=2),通过较大的感受野重构特征图的结构信息。此外,为了降低计算成本,通道数被缩减为 c 1 / 2 c_{1}/2 c1​/2。然后,应用全局平均pooling运算来获取特征地图在空间方向(垂直和水平)上的统计信息。具体而言,各方向的计算公式如下所示: v h i k = 1 w ∑ j = 0 w − 1 z ^ k ( i , j ) v_{h_{i}}^{k}=\frac{1}{w}\sum_{j=0}^{w-1}\widehat{z}^{k}(i,j) vhi​k​=w1​j=0∑w−1​z k(i,j) v w j k = 1 h ∑ i = 0 h − 1 z ^ k ( i , j ) v_{w_{j}}^{k}=\frac{1}{h}\sum_{i=0}^{h-1}\widehat{z}^{k}(i,j) vwj​k​=h1​i=0∑h−1​z k(i,j)其中, i , j , k i,j,k i,j,k分别是垂直方向、水平方向和通道的索引。这里, 0 ≤ i < h , 0 ≤ j < w , 0 ≤ k < c 1 2 0\leq i



【本文地址】


今日新闻


推荐新闻


CopyRight 2018-2019 办公设备维修网 版权所有 豫ICP备15022753号-3