SegRefiner: Towards Model-Agnostic Segmentation Refinement with Discrete Diffusion Process

Abstract

在这篇论文中，我们探讨了一种提高不同分割模型产生的目标掩模质量的主要方法。我们提出了一种名为SegRefiner的模型无关的解决方案，它通过将分割细化解释为数据生成过程，从而对这个问题提供了一个新颖的视角。因此，细化过程可以通过一系列去噪扩散步骤顺利实现。具体来说，SegRefiner以粗糙掩模作为输入，并使用离散扩散过程对其进行细化。通过为每个像素预测标签和相应的状态转移概率，SegRefiner以条件去噪的方式逐渐细化噪声掩模。为了评估SegRefiner的有效性，我们在各种分割任务上进行了全面的实验，包括语义分割、实例分割和二分图像分割。结果表明，我们的SegRefiner在多个方面均表现出了优越性。首先，它在不同类型的粗糙掩模上始终提高了分割指标和边界指标。其次，它在性能上明显优于先前的模型无关细化方法。最后，在细化高分辨率图像时，它展现出了捕捉极细节的强大能力。

1. Introduction

虽然在过去几十年中对图像[35, 23, 24, 14, 11]和视频[54, 37, 38, 16, 15]的分割进行了广泛研究，但获取准确和详细的分割掩模始终是一项挑战，因为高质量的分割需要模型捕捉高级语义信息和细粒度纹理信息以进行准确预测。对于分辨率为2K或更高的图像，这一挑战尤为突出，因为这需要大量的计算复杂性和内存使用。因此，现有的分割算法通常会在较小的尺寸上预测掩模，从而不可避免地由于在下采样过程中丢失了细粒度信息而导致准确性降低。

由于直接预测高质量的掩模具有挑战性，一些先前的工作已经将注意力转向了对先前分割模型获取的粗糙掩模的细化。一种流行的方法[29, 57, 27, 28]是通过为掩模校正增加分割模型（和特征）的新模块。然而，这种方法通常是特定于模型的，因此无法推广用于改进其他分割模型生成的粗糙掩模。另一部分分割细化工作[47, 56, 12, 44]则采用模型无关的方法，只需将图像和粗糙分割作为输入进行细化。这些方法具有更大的实用性，因为它们适用于改进不同的分割模型。然而，粗糙掩模中存在的各种类型的错误（例如，对象边界上的错误、在高分辨率图像中捕捉细粒度细节的失败以及由于错误的语义而导致的错误）给细化模型带来了巨大挑战，从而导致性能不佳（参见图1）。

2 Related Work

2.1 Segmentation Refinement

分割细化的目的是改进预先存在的分割模型中的掩模质量。一些工作集中于增强特定的分割模型。PointRend使用MLP来预测从Mask R-CNN输出的低置信度分数的像素的标签。RefineMask将语义头部作为额外指导融入到Mask R-CNN中。MaskTransfiner使用独立的FCN来检测不连贯区域，并使用Transformer来细化它们的标签。这些工作在其前身模型上展示了显着的性能提升。然而，它们的改进范围是特定于模型的，它们缺乏直接细化来自其他模型的粗糙掩模的能力。也有一些细化方法采用模型无关的方法，比如。这些策略强调利用各种形式的输入，包括整个图像、边界补丁、边缘条等。尽管这些技术可以细化来自不同模型的粗糙掩模，但它们的适用性仍然局限于特定的分割任务。作为一个特例，SegFix，它在Cityscapes数据集中用于语义分割和实例分割，学习了边缘像素和内部像素之间的映射函数，随后用相应的内部像素预测替换不准确的边缘预测。虽然SegFix的性能受到其准确识别图像内对象的能力的显著限制，从而导致在具有更广泛类别范围的数据集上（例如，COCO）的性能下降。

2.2 Diffusion Models for Detection and Segmentation

最近，扩散模型在研究中受到了广泛关注。最初的研究主要旨在增强和拓展扩散框架。随后的研究开始将扩散模型应用于更广泛的任务，并构建了综合性的条件生成框架。将扩散模型应用于检测和分割任务也成为越来越多研究的焦点。一些研究捕获扩散模型的中间激活，并利用MLP执行逐像素分类。DiffusionDet和DiffusionInst将扩散过程调整为在对象框和掩模滤波器中执行去噪。其他一些研究将图像分割任务视为掩模生成，并且利用额外的图像编码器提取图像特征作为生成掩模的条件。据我们所知，我们的SegRefiner是第一个将扩散模型应用于图像分割细化任务的工作，并且它还率先放弃了高斯假设，采用了新设计的离散扩散过程来进行基于扩散的分割任务。

3 Methodology

3.1 Preliminaries: Diffusion Models

扩散模型包括一个前向过程和一个后向过程。前向过程$q(x_{1:T}|x_0)$使用马尔可夫或非马尔可夫链逐渐将数据分布$x_0\sim q(x_0)$转化为完全的噪声$x_T$，而后向过程则采用逐渐去噪的过程$p_{\theta }(x_{0:T})$，将随机噪声逐渐转换回原始数据分布。

Continuous Diffusion Models
大多数现有的连续扩散模型[22, 46, 41]遵循高斯假设，并定义$p(x_T)=\mathcal{N}(x_T|0,1)$。前向过程的均值和方差由超参数${\beta }_t$定义，而后向过程利用模型预测的均值和方差，因此形式化为：

$$q(x_t|x_{t-1})=\mathcal{N}\left(x_t|(1-{\beta }_t)x_{t-1},{\beta }_{t}I\right)\text{\hspace{1em}(1)}$$
$$p_{\theta }(x_{t-1}|x_t)=\mathcal{N}\left(x_{t-1}|{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t)\right)\text{\hspace{1em}(2)}$$

Discrete Diffusion Models
与连续扩散模型相比，对于离散扩散模型的研究较少。Sohl-Dickstein等人[45]首次引入了二元扩散，用于重建一维噪声二进制序列。$x_T$被定义为遵循伯努利分布$\mathcal{B}(x_T|0.5)$。前向过程和后向过程表示为：
$$\begin{gathered} q(x_t|x_{t-1})=\mathcal{B}\left(x_t|x_{t-1}(1-{\beta }_t)+0.5{\beta }_t\right),(3) \\ {p}_{\theta }(x_{t-1}|x_t)=\mathcal{B}(x_{t-1}|f_b(x_t,t)),(4) \end{gathered}$$

其中${\beta }_t\in (0,1)$是一个超参数，$f_b(x_t,t)$是一个预测伯努利概率的模型。在 DDPM [22] 取得巨大成功之后，Austin 等人 [2] 将离散扩散模型的架构扩展到了一个更一般的形式。他们将前向过程定义为一个离散随机变量在多个状态之间转换，并使用状态转移矩阵$Q_t$来描述这个过程：
$$[Q_t{]}_{m,n}=q(x_t=n|x_{t-1}=m).\text{\hspace{1em}(5)}$$

3.2 SeRefiner

在这项工作中，我们提出了 SegRefiner，具有独特的离散扩散过程，可用于改进来自各种分割模型和任务的粗糙掩膜。SegRefiner 使用粗糙到精细的扩散过程进行细化。在前向过程中，SegRefiner 使用离散扩散过程，该过程被构建为单向随机状态转换，逐渐将地面真实掩膜退化为粗糙掩膜。在逆向过程中，SegRefiner 从提供的粗糙掩膜开始，逐渐将粗糙掩膜中的像素转换为精细状态，纠正粗糙掩膜中的错误预测区域。在接下来的段落中，我们将提供前向和逆向过程的详细描述。

前向扩散过程：在前向过程中，我们逐渐将真实掩膜/精细掩膜$M_{fine}$转换为粗糙掩膜$M_{coarse}$。换句话说，我们有$m_0=M_{fine}$和$m_T=M_{coarse}$。在任何中间时刻 t ∈ { 1 , 2 , . . . , T − 1 } t \in \{1, 2, ..., T - 1\} t∈{1,2,...,T−1}，因此中间掩膜$m_t$处于$M_{fine}$和$M_{coarse}$之间的过渡阶段。

我们定义$m_t$中的每个像素具有两个状态：精细和粗糙，因此前向过程被构建为这两个状态之间的状态转换。处于精细状态的像素将保留其来自$M_{fine}$的值，反之亦然。我们提出了一个新的转换采样模块来表达这个过程。如右图所示，该转换采样模块在前向过程中，采用上一掩膜$m_{t-1}$、粗糙掩膜$m_T$和状态转换概率作为输入，并输出转换后的掩膜$m_t$。状态转换概率描述了$m_{t-1}$中每个像素转换为粗糙状态的概率。该模块首先根据状态转换概率执行 Gumbel-max 抽样，得到转换的像素。然后，转换的像素将从$m_T$中获取值，而未转换的像素将保持不变。

请注意，转换采样模块表示的是一个单向过程，只发生“转换为粗糙状态”的过程。单向属性确保了前向过程将收敛到$M_{coarse}$，尽管每一步都是完全随机的。这是我们的 SegRefiner 与以前的扩散模型之间的一个重要区别，在以前的扩散模型中，前向过程收敛到随机噪声。

通过重新参数化技巧，我们引入了一个二进制随机变量$x$来表述上述过程。我们将$x_{i,j}^t$表示为一个 one-hot 向量，以表示$m_t$中像素$(i,j)$的状态，并设置$x_{i,j}^0=[1,0]$和$x_{i,j}^T=[0,1]$分别表示精细状态和粗糙状态。因此，前向过程可以表示为：
$$q(x_{i,j}^t|x_{i,j}^{t-1})=x_{i,j}^{t-1}{Q}_t,where{Q}_t=\left(\begin{array}{cc}{\beta }_t& 1-{\beta }_t\\ 0& 1\end{array}\right)$$

在我们的转换样本模块中，其中${\beta }_t\in [0,1]$，而$1-{\beta }_t$对应于我们转换样本模块中使用的状态转移概率。$Q_t$是一个状态转移矩阵。$Q_t$的形式明确表明了单向属性，即粗粒度状态中的所有像素永远不会回过头转换到细粒度状态，因为$q(x_t|[0,1])=[0,1]$。根据式（6），边缘分布可以表示为：
$$q(x_{i,j}^t|x_{i,j}^0)=x_{i,j}^0{Q}_1{Q}_2\cdots Q_t=x_0{\bar{Q}}_t=x_0\left(\begin{array}{cc}{\bar{\beta }}_t& 1-{\bar{\beta }}_t\\ 0& 1\end{array}\right),(7)$$

其中，$\overline{{\beta }_t}={\beta }_1{\beta }_2\dots {\beta }_t$。因此，我们可以在任何中间时间步$t$上获取中间掩码$m_t$，而无需逐步采样$q(x_t|x_{t-1})$，从而实现更快的训练。

逆扩散过程 逆扩散过程接受一个粗略的掩码$m^T$，并逐渐将其转换为一个细粒度掩码$m^0$。然而，由于细粒度掩码$m^0$和逆转状态转移概率未知，我们按照 DDPM [22] 的方法，训练一个由参数$\theta$参数化的神经网络$f_{\theta }$，用于预测每个时间步的细粒度掩码${\tilde{m}}_{0|t}$，表示为
$${\tilde{m}}_{0|t},p_{\theta }({\tilde{m}}^0|t)=f_{\theta }(I,m^T,t),(8)$$
其中 I 是相应的图像。 \tilde{m}^0|t 和 p_θ(\tilde{m}^0|t) 分别表示预测的二值细粒度掩码和其置信度分数。为了获得逆转状态转移概率，根据式(6)，式(7)和贝叶斯定理，我们首先在时间步 t-1 下构造后验概率如下：
$$q(x_{t-1}|x_t,x_0)=\frac{q(x_t|xt-1,x_0)q(xt-1|x_0)}{q(x_t|x_0)}=\frac{x_t{Q}_t^T\odot x_0{\bar{Q}}_{t-1}}{x_0{\bar{Q}}_t{x}_t^T},(9)$$

在训练期间，精细状态$x_0$被设定为 [1, 0]，表示真实情况。然而，在推断过程中，$x_0$是未知的，因为预测的${\tilde{m}}_{0|t}$可能不完全准确。由于置信度分数$p_{\theta }({\tilde{m}}_{0|t})$表示模型对每个像素预测正确的确定程度，$p_{\theta }({\tilde{m}}_{0|t})$也可以被解释为处于精细状态的概率。因此，直觉上，可以通过简单地设置阈值来获得${\tilde{m}}_0{|}_t$中每个像素的状态，就像 [2] 中所做的那样：
$$x_{i,j}^0{|}_t=\{\begin{array}{ll}[1,0]& \text{if }{p}_{\theta }({\tilde{m}}_0{|}_t{)}_{i,j}\ge 0.5\\ [0,1]& \text{otherwise}\end{array}$$

像素具有更高置信度分数的情况下，将具有$x_{i,j}^0{|}_t=[1,0]$，表示它们处于精细状态，反之亦然。然而，在这种一位有效形式下，状态转移概率的值仅由预先定义的超参数确定，导致信息严重损失。相反，我们保留软状态转移并制定$x_{i,j}^0{|}_t=[p_{\theta }({\tilde{m}}_0{|}_t{)}_{i,j},1-p_{\theta }({\tilde{m}}_0{|}_t{)}_{i,j}]$。通过这种设置，逆扩散过程可以重新制定为：

$p_{\theta }(x_{i,j}^{t-1}|x_{i,j}^t)=x_{i,j}^t{P}_{i,j}^{\theta ,t}$

其中，$$P_{i,j}^{\theta ,t}=\left(\begin{array}{cc}1& 0\\ \frac{p_{\theta }({\tilde{m}}_0,t{)}_{i,j}({\bar{\beta }}^{t-1}-{\bar{\beta }}^t)}{1-p_{\theta }({\tilde{m}}_0,t{)}_{i,j}{\bar{\beta }}^{t-1}}& \frac{1-p_{\theta }({\tilde{m}}_0,t{)}_{i,j}}{1-p_{\theta }({\tilde{m}}_0,t{)}_{i,j}{\bar{\beta }}^t}\end{array}\right)$$

其中，$P_{i,j}^{\theta ,t}$是反转的状态转移矩阵。通过上述反转的状态转移概率，以及输入$m_t$和${\tilde{m}}_0{|}_t$，过渡样本模块可以在每个时间步将一部分像素转换为精细状态，从而纠正错误的预测。

推断过程：给定一个粗糙的掩码 $m_T$ 和其对应的图像 $I$，我们首先初始化所有像素处于粗糙状态，即 $x_{i,j}^T=[0,1]$。然后我们在以下步骤之间进行迭代：（1）进行前向传递以获得 ${\tilde{m}}_{0|t}$ 和 $p_{\theta }({\tilde{m}}_{0|t})$（参见式（8））；（2）计算反转状态转移矩阵 $P_{i,j}^{\theta ,t}$ 并获得 $x_{t-1}$（参见式（11））；（3）基于 $x_{t-1}$，$m_t$ 和 ${\tilde{m}}_{0|t}$ 计算细化掩码 $m_{t-1}$。重复步骤（1）至（3），直到获得细化掩码 $m_0$。推断的可视化示例如图3所示。

4 Experiments

4.1 Implementation Details

模型架构 遵循 [39]，我们采用 U-Net 作为我们的去噪网络。我们修改了 U-Net 来接收 4 通道的输入（图像和相应的掩码 $m_t$ 的串联），并输出一个 1 通道的细化掩码。输入和输出的分辨率都设置为 256×256。除了上述修改之外，所有其他部分保持不变。

目标函数 遵循 [12]，我们采用二元交叉熵损失和纹理损失的组合来训练我们的模型，即 $L=L_{bce}+\alpha L_{texture}$，其中纹理损失被定义为预测掩码和地面真实掩码的分割梯度幅度之间的 L1 损失。$\alpha$ 设置为 5 以平衡两个损失的幅度。

噪声时间表 理论上，SegRefiner 的单向性质确保任何噪声时间表都可以使得前向过程在无限数量的时间步长下收敛到粗糙掩码。然而，在实践中，我们使用更少的时间步长（本文中 $T=6$）以确保高效的推断。我们将 ${\bar{\beta }}_T=0$ 设为0，使得所有像素的 $x_T=[0,1]$ 并且 $m_T=M_{coarse}$（参见式（7））。类似于 DDIM [46]，我们直接设置了一个从 0.8 到 0 的线性噪声时间表。

训练策略 SegRefiner 模型已经开发成两个版本，用于不同分辨率图像的细化：一个是低分辨率版本（以下简称 LR-SegRefiner），另一个是高分辨率版本（以下简称 HR-SegRefiner）。虽然这两个版本采用不同的训练数据集，但所有其他设置保持一致。LR-SegRefiner 在 LVIS 数据集上进行训练，而 HR-SegRefiner 则在两个高分辨率数据集 DIS5K 和 ThinObject-5K 合并后的复合数据集上进行训练。选择这些数据集是因为它们具有高度精确的像素级注释，从而有助于我们的模型更有效地捕捉细节。按照 [12] 的做法，用于训练的粗糙掩码是通过各种形态学操作获得的，比如随机扰动一些地面真实边缘点并进行膨胀、腐蚀等。在训练过程中，我们首先在低分辨率数据集上训练 LR-SegRefiner 直到收敛。随后，它在高分辨率数据集上进行微调，形成 HR-SegRefiner。

采用双随机裁剪作为主要的数据增强技术。它涉及对每个图像进行随机裁剪，以生成两个不同的模型输入：一个封装整个前景目标，另一个仅包含部分前景。然后，两个裁剪都被调整大小以匹配模型的输入大小。这个操作确保模型能够熟练地细化整个对象和不完整的局部区域，这在后续的实验中被证明是非常重要的能力。所有下面的实验都是在 8 个 NVIDIA RTX 3090 上进行的。