Python+AI给老照片上色

Posted 2022-06-17 机器学习社区

哈喽，大家好，今天继续给大家分享有趣的 AI 项目。

今天我们分享用NoGAN的图像增强技术给老照片着色。喜欢记得收藏、关注、点赞。

注：资料数据、文末提供技术交流群

效果如下：

原图

上色后

NoGAN是一种新型GAN，它能花费最少的时间进行GAN训练。

今天分享的这个项目已在GitHub开源项目，下面我们来运行一下。

1. 准备工作

首先，用git clone命令下载源码

git clone https://github.com/jantic/DeOldify.git

进入项目根目录，安装Python依赖包

pip3 install -r requirements.txt

编写代码运行项目之前，需要下载预训练好的模型。项目提供了三个模型

模型

区别如下：

• ColorizeArtistic_gen.pth：在有趣的细节和活力方面实现了最高质量的图像着色效果，该模型在 UNet 上使用 resnet34 为主干，通过 NoGAN 进行了 5 次评论家预训练/GAN 循环重复训练

• ColorizeStable_gen.pth：在风景和肖像方面取得了最佳效果，该模型在 UNet 上使用 resnet101 为主干，通过 NoGAN 进行了 3 次评论家预训练/GAN 循环重复训练

• ColorizeVideo_gen.pth：针对流畅的视频进行了优化，它仅使用初始生成器/评论家预训练/GAN NoGAN 训练。由于追求流畅的速度，它的色彩比前两者少。

将下载好的模型文件放在项目根目录的models目录下即可。

2. 编写代码

在项目根目录同级目录下创建Python文件，编写代码加载刚刚下载好的模型文件。

from DeOldify.deoldify.generators import gen_inference_wide
from DeOldify.deoldify.filters import MasterFilter, ColorizerFilter

# 指定模型文件
learn = gen_inference_wide(root_folder=Path('./DeOldify'), weights_name='ColorizeVideo_gen')

# 加载模型
deoldfly_model = MasterFilter([ColorizerFilter(learn=learn)], render_factor=10)

root_folder指定项目根目录，weights_name指定接下来使用哪个模型为照片上色。

读取老照片，进行上色

import cv2
import numpy as np
from PIL import Image

img = cv2.imread('./images/origin.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
pil_img = Image.fromarray(img)

filtered_image = deoldfly_model.filter(
    pil_img, pil_img, render_factor=35, post_process=True
)

result_img = np.asarray(filtered_image)
result_img = cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR)
cv2.imwrite('deoldify.jpg', result_img)

用cv2读取老照片，并用PIL.Image模块将图片转换成模型输入所需要的格式，送入模型进行上色，完成后保存。

上述代码是我从项目源码中抽取的，可以看到，运行代码还是非常简单的。

感兴趣的朋友可以自己运行一下，也可以试试其他模型的效果。

技术交流

目前已开通了技术交流群，群友已超过2000人，添加时最好的备注方式为：来源+兴趣方向，方便找到志同道合的朋友

• 方式①、发送如下图片至微信，长按识别，后台回复：加群；
• 方式②、微信搜索公众号：机器学习社区，后台回复：加群；
• 方式③、可以直接加微信号：mlc2060。加的时候备注一下：研究方向 +学校/公司+CSDN，即可。然后就可以拉你进群了。
  

研究方向包括：机器学习、数据挖掘、目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer等。

一定要备注：研究方向+学校/公司+昵称（如Transformer或者目标检测+上交+卡卡），根据格式备注，可更快被通过且邀请进群。