全局-局部的超分辨率图像分割方法

GLNet:Collaborative Global-Local Networks for Memory-Efficient Segmentation of Ultra-High Resolution Images

摘要:

超分辨率图像分割的要求越来越高,对算法的效率,尤其是显存占用提出了严峻的挑战。目前的方法要么是对超分图下采样,要么是切成小块处理。这两种方法中,局部细节或者全局信息的丢失限制了分割的精度。因此提出GLnet(Global-Local Network)高效地保存全局和局部地信息。GLnet由全局分支和局部分支组成,全局分支输入降采样地整个图像,局部分支输入裁剪地小块。对于分割任务,GLnet融合两个分支地特征图,既能够捕捉高分辨率地图像小块的精细信息,又能够捕捉下采样的上下文信息。为了进一步解决前景和背景的类别差异问题,提出了由粗到细的GLnet的变体,同时具有存储效率(memory-efficient)。在3个超分辨率数据集上进行了大量的实验,只需要一个1080Ti的GPU,2G的显存就能实现很好的效果。

一、介绍

随着摄影和传感器技术发展,超分辨率图像的普及为计算机视觉开辟新的领域,并对有效性分析增加了新的要求。目前,2K图像素至少为2048 \(\times\) 1080(大约2.2M),4K图至少为3840 \(\times\) 1080(大约4.1M),4K超高清至少为3840 \(\times\) 2160(大约8.3M)。

分割过程中,图像按像素级别分为不同的类别,如卫星图中的城市/森林/水域,或者皮肤区域中的病变区域。超分辨率图像分割在城市规划,遥感,疾病监测等领域有重要作用。

当图像的分辨率很大时,对效率和分割质量提出了很大的挑战,表1展示了几种超分数据集和经典分割数据集。

目前对超分图分割研究很少,经典操作如下采样或分块裁剪分别丢失精细信息和上下文信息。像FCN或者SegNet在分割6M图时需要5G到10G的显存,这些方法有较高的精度和显存占用,图1(a)的右上角。

图1:在DeepGlobe数据集的推断时显存占用和平均交并比(mIoU)。

(a):不同分割方法的mIoU和显存占用。

(b):不同图像大小的mIoU和显存占用。(降采样比例写在注释里了)

(c):不同patch大小的mIoU和显存占用。

最近,轻量分割法比如ICNet显存占用很小但是精度很差。对不同尺寸的全局图和不同数量的局部图的研究表明经典模型没办法在准确率和显存占用上取得良好的平衡。

👕贡献:

  • 首次提出memory-efficient(内存高效)的超分图分割方法,不仅准确率高,显存占用也小。
  • 能够有效使用全局上下文信息和局部精细结构,分割精度高,这两个模块都是必不可少的。
  • 针对超分图分割的类别不平衡问题,提出由粗到细的GLNet算法,进一步提高了性能。

GLNet有全局分支和局部分支,分别处理下采样的全局图和裁剪的局部图。它们通过feature map跨层共享,相互作用和调节。在准确率和显存占用上取得了很好的平衡。为了进一步解决类别不平衡问题,当对分割的小的前景区感兴趣时,提供了GLNet由粗到细的变体,其中全局分支提供了额外的边界框(bounding box)。它的设计使得能够综合使用全局信息和局部信息,确保精准地分割。同时极大地减少了GPU的使用,因为送到网络的是下采样的全局图和裁剪的局部图,原来整张的高清图是不会送进去的。

二、相关工作

2.1 语义分割的质量和精度

FCN是第一个高质量分割的CNN结构;U-net使用encoder-decoder结构,使用跳级连接将低级特征连接到高级特征,相似结构还有DeconvNet和SegNet;DeepLab使用膨胀卷积(也叫空洞卷积)扩大感受野,条件随机场用来建模空间关系。但是这些方法的显存占用都很高。

随着语义分割在很多实时/低延迟 任务(比如自动驾驶)中越来越重要,高效快速的分割获得越来越多的关注。ENet在早期使用下采样,使用非对称的encoder-decoder结构减少浮点运算;ICNet将多分辨率的分支的feature map在适当的标签下级联处理,并进行模型压缩。然而这些方法没有针对超分辨率图像,实验表明没有取得效率和精度的平衡(主要是精度不够)。

2.2 多尺度和上下文聚合

多尺度是集成高层和底层特征捕获不同粒度信息的模式,在分割上很有效。在RefineNet中,通过上采样低分辨率特征,构造多路径精化块来结合多尺度特征。文献32采用拉普拉斯金字塔结构利用高层特征来细化低分辨率图构造的边界。FPN逐步上采样不同尺度的特征图,并且自上而下地聚合它们。分级自动缩放网络(HAZN)采用两步自动缩放策略将粗级的边界框和预测分数送到细级。

上下文聚合在编码局部邻域空间甚至非局部信息方面也起着关键作用。ParseNet采用全局池化聚合不同级别的场景上下文。DeepLab中的空洞卷积和ASPP(空间金字塔池化)在不丢失feature map的情况下扩大感受野,从而将全局上下文聚合到局部信息中,PSPNet中的金字塔池实现了类似的效果。在ContextNet,BiSeNet和GUN中,深/浅分支被结合到全局上下文和高分辨率细节上。文献37把上下文信息看作是RNN建模的长期依存关系。需要注意在GLNet中,在输入级(全局/局部分支)和特征级都采用了上下文聚合。

2.3 超分辨率分割数据集

DeepGlobe陆地覆盖分类数据集是第一个提供在农村地区高分辨率亚米级卫星图像的公开benchmark。提供了7类像素级ground truth的mask:城市,农业,牧场,森林,水,荒地,和未知。包含1146张标注的卫星图,大小都是2448 \(\times\) 2448,比以往的陆地覆盖分类数据集有更高的分辨率和更大的挑战性。

ISIC收集大量的皮肤镜图像。它的子集ISIC Lesion Boundary Segmentation由2594张皮肤癌患者筛查的样本图像组成,所有图像都标注了ground truth的二值mask,指出了原发皮肤损伤的位置。超过64%的图像有超分辨率,最大的图像是6682 \(\times\) 4401像素。

Inria Aerial数据集涵盖了不同城市景观,从人口稠密的都市区到阿尔卑斯度假胜地。提供了180张5000 \(\times\) 5000像素的图片,都是二值mask,表示有建筑区和没有建筑区,与DeepGlobe不同的是,它将训练测试集按城市分类,而不是随机分类。

三、协同全局-局部网络

3.1 motivation:为什么不是单独的全局或者局部

对于有显存限制的超分辨率图像训练和推理,有两个想法:全局图下采样或者裁剪。然而它们性能很糟糕。

图3:DeepGlobe数据集的示例分割结果。

(1):原图

(2):Ground Truth

(3):只经过全局图下采样的预测

(4):只经过局部块的预测

(5):协同GLNet的预测

3.2 GLNet结构

3.2.1 全局局部分支

N张超分图像,图像片段\(\mathcal{D}=\{(\boldsymbol{I}_{i}, \boldsymbol{S}_{i})\}_{i=1}^{N}\),满足\(\boldsymbol{I}_{i}, \boldsymbol{S}_{i} \in \mathbb{R}^{H \times W}\) ,全局分支获得下采样低分辨率的图\(\mathcal{G}\) 并且\(\mathcal{D}^{\mathrm{lr}}=\{(\boldsymbol{I}_{i}^{\mathrm{lr}}, \boldsymbol{S}_{i}^{\mathrm{lr}})\}_{i=1}^{N}\),局部分支\(\mathcal{L}\)对于\(\mathcal{D}\)中的块图,分辨率是相同的,\(\mathcal{D}^{\mathrm{hr}}=\{\{(\boldsymbol{I}_{i j}^{\mathrm{hr}}, \boldsymbol{S}_{i j}^{\mathrm{hr}})\}_{j=1}^{n_{i}}\}_{i=1}^{N}\)\(\mathcal{D}\)中每个\(\mathcal{I}_i\)\(\mathcal{S}_i\)都由\(n_i\)个patch组成。注意\(\mathcal{I}_i\)\(\mathcal{S}_i\)被完全裁剪成patch,而不是随机裁剪,确保训练和推理是同分布的。\(\boldsymbol{I}_{i}^{\mathrm{Ir}}, \boldsymbol{S}_{i}^{\mathrm{lr}} \in \mathbb{R}^{h_{1} \times w_{1}}\)并且\(\boldsymbol{I}_{i}^{\mathrm{hr}}, \boldsymbol{S}_{i}^{\mathrm{hr}} \in \mathbb{R}^{h_{2} \times w_{2}}\),满足\(h_{1}, h_{2} \ll H\)\(w_{1}, w_{2} \ll W\),对于\(\mathcal{G}\)\(\mathcal{L}\),采用相同的backbone,两者都可以看作是从第1层到第\(L\)层的卷积级联块(如图5).

图4:GLNet概览。global分支输入下采样图,local分支输入裁剪图。深度feature map共享和feature map正则化确保global-local地融合。最终地分割结果是聚合两个分支的高层feature map得到的。

图5:全局和局部分支的深度feature map共享。在每一层,具有全局上下文的feature map和具有精细结构的feature map双向组合在一起,形成完整的基于patch的协同global-local模块。聚合结果的main loss和两个辅助分支的loss成为要优化的目标。

在分割过程中,各个分支的各层feature map与另外的分支共享,然后通过分支层\(f_{agg}\)对两组高层feature map聚合生成最终的分割mask。为了约束两个分支,稳定训练过程,局部分支训练采用弱耦合正则化(weakly-coupled)。

3.2.2 深度特征图共享

为了与local分支协作,需要将global分支中local patch对应空间位置裁剪,然后上采样到和local分支 feature map相同的大小,然后被同一层的local分支的feature map concatenate,通过对称的方式,local分支的feature map也是这么收集的。首先对local的feature map下采样,匹配从原图裁剪对应位置的patch,然后合并在一起(与局部patch被裁剪的顺序相同),形成完整的特征图,它的大小与全局分支的feature map相同。在送到下一层之前,这些局部特征也被concat到全局特征中。

图5展示了这个过程,除了最后一层外,它是layer-wise(逐层)的。共享的方向可以是单向的(\(\mathcal{G} \rightarrow \mathcal{L}\))也可以是双向(\(\mathcal{G} \rightleftarrows \mathcal{L}\))的。每一层中,全局上下文特征和局部精细特征相互融合借鉴。

3.2.3 带正则化的分支聚合

两个分支的聚合是通过聚合层\(f_{agg}\),它由一个3 \(\times\) 3的卷积核的卷积层实现。它从local分支的\(L_{th}\)\(\hat{X}_{L}^{\mathrm{Loc}}\)和对应的global分支的\(\hat{\boldsymbol{X}}_{L}^{\mathrm{Glb}}\),通过channel拼接(concat),输出层\(f_{agg}\)得到最后的分割结果\(\hat{S}^{\mathrm{Agg}}\)。另外为了enforce \(\hat{S}^{\mathrm{Agg}}\)的主要分割loss(main loss),采用了两种辅助的loss,为了enforce local分支的输出\(\hat{S}^{\mathrm{Loc}}\)和global分支的输出\(\hat{S}^{\mathrm{Glb}}\)靠近它们的分割map(local patch/global downsmaple),能够让训练过程更稳定。

在实践中发现local分支容易过拟合local detail(局部细节),从而overriding在global分支的学习,为了避免local分支学习速度比global分支快,在两个分支最后一层特征映射之间添加了弱耦合正则化(weakly-coupled)。具体的说是用2范数\(\lambda\|\hat{\boldsymbol{X}}_{L}^{\mathrm{Loc}}-\hat{\boldsymbol{X}}_{L}^{\mathrm{Glb}}\|_{2}\),防止\(\hat{\boldsymbol{X}}_{L}^{\mathrm{Loc}}\)\(\hat{\boldsymbol{X}}_{L}^{\mathrm{Glb}}\)之间有太大的变化,根据经验,\(\lambda\)固定为0.15。这种正则化主要是让local分支的训练变慢,与global分支的学习同步,它只更新local分支的参数。

3.3 由粗到细的GLNet

分割是为了分离前景与背景,(比如二值mask),前景通常在超分图中占据很小的空间。这种类别不平衡会严重影响性能。以ISIC数据集为例,99%的图像背景比前景多,超过60%的图像前景像素比例少于20%(图8中的蓝条),

图8:(1):ISIC 2018数据集前景背景像素比例直方图,蓝色表示边界框细化前的前景背景比,红色表示边界框细化后的前景背景比,后者更平衡。

(2):由粗到细分割的可视化结果。经过细化(从b到e),GLNet能捕捉更精确的边界。

很多local patch只包含背景,这会导致梯度不正常。

🎪两阶段细分解决方法

图6:两阶段分割。global 分支粗分割,local分支细分割并且只处理边界框的前景中心区域。

为了缓解类别不平衡,提出GLNet由粗到细的两阶段变体(如图6)。首先用全局分支对下采样的图像进行粗分割,然后为前景区创建一个边界框(动态调整边界框大小,以便边界区域前背景比例约为1,有助于第二步的类别平衡)。然后将原图的边界前景区送到local分支精细分割。

与并行的全局-局部分支不同,这种由粗到细的GLNet是两个分支的序列组合,在边界框细化过程中,只有边界框类的feature map才会由全局分支共享到局部分支,然后共享回来。

所有边界框之外的区域都会预测为背景。所以由粗到细的GLNet通过只选择细尺度处理 而降低了计算成本(因为粗尺度直接不考虑了)。

四、实验

在DeepGlobe,Inria Aerial数据集评估GLNet表现,在ISIC数据集上评估由粗到细的GLNet的表现。

4.1 实验设置

采用带有ResNet50的FPN(特征金字塔)作为backbone。深度特征共享策略在自底向上阶段应用于ResNet50的conv2到conv5 block,在自顶向下和smoothing阶段应用于FPN中。对于FPN中最后横向拼接阶段,采用特征图正则化,聚合到最终的分割结果中。

为了简单起见,下采样的全局图和裁剪的局部patch大小相同,都为500 \(\times\) 500像素,相邻的patch有50像素的重叠避免卷积层边界消失。使用\(\gamma=6\)的focal loss作为main loss和两个辅助loss。main loss和辅助loss都为1.0,特征图正则化项的\(\lambda\)设为0.15。

为了测量模型的GPU占用,使用命令行工具gpustatminibatch大小为1,避免计算任何梯度。训练和测试只用一张卡。

用Pytorch实现,优化器Adam(\(\beta_1=0.9,\beta_2=0.999\))global分支学习率0.0001,local分支学习率 0.00002,训练minibatch为6。

4.2 DeepGlobe

803张超分图,像素都为2448 \(\times\) 2448,随机划分训练验证测试为455,207,142张。包含7个类别,有1个类别是未知。

4.2.1 从浅到深的特征图共享

为了评估global-local的策略,逐步将模型的特征图共享从浅层到深层,对于下采样的全局图或者单独的局部图,每个分支的平均交并比分别只能达到57.3%和66.4%(表2)。

表2:在DeepGlobe测试集上评估不同层次的特征图共享策略的影响。Agg表示聚合层,Fmreg表示2范数的正则化。\(\mathcal{G} \rightarrow \mathcal{L}\)表示global共享到local,\(\mathcal{G} \rightleftarrows \mathcal{L}\)表示global和local双向共享,Shallow表示单层共享,Deep表示所有层共享。

通过两个分支高层特征的聚合以及正则化,mIoU可以到70.3%,当只共享global的一层到local时,聚合结果增加了0.2%,当所有层都共享时,mIoU到了70.9%,如果是双向共享,mIoU到了71.6%。

消融实验表明通过深度共享,正则化和聚合策略,全局分支和局部分支能够有效地协同,即使用双向共享,显存占用也只从1189MB略微增加到1865MB。

图7放大(a),(b)可视化改进的效果。全局结果(图7(3))有网状伪影,局部结果(图7(4))有不准确的边界,从聚合,浅层共享到深度共享,显著观察到误分类和不准确边界的情况明显减少。

图7:DeepGlobe数据集的示例分割结果。

(1):原图

(2):ground truth

(3):只通过全局图下采样的预测

(4):只通过裁剪的局部图的预测

(5):聚合全局图和局部图并且使用正则化的预测

(6):浅层共享

(7):双向深度共享

从(3)到(7)放大区域的(a)和(b)说明了局部区域的精细结构

4.2.2 准确性和显存占用比较

使用全局图和局部图进行训练和推理可能会产生不同的结果,主要是模型感受野,卷积核大小和填充策略不同,这使得训练推理有不同的选择,因此在消融实验中仔细比较了两种方法训练的模型。对一个模型训练和测试两次,然后选择最佳结果。

🎑固定图像/patch大小的粗略比较(所有模型都使用500\(\times500\)像素)

表3展示了所有模型全局推理下的最高mIoU,但是都占用了很多显存。在基于patch的推理中显存下降了,但是准确度也大幅度下降,只有GLNet实现了mIoU和显存占用的平衡,在图1(a)中画出了每种方法最佳的mIoU。

表3:在DeepGlobe测试集上的mIoU和显存占用。\(\mathcal{G} \rightarrow \mathcal{L}\)\(\mathcal{G} \rightleftarrows \mathcal{L}\)分别表示从全局共享到局部和双向共享,GLNet没有用全局图推理,如图1(a)。

🎁深入对比不同图像/patch大小

选择FCN和ICNet对不同图像/patch大小深入比较,因为它们分别具有较高的mIoU和较小的显存占用。在图1(b)和(c)中绘制了消融实验的细节。对于FCN和ICNet,精度更高因为着显存占用大,反之也是,这证明了经典模型无法平衡效率和精度。

4.3 ISIC

ISIC皮肤损害边界分割挑战数据集包含2596超分辨率图,随机切分训练,验证测试为2077, 260, 259张图。

4.3.1 由粗到细的分割

在严重不平衡的ISIC数据集上,global和local分支分别只能实现72.7%和48.5%的mIoU,当使用由粗到细的策略时,可以清楚看到更加平衡的前景背景比例(图8(1)中的红条)。

通过裁剪前景边界框,local分支只在更小的类别平衡的图像上训练,默认将裁掉的边界外区域作为背景。在类别平衡的图像总共,\(\mathcal{G} \rightarrow \mathcal{L}\)的mIoU是73.9%,\(\mathcal{G} \rightleftarrows \mathcal{L}\)的mIoU是75.2%,这种情况下,全局分支使用更加精确的全局上下文信息,因为对裁剪更小的图像进行下采样,丢失的信息会更少。由粗到细的方法表明它能较好地捕获上下文信息,解决类别不平衡问题。表4中是消融实验,图8(2)是可视化结果。

表4:在ISIC测试集上由粗到细的分割和深度特征图共享的有效性。Bbox表示由global分支的边界细化。

4.3.2 精度和显存比较

表5是ISIC数据集中GLNet的mIoU和显存占用情况,在精度和显存存上显著由于其他方法。

表5:在ISIC测试集上预测的显存占用和mIoU。

4.4 Inria Aerial

包含180张超分图,每个都是5000\(\times\) 5000像素,随机划分训练验证测试为126,27,27张。表6说明了深度特征图共享策略的有效性。测试结果如表7,GLNet的mIoU有71.2%,在显存占用和准确率上都比其他方法要好。

表6:GLNet在Inria Aerial测试集的评估

表7:在Inria Aerial测试集上几种方法的mIoU和显存占用对比。

五、总结

提出了GLNet做超分辨率图像分割,综合利用了全局上下文和局部精细信息提高了在超分辨率上的效果,并且没有多大的显存占用。同时也通过由粗到细的分割法解决类别不平衡问题。

在超分辨率中,效率和精度的平衡至关重要,这篇文章的工作是超分辨率分割效率和精度的先驱。