ARTICLE DETAIL

资讯详情

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

基于PyTorch的自动着色:Lab色彩空间与分类损失实战

基于PyTorch的自动着色:Lab色彩空间与分类损失实战 简介这份资源面向希望上手图像自动着色与深度先验学习的 Python 开发者与计算机视觉学习者核心是两套预训练着色模型eccv16 与 siggraph17的推理代码可对黑白照片实时上色并支持用户引导式着色。包内共 23 个文件以 py 脚本与 pyc 缓存为主辅以 jpg、jpeg、png 示例图片及 license、md、txt 等说明文档压缩包约 4.47MB体量轻便便于快速部署与二次开发。已有 680 人学习下载。读者可借助 demo_release.py 直接运行命令行着色流程理解从 RGB 转 Lab、缩放至 256×256、着色后与原分辨率拼接再转回 RGB 的完整预处理与后处理链路colorizers 模块中的 base_color、eccv16、siggraph17 与 util 等文件则展示了模型加载与推理封装方式配合 imgs 与 imgs_out 中的输入输出样例可对照验证着色效果适合作为深度着色方向的入门实践与实验基线。1. 黑白照片自动上色从“能跑”到“敢用”之间差了什么手里有一批老照片或者从档案库里扒出来一堆灰度扫描件想批量变成彩色——这个需求在修复、电商老图翻新、影视素材预处理里非常常见。使用深度神经网络的自动着色说白了就是让模型学会“看到灰度像素猜出对应的 a、b 色度值”把 L 通道保留、ab 通道预测出来再拼回 RGB。它解决的不是“艺术创作”而是批量化、可复现、色彩不飘的上色问题。适合两类人一类是刚配好vscode python 环境、想找一个完整python 代码练手的另一类是有实际修复需求、想评估这套方案到底能不能替代手工上色的从业者。彩色图像着色这件事模型选型、色彩空间、损失函数三处没对齐出来的结果就是“灰蒙蒙”或者“满脸橘色”这篇把这条链路拆开讲。2. 彩色图像着色的技术底座为什么是 Lab 而不是 RGB2.1 从“预测像素”到“预测色度”的范式转换早期做法是直接让网络输出 RGB 三通道结果普遍发灰。原因很直接RGB 三个通道高度耦合亮度信息混在三个通道里网络要同时学“这里多亮”和“这里什么颜色”任务太重。主流方案改用Lab 色彩空间把图像拆成 L亮度和 a、b两个色度轴。输入只给 L网络只负责预测 a、b亮度信息原样保留。这样任务从“重建整张图”降级成“预测两个通道”收敛快、颜色也更干净。这个思路在基于深度学习卷积神经网络的上色论文里基本是共识。网络结构上编码器负责把 L 压成特征解码器再上采样回原尺寸中间常接 skip connection 保留边缘。也有用 U-Net 变体的本质一样输入单通道输出双通道。2.2 损失函数L1 管结构分类管颜色只用 L1 或 L2 损失模型会倾向于输出“平均色”——也就是灰褐色因为它在像素级误差上最安全。解决办法是把 ab 空间量化成 313 个色块bin把回归问题转成分类问题用交叉熵训练再配合 L1 做结构约束。这就是经典实现里的做法criterion CrossEntropyLoss() λ * L1Loss()。参数上λ 一般取 0.5 到 1.0。太小颜色饱和但结构糊太大结构清晰但颜色灰。我一般从 0.5 起步看验证集的上色饱和度再调。2.3 数据准备ImageNet 就够但要注意灰度化方式训练数据不需要专门的“灰度-彩色”配对任意彩色图都能用读进来转 Lab取 L 当输入ab 当标签。常见做法是用 ImageNet 或 COCO 的子集。这里有个容易翻车的点灰度化必须用 Lab 的 L 通道不能用cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)后者是加权平均和 Lab 的 L 不完全等价会导致训练和推理时的输入分布不一致。import cv2 import numpy as np def load_image(path, size256): img cv2.imread(path) # BGR img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (size, size)) lab cv2.cvtColor(img, cv2.COLOR_RGB2LAB) # 转 Lab L lab[:, :, 0] # 亮度通道 ab lab[:, :, 1:] # 色度通道 L L.astype(np.float32) / 255.0 # 归一化到 [0,1] ab ab.astype(np.float32) / 128.0 # ab 范围约 [-128,127] return L, ab这段代码的关键在最后两行归一化L 除以 255 落到 [0,1]ab 除以 128 落到约 [-1,1]这是大多数实现的标准做法。size256是训练分辨率推理时可以改大但要注意网络是全卷积的尺寸不固定。如果显存吃紧降到 128 也能跑只是细节会丢。3. 用 PyTorch 搭一个能跑的上色网络结构、训练、推理3.1 网络结构编码器下采样 解码器上采样下面是一个可以直接跑的最小实现编码器用几层卷积加 stride 下采样解码器用转置卷积上采样最后输出 313 维的分类 logits。import torch import torch.nn as nn class ColorNet(nn.Module): def __init__(self, num_bins313): super().__init__() # 编码器输入 1 通道 L逐层下采样 self.encoder nn.Sequential( nn.Conv2d(1, 64, 3, stride2, padding1), nn.ReLU(), nn.Conv2d(64, 128, 3, stride2, padding1), nn.ReLU(), nn.Conv2d(128, 256, 3, stride2, padding1), nn.ReLU(), ) # 解码器上采样回原尺寸 self.decoder nn.Sequential( nn.ConvTranspose2d(256, 128, 3, stride2, padding1, output_padding1), nn.ReLU(), nn.ConvTranspose2d(128, 64, 3, stride2, padding1, output_padding1), nn.ReLU(), nn.ConvTranspose2d(64, num_bins, 3, stride2, padding1, output_padding1), ) def forward(self, x): feat self.encoder(x) out self.decoder(feat) # [B, 313, H, W] return outnum_bins313是 ab 空间量化后的色块数这个数字来自经典论文的量化方案不要随意改改了要和量化表对应。output_padding1是为了让上采样后的尺寸和下采样前对齐不加的话输出会比输入小一圈拼接时会报错。3.2 训练循环分类损失 L1 结构约束from torch.utils.data import DataLoader, Dataset class ColorDataset(Dataset): def __init__(self, paths, size256): self.paths paths self.size size def __len__(self): return len(self.paths) def __getitem__(self, idx): L, ab load_image(self.paths[idx], self.size) return torch.from_numpy(L).unsqueeze(0), torch.from_numpy(ab) def train(model, loader, epochs10, lr1e-4): device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) opt torch.optim.Adam(model.parameters(), lrlr) ce nn.CrossEntropyLoss() l1 nn.L1Loss() for epoch in range(epochs): for L, ab in loader: L, ab L.to(device), ab.to(device) logits model(L) # [B, 313, H, W] # 把 ab 量化成 bin 索引作为分类标签 target quantize_ab(ab).long().to(device) loss ce(logits, target) 0.5 * l1(logits.softmax(1), ab.permute(0,3,1,2)) opt.zero_grad() loss.backward() opt.step()quantize_ab是把连续的 ab 值映射到 313 个 bin 的函数实现时用最近邻查表即可。损失里ce管颜色分类l1管结构系数 0.5 是经验值。lr1e-4配 Adam 比较稳太大容易颜色崩太小收敛慢。batch size 视显存定256 分辨率下 16 到 32 都行。3.3 推理与回填把预测的 ab 拼回 RGBdef colorize(model, gray_path, size256): model.eval() L, _ load_image(gray_path, size) with torch.no_grad(): logits model(torch.from_numpy(L).unsqueeze(0).unsqueeze(0)) ab decode_ab(logits) # 从 313 bin 还原成 ab 两通道 lab np.concatenate([L[:, :, None] * 255, ab], axis2).astype(np.uint8) rgb cv2.cvtColor(lab, cv2.COLOR_LAB2RGB) return rgbdecode_ab是量化的逆过程取 logits 的 argmax 得到 bin 索引再查表还原成 ab 值。注意 L 要乘回 255 再和 ab 拼接因为 OpenCV 的 Lab 转换要求 L 在 [0,255]、ab 在 [0,255] 偏移后的范围。这一步范围搞错出来的图会偏色或者全黑是新手最常见的翻车点。4. 上色效果调优参数、后处理与批量落地4.1 三个必调参数温度、饱和度、分辨率推理阶段有三个参数直接决定观感。温度temperature作用在 softmax 上温度低颜色更确定、更饱和温度高更柔和但可能发灰一般 0.5 到 1.0。饱和度增益是在 ab 上乘一个系数1.0 是原样1.2 到 1.5 能让颜色更“跳”但超过 1.5 容易溢出。推理分辨率比训练分辨率大时细节更好但网络没见过大尺寸边缘可能出伪影常见做法是推理用 256 或 512再和原图做引导滤波融合。参数常用范围调大效果调小效果温度0.5 ~ 1.0颜色柔和、偏灰颜色饱和、对比强饱和度增益1.0 ~ 1.5颜色鲜艳颜色寡淡推理分辨率256 / 512细节多、伪影风险高稳定但糊4.2 后处理引导滤波保住边缘网络输出的 ab 是低分辨率的直接上采样会有颜色溢出到边缘外。常见做法是用原图的 L 通道做引导图对 ab 做引导滤波guided filter让颜色跟着边缘走。OpenCV 的ximgproc.guidedFilter可以直接用半径 8、eps 1e-2 是常用起点。这一步做完头发丝、树枝这类细结构的颜色会干净很多。4.3 批量处理与工程化注意批量跑的时候别一张张读、一张张推理用 DataLoader 或者自己写个队列把 IO 和计算重叠起来。另外输出格式统一成 sRGB如果下游是印刷或视频还要考虑色彩管理。文件命名建议保留原文件名加后缀方便回溯。如果要做成服务把模型加载一次常驻显存别每次请求都重新 load那个开销比推理本身还大。5. 避坑与排查上色翻车的五个真实场景5.1 输出全灰或全褐现象推理出来的图几乎没颜色或者整体偏土黄。原因损失函数里 L1 权重太大模型退化成输出平均色或者训练轮次不够分类头还没学会。解决把 L1 系数降到 0.1 到 0.5增加训练轮次检查量化表是否正确加载。5.2 颜色溢出到背景现象人物边缘、物体轮廓外出现色块。原因ab 上采样时用了双线性插值没有边缘约束。解决加引导滤波后处理或者网络里加 skip connection 把浅层特征引到解码器。5.3 推理尺寸和训练不一致导致报错现象换了大图推理拼接时维度对不上。原因网络下采样和上采样倍率不匹配output_padding没设对。解决确保输入尺寸是 2 的整数次幂倍或者推理前 resize 到训练尺寸再放大。5.4 显存爆了现象训练到一半 OOM。原因batch size 太大或者 313 维输出在 256 分辨率下显存占用高。解决降 batch size用混合精度训练torch.cuda.amp或者把分辨率降到 128。5.5 颜色在不同图片间跳变现象同一批图有的偏蓝有的偏黄。原因推理时没有固定随机种子或者温度参数在变。解决推理阶段关掉 dropout、固定温度确保每次前向一致。6. 把上色接进实际工作流一个可复用的验证习惯模型训完不是终点怎么判断它“能用”才是。我一般会固定一组回归测试图人像、风景、建筑、老照片各几张每次改完参数都跑一遍肉眼对比。比肉眼更靠谱的是算PSNR 和 SSIM但这两个指标对上色任务参考有限因为颜色预测本来就有多解。更实用的是看颜色分布的直方图和原图对比如果 ab 直方图整体偏移说明模型有系统性偏色。另一个习惯是保留中间结果L 通道、预测 ab、回填后的 RGB 都存下来。出问题时能快速定位是输入的问题、网络的问题还是后处理的问题。这个“后悔药”机制在调参阶段能省大量时间。如果要把这套东西产品化建议把模型导出成 ONNX推理端用 onnxruntimeCPU 上也能跑到可接受的速度。批量任务用队列加多进程别用单线程硬扛。最后一句血泪经验上色模型的泛化能力高度依赖训练数据的分布如果你的目标场景是特定年代的老照片最好拿一批同年代的图做微调否则再好的通用模型也会在肤色和天空色上翻车。希望帮到你。本文还有配套的精品资源点击获取
返回列表