[论文解读] SegRefiner: Towards Model-Agnostic Segmentation Refinement with Discrete Diffusion Process

Mengyu Wang, Henghui Ding+4 more|arXiv (Cornell University)|Dec 19, 2023

Advanced Neural Network Applications被引用 6

一句话总结

SegRefiner 提出了一种与模型无关的分割优化框架，将掩码优化视为一种离散扩散过程，通过单向像素状态转换迭代去噪粗分割掩码。该方法在语义分割、实例分割和二值图像分割任务中均达到最先进性能，相较于先前的与模型无关方法，IoU 提升 +3.42，Mask AP 提升 +0.9，同时在高分辨率图像中表现出优异的细节保留能力。

ABSTRACT

In this paper, we explore a principal way to enhance the quality of object masks produced by different segmentation models. We propose a model-agnostic solution called SegRefiner, which offers a novel perspective on this problem by interpreting segmentation refinement as a data generation process. As a result, the refinement process can be smoothly implemented through a series of denoising diffusion steps. Specifically, SegRefiner takes coarse masks as inputs and refines them using a discrete diffusion process. By predicting the label and corresponding states-transition probabilities for each pixel, SegRefiner progressively refines the noisy masks in a conditional denoising manner. To assess the effectiveness of SegRefiner, we conduct comprehensive experiments on various segmentation tasks, including semantic segmentation, instance segmentation, and dichotomous image segmentation. The results demonstrate the superiority of our SegRefiner from multiple aspects. Firstly, it consistently improves both the segmentation metrics and boundary metrics across different types of coarse masks. Secondly, it outperforms previous model-agnostic refinement methods by a significant margin. Lastly, it exhibits a strong capability to capture extremely fine details when refining high-resolution images. The source code and trained models are available at https://github.com/MengyuWang826/SegRefiner.

研究动机与目标

• 为解决从多样化模型生成的粗分割掩码进行优化的挑战，且无需重新训练。
• 提升分割质量，特别是在边界准确性和细粒度细节捕捉方面，适用于不同分割任务。
• 开发一种可泛化的、与模型无关的优化框架，适用于多种分割架构和任务。
• 通过引入一种新型离散扩散过程，克服单步优化的局限性，实现通过迭代去噪进行优化。
• 实现在高分辨率图像上的掩码优化，同时保持性能下降最小化，并在不同数据集和任务间具备强泛化能力。

提出的方法

• SegRefiner 将分割优化建模为一个条件数据生成任务，其中输入图像引导迭代式掩码优化过程。
• 它提出一种新型离散扩散过程，训练期间每个像素从真实标签向粗掩码进行单向、随机的状态转移。
• 推理阶段，粗掩码作为噪声起始点，模型通过逐像素预测标签转移和概率，执行迭代去噪。
• 模型采用基于 U-Net 的架构，并配备条件去噪头，基于图像上下文和当前掩码状态预测每个像素的下一状态。
• 支持全局与局部优化策略，两者结合在高分辨率图像上实现最优性能。
• 该方法设计为与模型无关，可无缝集成到任意分割模型的输出中，无需修改网络结构。

实验结果

研究问题

• RQ1 离散扩散过程是否能有效实现跨多样化分割任务的与模型无关的分割掩码优化？
• RQ2 与单步优化相比，通过逐步迭代去噪在纠正复杂错误（如边界不准确和细粒度细节缺失）方面表现如何？
• RQ3 结合全局与局部优化策略对高分辨率图像分割的影响是什么？
• RQ4 输入分辨率的选择如何影响精度、推理速度与计算成本之间的权衡？
• RQ5 与先前方法相比，该离散扩散过程在高分辨率图像中对细粒度细节的保留能力如何？

主要发现

• 与最佳的先前与模型无关方法相比，SegRefiner 在语义分割任务中将平均 IoU 提升 +3.42，mBA 提升 +2.21。
• 在实例分割任务中，相较于之前的最先进方法，Mask AP 提升 +0.9，Boundary AP 提升 +2.2。
• 该方法展现出强大的泛化能力，在最近提出的二值图像分割任务中无需结构修改即达到最先进性能。
• 消融实验表明，离散扩散过程显著提升性能，相较于非扩散基线，mBA 提升达 2.54 个百分点。
• 全局与局部优化策略的结合实现最佳平衡，在 BIG 数据集上达到 94.85 IoU 和 77.64 mBA，优于单一策略方法。
• 推理时间与计算成本随扩散步数线性增长，14 步时每张图像耗时 2.94 秒，计算量为 1492 GFLOPs（实例分割任务）。