基于深度学习的黑白照片自动着色:PyTorch代码包实战与避坑指南

发布时间:2026/10/5 5:39:56
基于深度学习的黑白照片自动着色:PyTorch代码包实战与避坑指南 简介这份资源面向希望上手图像自动着色的 Python 开发者与深度学习爱好者基于深度神经网络实现实时用户引导着色核心在于学习深度先验可将灰度照片还原为自然彩色图像。包内共 23 个文件以 6 个 py 脚本和 5 个 pyc 编译文件为主另含 jpg、jpeg、png 示例图片及 license、md、txt 等说明文档压缩包约 4.47MB结构紧凑便于快速运行与二次开发。代码提供 eccv16 与 siggraph17 两套预训练着色器并包含 Lab 空间转换、256x256 缩放、着色后与原始全分辨率拼接再转回 RGB 的完整预处理与后处理流程demo_release.py 可直接对 imgs 目录图片推理输出结果与 imgs_out 中的示例图对应。目前已有 680 人学习下载适合想理解着色模型加载、推理管线与效果验证的读者参考。1. 从黑白到彩色一份能直接跑的深度着色代码包翻到一张上世纪的黑白老照片想给它上色第一反应可能是丢进 Photoshop 手动涂但一张图能耗掉半小时。这份colorization-master.zip解决的就是这件事——用深度神经网络自动给灰度图着色输入一张黑白 JPEG输出一张彩色 PNG。它基于 ECCV16 和 SIGGRAPH17 两篇论文的官方实现作者是 Richard Zhang代码是 PyTorch 版本不是那个已经废弃的 Caffe 分支。包里带了预训练权重、示例图片、demo_release.py入口脚本以及colorizers目录下的模型定义。适合谁想快速验证着色效果的产品经理、需要批量处理老照片的开发者、想拿它当 baseline 做二次开发的研究生。不需要你从头训练加载权重就能出图这是它最大的价值。2. 环境搭建与模型加载把权重跑起来2.1 依赖安装与版本对齐拿到包先别急着python demo_release.py环境不对会直接报错。requirements.txt里列了核心依赖但没锁死版本这是第一个要留意的地方。我一般会建一个干净的虚拟环境避免和系统里的 PyTorch 冲突。# 创建虚拟环境Python 3.8 比较稳3.10 以上有些旧版 torchvision 会出问题 python -m venv color_env source color_env/bin/activate # Windows 用 color_env\Scripts\activate # 安装依赖torch 和 torchvision 建议单独指定版本 pip install torch1.13.1 torchvision0.14.1 pip install numpy opencv-python scikit-image逻辑说明torch和torchvision是模型推理的底座numpy做数组运算opencv-python负责图像读写和色彩空间转换scikit-image在部分后处理里会用到。参数上torch1.13.1是我实测比较稳的版本再新的版本在加载旧权重时偶尔会报state_dict键名不匹配。如果你用 GPU装对应 CUDA 版本的 torch 即可代码里会自动检测。提示requirements.txt里可能只写了torch没写版本直接pip install -r requirements.txt会拉最新版反而容易翻车。建议先看一遍文件内容再决定。2.2 加载 ECCV16 与 SIGGRAPH17 两个模型包里的colorizers目录下有两个模型文件eccv16.py和siggraph17.py分别对应两篇论文的网络结构。base_color.py是它们的基类util.py放了预处理和后处理函数。加载方式在摘要里已经给了但实际用的时候要注意.eval()和 GPU 迁移。import torch import colorizers # 加载 ECCV16 模型这是 2016 年那版速度更快色彩偏饱和 colorizer_eccv16 colorizers.eccv16(pretrainedTrue).eval() # 加载 SIGGRAPH17 模型2017 年改进版细节更稳但推理稍慢 colorizer_siggraph17 colorizers.siggraph17(pretrainedTrue).eval() # 如果有 GPU把模型挪过去 if torch.cuda.is_available(): colorizer_eccv16 colorizer_eccv16.cuda() colorizer_siggraph17 colorizer_siggraph17.cuda()逻辑说明pretrainedTrue会触发下载或加载本地权重包里的saved_eccv16.png和saved_siggraph17.png是示例输出不是权重文件权重通常在首次运行时从网络拉取或放在指定目录。.eval()把 BatchNorm 和 Dropout 切到推理模式忘了写这个输出会带随机性。参数上两个模型输入都要求是 Lab 空间的 L 通道尺寸 256x256这个在util.py的preprocess_img里已经封装好了。2.3 跑通 demo_release.py 的最小命令摘要里给了一条命令python demo_release.py -i imgs/ansel_adams3.jpg。这条命令背后做了四件事读图、转 Lab、调尺寸、着色、拼回原分辨率、转 RGB 保存。先跑通它再改参数。# 在项目根目录执行确保 imgs 文件夹和 demo_release.py 同级 python demo_release.py -i imgs/ansel_adams3.jpg # 如果想指定输出路径和模型常见做法是加参数 python demo_release.py -i imgs/ansel_adams.jpg -o output/ansel_color.png --model siggraph17逻辑说明-i指定输入图片路径-o指定输出路径--model切换模型。如果脚本没写--model参数那就默认用 ECCV16想换模型得改代码。跑完后去imgs_out文件夹看结果和saved_eccv16.png对比一下颜色分布应该接近。参数上输入图片最好是 JPEG 或 PNG灰度图直接喂彩色图会被转成灰度再着色这点要注意。3. 预处理与后处理Lab 空间、256 尺寸和全分辨率拼接3.1 为什么必须转 Lab 空间RGB 三个通道耦合太强直接让网络预测 RGB 容易偏色。Lab 空间把亮度 L 和色度 a、b 分开网络只需要根据 L 预测 a、b任务更干净。这是着色领域的标准做法ECCV16 和 SIGGRAPH17 都沿用。import cv2 import numpy as np def preprocess_img(img_path): # 读图opencv 默认 BGR img cv2.imread(img_path) # 转 RGB 再转 Lab注意 opencv 的 Lab 范围是 0-255 img_rgb cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img_lab cv2.cvtColor(img_rgb, cv2.COLOR_RGB2Lab) # 取 L 通道归一化到 0-1 L img_lab[:, :, 0] / 255.0 return L, img_rgb逻辑说明cv2.COLOR_RGB2Lab输出的 L 范围是 0-255除以 255 归一化。a、b 通道范围是 -128 到 127训练时网络输出会做相应缩放。参数上如果你用 PIL 读图色彩空间转换结果和 opencv 略有差异建议统一用 opencv。3.2 缩放到 256x256 的取舍网络输入固定 256x256但原图可能很大。直接缩放会丢细节所以流程是先缩到 256 着色得到低分辨率 a、b再上采样回原尺寸和原图 L 通道拼接。这样既保证网络输入合规又保留原图亮度细节。def resize_and_colorize(L, colorizer, devicecuda): # L 是 numpy转 tensor 并加 batch 维度 L_tensor torch.from_numpy(L).unsqueeze(0).unsqueeze(0).float() if device cuda: L_tensor L_tensor.cuda() # 缩放到 256x256 L_resized torch.nn.functional.interpolate( L_tensor, size(256, 256), modebilinear, align_cornersFalse ) # 着色输出 ab 通道 with torch.no_grad(): ab colorizer(L_resized) # 上采样回原尺寸 ab_upsampled torch.nn.functional.interpolate( ab, sizeL_tensor.shape[2:], modebilinear, align_cornersFalse ) return ab_upsampled逻辑说明interpolate做双线性插值align_cornersFalse是 PyTorch 推荐设置避免边缘偏移。torch.no_grad()省显存。参数上size(256, 256)是硬性要求改成其他尺寸网络会报维度错误。上采样回原尺寸时如果原图特别大显存不够可以分块处理但一般 4K 图单张推理没问题。3.3 拼接回全分辨率并转 RGB拿到 ab 后和原图 L 拼成 Lab再转 RGB 保存。这一步容易出问题的是数值范围ab 网络输出通常在 -128 到 127 附近但可能有溢出要 clip。def postprocess(L_original, ab, img_rgb): # L_original 是 0-1ab 是网络输出先转 numpy L_np (L_original * 255).astype(np.uint8) ab_np ab.squeeze().cpu().numpy().transpose(1, 2, 0) # 拼接成 Lab lab np.concatenate([L_np[:, :, np.newaxis], ab_np], axis2) # 转 RGB rgb cv2.cvtColor(lab, cv2.COLOR_Lab2RGB) return rgb逻辑说明L_np还原到 0-255ab_np调整维度顺序。cv2.COLOR_Lab2RGB要求输入是 uint8 或 float32这里用 uint8。参数上如果输出偏灰检查 ab 是否被过度平滑可以调后处理的锐化强度但代码里没暴露这个参数得自己加。4. 避坑与排查权重、显存和色彩偏移4.1 报错 “No module named colorizers”现象在项目根目录跑demo_release.py提示找不到colorizers模块。原因Python 路径没包含当前目录或者你在别的目录执行脚本。解决cd到项目根目录再跑或者export PYTHONPATH$PYTHONPATH:.。如果还不行检查colorizers下有没有__init__.py包里是有的别删。4.2 权重下载失败或加载报错现象首次运行卡在下载或者报KeyError: state_dict。原因权重文件没下全或者 torch 版本不匹配导致键名差异。解决手动下载权重放到指定目录常见做法是放在checkpoints文件夹然后改colorizers里的加载路径。如果键名不匹配用torch.load看权重结构手动映射。4.3 显存不足 CUDA out of memory现象跑大图时爆显存。原因上采样回原尺寸时ab 通道占显存和原图尺寸平方成正比。解决先缩图到 1024 长边再着色或者用torch.cuda.empty_cache()清缓存。如果还不行改 CPU 推理慢但能跑。4.4 输出颜色偏灰或偏绿现象着色结果像蒙了一层灰或者人脸发绿。原因ab 通道预测偏差或者 Lab 转 RGB 时数值溢出没 clip。解决在postprocess里加np.clip(ab_np, -128, 127)再转 RGB。如果还偏换 SIGGRAPH17 模型试试它对肤色更稳。4.5 输入彩色图被二次着色现象喂了一张彩色图输出颜色变了。原因代码默认把输入当灰度处理彩色信息被丢弃。解决如果只想给灰度图上色确保输入是单通道如果想保留原色得改代码跳过着色直接输出原图。这个坑我踩过批量处理时混入彩色图结果全被重着色了。5. 进阶技巧批量处理与效果验证批量处理是实际项目里最常用的。包里只给了单张 demo但改起来不难。我一般会写一个循环遍历文件夹每张图跑一遍输出到指定目录。注意显存释放每张图跑完del掉中间变量。import os import glob import torch from colorizers import eccv16, siggraph17 from util import preprocess_img, postprocess # 加载模型 colorizer siggraph17(pretrainedTrue).eval() if torch.cuda.is_available(): colorizer colorizer.cuda() input_dir imgs output_dir imgs_out os.makedirs(output_dir, exist_okTrue) for img_path in glob.glob(os.path.join(input_dir, *.jpg)): L, img_rgb preprocess_img(img_path) ab resize_and_colorize(L, colorizer) rgb postprocess(L, ab, img_rgb) # 保存用原文件名加 _color 后缀 base os.path.basename(img_path).replace(.jpg, _color.png) cv2.imwrite(os.path.join(output_dir, base), cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)) torch.cuda.empty_cache()逻辑说明glob匹配所有 jpgos.makedirs确保输出目录存在。torch.cuda.empty_cache()防止显存碎片。参数上如果图片多可以加进度条tqdm但这不是必须的。验证效果时别只看一张图。我习惯把原图、ECCV16 输出、SIGGRAPH17 输出并排拼成一张对比图肉眼扫一遍。如果某张图颜色特别怪大概率是原图对比度太低L 通道信息不足网络猜不出色度。这时候可以手动调一下原图对比度再喂进去或者换模型。注意批量处理时如果图片尺寸差异大建议先统一缩放到长边 1024再跑着色最后按需放大。这样速度和显存都可控。从那以后我每次拿到新的着色模型都强制先跑三张图一张人脸、一张风景、一张低对比度老照片。人脸看肤色风景看天空和植被低对比度看网络会不会摆烂。这三张过了才敢往批量流程里塞。希望帮到你。本文还有配套的精品资源点击获取