基于渐进生长Transformer Unet的遥感图像建筑物分割

叶 宽，杨 博，谢 欢，朱 戎，赵 蕾，张青月，赵 杰

(1.国网北京市电力公司电力科学研究院，北京 100075；2.国网新源控股有限公司检修分公司，北京 100067；3.北京大学 大数据分析与应用技术国家工程实验室，北京 100871)

0 引言

高分辨率遥感图像具有目标大小不一、纹理复杂多样和树木遮挡等不稳定因素，使得基于深度学习的目标分割任务具有较大挑战，精确的遥感目标分割对于军事、航天和城市规划都有重要的意义。近年来，随着深度卷积神经网络的不断发展，基于深度学习的目标检测和物体分割应用越来越广泛，如VGG[1]、基于残差连接的ResNet[2]、改进了编码解码结构的SegNet[3]、带有跳跃连接的编码-解码结构Unet[4]、带有空间搜索的NAS-Unet[5]、具有嵌套和密集链接的UNet++[6]、具有深度监督的UNet 3+[7]以及带有3D混合残余注意力的RA-UNet[8]等，这些研究使得深度学习技术对图像目标分割任务的结果不断提升。

许多专家利用深度学习方法对遥感图像的各种目标进行分割研究，不断提升分割效果。韩彬彬等[9]将基于残差密集空间金字塔的卷积神经网络应用于城市地区遥感图像分割任务，利用了遥感图像不同尺度下的特征。范自柱等[10]通过在Unet网络中添加特征金字塔结构，使用基于卷积神经网络的图像分割方法提取遥感图像中典型土地光谱信息和空间信息来识别遥感卫星图像。袁伟等[11]将不同尺度的Unet融合在一起提出一种多尺度自适应的遥感语义分割模型。刘航等[12]引入自适应感受野机制和通道注意力模块可以减少背景特征的干扰，提升遥感图像的分割精度。余帅等[13]提出一种基于多级通道注意力的遥感图像分割模型，有效解决目标遮挡和小目标难分割问题。何青等[14]利用多层次编码解码结构提取不同尺度特征，提出一种基于残差分组卷积的高分辨率遥感影像建筑物分割模型。

为进一步提高建筑物边缘提取效果，有效解决不同大小目标和被遮挡目标分割错误等问题，本文提出了一种基于渐进生长机制的Transformer Unet(PGT-Unet)遥感图像分割网络，整个流程包括4个训练阶段。首先对输入图像进行下采样处理，之后从(64 pixel×64 pixel)到(512 pixel×512 pixel)逐渐增大输入图像的分辨率，网络模型也从1层Transformer Unet逐渐增大至4层Transformer Unet，每个阶段模型收敛之后的参数直接迁移到下一阶段的同样卷积层中，使得模型在训练初期能够学习到大尺度的结构信息，后期的学习越来越集中到精确的细节特征，可以改善建筑物边缘提取效果，提升被遮挡目标和不同大小目标的分割效果。

本文工作的贡献总结如下：

① 提出了一个基于PGT-Unet的遥感图像分割网络来促进遥感图像的语义分割精度，并获得了显著的性能；

② 使用一种渐进生长的机制进行网络训练，使得神经网络可以由粗到细逐步学习不同尺度的目标结构信息；

③ 在模型的编码器、解码器和瓶颈模块中引入了Transformer Block进行多尺度的特征提取和特征融合，进一步增强模型提取不同大小目标特征的能力，使网络能够更好地区分特征之间的重要程度，从而聚焦有用特征。

1 基本原理

1.1 渐进式增长机制

传统的卷积核和可变形的卷积核都是同时提取所有尺度的目标信息，容易造成信息冗余，并且对于网络结构较深的模型，势必会造成很大的计算量和存储的浪费。在传统神经网络模型的基础上引入渐进式生长机制[15]的训练方式可以达到不同尺度多级特征提取的目的。Transformer Block结构如图1所示。

图1 Transformer Block结构Fig.1 Illustration of Transformer Block

通过设置4个训练阶段，在第1个阶段低分辨率(64 pixel×64 pixel)输入图像经过一层编码模块、瓶颈模块和解码模块得到建筑物粗分割结果；在第2个阶段增大输入图像的分辨率至(128 pixel×128 pixel)，给模型增加一层编码模块和一层解码模块，将上一阶段训练所得参数直接迁移到新模型中的对应层中，并给已经训练好的模型参数设置较小的学习速率，给新加的卷积层参数设置较大的学习速率，可以得到较大分辨率的建筑物精细分割结果；第3阶段和第4阶段以此类推，渐进式增大图像的分辨率(256 pixel×256 pixel)与(512 pixel×512 pixel)和模型的深度，最终得到越来越精细的建筑物分割轮廓。

这种训练方式允许网络先学习大尺度的图像粗结构信息，之后的阶段将注意力集中到相邻尺度的越来越细节的特征中，而不是让网络同时学习所有尺度的信息。在每个阶段模型接收不同大小的输入图像，从而可以分步地学习到不同大小目标区域的多尺度信息，在不增加额外参数量和计算量的情况下，使得模型更快收敛，具有更强的泛化能力和稳定性。

1.2 Transformer Unet模型

渐进生长机制下的每个阶段采用基于Transformer的U型架构——Transformer Unet模型进行遥感图像特征提取和建筑物分割，模型由编码器模块、瓶颈模块、解码器模块和跳过链接构成，基本单元是Transformer Block结构。编码器模块由一系列成组的下采样模块(步长为2的卷积层)和卷积模块Transformer Block构成，进行特征提取学习全局上下文信息和局部细节信息。瓶颈模块由2组卷积模块Transformer Block构成，降低模型参数量，增加模型的非线性表达能力。解码器模块由一系列成组的卷积模块Transformer Block和上采样模块(双线性插值)构成，进行图像重建和建筑物目标分割。跳过链接联通编码器和解码器，提取的上下文特征通过跳跃链接与编码器的多尺度特征融合，以弥补降采样造成的空间信息丢失。

Transformer Block主要由带残差连接和归一化层的多头注意力模块以及多层感知机组成，多头注意力模块是Transformer Block重要的组成部分，由多个自注意力连接组成，自注意力表示为：

(1)

式中，Q，K和V分别表示Query，Key和Value；dhead为通道维数，具有相同的维度(HW×C)。采用2个多层感知器和1个3×3的深度可分离卷积层来获得Transformer的Positional Encoding位置信息，表示如下：

Fout=MLP(GELU(Conv2D3×3(MLP(Fin))))+Fin，

(2)

式中，Fin是自注意力的特征图；GELU代表Gaussian Error Linear Unit激活函数。

1.3 模型整体架构

针对遥感图像建筑物分割任务，提出了一种基于PGT-Unet的卷积神经网络模型，如图2所示，渐进式的逐步增大模型的深度和输入图像的尺度，模型的输入为遥感图像，输出为建筑物分割结果图像。

图2 网络架构Fig.2 Illustration of network architecture

网络模型的第1阶段网络由一个下采样的编码模块和一个上采样的解码模块构成，编码模块和解码模块之间通过瓶颈模块连接，之后的第2阶段、第3阶段和第4阶段分别渐进式增加编码模块和解码模块的深度，使上采样和下采样的个数增加为2，3，4，输入图像的分辨率也从低分辨率(64 pixel×64 pixel)开始，逐渐增大至第2阶段(128 pixel×128 pixel)、第3阶段(256 pixel×256 pixel)和第4阶段(512 pixel×512 pixel)，最终第4阶段的Transformer Unet卷积神经网络结构如图3所示。每个阶段模型收敛之后的参数直接迁移到下一阶段的同样卷积层中。在训练过程的初始阶段，模型首先获得大尺度的结构信息和图像分布，然后在训练后期转移注意力到越来越精确的细节特征，相比之下，传统卷积神经网络是同时提取所有尺度的信息。训练过程中，为了避免每一次模型的变更对上个模型泛化能力的影响，对已训练好的、待迁移参数的低分辨率卷积层设置一个较小的学习速率(1×e-6)，对新加入的卷积层则设置了一个较大的学习速率(1×e-4)，之后开始新一轮的训练。

图3 最后阶段的Transformer Unet模型Fig.3 The last stage of Transformer Unet model

2 实验与结果

2.1 数据预处理

本文方法在Inria Aerial Image Labeling数据集上进行实验，该数据集是高分辨率遥感建筑物分割图像数据集，覆盖面积达810 km2，图像覆盖Austin，Chicago，Kitsap，Tyrol-w和Vienna等5个不同城市的建筑物，图像分辨率大小为5 000 pixel×5 000 pixel，每个城市有36张图像，该数据集共有180张图像。考虑到原始图像尺寸太大，将所有图像裁剪为512 pixel×512 pixel，最终得到14 400张训练样本和3 600张测试样本。由于原始图像存在细节模糊和颜色失真问题，对所有数据进行基于双边滤波和对数域MSR增强的去雾处理。图4(a)为原始分辨率为5 000 pixel×5 000 pixel的遥感图像和建筑物标注图像，图4(b)为裁剪之后一些小块的遥感图像用于训练和测试。

(a)原始图像

(b)切块图像图4 原始数据和切块图像Fig. 4 Raw data and patch images

2.2 损失函数

采用了Dice_loss和交叉熵损失相结合的综合损失函数对PGT-Unet模型进行训练学习，Dice_loss为一种集合相似度的度量函数，代表预测的分割结果与标注的建筑物的偏差，公式如下：

(3)

式中，TP表示神经网络模型的建筑物分割结果与手动标注结果的重叠区域；FP表示与手动标注结果相比神经网络模型的建筑物分割的错误区域；FN表示与手动标注结果相比神经网络模型未能自动分割出的建筑物区域。因此，Dice_loss的值越小，说明神经网络模型分割结果越准确。

交叉熵损失可以评估神经网络模型建筑物分割结果和手动标注结果2个分布之间的距离，使用交叉熵来评估当前训练得到的建筑物分割概率分布与真实手动标注结果分布的差异情况，公式如下：

(1-qi)×lg(1-pi)，

(4)

式中，q为真值概率；p为预测概率。

2.3 试验结果分析

在1个NVIDIA TESLA V100 GPU上进行模型的训练和测试，分辨率为(512 pixel×512 pixel)的去雾增强遥感图像作为模型的输入，输出为建筑物分割结果图像，batch size为12，训练的epoch为500。通过对比模型自动分割结果和建筑物标注图像，采用平均交并比(Intersection over Union, IoU)进行分割准确性评价，遥感图像建筑物分割的IoU为0.775。对比Inria Aerial Image Labeling数据集目前分割性能较好的模型，同样使用IoU作为评价指标，结果如表1所示。对比实验结果可以看出，本文提出的方法IoU值是最优的，表明此模型对不同大小目标、有遮挡目标的分割效果具有优势。

表1 现有方法在Inria Aerial Image Labeling数据集的建筑物分割结果对比Tab.1 Comparison of building segmentation results of existing methods in Inria Aerial Image Labeling dataset

2.4 消融实验

为了验证该模型的优越性，对提出的模型进行消融实验，所有模型采用相同的数据集和服务器环境，消融实验结果如表2所示。在Unet结构基础上增加Transformer模块的分割结果为0.760，在Unet结构基础上增加渐进生长机制的分割结果为0.754，同时引入渐进生长机制和Transformer的PGT-Unet分割结果为0.775。

表2 消融实验结果对比Tab.2 Comparison of ablation experimental results

消融实验的部分测试结果如图5和图6所示。图5和图6中，(a)为裁剪的分辨率为512 pixel×512 pixel的遥感图像，(b)为对应遥感图像手动标注结果，(c)为集成Unet和Transformer Block模型的建筑物分割结果，(d)为集成渐进生长机制的Unet模型建筑物分割结果，(e)为集成渐进生长机制和Transformer Block模块的PGT-Unet模型建筑物分割结果。

图5 PGT-Unet模型在树木遮挡样本的分割消融实验结果Fig.5 Segmentation and ablation experimental results of the PGT-Unet model in tree shelter samples

图6 PGT-Unet模型在不同大小目标样本的分割消融实验结果Fig.6 Segmentation and ablation experimental results of the PGT-Unet model in different size of object samples

由图5可以看出，第1行和第3行遥感图像主要为被树木遮挡的建筑物目标，第2行和第4行分别为上一行遮挡目标的局部放大，遮挡目标的分割是比较困难的，本文方法可以较好地学习遮挡目标的上下文特征信息，对树木遮挡的建筑物分割效果明显。

由图6可以看出，第1行和第2行遥感图像主要为成块建筑物，这样不规则的成块建筑物边缘轮廓较难分割，本文方法对不规则成块目标的分割效果更有优势。第3行和第4行遥感图像为小目标建筑物，小目标建筑物很容易被忽视，且目标边界很难准确分割，本文方法可以学习到小目标建筑物的细节特征。通过消融实验可以验证提出的基于PGT-Unet的卷积神经网络模型方法能提取到丰富的目标上下文特征，对有部分遮挡目标和不同大小目标建筑物的分割效果有很大提升。

3 结束语

本文提出了一种基于渐进生长机制的Transformer Unet卷积神经网络的遥感图像建筑物分割方法，在编码阶段和解码阶段之间引入Transformer Block模块进行多尺度特征提取和融合，获得更多目标互补特征信息，使网络能够更好地区分特征之间的重要程度，从而聚焦有用特征。通过渐进式生长机制渐进式地增大输入图像的分辨率和模型的深度，并不断迁移每个阶段的模型参数，使模型在训练初期先学习大尺度粗结构的特征信息，在后面的训练阶段逐渐学习越来越精细的细节结构信息，逐渐改善建筑物边缘分割效果，可以增强不同大小目标建筑物和遮挡目标建筑物边缘分割的完整性，对Inria Aerial Image Labeling数据集的建筑物分割具有很好的泛化能力。