【论文阅读】Reviving Iterative Training with Mask Guidance for Interactive Segmentation

论文题目：Reviving Iterative Training with Mask Guidance for Interactive Segmentation

论文地址：https://arxiv.org/abs/2102.06583

代码地址：https://github.com/saic-vul/ritm_interactive_segmentation#pretrained-models

文章贡献：

1. 提出了一个新的向主干发送编码点击的模型Conv1S，修改了论文[1]提出的迭代采样策略；

2. 对COCO和LVIS数据集进行合并，用于交互式语义分割，获得了先进的性能；

3. 证明了一个纯前馈模型与先进骨干架构可以达到甚至超过目前的最先进的结果。

1 背景和动机

与普通语义分割的区别：允许用户在几次迭代中使用交互式输入明确地控制预测，而普通的语义分割算法只能输入图像并在一次传递中输出分割掩码。

交互式分割的优点：

可以手动选择感兴趣的对象进行分割&使纠正分割预测错误成为可能；
加快数据标注过程，减少标记成本；
能够分割之前未见过的类；
可用在手机端照片编辑等方面。

交互式分割的类型：涂鸦、点击、边界点等。

其中，基于点击的方法最为广泛易于理解，而涂鸦对输入的判断不明确，边界点方法不够灵活，因此该论文主要研究基于点击的交互式语义分割方法。

论文实验认为交互式分割的性能主要取决于3个方面：baseline的强弱，mask的精细程度，以及数据集的大小。因此论文在LVIS和COCO数据集上进行训练。

论文扩展了交互式分割，使其能在现有的粗标签mask上进行修改，同时也扩展了模型的稳定性（避免了在加入新的点击后准确率下降）：

上图中红点表示负向点击，绿点表示正向点击。对于第一列的mask图，可以通过交互式分割方法进行修正，得到第二列的结果。

2 相关工作

现有的交互式语义分割方法旨在在不断提升分割精度的同时，减少推理运算时间，主要在于骨干网络架构的修改。而该论文不修改骨干网络架构，只修改一个纯前馈模型。

不同类型的互动反馈

bounding box：在选定区域内缺乏特定的对象参考，以及缺乏校正预测掩模的接口。--> 优化：将点击与边界框相结合，提供更具体的目标引导，并允许修正预测的蒙版。
extreme points（选定目标最左最右最上最下的4个点）：需要4次点击，类似于边界框，且不支持修正。
Scribbles：提供了更丰富的先验信息，相比于点击用户需进行更多的操作，且CNN很难模拟笔画，更多只能用在无需训练的方法中。
combines clicks with text：输入以更好地推断目标对象的属性，因此需要更少的单击。

分割掩码mask的细化

通常的方法是，在不改变全局mask情况下进行局部边界的细化，可作为网络框架的一部分，也可作为后处理（如CascadePSP）。

局部优化通常不需要提供额外的输入，而全局优化需要用户的额外反馈，交互式分割可以以点击作为反馈来指导优化mask（相当于是一种全局优化）。

3 交互式分割方法

（1）网络结构

交互式分割和语义分割类似，都是输入图片，然后输出预测mask。不同之处在于用户输入，即对网络内的编码输入进行处理。因此作者认为不需要重新设计网络架构，直接用效果好的语义分割网络就可以了，论文以DeepLabV3+和HRNet+OCR作为交互式分割模型的骨干。

点击编码

将坐标形式的点击操作编码成空间形式，论文对比了distance transform（DT）和disks with a small radius两种编码方法，其中disks效果更好。

如下图所示，当添加一个新点时，DT可能会产生较大的影响，特别是本来就只有几个点时；而disks产生的影响会较小，这种影响可能会混淆一个网络。

向骨干网络发送点击

语义分割骨干通常是以RGB图像作为输入，交互式分割处理额外输入的常用方法如下：

1. Conv1E：增加预训练模型第一卷积层的权值，使其接受n通道输入，而不是只接受RGB图像。将RGB图（3通道）和点击编码（2通道）合并为一个5通道的输入，传入一个增强了权值（？）的3x3卷积，使其通道数变为64，再继续后续的骨干网络。

2. Distance Maps Fusion (DMF)：距离映射融合模块，将连接了附加用户输入通道的图像转换为3通道输入，是该作者在他之前的论文[3]中提出的。如上图中DMF，将总的输入先后经过2个1x1卷积，先升维再降维，将5通道数据转化为3通道数据。

论文提出了一个新模块Conv1S来解决该问题，它不再将RBG图和点击编码合并，而是将点击编码经过一个1x1卷积升维到64，这与骨干网络中第一个3x3卷积的输出通道数相同。同时，输入图像经过3x3卷积得到一个结果，然后将这两者的结果进行逐元素相加的操作。

Conv1S虽然与Conv1E类似，但它允许为新的权重选择不同的学习速率，而不会影响预先训练的骨干的权重。

（2）迭代采样策略

迭代采样策略旨在模拟用户在训练期间的点击，论文介绍了2种方法：

1. 随机产生一组正的和负的点击，而不考虑它们之间的任何关系。但在实践中，每次新的点击都被放置在使用之前的点击集的网络所产生的预测的错误区域中。这使得无法将先前交互的mask集成到模型中，因为需要有序的交互和一系列相应的预测来成功地训练这样的模型。

2. [1]提出一个迭代训练过程，可以更加真实的模拟用户点击，集成之前交互的mask。但是完全的迭代采样计算量太大，因此一般采用随机采样来初始化，后续使用迭代采样过程来增加一些点击。

迭代训练过程如上图所示，具体步骤如下：

对比前一时刻预测的mask（如果之前没有mask那就是空）与真实值之间的区别，获取错误区域（图a的绿色部分）；
使用连通分量将这些差异像素分为不同的集群（图b中绿色黄色蓝色等）；
根据像素数确定这些集群中最大的群，在这个集群内产生一个点击（图c的红色点），如果之前在该集群中没有采集过点，那该点为此集群的中心点；如果之前采集过点，那么此次采集的点要确定跟之前采集的点以及该集群边界处的欧式距离最大，这主要是为了避免多个点之间距离过近；
如果该集群中对应的像素位置在目标对象上，则产生的点为正向点，否则为负向点。（图c中目标对象是人，而黄色集群在马身上，因此是红色负点）

实际操作中，选择中心点会使得对NoC评价指标过拟合，而选择集群边界的点会使得性能不稳定。因此该论文在[1]的基础上做了一些改进：

对获取的错误区域进行形态学腐蚀（morphological erosion）操作，使得区域面积减少4倍；
在训练期间，不为数据集样本保存模拟点击，而是分别为每个批次模拟用户点击。将采样的最大迭代次数限制为N次，每个批次可以一致地从0到N次迭代。

（3）合并前一步骤产生的mask

上图为[1]的网络结构图，将前一步骤的mask作为第三通道可选项，并分别为正、负编码点击提供两个通道。对于第一次交互，以及跳过迭代抽样的批次，第3通道的mask为空。该论文也采用[1]的形式。

（4）归一化局部损失

1. Binary cross entropy(BCE) 二值交叉熵损失

它平等地对待所有的例子，在后期，由于来自几乎正确分割区域的梯度传播类似于来自错误区域的梯度，因此减缓了训练。

2. Focal loss(FL) 焦点损失

表示网络的输出，p(i,j)表示预测在点(i,j)的置信度。当准确率增加时，p(i,j)接近1，则P(M)减小。因此FL也会随时时间推移而减缓训练。

3. normalized focal loss (NFL) 归一化焦点损失

NFL是作者在他之前的论文[2]中提出的。因为标准化，NFL的梯度不会而随着时间的推移而消失，并保持与BCE的总梯度相等。与BCE相比，收敛的更快准确率更高。

4 训练数据集

数据集选择

目前主流的基于CNN的交互式语义分割方法一般都使用Semantic Boundaries Dataset或者Pascal VOC dataset或两者的结合，共1w多张图，20个类，类别较少，对可预测对象的多样性有所限制。

数据集OpenImages有994k张图共350个类别，LVIS有100k张图超过1k个类别，且LVIS拥有很高的标注质量。因此LVIS数据集是最好的选择。

但LVIS存在类别数量不均衡的问题，由于COCO数据集和LVIS使用的是一组图像，因此从COCO中补充了相应类别的图片。对于这两者类别相交的部分，选择LVIS，最终实验数据集共104k图像。

参数设置

图像大小：将图像裁剪为320x480的大小，并按照0.75-1.4的倍率随机缩放。
数据增强：水平翻转，亮度、对比度、RGB值的随机抖动。

将整个数据集图像的一次传递作为一个epoch，共训练55epochs。

5 实验结果

测试数据集：GrabCut、Berkeley、DAVIS、Pascal VOC validation、SBD。
评价指标：NoC，IoU。NoC@85指IoU的阈值在85的情况下NoC的大小。
分割骨干：对比两种骨干网络，带有resnet的deeplabV3+和HRNet+OCR。

不同网络结构的参数量对比：

HRNet+OCR参数量小于DeepLab-ResNet。

收敛性分析

交互式分割算法的关键特性之一是随着点击次数的增加收敛到足够的精度。之前的研究优化了推理时间以达到收敛，该论文不改进任何方案，而是靠迭代训练过程来实现收敛。

下表为在SBD数据上训练后在3中数据集上测试的结果，第3列和第4列表示在点击了20次和100次后，没有达到IoU阈值90的图像数量。

网络架构消融实验

每个单元格包含两个结果“X/Y”，其中“X”和“Y”分别对应无f-BRS-B[3]和有f-BRS-B的评价，所有的模型都在SBD数据集上训练：

其中，点击编码策略Disks方法比DT效果好，且Disk原点半径为5时比半径为3时好；将编码点击喂入模型的策略中，Conv1S效果最好。

HRNet-18优于DeepLabV3+ ResNet-34。

训练数据集比较

上图是使用各种数据集的对比结果，在COCO数据集和论文组合的COCO+LVIS数据集上表现的性能相近，但COCO+LVIS数据量更多，模型泛化性能更好。

其中，在SBD上验证的结果中，以SBD和Pascal VOC+SBD为训练集的效果较好（第一个红框），这可能是因为训练集和验证集数据相似。而在其他验证集上，这两种训练数据集的效果没有COCO+LVIS好。

损失函数比较

评估了以COCO+LVIS训练，HRNet-18+OCR作为骨干网络时，使用不同损失函数的结果：

迭代训练消融

迭代采样策略中的最大迭代次数N=3时为最优，N>4时会导致模型不稳定以至于获得更差的结果：

下图为N=3时，4个验证集上mIoU随着点击次数增加的变化图：

蓝线是没有采用迭代策略的基线，橙/绿线是采用了迭代策略的结果，其中橙线没有前一步步骤的mask。

除了SBD数据集（比其他3个数据集图像数量大，因此相对平滑），当点击次数增加时，蓝线和橙线都出现了不稳定的情况。这表明，在交互式语义分割中，增加新的点击可能会导致分割结果更糟，而依赖于上一步骤的mask可以有效的解决该问题。

与以往工作的比较

带有IT-M的是指使用了迭代策略且设置N=3，否则是未使用迭代策略的结果。C+L表示数据集COCO+LVIS：

绿/红/紫线是该论文的结果，mIoU都更好，稳定性更好：

在Berkeley验证集上不同点击次数的IoU对比效果，大部分都收敛到了IoU90以上：

6 相关文献

[1] S. Mahadevan, P. Voigtlaender, B. Leibe, Iteratively trained interactive segmentation, in: British Machine Vision Conference 2018, BMVC 2018, Newcastle, UK, September 3-6, 2018, BMVA Press, 2018, p. 212.

[2] K. Sofifiiuk, O. Barinova, A. Konushin, O. Barinova, AdaptIS: Adaptive instance selection network, in: 2019 IEEE/CVF International Conference on Computer Vision (ICCV), IEEE, 2019, pp. 7355–7363. doi:10.1109/iccv.2019.00745.

[3] K. Sofifiiuk, I. Petrov, O. Barinova, A. Konushin, F-BRS: Rethinking backpropagating refifinement for interactive segmentation, in: 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), IEEE, 2020, pp. 8623–8632. doi:10.1109/cvpr42600.2020. 00865.