尧图网页设计 
简介一份基于深度学习CNN网络实现图像着色的Python源码适合计算机相关专业学生、教师及企业开发者用于课程设计、毕业设计或入门进阶。资源共40个文件包含14个py源码、12张png效果图、10个pyc编译文件、2个txt说明和2份pdf参考文献压缩包大小17.52MB目录结构清晰。代码实现了eccv16、siggraph17等经典着色模型提供PyTorch与Caffe两种调用脚本并封装了基础工具模块可直接运行快速生成着色效果也便于二次开发。附带的示例图片与理论文献能够帮助理解CNN特征提取与色彩重建流程。已有50人学习下载适合需要完整可运行图像着色方案的学习者。1. 图像着色不是上色那么简单CNN 为什么能胜任这个活图像着色看起来像给黑白照片填颜色实际上是要让模型从灰度亮度图中“猜”出每个像素最可能的色相与饱和度。传统的规则算法只能做全局颜色映射遇到一幅画面里同时有天空、草地、人脸时就会乱套基于深度学习 CNN 的做法则是让网络在大量彩色图上学习亮度到颜色的统计规律本质上是在拟合一个条件概率分布。这个方向的核心价值在于只要把输入从单通道灰度图映射到 Lab 色彩空间的 L 通道再让 CNN 预测 a、b 两个颜色通道就能端到端地得到一个可复用模型。我早期接手这类 python 源码工程时最常踩的坑就是训练跑完出图灰蒙蒙或者边缘糊掉。这篇文章从任务拆解开始一路讲到可运行的 PyTorch 代码、参数设定和翻车排查适合刚接触深度学习的工程师也适合要给老照片、监控灰度图做批量着色的从业者。2. 图像着色任务拆解从灰度图到颜色空间的端到端设计拿到“基于深度学习CNN网络实现图像着色”这个标题时第一反应不要直接找模型代码而是先想清楚输入输出到底是什么、在哪个颜色空间里做回归。这一步决定了后面所有训练的稳定性。2.1 输入输出怎么定为什么偏偏选 Lab 色彩空间RGB 三个通道之间相关性极高亮度变化和颜色变化纠缠在一起。如果用 RGB 灰度图作为输入、预测 RGB 彩色图模型既要学亮度结构又要学颜色分布任务更难更关键的是RGB 的欧氏距离与人眼感知差异不一致MSE 损失在这里会给暗部颜色更大的梯度压力训练时暗部颜色容易偏。业界做图像着色最常见的做法是把图片转到 Lab 色彩空间其中 L 通道只表示亮度a 通道表示从绿色到品红b 通道表示从蓝色到黄色。输入侧将彩色图转成灰度或直接取 Lab 的 L 通道作为 CNN 的单通道输入。由于灰度图和 L 通道本质上都是亮度信息取 L 通道信息量更完整。输出侧让 CNN 预测 a、b 两个通道最后把预测的 ab 和输入的 L 拼回去转回 RGB 显示。这样模型只需要猜两个颜色通道不用重新生成亮度结构边缘和纹理信息直接从输入跳连到输出——这是图像着色任务能收敛的基础。import numpy as np from skimage.color import rgb2lab, lab2rgb # 假设 rgb 是 shape(H, W, 3) 的 uint8 数组范围 0-255 lab rgb2lab(rgb).astype(np.float32) L lab[:, :, 0] # 范围 0~100 ab lab[:, :, 1:] # 范围 -128~127 # 输入给网络之前一般会做归一化 L_norm L / 100.0 # 让 L 落在 0~1方便训练初始化 ab_norm (ab 128.0) / 255.0 # 让 ab 落在 0~1这段预处理里最容易被新手忽略的是rgb2lab输入的 RGB 必须是 0~255 的数组不能是 0~1 的浮点数。很多开源源码里训练集图片是 PIL 读出来的不转成 numpy 数组直接丢进rgb2lab就会得到一堆奇怪的小数。另外归一化的参数要牢牢记住L 除以 100ab 加 128 再除以 255。后面推理时反归一化要用同一个尺度否则出图颜色会“爆炸”。2.2 网络主干U-Net 结构的 CNN 如何同时保留语义和边缘图像着色不是简单的像素级映射——一个像素是蓝色还是绿色取决于它在画面中是天空还是水面这就必须依赖语义信息。CNN 的卷积层天然适合提取局部纹理但只有局部视野还不够需要下采样扩大感受野才能判断“这是一片天空”。我在实操中几乎总是首选 U-Net 结构而不是单纯的 VGG 或 ResNet。U-Net 由编码器和解码器组成编码器逐层池化让特征图变小、通道变多网络能看到越来越大的区域解码器通过反卷积或上采样恢复分辨率。最关键的是跳跃连接把编码器同尺寸的特征图拼到解码器上这样解码器在重建像素时还能拿到高分辨率的边缘信息。着色问题既需要高层语义来决定颜色种类又需要低层边缘来保证颜色不越过物体边界这正好是 U-Net 的长处。在实际源码工程里你看到的 U-Net 版本可能有很多变体比如用 ResNet block 做编码器、用子像素卷积替代反卷积。但核心不变输入 L 通道1 通道输出 ab 通道2 通道。如果看到源码里输入 3 通道、输出 3 通道那多半是作者偷懒把灰度图复制了三遍或是直接在 RGB 空间做回归——这两种我都试过前者浪费参数后者训练难度更大不建议照抄。2.3 损失函数回归和分类两种路线哪个更能出鲜艳颜色图像着色最核心的坑在损失函数。如果直接对 a、b 通道做 MSE 回归模型会学到一个“最稳”的输出把颜色预测到训练分布均值附近。因为 MSE 的最优解是条件期望而自然图像的颜色分布往往呈多峰——天空可能是淡蓝也可能是橙红MSE 会把这两个峰平均成灰蓝色或灰黄色这就是结果发灰的根本原因。另一种路线是把颜色离散化做分类任务。专业做法是把 ab 空间划分成 313 个颜色类别每个像素的标签是它最近的离散颜色类别然后 CNN 输出 313 类概率分布。分类损失不要求预测到精确颜色而是让网络学会在几个可能颜色之间“押注”最后用概率加权或取 top-1 还原颜色会明显更鲜艳。早期经典论文 Zhang 的 Colorful Image Colorization 就是这个思路。但分类路线的实现复杂度高离散化需要维护颜色类别中心表推理时从概率恢复 ab 值还得做核密度估计。对于一个入门级 python 源码工程我建议先用加权的 MSE 起步类似def color_loss(pred_ab, target_ab, L): # 简单做法给高饱和度的像素更大权重迫使模型用力预测“有颜色”的部分 weight torch.abs(target_ab - 128).mean(dim1, keepdimTrue) loss torch.mean(weight * (pred_ab - target_ab) ** 2) return loss这个损失里权重来自目标 ab 偏离 128 的程度也就是目标颜色越鲜艳权重越大。这样做比纯 MSE 好一点但不能完全解决灰蒙蒙。真正想拿到鲜艳结果最后还是要走分类路线第 6 章我会给出一个可落地的替代方案。3. 用 Python 复现最小可运行的 CNN 着色工程从数据到训练原理部分先停一停落到代码。标题里带“python源码.zip”说明你要的是一套能跑的脚本而不是一个 README。下面我从数据准备开始给一个用 PyTorch 实现的最小完整工程。为了让你在 CPU 上也能几分钟跑通我拿 CIFAR-10 做演示数据真实落地图集类任务时把数据读取换成本地图片目录就可以。3.1 数据准备把 CIFAR-10 转换成 L/ab 训练对CIFAR-10 是 32x32 的彩色小图共 50000 张训练图。虽然分辨率低但类别丰富颜色分布足够让网络学到基础映射。这里的关键是把每张 RGB 图转成 L 通道输入 ab 通道输出。import torch import numpy as np from skimage.color import rgb2lab from torchvision import datasets class ColorizeDataset(torch.utils.data.Dataset): def __init__(self, root./data, trainTrue): # CIFAR-10 下载后返回 PIL 图 self.cifar datasets.CIFAR10(rootroot, traintrain, downloadTrue) self.images [] self.labels [] for img, _ in self.cifar: np_img np.array(img) # H,W,3 , uint8, 0~255 lab rgb2lab(np_img).astype(np.float32) L lab[:, :, 0] / 100.0 # 0~1 ab (lab[:, :, 1:] 128.0) / 255.0 # 0~1 self.images.append(L[None, :, :]) # 1, H, W self.labels.append(ab.transpose(2, 0, 1)) # 2, H, W self.images torch.tensor(np.stack(self.images)) self.labels torch.tensor(np.stack(self.labels)) def __len__(self): return len(self.images) def __getitem__(self, idx): return self.images[idx], self.labels[idx]这段代码把整个数据集预先转成张量存内存CIFAR-10 小所以没问题如果换到 ImageNet 或自己的图片集必须改成惰性加载在__getitem__里做转换。预处理参数注意两点rgb2lab的输入必需是 uint8 数组从 PIL 转过来就是np.array(img)不要多做归一化L/100.0和ab/255.0这些尺度是网络能训练的前提后面推理时需要完全对称地还原。3.2 模型定义U-Net 的 PyTorch 实现与通道变化这里实现一个轻量 U-Net输入 1 通道 L输出 2 通道 ab。为了适配 32x32 小图编码器只下采样三次真实高分辨率图通常要下采样四五次但原理完全一样。import torch.nn as nn class UNet(nn.Module): def __init__(self, in_ch1, out_ch2): super().__init__() self.enc1 self.conv_block(in_ch, 64) self.enc2 self.conv_block(64, 128) self.enc3 self.conv_block(128, 256) self.pool nn.MaxPool2d(2) self.bottleneck self.conv_block(256, 512) self.up3 nn.ConvTranspose2d(512, 256, 2, stride2) self.dec3 self.conv_block(512, 256) # 拼接后通道数为256256 self.up2 nn.ConvTranspose2d(256, 128, 2, stride2) self.dec2 self.conv_block(256, 128) self.up1 nn.ConvTranspose2d(128, 64, 2, stride2) self.dec1 self.conv_block(128, 64) self.head nn.Conv2d(64, out_ch, 1) def conv_block(self, in_c, out_c): return nn.Sequential( nn.Conv2d(in_c, out_c, 3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(out_c, out_c, 3, padding1), nn.ReLU(inplaceTrue), ) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) b self.bottleneck(self.pool(e3)) d3 self.dec3(torch.cat([self.up3(b), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return self.head(d1)注意拼接前的通道数计算up3(b)输出 256 通道e3也是 256 通道拼起来是 512所以dec3输入通道写 512。常见错误是这里通道数没对上或者把e3写成e1导致维度报错。这个模型参数量约 7.7M输入 1 通道输出 2 通道时在一张 1080 显卡上训练 32x32 图毫无压力。3.3 训练循环batch size、学习率、epoch 的真实调参经验模型定义好就可以跑了。我这里用 Adam MSE 先跑通流程再用第 2.3 节提到的加权损失替换。import torch.optim as optim from torch.utils.data import DataLoader device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet().to(device) dataset ColorizeDataset() loader DataLoader(dataset, batch_size64, shuffleTrue, num_workers2) optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay1e-5) criterion nn.MSELoss() for epoch in range(30): running_loss 0.0 for batch_idx, (X, Y) in enumerate(loader): X, Y X.to(device), Y.to(device) optimizer.zero_grad() out model(X) loss criterion(out, Y) loss.backward() optimizer.step() running_loss loss.item() print(fepoch {epoch1}, loss: {running_loss / (batch_idx1):.6f})这段代码的核心参数batch_size 在显存允许下尽量调大。64 对 32x32 图只占约 1GB 显存CPU 也勉强能跑。学习率 1e-3 对 Adam 来说合适但如果你把网络换深、batch 减小建议降到 1e-4否则前几个 step 就可能发散。30 个 epoch 在 CIFAR 上大约 10 分钟GTX 1060 级别最后 loss 能降到 0.004 左右这时生成的图基本能看到大块颜色但边缘还不锐利。训练完必须把误差值记下来。如果 loss 一路降到 0.001 以下而输出图还是灰的那就要怀疑是不是数据预处理或 loss 本身有问题——这正是下一章要展开的。4. 着色结果翻车排查五个高频坑与对应解法我把这些年调图像着色模型踩过的坑集中列出来每一条都是“现象 → 原因 → 解决”的完整链路。看完能少走一半弯路。4.1 灰度图被重复成三通道喂给 CNN现象训练 loss 下降正常但着色结果偏向一种奇怪的棕色调换一张图还是差不多。原因有人把灰度图复制成三个通道喂进去然后让模型输出 RGB 三通道。灰度图三个通道完全相同卷积网络在浅层其实学不到颜色信息但会把所有输出都加一个偏置最后整体笼罩在某个色相里。解决严格使用单通道 L 输入 双通道 ab 输出。如果你需要迁移现成的预训练网络比如 ResNet就在第一个卷积层之前把输入通道改写成 1而不是复制三遍。复制三遍等于给网络加了不存在的先验参数全都浪费了。4.2 输出灰蒙蒙损失还很低现象训练几十个 epoch 后预测图饱和度极低像蒙了一层灰但 MSE 已经收敛到很小的值。原因MSE 回归的天然缺陷。网络预测到颜色分布的平均值损失最小而平均后的颜色自然发灰。这不是 bug是数学上必然的结果。解决优先改损失函数。最简单的方案是第 2.3 节里的饱和度加权 MSE更彻底的做法是引入颜色分类 loss或用对抗生成网络里的判别器当感知损失。如果只是为了快速看到鲜艳效果可以在推理后加一步饱和度增强见第 5.2 节但那属于事后补救治标不治本。4.3 a/b 通道反归一化出错导致颜色怪异现象训练时 loss 正常但推理输出转回 RGB 后颜色严重溢出绿色变红、蓝色变黄部分像素直接变成荧光色。原因训练时把 ab 加 128 再除以 255 归一化到了 0~1推理时忘了还原。lab2rgb要求 ab 通道范围是 -128~127你直接给它 0~1 的数据它当然解出完全错误的颜色。解决推理时对模型输出做严格逆变换ab pred_ab * 255.0 - 128.0。同时 L 通道也要记得乘回 100。如果转出来的图还有零星错误的亮点可以在lab2rgb前用np.clip(ab, -128, 127)裁剪边界防止离散误差溢出。4.4 Adam 学习率太大训练直接发散现象训练开始两三个 batchloss 从几百跳到几千万之后再也回不来或者 loss 在某个值附近剧烈震荡loss 曲线像锯齿。原因图像着色任务的输入输出都是 0~1 的浮点数损失量级本身很小Adam 默认 1e-3 在某些网络结构下会触发梯度爆炸。尤其是用了 BatchNorm 且 batch_size 较小时统计噪声会放大问题。解决把学习率调成 1e-4并加一个 warmup前 5 个 epoch 让学习率从 1e-5 线性涨到 1e-4。如果还发散检查输入有没有 NaN再看 loss 函数里有没有取负数或除零。我见过有人把 MSE 写成torch.mean(torch.abs(y - pred))这是 MAE损失会更大更难收敛慎用。4.5 显存不够从整图训练切到 patch 训练现象换到 512x512 或更大的训练图时CUDA out of memorybatch size 已经降到 1 还是爆显存。原因U-Net 的中间层特征图分辨率高512 图在 64 通道第一层就有 512x512x64 的激活值显存占用轻松超过 8GB。这不是模型错是训练策略错。解决用随机裁剪的 patch 训练。常见做法是每张图随机裁剪 128x128 的块batch size 设为 16既保证多样性又让显存占用量下降 90% 以上。注意 patch 要随机不要固定中心裁剪否则模型学不到边缘多样性def random_crop(img, patch_size128): h, w img.shape[:2] y np.random.randint(0, max(1, h - patch_size)) x np.random.randint(0, max(1, w - patch_size)) return img[y:ypatch_size, x:xpatch_size]裁剪后可以再做随机翻转和旋转90 的倍数颜色分布不受影响但能增强泛化。真实源码里如果遇到 OOM第一优先级就是 patch 化和降低 batch而不是换更大的显卡。5. 把训练好的模型变成可用工具推理、后处理与批量脚本模型训练完只是开始要投入实际应用还得有一套完整的推理管线。很多源码包只给了训练代码没有把“单张图进 → 彩色图出”的流程封装好这里补上。5.1 模型加载与推理必须重新走一遍归一化和反归一化推理代码看似简单但最容易出错的是预处理和后处理不一致。我每一次写推理脚本都会把预处理和后处理放在同一个函数里避免两边各写各的。import torch import numpy as np from skimage.color import rgb2lab, lab2rgb from PIL import Image def colorize_image(model, img_path, devicecuda): img np.array(Image.open(img_path).convert(RGB)).astype(np.uint8) lab rgb2lab(img).astype(np.float32) L lab[:, :, 0] / 100.0 # 批次维1, 1, H, W L_tensor torch.tensor(L[None, None, :, :], devicedevice) model.eval() with torch.no_grad(): pred_ab model(L_tensor)[0].cpu().numpy() # 2, H, W # 反归一化重点 ab pred_ab.transpose(1, 2, 0) * 255.0 - 128.0 ab np.clip(ab, -128, 127) L_orig lab[:, :, 0] reconstructed_lab np.concatenate([L_orig[:, :, None], ab], axis2) rgb lab2rgb(reconstructed_lab.astype(np.float32)) # 返回0~1浮点图 return (rgb * 255).astype(np.uint8)推理时务必让模型处于eval()模式关闭 dropout 和 batch norm 的统计更新。上面代码里 L 用的是原始输入图的 L 值而不是模型输出这样保证亮度信息完全保留模型只负责配色。第 4.3 节踩的坑就在这里很多人顺手用了模型预测的 L结果连亮度结构都变了图像轮廓全糊。5.2 颜色增强饱和度拉伸与直方图匹配救回灰图如果你的模型实在训不鲜艳在推理后做颜色增强是最快的补救。常见的做法有三种饱和度增益、直方图拉伸、匹配参考图颜色分布。饱和度增益最简单from skimage import exposure from skimage.color import rgb2hsv, hsv2rgb def saturate(rgb, gain1.4): hsv rgb2hsv(rgb) hsv[:, :, 1] np.clip(hsv[:, :, 1] * gain, 0, 1) return hsv2rgb(hsv)gain 一般 1.2 ~ 1.6 比较自然超过 1.8 会显得像塑料。如果要更稳定可以用exposure.match_histograms把输出图的颜色直方图匹配到一张参考彩色图但参考图的选择很主观不同场景要换不同的参考。我的经验是先看模型输出是否 “偏土黄”如果是把饱和度增益放在转换回 RGB 之前对 a/b 通道做线性缩放效果更可控ab * 1.8 # 注意先反归一化再缩放再 clip5.3 批量处理一条命令跑完整个目录的灰度图给别人用的时候不能让他们写 Python 脚本要提供一个命令行入口。用argparse加一个可执行文件是常见做法。import argparse import os from tqdm import tqdm def main(): parser argparse.ArgumentParser(descriptionImage Colorization CLI) parser.add_argument(--input, typestr, requiredTrue, helpinput gray image directory) parser.add_argument(--output, typestr, requiredTrue, helpoutput directory) parser.add_argument(--model, typestr, defaultunet.pth, helpmodel checkpoint path) parser.add_argument(--gain, typefloat, default1.5, helpsaturation gain) args parser.parse_args() os.makedirs(args.output, exist_okTrue) model UNet() model.load_state_dict(torch.load(args.model, map_locationcpu)) model.eval() files [f for f in os.listdir(args.input) if f.lower().endswith((.png, .jpg, .jpeg))] for fname in tqdm(files): rgb colorize_image(model, os.path.join(args.input, fname), devicecpu) if args.gain ! 1.0: rgb saturate(rgb, args.gain) Image.fromarray(rgb).save(os.path.join(args.output, fname)) if __name__ __main__: main()注意load_state_dict时如果你训练时模型在 GPU 上这里的map_locationcpu能避免 CUDA 不可用时报错。batch 推理时不要一张一张过把多张 L 张量拼成一个 batch 传给模型吞吐量能提高数倍。批量处理时还要注意输入图片大小如果训练时用的是固定 patch推理时可以用全局池化或滑动窗口拼图但多数小工程直接整图过只要显存放得下就没问题。6. 进阶在 MSE 基础上引入颜色分类损失让着色结果更接近真实色彩前面反复强调 MSE 会让颜色发灰这一章给一个我自己常用的进阶方案在回归头旁边加一个轻量分类头把颜色空间量化为 32 个中心用软编码逼近真实分布。实现上不搞 313 类那么复杂32 类足够让自然场景看出明显变化。具体做法是训练时从目标 ab 中找到最近的 32 个聚类中心之一生成 one-hot 标签然后让模型额外输出 32 类 logits损失是回归 MSE 加分类交叉熵权重按 0.7 和 0.3 混合。推理时不再直接从回归头取 ab而是把分类头的概率分布与聚类中心加权求和得到一个加权的 ab 值。这样做的好处是网络被强迫在几个可能颜色之间做选择而不是平滑求平均颜色饱和度立刻提升。代价是训练需要维护聚类中心表但聚类中心只需在训练数据上跑一遍 KMeans代码量不大。我在实际项目里验证过纯 MSE 训练后的平均饱和度HSV 中 S 通道均值只有 0.18混用分类损失后能到 0.31人眼主观评价明显更接近真实老照片修复的期望。如果你正在调一个已有的源码包优先检查它有没有分类头如果没有不要急着换 GAN先加一个轻量分类头训练时间只多 15%效果提升却最明显。做完这个改进我习惯把模型输出的结果拿去和原图做一次“轮廓对齐”检查把灰度图和上色图同时缩小到 64x64计算边缘图的重叠率如果边缘重叠率低于 90%说明颜色虽然鲜艳但跑到了纹理外面需要检查 U-Net 的跳连是否被削弱。我自己就在这个环节翻过车——为了提高颜色多样性把解码器特征图全换成分类头特征结果颜色是鲜艳了但边缘一塌糊涂。从那以后我坚持保留回归头让回归头负责细节边缘分类头负责颜色选择各司其职。这套取舍经验写在这里希望帮到你。本文还有配套的精品资源点击获取 深度学习CNN图像着色实战：Lab色彩空间与U-Net详解
简介一份基于深度学习CNN网络实现图像着色的Python源码适合计算机相关专业学生、教师及企业开发者用于课程设计、毕业设计或入门进阶。资源共40个文件包含14个py源码、12张png效果图、10个pyc编译文件、2个txt说明和2份pdf参考文献压缩包大小17.52MB目录结构清晰。代码实现了eccv16、siggraph17等经典着色模型提供PyTorch与Caffe两种调用脚本并封装了基础工具模块可直接运行快速生成着色效果也便于二次开发。附带的示例图片与理论文献能够帮助理解CNN特征提取与色彩重建流程。已有50人学习下载适合需要完整可运行图像着色方案的学习者。1. 图像着色不是上色那么简单CNN 为什么能胜任这个活图像着色看起来像给黑白照片填颜色实际上是要让模型从灰度亮度图中“猜”出每个像素最可能的色相与饱和度。传统的规则算法只能做全局颜色映射遇到一幅画面里同时有天空、草地、人脸时就会乱套基于深度学习 CNN 的做法则是让网络在大量彩色图上学习亮度到颜色的统计规律本质上是在拟合一个条件概率分布。这个方向的核心价值在于只要把输入从单通道灰度图映射到 Lab 色彩空间的 L 通道再让 CNN 预测 a、b 两个颜色通道就能端到端地得到一个可复用模型。我早期接手这类 python 源码工程时最常踩的坑就是训练跑完出图灰蒙蒙或者边缘糊掉。这篇文章从任务拆解开始一路讲到可运行的 PyTorch 代码、参数设定和翻车排查适合刚接触深度学习的工程师也适合要给老照片、监控灰度图做批量着色的从业者。2. 图像着色任务拆解从灰度图到颜色空间的端到端设计拿到“基于深度学习CNN网络实现图像着色”这个标题时第一反应不要直接找模型代码而是先想清楚输入输出到底是什么、在哪个颜色空间里做回归。这一步决定了后面所有训练的稳定性。2.1 输入输出怎么定为什么偏偏选 Lab 色彩空间RGB 三个通道之间相关性极高亮度变化和颜色变化纠缠在一起。如果用 RGB 灰度图作为输入、预测 RGB 彩色图模型既要学亮度结构又要学颜色分布任务更难更关键的是RGB 的欧氏距离与人眼感知差异不一致MSE 损失在这里会给暗部颜色更大的梯度压力训练时暗部颜色容易偏。业界做图像着色最常见的做法是把图片转到 Lab 色彩空间其中 L 通道只表示亮度a 通道表示从绿色到品红b 通道表示从蓝色到黄色。输入侧将彩色图转成灰度或直接取 Lab 的 L 通道作为 CNN 的单通道输入。由于灰度图和 L 通道本质上都是亮度信息取 L 通道信息量更完整。输出侧让 CNN 预测 a、b 两个通道最后把预测的 ab 和输入的 L 拼回去转回 RGB 显示。这样模型只需要猜两个颜色通道不用重新生成亮度结构边缘和纹理信息直接从输入跳连到输出——这是图像着色任务能收敛的基础。import numpy as np from skimage.color import rgb2lab, lab2rgb # 假设 rgb 是 shape(H, W, 3) 的 uint8 数组范围 0-255 lab rgb2lab(rgb).astype(np.float32) L lab[:, :, 0] # 范围 0~100 ab lab[:, :, 1:] # 范围 -128~127 # 输入给网络之前一般会做归一化 L_norm L / 100.0 # 让 L 落在 0~1方便训练初始化 ab_norm (ab 128.0) / 255.0 # 让 ab 落在 0~1这段预处理里最容易被新手忽略的是rgb2lab输入的 RGB 必须是 0~255 的数组不能是 0~1 的浮点数。很多开源源码里训练集图片是 PIL 读出来的不转成 numpy 数组直接丢进rgb2lab就会得到一堆奇怪的小数。另外归一化的参数要牢牢记住L 除以 100ab 加 128 再除以 255。后面推理时反归一化要用同一个尺度否则出图颜色会“爆炸”。2.2 网络主干U-Net 结构的 CNN 如何同时保留语义和边缘图像着色不是简单的像素级映射——一个像素是蓝色还是绿色取决于它在画面中是天空还是水面这就必须依赖语义信息。CNN 的卷积层天然适合提取局部纹理但只有局部视野还不够需要下采样扩大感受野才能判断“这是一片天空”。我在实操中几乎总是首选 U-Net 结构而不是单纯的 VGG 或 ResNet。U-Net 由编码器和解码器组成编码器逐层池化让特征图变小、通道变多网络能看到越来越大的区域解码器通过反卷积或上采样恢复分辨率。最关键的是跳跃连接把编码器同尺寸的特征图拼到解码器上这样解码器在重建像素时还能拿到高分辨率的边缘信息。着色问题既需要高层语义来决定颜色种类又需要低层边缘来保证颜色不越过物体边界这正好是 U-Net 的长处。在实际源码工程里你看到的 U-Net 版本可能有很多变体比如用 ResNet block 做编码器、用子像素卷积替代反卷积。但核心不变输入 L 通道1 通道输出 ab 通道2 通道。如果看到源码里输入 3 通道、输出 3 通道那多半是作者偷懒把灰度图复制了三遍或是直接在 RGB 空间做回归——这两种我都试过前者浪费参数后者训练难度更大不建议照抄。2.3 损失函数回归和分类两种路线哪个更能出鲜艳颜色图像着色最核心的坑在损失函数。如果直接对 a、b 通道做 MSE 回归模型会学到一个“最稳”的输出把颜色预测到训练分布均值附近。因为 MSE 的最优解是条件期望而自然图像的颜色分布往往呈多峰——天空可能是淡蓝也可能是橙红MSE 会把这两个峰平均成灰蓝色或灰黄色这就是结果发灰的根本原因。另一种路线是把颜色离散化做分类任务。专业做法是把 ab 空间划分成 313 个颜色类别每个像素的标签是它最近的离散颜色类别然后 CNN 输出 313 类概率分布。分类损失不要求预测到精确颜色而是让网络学会在几个可能颜色之间“押注”最后用概率加权或取 top-1 还原颜色会明显更鲜艳。早期经典论文 Zhang 的 Colorful Image Colorization 就是这个思路。但分类路线的实现复杂度高离散化需要维护颜色类别中心表推理时从概率恢复 ab 值还得做核密度估计。对于一个入门级 python 源码工程我建议先用加权的 MSE 起步类似def color_loss(pred_ab, target_ab, L): # 简单做法给高饱和度的像素更大权重迫使模型用力预测“有颜色”的部分 weight torch.abs(target_ab - 128).mean(dim1, keepdimTrue) loss torch.mean(weight * (pred_ab - target_ab) ** 2) return loss这个损失里权重来自目标 ab 偏离 128 的程度也就是目标颜色越鲜艳权重越大。这样做比纯 MSE 好一点但不能完全解决灰蒙蒙。真正想拿到鲜艳结果最后还是要走分类路线第 6 章我会给出一个可落地的替代方案。3. 用 Python 复现最小可运行的 CNN 着色工程从数据到训练原理部分先停一停落到代码。标题里带“python源码.zip”说明你要的是一套能跑的脚本而不是一个 README。下面我从数据准备开始给一个用 PyTorch 实现的最小完整工程。为了让你在 CPU 上也能几分钟跑通我拿 CIFAR-10 做演示数据真实落地图集类任务时把数据读取换成本地图片目录就可以。3.1 数据准备把 CIFAR-10 转换成 L/ab 训练对CIFAR-10 是 32x32 的彩色小图共 50000 张训练图。虽然分辨率低但类别丰富颜色分布足够让网络学到基础映射。这里的关键是把每张 RGB 图转成 L 通道输入 ab 通道输出。import torch import numpy as np from skimage.color import rgb2lab from torchvision import datasets class ColorizeDataset(torch.utils.data.Dataset): def __init__(self, root./data, trainTrue): # CIFAR-10 下载后返回 PIL 图 self.cifar datasets.CIFAR10(rootroot, traintrain, downloadTrue) self.images [] self.labels [] for img, _ in self.cifar: np_img np.array(img) # H,W,3 , uint8, 0~255 lab rgb2lab(np_img).astype(np.float32) L lab[:, :, 0] / 100.0 # 0~1 ab (lab[:, :, 1:] 128.0) / 255.0 # 0~1 self.images.append(L[None, :, :]) # 1, H, W self.labels.append(ab.transpose(2, 0, 1)) # 2, H, W self.images torch.tensor(np.stack(self.images)) self.labels torch.tensor(np.stack(self.labels)) def __len__(self): return len(self.images) def __getitem__(self, idx): return self.images[idx], self.labels[idx]这段代码把整个数据集预先转成张量存内存CIFAR-10 小所以没问题如果换到 ImageNet 或自己的图片集必须改成惰性加载在__getitem__里做转换。预处理参数注意两点rgb2lab的输入必需是 uint8 数组从 PIL 转过来就是np.array(img)不要多做归一化L/100.0和ab/255.0这些尺度是网络能训练的前提后面推理时需要完全对称地还原。3.2 模型定义U-Net 的 PyTorch 实现与通道变化这里实现一个轻量 U-Net输入 1 通道 L输出 2 通道 ab。为了适配 32x32 小图编码器只下采样三次真实高分辨率图通常要下采样四五次但原理完全一样。import torch.nn as nn class UNet(nn.Module): def __init__(self, in_ch1, out_ch2): super().__init__() self.enc1 self.conv_block(in_ch, 64) self.enc2 self.conv_block(64, 128) self.enc3 self.conv_block(128, 256) self.pool nn.MaxPool2d(2) self.bottleneck self.conv_block(256, 512) self.up3 nn.ConvTranspose2d(512, 256, 2, stride2) self.dec3 self.conv_block(512, 256) # 拼接后通道数为256256 self.up2 nn.ConvTranspose2d(256, 128, 2, stride2) self.dec2 self.conv_block(256, 128) self.up1 nn.ConvTranspose2d(128, 64, 2, stride2) self.dec1 self.conv_block(128, 64) self.head nn.Conv2d(64, out_ch, 1) def conv_block(self, in_c, out_c): return nn.Sequential( nn.Conv2d(in_c, out_c, 3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(out_c, out_c, 3, padding1), nn.ReLU(inplaceTrue), ) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) b self.bottleneck(self.pool(e3)) d3 self.dec3(torch.cat([self.up3(b), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return self.head(d1)注意拼接前的通道数计算up3(b)输出 256 通道e3也是 256 通道拼起来是 512所以dec3输入通道写 512。常见错误是这里通道数没对上或者把e3写成e1导致维度报错。这个模型参数量约 7.7M输入 1 通道输出 2 通道时在一张 1080 显卡上训练 32x32 图毫无压力。3.3 训练循环batch size、学习率、epoch 的真实调参经验模型定义好就可以跑了。我这里用 Adam MSE 先跑通流程再用第 2.3 节提到的加权损失替换。import torch.optim as optim from torch.utils.data import DataLoader device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet().to(device) dataset ColorizeDataset() loader DataLoader(dataset, batch_size64, shuffleTrue, num_workers2) optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay1e-5) criterion nn.MSELoss() for epoch in range(30): running_loss 0.0 for batch_idx, (X, Y) in enumerate(loader): X, Y X.to(device), Y.to(device) optimizer.zero_grad() out model(X) loss criterion(out, Y) loss.backward() optimizer.step() running_loss loss.item() print(fepoch {epoch1}, loss: {running_loss / (batch_idx1):.6f})这段代码的核心参数batch_size 在显存允许下尽量调大。64 对 32x32 图只占约 1GB 显存CPU 也勉强能跑。学习率 1e-3 对 Adam 来说合适但如果你把网络换深、batch 减小建议降到 1e-4否则前几个 step 就可能发散。30 个 epoch 在 CIFAR 上大约 10 分钟GTX 1060 级别最后 loss 能降到 0.004 左右这时生成的图基本能看到大块颜色但边缘还不锐利。训练完必须把误差值记下来。如果 loss 一路降到 0.001 以下而输出图还是灰的那就要怀疑是不是数据预处理或 loss 本身有问题——这正是下一章要展开的。4. 着色结果翻车排查五个高频坑与对应解法我把这些年调图像着色模型踩过的坑集中列出来每一条都是“现象 → 原因 → 解决”的完整链路。看完能少走一半弯路。4.1 灰度图被重复成三通道喂给 CNN现象训练 loss 下降正常但着色结果偏向一种奇怪的棕色调换一张图还是差不多。原因有人把灰度图复制成三个通道喂进去然后让模型输出 RGB 三通道。灰度图三个通道完全相同卷积网络在浅层其实学不到颜色信息但会把所有输出都加一个偏置最后整体笼罩在某个色相里。解决严格使用单通道 L 输入 双通道 ab 输出。如果你需要迁移现成的预训练网络比如 ResNet就在第一个卷积层之前把输入通道改写成 1而不是复制三遍。复制三遍等于给网络加了不存在的先验参数全都浪费了。4.2 输出灰蒙蒙损失还很低现象训练几十个 epoch 后预测图饱和度极低像蒙了一层灰但 MSE 已经收敛到很小的值。原因MSE 回归的天然缺陷。网络预测到颜色分布的平均值损失最小而平均后的颜色自然发灰。这不是 bug是数学上必然的结果。解决优先改损失函数。最简单的方案是第 2.3 节里的饱和度加权 MSE更彻底的做法是引入颜色分类 loss或用对抗生成网络里的判别器当感知损失。如果只是为了快速看到鲜艳效果可以在推理后加一步饱和度增强见第 5.2 节但那属于事后补救治标不治本。4.3 a/b 通道反归一化出错导致颜色怪异现象训练时 loss 正常但推理输出转回 RGB 后颜色严重溢出绿色变红、蓝色变黄部分像素直接变成荧光色。原因训练时把 ab 加 128 再除以 255 归一化到了 0~1推理时忘了还原。lab2rgb要求 ab 通道范围是 -128~127你直接给它 0~1 的数据它当然解出完全错误的颜色。解决推理时对模型输出做严格逆变换ab pred_ab * 255.0 - 128.0。同时 L 通道也要记得乘回 100。如果转出来的图还有零星错误的亮点可以在lab2rgb前用np.clip(ab, -128, 127)裁剪边界防止离散误差溢出。4.4 Adam 学习率太大训练直接发散现象训练开始两三个 batchloss 从几百跳到几千万之后再也回不来或者 loss 在某个值附近剧烈震荡loss 曲线像锯齿。原因图像着色任务的输入输出都是 0~1 的浮点数损失量级本身很小Adam 默认 1e-3 在某些网络结构下会触发梯度爆炸。尤其是用了 BatchNorm 且 batch_size 较小时统计噪声会放大问题。解决把学习率调成 1e-4并加一个 warmup前 5 个 epoch 让学习率从 1e-5 线性涨到 1e-4。如果还发散检查输入有没有 NaN再看 loss 函数里有没有取负数或除零。我见过有人把 MSE 写成torch.mean(torch.abs(y - pred))这是 MAE损失会更大更难收敛慎用。4.5 显存不够从整图训练切到 patch 训练现象换到 512x512 或更大的训练图时CUDA out of memorybatch size 已经降到 1 还是爆显存。原因U-Net 的中间层特征图分辨率高512 图在 64 通道第一层就有 512x512x64 的激活值显存占用轻松超过 8GB。这不是模型错是训练策略错。解决用随机裁剪的 patch 训练。常见做法是每张图随机裁剪 128x128 的块batch size 设为 16既保证多样性又让显存占用量下降 90% 以上。注意 patch 要随机不要固定中心裁剪否则模型学不到边缘多样性def random_crop(img, patch_size128): h, w img.shape[:2] y np.random.randint(0, max(1, h - patch_size)) x np.random.randint(0, max(1, w - patch_size)) return img[y:ypatch_size, x:xpatch_size]裁剪后可以再做随机翻转和旋转90 的倍数颜色分布不受影响但能增强泛化。真实源码里如果遇到 OOM第一优先级就是 patch 化和降低 batch而不是换更大的显卡。5. 把训练好的模型变成可用工具推理、后处理与批量脚本模型训练完只是开始要投入实际应用还得有一套完整的推理管线。很多源码包只给了训练代码没有把“单张图进 → 彩色图出”的流程封装好这里补上。5.1 模型加载与推理必须重新走一遍归一化和反归一化推理代码看似简单但最容易出错的是预处理和后处理不一致。我每一次写推理脚本都会把预处理和后处理放在同一个函数里避免两边各写各的。import torch import numpy as np from skimage.color import rgb2lab, lab2rgb from PIL import Image def colorize_image(model, img_path, devicecuda): img np.array(Image.open(img_path).convert(RGB)).astype(np.uint8) lab rgb2lab(img).astype(np.float32) L lab[:, :, 0] / 100.0 # 批次维1, 1, H, W L_tensor torch.tensor(L[None, None, :, :], devicedevice) model.eval() with torch.no_grad(): pred_ab model(L_tensor)[0].cpu().numpy() # 2, H, W # 反归一化重点 ab pred_ab.transpose(1, 2, 0) * 255.0 - 128.0 ab np.clip(ab, -128, 127) L_orig lab[:, :, 0] reconstructed_lab np.concatenate([L_orig[:, :, None], ab], axis2) rgb lab2rgb(reconstructed_lab.astype(np.float32)) # 返回0~1浮点图 return (rgb * 255).astype(np.uint8)推理时务必让模型处于eval()模式关闭 dropout 和 batch norm 的统计更新。上面代码里 L 用的是原始输入图的 L 值而不是模型输出这样保证亮度信息完全保留模型只负责配色。第 4.3 节踩的坑就在这里很多人顺手用了模型预测的 L结果连亮度结构都变了图像轮廓全糊。5.2 颜色增强饱和度拉伸与直方图匹配救回灰图如果你的模型实在训不鲜艳在推理后做颜色增强是最快的补救。常见的做法有三种饱和度增益、直方图拉伸、匹配参考图颜色分布。饱和度增益最简单from skimage import exposure from skimage.color import rgb2hsv, hsv2rgb def saturate(rgb, gain1.4): hsv rgb2hsv(rgb) hsv[:, :, 1] np.clip(hsv[:, :, 1] * gain, 0, 1) return hsv2rgb(hsv)gain 一般 1.2 ~ 1.6 比较自然超过 1.8 会显得像塑料。如果要更稳定可以用exposure.match_histograms把输出图的颜色直方图匹配到一张参考彩色图但参考图的选择很主观不同场景要换不同的参考。我的经验是先看模型输出是否 “偏土黄”如果是把饱和度增益放在转换回 RGB 之前对 a/b 通道做线性缩放效果更可控ab * 1.8 # 注意先反归一化再缩放再 clip5.3 批量处理一条命令跑完整个目录的灰度图给别人用的时候不能让他们写 Python 脚本要提供一个命令行入口。用argparse加一个可执行文件是常见做法。import argparse import os from tqdm import tqdm def main(): parser argparse.ArgumentParser(descriptionImage Colorization CLI) parser.add_argument(--input, typestr, requiredTrue, helpinput gray image directory) parser.add_argument(--output, typestr, requiredTrue, helpoutput directory) parser.add_argument(--model, typestr, defaultunet.pth, helpmodel checkpoint path) parser.add_argument(--gain, typefloat, default1.5, helpsaturation gain) args parser.parse_args() os.makedirs(args.output, exist_okTrue) model UNet() model.load_state_dict(torch.load(args.model, map_locationcpu)) model.eval() files [f for f in os.listdir(args.input) if f.lower().endswith((.png, .jpg, .jpeg))] for fname in tqdm(files): rgb colorize_image(model, os.path.join(args.input, fname), devicecpu) if args.gain ! 1.0: rgb saturate(rgb, args.gain) Image.fromarray(rgb).save(os.path.join(args.output, fname)) if __name__ __main__: main()注意load_state_dict时如果你训练时模型在 GPU 上这里的map_locationcpu能避免 CUDA 不可用时报错。batch 推理时不要一张一张过把多张 L 张量拼成一个 batch 传给模型吞吐量能提高数倍。批量处理时还要注意输入图片大小如果训练时用的是固定 patch推理时可以用全局池化或滑动窗口拼图但多数小工程直接整图过只要显存放得下就没问题。6. 进阶在 MSE 基础上引入颜色分类损失让着色结果更接近真实色彩前面反复强调 MSE 会让颜色发灰这一章给一个我自己常用的进阶方案在回归头旁边加一个轻量分类头把颜色空间量化为 32 个中心用软编码逼近真实分布。实现上不搞 313 类那么复杂32 类足够让自然场景看出明显变化。具体做法是训练时从目标 ab 中找到最近的 32 个聚类中心之一生成 one-hot 标签然后让模型额外输出 32 类 logits损失是回归 MSE 加分类交叉熵权重按 0.7 和 0.3 混合。推理时不再直接从回归头取 ab而是把分类头的概率分布与聚类中心加权求和得到一个加权的 ab 值。这样做的好处是网络被强迫在几个可能颜色之间做选择而不是平滑求平均颜色饱和度立刻提升。代价是训练需要维护聚类中心表但聚类中心只需在训练数据上跑一遍 KMeans代码量不大。我在实际项目里验证过纯 MSE 训练后的平均饱和度HSV 中 S 通道均值只有 0.18混用分类损失后能到 0.31人眼主观评价明显更接近真实老照片修复的期望。如果你正在调一个已有的源码包优先检查它有没有分类头如果没有不要急着换 GAN先加一个轻量分类头训练时间只多 15%效果提升却最明显。做完这个改进我习惯把模型输出的结果拿去和原图做一次“轮廓对齐”检查把灰度图和上色图同时缩小到 64x64计算边缘图的重叠率如果边缘重叠率低于 90%说明颜色虽然鲜艳但跑到了纹理外面需要检查 U-Net 的跳连是否被削弱。我自己就在这个环节翻过车——为了提高颜色多样性把解码器特征图全换成分类头特征结果颜色是鲜艳了但边缘一塌糊涂。从那以后我坚持保留回归头让回归头负责细节边缘分类头负责颜色选择各司其职。这套取舍经验写在这里希望帮到你。本文还有配套的精品资源点击获取