SOLO论文笔记

    论文:Xinlong Wang,Tao Kong,Chunhua Shen1,Yuning Jiang,Lei Li.SOLO: Segmenting Objects by Locations.阿莱德大学和字节跳动。


摘要

    这篇论文提出了一个新的实例分割模型SOLO。与语义分割相比,实例分割任务由于实例数量的不确定,因此更加困难。为了预测实例掩码,主流方法要么“detect-then-segment”,比如Mask R-CNN;要么先预测每个像素的类别,再使用聚类技术划分实例。本文从新角度看待实例分割问题,引入了“实例类别(instance categories)”的概念,根据实例的位置和大小,对每个实例的像素分配一个实例类别,将实例分割问题转换为分类问题。由此,实例分割任务被解耦成两个并行的分类任务。本文展示简单、灵活的SOLO实例分割框架可以获得很强力的性能,达到Mask R-CNN的精度,胜过最近提出的单阶段实例分割算法。



1. 引言

    实例分割是四大计算机视觉基础任务中最难的一个,它既需要正确的定位图像中的所有实例,同时需要在像素级分割出每个实例。对于一张图像来说,语义类别的数量是固定的,因此,语义分割可以很容易的被建模为密集的逐像素分类问题。但是每个图像中的实例数量是不确定的,因此无法直接按照语义分割的思路进行实例分割。

    为了解决这个问题,近几年提出的实例分割方法可以分为两大类:自顶向下(Top-down)自底向上(Bottom-up)的方法。自顶向下方法又称为“detect-then-segment”,先进行目标检测,然后在每个包围框内分割实例掩码。自底向上的方法学习对每个像素学习一个嵌入向量,然后分开不属于同一个实例的像素,而拉近属于同一个像素的实例,然后进行后处理分离出每个实例。两种方法都是逐步的、不直接的,要么严重依赖包围框的精度,要么依赖像素的嵌入学习和后处理。

    本文的目标是使用完整的实例分割标注作为监督,直接分割实例掩码。因此,重新思考实例分割问题:图像中实例之间的本质差异是什么?以COCO数据集为例,验证集共有36780个实例,98.7%的实例对之间的中心距离超过30个像素,剩下的1.7%的实例中40.5%的实例大小比例超过1.5x。因此,大部分图像中的两个实例之间具有不同的位置和大小。这个发现使得我们考虑是否可以根据实例的中心位置和大小区分它们?

    当前语义分割的主流方法是使用FCN预测N个通道的分割掩码。每个通道负责预测一个语义类别。语义分割的目标是区分出不同的语义类别,类似的,本文通过引入“实例类别”区分不同的实例,称为SOLO

    SOLO的核心思路是通过位置和尺寸从图像中分离出实例

位置:

    一个图像可以被划分为SxS的网格,共有S^2个位置类别。通过实例的中心区域,实例被分配到某个网格,作为它的位置类别。与DeepMask和TensorMask不同,这里编码位置类别作为通道轴,像语义分割那样输出张量的每个通道表示一个位置类别预测的实例掩码,这个实例掩码是中心落在该位置的实例的掩码。因此,结构几何信息自然地被保存在矩阵中。

    本质上,实例类别近似实例中心的位置。因此,通过将每个像素分类到对应的实例类别,等价于使用回归从每个像素预测实例的中心。这里将实例的位置预测转换为分类而不是回归,通过分类,可以通过固定数量的通道 对 每张图像中不同数量的实例 进行建模,同时不依赖于后处理(比如聚类或学习嵌入)。

尺寸

    为了区分不同大小的实例,这里采用特征金字塔网络(FPN),以便将不同大小的实例分配到不同层级的特征图,作为尺寸类别。由此,所有的实例被有规律地分离,从而根据实例类别进行分类。在本文的后面将展示FPN是SOLO的核心组件,对分割效果有巨大影响。

    基于SOLO框架,可以端到端的方式优化实例分割网络,使用掩码标注,执行像素级实例分割,而不需要包围框检测和像素聚合。而当前大多数实例分割方法需要包围框标注作为监督。首先,本文将展示SOLO可以与Mask R-CNN在COCO数据集上达到相同的表现。接着通过实例边界任务展示了SOLO框架的通用性。

    SOLO仅需要解决两个像素级分类问题,类似于语义分割,因此可以通过语义分割的tricks来提高SOLO。本质上,SOLO通过离散量化将边界回归万转换为分类问题。这么做避免了检测算法中常见的heuristic coordination normalization和log-transformation。SOLO具有令人尴尬的简单性(其它方法这么复杂,但实例分割可以这么简单)和强大的性能。



2.相关工作

自顶向下实例的分割方法

    在先验包围框中分割实例的方法是典型的Top-down方法。FCIS在RPN生成的RoI中组装位置敏感的得分图(position-sensitive score map),用于预测实例掩码。Mask R-CNN在Faster R-CNN的基础上进行拓展,增加了一个分枝用于在预测的包围框内进行实例分割。基于Mask R-CNN,PANet通过增强特征的表示来提高精度,Mask Scoring R-CNN增加了mask-IoU分枝来预测每个掩码的质量并打分来提高精度。TensorMask对每个像素采用稠密地滑动窗口的方法在局部窗口中分割实例,每个像素预定义窗口的数量和尺度。

图片1.jpg

    与上面这些top-down方法相反,SOLO是box-free的,因此不受包围框位置和尺度的约束,因此可以利用FCN的内在优势。

自底向上的分割方法

    这类方法通过像素聚合来生成实例掩码。[Associative embedding: End-to-end learning for joint detection and grouping.]使用关联嵌入将像素分组为实例。[Semantic instance segmentation with a discriminative loss function.]提出的判别损失函数通过推开不同实例的像素同时拉近同一实例的像素来有效的学习像素级实例嵌入。SGN将实例分割问题分解为一系列子分组问题。SSAP学习一个像素对亲和力金字塔,即两个像素属于同一实例的概率,然后通过级联图分区顺序生成实例。

    典型的bottom-up方法落后于top-down方法,尤其多样场景和类别的数据集中。与bottom-up方法不同,SOLO仅在训练时直接使用实例掩码标注进行学习,无需进行后处理,即可端对端预测实例掩码和语义类别。从这个意义上看,SOLO属于直接的端到端的实例分割方法。

直接的分割方法

    SOLO可能是第一个’direct’的方法,这里介绍记忆个’semi-direct’方法。AdaptIS首先预测point proposals,然后为了检测到的点出的实例生成掩码。PolarMask提出使用极坐标来编码掩码,并将每个像素的掩码转换为距离回归。他们不需要包围框进行训练,但是要么是逐步的,要么是基于compromise的,比如掩码的粗略参数表示。

    SOLO以图像作为输入,通过FCN以box-free和grouping-free的方式直接输出实例掩码和对应的类别概率。



3.SOLO

3.1问题表述

    实例分割:给定任意一张图像,一个实例分割系统需要确定是否图像中存在实例,若存在实例,则输出实例的掩码。SOLO的中心思想是将实例分割重新形式化为两个同时的类别感知(category-aware)预测和实例感知(instance-aware)掩码生成的问题。具体地说,SOLO将输入图像划分为SxS的网格,如果实例的中心落在某个网格,那么该网格负责:(1)预测语义类别,(2)分割实例。(这就是YOLO的思想阿),这两个是并行的预测分枝:

图片2.jpg

3.1.1 语义类别

    SOLO对于每个网格预测C维的语义类别的概率,C是语义类别的数量。如果将输入图像划分为SxS的网格,那么输出的语义类别空间是SxSxC,如Figure 2(Top)所示。这个设计基于每个网格最多有一个实例的假设。

3.1.2 实例掩码

    除了预测语义类别之外,每个网格还并行的预测实例掩码。对于输入图像I,如果将其划分为SxS个网格,那么总共将预测S^2个实例掩码。这里将S^2个掩码编码成一个3D掩码张量,具体地说,输出的实例掩码张量的维度为HxWxS^2。掩码张量的第k个通道负责分割(i,j)网格的实例,k=i*S+j,i,j∈[0,S-1]。由此,在语义类别和实例掩码之间建立了一对一的对应关系,如Figure 2所示。

    一个预测实例掩码的直接方法是采用全卷积网络,比如FCNs。然后传统的卷积运算具有空间不变性(spatially invariant),对某些任务来说空间不变性是有益的,然而这里更想要的是空间变动性(spatially variant),具体地说是位置敏感(position sensitive)。因为这里的分割掩码是根据网络进行排列的,每个网格在不同的输出通道中预测掩码。

    这里的解决方案是:在网络的初始阶段直接输入归一化的像素坐标到网络中,即CoordConv。具体地,创建一个与输入图像相同大小的张量,该张量包含归一化到[-1,1]的像素坐标。然后这个张量与输入特征进行concatenate,传到后面的层。通过这样简单的将坐标输入到卷积中,我们可以在传统的FCN中加入空间功能。当然CoordConv并不是唯一的选择,比如也可以使用半卷积操作,但是 CoordConv 更简洁、容易实现。如果原来的特征张量大小是H×W×D,现在新的张量就是H×W×(D+2),最后两个通道为xy像素坐标。更多的信息可参考CoordConv。

    形成实例分割:在SOLO中,类别预测和掩码预测的对应关系为:k=i*S+j。基于此,可以直接生成最终的实例分割结果。原始的实例分割结果通过聚合各个网格的预测结果生成的。最终,执行NMS获得最终的实例分割结果。


3.2 网络结构

    这里使用FPN作为backbone。FPN生成多个不同分辨率但是通道数固定(通常是256d)的特征图,这些特征图输入到后面预测头中,预测头包含两个并行的分枝:语义类别预测和实例掩码预测中。预测头的权重在不同层级的特征图是共享的,但是不同层级划分的网格数量不同,因此最后的预测头最后的1x1卷积层不共享参数。

    为了展示SOLO的通用型和高效,这里将SOLO搭配多个架构,它们之间的区别包括:(a)抽取特征的backbone架构,(b)计算实例分割记过的网络头,(c)损失函数。大多数实验是使用Figure 3所示的预测头架构:

图片3.jpg

    注意到本文使用的实例分割头的架构很简单,更复杂的设计可能提高性能,但这并非是我们工作的重点(句式学到了)。


3.3 SOLO学习

3.3.1 标签对齐

    对于类别预测分枝来说,需要给每个网格分配一个类别概率。具体地说,如果某个实例的ground truth的中心区域与网络(i,j)的重叠区域大于阈值时,该网格被视作正样本,否则视作负样本。这里采用了和Foveabox、FCOS类似的中心采用技术用于掩码分类。给定ground truth掩码的质心(cx,cy),宽w,高h,那么中心区域的为(cx,cy,aw,ah),其中a是常量尺度系数,这里设置为0.2,结果平均每个ground truth有三个正样本。

    除了实例类别的标签之外,每个正样本还需要二进制分割掩码标签。因为有S^2个网格,因此实例掩码分枝会输出S^2通道的张量。对于其中的每个正样本,会有对应的标签。理论上掩码的顺序很重要,但这里仅仅使用简单的行优先排序,结果表明也能工作的很好。

3.3.2 损失函数

    SOLO的损失函数为:

图片4.jpg

    Lcate是实例分类的Focal Loss,lambda系数设置为3,Lmask是掩码预测的损失函数,公式如下:

图片5.jpg

    索引i=floor(k/S),j=k mod S,即在网格上 从左到右,从上到下索引。Npos是正样本的数量,p*和m*分别表示预测的类别概率和掩码。1{p*i,j>0}是指示函数(indicator function),若1{p*i,j>0}则该函数值为1,否则为0。

    本文尝试使用不同的dmask损失函数:BCE、Focal Loss、Dice Loss等,最终采用Dice Loss,因为它即高效又稳定。Dice Loss定义如下:

图片6.jpg

    D是dice系数,定义如下:

图片7.jpg

    其中Px,y和Qx,y分别是预测的掩码和ground truth掩码在位置(x,y)上的像素值。


3.4 推断

    SOLO的推断过程很直接,给定一张输入图像,经过backbone和FPN,然后获得各个网格(i,j)的类别置信度Pi,j和对应的掩码mk,其中k=i*S+j。根据根据置信度阈值0.1,过滤掉低置信度的预测,然后选择top-500个掩码执行NMS。最后使用阈值0.5,将预测的掩码进行二进制阈值化。评估时保留top-100个实例mask。



4.实验

    在MS COCO数据集上进行实验,训练配置:在8GPUs上使用SGD训练,batch_size=16张图像,每个GPUs 2张图像。共训练36 epochs,初始学习率0.01,在27 epoch和33 epoch时学习率除以10。权重衰减系数为0.0001,动量系数为0.9,所有模型都爱ImageNet上进行预训练。输入图像的短侧被随机缩放至640-800之间。

4.1主要结果

    下表是SOLO与其它SOTA模型的比较,可以看到,在使用Res-DCN-101-FPN作为backbone的情况下,SOLO是达到了SOTA:40 AP。

图片8.jpg

    下图是分割结果示例:

图片9.jpg

图片10.jpg


4.2 SOLO的工作原理

    当S=12时,网络输出如下:

图片11.jpg

    从上图中可以看到,不同位置预测的实例掩码不同。这种显式地在不同位置进行实例分割的方式,SOLO将实例分割问题转换为位置感知(position-aware)的分类问题。每个网格仅预测一个实例,而多个相邻的网格可能会预测同一个实例。因此推断时通过NMS解决冗余的问题。


4.3消融实例

    网格数量:网格数量对性能的影响如下:

图片12.jpg

    在未使用FPN时,使用ResNet的C3、C4、C5输出来预测,可以看到S=12就可以达到27.2的AP。

    多层级预测:如Table 2所示,最后加入了FPN进行对比。FPN共5层,将不同大小的ground-truth分配到不同层级上。加入FPN后,SOLO可以处理多尺度的实例,因此AP升至35.8。

图片13.jpg

    CoordConv:如下表所示,增加了CoordConv前后的对比如下表:

图片14.jpg

    加入了CoordConv后,增加了3.6的AP,而更多的CoordConv没有带来明显的提升,单个CoordConv足够了。

    损失函数:下面是不同的损失函数的对比结果:

图片15.jpg

    结果表明Dice loss带来的提升最大。Dice Loss将像素看作整个目标,可以自动在前景和背景像素之间建立正确的平衡。

    类别分枝的对齐:在分类预测分支中,必须将HxW的特征匹配到SxS,这里比较三个不同的实现:interpolation、adaptive-pool、region-grid-interpolation。

    Interpolation:直接将HxW的特征图使用双线性插值法缩放到SxS。

    Adaptive-pool:对HxW应用2D适应性的最大池化,输出SxS。

    Region-grid-interpolation:对于每个网格,使用基于密集采样点的双线性插值,并将结果与平均值进行聚合。

    结果表明这三个方法的结果差不多。

    不同的预测头深度:SOLO的预测头使用FCN来预测分类和掩码。下图比较了不同的深度对结果的影响:

图片16.jpg

    可以看到7层后性能提升不明显,因此这里设置的预测头的深度为7层。SOLO中掩码根据空间位置来优化,而坐标是在预测头的开始加入的。因此掩码分枝必须有足够的表示能力来学习这种变换。


4.4 SOLO-512

    本文还训练了一个轻量级SOLO版本:SOLO-512,除了输入分辨率是512之外,其它的配置与SOLO一致。下表是SOLO和SOLO-512的对比:

图片17.jpg

    可以看到SOLO-512的精度略低于SOLO,但是速度快了接近一倍,拥有实时运行的潜力(这是在V100 GPU上测试的)。



5.Decoupled SOLO

    当S=20的时候,SOLO会输出400个通道的特征图,然而,这么多通道在大多数情况下是冗余的,因为一张图像中的实例数量通常不多。这里提出一个高效的SOLO版本:Decoupled SOLO,如下图所示:

图片18.jpg

    Decoupled SOLO中,原始的输出张量维度为HxWxS^2,可以替换为两个维度为HxWxS的张量,对应两个维度。因此输出空间由HxWxS^2下降到HxWx2S。对于在网格(i,j)中的实例,其本来的预测的掩码在输出张量的第k个通道,k=i*S+j。而对于Decoupled SOLO来说,这个预测的张量是两个输出张量对应通道的逐像素的乘积。

    下面是SOLO和Decouple SOLO的对比:

图片19.jpg

    可以看到Decouple SOLO与SOLO的精度差不多,而所需的内存将大大减小。



6.误差分析

    为了定量分析SOLO的掩码预测,这里将预测的掩码替换为ground truth,结果如下:

图片20.jpg

    这表明掩码分枝仍有较大的提升空间。



7.SOLO实例曲线检测

    结果如下:

图片21.jpg





上一篇:

首页 所有文章 机器人 计算机视觉 自然语言处理 机器学习 编程随笔 关于