[论文阅读] SegDiff: Image Segmentation with Diffusion Probabilistic Models

[论文地址] [代码] [arXiv 21.12]

Abstract

扩散概率方法被用于最先进的图像生成。在这项工作中，我们提出了一种扩展此类模型的方法，用于执行图像分割。该方法进行端到端的学习，而不依赖于预先训练的骨干。输入图像中的信息和当前估计的分割图中的信息是通过两个编码器的输出相加而合并的。然后，额外的编码层和解码器被用来迭代完善分割图，使用一个扩散模型。由于扩散模型是概率性的，它被多次应用，其结果被合并成最终的分割图。新方法在Cityscapes验证集、Vaihingen建筑分割基准和MoNuSeg数据集上产生了最先进的结果。

Method

本文是第一篇将diffusion模型用于分割任务的文章。结构如下：

对于原始的diffusion来说，输入是一个高斯噪声，输出是一张真实图像$x_0$；那么由于本文是分割任务，因此输出的$x_0$是分割结果。现在问题来了，分割结果是从分割图像里来的，那么图像是怎么输进去呢？类似于条件GAN的思想，我们怎么去构建一个"条件diffusion"的模型呢？

本文的具体做法是，把diffusion中所使用的UNet进一步解耦，对于其编码器，将其进一步拆分为E，F，G，其中E负责与解码器D进行连接，而G负责编码原始图像，F负责编码上一步的噪声。写成公式就是：$${ϵ}_{\theta }(x_t,I,t)=D(E(F(x_t)+G(I),t),t)$$

这个基本就是本文的核心创新所在了，即怎么修改网络结构，把图像给塞到diffusion过程中去。此外，本文在测试阶段还有个小技巧，因为diffusion作为一种生成模型，每次生成的结果肯定都是不同的，而分割任务的答案是唯一的；而且，生成结果不一致，也不利于结果复现。本文的解决方案是测试多次取平均值。