深度学习CNN实现灰度图像着色:U-Net与PyTorch实战解析 简介基于深度学习CNN网络实现图像着色的Python源码包专为人工智能、计算机视觉方向的学生和开发者设计适用于计科、数据科学、物联网等计算机相关专业。项目完整实现了ECCV16与SIGGRAPH17两种经典着色模型涵盖数据预处理、模型加载、特征提取、颜色预测与图像输出等完整环节包含模型定义、工具函数、预测脚本与示例图片代码结构清晰可直接运行体验彩色化效果也适合用于课程作业、毕业设计或深度学习入门进阶。压缩包共40个文件以14个Python脚本为主配合12张效果对比图、10个pyc缓存文件、2个说明文档和2篇参考文献PDF整体大小17.52MB目录划分明确便于按需查阅。目前已有50人学习下载资源兼具教学演示与二次开发价值若基础较好还可基于现有模块DIY扩展其他图像处理功能也适合作为项目立项演示。1. 从一张灰度图说起这个标题到底让你解决什么问题手头有一批黑白老照片、扫描档案或者视频抽帧想批量变成彩色人工上色一张人像可能要抠半小时而「基于深度学习CNN网络实现图像着色python源码.zip」这个标题说的就是把这件事交给cnn卷积神经网络输入端只给灰度图的亮度信息输出端由网络预测合理的颜色整套训练与推理逻辑用python写成源码打包成zip交付。它解决的问题很具体——批量、自动、可复现的灰度图着色。它适合刚学完cnn基础、想完整跑通一个视觉项目的初学者也适合做老照片修复、档案数字化的工程师。接下来从任务本质讲起一路到数据、模型、训练、推理每一步都有能直接复现的代码和踩过的坑。2. 图像着色任务拆解为什么CNN能把灰色图变彩色2.1 着色的本质是分布预测不是像素回归图像着色和图像分类有一个根本差别分类是单选着色是多选。灰度图上一个像素的亮度值无法唯一决定它的颜色。白衬衫在暖光下偏黄、在冷光下偏蓝投影到灰度图里可能一模一样。网络如果被要求逼近唯一答案它会对同一个输入学到多个合理颜色的平均平均的结果就是饱和度极低、灰蒙蒙一片。这是无数着色项目在第一天就翻车的原因所以第一步要先接受颜色预测本质上是一个分布预测问题。常见做法是在Lab色彩空间里做而不是RGB。把RGB转成Lab之后L通道表示亮度范围0到100a通道表示从绿到洋红b通道表示从蓝到黄。模型输入L通道预测a、b两个通道最后再把Lab转回RGB显示。这样设计有两个直接好处第一L通道与颜色通道解耦模型不需要自己学亮度恒等映射学习目标更单纯第二ab通道的欧氏距离与人眼感知差异更接近比在RGB空间里直接回归更容易得到自然的过渡色。还需要理解一个信息论上的现实训练时输入L通道、监督ab通道推理时只有灰度图模型要补的信息在输入端已经被丢弃了。所以模型输出的颜色永远只是「一个合理解」不是「唯一解」。这解释了为什么着色模型偶尔会画出与原图颜色不同的结果——那不是bug是任务本身的不确定性。我见过有人拿着一张彩色原图对比模型输出发现草地被画成红色就断定模型坏了其实比较合理的做法是看它画的是不是「符合条件的场景色」。2.2 为什么是CNN局部纹理、感受野与语义先验传统做法里最典型的是颜色传递用户给一张参考彩色图算法把参考图的颜色统计搬到灰度图上。它的问题在于需要一张与目标内容高度相似的参考图老照片一般找不到对应参考这个方法直接失效。人工规则也不可行蓝天在上、草地在下、肤色偏橙这些是语义知识不是几行if-else能枚举的——这也是为什么标题会把重心放在深度学习CNN上。cnn卷积神经网络的价值在于卷积核的局部连接和权值共享让它天然擅长提取局部纹理而颜色判断恰恰依赖这些纹理线索。一个区域纹理细碎又偏高频像树叶颜色倾向绿色一个区域纹理平滑且位于画面上方倾向天空蓝。网络通过层层卷积隐式学到的其实是「物体是什么」的语义先验这是传统方法不具备的核心能力。你用ImageNet预训练模型做编码器相当于把这些语义先验直接搬过来后面训练只是做微调收敛会快很多。还要看感受野。着色不是抠图判断单个像素的颜色需要大范围上下文它属于什么物体、处于画面什么位置。浅层特征图分辨率高、细节足但感受野小深层特征图经过多次下采样感受野大携带语义信息。所以主流着色模型几乎都是编码器-解码器结构编码器逐步下采样扩大感受野解码器逐步上采样恢复分辨率。U-Net在这个结构上加了跳跃连接把浅层细节直接拼到解码器后面讲模型时我会重点展开这一点。2.3 评估指标PSNR高不等于好看别被数字骗了着色没有唯一正确答案所以评估特别容易被数字骗。初学者最容易拿PSNR当标准但PSNR在着色任务里对灰色图特别友好灰色图的每个像素都接近所有合理答案的平均值像素误差天然小PSNR虚高可画面毫无生气。SSIM会看结构相似性比PSNR靠谱一点但SSIM也不懂「草应该是绿的」这种语义正确性。更麻烦的是还有人把着色当分类任务看Top-1准确率这在313个bin的分类设定下意义有限因为相差一个格子的颜色错误远比把天空画成草地轻。我自己的习惯是三个指标一起看PSNR看亮度保真SSIM看结构再加一个ab通道标准差看色彩丰富度。ab通道标准差这个指标很简单把预测图的ab通道取标准差如果明显低于真实图说明颜色集中在中性色附近模型在偷懒。论文里通常还会做用户研究找一群人来盲评哪张更像照片工程上就是让最终用户看一眼。到这一步你会发现指标只是帮助定位问题的工具最终拍板的是人眼。指标它衡量什么着色任务里的坑PSNR像素级误差对灰蒙蒙的结果虚高不能单独用SSIM局部结构相似性不评估颜色语义是否正确ab通道标准差色彩丰富度不能反映颜色是否对要配合前两个用3. 数据与预处理先学会把RGB拆成Lab再谈模型3.1 数据集选型与目录结构拿什么练、怎么摆图像着色需要「彩色图-转灰度-预测颜色」的监督对所以训练集必须是彩色图转灰度可以在数据加载时现做也可以预处理成npy缓存。数据集常见选择有三个ImageNet类别多、语义丰富Places365场景多、颜色分布和真实摄影接近COCO目标密集、适合验证物体边界。我自己做老照片方向时最常用Places365的验证集因为老照片大多是风景、建筑、人像这类场景色和Places365的分布很接近。工程目录的常见摆法是这样的这个结构也适用于绝大多数视觉训练项目colorization_project/ ├── data/ │ ├── raw/ # 原始彩色图片按类别分目录 │ └── processed/ # 转换后的 npy 缓存L 和 ab 分开存 ├── checkpoints/ # 训练权重 ├── logs/ # TensorBoard 日志 ├── train.py ├── model.py ├── dataset.py └── requirements.txt注意原始图按类别分目录不是为了好看是为了训练时能做按类别的采样和评估比如只看人像类别的效果、只看建筑类别的效果便于定位是数据问题还是模型问题。环境搭建的话我不建议跟着网上零散的python安装教程一条条装直接用conda建一个python 3.8以上的虚拟环境再装requirements.txt里固定的torch、torchvision、opencv-python、scikit-image、numpy这几个关键依赖版本锁死避免后人复现时被torch版本差异坑到。3.2 预处理脚本RGB转Lab并缓存成npy下面这段脚本做的是把一张彩色图变成模型的训练样本裁剪、缩放、RGB转Lab、归一化、存npy。这是整个项目里最不起眼但最容易出错的一步。import numpy as np import cv2 from skimage import color from pathlib import Path INPUT_SIZE 256 def preprocess_image(img_path: str, out_l: Path, out_ab: Path): # 读图先用中心裁剪切成正方形再缩放到统一尺寸 img cv2.imread(str(img_path)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) h, w img.shape[:2] m min(h, w) top, left (h - m) // 2, (w - m) // 2 img img[top:top m, left:left m] img cv2.resize(img, (INPUT_SIZE, INPUT_SIZE)) # RGB - LabL 作为输入ab 作为监督 lab color.rgb2lab(img).astype(np.float32) L lab[:, :, 0] # 0~100 ab lab[:, :, 1:] # a: -128~127, b: -128~127 # 统一归一化L 映射到 [-1, 1]ab 缩放到 [-1, 1] L_norm (L / 50.0) - 1.0 ab_norm ab / 128.0 np.save(out_l, L_norm) np.save(out_ab, ab_norm)逻辑上只有三个关键点但每个都会影响训练。第一用中心裁剪而不是直接resize因为直接resize会把主体长宽比拉歪解码器后面要恢复的细节边界是扭曲的颜色会顺着错误的边界晕开。第二L除以50再减1是把0到100线性映射到-1到1ab除以128是映射到-1到1。这样所有输入输出共享同一量纲最后一层用tanh或者线性输出都比较好对齐。第三skimage的rgb2lab要求输入是0到255的RGB数组类型可以是uint8或float32千万别送0到1的归一化数据进去算出来的Lab全是错的。INPUT_SIZE这个参数直接影响显存和画质。256是精度和显存的常见折中训练速度快颜色边界也够用。如果项目要求高分辨率输出我的做法是先用256训练到收敛再用512分辨率对模型做少量step的微调而不是一开始就上大图否则显存和训练时间都会涨得很离谱。3.3 数据增强与类别均衡别让网络只会画灰色数据增强以几何为主随机水平翻转、随机裁剪、随机旋转正负10度。这些增强不改变颜色分布安全又有效。有一点要特别注意不要在Lab空间里做类似RGB空间里那种对通道做的色彩抖动因为ab通道本来就是学习目标抖动会直接扭曲监督信号。如果你确实想做颜色类的增强正确的顺序是在RGB空间做完增强再转Lab拆通道这样既增加了颜色多样性又不污染监督。更隐蔽的问题是ab通道的分布极度不平衡。自然图像里灰色、棕色、肤色这类低饱和颜色出现的频率远高于荧光粉、亮紫这种高饱和颜色ab空间里低饱和区域的样本量可能是高饱和区域的上百倍。如果不做处理网络只要学会输出灰调就能把loss压得很低训练结束时色彩丰富度永远上不去。常见做法是给每个量化bin设置权重稀有颜色的样本权重放大公式是这样# 训练前统计每个 bin 的经验概率 p_q然后做平滑加权 epsilon 1e-6 lam 0.5 Q 313 weights (1.0 - lam) * (p_q epsilon) lam / Q weights weights / weights.sum() * Q # 归一化保持量纲这段代码里的p_q来自训练集统计每个bin对应ab平面上一个量化格子统计所有训练样本的ab值落在哪个格子。lam取0.5是论文里常用的平滑系数意思是稀有颜色最多能获得接近均匀分布的权重不会因为过采样把某些极端颜色学爆。这个weights在损失函数里作为每个类的加权系数传入后面第5章会再提到「loss不降先查权重有没有传进去」指的就是它。4. 模型与损失设计用PyTorch搭一个能画彩色的U-Net4.1 编码器-解码器与跳跃连接U-Net为什么是默认选编码器部分我一般直接复用ResNet18的卷积骨架把输入从3通道改成1通道L。为什么要用预训练的ResNet因为ImageNet上的分类预训练已经让编码器具备很强的语义先验着色训练时收敛速度快很多。不过预训练模型的conv1原本接收3通道需要把权重按通道平均后复制到1通道这是常见做法不这么做的话随机初始化会让语义先验白白浪费。跳跃连接为什么关键要看编码器和解码器的信息差异。编码器浅层分辨率高、细节足但语义弱深层分辨率低、语义强但细节丢得多。如果没有跳跃连接解码器只能靠深层低分辨率特征上采样物体边缘会出现颜色溢出树叶的边界糊成一团。跳跃连接把高分辨率的局部纹理直接传给解码器让颜色严格贴合物体边界。所以U-Net结构在像素级预测任务里几乎是默认选着色、分割、去噪、超分都用它打底。4.2 分类还是回归把ab通道量化成313个bin直接回归ab通道用L1或L2损失实现简单但前面章节反复提过回归有均值陷阱颜色会偏灰。实践中效果更稳的做法是分类把ab平面按步长10量化剔除落在有效色域外的格子得到约313个有效bin。网络最后一层输出313维概率分布训练用加权交叉熵推理时按概率期望还原ab值而不是直接argmax。这里有个容易误解的点为什么不干脆argmax取概率最大的那个bin因为量化是有台阶的argmax会让相邻像素在bin边界处出现跳变输出图会有明显的色带伪影。用概率期望也叫soft-encoding相当于对bin中心坐标做加权平均颜色过渡更平滑。这也是标题里「深度学习CNN」和传统查表法最大的区别——网络学的是颜色分布不是一张映射表。4.3 最小可跑通的模型代码model.py与损失函数下面是一个精简但完整的U-Net实现输入单通道L输出313通道的概率图。通道数量比标准U-Net小是为了单卡能跑也方便你改回2通道做回归对比。import torch import torch.nn as nn import torch.nn.functional as F class ConvBlock(nn.Module): def __init__(self, ch_in, ch_out): super().__init__() self.conv nn.Sequential( nn.Conv2d(ch_in, ch_out, 3, padding1, biasFalse), nn.BatchNorm2d(ch_out), nn.ReLU(inplaceTrue), nn.Conv2d(ch_out, ch_out, 3, padding1, biasFalse), nn.BatchNorm2d(ch_out), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.conv(x) class ColorNet(nn.Module): def __init__(self, n_bins313): super().__init__() self.enc1 ConvBlock(1, 64) self.enc2 ConvBlock(64, 128) self.enc3 ConvBlock(128, 256) self.enc4 ConvBlock(256, 512) self.bridge ConvBlock(512, 512) self.dec4 ConvBlock(1024, 256) self.dec3 ConvBlock(512, 128) self.dec2 ConvBlock(256, 64) self.head nn.Conv2d(64, n_bins, 1) def forward(self, L): s1 self.enc1(L) # 分辨率 256 s2 self.enc2(F.max_pool2d(s1, 2)) # 128 s3 self.enc3(F.max_pool2d(s2, 2)) # 64 s4 self.enc4(F.max_pool2d(s3, 2)) # 32 x self.bridge(F.max_pool2d(s4, 2)) # 16 # 解码器上采样后与同层浅层特征拼接再卷积 x F.interpolate(x, scale_factor2, modebilinear, align_cornersFalse) x self.dec4(torch.cat([x, s4], dim1)) x F.interpolate(x, scale_factor2, modebilinear, align_cornersFalse) x self.dec3(torch.cat([x, s3], dim1)) x F.interpolate(x, scale_factor2, modebilinear, align_cornersFalse) x self.dec2(torch.cat([x, s2], dim1)) x F.interpolate(x, scale_factor2, modebilinear, align_cornersFalse) return self.head(x) # (B, 313, H, W)forward里的流程是四次下采样把分辨率降到原图的1/16bridge层在最低分辨率做特征变换然后三次上采样每次上采样后都和对应层的浅层特征拼接最后head卷积输出313通道的概率图。注意dec4的输入通道是bridge的512加上s4的512所以ConvBlock第一层写成1024dec3是256加128得512dec2是128加64得256这个通道数对应关系不能错一错维度就对不上。配合分类损失的代码如下class WeightedCrossEntropyLoss(nn.Module): def __init__(self, weight): super().__init__() self.weight weight # 形状 (313,)训练前统计得到 def forward(self, pred, target): # pred: (B, 313, H, W), target: (B, H, W)存的是 bin 索引 B, C, H, W pred.shape logp F.log_softmax(pred, dim1) target target.unsqueeze(1) one_hot torch.zeros_like(pred).scatter_(1, target, 1.0) loss -(one_hot * logp * self.weight.view(1, C, 1, 1)).sum(dim1).mean() return losstarget在dataset里生成把ab通道的连续值量化成bin索引ab归一化到-1到1之后先用ab乘以128复原坐标再按步长10的量化表查索引。如果你只是想先跑通训练也可以把head输出改成2通道、用L1损失回归ab网络结构不用变只是效果大概率偏灰。5. 训练避坑排查颜色偏灰、loss震荡、显存爆掉的三个现场5.1 训练超参数参考batch、lr、epoch怎么设训练超参是这类项目里最吃经验的部分给一组我常用的起点值超参数参考值说明输入分辨率256×256显存紧张先降到192×192batch size32单卡不够就用梯度累积凑等效batch优化器AdamWlr2e-4weight_decay1e-4学习率调度cosine decay 5个epoch warmup前5个epoch从0线性升到2e-4总epoch数60到80Places365验证集规模足够混合精度AMP显存省30%到40%速度提升明显warmup不是可选项。分类头是随机初始化的一开始就用大学习率313维输出很容易在第一个batch炸出nan或者把预训练的编码器权重冲坏。cosine decay最后衰减到1e-6左右让模型在后期稳定微调。AMP的话PyTorch自带的torch.cuda.amp就可以别自己手写fp16转换容易踩数值稳定的坑。5.2 现象一loss不降或震荡问题多半不在模型现象是前几个epoch loss下降很快到某个点后开始震荡ab标准差始终上不去。我排查这类问题有个经验先别改网络结构着色任务的结构容错率很高真正的问题往往在数据和训练配置上。常见原因有三个。第一类别权重没传进loss里或者weight被默认归一化成了均匀分布稀有颜色完全没被照顾模型当然选择躺平输出灰色。检查方法很简单训练脚本里print一下loss模块里的weight.sum()如果等于Q说明是均匀的等于Q才是你要的加权。第二learning rate偏高loss在最优值附近来回弹跳用lr1e-4重新跑一个短实验对比。第三batch size太小ab分布采样不稳定每个step看到的颜色分布差异过大把batch提到32以上或用梯度累积。还有一个比较玄学但真实存在的坑不要用torchvision自带的ImageFolder直接当dataset它不会做Lab拆分也不会做bin量化很多人在这上面花了两天时间最后发现模型学到的是RGB到RGB的恒等映射。5.3 现象二输出全是灰色均值陷阱与分类反转现象是训练正常结束输出图结构清晰但颜色像褪色ab标准差只有真实图的一半不到。这是着色项目最常见的翻车现场。原因分三种对应三个不同的解决方向。第一种用了回归损失L1或L2模型学到了多个合理答案的平均值这种情况把损失换成加权交叉熵基本能解决。第二种用的是分类方案但推理时用了argmax量化误差在低饱和区域尤其明显解决方法是改成soft-encoding就是第4章说的按概率期望还原ab值。第三种反归一化写错了ab还原时没有乘以128输出颜色被压缩在一个极小的范围内。检查方法是拿一张测试图统计ab通道的标准差如果预测图的标准差只有真实图的1/3甚至更低先怀疑后处理再怀疑损失函数。我之前交付过一版颜色偏灰的模型客户看了直摇头后来发现就是反归一化少乘了128改完ab标准差从2.x直接回到8.x。5.4 现象三显存不足与训练卡死先查数据加载再查GPU现象是在单卡上跑256分辨率加batch32居然OOM或者训练几轮后程序卡住不动GPU利用率掉到0。OOM的常见原因不是编码器而是解码器的通道数太大。bridge层输出512通道解码器第一个卷积块要处理1024通道的拼接特征这是显存高峰。解决方法是把解码器通道数整体砍半dec4从256降到128dec3从128降到64dec2从64降到32效果损失不大但显存占用明显下降。另一个方案是梯度累积batch32拆成4个step各8张等4个step的梯度加起来再更新一次参数等效batch还是32。训练卡死这个坑更隐蔽。先看GPU利用率如果利用率是0、CPU占用很高问题在DataLoadernum_workers设成了0图像解码全在主进程里排队加图片缓存文件损坏时cv2.imread返回Nonedataset里没做过滤模型forward到一半数据是空张量直接崩。解决方法是num_workers设为4到8并在dataset里检查cv2.imread结果是否为None为None直接跳过同时提前扫描一遍数据目录把损坏文件隔离出来。6. 推理脚本与进阶把权重用起来再把思路迁移出去6.1 推理脚本把权重用到一张新灰度图上训练好的模型最终要落到一张灰度图上推理脚本的核心是读灰度图、归一化L、进网络、按概率期望还原ab、合并Lab、转回RGB。def infer_color(model, gray_path, out_path, bin2ab, devicecuda): # bin2ab: (313, 2) 的映射表训练时保存用于把 bin 索引换成 ab 坐标 model.eval().to(device) gray cv2.imread(gray_path, cv2.IMREAD_GRAYSCALE) gray cv2.resize(gray, (256, 256)) L_norm torch.from_numpy((gray.astype(np.float32) / 50.0 - 1.0))[None, None].to(device) with torch.no_grad(): logits model(L_norm) # (1, 313, 256, 256) prob torch.softmax(logits, dim1) grid torch.arange(313, devicedevice).view(1, 313, 1, 1) quant_idx (prob * grid).sum(dim1, keepdimTrue) # soft-encoding ids quant_idx[0, 0].long().cpu().numpy() a bin2ab[ids, 0] b bin2ab[ids, 1] L cv2.resize(gray, (256, 256)) lab np.stack([L, a * 128.0, b * 128.0], axis-1) rgb color.lab2rgb(lab) cv2.imwrite(out_path, cv2.cvtColor((rgb * 255).astype(np.uint8), cv2.COLOR_RGB2BGR))注意gray先读成0到255的uint8L的范围是0到100但OpenCV读灰度图的范围是0到255两者差了一个尺度。我在上面的代码里直接用gray做L通道配合ab乘以128Lab合并时容易对不上。更稳的做法是推理时重新读原图并转Lab取它的L通道再缩放L color.rgb2lab(cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB))[:, :, 0]然后L_norm (L / 50.0 - 1.0)。bin2ab这个映射表必须在训练时随权重一起保存否则推理端只能靠训练时的量化表反推容易错位。6.2 进阶领域微调、Transformer与像素级任务迁移模型跑通只是起点。如果你想把它用在老照片修复上常见做法是收集少量目标风格的彩色照片合成灰度做微调让模型适应老照片的色调和噪点分布而不是直接在自然图模型上硬推理。网络结构方面Swin-UNet这类Transformer结构在颜色一致性上比纯CNN更强但需要更多训练数据数据量不够时反而会退化。更重要的是思路迁移Lab拆通道加分类损失这套骨架在视频插帧、超分辨、图像修复这些像素级任务里都能复用换任务时损失函数和数据均衡的逻辑不用大改。我养成的习惯是训练完先看ab标准差再看PSNR和SSIM最后直接输出效果图让眼睛判断指标只是定位问题的工具人眼才是最终验收。希望帮到你。本文还有配套的精品资源点击获取