深入浅出CNN图像着色：Lab色彩空间与VGG16实战解析

简介基于深度学习CNN网络实现图像着色的Python源码包聚焦自动化图像上色任务适合计算机视觉、人工智能方向的在校学生、教师或开发者用于课程作业、毕业设计及项目演示。代码涵盖PyTorch与Caffe两种主流框架分别对应ECCV16与SIGGRAPH17经典着色模型并附有现成测试图片与输出样例可作为入门学习或二次开发的实用起点。资源包共40个文件整体约17.52MB。其中包含14个py源码文件用于定义网络结构、训练与推理流程12张png图片作为输入或效果展示10个pyc缓存文件方便直接导入运行另有2个txt说明文档和2篇pdf参考文献帮助理解算法原理与环境配置。目录结构较为清晰模型代码集中在colorizers目录下。目前已有50人学习下载。解压后将项目路径改为英文名即可运行适合希望快速复现深度学习图像着色效果、进一步DIY功能的用户。借助完整源码、示例图片和参考文档读者可深入理解经典着色网络的设计思路并基于自身需求扩展新的视觉应用。1. 图像着色不是PS滤镜CNN方案解决的是什么问题灰度图着色这个题目在深度学习流行之前基本是靠人工或者颜色统计先验硬猜效果经不起放大看。而基于深度学习CNN网络的图像着色核心是把“这是一个什么物体”和“这个物体通常是什么颜色”两件事绑在一起学让网络从语义层面猜测颜色而不是靠滤镜糊一层上去。这个python源码包对应的就是一条完整链路图像预处理、VGG编码器、解码器、颜色分类损失、推理上色。它既是CNN相关技术的综合练手项目也是一个能直接接到老照片修复、影视素材着色、设计辅助上的可用方案。适合刚学完卷积神经网络想找完整项目经验的初学者也适合需要给灰度素材批量上色的从业者。接下来我从原理讲到跑通再把训练时最容易翻车的地方一个个说清。2. 原理与选型Lab色彩空间、VGG编码器与分类损失为什么是标配图像着色这个任务最反直觉的一点是直接预测 RGB 三通道是必败路线。这个结论不是谁拍脑袋定的而是被大量实验反复验证过的。要理解为什么先得明白着色任务和普通回归任务在数学本质上的差别。2.1 为什么直接预测 RGB 会糊成一团灰RGB 三个通道之间存在很强的线性相关性。自然图像里一个像素的 R、G、B 值往往同步变化这意味着网络在输出层等于要同时回归三个耦合的目标。更麻烦的是L2 损失函数在遇到“多模态输出”时会自动取平均值。什么叫多模态一张灰度图里某个像素可能对应天空的亮蓝也可能对应水泥墙的浅灰——同一个输入多个合理答案。L2 回归会把所有合理答案平均在一起蓝色加灰色加白色平均完就是一片灰蒙蒙的脏颜色饱和度全丢。所以常见的着色源码包里第一步一定是把图像从 RGB 转到 Lab 色彩空间。Lab 是一个把亮度和颜色拆开的色彩模型L 通道只表示明暗a 通道表示从绿到红的渐变b 通道表示从蓝到黄的渐变。做训练时把 L 通道作为网络输入目标就只剩预测 a、b 两个通道且 a、b 通道之间的相关性远低于 RGB 三通道。网络不用再同时猜亮度和颜色专心想“这块区域该偏红还是偏蓝”就行。这个改动本身不增加任何计算量但能让收敛速度和最终饱和度上一个台阶。还有一个容易被忽略的细节图像从 RGB 转到 Lab 再转回来本身是有信息损耗的加上浮点取整如果来回转几次暗部会出现色带。所以源码里一般只做一次转换全程在 Lab 空间里训练和推理最后输出时才转回 RGB。这个习惯在调代码时能少踩很多坑。2.2 编码器选型VGG16 的预训练权重为什么是首选确定了 Lab 空间之后下一步是选网络骨干。当前着色源码最常见的结构是“编码器-解码器”编码器负责把灰度图压缩成高层次的语义特征解码器负责把特征图上采样回原尺寸同时预测每个像素的颜色。编码器这一侧VGG16 出现在大多数 python 实现里不是因为它结构最新而是因为它有三个实在的好处。第一VGG16 的 ImageNet 预训练权重在 PyTorch 和 TensorFlow 里都是官方维护、一行代码就能加载不需要自己费劲去找权重文件。第二VGG16 做的是 5 次下采样每次把特征图尺寸缩小一半总共缩小 1/32。这个下采样倍数意味着网络在最后层能看到足够大的感受野能理解“这是一只鸟”而不是“这是一个灰色的小块”。第三VGG16 没有残差连接和多分支结构就是一条直路Feature Map 的尺寸变化极其规则解码器做上采样时不用费心对齐。对比之下ResNet 的残差结构在语义特征提取上确实更强特备是深层语义但作为编码器时残差分支的存在会让解码器在融合特征时多一层设计成本。MobileNet 系列胜在轻量适合部署但作为编码器时首层卷积的 stride 和通道数都要跟着改改动量比换 ResNet 还大。所以对于一个以跑通为主、快速验证效果的源码包VGG16 是性价比最高的选择这也解释了为什么各类着色开源项目里 VGG 系骨干的出现频率远高于其他网络。当然VGG16 的短板也要心里有数第一个全连接层是 4096 维参数量很大显存占用高。因此源码里通常会做一件事——把全连接层整体丢掉只保留前面的卷积层。编码器输出的是一组特征图而不是一维向量正好满足解码器逐像素上采样的需求。这一刀砍掉的是全连接层的千万级参数精度几乎不受影响。2.3 训练目标的变换分类损失比回归损失更能保住颜色饱和度很多人第一次看到着色源码里的损失函数时会愣一下为什么算的是交叉熵而不是常见的均方误差这就要回到着色任务的多模态本质。前面说过L2 回归会把多个合理颜色平均成灰色那换个思路——不把颜色当连续值而是当成类别来分类。这就是分类损失的核心逻辑。具体做法分三步。先把 a、b 两个通道的取值范围量化成若干个离散的类别常见的是 313 类这个数字是把 a 通道从 -110 到 110、b 通道从 -110 到 110 按 10 左右的间隔切分后组合出来的。训练时每个像素的真实颜色被映射到最近的颜色类别上网络要对每个像素输出一个 313 维的概率分布然后用交叉熵损失去逼近真实类别。第二步是类别权重。自然图像里的颜色分布极不均匀地面、墙壁、皮肤的灰黄色系占了大多数像素而紫色、橙色、青色这类高饱和颜色占比极低。如果直接按普通交叉熵算网络学到的会是“全都猜棕色和灰色损失最小”。所以源码里一般会给稀有颜色更高的权重常见做法是用颜色分布的核密度估计取倒数做权重。这个权重表是在训练开始前统计整个数据集得到的和模型参数一起参与训练但不参与梯度更新。第三步是退火温度。推理时模型的最后输出是一个概率分布直接取最大概率的颜色类别会导致颜色块状化太硬取概率加权平均又会出现回归损失那种泛灰问题。折中方案是引入温度系数 T预测概率除以 T 之后再归一化。T 大于 1 时分布变平颜色过渡自然适合风景类大色块区域T 小于 1 时分布变尖颜色浓烈适合需要凸显饱和度的场景。源码里通常默认 T 在 0.38 附近这个值是原作者在大量实验里试出来的一般不用改。3. 跑通源码环境准备、数据集组织与训练命令原理部分理解了接下来进入动手环节。拿到 python 源码包之后正确顺序不是先打开模型定义文件而是先把环境和数据准备好然后让训练脚本跑起来最后再回头改结构。这样每次改动都有明确的验证基准。3.1 环境准备Python 版本、CUDA 与依赖安装这个项目用 PyTorch 实现是最常见的因为预训练权重获取方便、调试灵活改网络结构也省事。Python 版本建议 3.8 到 3.10太老的新版 PyTorch 装不上太新的版本又容易出现依赖冲突。CUDA 部分训练阶段建议 11.7 以上算力小于 6.0 的老显卡会踩到算子兼容的坑但不影响代码运行只是慢。# 创建独立虚拟环境避免污染系统 Python conda create -n colorizer python3.9 -y conda activate colorizer # 安装 PyTorch按 CUDA 版本选择对应命令 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 # 安装图像处理和训练辅助依赖 pip install numpy opencv-python pillow tqdm scikit-image matplotlib这段命令里--index-url指定的是 PyTorch 官方 CUDA 11.8 版本的下载源如果你的机器 CUDA 版本不同要到官网查对应命令。opencv-python是主力图像库scikit-image用来做 Lab 颜色空间转换——注意这里用skimage.color.rgb2lab比 OpenCV 的转换更规范因为 OpenCV 的 Lab 范围定义和论文里用的不完全一致后续量化 313 类颜色时容易差一个偏移量。依赖装完后先跑一个一句话验证确认 PyTorch 能正常调用 GPU同时验证 Lab 转换函数是否安装成功。import torch import numpy as np from skimage.color import rgb2lab, lab2rgb # 验证 GPU 可用性 assert torch.cuda.is_available(), CUDA 不可用检查驱动和 PyTorch 版本 # 验证 Lab 转换接口 fake_rgb np.random.randint(0, 255, (64, 64, 3), dtypenp.uint8) lab rgb2lab(fake_rgb) print(GPU:, torch.cuda.get_device_name(0)) print(Lab 转换正常a通道范围:, lab[..., 1].min(), ~, lab[..., 1].max())这段验证脚本里rgb2lab输出的 a 通道正常人眼可见颜色范围大约在 -128 到 127 之间但实际图像基本落在 -110 到 110。如果看到极端值就说明输入图像本身有问题不是转换函数的问题。同时脚本里显式断言了 CUDA 可用性如果不加这一步后面训练到一半才发现模型被放到了 CPU 上那才是真浪费时间。3.2 数据准备从自定义图片到可训练张量着色任务的训练数据不需要人工标注直接用 ImageNet 或者其他自然图像数据集就能训练因为输入 L 通道、目标 ab 通道都是自动生成的。这也让这个项目在数据准备环节比其他视觉任务省事得多。数据集目录结构一般这么组织dataset/ ├── train/ │ ├── 000001.jpg │ ├── 000002.jpg │ └── ... └── val/ ├── 000101.jpg └── ...目录本身没有任何特殊要求关键是训练脚本里要写一个通用的图片读取和预处理函数。下面这段是我常用的预处理流程输入一张任意尺寸的 RGB 图输出一个可以直接喂给网络的张量import cv2 import torch import numpy as np from skimage.color import rgb2lab def load_image_as_lab(image_path, target_size224): # 读取图片OpenCV 默认读成 BGR要转回 RGB img_bgr cv2.imread(image_path) img_rgb cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) # 短边缩放到 target_size中心裁剪成正方形 h, w img_rgb.shape[:2] scale target_size / min(h, w) new_h, new_w int(h * scale), int(w * scale) img_resized cv2.resize(img_rgb, (new_w, new_h), interpolationcv2.INTER_AREA) start_h (new_h - target_size) // 2 start_w (new_w - target_size) // 2 img_cropped img_resized[start_h:start_h target_size, start_w:start_w target_size] # RGB 转 Lab提取 L 通道作为输入ab 通道作为目标 img_lab rgb2lab(img_cropped) l_channel img_lab[:, :, 0] / 100.0 # L 范围 0~100归一化到 0~1 ab_channels img_lab[:, :, 1:] # ab 通道保持原始范围 # 转成 PyTorch 张量形状 (1, H, W) 和 (2, H, W) l_tensor torch.from_numpy(l_channel).unsqueeze(0).float() ab_tensor torch.from_numpy(ab_channels.transpose(2, 0, 1)).float() return l_tensor, ab_tensor这段代码有几个参数值得细说。target_size224不是随意选的它对齐了 VGG16 预训练模型的输入尺寸特征图在 5 次下采样后会变成 7×7这个尺寸在解码阶段正好是 2 的完整幂次倍上采样时不会出现尺寸对不齐。scale target_size / min(h, w)是等比例缩放保证图像不变形直接用cv2.resize粗暴缩到 224×224 会让画面变形颜色语义也跟着歪。还有l_channel / 100.0这步归一化很多初次接触的人会漏掉。L 通道的范围是 0 到 100如果不归一化直接喂进网络VGG16 预训练权重的第一个卷积层是按照 ImageNet 标准化后的输入分布设计的数值范围不对会导致通道激活异常看起来像是“模型没收敛”其实是输入分布出了偏差。3.3 训练启动超参设置与日志里到底要看什么环境就绪、数据函数写完接下来就是核心的训练命令。需要说明的是这里给的是我落地这个源码包时常用的一套配置如果你的 GPU 显存不是很大按注释调整即可。# train.py 核心训练循环示意 import torch import torch.nn as nn from torch.utils.data import DataLoader, Dataset # 关键超参数 batch_size 32 epochs 30 learning_rate 3e-4 temperature 0.38 # 推理时用的温度系数 num_color_classes 313 # ab 通道量化后的颜色类别数 # 损失函数带类别权重的交叉熵 # class_weights 在数据集统计阶段预先算好形状是 (313,) criterion nn.CrossEntropyLoss(weightclass_weights) # 优化器Adam 在这类任务上比 SGD 好调收敛更快 optimizer torch.optim.Adam(model.parameters(), lrlearning_rate) # 学习率调度每 10 个 epoch 衰减一半让训练后期微调更稳定 scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.5) for epoch in range(epochs): model.train() running_loss 0.0 for l_tensor, ab_tensor in train_loader: l_tensor l_tensor.cuda() # ab_tensor 是连续值需要按量化表映射到颜色类别索引 color_labels quantize_ab_to_class(ab_tensor).cuda() # 前向传播输入 L 通道输出每个像素在 313 个颜色类上的概率 logits model(l_tensor) # 输出形状: (B, 313, H, W) # 计算交叉熵注意 PyTorch 的 CrossEntropyLoss 需要把通道维放在第二维 loss criterion(logits, color_labels) optimizer.zero_grad() loss.backward() optimizer.step() running_loss loss.item() * l_tensor.size(0) avg_loss running_loss / len(train_loader.dataset) print(fEpoch {epoch1}/{epochs}, Loss: {avg_loss:.4f}) scheduler.step()训练脚本里最关键的参数是num_color_classes。在量化 ab 通道时a、b 取值范围通常截断到 [-110, 110]按 10 为步长量化就是 22×22 个格子但其中只在色度图上有一个类似椭圆形状的区域覆盖了真实颜色实际只保留其中 313 个点。如果你改了这个量化参数模型输出层维度、损失函数权重、推理时的反量化表都要同步改任何一个没对齐都会立刻报错。日志部分每轮打印的平均 Loss 值要注意它的正常走势前 3 个 epoch 损失会从 5.0 以上快速掉到 4.0 以下这个过程对应的是网络正在学会把图像分割成语义区域后面 10 个 epoch 损失下降变缓进入颜色细化的阶段。如果第 3 个 epoch 损失还在 5.0 以上先不要急着调模型检查是不是学习率设置过高、数据加载时 Lab 归一化出了问题或者类别权重没有生效。4. 避坑训练翻车的五个高频点这个项目跑通不难跑好有门槛。我把实践过程中遇到最多的问题按“现象→原因→解决”整理在这里每一个都是真实踩过的顺序按出现频率从高到低排。4.1 现象Loss 下降正常但输出图泛灰、像蒙了一层雾原因这个是最常见的情况占到了问题的一半。网络“学会”了输出所有类别的平均概率相当于每个像素都猜一个灰色交叉熵损失确实降低了但色彩饱和度全无。本质上是模型收敛到了一个次优解多模态问题里这个坑极其常见。解决分两步排查。第一步确认类别权重是否加上了如果数据集里某个颜色占比极高比如超过 20%不加权重网络一定会偏向它。第二步降低温度系数的初值训练最后几个 epoch 时把 T 从 0.38 调到 0.25让概率分布更尖锐强迫模型做更明确的颜色选择。4.2 现象训练 Loss 一直在 5.0 以上不下来原因Lab 通道预处理和网络输入对不上最常见的是 L 通道没有归一化到 [0,1]或者输入张量的通道维顺序不对。还有一种是数据加载函数返回的张量在 CPU 和 GPU 之间反复搬运每一步都在做数据拷贝效率极低且梯度传播被打断。解决在训练循环里加几行调试代码打印输入张量的形状、数值范围和是否要求梯度print(输入形状:, l_tensor.shape) # 期望 (B, 1, 224, 224) print(输入范围:, l_tensor.min().item(), l_tensor.max().item()) print(输入是否要求梯度:, l_tensor.requires_grad) # 期望 False print(模型输出形状:, logits.shape) # 期望 (B, 313, 7, 7) 或上采样后的尺寸这段调试代码能快速区分问题出在数据层还是模型层。如果输入范围不对去查预处理函数如果输出形状不对去查解码器的上采样倍数如果都正常但 Loss 还高再考虑是模型容量的问题。4.3 现象训练时显存溢出或者 batch size 稍微一调就 OOM原因VGG16 的参数量摆在那里加上解码器的转置卷积和上采样操作整张计算图很占显存。很多人拿到源码后直接照着 README 里的 batch size 跑完全忽略了自己的显卡只有 8G 显存这一现实。解决不做架构改动的前提下先把输入尺寸从 224 缩到 192。虽然 VGG16 预训练权重是为 224 设计的但把权重加载时用torchvision.models自带的transform_input处理一下192 输入也能正常跑显存占用直接降为原来的 73% 左右。还不行就启用梯度累积# 梯度累积batch_size8累积 4 步等效于 batch_size32 accumulation_steps 4 scaler torch.cuda.amp.GradScaler() # 混合精度训练 for i, (l_tensor, ab_tensor) in enumerate(train_loader): with torch.cuda.amp.autocast(): logits model(l_tensor) loss criterion(logits, color_labels) / accumulation_steps scaler.scale(loss).backward() if (i 1) % accumulation_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()这段代码里混合精度训练用torch.cuda.amp的两个关键 APIautocast自动把部分算子切到 FP16GradScaler防止梯度下溢。要特别注意loss除以accumulation_steps是必须的否则梯度累积后等效 batch size 变大学习率不调整就会让训练不稳定。4.4 现象训练时指标正常但推理时颜色整体偏移原因这是训练和推理代码不一致导致的问题。预处理时把 L 通道除以 100 做了归一化但推理脚本里忘了对输入做同样的归一化或者训练时用了数据增强的随机裁剪推理时直接缩放导致图像内容比例变化。解决把预处理逻辑抽成一个公共函数训练和推理共用一份代码。然后在保存模型时把图像的均值、标准差、Lab 通道归一化系数一并写进 checkpoint 文件里这样即使一个月后回来用这个模型也能准确还原预处理参数不会犯“当时好像没加这一步”的迷糊# 保存训练状态时附带预处理参数 preprocess_cfg { l_mean: 50.0, l_std: 50.0, ab_quantization_step: 10, temperature: 0.38, } torch.save({ model_state: model.state_dict(), preprocess: preprocess_cfg, }, checkpoints/model_epoch30.pth)4.5 现象训练很慢GPU 利用率只有 30%原因数据加载环节成了瓶颈。图像从磁盘读取、缩放、转 Lab、裁剪这些操作如果都在主进程里串行执行GPU 就只能干等着。哪怕数据集只有几千张图这个问题也会暴露得很明显。解决在DataLoader里开多进程加载并调大预取数量train_loader DataLoader( train_dataset, batch_size32, shuffleTrue, num_workers8, # 数据加载进程数建议和 CPU 核数相当 pin_memoryTrue, # 锁页内存加速 CPU 到 GPU 的数据拷贝 persistent_workersTrue, # 保持 worker 存活避免每个 epoch 重复创建 )前两个参数是显性提速第三个persistent_workers容易被忽略但它是训练提速最明显的一个。在 PyTorch 1.10 之前的版本里每个 epoch 结束时 worker 会退出下个 epoch 重新启动这部分启动开销在数据量大时能达到总训练时长的 10%。开启这个参数之后worker 常驻内存代价是占用一点内存空间对训练速度的提升非常明显。5. 一个技巧把“能着色”变成“可落地”的最后一公里源码跑通只代表你拿到了一台能出图的机器距离真正把它用到产品里还有一段距离。这里分享一个我最近在做的事也是我认为这个方向接下来值得投入的点。着色模型和超分辨率、去噪这类修复模型有一个共同特点它们都在生成不存在的细节。用户对皮肤、天空、植物这些常见物体的颜色有极强的先验判断力模型一旦把草地涂成紫色一眼就会被发现。所以落地时不能只靠模型输出得给结果加一道可控的约束。我一般会在模型推理后加一个低通融合把灰度原图的亮度通道和模型生成的 ab 通道分开处理——L 通道永远用原始图的只取模型生成的 ab 通道最后转为 RGB 时再附加上原始对比度信息。这样既能保留模型的语义颜色又不会因为模型对亮度通道的预测偏差导致整张图发灰或者过曝。如果你有足够的机器资源另外一个值得尝试的方向是把 VGG16 编码器替换成轻量化骨干比如 MobileNetV3。同一套数据上替换后模型体积能压缩到原来的四分之一左右推理速度翻倍在 CPU 上也能实时跑部署价值会高很多。但注意 MobileNetV3 的下采样倍数和 VGG16 不一样解码器的上采样层要做对应调整。验证这一步我强烈建议每次训练完都挑十张不同场景的图包括人像、风景、室内、夜景观察它们各自的着色倾向而不是只看一张图的效果。因为着色任务里Lake 的数值指标永远没有人的眼睛准数据增强翻车、颜色分布偏斜这类问题几乎都能在 10 张图里肉眼抓到。这是我的习惯也是希望你能坚持下去的方向。希望帮到你。本文还有配套的精品资源点击获取