DiffBIR论文阅读笔记

在这里插入图片描述

这篇是董超老师通讯作者的一篇盲图像修复的论文，目前好像没看到发表在哪个会议期刊，应该是还在投，这个是arxiv版本，代码倒是开源了。本文所指的BIR并不是一个single模型对任何未知图像degradation都能处理，而是用同一个framework解决BSR,BID和BFR等blind image restoration任务，每个任务对应一个模型参数。
文章提出，BIR可以分解为两个阶段：去除 image-dependent degradation （由一个restoration module来完成）和生成丢失的图像内容（由IRControlNet来完成，用现有的train好的text-to-image latent diffusion model）。现有利用diffusion模型来进行IR的方法，基本都是把LQ作为diffusion的condition，和xt一起送进denoising network。但是本文提出，使用有noise和artifact的图片作为condition会影响生成，导致生成结果有artifact。因此，本文先用一个restoration模型来去除degradation，再把去除了degradation的图片作为condition送进diffusion模型。
文章提出一个training-free region-adaptive restoration guidance，有点类似classifer-based的方法控制diffusion生成图片类别的做法那样，在diffusion的中间加梯度来影响diffusion的方向，对图像中低频的部分要求和guidance image（即一阶段restore的结果）更接近（MSE Loss），而高频区域则受MSELoss影响小以图取更好的细节生成能力。
一二阶段的训练是分开的。一阶段用的是现有的BID BSR BFR网络，用MSELoss来train。而且选的是这三个可以从HQ合成LQ的任务，用了大范围的参数来合成各种各样的LQ图片来训练。二阶段如图所示，首先把restored image用一个pretrained fixed VAE encoder提取特征cRM，然后和zt送进stable diffusion模型做reverse预测zt-1。黄色的部分则是刚刚提的guidance。

在这里插入图片描述

reverse用的是pretrained fixed stable diffusion model，加一个自己的controlnet，train的是control的部分，如下图所示：
guidance是这样做的，其实没什么知识点，知道classifier-based 的diffusion model是怎么控制生成图片的类别的就知道了。首先，原来的diffusion是从 $z_t$ 预测噪声epsilon减去得出 $z_0$ ，然后从 $z_t$ 和 $z_0$ 推 $z_{t-1}$ ，这就是一个reverse的step。那么guidance就是，在 $z_0$ 的基础上，计算一个损失，然后计算损失对于对于 $z_0$ 的梯度，乘以一个scale，加上去，相当于对 $z_0$ 做了一步梯度回传加更新，得到 $\hat z_0$ 。然后以更新后的 $\hat z_0$ 和 $z_t$ 来推 $z_{t-1}$ ，作为一个reverse的step。关键就是这个损失怎么算，一般guidance就是计算 $z_0$ 和guidance的相似度，本文也一样，guidance image就是一阶段restored 的image，但是我们希望只有平滑区域学restored image，有纹理区域不要学，让diffusion生成，所以就在算loss的时候乘了个平滑度图。
如上所示， $D(z_0)$ 就是stable diffusion的decoder，stable diffusion不是直接对图片进行reverse step，而是对特征进行diffusion，所以需要先用decoder把特征转成图像域，再和restored image算差，差还是CHW的，还没取平均，这时候先element-wise地乘以一个权重图 $W$ ，这个 $W$ 表示的是区域的平滑度，越高则越接近1，否则越接近0。从而这个损失对有纹理细节的区域影响相对小，而对平滑区域影响相对大，所以实现了前面的目的。 $W$ 怎么来的，很简单，先用sobel算子对restored image算梯度的模，然后切不相交的patch，算patch内部的平均，然后用1来减。
评价：guidance的loss的算法其实在很多无监督半监督的restoration方法中有用到，很常见。guidance也是标准的diffusion模型的guidance算法。加condition的方式也是现有的controlnet的方式。相比现有diffusion做restoration的方法，其实就是多了一步restoration，相当于拼起来一个restoration模型加一个diffusion模型，感觉有点臃肿。核心动机是两个，一个是改善diffusion的条件，不要拿LQ作为条件；第二个是先restore再接diffusion，并且把restore的结果作为guidance来使得整个方法把restore和generation解耦出来，diffusion专门做generation，要说diffusion的作用应该是对前面restore方法的结果的锦上添花，起添加一些细节上去的作用。因此从实验结果可以看到，经过两阶段后，确实是很多细节精细了很多。