人工智能 - EXITS：基于极值点低标注成本弱监督实例分割 | CVPR 2024 - 晓飞的算法工程笔记

EXITS将极值点视为真实实例掩模的一部分并将它们传播以识别潜在前景和背景点，所有这些都用于训练伪标签生成器，然后由生成器给出的伪标签反过来用于最终模型的监督学习。
来源：晓飞的算法工程笔记公众号

论文: Extreme Point Supervised Instance Segmentation

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

Introduction

实例分割是一项同时预测个体对象的类别和掩模的任务，由于深度神经网络的监督学习取得了显著进展。然而，手动为每个实例注释像素级掩模成本过高，这通常导致类别多样性和训练数据量的不足。这一问题引导研究界转向标签高效学习方法，如弱监督学习和半监督学习。

在这个势头上，最近使用box监督学习实例分割引起了相当大的关注。为了使用box监督训练实例分割模型，这些方法采用了边界框紧密性先验，即垂直（或水平）穿过边界框的线必须包含至少一个属于对象的像素（见图1）。这个先验通过各种损失函数来表达。尽管box监督已经被证明在保持注释成本低的同时对学习实例分割非常有效，但在这个方向上还有进一步改进的空间，特别是由于它忽视了极值点，这是常见的box注释过程的副产品，提供了一个有助于估计实例掩模的强有力线索。

如今，极值点在边界框注释过程中是免费提供的，人类注释者被指示点击目标对象的四个极值点，即最顶部、最左侧、最底部和最右侧的点，而不是点击边界框的两个角点。这是因为前者通常需要较少的注释时间，而后者通常需要多次调整初始框标签。此外，由于它们绝对是目标真实掩模的一部分，极值点为分割提供了一个在box监督中缺失的强有力线索。

受此启发，论文研究了使用极值点进行弱监督学习以进一步提高性能而不增加注释成本的实例分割，提出了EXtreme point supervised InsTance Segmentation（EXITS）框架，整体流程如图2所示。该框架将极值点视为真实实例掩模的一部分，并利用它们作为训练伪标签生成器的监督。然后，生成器产生的伪分割标签反过来用于最终模型的监督学习，该模型可以是任何用于实例分割的任意网络。

EXITS的成功关键在于如何使用极值点训练伪标签生成器。一种直接的方法是将极值点视为前景，边界框外的点视为背景，然后利用它们进行监督学习。然而，以这种方式训练的伪标签生成器无法生成清晰的对象掩模，因为由于极值点的稀疏性，在训练过程中大多数对象区域保持未标记状态。为了解决这个问题，EXITS通过传播box外的极值点和背景点来估计边界框内的潜在前景和背景点。传播过程基于预训练的Transformer编码器导出的点之间的成对语义相似性，从而分别揭示与极值点和附近背景语义上相似的前景和背景候选点。检索到的点与极值点和明确的背景点一起作为训练伪标签生成器的监督。

如图1所示，伪标签生成器生成了高质量的伪掩模，特别是当目标被分成多个部分时，伪分割标签的增强质量导致了我们最终模型性能的提升。这一成功之处在于标签传播是在所有点的全连接图上进行的，因此可以将极值点传播到空间上相距较远的点，减轻了边界框紧密性先验在遮挡情况下被违反的副作用。传统的基于边界框监督的方法在这种情况下往往失败，因为它们严重依赖于先验信息。

为了定量比较分开对象的伪标签质量，在Separated COCO上测量了伪标签质量，这是COCO的一个子集，仅包含分开的对象。在该数据集上，论文的方法在mIoU上超过了先前最佳方法7.3个百分点。论文进一步在三个公共基准数据集PASCAL VOC、COCO和LVIS上评估了EXITS，在这些数据集上，EXITS优于所有先前基于边界框监督的方法。

简而言之，论文的主要贡献有三个：

使用极值点来解决弱监督实例分割问题，这些极值点可以在边界框标注过程中获得，无需额外成本。
引入了一种点检索算法，该算法有效地利用极值点来估计边界框内点的标签。具体而言，该算法基于传播到极值点和背景点的概率来估计点的标签。
在三个公共基准测试中达到了最先进的水平，结果表明论文的方法能够生成高质量的伪掩码，尤其适用于分离的物体。

Proposed Method

EXITS包括两个阶段：（1）学习一个模型，使用极值点标签生成训练图像的伪分割标签，（2）使用伪标签训练实例分割模型。在第一阶段，使用围绕每个对象的极值点裁剪的对象图像作为伪标签生成器的输入，使得模型学习在裁剪图像内预测对象的二进制掩码。在第二阶段中的实例分割模型，即最终的模型，学习检测和分割多个对象。需要注意的是，伪标签生成器处理一个更容易的任务，即单个对象图像上的实例分割任务，这有助于提高它生成伪标签的质量。整个EXITS流程如图2所示。

由于第二阶段是传统的监督学习，可以应用于任何实例分割模型，因此主要阐述第一阶段，特别是EXITS如何为分割提供有效的监督学习。第一阶段的整体流程如图3所示。EXITS的关键思想是在给定极值点的情况下检索可能属于对象的像素，并将它们作为伪标签生成器的监督。这个想法是通过将极值点传播到输入对象图像中的其他像素来实现的，同时将极值点视为对象的真实像素的子集。

Motivation for Using Extreme Points

极值点被定义为对象在基本方向上最外层的像素：最顶部的点 $(x^{(t)}, y^{(t)})$ ，最左侧的点 $(x^{(l)}, y^{(l)})$ ，最底部的点 $(x^{(b)}, y^{(b)})$ ，最右侧的点 $(x^{(r)}, y^{(r)})$ 。有研究证明了标记这些点相对于传统方法—标记框的左上角 $(x^{(l)}, y^(t))$ 和右下角 $(x^(r),y^(b))$ 处更高效，因为这样的角落往往难以确定是否属于目标对象区域，注解者经常不得不多次调整他们初始标记的角落。另一方面，相对于角落，在目标对象边界上极值点可以轻松地进行标记并直接转换成一个包围框。此外，它们本身提供了更多有关目标对象形状和外观信息，因为它们位于对象边界上。

Learning Pseudo Label Generator

伪标签生成器的目标是在给定围绕对象裁剪的图像的情况下预测对象的二进制掩码，由一个视觉变换器（ViT）编码器和一个掩码解码器组成。检索可能属于对象（即前景）或背景的点，并使用检索到的点以及box外的极值点和明确的背景点作为监督来训练生成器。

具体来说，前景点的初始集合是从极值点派生的，如下所示： $\mathcal{P}_{\textrm{FG}}:= \big\{ (x^{(t)}, y^{(t)}-\delta), (x^{(l)}+\delta, y^{(l)}), (x^{(b)}, y^{(b)}+\delta), (x^{(r)}-\delta, y^{(r)}) \big\}$ ，其中 $\delta$ 是引入的一个小边距，用于将极值点推向对象中心，使得 $\mathcal{P}_{\textrm{FG}}$ 中的点更向内部，并更可靠地代表对象。另一方面，背景点的初始集合 $\mathcal{P}_{\textrm{BG}}$ 由位于由极值点定义的边界框之外的点组成。为了给边界框内未标记的点分配伪标签，记为 $\mathcal{P}_{\textrm{Box}}$ ，从 $\mathcal{P}_{\textrm{FG}}$ 和 $\mathcal{P}_{\textrm{BG}}$ 中传播的初始标签通过随机游走传播到它们，使用一个转移概率矩阵，即输入图像中点之间的成对语义相似性矩阵。具体而言，更可能从 $\mathcal{P}_{\textrm{FG}}$ 传播而不是从 $\mathcal{P}_{\textrm{BG}}$ 传播的点被视为伪前景，而更可能从$\mathcal{P}_{BG}$ 传播而不是从 $\mathcal{P}_{\textrm{FG}}$ 传播的点则被视为伪背景。

Constructing Transition Probability Matrix

为了捕捉点之间的语义相似性，EXITS利用从ViT编码器的多头自注意力（MHSA）获得的注意力矩阵。由于随机初始化或ImageNet预训练的ViT的注意力矩阵无法区分前景和背景，引入了一个额外的预训练ViT编码器，称为相似性提取器，该编码器在目标数据集上进行了几个时期的额外训练，使用多实例学习（MIL）损失：

$$ \begin{equation} \label{eq:mil_loss} \begin{aligned} \mathcal{L}_{\text{mil}} = \mathcal{L}_{\text{dice}}\big(&\text{Proj}_{x}(\mathbf{M}), \text{Proj}_{x}(\hat{\mathbf{Y}}_\text{box})\big) \\ &+ \mathcal{L}_{\text{dice}}\big( \text{Proj}_y(\mathbf{M}), \text{Proj}_y(\hat{\mathbf{Y}}_\text{box}) \big) \;, \end{aligned} \end{equation} $$

其中， $\mathbf{M}\in [0,1]^{H\times W}$ 是一个掩模预测， $\hat{\mathbf{Y}}_\text{box} \in \{0,1\}^{H\times W}$ 是边界框的区域， $\mathcal{L}_{\text{dice}}$ 表示Dice损失，并且 $\text{Proj}_x: \mathbb{R}^{H\times W} \mapsto \mathbb{R}^{W}$ 和 $\text{Proj}_y: \mathbb{R}^{H\times W} \mapsto \mathbb{R}^H$ 是投影运算，分别在输入矩阵的每列和每行向量上应用最大化操作。一旦训练完成，相似性提取器被冻结并用于计算伪标签生成器训练期间的转移概率矩阵。

将每个点视为完全连接图中的一个节点，并使用它们的语义相似性构建这些节点之间的转移概率。为了计算转移概率矩阵，裁剪后的图像被分成 $N \times N$ 个块并展平，然后输入相似性提取器。然后通过对来自变压器层多个注意力头的自注意力矩阵进行平均，得到相似性矩阵 $\mathbf{S} \in \mathbb{R}^{N^2 \times N^2}$ 。为了构建双随机形式的转移概率矩阵 $\mathbf{T}$ ，对 $\mathbf{S}$ 应用Sinkhorn归一化，该归一化由以下公式计算

$$ \begin{equation} \mathbf{T} = \frac{\mathbf{A} + \mathbf{A}^\top}{2}, \text{ where } \mathbf{A} = \text{Sinkhorn}(\mathbf{S}) \;, \end{equation} $$

其中， $\text{Sinkhorn}(\cdot)$ 是Sinkhorn-Knopp算法。

使用MHSA构建转移概率矩阵具有两个优势。首先，由于MHSA捕获了点之间的高级语义关系，因此得到的转移概率矩阵可以防止将点传播到外观相似但语义不同的其他点。其次，MHSA计算所有点对的相似性，从而自然地产生一个完全连接图的转移概率矩阵。这允许标签在对象的分隔段之间传播，增强了标签分配过程的准确性

Generating Pseudo Point Supervision

通过从 $\mathcal{P}_{\text{FG}}$ 和 $\mathcal{P}_{\text{BG}}$ 的每个成员到 $\mathbf{p}_i$ 的转移概率进行随机游走计算，为 $\mathbf{p}_i \in \mathcal{P}_{\text{Box}}$ 分配一个伪标签。将 $\mathbf{p}_i$ 的前景传播得分定义为 $\pi_i^{\text{(f)}}$ ，其计算方式为

$$ \begin{equation} \pi_i^{\text{(f)}} = \frac{1}{|\mathcal{P}_{\text{FG}}|} \sum_{\mathbf{p}_j \in \mathcal{P}_{\text{FG}}} \mathbf{T}^\alpha(j, i) \;, \end{equation} $$

其中， $\mathbf{T}^\alpha(j, i)$ 表示点 $\mathbf{p}_j$ 经过 $\alpha$ 次随机游走传播到点 $\mathbf{p}_i$ 的转移概率。类似地， $\mathbf{p}_i$ 的背景传播得分定义方式如下：

$$ \begin{equation} \pi_i^{\text{(b)}} = \frac{1}{|\mathcal{P}_{\text{BG}}|} \sum_{\mathbf{p}_j \in \mathcal{P}_{\text{BG}}} \mathbf{T}^\alpha(j, i) \;. \end{equation} $$

利用这些得分，定义伪前景点集 $\hat{\mathcal{P}}_{\text{FG}}$ 和伪背景点集 $\hat{\mathcal{P}}_\text{BG}$ 如下：

$$ \begin{equation} \begin{aligned} &\hat{\mathcal{P}}_\text{FG} := \big\{ (x_i, y_i): \exists (x_i, y_i) \in \mathcal{P}_\text{Box}, \, \pi_i^{\text{(f)}}-\pi_i^{\text{(b)}} \geq \tau_{\text{FG}} \big\} \\ &\hat{\mathcal{P}}_\text{BG} := \big\{ (x_i, y_i): \exists (x_i, y_i) \in \mathcal{P}_\text{Box}, \, \pi_i^{\text{(f)}}-\pi_i^{\text{(b)}} \leq \tau_{\text{BG}} \big\} \;, \end{aligned} \end{equation} $$

其中， $\tau_\text{FG}$ 和 $\tau_\text{BG}$ 是阈值超参数。

Point dropout

为增强伪点监督的多样性并防止过拟合，采用一种称为点丢弃的增强技术。每个时期，点丢弃独立地从 $\hat{\mathcal{P}}_\text{FG}$ 和 $\hat{\mathcal{P}}_\text{BG}$ 中随机消除一个子集，并在该时期的训练过程中排除已删除的子集。

Training Objective

Point loss

设 $(x_i, y_i)$ 表示点 $\mathbf{p}_i$ 的二维坐标，按以下方式构造稀疏二进制掩码 $\hat{\mathbf{Y}}\in \mathbb{R}^{N \times N}$ ：

$$ \begin{equation} \hat{\mathbf{Y}}(x_i, y_i) = \begin{cases} \mathsf{1} & \text{if } \mathbf{p}_i \in \mathcal{P}_{\text{FG}} \cup \hat{\mathcal{P}}_\text{FG} \\ \mathsf{0} & \text{otherwise} \;. \end{cases} \end{equation} $$

此外，按以下方式构造一个掩码矩阵 $ \mathbf{K} \in \mathbb{R}^{N \times N}$ ，该矩阵编码了具有点监督的区域：

$$ \begin{equation} \mathbf{K}(x_i, y_i) = \begin{cases} \mathsf{1} & \text{if } \mathbf{p}_i \in \mathcal{P}_\text{FG} \cup \hat{\mathcal{P}}_\text{FG} \cup \mathcal{P}_\text{BG} \cup \hat{\mathcal{P}}_\text{BG} \\ \mathsf{0} & \text{otherwise} \;. \end{cases} \end{equation} $$

采用 $\hat{\mathbf{Y}}$ 和预测的掩码概率 $\mathbf{M}$ 之间的Dice损失。在计算损失之前，将 $\mathbf{M}$ 下采样为 $\tilde{\mathbf{M}} \in [0,1]^{N\times N}$ 与 $\hat{\mathbf{Y}}$ 的大小匹配。此外，对 $\tilde{\mathbf{M}}$ 和 $\mathbf{K}$ 进行逐元素乘法，以便损失信号仅应用于标记点。在没有使用点检索算法检索到任何点的情况下，即 $ | \hat{\mathcal{P}}_\text{FG} \cup \hat{\mathcal{P}}_\text{BG} | = 0 $ ，还额外应用MIL损失（见公式\`1`）。点损失定义如下：

$$ \begin{equation} \mathcal{L}_{\text{point}} = \mathcal{L}_{\text{dice}}(\tilde{\mathbf{M}} \odot \mathbf{K},\hat{\mathbf{Y}}) + \mathbb{1}_{\{|\hat{\mathcal{P}}_\text{FG} \cup \hat{\mathcal{P}}_\text{BG}| = 0\}} \lambda_{\text{mil}} \, \mathcal{L}_{\text{mil}} \;, \end{equation} $$

其中， $\odot$ 是Hadamard积运算符， $\mathbb{1}$ 是指示函数， $\lambda_{\text{mil}}$ 是一个平衡超参数。

Conditional random field loss

为了进一步优化预测的掩码，EXITS采用CRF损失。具体来说，EXITS利用训练网络的指数移动平均得到一个教师网络，即伪标签生成器参数中的ViT编码器和mask解码器。随后，来自训练网络和教师网络的掩码预测被平均以获得 $\mathbf{M}^{\text{avg}}$ 。然后，通过使用平均场算法对 $\mathbf{M}^{\text{avg}}$ 进行CRF精化，并将其作为伪真实掩码使用Dice损失如下：

$$ \begin{equation} \begin{aligned} \mathcal{L}_{\text{crf}} = \mathcal{L}_{\text{dice}}(\mathbf{M},\textrm{CRF}(\mathbf{M}^{\text{avg}})) \; , \end{aligned} \end{equation} $$

其中， $\textrm{CRF}(\cdot)$ 是CRF操作。这种方法使网络能够逐步生成更详细的目标掩码。总之，EXITS的整体损失函数为

$$ \begin{equation} \begin{aligned} \mathcal{L}_{overall} = \lambda_{\text{point}} \mathcal{L}_{\text{point}} + \lambda_{\text{crf}} \mathcal{L}_{\text{crf}} \;, \end{aligned} \end{equation} $$

其中， $\lambda_{\text{point}}$ 和 $\lambda_{\text{crf}}$ 是平衡超参数。

Learning a Fully Supervised Model

在第二阶段，EXITS利用训练好的伪标签生成器生成伪掩码标签，作为全监督实例分割模型的真实标签进行训练。为了生成伪掩码标签，包含 $k$ 个实例的图像被裁剪到相应的极端点注释周围，并输入到生成器中，得到每个对象的伪掩码。实例分割和伪标记模型的解耦设计允许将伪标签无缝地集成到任何全监督实例分割模型中。