人工智能

BrushNet:具有双分支扩散的图像修复

mm
BrushNet: Plug and Play Image Inpainting with Dual Branch Diffusion

图像修复是计算机视觉领域的一个经典问题,旨在恢复图像中的缺失区域,同时保持整体的连贯性。现有的使用传统图像修复技术的工作,例如生成对抗网络(GANs)和变分自编码器(VAEs),通常需要辅助手工工程特征,但同时也不能提供令人满意的结果。过去几年中,基于扩散的方法在计算机视觉社区中获得了普遍的欢迎,主要是因为它们具有生成高质量图像、输出多样性和细粒度控制的能力。最初尝试使用扩散模型进行文本引导图像修复的方法,通过从预训练的扩散模型中采样掩码区域,并从给定的图像中复制未掩码区域,修改了标准的去噪策略。虽然这些方法在简单的图像修复任务中表现出令人满意的性能,但它们在处理复杂的掩码形状、文本提示和图像内容时却遇到了困难,导致整体缺乏连贯性。这种不连贯性主要是由于这些方法对掩码边界和未掩码图像区域上下文的感知知识有限所致。

尽管这些模型在过去几年中取得了进展和发展,图像修复仍然是计算机视觉开发人员面临的重大挑战。当前对扩散模型进行图像修复任务的改进,通常涉及修改采样策略或开发专门用于图像修复的扩散模型,但这些方法往往会降低图像质量和语义的一致性。为了应对这些挑战并为图像修复模型铺平道路,我们将在本文中讨论BrushNet,一种新颖的即插即用双分支工程框架,它将像素级掩码图像特征嵌入到任何预训练的扩散模型中,从而保证连贯性和增强的图像修复结果。BrushNet框架引入了一种新范式,即将图像特征和噪声潜变量分成独立的分支。这种分离减少了模型的学习负担,并在层次化的方式中促进了对掩码图像信息的细致纳入。在BrushNet框架之外,我们还将讨论BrushBench和BrushData,它们分别用于分割基于的性能评估和图像修复训练。

本文旨在深入探讨BrushNet框架,我们将探索其机制、方法、架构以及与最先进框架的比较。让我们开始吧。

BrushNet:具有双分支扩散的图像修复

图像修复是一种尝试恢复图像中缺失区域同时保持整体连贯性的方法,已经成为计算机视觉领域的一个长期问题,困扰了开发人员和研究人员多年。图像修复在图像编辑和虚拟试穿等多种计算机视觉任务中都有应用。最近,扩散模型如Stable Diffusion和Stable Diffusion 1.5,展示了生成高质量图像的卓越能力,并为用户提供了控制语义和结构的灵活性。扩散模型的卓越潜力促使研究人员使用扩散模型进行高质量的图像修复任务,以满足输入文本提示的要求。

传统的基于扩散的文本引导图像修复框架的方法可以分为两类,采样策略修改专用图像修复模型。采样策略修改方法通过修改标准的去噪过程,从预训练的扩散模型中采样掩码区域,并从给定的图像中复制未掩码区域。虽然采样策略修改方法可以在任意扩散模型中实现,但它们通常会导致不连贯的图像修复结果,因为它们对掩码边界和未掩码图像区域上下文的感知知识有限。另一方面,专用图像修复模型通过扩展基础扩散模型的输入通道来融合损坏的图像和掩码,从而实现更令人满意的结果,但可能不是图像修复模型的最佳架构设计。

如以下图所示,专用图像修复模型在早期阶段融合了掩码图像潜变量、噪声潜变量、文本和掩码。这种专用图像修复模型的架构设计容易影响掩码图像特征,并阻止UNet架构中的后续层获得纯粹的掩码图像特征,因为文本的影响。此外,在单个分支中处理生成和条件增加了UNet架构的负担,并且这些方法还需要在不同扩散骨架的变体中进行微调,因此这些方法通常耗时且具有有限的可转移性。

可能看起来添加一个额外的分支来提取掩码图像特征是一个合适的解决方案,但现有的框架通常在直接应用于图像修复时会提取和插入不充分的信息。因此,现有的框架如ControlNet在与专用图像修复模型相比时会产生不令人满意的结果。为了有效地解决这个问题,BrushNet框架在原始扩散网络中添加了一个额外的分支,从而创建了一个更适合图像修复任务的架构。BrushNet框架的设计和架构可以总结为三个要点。

  1. BrushNet框架使用VAE编码器来处理掩码图像,而不是随机初始化卷积层。这样,BrushNet框架可以更有效地提取图像特征以适应UNet分布。
  2. BrushNet框架逐渐将完整的UNet特征层融入预训练的UNet架构中,采用一种分层的方法,实现了密集的像素级控制。
  3. BrushNet框架从UNet组件中删除了文本交叉注意力,以确保在额外的分支中仅考虑纯图像信息。另外,BrushNet模型还提出了一种模糊混合策略,以获得更好的一致性和更高的可控性,在图像的未掩码区域。

BrushNet:方法和架构

以下图表给出了BrushNet框架的简要概述。

如图所示,框架采用双分支策略来插入掩码图像引导,并使用模糊掩码的混合操作,以确保未掩码区域的更好保存。值得注意的是,BrushNet框架能够调整添加的比例,以实现灵活的控制。对于给定的掩码图像输入和掩码,BrushNet模型输出一个修复的图像。模型首先将掩码下采样以适应潜变量的大小,并将掩码图像作为输入传递给VAE编码器,以使潜变量空间的分布对齐。然后,模型将掩码图像潜变量、噪声潜变量和下采样的掩码连接起来,并将其用作输入。模型提取的特征然后添加到预训练的UNet层中,经过零卷积块。去噪后,模型使用模糊掩码将掩码图像和生成的图像混合。

掩码图像引导

BrushNet框架通过一个额外的分支将掩码图像特征插入到预训练的扩散网络中,该分支明确地将掩码图像特征的提取与图像生成过程分开。输入是通过连接掩码图像潜变量、噪声潜变量和下采样的掩码形成的。更具体地说,噪声潜变量提供了图像生成过程中的信息,并帮助框架增强掩码图像特征的语义连贯性。BrushNet框架然后使用变分自编码器从掩码图像中提取掩码图像潜变量。此外,框架采用立方插值来下采样掩码,以确保掩码大小与掩码图像潜变量和噪声潜变量对齐。为了处理掩码图像特征,BrushNet框架实现了预训练扩散模型的一个副本,并排除了扩散模型的交叉注意力层。预训练的扩散模型的权重作为强先验,用于提取掩码图像的特征,而排除交叉注意力层确保模型仅在额外的分支中考虑纯图像信息。BrushNet框架将特征层层插入到冻结的扩散模型中,从而实现了分层的密集像素级控制,并且采用零卷积层来建立可训练的BrushNet模型和锁定的模型之间的连接,确保在初始训练阶段,噪声不会影响可训练副本中的隐藏状态。

混合操作

如前所述,在潜变量空间中进行混合操作会导致掩码被重塑,通常会产生不准确的结果,BrushNet框架也遇到了类似的问题,即需要将掩码重塑以匹配潜变量空间的大小。另外,值得注意的是,变分自编码器中的编码和解码操作具有固有的局限性,可能无法确保完全的图像重构。为了确保框架能够重构未掩码区域的一致图像,现有工作已经实施了不同的技术,例如从原始图像中复制未掩码区域。虽然这种方法有效,但它通常会导致生成的最终结果缺乏语义连贯性。另一方面,像采用潜变量混合操作的方法,在保存未掩码区域的期望信息方面遇到了困难。

灵活控制

BrushNet框架的架构设计使其成为一种适合即插即用集成的选择,可以与各种预训练的扩散模型灵活地保存比例。由于BrushNet框架不会修改预训练扩散模型的权重,开发人员可以将其作为即插即用组件与微调的扩散模型集成,允许轻松采用和实验预训练模型。另外,开发人员还可以通过将BrushNet模型的特征融入冻结的扩散模型中,使用给定的权重w来控制未掩码区域的保存比例,从而提供对保存级别的期望控制。最后,BrushNet框架允许用户调整模糊比例,并决定是否实施模糊操作,从而轻松地自定义未掩码区域的保存比例,为图像修复过程提供了灵活的调整和细粒度的控制。

BrushNet:实现和结果

为了分析其结果,BrushNet框架提出了BrushBench,一种基于分割的图像修复数据集,包含超过600张图像,每张图像都有一个手动注释的掩码和字幕注释。基准数据集中的图像均匀地分布在自然图像和人工图像之间,并且在不同类别中也均匀分布,允许在不同类别中进行公平的评估。为了进一步增强图像修复任务的分析,BrushNet框架将数据集分为两个不同的部分,基于使用的方法:基于分割的方法和刷掩码方法。

定量比较

以下表格比较了BrushNet框架与现有的基于扩散的图像修复模型在BrushBench数据集上的性能,使用稳定扩散作为基础模型。

如图所示,BrushNet框架在掩码区域保存、文本对齐和图像质量方面表现出显著的效率。此外,像Stable Diffusion Inpainting、HD-Painter、PowerPaint等模型在图像内部修复任务中表现出强大的性能,但在外部修复任务中,尤其是在文本对齐和图像质量方面,却难以复制其性能。总体而言,BrushNet框架提供了最强的结果。

另外,以下表格比较了BrushNet框架与现有的基于扩散的图像修复模型在EditBench数据集上的性能,性能与在BrushBench数据集上观察到的性能相似。结果表明,BrushNet框架在具有不同掩码类型的广泛图像修复任务中提供了强大的性能。

定性比较

以下图表定性地比较了BrushNet框架与现有的图像修复方法,结果涵盖了人工智能和自然图像,跨越了不同图像修复任务,包括随机掩码修复、分割掩码内部修复和分割掩码外部修复。

如图所示,BrushNet框架在未掩码区域的连贯性和连贯区域方面提供了显著的结果,并成功地实现了对背景信息的感知,这得益于双分支解耦的实施。此外,预训练扩散模型的未修改分支还提供了在不同数据域(如动漫和绘画)中表现更好的优势,导致在不同场景中的性能更好。

最终思考

在本文中,我们讨论了BrushNet,一种新颖的即插即用双分支工程框架,它将像素级掩码图像特征嵌入到任何预训练的扩散模型中,从而保证连贯性和增强的图像修复结果。BrushNet框架引入了一种新范式,即将图像特征和噪声潜变量分成独立的分支。这种分离减少了模型的学习负担,并在层次化的方式中促进了对掩码图像信息的细致纳入。在BrushNet框架之外,我们还将讨论BrushBench和BrushData,它们分别用于分割基于的性能评估和图像修复训练。

专业为工程师,心为作家。 Kunal是一名技术作家,对AI和ML有着深厚的热爱和理解,致力于通过其引人入胜和信息丰富的文档来简化这些领域中的复杂概念。