用于圖像著色的五個開源Python工具
?成千上萬的老式照片和電影是在彩色圖像還沒有出現(xiàn)的年代拍攝的。如今,在深度學(xué)習(xí)的幫助下,我們可以為這些圖片著色,使它們更接近原來的樣子。
作為測試,我將使用兩張圖像。
本文將使用開源工具,這些工具可以從GitHub下載。
DeOldify
DeOldify是基于SA-GAN (Self-Attention - generate對抗網(wǎng)絡(luò))。一般來說,GAN由兩個獨立的神經(jīng)網(wǎng)絡(luò)組成——生成器和判別器。這兩個網(wǎng)絡(luò)都是由大量的圖像訓(xùn)練而成,在訓(xùn)練過程中,生成器學(xué)會了制作似是而非的圖像,而判別器學(xué)會將生成的圖像與真實的圖像區(qū)分開來。
為了在本地運行DeOldify,我們需要從GitHub中獲取項目,并將預(yù)訓(xùn)練好的神經(jīng)網(wǎng)絡(luò)權(quán)重下載到“models”文件夾中(鏈接來自項目頁面,未來作者可能會更改):
git clone https://github.com/jantic/DeOldify
cd DeOldify
mkdir models
cd models
wget https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth
wget https://www.dropbox.com/s/usf7uifrctqw9rl/ColorizeStable_gen.pth?dl=1 -O ColorizeStable_gen.pth
cd ..
在這里,我將“device”設(shè)置為CPU - 如果您沒有好的顯卡,則很可能會收到“內(nèi)存不足”錯誤(CPU上的處理時間約為3-5s,GPU上的處理時間約為0.5s,因此CPU計算也運行良好)。如果您希望運行獨立的Python代碼,則可以使用以下Python代碼:
from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import *
torch.backends.cudnn.benchmark=True
device.set(device=DeviceId.CPU)
colorizer = get_image_colorizer(artistic=True)
img_out = colorizer.get_transformed_image(path="anna_bw.jpg",
render_factor=15,
watermarked=True)
img_out.save("anna_color.jpg")
至于結(jié)果,相當不錯:
大家還可以嘗試更改模型(有兩種模型,具有更鮮艷顏色的“artistic”模型和“stable”模型可用)和影響輸出顏色的“render_factor”變量。要去除水印可以將參數(shù)watermarked設(shè)置為False。
彩色圖像著色(Colorful Image Colorization)
這個項目使用卷積神經(jīng)網(wǎng)絡(luò)(CNN)來生成彩色圖像。在體系結(jié)構(gòu)上,它比DeOldify簡單得多,但正因為如此,它可能更方便大家理解它的工作原理。
最后一次項目更新是在2020年,但代碼仍然可以工作,并且可以很容易地在本地運行。
有兩種模型可用,Python代碼如下:
import colorizers as c
import torch
import matplotlib.image as plt
img = c.load_img("image_bw.jpg")
tens_l_orig, tens_l_rs = c.preprocess_img(img, HW=(256, 256))
img_bw = c.postprocess_tens(tens_l_orig,
torch.cat((0*tens_l_orig, 0*tens_l_orig), dim=1))
colorizer_eccv16 = c.eccv16(pretrained=True).eval()
out_img_eccv16 = c.postprocess_tens(tens_l_orig,
colorizer_eccv16(tens_l_rs).cpu())
plt.imsave('image_eccv16.jpg', out_img_eccv16)
colorizer_siggraph17 = c.siggraph17(pretrained=True).eval()
out_img_siggraph17 = c.postprocess_tens(tens_l_orig,
colorizer_siggraph17(tens_l_rs).cpu())
plt.imsave('image_siggraph17.jpg', out_img_siggraph17)
結(jié)果如下:
ChromaGAN
顧名思義,ChromaGAN的作者也在使用生成對抗網(wǎng)絡(luò)給圖像上色。
這個項目可能只是作為研究的一個演示,操作并不是太友好。在使用項目之前,應(yīng)下載“my_model_colorization.h5”文件(鏈接在GitHub頁面上提供)并放入“MODEL”文件夾中。源圖像和輸出圖像應(yīng)該分別放在“chromagan_images”和“chromagan_results”文件夾中,然后可以參考作者jupyter notebook進行處理。要在本地PC上運行代碼,“from google.colab.patches import cv2_imshow”和“cv2_imshow(…)”行應(yīng)該刪除。該項目正在使用Keras,如果出現(xiàn)“內(nèi)存不足”的錯誤,建議在文件開頭添加os.environ["CUDA_VISIBLE_DEVICES"] = "-1"。
Google Colorization Transformer (ColTran)
這個項目可以從google-research GitHub頁面下載,更詳細的研究論文也可以下載。作者使用的是具有自注意力架構(gòu)的軸向transformer,而不是GAN。在使用ColTran之前,我們需要下載預(yù)訓(xùn)練模型,這些模型位于ColTran.zip歸檔文件中。該歸檔文件包含3個模型檢查點:colorizer、color_upsampler和spatial_upsampler。然后我們可以運行3個Python命令:
python3 custom_colorize.py --config=configs/colorizer.py --mode=colorize --accelerator_type=CPU --logdir=colorizer --img_dir=img_dir --store_dir=target_dir
python3 custom_colorize.py --config=configs/color_upsampler.py --mode=colorize --accelerator_type=CPU --logdir=color_upsampler --img_dir=img_dir --store_dir=target_dir --gen_data_dir=target_dir/stage1 --mode=colorize
python3 custom_colorize.py --config=configs/spatial_upsampler.py --mode=colorize --accelerator_type=CPU --logdir=spatial_upsampler --img_dir=img_dir --store_dir=target_dir --gen_data_dir=target_dir/stage2
這里的img_dir是一個文件夾,包含源圖像,store_dir是輸出文件夾,colorize是一種處理模式,而logdir是到預(yù)訓(xùn)練模型的路徑。我們有3個處理步驟:我們有3個處理步驟:“colorizer”只使用512種輸出顏色和64x64的輸出圖像進行粗著色,“color upsampler”改善顏色,“spatial upsampler”將圖像提升到256x256的分辨率。
結(jié)果是顏色也很準確:
這個工具可能只是作為研究論文的演示而制作的,與以前的項目相比,沒有現(xiàn)成的方法來處理任意分辨率的圖像。輸出僅限于 256x256 大小。
BigColor
BigColor項目是由作者在2022年提出的。作者還在他們的論文中寫道:“我們將BigColor與最近的自動著色方法進行了比較,包括CIC、ChromaGAN、DeOldify、InstColor、ColTran和ToVivid。在六張具有挑戰(zhàn)性的圖像上,BigColor的優(yōu)于所有方法。”
該項目本身可以從GitHub頁面下載(https://github.com/KIMGEONUNG/BigColor)。使用代碼很簡單。在進行轉(zhuǎn)換之前,應(yīng)該執(zhí)行兩個腳本download- pretraining.sh和download-bigcolor.sh。之后,可以使用一個命令完成轉(zhuǎn)換:
python3 colorize_real.py --path_ckpt=ckpts/bigcolor --path_input=images_gray --epoch=11 --type_resize=powerof --seed=-1 --device=cpu
此處的path_ckpt是指向預(yù)訓(xùn)練模型的路徑,images_gray是包含源圖像的文件夾。結(jié)果如下:
最后
圖像著色是一個有趣的話題,正如我們所看到的,不同的方法和架構(gòu)是可能的。從準確性的角度來看,事情也很復(fù)雜。通常,黑白照片不再有顏色信息,因此神經(jīng)網(wǎng)絡(luò)只能根據(jù)之前訓(xùn)練的圖像來猜測輸出結(jié)果。例如,這是我用來測試的原始圖像:
這是相同的圖像,轉(zhuǎn)換為黑白:
這是使用DeOldify制作的圖像:
樹是綠色的,天空是藍色的,這已經(jīng)挺不錯了。但是,不僅DeOldify,而且其他經(jīng)過測試的項目都無法正確確定百葉窗的顏色。在大多數(shù)情況下,這些結(jié)果已經(jīng)足夠好了。