全国人工智能大赛 AI+4K HDR赛项冠军团队方案分享

冠军团队由三名博士生组成，目前此方案在投论文中，之后方案代码会开源。\

团队名称：没学可上

团队成员：

曾辉博士研究生（第4年）

杨熹博士研究生（第1年）

梁杰博士研究生（第1年）

赛题任务

视频四倍超分+HDR，要求将低质量540p SDR视频重建为高质量、已调色的4K HDR视频。

评测方法：最终得分 = 25 * PSNR 项 + 25 * SSIM 项 + 50 * VMAF 项

关于赛题的详细信息可以查看和鲸科技赛事页面：

www.kesci.com/home/compet…

赛题分析

此次4K HDR比赛包含众多子任务，初步按照高频和低频两个方向划分为：

高频方面，需要去噪和超分

低频方面，需要做亮度矫正、色彩增强和SDR转HDR (即inverse tone mapping)

给定时间和硬件环境下同时解决上述多个任务，需要重点考虑两个方面的矛盾

矛盾一：高频和低频

亮度矫正和色彩增强等低频任务需要足够多的图像内容才能进行准确判断，因此训练和测试都依赖于较大的patchsize和感受野。在显存有限的情况下，大patch限制了模型的复杂度。

而去噪和超分等高频任务pattern更不固定，更难学，需要较深的模型和比较长的训练时间才能保证性能。

矛盾二：去噪和超分

去噪和超分关注图像高频细节，其中去噪不可避免地会损失高频信息，而超分又需要恢复高频，两个任务本身存在矛盾之处。

基于以上对赛题的分析，需要综合考虑亮度、色彩、去噪、超分等几个方面，在相互矛盾的子任务之间找到最合适的平衡点，设计最优的解决方案。同时由于本次任务所给数据退化较严重，需要仔细分析退化模式并相应地设计出有效的数据处理方案。

整体方案

数据预处理

这一部分的目的是修正数据分布，简化学习任务。处理步骤基于对数据集的观察及实验，人工得出的经验。

具体操作：

计算YUV空间亮度均值。这一步可以筛出严重过曝的视频（直接排除），以及视频的黑边（直接输出结果0(SDR)->4099(HDR)
计算HSV空间直方图距离。这一步可以检测出场景的转换，提高多帧融合稳定性。
中值滤波去除椒盐噪声点，如果滤波后与滤波前像素值大于某个阈值，则用非椒盐噪声点替换原始值，反之不处理。
Non-local mean滤波去除部分高斯噪声

模型学习

特征提取

1.采用大patch（384x768）输入，确保低频任务获得足够大的感受野；

2.采用一层instance normalization对亮度做归一化；

3.提前下采样，降低显存消耗，节省计算资源。

特征对齐与融合

1. 采用EDVR中的PCD模块多帧对齐；

2. HardTanh约束offset越界,(-32,+32)；

3. TV约束offset平滑,利用图像空域的连续性先验。

低频任务学习

方案一

1.UNet结构进一步增大模型感受野，促进低频内容学习;

2.采用pixel shuffle的方式上采样，减少上采样产生的halo。

方案二

学习多个3D LUTs进行色彩增强

此模型优点为：仅0.6M参数；SOTA性能；4K分辨率上200+FPS；有明确物理意义。

高频细节恢复

1.采用RCAN中的deep residual channel attention模块重构细节；

2.不同尺度采用不同的channel维度，兼顾性能与效率。

训练&预测

将前面几个模块串联起来之后，计算多尺度损失，这样可以提高训练效率。模型训练好之后，输出中间的预测结果。

结果后处理

这一部分会对前一步的预测结果进行进一步处理，改善遗留问题，并且提升分数视觉效果。

具体操作：

亮度修正，帧间平滑
细节增强
黑边补偿
模型融合

理论分析及实验发现，本任务中存在一些相互矛盾的子任务，需要合理设计两者之间的trade-off以达到最佳的结果。比如，单个模型对不同场景亮度的恢复有时候偏亮有时候偏暗，模型的去噪能力与超分后的细节保留能力之间存在着明显的冲突等等。基于以上分析，我们设计了一个双模型融合的框架，其中，第一个模型Model_1采用未经处理的原始数据训练，采用512x512的大patch输入，重点恢复图像低频内容，第二个模型Model_2使用预处理过的数据，训练采用384x384的patch训练恢复尽可能多的细节。实验发现，两个模型的预测结果在去噪、细节保留、亮度估计、色彩估计等方面均有一定的互补性，融合之后能够取得最佳效果。