NeruIPS 2023 | SegRefiner:通过扩散模型实现高精度图像分割

点击下方卡片,关注“CVer”公众号

AI/CV重磅干货,第一时间送达

点击进入—>【视觉和Transformer】微信交流群

扫码加入CVer学术星球可以获得最新顶会/顶刊上的论文idea和CV从入门到精通资料,以及最前沿项目和应用!发论文搞科研,强烈推荐!

addbc5a871739f1316928af447dbe6ca.jpeg

在CVer微信公众号后台回复:论文,即可下载论文pdf和代码链接!快学起来!

背景介绍

尽管图像分割在过去得到了广泛研究和快速发展,但获得细节上非常准确的分割 mask 始终十分具有挑战性。因为达成高精度的分割既需要高级语义信息,也需要细粒度的纹理信息,这将导致较大的计算复杂性和内存使用。而对于分辨率达到2K甚至更高的图像,这一挑战尤为突出。由于直接预测高质量分割 mask 具有挑战性,因此一些研究开始集中于 refine 已有分割模型输出的粗糙 mask。为了实现高精度的图像分割,来自北京交大、南洋理工、字节跳动等的研究者们引入了一种基于扩散模型Diffusion去逐步提高mask质量的方法。a1eed5a19ee0f4543cc520835ae8f023.png

论文:https://arxiv.org/abs/2312.12425

代码:github.com/MengyuWang826/SegRefiner

现有方法    

Model-Specific

一类常见的 Refinement 方法是 Model-Specific 的,其通过在已有分割模型中引入一些新模块,从而为预测 Mask 补充了更多额外信息,从而增强了已有模型对于细节的感知能力。这一类方法中代表性的工作有 PointRend,RefineMask,MaskTransfiner等。然而,这些方法是基于特定模型的改进,因此不能直接用于 refine 其他分割模型输出的粗糙 mask。

Model-Agnostic

另一类 Refinement 方法是 Model-Agnostic 的,其只使用原始图像和粗糙mask作为输入信息,如 BPR,SegFix,CascadePSP,CRM 等。由于这类方法在 Refinement 过程中未使用已有模型的中间特征,因此不依赖于特定分割模型,可以用于不同分割模型的 Refinement。然而,尽管这类方法能够有效地提升分割准确度,但由于粗糙 mask 中存在多种多样的错误预测(如下图所示),导致模型无法稳定地修正粗糙 mask 中的全部预测错误。

3c3bed2a07562efe0e36f3db24c94b97.png

实现目标

相比于 Model-Specific 的方法,Model-Agnostic 的方法能够直接应用于不同分割模型的 Refinement,从而有着更高的实用价值。更进一步地,由于不同分割任务(语义分割,实例分割等)的结果都可以被表示为一系列 binary mask,具有相同的表征形式,在同一个模型中统一实现不同分割任务的 Refinement 同样是可能的。因此,我们希望实现能够应用于不同分割模型和分割任务的通用 Refinement 模型。

如前所述,已有分割模型产生的错误预测是多种多样的,而想要通过一个通用模型一次性地更正这些多样性的错误十分困难。面对这一问题,在图像生成任务中取得巨大成功的扩散概率模型给予了我们启发:扩散概率模型的迭代策略使得模型可以在每一个时间步中仅仅消除一部分噪声,并通过多步迭代来不断接近真实图像的分布。这大大降低了一次性拟合出目标数据分布的难度,从而赋予了扩散模型生成高质量图像的能力。直观地,如果将扩散概率模型的策略迁移到 Refinement 任务中,可以使得模型在进行 Refinement 时每一步仅关注一些“最明显的错误”,这将降低一次性修正所有错误预测的难度,并可以通过不断迭代来逐渐接近精细分割结果,从而使得模型能够应对更具挑战性的实例并持续纠正错误,产生精确分割结果。   

在这一想法下,我们提出了一个新的视角:将粗糙 mask 视作 ground truth 的带噪版本,并通过一个去噪扩散过程来实现粗糙 mask 的 Refinement,从而将 Refinement 任务表示为一个以图像为条件,目标为精细 mask 的数据生成过程。

算法方案

扩散概率模型是一种由前向和反向过程表示的生成模型,其中前向过程通过不断加入高斯噪声得到不同程度的带噪图像,并训练模型预测噪声;而反向过程则从纯高斯噪声开始逐步迭代去噪,最终采样出图像。而将扩散概率模型迁移到 Refinement 任务中,数据形式的不同带来了以下两个问题:

(1) 由于自然图像往往被视作高维高斯变量,将图像生成的过程建模为一系列高斯过程是十分自然的,因此已有的扩散概率模型大多基于高斯假设建立;而我们的目标数据是 binary mask,通过高斯过程拟合这样一个离散变量的分布并不合理。

(2) 作为一种分割 Refinement 方法,我们的核心思想是将粗糙 mask 视为带有噪声的 ground truth,并通过消除这种噪声来恢复高质量的分割结果。这意味着我们扩散过程的结尾应当收敛到确定性的粗糙 mask(而非纯噪声),这也与已有的扩散概率模型不同。

针对上述问题,我们建立了如下图所示的基于“随机状态转移”的离散扩散过程。其中,前向过程将 ground truth 转换为“不同粗糙程度”的 mask,并用于训练;而反向过程用于模型推理,SegRefiner 从给出的粗糙 mask 开始,通过逐步迭代修正粗糙 mask 中的错误预测区域。以下将详细介绍前向和反向过程。

98cc21661b00f0e5445ab46d280b35e2.png

前向过程

前向过程的目标是将 ground truth 提供的精细 mask 逐步消融为粗糙的 mask,记前向过程每一步的变量为 90b3e32d845f6be3d8e9e3fc3c594940.jpeg,则前向过程应当满足:

(1) d926a25795d10c633f40c2c47be8c8d2.jpeg为 ground truth

(2) 74c6449562a7181b30933f6b2127f2de.jpeg为粗糙 mask

(3) 927829807887cf5d65eeebc3e78aaae0.jpeg介于023357a61dd4d12b202bbc71ef83b457.jpeg6570cad0abccc6bb9238a10b977de04f.jpeg之间,随 t 增大逐渐向粗糙 mask 演变    

基于这些限制条件,我们用随机状态转移来表述前向过程:假设变量37ca7d709918741397469fb9c1770d51.jpeg中的每一个像素都有两种可能的状态:精细和粗糙,处于精细状态的像素值与50c1e465b49179edb3e8d8e685f15dd3.jpeg保持一致,处于粗糙状态的像素则取 02903f4562e109e56435e8a3211b6ae9.jpeg的值(即使二者一致)。我们提出了一个“转移采样”模块来进行这一过程,如上图右侧所示。在每一个时间步,其以当前 mask a0b923667f2428bf087e6919f1d9125b.jpeg,粗糙 mask f6ac3690592c29a0e28a2ccda2c3db1d.jpeg以及状态转移概率作为输入。在前向过程中,状态转移概率描述了当前 mask c0066a8e0a709786e012df50eeded92b.jpeg中每个像素转移到980062dfbc8337422d7f8fc5859acec1.jpeg中的状态的概率。根据状态转移概率进行采样,可以得到后一个时间步 ee10c21342620eeb43cc3062da497e60.jpeg 每个像素的状态,从而确定其取值。这一模块确定了一个“单向”过程,即只会发生“转移到目标状态”的情况。这一单向性质确保了前向过程会收敛到136af63196db5058dfa1125fff0d2f8d.jpeg(尽管每一步都是完全随机的),从而满足了上述限制 条件(2),(3)。

通过重参数技巧,我们引入了一个二元随机变量e5a000a16850f4e053f73c94cbd4d99d.jpeg来描述上述过程:我们将 5d605f092373448f4603f9c54e643c11.jpeg表示为一个 one-hot 向量,用来表示中间掩模 12fc9057a1c2ad2230834c201f8d4380.jpeg中像素 090039a834bb210d3ef58c2fb4a44a8e.jpeg的状态,同时设置 eeecc1aadcba4ad453d6cd8743b23ec9.jpeg37ec7cb7876a55e0bb906ed8eaed2851.jpeg来表示精细状态和粗糙状态。因此,前向过程可以被表示为:

3d1ded8c90cfed7d4fa69f7374d031d8.png

其中 6844542582d771654148b58338ccad88.jpeg为超参数,而 7f46f26b55464ba51442249d951a16c2.jpeg对应了上述的状态转移概率,61b957e9194d2684fac0c72ab4fc6c5d.jpeg是状态转移矩阵。则前向过程的边缘分布可以表示为:

45bc32ba702383896915fa227227bee6.png

其中 6d30e46db2b0e16ee588ebcc03f8a839.jpeg。从而我们可以直接获取任何中间时间步 271bb0e0051300f4c9afb4d9a3dcaa90.jpeg的 mask 7013fd22581a547f167eb05165bce909.jpeg并用于训练,而无需逐步采样 ffe517b3e63d00d4de7a90e19865bd70.jpeg

反向过程

反向扩散过程用于模型推理,目标是将粗糙 mask 9be093f3f331fab4ab09c08467b28292.jpeg逐渐修正为精细 mask 6b9488619d9c77d011b9b0245fea9535.jpeg。由于此时精细 mask 9d8a87e849085979ea0129fd675479fe.jpeg和状态转移概率未知,类似 DDPM 的做法,我们训练一个神经网络 3c3535b2e9502c37a459fe96becf4956.jpeg,来预测精细 mask cb68b7b3e72c1fa420bdcf3dfa9c3ff3.jpeg,表示为:

57b46048ee228836c1b337091403bcaa.png

其中 d40770039b5dcbcf06e5173174298320.jpeg是相应的图像。a4a0bb079ebe83d87d1e00ba9b5424c5.jpeg2c205c9725dc2e78270149643fd4933c.jpeg分别表示预测的精细 mask 和其置信度分数。这里2ea5ba2e4480f37e2b65e05a1de756d6.jpeg表示了网络对于2d5700afe1b8ef2eac0bfe53def35f46.jpeg预测准确与否的置信度,故同样可以被视作8d5da4b3c0ac0b9eefeb0ae4a42c8ec3.jpeg中每个像素处于“精细状态”的概率。为了获得反向状态转移概率,根据前向过程的设定和贝叶斯定理,延续 DDPM 的做法,我们可以由前向过程的后验概率和预测的62f77e1cbf1a90d36db440999e58d3cb.jpeg得到反向过程的概率分布,为:   

3dc99c9802a445dcb63d64c29569f3cc.png

其中 8c462c72950949697e940ee9f6e0f76c.jpeg为反向过程的状态转移概率。给定粗糙 mask 7e136c87000216148fee184f7ab6fe58.jpeg以及相应的图像ebc4e734b57486ab37aec9ef171b7685.jpeg,我们首先将所有像素初始化为粗糙状态 63d8015e41930063d12e7b8f6c1d3611.jpeg,然后通过不断迭代地状态转移,逐渐修正 797023f4630cdc549499b26d12c01fef.jpeg中预测值。下图为一个推理过程的可视化展示。

02b87d2371d39c593c3fb46e007d559b.png

模型结构

任意满足 8afa5c77eab20a577a8efea14de388bf.jpeg形式的网络均可满足我们的要求,这里我们延续了之前工作的做法,采用 U-Net 作为我们的去噪网络,将其输入通道数修改为4(图像和 34be501da8d21162eceaa3cda207e428.jpeg在通道维度上串联),并输出1通道的改进掩模。

算法评估

由于 Refinement 任务的核心是获取细节精确的分割结果,在实验中我们选取了三个代表性的高质量分割数据集,分别对应Semantic Segmentation,Instance Segmentation 和 Dichotomous Image Segmentation。

Semantic Segmentation

如表1所示,我们在 BIG 数据集上将提出的 SegRefiner 与四种已有方法:SegFix,CascadePSP,CRM 以及 MGMatting 进行了对比。其中前三个为语义分割的 Refinement 方法,而 MGMatting 使用图像和 mask 进行 Matting 任务,也可以用于 Refinement 任务。结果表明,我们提出的 SegRefiner 在 refine 四个不同语义分割模型的粗糙 mask 时,都在 IoU 和 mBA 两项指标上获得了明显提升,且超越了之前的方法。

e8b2a0617f72fff62c26668af8871fa4.png    

Instance Segmentation

实例分割中,我们选择了之前的工作广泛使用的 COCO 数据集进行测试,并使用了 LVIS 数据集的标注。与原始 COCO 标注相比,LVIS 标注提供了更高质量和更详细的结构,这使得 LVIS 标注更适合评估改进模型的性能。

首先,在表2中,我们将提出的SegRefiner与两种 Model-Agnostic 的实例分割 Refinement 方法 BPR 和 SegFix 进行了比较。结果表明我们的 SegRefiner 在性能上明显优于这两种方法。

69ac279360fb842a99bcfe082095057b.png

然后在表3中,我们将 SegRefiner 应用于其他7种实例分割模型。我们的方法在不同准确度水平的模型上都取得了显著的增强效果。值得注意的是,当应用于三种 Model-Specific 的实例分割 Refinement 模型(包括PointRend、RefineMask 和 Mask TransFiner)时,SegRefiner 依然能稳定提升它们的性能,这说明 SegRefiner 具有更强大的细节感知能力。

2f3a48e4f8fe33c123afae7bf3a6f4de.png

Dichotomous Image Segmentation

Dichotomous Image Segmentation 是一个较新提出的任务,如下图所示,其数据集包含大量具有复杂细节结构的对象,因此十分适合评估我们 SegRefiner 对细节的感知能力。   

276ea40dc51433a5e8a13d8deddd57d5.png

在本实验中,我们将 SegRefiner 应用于6种分割模型,结果如表4所示。可以看到,我们的SegRefiner在 IoU 和 mBA 两项指标上都明显提升了每个分割模型的准确度。

9dba95509ec53058eaf0231b8708e7ac.png

可视化展示

a9f2d6476ae42166d2bdb4311731af10.png

3ad53beb57d50dc1fe92e5eaa9aee4c1.png

在CVer微信公众号后台回复:论文,即可下载论文pdf和代码链接!快学起来!

CVPR / ICCV 2023论文和代码下载

 
 

后台回复:CVPR2023,即可下载CVPR 2023论文和代码开源的论文合集

后台回复:ICCV2023,即可下载ICCV 2023论文和代码开源的论文合集
计算机视觉和Transformer交流群成立
扫描下方二维码,或者添加微信:CVer444,即可添加CVer小助手微信,便可申请加入CVer-计算机视觉或者Transformer 微信交流群。另外其他垂直方向已涵盖:目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF等。
一定要备注:研究方向+地点+学校/公司+昵称(如目标检测或者Transformer+上海+上交+卡卡),根据格式备注,可更快被通过且邀请进群

 
 
▲扫码或加微信号: CVer444,进交流群
CVer计算机视觉(知识星球)来了!想要了解最新最快最好的CV/DL/AI论文速递、优质实战项目、AI行业前沿、从入门到精通学习教程等资料,欢迎扫描下方二维码,加入CVer计算机视觉(知识星球),已汇集近万人!

▲扫码加入星球学习
 
 
▲点击上方卡片,关注CVer公众号
整理不易,请点赞和在看1e6a103ae0690f83d58596a242b2b977.gif
  • 1
    点赞
  • 7
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值