简介这份资源面向希望上手图像自动着色、了解深度先验应用的 Python 开发者与计算机视觉学习者核心是两套预训练着色模型eccv16 与 siggraph17可在实时用户引导下为黑白照片上色。包内共 23 个文件以 py 脚本与 pyc 缓存为主辅以 jpg、jpeg、png 示例图另有 license、md 说明与 txt 依赖清单压缩包约 4.47MB结构紧凑、开箱即用。已有 680 人学习下载。读者可借助 demo_release.py 直接运行推理观察从 Lab 空间转换、256×256 缩放、着色再融合回全分辨率并转回 RGB 的完整预处理与后处理链路colorizers 模块封装了模型加载与调用方式示例图与输出结果便于对照验证。对于想复现经典着色网络、理解深度先验在着色任务中作用或以此为基线做二次开发的读者这份代码提供了清晰可读的参考实现与可运行入口。1. 彩色图像着色从灰度图到彩色图的深度神经网络方案手里有一批老照片、医学影像或者监控截图全是灰度的想批量上色又不想一张张丢进在线工具里等半天——这个场景下用深度神经网络做自动着色Colorization就是最直接的解法。它的核心思路不复杂把灰度图当作 L 通道输入让网络去预测对应的 a、b 两个色度通道拼回 Lab 空间再转 RGB就得到彩色图。整套流程用 Python 就能跑通代码量不大但选型、损失函数和数据准备这几步决定了最终效果是「能看」还是「翻车」。这篇笔记面向想自己动手复现的开发者从环境搭建、数据准备、模型搭建到训练调参和避坑一步步拆开讲新手能跟着跑熟手能看到参数边界和常见陷阱。2. 自动着色的技术选型为什么是 Lab 空间加卷积神经网络2.1 着色问题的本质是回归而非分类灰度图着色的数学本质是给定亮度 L预测色度 a 和 b。这是一个一对多的映射问题——同一张灰度图天空可以是蓝色也可以是橙色草地可以是绿色也可以是枯黄色。正因为存在多种合理答案如果用逐像素的 L2 损失去训练网络会倾向于输出所有可能颜色的均值结果就是饱和度极低、灰蒙蒙的「安全色」。常见做法是把 a、b 通道离散化成网格比如量化成 313 个色块把回归问题转成分类问题再用 softmax 输出概率分布。推理时取概率最高的色块或者做退火均值。这样训练更稳定颜色也更鲜艳。我一般会推荐这种方式尤其是数据量不大的时候。2.2 为什么选 Lab 而不是 RGB 或 HSVRGB 三个通道高度耦合直接预测 RGB 意味着网络要同时学亮度和颜色训练难度大。HSV 的色相 H 在低饱和度区域不稳定灰色像素的 H 值几乎是噪声。Lab 空间把亮度L和色度a、b解耦灰度图直接就是 L 通道网络只需要专注预测 a、b输入输出关系清晰。这是目前主流着色方案的标准做法。2.3 网络结构编码器-解码器加跳跃连接基础结构用编码器-解码器Encoder-Decoder就够了。编码器逐层下采样提取语义特征解码器逐层上采样恢复空间分辨率。中间加跳跃连接Skip Connection把浅层的高频细节传到深层避免上采样后边缘模糊。编码器可以用几层卷积加池化解码器用转置卷积或上采样加卷积。如果追求更好的效果可以把编码器换成预训练的分类网络比如 ResNet 的前几层利用 ImageNet 上学到的语义特征来指导着色——网络见过「天空」「草地」「人脸」这些概念后上色会合理得多。提示编码器用预训练权重时注意输入通道数要改成 1只接 L 通道或者把 L 复制成 3 通道再送入。前者需要改第一层卷积后者不用改结构但计算量稍大。3. 用 Python 跑通着色模型环境、数据与训练代码3.1 环境搭建与依赖安装先确认 Python 版本建议 3.8 及以上。核心依赖就几个PyTorch、NumPy、Pillow、scikit-image。安装命令如下pip install torch torchvision numpy pillow scikit-image如果要用 GPU 训练去 PyTorch 官网对照 CUDA 版本选对应的安装命令。装完后验证一下import torch print(torch.__version__) print(torch.cuda.is_available()) # 有 GPU 应返回 True逻辑说明torch.cuda.is_available()返回 False 时检查显卡驱动和 CUDA 版本是否匹配。CPU 也能跑只是训练时间会从几小时变成几天。3.2 数据准备把彩色图转成 Lab 并提取 L 通道训练数据就是一批彩色图片。用 scikit-image 做 RGB 到 Lab 的转换import numpy as np from skimage import color, io, transform def load_and_preprocess(img_path, size256): 读取彩色图转 Lab返回 L 通道和 ab 通道 rgb io.imread(img_path) # 统一尺寸 rgb transform.resize(rgb, (size, size), anti_aliasingTrue) # RGB 转 Lab注意 skimage 要求输入为 [0,1] 浮点 lab color.rgb2lab(rgb) L lab[:, :, 0] # 亮度通道范围约 [0, 100] ab lab[:, :, 1:] # 色度通道范围约 [-128, 127] # 归一化到 [-1, 1]方便网络训练 L_norm L / 50.0 - 1.0 ab_norm ab / 128.0 return L_norm, ab_norm参数说明size控制输入分辨率256×256 是常见起点显存够可以上 512。L / 50.0 - 1.0把 L 从 [0,100] 映射到 [-1,1]。ab / 128.0把色度压到 [-1,1]。归一化这步不能省否则损失值波动大收敛慢。3.3 模型定义一个最小可用的着色网络import torch import torch.nn as nn class ColorNet(nn.Module): def __init__(self): super().__init__() # 编码器输入 1 通道L逐层下采样 self.encoder nn.Sequential( nn.Conv2d(1, 64, 3, stride2, padding1), # 256 - 128 nn.ReLU(inplaceTrue), nn.Conv2d(64, 128, 3, stride2, padding1), # 128 - 64 nn.ReLU(inplaceTrue), nn.Conv2d(128, 256, 3, stride2, padding1), # 64 - 32 nn.ReLU(inplaceTrue), ) # 解码器逐层上采样输出 2 通道a, b self.decoder nn.Sequential( nn.ConvTranspose2d(256, 128, 3, stride2, padding1, output_padding1), nn.ReLU(inplaceTrue), nn.ConvTranspose2d(128, 64, 3, stride2, padding1, output_padding1), nn.ReLU(inplaceTrue), nn.ConvTranspose2d(64, 2, 3, stride2, padding1, output_padding1), nn.Tanh() # 输出范围 [-1, 1]对应归一化后的 ab ) def forward(self, L): feat self.encoder(L) ab self.decoder(feat) return ab逻辑说明编码器三层卷积把 256×256 降到 32×32解码器再升回 256×256。output_padding1是为了让输出尺寸和输入对齐。最后一层用Tanh把输出限制在 [-1,1]和前面 ab 的归一化范围一致。参数说明卷积核统一用 3×3通道数 64→128→256 是常见配置。显存不够就把通道数减半。stride2代替池化做下采样保留更多信息。3.4 训练循环与损失函数from torch.utils.data import DataLoader, Dataset import os class ColorDataset(Dataset): def __init__(self, img_dir, size256): self.files [os.path.join(img_dir, f) for f in os.listdir(img_dir) if f.lower().endswith((.jpg, .png, .jpeg))] self.size size def __len__(self): return len(self.files) def __getitem__(self, idx): L, ab load_and_preprocess(self.files[idx], self.size) # 增加通道维度: (H,W) - (1,H,W) 和 (2,H,W) L torch.tensor(L, dtypetorch.float32).unsqueeze(0) ab torch.tensor(ab, dtypetorch.float32).permute(2, 0, 1) return L, ab # 训练 device torch.device(cuda if torch.cuda.is_available() else cpu) model ColorNet().to(device) optimizer torch.optim.Adam(model.parameters(), lr1e-3) criterion nn.MSELoss() dataset ColorDataset(path/to/your/images) loader DataLoader(dataset, batch_size16, shuffleTrue, num_workers2) for epoch in range(50): total_loss 0 for L, ab in loader: L, ab L.to(device), ab.to(device) pred_ab model(L) loss criterion(pred_ab, ab) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss/len(loader):.4f})逻辑说明ColorDataset负责读图、转 Lab、归一化、转 tensor。训练循环就是标准的 PyTorch 流程。损失函数先用 MSE 跑通后面再换更好的。参数说明batch_size16是 8GB 显存的保守值显存大可加到 32 或 64。lr1e-3是 Adam 的常用起点loss 震荡就降到 1e-4。num_workers设成 CPU 核心数的一半左右太多反而拖慢。3.5 推理把预测的 ab 拼回彩色图from skimage import color def colorize(model, img_path, size256): model.eval() L_norm, _ load_and_preprocess(img_path, size) L_tensor torch.tensor(L_norm, dtypetorch.float32).unsqueeze(0).unsqueeze(0).to(device) with torch.no_grad(): ab_pred model(L_tensor).squeeze(0).permute(1, 2, 0).cpu().numpy() # 反归一化 L_orig (L_norm 1.0) * 50.0 ab_orig ab_pred * 128.0 lab np.concatenate([L_orig[:, :, None], ab_orig], axis2) rgb color.lab2rgb(lab) return (rgb * 255).astype(np.uint8)逻辑说明推理时不需要计算梯度用torch.no_grad()省显存。预测出的 ab 反归一化后和原始 L 拼成 Lab再转 RGB。注意lab2rgb输出是 [0,1] 浮点乘 255 转成 uint8 才能保存。参数说明推理的size要和训练时一致否则网络看到的分布不匹配颜色会偏。如果原图不是正方形先做中心裁剪或 padding 再 resize。4. 着色效果翻车的五个坑从灰蒙蒙到颜色溢出4.1 输出全是灰色或低饱和度现象训练 loss 降下去了但推理结果几乎看不出颜色整体灰蒙蒙。原因用 MSE 损失直接回归 ab 值时网络学到的是条件均值。对于一张图里某个像素可能是蓝也可能是绿的情况均值就是灰色。解决改用分类方案。把 ab 空间量化成 313 个色块这是常见做法网络输出每个色块的概率损失用交叉熵。推理时取概率最高的色块对应的 ab 值或者用温度参数做 softmax 退火。颜色饱和度会明显提升。4.2 颜色溢出到相邻区域现象天空的蓝色渗到了建筑物边缘或者人物衣服的颜色糊到了背景上。原因编码器下采样太狠空间信息丢失严重解码器上采样时无法精确恢复边界。另外感受野太大也会导致颜色「漏」到不该去的地方。解决加跳跃连接把编码器浅层的高分辨率特征直接拼到解码器对应层。或者用 U-Net 结构它的跳跃连接是标配。另一个办法是减小下采样倍数比如只下采样到 1/4 而不是 1/8。4.3 训练 loss 不下降或震荡剧烈现象loss 在某个值附近来回跳或者一直不降。原因学习率太大、batch size 太小、数据归一化没做对或者网络输出没有做值域限制。解决先检查归一化。L 必须在 [-1,1]ab 也必须在 [-1,1]。然后降学习率从 1e-3 降到 1e-4 试试。batch size 至少 8太小梯度噪声大。最后确认最后一层有没有 Tanh没有的话输出可能跑到几百loss 直接爆炸。4.4 推理时颜色和训练时不一致现象训练集上的图着色正常换一张新图颜色就偏了。原因推理时的预处理和训练时不一致。比如训练用了 anti_aliasing 的 resize推理用了最近邻或者训练时做了归一化推理忘了。解决把预处理封装成一个函数训练和推理都调同一个。检查 resize 方法、归一化参数、通道顺序是否完全一致。这个坑很隐蔽血泪经验是写个单元测试对比训练和推理的预处理输出。4.5 显存不够导致训练中断现象跑几个 batch 就报 CUDA out of memory。原因输入分辨率太高、batch size 太大、模型通道数太多或者没及时释放中间变量。解决先把 batch size 降到 4 试试。不够就把输入从 256 降到 128。还不够就砍通道数把 256 改成 128。另外训练循环里用del及时删掉不用的中间 tensor配合torch.cuda.empty_cache()。如果这些都不行用梯度累积模拟大 batch。5. 进阶技巧用感知损失和类别平衡把颜色做自然基础版跑通后想让颜色更自然、更符合人类审美有两个方向值得试。感知损失Perceptual Loss。MSE 只关心像素值差异不关心颜色看起来是否合理。感知损失的做法是把预测图和真实图都送进一个预训练的分类网络比如 VGG取中间某层的特征图算它们的 L2 距离。这样损失函数约束的是「语义特征要像」而不是「每个像素要一模一样」。颜色会更符合语义——天空像天空草地像草地。实现上把 VGG 的前几层冻结接在着色网络后面训练时只更新着色网络的参数。import torchvision.models as models vgg models.vgg16(pretrainedTrue).features[:16].to(device).eval() for p in vgg.parameters(): p.requires_grad False def perceptual_loss(pred_rgb, target_rgb): # pred_rgb, target_rgb: (B, 3, H, W), 范围 [0,1] feat_pred vgg(pred_rgb) feat_target vgg(target_rgb) return nn.functional.mse_loss(feat_pred, feat_target)参数说明features[:16]取的是 VGG 前几层感受野小关注局部纹理和颜色。取太深层会丢失空间信息。权重上感知损失和 MSE 按 1:10 或 1:100 混合具体看效果调。类别平衡Class Rebalancing。ab 空间量化成 313 个色块后分布极不均匀——灰色、棕色、蓝色占了大头红色、紫色、亮绿色很少。直接训练的话网络会偏向高频色块稀有颜色永远预测不出来。做法是给每个色块按出现频率的倒数加权频率越低权重越高。这样网络会被迫关注稀有颜色输出更多样化。# 假设 class_freq 是 313 个色块的频率统计 weights 1.0 / (class_freq 1e-6) weights weights / weights.sum() * len(weights) # 归一化 criterion nn.CrossEntropyLoss(weighttorch.tensor(weights, dtypetorch.float32).to(device))参数说明1e-6防止除零。归一化让权重均值为 1避免整体 loss 尺度变化太大。权重太极端会导致训练不稳定可以开根号或取对数缓和一下。验证方法除了看 loss更直观的是固定几张测试图每个 epoch 存一次着色结果拼成网格图观察颜色变化。另外可以算 PSNR 和 SSIM但这两个指标和人类感知的相关性有限只能做参考。真正靠谱的还是人眼看。我自己踩过的坑是一开始只盯着 loss 调结果 loss 很低但颜色一塌糊涂。后来养成习惯每训几个 epoch 就导出一批测试图的着色结果肉眼过一遍。颜色偏了、溢出了、灰了一眼就能看出来比看数字快得多。希望帮到你。本文还有配套的精品资源点击获取