图像增强领域大突破！以1.66ms的速度处理4K图像，港理工提出图像自适应的3DLUT

作者丨Happy
编辑丨极市平台

极市导读

该文是香港理工大学张磊老师及其学生在图像增强领域的又一颠覆性成果。它将深度学习技术与传统3DLUT图像增强技术结合，得到了一种更灵活、更高效的图像增强技术。所提方法能够以1.66ms的速度对4K分辨率图像进行增强(硬件平台：Titan RTX GPU)。

paper: https://www4.comp.polyu.edu.hk/~cslzhang/paper/PAMI_LUT.pdf
code: https://github.com/HuiZeng/Image-Adaptive-3DLUT（即将开源）

Abstract

最近几年基于学习的方法已成为图像增强(增强图像的色调、对比度等)的主流方案，然而这些图像增强方法或者产生不好的效果或者需要大量的计算量以及内存占用，严重影响了这些方法在高分辨率图像方面的实际应用(比如12M像素)。
该文提出了一种图像自适应的3D查找表(3D LUT)方法用于图像增强，在具有超快速度的同时具有鲁棒的增强效果。3DLUT是传统图像编辑领域常用方案，但这些查找表往往需要人工调节并固化。该文首次提出采用成对/不成对学习的方式从标注数据集中学习3DLUT。更重要的是，所学习到的3DLUT具有图像自适应特性，可以灵活的进行图像增强。所提方法采用端到端的方式同时学习3DLUT以及一个小网络，小网络用于根据输入图像学习查找表的融合权值，将融合后的3DLUT作用于输入图像即可得到期望的输出。所提方法仅仅包含不到600K参数，能够以不超过2ms的速度处理4K分辨率图像(硬件平台：Titan RTX GPU)。在具有超快的推理速度同时，所提方法同样以极大的优势(PSNR，SSIM以及颜色差异度量指标)超越其他图像增强方法。
该文的贡献主要包含以下三点：
  • 首个采用深度学习方法学习3DLUT并用于自动图像增强的方案，更重要的是，所提方法学习到的3DLUT具有图像自适应性，可以更灵活的进行图像增强；
  • 所提方法仅有不超过600K参数量，且能够以不超过2ms的速度处理4K分辨率图像(GPU)；
  • 在两个公开数据集上以极大优势超越其他SOTA图像增强方案。

Method

上图给出了本文所提出方案的框架图，它包含多个基准3DLUT与一个小的CNN网络。基准3DLUT可以确保输入到输出的颜色变换可以被更好的包含，而CNN则用于学习3DLUT的融合权值并将多个3DLUT融合为一个新的3DLUT，最后所得3DLUT用于图像变换。

3DLUT and trilinear interpolation

3DLUT是一种经典且有效并被广泛应用的图像增强技术，下图给出了3DLUT的示意图。3DLUT定义了一个包含个元素的晶格（注：M表示每个颜色通道的采样数，实际应用中一般为33）。每个定义了输入RGB颜色及对应的输出RGB颜色,而变换的精度则由M控制(默认33)。
给定M以及索引 RGB颜色可以通过均匀离散RGB颜色空间得到。不同的3DLUT具有不同的输出RGB颜色,而这在所提方法中是可学习的参数。当M=33时，每个3DLUT包含108K参数量。
基于3DLUT的颜色变换可以通过两个操作达成：lookup与trilinear interpolation。给定输入RGB颜色，lookup操作可以通过如下公式进行：
其中,而表示颜色值。在完成输入RGB颜色定位后，它的8个最近邻元素将被用于trilinear interpolation并得到输出RGB。trilinear interpolation可以描述如下：假设，,那么三次插值公式如下：
其中。上述三次插值过程是可微的，可以轻易的推导出的梯度。由于每个输入的三次插值结果完全独立，故而可以采用GPU进行并行加速。

Learning image-adaptive 3D LUTs

传统的基于3DLUT的图像增强方法具有两个主要局限：(1)3DLUT大多人工设计，需要设计人员具有丰富经验，同时消耗大量时间；(2)一个3DLUT仅能处理一种固定的变换，难以适应全场景的变换。尽管现有相机、图像编辑工具中提供了多个LUTs集合供用户手动选择以大臣个不同的增强效果，但是这种交互方式不够灵活。
为解决上述第一个局限，作者提出采用数据驱动方法生成3DLUT；为解决上述第二个局限，模型应当具有内容自适应特性，因此作者提出学习图像自适应的3DLUT。
一个很自然的想法是通过场景分类的方式选择不同的3DLUT，这也是诸多相机、图像编辑工具常用方案。然而这种硬决策方案具有这样几个缺陷：(1) 它难以将多个可能场景正确的划分到预定义类别中，因为不同类别具有不同的3DLUT；(2)需要大量的3DLUT以尽可能的覆盖全场景；(3) 分类器的训练独立于3DLUT，导致组合方案并非最优。如果场景分类错误，那么所选择的3DLUT大概率生成非常差的结果。
在该文中，作者采用一种更高效的方式学习图像自适应的3DLUT：同时学习N个3DLUT与一个小CNN模型(用于预测3DLUT的融合权值)。对于输入图像，其最终的增强输出可以描述如下：
而表示通过CNN模型预测的内容自适应权值。所提方法仅仅需要几个3DLUT(N=3)进行图像变换，同时采用软加权策略以获得图像内容自适应的变换效果。更进一步，考虑到像素变换的独立性，上述操作可以进一步简化如下：
也就是时候，可以先进行3DLUT融合，再进行图像变换。这种处理方式进一步提升的所提方法的推理速度。下表给出了所提方法中CNN模型的参数配置，它包含5个卷积(后接LeakyReLU、IN)、一个dropout以及一个全连接层。整个CNN模型包含269参数量(N=3)。

Learning Criteria

接下来就要介绍了优化的目标函数了，目标函数定义如下：
其中表示CNN模型与3DLUT，表示某些损失函数与正则项。
如前所述，所提方法可以采用成对数据进行训练，也可以采用非成对数据进行训练。我们先来看一下成对训练的损失函数，这个就比较简单的常规的损失函数即可，比如L1、MSE等。在这里作者选用了简单的MSE损失函数，定义如下
接下来，我们再看一下非成对数据的损失函数。这里借鉴的是GAN的思想进行非成对数据的训练，前述的3DLUT以及CNN可以视作生成器，只需要再引入一个判别器即可。此时生成器通过如下损失进行训练：
其中。判别器的损失函数定义如下：
其中，,前两项为标准的判别损失，最后一项为梯度惩罚项用于稳定训练。那么总的损失函数则定义为：

Regularization

我们可以采用前述损失以及SGD/Adam优化器进行训练3DLUT以及CNN。然而得到的3DLUT具有非平滑现象，见下图a。这种局部陡峭现象会导致变换后图像的伪影问题，为使得所学习的3DLUT具有更好的稳定性和鲁棒性，作者引入了两种正则项。
  • Smooth Regularization。为了稳定的将输入RGB变换到期望的颜色空间，而不引入伪影，3DLUT的输出RGB应当具有局部平滑性。全变差Total Variation, TV是图像复原领域经典的平滑正则化技术，作者将其引入到3DLUT，此时TV正则定义如下：
除了3DLUT应具有平滑性外，CNN预测的内容自适应权值同样非常重要，作者也为其引入了正则以提升所得3DLUT的平滑性。总的平滑正则项描述如下：
  • Monotonicity regularization。除了平滑性外，单调性也适合3DLUT应当具有的一个属性。为此作者设计了如下的单调正则项：
其中表示标准ReLU操作。

Final training losses and implementation

通过集成前述两个正则项，最终的损失函数定义如下：
其中。
在实现方面，基本框架为Pytorch，trilinear interpolation通过CUDA并行实现。训练的优化器为Adam，Batch=1，学习率固定0.0001(pair)和0.0002(unpair)，随机裁剪、随机镜像等数据增强。而这N个3DLUT的初始化方式为：第一个为恒等映射，其他则初始化为0。CNN的全连接层的bias初始化为1，确保起始阶段的权值预测为1。

Experiments

训练数据：MIT-Adobe FiveK， HDR+。评价准则：PSNR、SSIM。下面直接看结果咯。
下图给出了三个不同图像采用所提方法生成的3DLUT以及最终的增强结果。可以看到：所提方法生成的结果比固定3DLUT的方式更灵活，效果更佳。
下图给出了所提方法与其他SOTA方法在FiveK数据集上的指标对比以及视觉效果对。可以看到所提方法不仅具有更好的指标，同样具有更好的视觉效果。
下面给出了所提方法与其他SOTA方法在HDR+数据集上的指标与视觉效果对比。
与此同时，作者也给出了不同图像编辑工具与所提方法的处理效果对比，见下图。
最后，作者给出不同的图像增强方法的推理速度对比，见下图。还有什么比这个速度更惊人的吗？以1.66ms的速度处理4K图像，只能感慨一下，快到没反应。
全文到此结束，对该感兴趣的同学建议去看一下原文。最后期待作者尽快开源代码。
◎作者档案
Happy，一个爱“胡思乱想”的AI行者
个人公众号：AIWalker
欢迎大家联系极市小编（微信ID:fengcall19）加入极市原创作者行列
赞 (0)

相关推荐