Paper:https://arxiv.org/pdf/2503.24357
Code:https://github.com/shuaizhengliu/InstructRestore
Baseline:SD2.1
文章目录
前言
一、整体框架
3.1. Dataset Construction Pipeline
4.1. Training Framework
二、实验
实验设置
三、实验
5.2. Results on Localized Enhancement
5.3. Results on Images with Bokeh Effects
前言
问题:通过生成与图像底层内容语义一致的细节,基于 SD 的方法在合成真实图像细节方面比基于 GAN 的方法有显著提高。然而,现有的方法对整个输入图像进行统一的恢复,缺乏根据图像内容和用户指令调整区域自定义恢复级别的能力。
目前还没有一个公开可用的数据集可以提供这样的高质量图像、参考描述和相应的区域 mask 的三元组。与我们的任务最相关的数据集可以是参考分割数据集,如RefCOCO[46]。然而,它的图像质量和分辨率不足以支持 IR 任务。
动机:现有的基于生成先验的 IR 方法在处理此类图像时,可能会在背景区域产生许多不必要的纹理/细节,如图 1(a)所示。此外,图像的不同区域需要不同程度的恢复,以达到最佳的整体感知质量。例如,纹理区域通常需要更高层次的生成细节来实现更清晰的视觉外观,而平面区域需要更低程度的生成来实现视觉平滑,如图 1(b)所示。
遗憾的是,现有的方法很难实现不同区域的定制恢复。
贡献:数据集:我们开发了一个数据生成引擎。利用Semantic-Sam[16]和Osprey[48]模型,我们从一组选定的高质量图像中获得 mask 和初始描述。然后,我们使用大型语言模型(LLMs),更具体地说是Qwen[42],来迭代地解析和细化这些描述,格式化它们以满足IR任务的教学需求。最后,我们构建了一个包含536,945个三胞胎的数据集,涵盖了植物,建筑,动物等多种场景。
网络架构:为了确保模型能够准确识别人类指定区域并适当增强指定区域,我们提出将低质量输入图像的条件特征集成到ControlNet-like的架构中。指令作为文本提示符用于控制网络的控制分支。
在我们策划的数据集上训练,ControlBranch 可以同时生成区域 Mask 和条件特征。通过对用户自定义区域及其周围环境的条件特征应用不同的集成尺度,我们的 InstructRestore 模型实现了与用户意图一致的局部控制恢复。
一、整体框架
3.1. Dataset Construction Pipeline
(1)High-quality ground-truth image collection
高分辨率和高质量的真实值(GT)图像是训练红外模型的关键。因此,我们从LSDIR [18], EntitySeg[28]火车集和EBB![12]散景列车集,短边大于512像素,MUSIQ[15]得分大于60。
(2)Annotation pipeline
为了获得高质量的图像,我们设计了一个自动标注 pipeline,结合最先进的模型提取语义区域 Mask 及其相应的描述。
在 mask 提取阶段,我们首先利用最先进的分割模型,例如semantic - sam[16],为图像的语义区域生成粗粒度的语义分割掩码。对于来自EntitySeg[28]的图像,我们直接重用其预注释的 mask。一旦获得 mask,我们将每个图像与其 mask 配对,并将其输入多模态大语言模型Osprey[48]中,以生成区域级描述。这些描述作为指令的一部分,用于指定要处理或恢复的区域。
在这个阶段,虽然我们获得了初步的掩码和描述,但由于以下两个关键问题,它们作为我们的训练数据还远远不够完善:(1)semantic - sam[16]偶尔会为一个语义产生多个掩码片段,导致区域模糊,不利于区域定制学习; (2)由于回复的随意性,描述并不总是采用名词短语格式,不适合嵌入指令。
为了解决这些问题,我们首先利用Qwen-7B[2],一个大型语言模型,通过提示调优执行以下两项任务: (1)从描述中解析主题,(2)将其重新格式化为名词短语。
由于大型语言模型输出的不稳定性,例如,重复或拼写错误,我们迭代地执行改进过程。具体来说,我们通过一个更大的大型语言模型 Qwen-72B[2]识别错误案例并重新执行上述过程。此循环重复3次,以确保高质量的输出。
(3)Dataset Statistics
我们的数据集提供了高质量的GT图像、区域 Mask 和描述性说明的三联体。
为了强调我们数据集的相关性和实用性,我们将其与表1中最相关的参考分割数据集进行了比较,包括RefClef [14], RefCOCO [46], RefCOCO+[46]和RefCOCOg[24]。这些数据集还提供了语义区域的 mask 和标题。
我们的数据集不仅提供了536,945个注释区域(在规模上超过其他数据集),而且还提供了具有优越感知质量的更高分辨率图像,满足了 IR 任务的需要。
在我们的数据集中,我们基于从注释描述中提取的主题的频率绘制了一个词云。云中的每个单词的大小与其相对频率相对应,较大的单词代表更流行的主题。
4.1. Training Framework
我们观察到,在推理期间通过系数 α 缩放 ControlNet 特征可以灵活地控制数据保真度和语义增强。基于这一见解,我们提出了InstructRestore框架,它首先识别用户指令中指定的区域,然后执行区域自定义恢复。
(1)Architecture design
如图 4 所示,我们的 InstructRestore 模型由预训练的 SD 主干、ControlNet 和轻量级 Mask 解码器(为了准确定位用户指令中的目标区域)组成。预训练的SD模型在整个训练阶段被冻结。从用户指令中提取的区域标题作为SD模型的文本提示,提供语义指导并生成语义细节。
ControlNet 复制预训练 SD UNet 的编码器和中间块作为可训练的副本。它接收从 LQ 图像和用户指令中提取的特征作为输入,然后从输入中提取分层条件特征,并在多个尺度上注入到 SD UNet 解码器块中。
(2)Training proces
HQ图像首先由预训练的VAE编码器编码到潜在空间,得到
。扩散过程在随机采样的时间步长
上逐渐用高斯噪声破坏
,导致有噪声的潜在状态
其中,和
遵循余弦噪声调度。
ControlNet以和
作为输入,产生条件特征
,在训练时加入到缩放系数
的冻结 SD UNet 解码器中。训练流程在图中用红色突出显示。Mask 解码器以来自ControlNet的交叉注意特征
作为输入,生成目标区域掩码
,由具有 CrossEntropy损失的 GT Mask
监督。表示为
的 InstructRestore 网络以噪声潜在
、LQ图像
、指令
和区域标题
为条件,训练目标
结合了标准扩散损失和 Mask 监督
:
(3)Region-customized Inference
对于一般恢复,我们将模板设置为“make {region caption} clear with {s1}, and make other parts with {s2}”;
对于散景感知复原,模板设计为“{“make {region caption} clear with {s1}, and keep other parts bokeh blur with {s2}”。
{region caption} 指定感兴趣区域的文本标题(例如“The dog on The sand beach”),分别定义目标区域和背景区域的增强尺度。
指令解析如下:
1) 区域自定义恢复以“{region caption}”作为SD条件,感知散景恢复以“{region caption} in front of bokeh background”;
2) 指令主体与任务相适应:对于区域自定义恢复,将“make {region caption} clear”输入ControlNet的文本编码器,而对于散景感知恢复,则使用“make {region caption} clear and keep other parts bokeh blur”;
3) 通过正则表达式解析提取 Mask 调制的保真度尺度s1、s2。
训练后的 ControlNet 分支对降级图像和解析后的指令进行处理,生成一个恢复 Mask,该 Mask 根据U-Net上采样解码器各层的空间维度动态调整大小,生成多个尺度的 Mask
。在每一层1,调制映射计算如下:
其中,和
分别控制目标区域和背景区域的增强强度。调制的 ControlNet 特征
通过元素乘法与基本SD特性
融合:
二、实验
实验设置
我们的模型首先在通用退化数据集上进行 120K 次迭代训练,指导指令模板“Make the {region caption} clear”。通过将散景数据集与一般退化数据集相结合进行 14k 次迭代,继续进行训练。
在此阶段,一般退化数据集的采样概率设置为25%,散景数据集的采样概率设置为75%,并与“Make the { region caption } clear and keep other parts bokeh blur.”的指令模板配对。
训练在两个 A100GPU上进行,batch大小为64,初始学习率为5e−5。采用AdamW作为网络训练的优化器。
三、实验
5.2. Results on Localized Enhancement
(1)Localized enhancement with user-instruction
(2)Comparison with other methods
5.3. Results on Images with Bokeh Effects
(1)Control of bokeh effect
(2)Comparison with other methods
声明:若论文中有理解有误的地方,欢迎大家批评指正。