文章标题:《DiffRAW: Leveraging Diffusion Model to Generate DSLR-Comparable Perceptual Quality sRGB from Smartphone RAW Images》
原文链接:DiffRAW
本文是清华大学深圳研究院 联合华为发表在AAAI-2024上的论文(小声bb:华为这个方向做的不少啊...)。
【1】论文的主要motivation 是解决Smartphone Raw to DSLR images的难点(小声bb:工业界的论文就是切实)。
- 智能手机拍摄相关的硬件本身就不如数码单反相机,硬件相关的限制导致了RAW图采集时在细节上必定会有损失。所以从智能手机感知到的RAW图推导出单反相机画质的图像本身就是一个困难的任务;
- 用手机采集RAW图,单反相机采集RGB图像,不同的设备拍摄同一个场景势必导致RAW-sRGB图像对空间位置不精确对齐的问题;
- 由于RAW-sRGB数据对是在不同的环境条件和相机参数下分别采集的,这会导致颜色差异以及不稳定的颜色映射关系。
【2】论文的主要贡献在于以下四点:
- 设计了一个新颖高效的前馈和逆向过程 ,即DTDM(Domain Transform Diffusion Method)。该方法减少了推理期间所用的迭代步数 。DTDM是一个对于基于扩散模型的图像恢复/增强方法通用的加速算法,可以灵活 的被迁移到其他基于扩散模型的图像恢复/增强方法中;
- 第一次使用RAW图作为扩散条件,在生成图像中保留了RAW图的轮廓和纹理信息;
- 通过设计的颜色-位置保留条件,本文缓解了由于颜色/空间不对齐导致的训练干扰;
- DiffRAW拥有一个颜色即插即用特征,使用不同的颜色条件作为颜色信息注入可以灵活的调整生成图像的颜色风格。
Preliminary
扩散模型
关于扩散模型的相关知识这里就不展开了,推荐看一下b站上李宏毅老师的扩散模型课程:扩散模型-李宏毅。
基于扩散模型的图像修复/增强
定义 x x x和 y y y分别表示低质量和高质量的图像。经典的图像修复/增强算法都会先构建一个关于高质量图像的图像序列作为扩散模型的前馈过程: { y t = α ‾ t y + 1 − α ‾ t ϵ } t = 1 T \{y_t = \sqrt{\overline{{\alpha}}_t}y + \sqrt{1-\overline{\alpha}t}\epsilon\}{t=1}^T {yt=αt y+1−αt ϵ}t=1T。
在训练过程中,低质量图像的信息会作为条件被注入到U-Net网络 f θ ( y t , x , t ) f_{\theta}(y_t,x,t) fθ(yt,x,t)中,网络最后会学到未知的条件分布 p ( y ∣ x ) p(y|x) p(y∣x)。
p θ ( y t − 1 ∣ y t ) = N ( y t − 1 ; μ θ ( y t , x , t ) , σ t 2 I ) p_\theta(y_{t-1}|y_t) = \mathcal{N}(y_{t-1};\mu_\theta(y_t,x,t),\sigma_t^2I) pθ(yt−1∣yt)=N(yt−1;μθ(yt,x,t),σt2I)
μ θ ( y t , x , t ) = 1 α t ( y t − 1 − α t 1 − α ‾ t f θ ( y t , x , t ) ) \mu_\theta(y_t,x,t) = \frac{1}{\sqrt{\alpha_t}}(y_t-\frac{1-\alpha_t}{\sqrt{1-\overline{\alpha}t}}f\theta(y_t,x,t)) μθ(yt,x,t)=αt 1(yt−1−αt 1−αtfθ(yt,x,t))
在推理过程中, y T ∼ N ( 0 , I ) y_T\sim\mathcal{N(0,I)} yT∼N(0,I)通常被用于生成过程的起始点,通过应用上述公式,目标图像 y y y会在T个迭代步以后被生成。
Abstract
由于明显的细节差异(discernible detail disparity) 、不稳定的颜色映射关系 和RAW-sRGB pair的空间尺寸不对齐 ,从智能手机RAW图推导出数码单反相机(DSLR,Digital Signal Lens Reflex Camera)质量的RGB图像 已经成为了一个很吸引人的挑战。本文提出了DiffRAW,一个第一次把扩散模型用于学习RAW-to-sRGB映射 的新颖的方法。本方法利用扩散模型学习到了DSLR图像的高质量细节分布 ,进而增强了输出图像的细节。同时,本文使用RAW图作为扩散模型的输入条件 以保持图像结构信息(例如,轮廓和纹理)。为了减轻由训练数据对中颜色和空间不对齐带来的干扰,本文在DiffRAW中嵌入了一个颜色-位置保持条件 ,以确保输出图像没有颜色和像素偏差问题。为了加速DiffRAW的推理速度,本文设计了域变换扩散方法 ,一个高效的扩散过程及其逆过程(reverse process)。域变换扩散方法在减少基于扩散模型的图像恢复/增强算法所需推理步数的同时,还提高了生成图像的质量。通过在ZRR数据集上的评估,DiffRAW在所有的感知质量评估指标(例如,LPIPS,FID,MUSIQ)上都展示了sota的性能,并且在SSIM和PSNR指标上达到了还不错的结果。
Methodology
图1. DiffRAW模型架构
颜色-位置保持条件
定义 w w w 为智能手机捕获到的RAW图,定义 y y y为数码单反相机捕捉到的目标RGB图像,定义颜色-位置保持条件为符号 c c c 。
在训练过程中, c t r a i n c^{train} ctrain是使用一个高阶退化模型 D 2 D^2 D2 (Wang et al. 2021)对 y y y做退化处理得到的结果,即 c t r a i n = D 2 ( y ) c^{train}=D^2(y) ctrain=D2(y)。本方法对该模型进行了微调,以确保 c t r a i n c^{train} ctrain 和 y y y的颜色一致性。
在测试过程中,使用lightweight ISPNet (Zhang et al. 2021)处理 w w w 得到 c t e s t c^{test} ctest 。
图2. 实验结果可视化
这里把实验结果可视化搬过来,以更好说明颜色-位置保持条件的作用。可以看到,图2 (a )是用简单isp pipeline处理后的rgb图像;图2 (b )没有加入任何条件,失去了轮廓、纹理、颜色信息(颜色映射和空间映射都学乱了 );图2 (c )加入了 w w w ,保持了原RAW图的轮廓和纹理信息;图2 (d )基本保持了图2(a)的颜色信息。这里注意,图2(d)不要跟图2(f)比,因为这是测试过程,是要看跟 c t e s t c^{test} ctest,即图2(a)的关系;图2 (e)是单反相机图像退化后得到的可视化结果。
DTDM (Domain Transform Diffusion Method)
这一小节的逻辑比较复杂,作者写得很乱,我尽量讲的详细一点。
在这一小节,作者指出,在forward过程中,其实没必要把 y y y 一直加噪到纯噪声,只要加噪到 y s y_s ys 就好了。这样在reverse过程中,从 y s y_s ys 往回倒,推出 y y y所需要的步数自然就少了。
那么这里有一个问题:在测试集中, y s y_s ys 怎么获得?(注:因为在测试集里,我们只有手机拍的raw图 w w w ,需要得到DSLR-qualified images y y y )
你可能已经想到了,我们既然有 w w w ,就可以用一个神奇的数学魔法 x t e s t = I S P ( w ) x^{test}=ISP(w) xtest=ISP(w)【文中用一个函数G替代ISP,但其实是一样的】,如果我们大概能证明 x s t e s t x_s^{test} xstest 和 y s y_s ys差不多,是不是就可以了?
不太令人满意的是, x s t e s t x_s^{test} xstest 和 y s y_s ys稍微有点区别。因为前者在图像质量差的域里面,后者在图像质量好的域里面,两者并不能完全等价。讲到这里,你或许已经发现了为什么这个方法叫做"域变换扩散方法"。
没错!聪明的作者朋友们指出了如下公式:
与 x t e s t x^{test} xtest相对应,有 x t r a i n = D 2 ( y ) x^{train}=D^2(y) xtrain=D2(y),作者试图在forward过程中,把 y y y变成 x s t r a i n x_s^{train} xstrain 。这不仅涉及到加噪过程,还在每一个微小步中涉及到一点点的域移动,域总体移动量大概可以认为是 ( y − x t r a i n ) (y−x_{train}) (y−xtrain) 。继而又有 x s t r a i n = α ‾ s x t r a i n + 1 − α ‾ s ϵ x_s^{train}=\sqrt{\overline{\alpha}_sx^{train}}+\sqrt{1-\overline{\alpha}_s}\epsilon xstrain=αsxtrain +1−αs ϵ。
到此为止,DTDM的核心思想就算是讲完了。总结如下:作者通过精巧的设计,令高质量的图像 y y y 在forward过程中,不仅做了加噪退化,还做了域偏移处理,最后得到了位于另一个域中的加噪数据点 x s t r a i n x_s^{train} xstrain (而非高斯纯噪声)。
因此,在测试过程中,作者使用了两个数学魔法, x t e s t = I S P ( w ) x^{test}=ISP(w) xtest=ISP(w) , x s t e s t = α ‾ s x t e s t + 1 − α ‾ s ϵ x_s^{test}=\sqrt{\overline{\alpha}_s}x^{test}+\sqrt{1-\overline{\alpha}_s}\epsilon xstest=αs xtest+1−αs ϵ,直接得到了另一个域中的数据点,再通过上述过程的逆过程,即可转移回数据点 y y y 。
本小节的剩余部分,作者对forward、train、reverse三个过程做了精细的数学推导。我们大致可以理解为:经典的扩散都是只加噪声,DTDM却还做了域偏移,那每次的微小偏移是偏移多少?作者要给出证明推导。此外,在新的模型框架下,原有的关于噪声预测的Loss也需要修正。总而言之,在推理阶段,DTDM不仅对 x s x_s xs 做了去噪,还执行了一个从 x x x 到 y y y 的域迁移,允许DTDM在很短的步数内把 x s t e s t x_s^{test} xstest 转换成 y y y ,同时提高了生成图像的质量。
Conclusion
本文提出了一个新颖的可以自适应解决smartphone raw to DSLR-quality sRGB图像的方案。融合RAW图信息保持轮廓,提出了高效的diffusion过程加强图像质量、减少推理步骤。
【思考1】思考了一下本文提出颜色注入特征 的过程,可能是如下所述的这种情况。①CNN版本 : 相关研究者发现训练数据集中的pair对颜色是不对齐的,这可能导致网络需要学习对应的颜色映射关系,但是映射关系也不稳定,这导致网络对于颜色的拟合效果不好,用CNN做这件事情就不太行了。②扩散模型基础版本 : 用raw图信息注入,消除纯高斯噪声,得到对应的RGB图像,但是也存在映射不稳定的问题(因为raw图注入也会产生一些颜色信息的注入),会干扰网络训练。③DiffRAW : 干脆将颜色信息剥离出来,直接作为一个输入,模型学raw输入的轮廓和纹理,学注入的rgb信息的颜色特征,而且还能通过这个模块改变颜色风格,一举两得。空间注入特征的推导同理。
【思考2】感觉本文公式好多,有些变量对齐起来比较麻烦。比如在原文Methodology的第二小节中的 c t r a i n = D 2 ( y ) c^{train}=D^2(y) ctrain=D2(y) 到了第三小节又变成了 x x x 。
【思考3】关于DTDM小节,本文写得不甚详细,望广大读者朋友批评指正。