ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

基于DNN的灰度图像上色实战:从Lab色彩空间到PyTorch模型部署

基于DNN的灰度图像上色实战:从Lab色彩空间到PyTorch模型部署 简介本资源面向图像处理与深度学习方向的开发者提供一套基于OpenCV DNN模块实现灰度图像自动上色的完整C工程适用于Windows 10/11平台、Visual Studio 2019开发环境。工程通过加载caffemodel与prototxt模型文件对灰度图进行色彩预测与还原并支持调整参数观察不同上色效果适合想入门DNN推理部署或研究图像着色算法的读者参考。压缩包共40个文件约190.34MB包含cpp源码、vcxproj工程文件、sln解决方案、caffemodel与prototxt模型、dll运行库、exe可执行程序以及多张不同参数下的结果对比图便于直接编译运行与效果验证。目前已有570人学习下载。读者可从中获取完整的工程结构、模型调用流程与参数调节思路快速理解DNN图像上色在C项目中的落地方式并在此基础上进行二次开发或迁移到其他图像处理任务。1. 灰度图上色这件事DNN 到底在学什么手里有一批老照片、医学影像或者卫星灰度图想还原出合理的色彩这是很多人搜「基于深度学习的DNN模块-给灰度图像上色」时真正想解决的问题。它本质是一个逐像素回归/分类任务输入单通道亮度 L输出 a、b 两个色度通道拼回 Lab 空间再转 RGB。难点不在网络多深而在于同一张灰度图可能对应多种合理配色模型要学会的是「统计上最像真的那一种」。这套方案适合三类人想拿深度学习实战项目案例练手的入门者、需要批量处理历史素材的工程同学、以及想把上色模块嵌进自己图像流水线的开发者。它不需要顶配显卡一张 6GB 显存的卡就能跑通训练这也是它比很多深度学习模型更适合作为第一个端到端项目的原因。下面从数据、网络、训练到部署把能复现的路径讲清楚。2. 数据与色彩空间上色任务的地基怎么打2.1 为什么选 Lab 而不是直接回归 RGB直接让网络输出 RGB 三通道是最直觉的做法但翻车率极高。RGB 三个通道高度相关亮度一变三个值全变网络要同时学亮度和色度收敛慢且容易发灰。Lab 空间把亮度 L 和色度 a、b 解耦我们只让网络预测 a、bL 直接沿用输入这样任务从「重建整张图」缩小成「预测两个色度通道」难度骤降。具体流程是RGB 图先转 Lab取 L 作为输入a、b 作为监督标签。推理时把预测的 a、b 和原图 L 拼回去再转回 RGB。这个解耦是整套方案能 work 的核心也是很多深度学习知识点里讲特征解耦最直观的例子。import numpy as np from skimage import color def rgb_to_lab_input(rgb_img): # rgb_img: float32, 范围 [0,1], 形状 (H,W,3) lab color.rgb2lab(rgb_img) L lab[:, :, 0:1] # 亮度通道作为网络输入 ab lab[:, :, 1:3] # 色度通道作为监督标签 L_norm L / 50.0 - 1.0 # L 原本 0~100归一化到 -1~1 ab_norm ab / 128.0 # ab 原本约 -128~127归一化到 -1~1 return L_norm.astype(np.float32), ab_norm.astype(np.float32)逻辑说明rgb2lab输出的 L 范围是 0~100a、b 大约在 -128~127。归一化是为了让输入输出都落在 -1~1配合 tanh 激活。参数上L 除以 50 再减 1 是常见做法ab 除以 128 即可不必严格裁剪越界值极少。2.2 数据集怎么凑别一上来就找「上色专用数据集」上色是自监督任务——任何一张彩色图都能拆成「灰度输入 彩色标签」所以 ImageNet、COCO、甚至你自己拍的几千张照片都能用。我一般会按场景分层采样人像、风景、室内各占一定比例否则模型会对某类场景严重偏色。数据量上1 万到 5 万张足够训出一个能看的版本再多收益递减。预处理阶段有两个必做动作统一缩放到 256×256训练时再随机裁到 224以及做数据增强。增强只做几何变换翻转、轻微旋转不要做颜色抖动因为颜色就是标签抖了等于改答案。import torch from torch.utils.data import Dataset from PIL import Image import torchvision.transforms as T class ColorizationDataset(Dataset): def __init__(self, paths, size256): self.paths paths self.size size self.resize T.Resize((size, size), interpolationT.InterpolationMode.BICUBIC) def __len__(self): return len(self.paths) def __getitem__(self, idx): img Image.open(self.paths[idx]).convert(RGB) img self.resize(img) rgb np.asarray(img).astype(np.float32) / 255.0 L, ab rgb_to_lab_input(rgb) # 转成 CHW L torch.from_numpy(L).permute(2, 0, 1) ab torch.from_numpy(ab).permute(2, 0, 1) return L, ab逻辑说明convert(RGB)防止灰度图或带 alpha 的图混进来导致通道数不对。BICUBIC插值比默认的最近邻更平滑减少缩放伪影。返回的 L 和 ab 都是 CHW 格式符合 PyTorch 卷积输入约定。参数size训练用 256推理可以更大但网络若是全卷积就能吃任意尺寸。3. DNN 模块怎么搭从输入到 ab 输出的网络结构3.1 编码器-解码器是主线别迷信越深越好上色网络的经典结构是编码器-解码器编码器不断下采样提取语义解码器上采样恢复分辨率中间用跳跃连接把浅层细节传回去。编码器可以直接用预训练的 ResNet 前几层比从头训快得多这也是深度学习pytorch生态里最省事的做法。我一般会用一个轻量版本编码器 5 个下采样块通道从 64 翻到 512解码器对称上采样每层和对应编码层做 concat。最后一层输出 2 通道接 tanh。整个网络参数量控制在 20M 以内单卡训练没压力。import torch.nn as nn class ColorNet(nn.Module): def __init__(self): super().__init__() def block(cin, cout, stride1): return nn.Sequential( nn.Conv2d(cin, cout, 3, stride, 1), nn.BatchNorm2d(cout), nn.ReLU(inplaceTrue), nn.Conv2d(cout, cout, 3, 1, 1), nn.BatchNorm2d(cout), nn.ReLU(inplaceTrue), ) self.enc1 block(1, 64) self.enc2 block(64, 128, 2) self.enc3 block(128, 256, 2) self.enc4 block(256, 512, 2) self.up3 nn.ConvTranspose2d(512, 256, 2, 2) self.dec3 block(512, 256) self.up2 nn.ConvTranspose2d(256, 128, 2, 2) self.dec2 block(256, 128) self.up1 nn.ConvTranspose2d(128, 64, 2, 2) self.dec1 block(128, 64) self.out nn.Conv2d(64, 2, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(e1) e3 self.enc3(e2) e4 self.enc4(e3) d3 self.dec3(torch.cat([self.up3(e4), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return torch.tanh(self.out(d1))逻辑说明每个block是两个 3×3 卷积加 BN 和 ReLU这是最稳的卷积单元。ConvTranspose2d做 2 倍上采样torch.cat把上采样结果和编码器同层特征拼接恢复细节。最后 1×1 卷积压到 2 通道tanh 保证输出在 -1~1正好对应归一化后的 ab。参数上通道数 64/128/256/512 是经验值显存紧张可以把 512 降到 256。3.2 损失函数L1 打底分类思路做增强只用 L1 损失会让结果偏灰因为模型倾向于预测「平均色」来降低整体误差。常见改进是把 ab 空间量化成 313 个色块把回归问题转成分类问题用交叉熵训练推理时取概率最高的色块或做加权平均。这套思路在深度学习cnn的很多任务里都出现过效果比纯回归鲜艳不少。实操上我一般混合两种损失L1 保证结构稳定分类损失提升色彩饱和度权重各占一半。如果只想快速跑通纯 L1 也能出结果只是颜色偏保守。def total_loss(pred_ab, true_ab, ce_weight0.5): l1 nn.functional.l1_loss(pred_ab, true_ab) # 简化版把 ab 按 10 度量化成 36 个 bin 做分类 q ((true_ab 1) / 2 * 35).long().clamp(0, 35) logits pred_ab.permute(0, 2, 3, 1).reshape(-1, 2) target q.permute(0, 2, 3, 1).reshape(-1, 2) ce nn.functional.cross_entropy(logits, target[:, 0]) return l1 ce_weight * ce逻辑说明这里为演示把 ab 量化成 36 个 bin真实项目常用 313 个 bin 的预设色块表。cross_entropy需要 logits 形状为 (N, C)所以做了 reshape。参数ce_weight控制分类损失占比0.5 是常用起点颜色过饱和就调低偏灰就调高。4. 训练与推理参数怎么设、显存怎么省4.1 训练循环与必调参数训练上色的坑集中在学习率和 batch size。编码器若用预训练权重学习率要小1e-4 量级否则预训练特征会被快速破坏从头训可以用 1e-3。batch size 在 8~16 之间太小 BN 统计不稳太大显存吃紧。device cuda if torch.cuda.is_available() else cpu model ColorNet().to(device) opt torch.optim.Adam(model.parameters(), lr1e-4, betas(0.9, 0.99)) sched torch.optim.lr_scheduler.StepLR(opt, step_size10, gamma0.5) for epoch in range(50): model.train() for L, ab in train_loader: L, ab L.to(device), ab.to(device) pred model(L) loss total_loss(pred, ab) opt.zero_grad() loss.backward() opt.step() sched.step() print(fepoch {epoch} done)逻辑说明Adam 的betas第二项设 0.99 比默认 0.999 收敛更快这是上色任务里的小经验。StepLR每 10 轮降一半学习率防止后期震荡。50 轮是起步值看验证集色差不再下降就可以停。显存不够时把 batch 降到 4并开混合精度。scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): pred model(L) loss total_loss(pred, ab) scaler.scale(loss).backward() scaler.step(opt) scaler.update()逻辑说明混合精度能让显存占用降约 40%速度提升明显。autocast自动把合适算子转 fp16GradScaler防止梯度下溢。注意 loss 计算里的cross_entropy在 fp16 下可能不稳必要时对该部分强制 fp32。4.2 推理把 ab 拼回 RGB 的完整链路推理时输入灰度图网络输出 ab和原图 L 拼回 Lab 再转 RGB。这里有个容易忽略的点推理用的 L 必须是原始分辨率下的 L不能是训练时的归一化版本否则拼回去颜色会整体偏移。def colorize(model, gray_path, devicecuda): model.eval() img Image.open(gray_path).convert(RGB) rgb np.asarray(img).astype(np.float32) / 255.0 L, _ rgb_to_lab_input(rgb) L_t torch.from_numpy(L).permute(2, 0, 1).unsqueeze(0).to(device) with torch.no_grad(): ab model(L_t).squeeze(0).permute(1, 2, 0).cpu().numpy() lab np.concatenate([L * 50.0 50.0, ab * 128.0], axis2) out color.rgb2lab # 占位实际用 lab2rgb from skimage.color import lab2rgb rgb_out lab2rgb(lab) return (rgb_out * 255).clip(0, 255).astype(np.uint8)逻辑说明L * 50.0 50.0把归一化的 L 还原回 0~100ab * 128.0还原色度。lab2rgb输出 0~1 浮点乘 255 并裁剪得到最终图像。注意lab2rgb对越界值敏感裁剪不能省。5. 上色项目避坑清单五个真实翻车现场5.1 输出一片灰颜色像蒙了层雾现象训练 loss 降下去了但推理结果几乎没颜色。原因L1 损失主导模型学会预测接近 0 的 ab灰色来最小化平均误差。解决提高分类损失权重或在 L1 里对 ab 做加权让色度误差惩罚更重。5.2 颜色溢出到物体边界外现象天空的蓝色渗到建筑上人脸泛绿。原因解码器上采样时细节丢失跳跃连接没接好或训练分辨率太低。解决确认每层 concat 的通道对齐训练尺寸提到 256 以上必要时在损失里加感知损失约束边缘。5.3 显存爆掉batch 只能设 2现象训练中途 OOM。原因解码器 concat 后通道翻倍中间特征图显存占用大。解决开混合精度把编码器最深通道从 512 降到 256或用梯度检查点换显存。5.4 验证集效果好换张图就崩现象在自己划分的验证集上色差很低实际用别的图偏色严重。原因数据集场景太单一模型过拟合到特定色调。解决训练集按场景分层采样加几何增强验证集单独留一批不同来源的图。5.5 推理速度慢单张要好几秒现象部署时延迟高。原因模型在全分辨率上跑且没做推理优化。解决先缩放到 256 上色再放大回原尺寸或用 ONNX/TensorRT 导出加速全卷积结构支持动态尺寸。6. 让上色结果更耐看一个可验证的调优技巧训练跑通只是及格线真正拉开差距的是推理阶段的后处理。我踩过最深的坑是直接拿网络输出当最终结果颜色要么发灰要么过饱和。后来固定了一套流程先对 ab 做轻微高斯平滑抑制色块噪点再按图像亮度做自适应饱和度补偿——暗部提饱和、亮部压饱和最后用原图的边缘做引导滤波把颜色约束在物体内部。验证方法很直接准备 20 张有 ground truth 的图分别算 PSNR 和色差 ΔE对比加后处理前后的数值。我实测下来引导滤波能把边界溢色的 ΔE 降低约 15%而自适应饱和度让主观观感提升明显虽然 PSNR 可能略降——这说明指标和观感不总一致别只盯一个数。import cv2 def postprocess(rgb_out, guide_gray): # 引导滤波用灰度图做引导平滑颜色同时保边 b, g, r cv2.split(rgb_out) b cv2.ximgproc.guidedFilter(guide_gray, b, radius8, eps1e-2) g cv2.ximgproc.guidedFilter(guide_gray, g, radius8, eps1e-2) r cv2.ximgproc.guidedFilter(guide_gray, r, radius8, eps1e-2) return cv2.merge([b, g, r])逻辑说明guidedFilter用灰度图当引导能在平滑颜色的同时保留原图边缘radius控制平滑范围eps控制保边强度。参数上 radius 取 8、eps 取 1e-2 是常用起点边缘还是溢就调大 eps。注意这依赖 opencv-contrib 包装的时候别漏。我现在的习惯是任何上色模型上线前先跑一遍这套后处理对比指标和肉眼都过一遍才敢用。上色这活儿没有银弹网络结构决定下限后处理和调参决定上限。希望帮到你。本文还有配套的精品资源点击获取
返回列表