基于CNN的图像着色：从U-Net到损失函数实战

简介这份资源面向计算机相关专业的在校学生、教师及企业员工提供一套基于深度学习CNN网络实现图像着色的完整Python源码可用于课程设计、毕业设计、大作业或初期项目立项演示。压缩包共40个文件约17.52MB包含14个py源码文件、12张png效果图、10个pyc编译文件以及txt说明和pdf参考文献源码涵盖eccv16、siggraph17等经典着色模型及工具模块结构清晰便于按模块学习。目前已有50人学习下载。读者可借此掌握CNN图像着色从数据到推理的完整流程理解模型组织与调用方式并参考效果图与文献进行二次开发或功能扩展适合入门进阶与项目借鉴。1. 从一张泛黄老照片说起CNN图像着色到底在做什么手里有一张上世纪的黑白照片人脸模糊、天空灰白、衣服看不出颜色。你想把它变成彩色但手工上色一张要几个小时还未必自然。基于深度学习CNN网络实现图像着色解决的就是这件事输入一张灰度图或L通道输出对应的a、b两个色度通道再和原L通道拼回RGB彩色图。它适合三类人想入门深度学习但不想只跑MNIST的Python开发者、需要批量处理老照片的影像从业者、以及想拿一个完整源码项目练手CNN回归任务的学生。这个方向不算新但它是极好的练手项目——输入输出都是图像损失函数直观训练过程肉眼可见地在变好。更关键的是它把CNN从“分类”拉到了“回归生成”的领域你会第一次认真思考为什么预测颜色这么难因为同一张灰度图天空可以是蓝色也可以是橙色模型必须学会“合理”而不是“唯一”。2. 图像着色的CNN架构选型从编码器-解码器到U-Net2.1 为什么着色任务不适合直接套用分类网络图像着色本质上是一个逐像素回归问题。输入是H×W×1的灰度图输出是H×W×2的ab色度图。分类网络如ResNet、VGG最后接全连接层输出类别概率这个结构对着色毫无用处——你需要的是和输入同分辨率的空间输出而不是一个类别标签。常见做法是采用编码器-解码器结构编码器用卷积池化逐步降采样提取语义特征解码器用转置卷积或上采样逐步恢复分辨率。但朴素编码器-解码器有个致命问题降采样过程中空间细节丢失上色结果边缘模糊、颜色溢出。我一般会直接上U-Net因为它的跳跃连接skip connection把编码器的高分辨率特征直接拼到解码器对应层边缘和纹理信息得以保留。对于着色任务这个改进不是锦上添花是刚需。另一个选型点是输入输出空间。不要在RGB空间直接回归因为RGB三个通道高度相关模型容易学成“灰色平均”。标准做法是转到Lab色彩空间L是亮度a和b是色度。输入只给L让模型预测a和b任务解耦更干净。这也是学术界和工业界最通用的方案。2.2 用PyTorch搭一个最小可用的着色U-Net下面这段代码是一个可以直接跑通的U-Net着色网络定义。我刻意把通道数压小方便你在单卡甚至CPU上先验证流程。import torch import torch.nn as nn import torch.nn.functional as F class ColorizationUNet(nn.Module): def __init__(self): super().__init__() # 编码器输入1通道(L)逐层升到256通道 self.enc1 self._block(1, 64) self.enc2 self._block(64, 128) self.enc3 self._block(128, 256) self.enc4 self._block(256, 256) self.pool nn.MaxPool2d(2) # 瓶颈层 self.bottleneck self._block(256, 512) # 解码器转置卷积 跳跃连接拼接 self.up4 nn.ConvTranspose2d(512, 256, 2, stride2) self.dec4 self._block(512, 256) # 256(上采样)256(skip)512 self.up3 nn.ConvTranspose2d(256, 128, 2, stride2) self.dec3 self._block(256, 128) self.up2 nn.ConvTranspose2d(128, 64, 2, stride2) self.dec2 self._block(128, 64) self.up1 nn.ConvTranspose2d(64, 32, 2, stride2) self.dec1 self._block(64, 32) # 输出层2通道对应a、b用tanh限制到[-1,1] self.out_conv nn.Conv2d(32, 2, 1) self.tanh nn.Tanh() def _block(self, in_ch, out_ch): return nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): # 编码路径 e1 self.enc1(x) # [B,64,H,W] e2 self.enc2(self.pool(e1)) # [B,128,H/2,W/2] e3 self.enc3(self.pool(e2)) # [B,256,H/4,W/4] e4 self.enc4(self.pool(e3)) # [B,256,H/8,W/8] # 瓶颈 b self.bottleneck(self.pool(e4)) # [B,512,H/16,W/16] # 解码路径 跳跃连接 d4 self.up4(b) d4 self.dec4(torch.cat([d4, e4], dim1)) d3 self.up3(d4) d3 self.dec3(torch.cat([d3, e3], dim1)) d2 self.up2(d3) d2 self.dec2(torch.cat([d2, e2], dim1)) d1 self.up1(d2) d1 self.dec1(torch.cat([d1, e1], dim1)) return self.tanh(self.out_conv(d1)) # [B,2,H,W]逻辑说明编码器每经过一个_block后做一次2倍下采样特征图尺寸减半、通道翻倍。瓶颈层在H/16分辨率上提取全局语义。解码器每一步先转置卷积放大2倍再和编码器同分辨率的特征在通道维拼接最后过卷积块融合。输出层用1×1卷积压到2通道tanh把值域限制在[-1,1]和Lab空间中a、b的归一化范围对齐。参数说明_block里的两个3×3卷积是标准配置padding1保证尺寸不变。BatchNorm加速收敛如果显存吃紧可以换成InstanceNorm。转置卷积的kernel_size2、stride2是精确2倍上采样的常用组合。输出通道数必须是2这是Lab空间决定的不要改成3。2.3 损失函数怎么选L1、L2还是分类式着色任务最常用的损失是L1损失平均绝对误差直接对预测的a、b和真实的a、b做逐像素比较。L2MSE也可以但L2对异常值更敏感容易让模型偏向“安全”的灰色。实践中L1收敛更稳颜色更饱和。但纯L1有个问题它假设每个像素的颜色是确定的而实际上同一灰度值可能对应多种合理颜色。进阶做法是把a、b空间量化成313个色块bin把回归问题转成分类问题用交叉熵训练推理时取概率最高的bin或做软编码。这个方案来自2016年的一篇经典工作效果比纯回归好但实现复杂度高不少。我的建议是先用L1跑通全流程确认数据管道和网络结构没问题再考虑升级到分类式。# L1损失直接比较预测ab和真实ab criterion nn.L1Loss() # 训练循环中的关键一行 pred_ab model(L) # [B,2,H,W] loss criterion(pred_ab, true_ab)参数说明L1Loss没有超参数直接调用即可。如果发现颜色偏灰可以尝试对损失加权比如对a、b通道分别给不同权重或者加入梯度惩罚项让边缘更锐利。3. 数据管道与训练流程从ImageNet到自定义照片3.1 数据从哪来、怎么转Lab、怎么归一化着色模型的训练数据就是彩色图像本身。你不需要标注任何彩色图都可以ImageNet、COCO、你自己的相册。流程是读入RGB图 → 转Lab → 取L作为输入、ab作为标签 → 归一化。import numpy as np from PIL import Image import torch from torch.utils.data import Dataset, DataLoader import torchvision.transforms as T class ColorizationDataset(Dataset): def __init__(self, image_paths, size256): self.paths image_paths self.size size self.resize T.Resize((size, size), interpolationT.InterpolationMode.BICUBIC) def __len__(self): return len(self.paths) def __getitem__(self, idx): img Image.open(self.paths[idx]).convert(RGB) img self.resize(img) img np.array(img).astype(np.float32) / 255.0 # RGB - Lab注意skimage的Lab范围 from skimage.color import rgb2lab lab rgb2lab(img) # L:[0,100], a:[-128,127], b:[-128,127] L lab[:, :, 0] / 50.0 - 1.0 # 归一化到[-1,1] ab lab[:, :, 1:] / 128.0 # 归一化到约[-1,1] L torch.from_numpy(L).unsqueeze(0).float() # [1,H,W] ab torch.from_numpy(ab).permute(2,0,1).float() # [2,H,W] return L, ab # 使用示例 dataset ColorizationDataset([photo1.jpg, photo2.jpg]) loader DataLoader(dataset, batch_size16, shuffleTrue, num_workers4)逻辑说明rgb2lab把RGB转到Lab空间。L通道除以50再减1映射到[-1,1]ab通道除以128也大致落在[-1,1]。这样输入输出都在同一量级训练更稳定。参数说明size256是常用训练分辨率太小丢失细节太大显存吃不消。batch_size16是8GB显存下的保守值可以按需调整。num_workers4加速数据加载Windows下如果报错就改成0。3.2 训练循环与关键超参数import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model ColorizationUNet().to(device) optimizer optim.Adam(model.parameters(), lr1e-4, betas(0.5, 0.999)) criterion nn.L1Loss() for epoch in range(50): model.train() total_loss 0 for L, ab in loader: L, ab L.to(device), ab.to(device) pred model(L) loss criterion(pred, ab) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss/len(loader):.4f})参数说明Adam的lr1e-4是着色任务的常用起点太大颜色会震荡太小收敛慢。betas(0.5, 0.999)是GAN训练中的经典设置纯回归任务用默认(0.9, 0.999)也行。训练50轮是起步实际到100轮以上颜色才比较自然。每轮打印平均损失如果损失卡在0.05左右不降说明模型容量或数据量到瓶颈了。3.3 推理把预测的ab拼回彩色图def colorize(model, img_path, device): model.eval() img Image.open(img_path).convert(RGB) img img.resize((256, 256)) arr np.array(img).astype(np.float32) / 255.0 from skimage.color import rgb2lab, lab2rgb lab rgb2lab(arr) L lab[:, :, 0] / 50.0 - 1.0 L_tensor torch.from_numpy(L).unsqueeze(0).unsqueeze(0).float().to(device) with torch.no_grad(): pred_ab model(L_tensor).cpu().numpy()[0] # [2,H,W] # 反归一化 pred_ab pred_ab.transpose(1, 2, 0) * 128.0 lab_out np.zeros((256, 256, 3)) lab_out[:, :, 0] (L 1.0) * 50.0 lab_out[:, :, 1:] pred_ab rgb_out lab2rgb(lab_out) return (rgb_out * 255).astype(np.uint8)逻辑说明推理时只取L通道送入模型得到预测的ab后反归一化和原始L拼成完整Lab再转回RGB。注意lab2rgb输出是[0,1]浮点乘255转成uint8才能保存。参数说明推理分辨率必须和训练一致否则跳跃连接的尺寸对不上。如果要用不同尺寸需要把网络改成全卷积结构并处理padding。4. 避坑与排查着色模型翻车的五个血泪现场4.1 输出全是灰色颜色去哪了现象训练几十轮后推理结果几乎还是灰度图只有极淡的色偏。原因最常见的是损失函数用错。如果用MSE且学习率偏大模型很快学会“预测ab接近0”这个安全策略因为灰色在所有样本上平均损失最低。另一个原因是ab归一化范围不对导致梯度消失。解决换L1损失检查ab是否真的落在[-1,1]。如果还不行在损失里给ab通道加权比如loss 1.0*L1(pred, target) 0.5*L1(pred[:,0], target[:,0])强制模型关注色度。4.2 颜色溢出到不该有的区域现象天空的蓝色渗到建筑上人脸的颜色糊到背景。原因编码器降采样太狠空间信息丢失严重。或者跳跃连接没加解码器只能靠低分辨率特征猜颜色。解决确认U-Net的skip connection正确拼接。如果还溢出减少下采样次数比如从4次降到3次。另一个技巧是在损失里加入感知损失用预训练VGG提取特征做比较能显著改善边缘。4.3 训练损失降了但视觉效果没变好现象Loss从0.1降到0.03但生成的彩色图看起来还是脏脏的。原因L1/L2损失优化的是像素平均误差和人类视觉感知不一致。模型可能学会了“平均颜色”但缺乏饱和度和对比度。解决引入对抗损失加一个判别器判断生成图是真是假。这就是着色GAN的思路。判别器用PatchGAN输出N×N的patch真假概率。对抗损失让颜色更鲜艳、更真实。代价是训练不稳定需要调GAN的权重。4.4 显存爆炸batch size只能设1现象训练时OOM只能把batch size降到1训练极慢。原因U-Net在256×256分辨率下第一层64通道的特征图占大量显存。加上跳跃连接保留所有中间特征显存占用是普通CNN的好几倍。解决用混合精度训练AMP显存直接减半。代码里加torch.cuda.amp即可。另一个办法是把训练分辨率降到128或者用梯度累积模拟大batch。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): pred model(L) loss criterion(pred, ab) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()4.5 推理速度慢一张图要好几秒现象模型在GPU上推理一张256×256的图要2-3秒批量处理老照片时等不起。原因U-Net的转置卷积和跳跃连接拼接在推理时也有大量计算。如果没开torch.no_grad()还会额外占用显存和计算。解决推理时务必加with torch.no_grad():。导出ONNX或TorchScript做图优化。如果还慢把模型通道数减半或者用MobileNet作为编码器骨干。5. 进阶技巧让着色结果从“能看”到“好看”跑通基础版本后你大概会发现颜色是有了但总差一口气。天空不够蓝草地不够绿人脸偏黄。这不是模型不够大而是损失函数和训练策略的问题。我自己的经验是从L1升级到L1对抗损失感知损失视觉质量会有质的飞跃。具体做法保留U-Net作为生成器加一个PatchGAN判别器。判别器输入是ab通道或Lab三通道输出是N×N的真假概率图。生成器损失 L1损失 0.1×对抗损失 0.01×感知损失。感知损失用预训练VGG16的relu3_3层特征做比较。这个组合在多个公开数据集上都被验证有效。另一个技巧是训练数据增强。着色模型对数据量很敏感ImageNet的128万张图是底线。如果只有几千张照片必须做增强随机裁剪、水平翻转、颜色抖动对输入L做轻微亮度扰动。注意不要对ab做颜色抖动那会破坏标签。还有一个容易被忽略的点推理时的后处理。模型输出的ab可能有噪声可以用双边滤波或导向滤波平滑一下边缘保持的同时去掉色斑。如果做视频着色还要在时间维度上做一致性约束否则相邻帧颜色会跳。最后说一个我踩过的坑不要用太小的数据集硬训。我曾经用500张风景照训了200轮结果模型把所有天空都涂成紫色。后来换成ImageNet子集5万张同样的网络结构颜色立刻正常了。着色任务对数据多样性的要求比分类任务还高因为模型要学的不是“这是什么”而是“这应该是什么颜色”——后者需要见过足够多的场景组合。希望帮到你。本文还有配套的精品资源点击获取