Python+AI給老照片上色
哈嘍,大家好。
今天繼續(xù)給大家分享有趣的 AI 項(xiàng)目。
上次我們分享了用GAN(生成對(duì)抗網(wǎng)絡(luò))??讓靜態(tài)圖片動(dòng)起來(lái)??。
今天我們分享用NoGAN的圖像增強(qiáng)技術(shù)給老照片著色。效果如下:
原圖
上色后
NoGAN是一種新型GAN,它能花費(fèi)最少的時(shí)間進(jìn)行GAN訓(xùn)練。
今天分享的這個(gè)項(xiàng)目已在GitHub開(kāi)源項(xiàng)目,下面我們來(lái)運(yùn)行一下。
1. 準(zhǔn)備工作
首先,用git clone命令下載源碼。
git clone https://github.com/jantic/DeOldify.git
進(jìn)入項(xiàng)目根目錄,安裝Python依賴(lài)包。
pip3 install -r requirements.txt
編寫(xiě)代碼運(yùn)行項(xiàng)目之前,需要下載預(yù)訓(xùn)練好的模型。項(xiàng)目提供了三個(gè)模型:
模型
區(qū)別如下:
- ColorizeArtistic_gen.pth:在有趣的細(xì)節(jié)和活力方面實(shí)現(xiàn)了最高質(zhì)量的圖像著色效果,該模型在 UNet 上使用 resnet34 為主干,通過(guò) NoGAN 進(jìn)行了 5 次評(píng)論家預(yù)訓(xùn)練/GAN 循環(huán)重復(fù)訓(xùn)練。
- ColorizeStable_gen.pth:在風(fēng)景和肖像方面取得了最佳效果,該模型在 UNet 上使用 resnet101 為主干,通過(guò) NoGAN 進(jìn)行了 3 次評(píng)論家預(yù)訓(xùn)練/GAN 循環(huán)重復(fù)訓(xùn)練。
- ColorizeVideo_gen.pth:針對(duì)流暢的視頻進(jìn)行了優(yōu)化,它僅使用初始生成器/評(píng)論家預(yù)訓(xùn)練/GAN NoGAN 訓(xùn)練。由于追求流暢的速度,它的色彩比前兩者少。
- 將下載好的模型文件放在項(xiàng)目根目錄的models目錄下即可。
2. 編寫(xiě)代碼
在項(xiàng)目根目錄同級(jí)目錄下創(chuàng)建Python文件,編寫(xiě)代碼加載剛剛下載好的模型文件。
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指定項(xiàng)目根目錄,weights_name指定接下來(lái)使用哪個(gè)模型為照片上色。
讀取老照片,進(jìn)行上色;
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模塊將圖片轉(zhuǎn)換成模型輸入所需要的格式,送入模型進(jìn)行上色,完成后保存。
上述代碼是我從項(xiàng)目源碼中抽取的,可以看到,運(yùn)行代碼還是非常簡(jiǎn)單的。