[读论文] Text-Driven Image Editing via Learnable Regions

Abstract

语言已经成为图像编辑的一种自然界面。在本文中，我们提出了一种基于文本提示的区域图像编辑方法，无需用户提供遮罩或草图。具体而言，我们的方法利用了现有的预训练文本到图像模型，并引入了一个边界框生成器，用于识别与文本提示对齐的编辑区域。我们证明了这种简单的方法可以实现与当前图像生成模型兼容的灵活编辑，并能够处理包含多对象、复杂句子或长段落的复杂提示。我们进行了广泛的用户研究，将我们的方法与最先进的方法进行比较。实验表明，我们的方法在根据语言描述操作图像时，能够以高保真度和真实性表现出竞争性性能。

1. Introduction

paper: https://arxiv.org/abs/2311.16432
code: https://github.com/yuanze-lin/Learnable_Regions

2. Related Work

3. Proposed Method

基于文本驱动的图像编辑通过调整输入图像的视觉内容，使其与文本中指定的上下文或修改保持一致。我们的目标是使文本到图像模型能够实现无需遮罩的局部图像编辑。为此，我们提出了一种区域生成网络，可以生成用于图像编辑的高质量区域。

3.1 Edit-Region Generation

给定输入图像 $X\in\R^{3\times H\times W}$ 和包含 $p$ 个单词的文本 $T\in\Z^p$ ，我们首先使用预训练的 vision transformer 模型 ViT-B/16 进行特征提取。该模型通过 DINO 自监督学习目标进行预训练。最后一层生成的特征 $F\in\R^{d\times h\times w}$ 已被证明包含对象的语义分割信息，这可以作为我们问题中的先验知识。

我们初始化了 $K$ 个锚点 $\{C_i\}^K_{i=1}$ ，这些锚点位于基于 DINO 预训练 Transformer 的 [CLS] token 查询生成的自注意力图中得分最高的 $K$ 个区域。其中，[CLS] token 用于指导定位对象的语义信息部分。

接下来，我们为每个锚点 $C_i$ 定义了一组边界框提议 $B_i=\{B_j\}^M_{j=1}$ ，每个边界框以其对应的锚点为中心。为了简化，我们通过单一参数对边界框进行参数化，使得每个 $B_j$ 为边长为 $j$ 的正方形。随后，我们训练一个区域生成网络，显式考虑由同一锚点生成的所有不同的边界框。对于给定的锚点 $C_i$ ，我们将得到：

$\begin{aligned} f_j&=\text{ROI-pool}(F,Bj)\\ S([f_1,\cdots,f_M])&=[\pi_1,\cdots,\pi_M] \end{aligned}$

其中 $[\cdot]$ 表示在通道维度上拼接特征，ROI 池化操作用于对特征 $F\in\R^{d\times h\times w}$ 进行池化，池化的区域由边界框 $B_j$ 决定，结果为一个特征张量 $f_j\in\R^{d\times l\times l}$ 。在实验中，我们将 $l$ 设置为 7。 $S$ 是所提出的区域生成网络，由两个卷积层和两个全连接层组成，连续层之间包含 ReLU 激活层。最终全连接层的输出被输入 Softmax 函数，用于预测边界框提议的得分，即 $\text{Softmax}([\pi_1,\cdots,\pi_M])$ 。

为了学习区域生成网络的参数，我们使用 Gumbel-Softmax 技巧，通过添加一个小的 Gumbel 噪声对 $\pi_j$ 进行重新参数化，其中 $g_j=-\log(-\log(u_j))$ ，且 $u_j\sim\text{Uniform}(0,1)$ 。在训练过程中，我们采用直通梯度估计（Straight-Through Gradient Estimation），即反向传播使用可微分变量（如 softmax），而前向传播仍然选择 $\arg\max$ ，将 $\pi$ 视为分类变量。

对于每个锚点，当我们获得具有最高 Softmax 分数的编辑区域后，我们首先生成一个对应的框状掩码图像，然后将该掩码图像、输入图像和编辑提示输入文本到图像模型，生成编辑后的图像。因此，我们可以针对所有锚点生成 $K$ 张编辑图像。

3.2 Training Objectives

由于 CLIP 模型可以估计图像与文本之间的相似性，我们利用其功能根据用户指定的提示词来指导图像编辑。为了训练我们的模型，我们提出了一种复合编辑损失函数，该损失由以下三个部分组成：

CLIP guidance loss： $\mathcal L_{\text{Clip}}$ 表示从 CLIP 编码器最后一层提取的生成图像与文本特征之间的余弦距离。
directional loss： $\mathcal L_{\text{Dir}}$ 用于控制在 CLIP 空间内应用编辑的方向
structural loss： $\mathcal L_{\text{Str}}$ 考虑了源图像与生成图像之间特征的自相似性，在保留源图像中物体原始空间布局的同时，支持纹理和外观的编辑。

$\begin{aligned} \mathcal L_{\text{Clip}}&=D_{\cos}(E_v(X_o),E_t(T))\\ \mathcal L_{\text{Str}}&=\parallel Q(f_{X_o})-Q(f_X)\parallel_2\\ \mathcal L_{\text{Dir}}&=D_{\cos}(E_v(X_o)-E_v(X),E_t(T)-E_t(T_{ROI})) \end{aligned}$

其中， $E_v$ 和 $E_t$ 分别表示 CLIP 模型的视觉编码器和文本编码器。 $X$ 、 $T$ 和 $X_o$ 分别表示输入图像、文本提示和由所提出的区域生成的编辑后图像。 $f_{X_o}$ 和 $f_X$ 分别表示由 CLIP 视觉编码器最后一层提取的视觉特征， $Q$ 表示相似性矩阵。需要注意的是， $T_{ROI}$ 表示源图像中指定用于编辑的感兴趣区域。

3.3 Inference

在推理过程中，我们定义了一个质量评分，用于对从不同锚点生成的编辑图像进行排序，并选择得分最高的图像呈现给用户。虽然存在更先进的方法，我们使用一个简单的加权平均来计算质量评分：

$S=\alpha\cdot S_{t2i}+\beta\cdot S_{i2i}$

其中， $S_{t2i}$ 估计给定文本描述与编辑图像之间的余弦相似性评分， $S_{i2i}$ 评估源图像与编辑图像之间的余弦相似性评分，我们使用从 CLIP 编码器最后一层提取的特征来进行相似性计算。

3.4 Compatibility with Pretrained Editing Models

区域生成器可以与各种图像编辑模型集成，用于根据提示修改源图像的内容。为了展示其多功能性，我们将其应用于两种不同的图像合成模型：MaskGIT 和 Muse 中使用的非自回归 Transformer，以及 Stable Diffusion 中使用的扩散 U-Net。

Transformer 和扩散模型代表了不同的基础编辑模型，用于验证所提方法的适用性。需要注意的是，MaskGIT 和 Muse 是基于离散 token 的 Transformer，它们的 token 是通过 VQ 自动编码器生成的，而扩散模型则是在连续空间中操作。因此，MaskGIT 和 Muse 的潜在空间只适用于盒状的掩码，不具备像素级掩码的精确度，这在图像编辑中是一个限制。

实验使用了官方的 MaskGIT 模型，而没有使用尚未公开的 Muse 模型。此外，我们将文本提示限制在模型训练时使用的类别词汇范围内。