超详细！！！一文带你理解MosaicFusion: Diffusion Models as Data Augmenters for LargeVocabulary Instance Segmenta

一篇用扩散模型进行数据增强的文章，对视觉各种下游任务都有很多启发！

论文链接：https://arxiv.org/abs/2309.13042

代码链接：https://github.com/Jiahao000/MosaicFusion

一种名为MosaicFusion的基于扩散模型的数据增强方法，用于大词汇实例分割。该方法不需要额外的训练和标签监督，通过以下两个关键设计，使现成的文本到图像扩散模型成为有用的数据集生成器：

1. 图像生成：将图像画布分为多个区域，并在每个区域内同时运行单轮扩散过程，条件是不同的文本提示，从而在单张图像中生成多个对象。

2. 掩码生成：通过聚合与对象提示相关的跨注意力图，在不同层和扩散时间步长上生成相应的实例掩码，然后通过简单的阈值处理和边缘感知的细化处理来生成最终的掩码。

实验结果表明，MosaicFusion在LVIS长尾和开放词汇基准测试中，显著提高了现有实例分割模型的性能，尤其是在稀有和新颖类别上。

丰富数据多样性：MosaicFusion能够在单个图像中生成多个对象，显著增加了训练数据的多样性。这种方法可以生成许多不同类别的对象，特别是那些在真实数据集中稀少出现的类别。通过增加这些类别的出现频率，模型能够更好地学习到它们的特征，从而提升对这些稀有类别的识别能力。

写作动机：

1. 挑战与现状：实例分割任务需要手动标注大规模数据集，这一过程非常繁琐且昂贵，特别是在大词汇量的情况下。因此，扩展实例分割数据集的词汇量是非常困难的。现有的方法在处理长尾分布和开放词汇场景时性能较差。

2. 现有方法的局限性：尽管一些方法通过深度生成模型合成大量带有像素级标签的图像，但这些方法依赖于额外的分割架构训练，限制了它们的应用范围。

3. 研究目标：为了解决大词汇实例分割训练数据有限的问题，特别是在长尾和开放集场景下，提出了MosaicFusion，通过利用大规模文本到图像扩散模型生成大量合成数据，从而增强实例分割训练。

主要贡献：

1. 自动化数据增强管道：提出了MosaicFusion，这是一种基于扩散的自动化数据增强管道，可以扩展现有的实例分割数据集，生成图像和掩码，不需要额外的检测和分割模型来标注数据。

2. 多对象图像生成：提出了在单张图像中生成多个对象的技术，研究表明，生成多对象图像比单对象图像更有利。

3. 广泛实验验证：在两个具有挑战性的基准测试中（LVIS长尾和开放词汇实例分割）进行了广泛的实验，证明了MosaicFusion方法能够显著提高现有检测器和实例分割器的性能，尤其是在稀有和未见类别上。

实验框架

1. Variational Autoencoder (VAE)：

   1. VAE负责将图像编码为潜在向量，并将潜在向量解码为图像。这个过程在压缩的潜在空间中进行，而不是在像素空间中进行，以提高效率。

2. Time-conditional U-Net：

   1. 这个U-Net模型在不同的时间步长上对潜在向量进行去噪。通过在不同时间步上迭代地去噪，从随机噪声中逐渐生成清晰的图像。

3. Text Encoder (例如CLIP)：

   1. 文本编码器将文本输入编码为文本嵌入，用于在图像生成过程中指导扩散模型。CLIP模型是一个典型的文本编码器，通过将文本和图像嵌入到同一潜在空间中，实现文本到图像的生成。

MosaicFusion框架中的主要步骤

1. 图像生成（Image Generation）：

   1. 图像画布划分：首先将一个图像画布分成多个区域。

   2. 扩散过程：在每个区域内同时运行单轮扩散过程，条件是不同的文本提示。

   3. 生成多对象图像：通过控制扩散模型，在单张图像中生成多个特定位置的对象。

2. 掩码生成（Mask Generation）：

   1. 跨注意力图聚合：在扩散过程中，通过U-Net的跨注意力层聚合文本令牌对应的跨注意力图。

   2. 阈值处理和边缘细化：对聚合的注意力图进行阈值处理和边缘感知的细化处理，生成精细的实例掩码。

实验设置

• 使用的Stable Diffusion版本：论文中采用了开源的Stable Diffusion v1.4模型，并使用LMS（Laplacian Pyramid Sampling）调度器。

• 指导尺度因子：设定为7.5。

• 扩散步数：所有实验均进行50步推理。

• 区域大小：平均区域大小为384×512像素。

Stable Diffusion模型

我们构建的数据生成框架基于最先进的文本到图像潜在扩散模型，即Stable Diffusion（SD）（Rombach et al, 2022）。SD在压缩的潜在空间中运行扩散过程，而不是在像素空间中，以提高效率。它由三个组件组成：

1. 变分自编码器（VAE）（Kingma和Welling, 2014）：用于对图像进行编码和解码，将图像转换为潜在向量。

2. 时间条件U-Net（Ronneberger et al, 2015）：在不同时间步长上对潜在向量进行去噪。

3. 文本编码器（如CLIP，Radford et al, 2021b）：将文本输入编码为文本嵌入。

预训练的VAE将图像编码为潜在向量用于扩散训练。在推理过程中，生成过程从潜在空间中的随机高斯噪声z_T ∼ N (0, 1)开始，并在文本提示p的条件下逐步去噪。具体来说，在每个去噪步骤t = 1, ..., T中，给定zt和p，U-Net预测噪声估计项ϵ并从zt中减去它以获得zt−1。最终的去噪潜在向量z0被传递给VAE解码器D，以获得生成的图像I = D(z0)。

扩散过程

由于像Stable Diffusion这样的潜在扩散模型在潜在空间中运行扩散过程，我们需要将图像画布从像素空间适配到潜在空间。以下是具体步骤：

1. 像素空间到潜在空间的映射：

   1. 由于U-Net（和VAE）的全卷积特性，潜在空间中的坐标可以通过乘以一个上采样因子U来简单地映射到像素空间中。这个上采样因子的值等于Stable Diffusion中U-Net的上采样因子，即U = 8。

   2. 因此，给定像素空间中的区域坐标Ri = (Xi, Yi, Wi, Hi)，我们可以得到潜在空间中的对应区域坐标ri = (xi, yi, wi, hi) = (Xi/U, Yi/U, Wi/U, Hi/U)。

   3. 为了避免小数坐标，我们确保像素空间中的(Xi, Yi, Wi, Hi)可以被U整除。

2. 扩散过程的运行：

   1. 在将图像画布映射到潜在画布之后，我们在不同的潜在区域ri上同时运行N个扩散过程，这些扩散过程共享同一个噪声预测网络，初始噪声相同，但条件是不同的文本提示pi。

   2. 由于所有扩散过程共享同一个网络，因此内存消耗基本上与单个扩散过程影响的最大区域相同。

3. 生成多对象图像：

   1. 通过这种方法，我们可以在每个区域内生成多个指定对象的图像，从而获得一张包含多个对象的生成图像。

这部分内容解释了如何通过在潜在空间中运行扩散过程来生成图像。首先将图像画布从像素空间映射到潜在空间，然后在潜在空间的不同区域同时运行扩散过程。这些过程共享同一个噪声预测网络，并在初始噪声和不同文本提示的条件下生成图像。最终，我们可以在每个区域内生成多个对象，从而得到一张包含多个对象的图像。

创新1（分画布进行扩散）

1. 同时在每个区域运行扩散过程：

   1. 扩散过程（Diffusion Process）：这是一个生成图像的过程，通常通过逐步去噪从初始噪声图像生成目标图像。在这个方法中，扩散过程在每个区域上同时进行。

   2. 每个区域（Region）：图像画布被划分为多个区域，每个区域将生成一个特定的对象。例如，图像被划分为四个区域，每个区域生成不同的对象（如猫、狗、鸟等）。

2. 使用共享的噪声预测模型：

   1. 噪声预测模型：在扩散模型中，噪声预测模型用于预测并逐步去除图像中的噪声，最终生成清晰的图像。这里提到的共享噪声预测模型意味着在所有区域中使用同一个模型，这有助于保持生成过程的一致性和效率。

3. 根据不同的文本提示进行条件化：

   1. 文本提示（Text Prompt）：每个区域的图像生成是基于特定的文本提示进行的。文本提示描述了希望在该区域生成的对象，例如“a photo of a single cat”用于生成猫的图像，“a photo of a single dog”用于生成狗的图像。

   2. 条件化（Conditioning）：扩散模型根据这些文本提示来控制生成过程，以确保在每个区域生成预期的对象。

具体解释

假设我们有一个图像画布被分为四个区域：R1, R2, R3, R4。

1. 初始化：

   1. 在整个画布上初始化一个随机噪声图像。

2. 分区域扩散：

   1. 对每个区域分别提供一个特定的文本提示。

   2. 例如，区域R1R_1R1的文本提示为“a photo of a single cat”，区域R2R_2R2的文本提示为“a photo of a single dog”，区域R3R_3R3的文本提示为“a photo of a single bird”，区域R4R_4R4的文本提示为“a photo of a single tree”。

3. 同时运行扩散过程：

   1. 在同一时间，扩散模型在每个区域内运行，根据对应的文本提示逐步去噪，生成目标对象的图像。

   2. 使用共享的噪声预测模型确保每个区域的生成过程一致且高效。

最终结果

• 多对象图像：通过在每个区域生成特定的对象，最终得到一张包含多个不同对象的图像。每个对象位于预定义的特定位置。

• 控制生成位置：这种方法允许精确控制每个对象在图像中的生成位置，从而可以生成复杂的多对象场景，满足实例分割任务的需要。

总结

通过在每个区域同时运行扩散过程并使用不同的文本提示，MosaicFusion方法能够在单张图像中生成多个特定位置的对象。这种方法不仅提高了图像生成的多样性和质量，还能有效地扩展实例分割数据集，提升模型的性能。

创新2（生成对象掩码，包括交叉注意力）

1. 聚合跨注意力图：

   1. 跨注意力图（Cross-Attention Maps）：在扩散过程中，模型生成跨注意力图，这些图显示了每个文本令牌（例如“cat”或“dog”）对图像中每个像素的影响力。

   2. 聚合注意力图：针对每个对象的文本令牌（如“cat”），在不同的层（layers）和时间步（time steps）上生成的注意力图进行聚合。也就是说，将这些注意力图叠加在一起，得到一个综合的注意力图。这一步可以确保捕捉到对象在生成过程中的所有细节和变化。

2. 阈值处理：

   1. 阈值处理（Thresholding）：将聚合的注意力图转换为二值图像。通过选择一个阈值，将高于阈值的像素设置为1（表示属于对象的区域），低于阈值的像素设置为0（表示背景区域）。这一过程可以提取出对象的大致轮廓。

3. 边界优化：

   1. 边界优化算法（Edge-Aware Refinement Algorithms）：使用如双边解算器（Bilateral Solver, BS）等标准的边界优化技术，进一步细化掩码的边界。这些算法可以结合图像的颜色和空间信息，平滑掩码边缘，使其更准确地匹配对象的实际边界。

   2. 双边解算器（Bilateral Solver）：一种常用的边界优化技术，它通过考虑像素之间的相似性和空间距离，优化掩码的边缘，使得掩码边界更自然、更精确。

4. 生成合成标注数据集：

   1. 合成标注数据集：最终生成的图像和对应的精确实例掩码组成了一个合成的标注数据集。这个数据集可以用于训练现成的实例分割模型，帮助模型学习如何识别和分割图像中的对象。

具体解释

假设我们有一个文本提示：“a photo of a single cat”，并且我们已经生成了一张包含一只猫的图像。我们希望为这只猫生成一个精确的掩码。

1. 聚合跨注意力图

• 生成跨注意力图：在扩散模型的每个层和每个时间步，都会生成一个注意力图，显示“cat”文本令牌对图像每个像素的影响。

• 聚合这些图：将所有层和时间步的注意力图进行上采样到原始图像尺寸，并进行平均处理，得到一个综合的注意力图A^。

1. 阈值处理

• 归一化处理：将综合注意力图的值范围调整到0到1之间。

• 二值化处理：使用大津法（Otsu’s method）自动选择一个阈值，将高于阈值的像素设为1，低于阈值的像素设为0，得到一个初步的二值掩码M^。

1. 边界优化

• 双边解算器：对初步的二值掩码M^应用双边解算器，使掩码边缘更加平滑和准确，生成优化后的掩码M~。

1. 过滤和扩展掩码

• 过滤低质量掩码：通过连接组件分析，移除面积过小或过大的掩码，只保留高质量的掩码。

• 扩展掩码：将优化后的掩码扩展到整个图像画布的尺寸，得到最终的实例掩码M。

结果

通过这些步骤，最终得到包含猫的图像和对应的高质量掩码。这些合成的图像和掩码可以用于训练实例分割模型，提高模型在实际应用中的性能。

总结

通过聚合跨注意力图、阈值处理和边界优化，MosaicFusion方法能够生成高质量的对象掩码。这些掩码与生成的图像一起，组成了合成标注数据集，可以显著提升实例分割模型的训练效果。

假设我们要生成一张包含猫的图像，文本提示为“a photo of a single cat”：

1. 初始噪声图像：从随机噪声图像 zTz_TzT 开始。

2. 逐步去噪：

   1. 在第一个时间步 t=Tt=Tt=T，模型根据当前噪声图像 zT 和文本提示 p 计算跨注意力图 AT。

   2. 使用 AT 调整噪声预测过程，生成去噪后的图像 zT−1。

3. 重复上述步骤，在每个时间步 ttt 生成新的跨注意力图 AtA_tAt，并使用这些注意力图指导去噪过程。

4. 聚合注意力图：将所有时间步和层的跨注意力图进行聚合，得到一个综合的注意力图 A^。

5. 生成掩码：

   1. 对综合注意力图 A^ 进行二值化处理，生成初步的掩码。

   2. 使用边界优化算法细化掩码，得到最终的高质量掩码。

通过这种方式，跨注意力图在每个去噪步骤中都发挥作用，确保生成的图像与文本提示匹配，同时生成高质量的实例掩码。

具体例子

假设我们要生成一张包含猫的图像，并生成对应的掩码，文本提示为“a photo of a single cat”。

1. 初始噪声图像：从随机噪声图像 zT 开始。

2. 逐步去噪：

   1. 在每个时间步 t，生成当前图像状态 zt 和文本提示 p 的跨注意力图 At。

   2. 使用这些跨注意力图指导去噪过程，逐步生成更清晰的图像。

3. 不同层的跨注意力图：

   1. 在每个时间步 t，模型的不同层生成的跨注意力图捕捉对象的不同特征。

   2. 低层次的注意力图可能高亮图像中猫的大致位置，而高层次的注意力图则提供猫的细节特征。

4. 聚合跨注意力图：

   1. 将所有时间步和层的跨注意力图上采样到原始图像尺寸。

   2. 在所有时间步和层上对注意力图进行平均处理，得到一个综合的注意力图 A^。

5. 生成掩码：

   1. 对综合注意力图 A^ 进行阈值处理和边界优化，生成最终的高质量掩码。