ShadowDiffusion: When Degradation Prior Meets Diffusion Model for Shadow Removal (Paper reading)

ShadowDiffusion: When Degradation Prior Meets Diffusion Model for Shadow Removal

Lanqing Guo, Nanyang Technological University, Singapore, arXiv, Cited:1, Code: 无, Paper.

1 前言

最近的深度学习方法在去除图像阴影方面取得了令人满意的结果。然而，由于嵌入前的退化和建模能力的不足，它们恢复的图像仍然存在不满意的边界伪影。我们的工作通过提出一个统一的扩散框架来解决这些问题，该框架集成了图像和退化先验知识，以实现高效的阴影去除。具体而言，我们首先提出了一个阴影退化模型，该模型启发我们建立了一个新的展开扩散模型，称为ShandowDiffulation。它通过在退化先验和扩散生成先验的情况下逐步细化期望输出，显著地提高了模型的阴影去除能力，本质上可以作为图像恢复的新的强基线。此外，阴影扩散作为扩散生成器的辅助任务，逐步细化估计的阴影掩模，从而产生更准确和更稳健的无阴影图像生成。我们在ISTD、ISTD+和SRD三个流行的公共数据集上进行了大量的实验，以验证我们的方法的有效性。与现有的方法相比，我们的模型在峰值信噪比方面有了显著的提高，在SRD数据集上从31.69dB提高到34.73dB。

2 整体思想

去阴影任务首先是想获得准确的mask，本文认为估计的阴影不够精细，因为阴影与光照有关，所以本文在使用扩散模型的迭代思想逐步细化阴影。首先训练去噪器U-Net，使其可以估计当前步的噪声图和mask。然后使用HQS思想对当前步生成的x_{t}进行更新，获得无阴影图和mask，不断迭代直到收敛。

3 方法

假设阴影图像$y$是由无阴影图像$x$和光照$a$形成$y=ax$，这里严格假设亮区域的$a$值为1，阴影区域的$a$值为常数$a\in (0,1)$。通常来说，上式过于严格，因为现实光照是不平均的。一般来说，照明退化应该在空间上是可变的，并且高度依赖于阴影信息。因此，本文提出了一种新的阴影退化模型：
$$y=hx=wmx+(1-m)x\text{\hspace{1em}(1)}$$
其中，$h$是逐像素照明退化模型，其可以被分解为阴影mask $m$ 和照明权重$w$。$m$中阴影区域为1，其余为0。那么式1中$(1-m)x$可以理解为对五阴影区域的抠图，前项理解为带权重的阴影区域，则$y$为空间可变的阴影图像。

3.1 动态掩码感知扩散模型

阴影mask对去除阴影至关重要。本文将阴影去除重塑为一个联合任务，以追求无阴影图像和精化掩模，其中掩模细化将是扩散生成器的辅助任务，在生成无阴影图像的同时逐步精化阴影掩模。考虑条件扩散模型：
$$p_{\theta }(x_{0:T}|y,m_{0:T})=p(x_T)\prod _T^{t=1}{p}_{\theta }(x_{t-1}|x_t,y,m_t)\text{\hspace{1em}(2)}$$
这里可以理解为，以$y,m_t$作为条件，$m_{0:T}$是通过扩散模型逐步细化。去噪器${ϵ}_{\theta }$以阴影图像$y$，中间变量$x_t$，初始mask $\tilde{m}$和时间步$t$作为输入来预测噪声图$e_t$和细化mask $m_t$：
$$e_t,m_t={ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,y,\tilde{m},t)\text{\hspace{1em}(3)}$$
当然，去噪器的损失函数需要加上预测$m_t$的，这里使用的方法是在去噪网络的最后一层之后增加一个mask预测头，用一个1×1卷积层和一个Sigmoid函数来预测改进后的mask。损失函数为：
$$\mathcal{L}={\mathbb{E}}_{x_0,t,ϵ}||e_t-ϵ|{|}_F^2+{\mathbb{E}}_{t\sim [1,T]}||m_t-m_{gt}|{|}_F^2\text{\hspace{1em}(4)}$$
真实mask可以通过阴影图像和无阴影图像的差获得。去噪器的训练过程如图：

3.2 Unrolling-Inspired 扩散采样

上一小节我们重新训练了去噪器，可以得到一个额外的输出。本节介绍如何采样获得清晰图片和mask。我们利用下面的图形-mask联合能量函数，在优化无阴影图像的同时迭代细化mask：
$$\underset{x,m,z,v}{\min }\frac{1}{2}||hz-y|{|}_F^2+\psi \mathcal{R}([x|m])+\frac{\varphi }{2}||v-\tilde{m}|{|}_F^{2}s.t.x=\mathrm{z},\mathrm{m}=\mathrm{v}\text{\hspace{1em}(5)}$$
其中，$z$和$v$是将(5)转换为约束问题的辅助变量，而$\mathcal{R}(\cdot )$是捕获假定的联合图像和掩模先验的正则化。[·|·]表示串联操作。ψ和φ是权衡参数。这里，我们假设退化矩阵h可以由预先训练的退化估计网络估计为$h=\mathcal{D}(y,\tilde{m})$。为了处理等式约束，引入了两个二次惩罚项，并将问题重写为：
$$\begin{array}{c}\underset{x,m,z,v}{\min }\frac{1}{2}||hz-y|{|}_F^2+\psi \mathcal{R}([x|m])+\frac{\varphi }{2}||v-\tilde{m}|{|}_F^2+\frac{{\rho }_1}{2}||x-z|{|}_F^2+\frac{{\rho }_2}{2}||m-v|{|}_F^{2}s.t.x=\mathrm{z},\mathrm{m}=\mathrm{v}\end{array}\text{\hspace{1em}(6)}$$
其中ρ1和ρ2是惩罚参数(对于更简单的解，本文设置了ρ=ρ1=ρ2)。通过采用诸如半二次分裂(HQS)的变量分裂算法，可以通过迭代地解决如下三个子问题来解决优化问题：
$$\begin{array}{rl}[x|m{]}_{t-1}& =\arg \underset{x,m}{\min }\psi \mathcal{R}([x|m])+\frac{\rho }{2}||[x|m]-[z_t|v_t]|{|}_F^2\\ z_{t-1}& =\arg \underset{z}{\min }\frac{1}{2}||hz-y|{|}_F^2+\frac{\rho }{2}||x_{t-1}-z|{|}_F^2\\ v_{t-1}& =\arg \underset{v}{\min }\frac{\varphi }{2}||v-\tilde{m}|{|}_F^2+\frac{\rho }{2}||m_{t-1}-v|{|}_F^2\end{array}\text{\hspace{1em}(7)(8)(9)}$$
这里(8)和(9)是带有二次罚项的最小二乘问题，它们有闭合形式的解：
$$\begin{array}{rl}{z}_{t-1}& =(hy+\rho x_{t-1})/(hh+{\rho }_1)\\ v_{t-1}& =(\varphi \tilde{m}+\rho m_{t-1})/(\varphi +\rho )\end{array}\text{\hspace{1em}(10)(11)}$$
这里我们获得了$z$和$v$的迭代公式，$h$同样是用神经网络估计得出，算法流程如下：

总结

可以看到该方法的性能非常好，指标遥遥领先。我认为缺点就是比较复杂，需要一个预训练模型估计退化，然后是采样过程。相对于端到端的结构可能慢很多。但是，作为生成模型，per-pixel指标很高，这很厉害。相对于guide-diffusion，这篇文章的方法确实很新颖。