本記事では、機械学習手法Time-Travel Rephotographyを用いて、昔の白黒写真をカラー化する方法をご紹介します。

出典: Time-Travel-Rephotography.github.io

Time-Travel Rephotography

概要

従来のColorizationタスクを実現する技術は、ノイズ除去、カラー化、超解像などサブタスクがそれぞれ独立しており、それらを組み合わせ実現していました。

Time-Travel Rephotographyでは、StyleGAN2フレームワークを活用し、古い写真を高解像度写真の空間に投影することで、独立していたサブタスクを統一したフレームワーク上で実現しています。従来技術と異なり、元の写真の特徴やポーズを維持しながら、高品質でカラー化された写真を生成することでColorizationを実現しています。

出典: Time-Travel Rephotography

詳細はこちらの論文をご参照ください。

本記事では上記手法を用いて、任意の写真のカラー化を行います。

デモ(Colaboratory)

それでは、実際に動かしながら写真のカラー化を行っていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。

また、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

初めにGithubからソースコードを取得します。

%cd /content

!git clone  --depth 1 --recurse-submodules --shallow-submodules \
  https://github.com/Time-Travel-Rephotography/Time-Travel-Rephotography.github.io.git Time-Travel-Rephotography

# for face align
!git clone https://github.com/adamian98/pulse.git

次にライブラリをインストールします。

%cd /content/Time-Travel-Rephotography

!pip3 install -r requirements.txt
!pip3 install --upgrade gdown

最後にライブラリをインポートします。

%cd /content/Time-Travel-Rephotography

from pathlib import Path
import os

from PIL import Image
from IPython.display import display

from argparse import Namespace
from projector import (
    ProjectorArguments,
    main,
)

以上で環境セットアップは完了です。

学習済みモデルのセットアップ

続いて、学習済みモデルをダウンロードしていきます。

%cd /content/Time-Travel-Rephotography
!mkdir -p ./checkpoint/encoder

if not os.path.exists('checkpoint/e4e_ffhq_encode.pt'):
  !gdown https://drive.google.com/uc?id=1YLB-3ZCv6FRAWwCHkTvUou_CAvDcb5Pr -O checkpoint/e4e_ffhq_encode.pt
if not os.path.exists('checkpoint/stylegan2-ffhq-config-f.pt'):
  !gdown https://drive.google.com/uc?id=1aSnTVHGNzQ-Eh5rNDOrOc_4SjQbNMq3B -O checkpoint/stylegan2-ffhq-config-f.pt
if not os.path.exists('checkpoint/vgg_face_dag.pt'):
  !gdown https://drive.google.com/uc?id=12BHlsSQM0D8KyXprIc7WmJRLWBWdtF56 -O checkpoint/vgg_face_dag.pt

if not os.path.exists('checkpoint/checkpoint_b.pt'):
  !gdown https://drive.google.com/uc?id=1fXBCyBKNEZfeiI6LEHTEnUg_DjWHJy5r -O checkpoint/encoder/checkpoint_b.pt
if not os.path.exists('checkpoint/checkpoint_g.pt'):
  !gdown https://drive.google.com/uc?id=1YnQEPf7FyfZxAVQg-CpbvdF4otZeTOZh -O checkpoint/encoder/checkpoint_g.pt
if not os.path.exists('checkpoint/checkpoint_gb.pt'):
  !gdown https://drive.google.com/uc?id=1kelQK3pUdeHwu8uVE7A4WZH_EjHxGE6v -O checkpoint/encoder/checkpoint_gb.pt

%cd /content/Time-Travel-Rephotography/third_party/face_parsing
!mkdir -p res/cp

if not os.path.exists('checkpoint/res/cp/79999_iter.pth'):
  !gdown https://drive.google.com/uc?id=1vwm4BcAKISQgcJLvTUcvesk73UIYdDMF -O res/cp/79999_iter.pth

モデルの数は多いですが、1～2分でダウンロードは完了します。

テスト画像のセットアップ

続いて、カラー化したい任意の画像をダウンロードしておきます。
今回はwgetで白黒画像をGoogle Colaboratory環境下にダウンロードして使用します。

%cd /content/Time-Travel-Rephotography
!mkdir test_imgs

# マリリンモンロー
!wget https://www.allcinema.net/img/6/6476/p_40178_01_01_02.jpg -O ./test_imgs/test.jpg

画像から顔部分を切り取ります。

%cd /content/pulse/

!python align_face.py \
  -input_dir /content/Time-Travel-Rephotography/test_imgs \
  -output_dir  /content/Time-Travel-Rephotography/test_aligns \
  -output_size 512 \
  -seed 12 \

入力画像は以下の通りです。

Colorization

それでは、カラー化を実施していきます。
spectral_sensitivityは白黒写真のタイプによって切り替えるとより良い結果が得られます。

#@markdown 'b' (blue-sensitive), 'gb' (orthochromatic), 'g' (panchromatic)
spectral_sensitivity = "g" # @param ["b", "gb", "g"]
#@markdown Estimated blur radius of the input photo
gaussian_radius =   0.75 # @param {type:"number"}

args = ProjectorArguments().parse(
    args=[str(input_path)], 
    namespace=Namespace(
        spectral_sensitivity=spectral_sensitivity,
        encoder_ckpt=f"checkpoint/encoder/checkpoint_{spectral_sensitivity}.pt",
        encoder_name=spectral_sensitivity,
        gaussian=gaussian_radius,
        log_visual_freq=1000,
        log_dir="log/",
        results_dir="results/"
))
main(args)

生成された画像を表示します。

def get_concat_h_multi_resize(im_list, resample=Image.BICUBIC):
    min_height = min(im.height for im in im_list)
    im_list_resize = [im.resize((int(im.width * min_height / im.height), min_height),resample=resample)
                      for im in im_list]
    total_width = sum(im.width for im in im_list_resize)
    dst = Image.new('RGB', (total_width, min_height))
    pos_x = 0
    for im in im_list_resize:
        dst.paste(im, (pos_x, 0))
        pos_x += im.width
    return dst

im1 = Image.open(f"/content/Time-Travel-Rephotography/results/test_0_{spectral_sensitivity}.png")
im2 = Image.open('/content/Time-Travel-Rephotography/results/test/skin_mask/face_parsing/input.jpg')
im3 = Image.open(f"/content/Time-Travel-Rephotography/results/test_0-{spectral_sensitivity}-G{gaussian_radius}-init(10,18)-s256-vgg1-vggface0.3-eye0.1-color1.0e+10-cx0.1(relu3_4,relu2_2,relu1_2)-NR5.0e+04-lr0.1_0.01-c32-wp(250,750).png")

concat_img = get_concat_h_multi_resize([im1, im2, im3])

display(concat_img)

出力結果は以下の通りです。

まとめ

本記事では、機械学習手法Time-Travel Rephotographyを用いて、昔の白黒写真をカラー化する方法をご紹介しました。
GANを使用しているため時折不安定さを感じますが、自然なカラー写真が生成されます。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - Time-Travel Rephotography

2. GitHub - Time-Travel-Rephotography.github.io