056、深度学习去马赛克：基于CNN的RAW域重建与边缘保持技术

056、深度学习去马赛克基于CNN的RAW域重建与边缘保持技术去年在调试某款旗舰手机的主摄时遇到一个让人头疼的问题——暗光场景下画面边缘的彩色条纹像彩虹一样铺开尤其是拍树叶和金属栏杆时简直没法看。当时团队里有人怀疑是镜头色散有人说是ISP的demosaic算法参数没调好。我拿着RAW图一层层扒下去发现罪魁祸首就是传统的去马赛克算法在边缘处“猜”错了颜色。这事让我下定决心必须把深度学习去马赛克这块啃透。为什么传统算法搞不定边缘先说说传统去马赛克为什么在边缘处翻车。Bayer阵列每个像素只记录一种颜色缺失的两种颜色要靠周围像素插值。最基础的最近邻插值边缘处直接糊成一团。双线性插值稍微好点但遇到高频纹理——比如头发丝、金属拉丝——就会产生伪彩色也就是我们常说的“拉链效应”或“假色”。后来有了方向性插值先判断边缘方向再沿着边缘插值。这个方法在理想边缘上表现不错但真实场景哪有那么多“理想边缘”45度、135度、不规则纹理算法一多就容易误判。更致命的是传统方法都是基于局部像素的线性操作无法理解图像的高层语义。比如一块红色布料上的黑色条纹传统算法会认为红色区域里突然出现黑色是“异常”拼命往红色方向拉结果边缘就变成了紫红色。CNN凭什么能重建RAW域深度学习去马赛克的核心思路是把Bayer阵列的RAW图当成一个“残缺”的输入让CNN学习从残缺到完整的映射。这里有个关键点——必须在RAW域操作而不是在RGB域。为什么因为RAW域的数据是线性的没有经过gamma校正和白平衡保留了最原始的物理信息。一旦转到RGB域非线性变换会把颜色关系搞复杂CNN学起来事倍功半。我最早尝试的方案很简单把Bayer图拆成4个通道R、G1、G2、B每个通道只有1/4的像素有值其余填0然后扔进一个U-Net。结果训练出来的模型在测试集上PSNR很高但实际拍出来的照片边缘总有轻微的“水波纹”。后来才意识到这种简单的“填0”方式破坏了Bayer阵列的相位信息——G1和G2虽然都是绿色但它们的位置不同对边缘的响应也不同。正确的做法是保持Bayer的原始排列用卷积核去“感受”这种相位差异。边缘保持的关键多尺度与注意力在调试过程中我发现一个规律CNN对低频区域的重建效果很好但高频边缘总是差一口气。原因在于卷积核的感受野是有限的大感受野能理解全局结构但丢失细节小感受野能保留细节但容易受噪声干扰。我的解决方案是引入多尺度特征融合。具体来说在网络的浅层用3x3卷积捕捉局部纹理深层用空洞卷积扩大感受野最后把不同尺度的特征图通过跳跃连接合并。这里踩过坑——直接拼接不同尺度的特征图会导致维度爆炸参数量翻倍但效果提升有限。正确的做法是用1x1卷积先降维再拼接。另一个关键点是注意力机制。传统的去马赛克算法在平坦区域和纹理区域应该采用不同的策略。平坦区域需要平滑纹理区域需要锐利。我尝试在网络的中间层加入空间注意力模块让网络自己学会“哪里该平滑哪里该保持”。别这样写——不要用那种复杂的self-attention计算量太大手机端跑不动。用简单的通道注意力加空间注意力就够了参数量增加不到10%但边缘的PSNR能提升0.3-0.5dB。训练数据的“坑”与“道”训练数据是另一个大坑。很多人直接用公开数据集比如Kodak、McMaster但这些数据集的分辨率太低只有几百万像素而且场景单一。实际手机拍摄的RAW图有1200万甚至4800万像素场景复杂度远超这些数据集。我的做法是自建数据集。用实验室的相机拍摄各种场景的RAW图包括室内、室外、夜景、逆光、微距等。拍摄时注意保持相机稳定避免运动模糊。然后对每张RAW图做两次采样——一次是原始分辨率一次是降采样到1/4分辨率再通过双线性插值放大回原始分辨率。这样做的目的是让网络学会区分“真实细节”和“插值伪影”。训练时还有一个容易忽略的点——损失函数。L1损失和L2损失在平坦区域表现不错但在边缘处容易产生模糊。我试过加入感知损失效果有提升但训练不稳定。最后用的是L1损失加边缘梯度损失让网络在边缘处保持锐利。梯度损失的计算很简单对重建图和真值图分别求Sobel梯度然后计算梯度图的L1距离。别这样写——不要用那种复杂的拉普拉斯金字塔损失调参太麻烦效果也不一定好。实际部署的“血泪史”模型在PC上跑得飞起一上手机就翻车。首先是计算量问题一个U-Net在手机上跑一次要几百毫秒根本达不到实时预览的要求。解决方案是模型剪枝和量化。剪枝时注意不要剪掉浅层的卷积核那些是保留边缘细节的关键。量化时用INT8精度损失在0.1dB以内但速度提升4倍。另一个问题是内存带宽。手机的内存带宽有限频繁读写大尺寸的特征图会导致卡顿。我的做法是把网络分成几个阶段每个阶段处理完后立即输出中间结果避免一次性加载所有特征图。这里踩过坑——不要用那种“先下采样再上采样”的对称结构下采样后的特征图尺寸太小上采样时信息丢失严重。改用“下采样-处理-上采样”的非对称结构下采样只做一次上采样用双线性插值加卷积效果更好。经验性建议别迷信公开数据集。Kodak和McMaster都是实验室环境拍的噪声低、光照均匀。实际场景的噪声、暗角、镜头畸变都会影响去马赛克效果。最好用自己拍摄的RAW图做微调哪怕只有几十张效果也比纯公开数据集好。边缘保持是“系统工程”。不要指望一个模块解决所有问题。从Bayer排列的相位保持到多尺度特征融合再到损失函数的设计每个环节都要为边缘服务。我见过有人花大量时间调注意力模块但Bayer输入的处理方式却是错的结果事倍功半。硬件协同优化。如果芯片支持硬件加速比如高通平台的HVX或联发科的APU尽量把卷积操作映射到硬件上。但要注意硬件加速对卷积核大小和步长有限制设计网络时要提前了解硬件规格。调试时多看RAW图。不要只看RGB输出。把重建的RAW图转成RGB后用差值图重建图减真值图观察误差分布。如果误差集中在边缘说明网络对高频信息的学习不够如果误差集中在平坦区域说明噪声抑制过强。最后一条也是最重要的一条——不要为了追求PSNR而牺牲视觉质量。PSNR高0.1dB但边缘出现伪影用户一眼就能看出来。我宁愿PSNR低0.2dB也要保证边缘干净。毕竟用户看的是照片不是指标。