这项工作的目标是在不过载GPU内存使用或丢失输出分割图中的细节的情况下对高分辨率图像进行分割。内存限制意味着我们必须对大图像进行降采样,或者将图像分为局部补丁进行分离处理。然而,前一种方法会丢失细节,而后一种方法由于缺乏全局图片可能会产生歧义。在这项工作中,我们提出了MagNet,一个多尺度框架,通过查看多个放大级别的图像来解决局部模糊问题。MagNet有多个处理阶段,每个阶段对应一个放大级别,一个阶段的输出被馈送到下一个阶段进行从粗到细的信息传播。每个阶段在比前一个阶段更高的分辨率下分析图像,恢复由于有损降采样步骤而丢失的细节,并且通过处理阶段逐渐优化分割输出。在城市景观、航拍场景和医学图像的三个高分辨率数据集上的实验证明,MagNet始终以显著的优势超越了最先进的方法。代码可在https://github.com/VinAIResearch/MagNet上找到。
当前最先进的语义图像分割技术[1,4,16,19,21,23,26]基于深度学习,其中卷积神经网络(CNN)接收输入图像并输出分割图。然而,大多数现有技术都假设整个分割过程可以通过一次前向传递输入图像来完成,并且整个过程可以适应GPU内存。不幸的是,由于内存和其他计算约束,大多数现有技术无法处理高分辨率的输入图像。处理大型输入图像的一种方法是对其进行降采样,但这会导致低分辨率的分割图,这对于需要具有细节的高分辨率输出的应用程序来说是不足够的[12,24],例如跟踪恶性病变的进展[8]。处理大型输入图像的另一种方法是将图像分成小块并独立处理每个块。然而,这种方法没有考虑到解决局部块中的歧义所需的全局信息[22]。这两种方法的局限性如图1(c)&(d)所示。
解决上述两种方法的局限性的一种方法是将它们结合起来,即融合全局和局部分割过程。一方面,整个图像的全局视图可以用来解决局部块外观上的歧义。另一方面,通过分析局部块,我们可以改善分割边界并恢复由于全局分割过程的降采样而丢失的细节。这种方法最近已经成功地通过全局局部网络(GLNet)进行了展示[5]。然而,对于具有超高分辨率的输入图像,整个图像的尺度与局部块的尺度之间存在巨大差距。这将导致对比鲜明的输出分割图,并且很难在单个前向处理阶段中结合和调和差异(见图1e);这种组合任务的困难类似于在宽阔的河流上建造一座单跨桥梁。
为了弥合尺度空间两个极端端点之间的差距,我们提出在其中考虑多个尺度。我们引入了一个新颖的多尺度框架,其中输出的分割图将在从粗到细的尺度上逐渐改进,随着图像的分析而逐渐改进。我们框架的核心是一个细化模块,可以使用一个分割图来改进另一个分割图。这个细化模块在我们的多尺度处理流程的每个阶段都被用来在最不确定的位置改进输出的分割图。我们的框架可以整合全局上下文提示以产生更准确的分割,并且可以在内存约束下输出高分辨率的详细分割图。图1显示了MagNet的结果,并将其与其他分割方法进行了比较,包括最近提出的PointRend [14]方法,该方法仅在最不确定的像素上进行改进。
图1:在高分辨率输入图像上比较几种语义分割和细化方法。降采样会丢失细节,而块处理由于缺乏全局上下文而错误地分类局部块。协作的全局局部网络GLNet由于全局和局部分支之间的巨大差异而失败。后处理和细化方法,如DenseCRF和PointRend,只能纠正由于局部不一致性而产生的小错误。MagNet通过一种新颖的多尺度分割和细化框架,表现优于其他方法。最好以彩色查看。
多尺度、多阶段、上下文聚合。多尺度级别的组合帮助网络聚合不同的视野,并为每个像素提供更多上下文信息[3,11]。ICNet [38]使用级联架构处理不同下采样输入的特征图,而RefineNet [18]融合了处理不同低分辨率输入的分支的上采样输出。特征金字塔网络(FPN)[13]在不同尺度上采样特征图,并将其与低层输出进行聚合。DeepLab [4]中的扩张卷积和空洞空间金字塔池化(ASPP)模块扩大了感受野,创建了远离的像素之间的连接。PSPNet [37]通过组合不同缩放的特征图来扩大感受野。High-resolution Net (HRNet)[30]提出了另一种尺度融合模式,在每个阶段之后添加一个具有更大感受野的新分支。许多最近的方法[2,34,35]还使用了注意力技术,以向每个点添加更多的全局信息。
处理高分辨率图像的另一种方法是使用多阶段网络,其中图像在几个阶段或子网络之后进行分割。Xia等人[32]提出了分层自动缩放网络,一种在大图像中滑动视图时缩放视野的策略。对于超高分辨率图像,Takahama等人[27]通过在分割之前预测整个块是否包含前景像素来解决背景和前景之间的不平衡问题。
将全局信息传播到局部块是处理高分辨率图像的一种有希望的方法。ParseNet [20]将全局上下文汇聚到局部视野中以获取更多信息。BiSeNet [33]增加了一个用于全局池化的分支,并且全局上下文会在最后一个阶段被添加到特征图中。虽然这些方法很有效,但对于超高分辨率图像来说,它们消耗了大量的GPU内存。Tokunaga等人[28]提出了一种针对超高分辨率图像的方法,使用独立的多尺度网络和自适应权重生成器。网络成员的输出与相应的训练权重结合,产生最终的输出,但网络分支之间没有知识共享。与[28]不同的是,[5]包含了两个具有共享信息的子网络,其中全局分支使用下采样的图像提取全局上下文,而局部分支使用块和相应的全局特征来改善高分辨率图像的细节。然而,由于全局和局部分支的临时组合,很难扩展到超过两个尺度。此外,根据我们的经验,由于强大的全局分支的主导作用,训练局部网络是困难的。
分割细化。有几种方法可以通过后处理来改进分割结果。一种方法是在深度学习网络生成的分割掩模上使用传统方法,如条件随机场(CRF)[15]或引导滤波器(GF)[10]。然而,这些方法速度较慢,改进效果渐进。可以使用深度学习版本的引导滤波器(DGF)[31]来提高推理速度。另一种后处理方法是使用深度网络。迭代实例分割(ISS)[17]通过多次将输入图像和分割图通过一个细化模块进行传递来改进输出。该方法基于自我反思,每个细化阶段的输入图像都相同。类似于ISS,级联PSP [6]使用相同的细化方案,但每个细化阶段的输入分辨率不同。然而,任何中间阶段的错误预测都可能显著影响后续步骤的性能。一些方法旨在仅对输出的部分进行细化,例如边界上的像素[36,39]或不确定位置的像素(PointRend [14])。然而,边界细化方法[36,39]未能恢复微小的对象,而PointRend [14]仅使用局部上下文进行细化。此外,由于PointRend的输入是深度网络的高级特征,它必须针对每个分割骨干进行特定训练。在本文中,我们提出了一个模块化框架,可以具有任意数量的尺度级别。它简单而有效地对粗糙的分割输出进行细化,能够在每个阶段之后添加更多细节而保持粗糙分割输出的整体结构,而不会出现主导问题。表1比较了不同方法的关键特点。
我们现在介绍MagNet,这是一个用于超高分辨率图像的多尺度分割框架。它是一个多阶段的网络架构,每个阶段对应于特定的尺度。输入图像将在多个尺度上进行检查,从最粗糙到最细致。我们框架的核心是一个分割模块和一个细化模块,它们在每个处理阶段都被使用。在每个阶段,细化模块的输入是两个分割图:
(1)来自前几个阶段的累积结果,和
(2)仅在当前尺度上运行分割模块得到的结果。
细化模块的目标是使用后者的分割图在基于两个估计分割图的不确定性的选择位置上对前者进行细化。在我们的框架中,分割模块可以是任何分割骨干,只要它能够输出带有不确定性估计的分割图。细化模块对分割骨干是不可知的,它可以用一个骨干进行训练并用另一个骨干进行使用。接下来,我们将详细描述多阶段处理流程和细化模块。
MagNet的架构和处理流程如图2所示。其中有一个分割模块和一个细化模块,在m个处理阶段中重复使用,其中m是我们要分析的尺度数量的超参数。我们用s表示处理阶段,其中s = 1对应于最粗糙的尺度,s = m对应于最细致的尺度。假设X∈RH×W×3是一个输入图像,其中H和W是图像的高度和宽度。我们考虑当H和W太大以至于无法在不进行降采样的情况下处理图像X的情况,并且让h和w是分割模块可以处理的最大(或理想)尺寸。我们用hs和ws表示尺度级别s的高度和宽度。我们确定尺度级别,使其跨越整个尺度空间:H = h1 > ... > hm = h和W = w1 > ... > wm = w.
对于特定的尺度级别s,我们将输入图像X划分为大小为hs×ws的补丁,并对这些补丁进行语义分割。这些补丁的位置由一组矩形窗口定义,设Ps表示这些窗口的集合:Ps = {p = (x, y, ws, hs)},其中每个窗口由左上角的坐标、宽度和高度来确定。随着尺度级别s的增加,矩形窗口的宽度和高度减小,但Ps的基数增加。对于特定的窗口p,我们将使用Xp来表示在窗口p处提取的图像补丁。
我们的网络将接收一个图像X∈RH×W×3,并生成一系列的分割图Y1,···,Ym∈RH×W×C,其中C是考虑的语义类别数量。在第s个阶段,我们首先确定用于补丁划分的矩形窗口集合Ps,并对每个补丁的分割图进行细化。具体来说,对于每个窗口p∈Ps,执行以下步骤:
1.提取图像补丁Xps和窗口p定义的先前分割输出Yps−1。这些张量的高度和宽度分别为hs和ws。
2.对Xps和Yps−1进行降采样,使得新的高度和宽度为h和w,这是可以适应GPU内存并由分割和细化模块处理的尺寸。用X¯sp和Y¯ps−1表示降采样后的张量。
3.以X¯sp作为输入,使用分割模块获得尺度特定的分割图O¯ps。
4.以Y¯s−1p和O¯ps作为输入,使用细化模块对Y¯s−1p进行细化,得到Y¯ps(详见第3.2节)。
5.上采样Y¯ps以得到尺寸为hs×ws×C的Yps。这些处理步骤在图2中有详细说明。
图2:我们提出的MagNet的概述。分割网络产生尺度特定的预测,而细化模块根据局部预测有选择地细化前几个阶段的粗糙预测。
图3:细化模块的概述。累积分割Y部分地被基于分数Q的尺度特定分割图O所替换。
细化模块是我们框架的核心组件,用于在我们流水线的每个处理阶段对分割图的个别补丁进行细化。该模块的输入是两个尺寸为h×w×C的分割图:(1)来自所有先前尺度的累积分割图Y,和(2)来自当前尺度的尺度特定分割图O。模块的输出是更新后的尺度累积分割图。
图3描述了细化过程,包括以下步骤。首先,使用一个小网络,以Y和O作为输入,我们获得一个初始的组合分割图R。然后,我们计算Y和R的预测不确定性图。具体来说,对于Y的每个像素,该位置的预测置信度被定义为最高概率值与第二高概率值之间的绝对差异(在C个类别的概率值中)。然后根据置信度得分计算不确定性分数,使得这两个分数相加等于1.类似地,我们可以计算R的预测不确定性图。令Yu和Ru分别表示Y和R的预测不确定性图。
接下来,我们将描述如何使用两个不确定性图来选择Y的k个位置进行细化。在之前的工作PointRend中,只使用了一个不确定性预测,我们将这种方法扩展到更好的选择策略。这些位置是Y对其预测不确定的位置,而R对其预测是确定的。用于对像素进行排名的得分图计算公式为Q = F(Yu(1-Ru)),其中表示逐点乘法,F表示中值模糊以平滑得分图。在第4.2.3节中可以找到公式中每个元素影响的实证比较。
我们的框架在准确性和运行时效率之间存在权衡。降低运行时间的一种方法是减少要处理的尺度数量。另一种方法是在每个尺度级别上对图像补丁的子集进行分割和细化。当MagNet-Fast在较少的尺度上运行时,它将这两种方法结合起来,并且只选择具有最高预测不确定性Yu的补丁进行细化。在MagNet-Fast中,需要输入到分割模块中的图像补丁的总数甚至可能小于补丁处理方法中使用的图像补丁数量。此外,MagNet-Fast可以利用全局上下文和详细信息进行分割,从而产生出色的结果,这将在我们的实验证明中看到。
图4:这两个残差块被训练用于在每个尺度上细化分割。该模块的输出与输入具有相同的尺寸h×w。
我们在三个高分辨率数据集上评估了MagNet的性能:Cityscapes [7],DeepGlobe [9]和Gleason [29]。关于这些数据集的一些信息列在表2中。每个图像的像素数在200万到2500万之间。在本节中,我们将展示与其他最先进的语义分割框架进行比较的实验,并描述对Cityscapes进行的一些消融研究。
表2:用于评估我们框架的高分辨率数据集的详细信息。所有图像都具有从200万到2500万个像素,并且具有大量细节。
细化模块的架构。图4描述了在所有实验中使用的细化模块的架构。主要组件是两个残差块。输入大小为h×w×2C的细化模块会产生大小为h×w×C的输出。
训练。对于每个数据集,我们在降采样的图像上训练了一个最先进的分割模型,并训练了一个细化模块来改善切片图像的粗糙输出。在训练细化模块时,我们随机提取图像补丁,并应用以下数据增强:旋转、水平和垂直翻转。我们使用带有动量0.9的SGD优化器,衰减权重为5×10^-4,初始学习率为10^-3。我们训练细化模块50个epochs,并在第20、30、40和45个epoch时将学习率降低十倍。我们使用交叉熵作为训练分割和细化模块的损失函数。我们使用PyTorch [25]实现了MagNet,从HRNet-OCR [35]的公共实现开始。我们在一台搭载Tesla V100 GPU的DGX-1工作站上使用批处理大小为16进行训练。
测试。在推理过程中,我们在每个尺度上提取非重叠的补丁进行处理。评估指标是平均交并比(mIoU)。为了进行内存和速度比较,我们在一台搭载Intel i7 CPU和RTX2080Ti GPU的机器上进行了基准测试实验。
表3:MagNet在Cityscapes上的性能(包括和不包括中间尺度级别)。拥有中间尺度级别是必不可少的。
Cityscapes是一个高分辨率城市场景的数据集,包含2048×1024像素大小的图像。任务是对由汽车摄像头捕捉的视频中的对象进行分割。数据分为粗糙和细致两种标签。在我们的实验中,我们使用了带有细致标签的数据集,包括2975个训练图像和500个测试图像。我们考虑了四个可能的尺度级别(m=4),对应于2048×1024、1024×512、512×256和256×128的补丁大小。分割模块的输入尺寸始终为256×128,以满足内存限制,因此任何更大的补丁都将被降采样。
表3展示了MagNet在不同尺度数量下的结果。从最低尺度到最高尺度的直接细化提高了约2%的mIoU,从63.23%提高到65.23%。在最小尺度和最大尺度之间添加两个中间尺度将性能提高了4.34%的mIoU。可以在图5中观察到每个处理阶段的定性改进。在每个阶段之后,错误减少,分割掩码变得更精细。
表4:MagNet和其他分割细化方法在Cityscapes数据集上的性能。对于所有细化方法,我们使用HRNetV2-W18+OCR [35]作为分割模块的主干网络。
表4和表5比较了MagNet与几种最先进的语义分割细化方法的性能。所有方法都是在预训练的HRNet+OCR [35]模型的输出上进行训练和应用的,该模型是该数据集上的领先方法之一。尽管有各种各样的HRNet模型[30],但我们选择了HRNetV2-W18,因为它的复杂度适合我们的实验要求,尤其是在内存约束方面。MagNet-Fast是MagNet的高效版本,在每个尺度上只对最不确定的补丁进行细化。在这个实验中,我们选择了补丁的数量,使得MagNet-Fast的处理速度与Seg Fix [36]和DGF [31]相当。该模型仅在三个尺度256×512×2048上运行,对于每个尺度选择了三个最不确定的补丁。总共,MagNet-Fast需要在7个补丁上运行推断,而MagNet需要在1+4+16+64=85个补丁上运行,而补丁处理方法则需要在64个补丁上运行推断。
在这个实验中,除了DenseCRF [15]之外,我们对其他框架进行了微调,以获得与分割主干网络最佳结果。对于SegFix [36],使用作者发布的偏移预测,最佳结果是在偏移宽度为10时获得的。DGF [31]和PointRend [14]都是使用分割主干网络的输出进行训练的。此外,DGF在补丁上运行,以便在速度和内存上与我们的方法进行公平比较。可以观察到,DenseCRF [15]是唯一不能改善粗糙分割的方法。PointRend [14]是最快的方法,但改进很小。MagNet-Fast、SegFix [36]和DGF [31]的推断时间相似,但MagNet-Fast明显优于其他方法。MagNet的速度较慢,但具有最高的mIoU值。
图6a显示了这些方法在我们的实验中的累积IoU分布。MagNet和其他方法之间存在很大的差距,特别是在放大窗口时。图6b显示了几种方法在两个Cityscapes图像上的结果。MagNet和MagNet-Fast都能够提供最好的细化效果。SegFix [36]不能恢复小物体,比如与较大区域错误合并的标志杆,而PointRend [14]由于缺乏全局上下文而表现不佳。
图5:在Cityscapes数据集上通过每个处理阶段对分割输出进行可视化。第一张是原始图像。第一行中的其他图像是要进行细化的选择点(红色)。第二行是分割输出。第三行是与真实标注进行比较的错误。最好以彩色查看。
表5:Cityscapes数据集上某些特定类别的IoU值。最佳和先前最佳方法分别以红色和蓝色突出显示,并且它们之间的差异显示在最后一行。
图6:我们的方法在Cityscapes数据集上优于其他细化框架。(a)数据集上每个图像的mIoU的累积分布(较低为更好)。MagNet和MagNet-Fast在其他方法中取得了最佳结果。(b)细化框架和我们的MagNet的一些分割结果。图像下方显示的是mIoU的数值。MagNet和MagNet-Fast能够更好地识别更多的小物体,并对边界进行更好的细化。在放大的数字设备上查看效果最佳。
表6:MagNet在Cityscapes数据集上使用不同排名得分时的性能。初始分割的mIoU为63.23%。当使用Y u和(1 −Ru)进行平滑操作时,当k=216时,该框架达到了最佳性能。
图7:Cityscapes数据集上每个尺度选择点数量与平均IoU之间的相关性。当选择点数量增加时,MagNet的性能持续增长,而PointRend的mIoU则下降。
表6展示了我们在选择细化点时使用预测不确定性图Y u、Ru和中值滤波函数F的重要性的割除研究。当同时使用两个不确定性图时,获得了最佳性能。此外,在每种情况下使用中值滤波进行平滑改善了结果。我们还研究了细化点数量与准确性之间的相关性。图7展示了MagNet和PointRend [14]的结果。可以看出,当细化点数量增加时,MagNet的性能提高了。当点数达到216个后,性能停止增加,如果选择了所有点进行细化,性能下降到了66.86%。与此同时,当选择的点数超过2048时,PointRend的性能显著下降,甚至低于未进行细化的初始值。
我们还使用了两种不同的分割主干网络对MagNet框架进行了测试,结果显示在表7中。在两种情况下,MagNet显著改善了原始网络的分割结果,提高了2%到5%。在这个实验中,我们使用了四个尺度级别:256、512、1024、2048,并且每个补丁的细化位置数量为k=216。
表7:使用两种主干网络的MagNet结果。MagNet可以与不同的分割主干网络一起使用,并改善它们的分割结果。有关使用HRNetV2-W18主干的MagNet结果,请参见表4。
表8:DeepGlobe数据集上的分割结果。我们在最后部分使用了相同的分割主干(FPN)来进行所有细化方法的比较。
DeepGlobe是一个高分辨率卫星图像的数据集。该数据集包含803张图像,标注了七个景观类别,包括未知类别。按照[5]的评估协议,mIoU计算中会忽略未知类别,因此只需要考虑六个类别。图像的尺寸为2448×2448像素。我们使用了与[5]报告中相同的训练/验证/测试划分,分别使用455、207和142张图像进行训练、验证和测试。
在先前的工作GLNet [5]中,我们使用了具有ResNet-50主干的特征金字塔网络(FPN)作为分割网络。我们还使用了与GLNet相同的输入尺寸508×508。我们使用了三个细化阶段,分别是三个尺度612、1224和2448,并在每个尺度上选择了216个点进行细化。结果显示在表8中。对于MagNet-Fast,在每个尺度级别上,我们选择了预测不确定性最高的前三个补丁进行细化。PointRend [14]也使用了相同的分割主干进行训练,并获得了比GLNet更高的准确性。MagNet和MagNet-Fast都优于其他方法。
Gleason [29]是一个包含前列腺癌组织病理学图像的医学数据集。任务是在超高分辨率图像上进行病变的分割和分级。数据集中有四个需要进行分割的类别:良性、3级、4级和5级。数据集包含244张图像,每张图像的尺寸为5000×5000像素,并由六名临床专家提供了分割标签。最终的标签是基于多数投票的结果。我们将数据集随机分为195个训练图像和49个测试图像。
在Gleason数据集中,PSPNet [37]是在排行榜上排名最高的方法,我们使用了Resnet-101作为分割网络的主干。我们使用了512×512的输入尺寸,并使用了四个细化阶段和四个尺度:625、1250、2500、5000。MagNet-Fast也在四个尺度上运行,但每个尺度只选择了预测不确定性最高的三个补丁。每个尺度有216个细化点。MagNet和MagNet-Fast的结果,以及获胜解决方案PSPNet的结果和医学专家之间的mIoU一致性,显示在表9中。MagNet使用了PSPNet的分割主干,并将PSPNet的性能提升了1.7%。
表9:MagNet和其他框架在Gleason数据集上的性能,使用PSPNet [37]作为主干。
我们提出了MagNet,一个用于高分辨率图像的多尺度分割框架。MagNet可以在不爆炸GPU内存使用的情况下生成高分辨率的分割输出,方法是将输入图像划分为补丁。为了避免过于全局或局部的问题,我们考虑了多个尺度的补丁,从最粗糙到最精细的级别。MagNet有多个分割阶段,其中一个阶段的输出将被用作下一个阶段的输入,并且分割输出将逐步细化。我们在三个具有挑战性的高分辨率图像数据集上展示了MagNet的优势,在平均交并比(mIoU)方面,MagNet的性能较之前的最先进方法提高了1%至2%。