HDRUNet | 深圳先进院董超团队提出带降噪与反量化功能的单帧HDR重建算法

AIWalker

发布于 2023-09-09 10:22:15

6520

发布于 2023-09-09 10:22:15

文章被收录于专栏：AIWalker

标题&作者团队

Paper: https://arxiv.org/pdf/2105.13084.pdf

本文是深圳先进技术研究院董超团队在HDR领域的最新力作，取得了NTIRE2021 单帧HDR竞赛第二成绩，指标仅低0.07dB，但推理速度快116倍。针对HDR重建伴随的高光区域细节重建、低光区域噪声抑制以及正常曝光区域的量化损失、噪声抑制问题，HDRUNet提出了(1) 条件网络提供空域可变调制；(2) 加权网络提供正常曝光区域信息保留；(3) Tanh_L1用平衡不同曝光区域的贡献影响。

Abstract

受限于传感器，现有消费级数码相机在真实场景中仅能捕获有限范围的亮度。此外，成像过程中还会引入噪声与量化误差。为获得具有优异视觉质量的高动态范围图像，现有方案往往采用多曝光图像合成方式。然而，相同场景的多曝光图像的获取难度极大，而且现有HDR重建方法往往忽略了噪声与量化损失。

为解决上述问题，我们提出了一种新的基于学习的方法：HDRUNet，它采用空域动态编解码网络学习端到端的映射，同时考虑的降噪与反量化。该网络包含一个UNet风格的基础网络以充分利用多尺度分层信息、一个条件网络用于执行模式相关的调制以及一个加权网络用于自适应选择保留信息。

此外，我们还提出了Tanh_L1损失用于平衡过曝与正常曝光的影响。所提方法在量化性能与视觉质量方法取得了SOTA性能，所提HDRUNet取得了NTIRE2021单帧高动态范围竞赛第二名的好成绩。

Method

Observations

图像HDR重建往往会伴随降噪与反量化(dequantization)，为说明该问题，我们对LDR-HDR图像对的梯度图进行可视化，见下图。

从上图可以看到，相比HDR图像，

由于动态范围压缩与量化，LDR图像高光区域的梯度几乎不可见；
在正常曝光区域，LDR与HDR图像的噪声梯度均清晰可见，这意味着两种类型图像中均存在噪声。尽管如此，由于噪声水平不同导致两种类型图像中的噪声模式存在显著差异。

此外，这些图像中的噪声并非均匀分布，即与常见高斯白噪声不同。因此，噪声模式差异不仅仅存在于高亮与暗部区域，同时还在于正常曝光区域。这就启发我们设计一种空域可变调制模块。

Network Structure

基于前述发现与分析，我们设计了一种带空域调制的UNet风格的网络用于单帧HDR重建，所设计网络结构见下图，它包含三个主要成分，接下来，我们将针对这三个成分进行更详细介绍。

Base Network BaseNet采用了UNet风格结构，它以8bit带噪LDR图像作为输入，重建16bit的HDR图像，所预测的HDR图像应当在过曝与欠曝区域包含更多的细节，同时具有更少的噪声。

Condition Network HDR重建的关键：对输入LDR图像的过曝与欠曝区域进行遗失细节重建。不同的区域具有不同的曝光核亮度，不同的凸显更具有不同的整体亮度和对比度信息。因此，很有必要采用位置相关、图像相关的方式进行图像处理。此外，还要求网络能够很好的处理非均匀分布噪声。然而，传统CNN具有空域不变形，即所有图像、所有位置均采用相同的滤波器权值处理。受启发于CSRNet、SFTGAN，我们提出了带SFT(Spatial Feature Transform)的条件网络用于提供空域可变调制。具体来说，条件网络以LDR图像作为输入，预测对应的额条件映射图并用于对BaseNet的中间特征进行调制。条件网络与SFT的机制见上面的Figure，公式定义如下:

SFT(x) = \alpha \odot x + \beta

通过利用该调制策略，所提方法可以达到位置相关、图像相关的调制处理。

Weighting Network HDR重建最大的挑战在于：过曝与欠曝区域的细节重建，而大部分的正常曝光内容对于学习阶段的贡献非常小。为此，我们提出了一种甲醛估计网络为需要保留的正常曝光区域预测一个软加权图W。于是，整个网络将主要主要聚焦于过曝区域细节重建：

\hat{Y} = W \odot I +　\mathcal{G}(I)

注：I表示LDR输入，

\hat{Y}

表示最终重建的HDR图像，

\mathcal{G}(I)

表示BaseNet的输出。

Loss Function

在真实场景图像HDR重建中，我们不仅要考虑动态范围的重建，同时还需要考虑噪声与量化伪影的降低。然而，现有损失函数(比如

L_1, L_2

)并不适用于同时处理上述问题。直接作用于HDR数值上的损失函数会聚焦于高亮区域而低估低量区域，导致比较差的量化性能与视觉质量。

为解决上述问题，我们提出了一种定义损失函数：

Tanh\_L_1

，定义如下：

Tanh\_L_1(\hat{Y}, H) = |Tanh(\hat{Y}) - Tanh(H)|

Experiments

Experimental Setup

Dataset 训练数据采用NTIRE2021 HDR竞赛的数据，它包含1494LDR/HDR对用于训练，60张LDR用于验证，201张LDR用于测试。注：LDR/HDR图像对在时间轴、曝光等级方面进行了对齐并进行伽马校正后保存。

Metrics 度量准则选择了PSNR-L与PSNR-

\mu

，前者更倾向于高亮值，而后置则更倾向于视觉相似性。因此，主要度量准则为后者。

Details 在后续实验中，残差模块数为8，stride=2的卷积用于下采样，pixelshuffle用上采样。在训练开始前，我们对数据预处理：裁剪为

480 \times 480

的图像块，间隔步长为240。训练过程中，batch=16，Adam优化器+Kaiming初始化，初始学习率为

2\times 10^{-4}

，每

2\times 10^5

次迭代折半。当训练patch为

256\times 256

时，总计训练花费5天。

Ablation Study

我们先来看一下不同配置的影响，包含训练块尺寸、损失函数以及调制策略。

Patch Size 实际上，训练块尺寸对于该任务影响非常大。由于HDR重建不仅仅是简单的局部处理，它还包含更重的全局与整体调制，不同的区域需要不同对待。此外，由于过曝区域严重的信息损失，我们认为这些区域的细节重建需要更大的感受野。上表给出了不同尺寸的影响，可以看到：随着块尺寸的提升，模型的量化性能逐步提升。考虑到性能与计算消耗，我们选择

256\times 256

作为建议块尺寸。

Loss Function 在前面的章节中，我们提出了用于带降噪与反量化的HDR重建的Tanh_L1损失。为加速训练，我们固定块尺寸为

160 \times 160

。上表对比了不同损失的性能差异，从中可以看到：

相比L2损失，L1损失取得了更佳的量化性能，即更高的PSNR-L与PSNR-u指标；
通过引入Tanh操作，PSNR-u指标进一步提升，而PSNR-L出现了一定下降。具体来说，采用Tanh_L1损失可以提升PSNR-u指标高达0.5dB。这是因为：当直接使用L1与L2损失时，高亮区域的损失具有更大的权重，因而具有更高的PSNR-L指标；正如前面所提到的PSNR-u更能反应视觉相似性，且该竞赛也以此作为主要参考指标，因此我们采用Tanh_L1作为损失函数。