当前位置：首页 > ds >正文

用Python做有趣的AI项目 3：黑白图像自动上色（AI 上色器）

ds 2025/5/6 14:52:14

这个项目的目标是使用深度学习模型，将黑白照片自动“还原”出彩色版本，可以用在修复老照片、创意设计等场景。

🧠 项目概述

我们将使用：

PythonOpenCV：读取和处理图像

预训练的深度学习模型（基于 Caffe 或 TensorFlow）来实现图像上色

如果你不打算训练自己的模型（因为需要很多数据和时间），我们可以使用 OpenCV 官方集成的预训练模型。

✅ 项目准备

第一步：安装依赖

bashpip install opencv-python opencv-contrib-python numpy

📦 第二步：下载预训练模型（OpenCV 提供）

我们将使用由 Richard Zhang 等人提供的模型，OpenCV 已集成支持。

需要下载这三个文件（约 300MB）：

colorization_deploy_v2.prototxtcolorization_release_v2.caffemodelpts_in_hull.npy（色彩分布点）

🧰 第三步：上色代码实现
python

import cv2
import numpy as np# 模型路径（需要你提前下载好）
proto_file = 'models/colorization_deploy_v2.prototxt'
model_file = 'models/colorization_release_v2.caffemodel'
pts_file = 'models/pts_in_hull.npy'# 加载网络
net = cv2.dnn.readNetFromCaffe(proto_file, model_file)
pts = np.load(pts_file)# 加载颜色点
class8 = net.getLayerId("class8_ab")
conv8 = net.getLayerId("conv8_313_rh")
pts = pts.transpose().reshape(2, 313, 1, 1)
net.getLayer(class8).blobs = [pts.astype("float32")]
net.getLayer(conv8).blobs = [np.full([1, 313], 2.606, dtype="float32")]def colorize_image(image_path):# 读取灰度图像并转换为 LABbw_image = cv2.imread(image_path)h, w = bw_image.shape[:2]img_rgb = (bw_image.astype("float32")) / 255.0img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2Lab)l_channel = img_lab[:, :, 0]l_resized = cv2.resize(l_channel, (224, 224))l_resized -= 50  # 中心化net.setInput(cv2.dnn.blobFromImage(l_resized))ab_channels = net.forward()[0, :, :, :].transpose((1, 2, 0))ab_channels = cv2.resize(ab_channels, (w, h))lab_output = np.concatenate((l_channel[:, :, np.newaxis], ab_channels), axis=2)color_image = cv2.cvtColor(lab_output, cv2.COLOR_Lab2BGR)color_image = (np.clip(color_image, 0, 1) * 255).astype("uint8")return color_image

🖼️ 第四步：测试和显示结果

pythonimport matplotlib.pyplot as pltdef show_result(image_path):color_img = colorize_image(image_path)bw_img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)plt.figure(figsize=(10, 5))plt.subplot(1, 2, 1)plt.imshow(bw_img, cmap='gray')plt.title("原始黑白图")plt.axis('off')plt.subplot(1, 2, 2)plt.imshow(cv2.cvtColor(color_img, cv2.COLOR_BGR2RGB))plt.title("上色后的图像")plt.axis('off')plt.show()