船木俊介「デジタル進化論」

スーパーソフトウエア東京オフィス代表&キッズラインCTO

人工知能でモノクロ写真を自動でカラーにする方法

ディープラーニング(Convolutional Neural Networks(畳み込みニューラルネットワーク))を使ってモノクロ写真を自動でカラー化するよ、という面白そうなプロジェクトAutomatic Colorizationがあったので試してみた。

Rather than reconstructing the entire color RGB image, I trained models to produce two color channels which I concatenate with the grayscale input channel to produce a YUV image. The Y channel is intensity. This ensures that the intensity of the output will always be the same as the input.
輝度信号Yと色差信号のYUVを使うところがポイントで、グレースケールの入力から2つの色差信号を取り出すようにトレーニングされたモデル、とのことです。

Models were trained on the ILSVRC 2012 classification training dataset. The same training set used for the pretrained VGG16. It's 147GB and over 1.2 million images!
ImageNet LSVRC 2012の学習済みモデルが配布されていて、TensorFlow環境さえあればスクリプトの実行だけで動かせるのでいくつかやってみると、

まじか!?というくらい、草原、空など最も上手くカラー復元できた写真がこれ。教師データに多かったであろう自然の風景は、驚くほど再現されている。モノクロ画像の情報だけからここまで色彩を復元できるのは驚き。 m-c
ゴダール「勝手にしやがれ」のジャンポールベルモント。 m2-c
フリー素材の女子高生 m3-c
「ローマの休日」のオードリー・ヘップバーン m4-c

静止画でできるのであれば動画もできるはずということで、ffmpegで切り出して処理して結合という定番の方法でやってみました。

ゾンビ映画の名作、ジョージ・A・ロメロのNight of the Living Dead(ナイト・オブ・ザ・リビングデッド)をカラーにしてみたもの。やっぱり、冒頭の墓地の風景などは、とても上手くカラーになっています。



How to

TensorFlowの学習済みモデルをAutomatic Colorizationからダウンロードしておきます。torrent形式なのでクライアントソフトが必要です。

まず、GoogleのTensorFlowが動く環境を作ります。TensorFlowは2015年11月に無償公開されたばかりのライブラリです。(参考 グーグル、重要な「ヒトの数百万倍の能力」技術を無償公開 革新的分野で圧倒的主導権)

tensorflowのインストールガイドに従って、virtualenv環境でインストールします。
# Mac OS X
$ sudo easy_install pip
$ sudo pip install --upgrade virtualenv

tensorflowをアクティベート
$ virtualenv --system-site-packages ~/tensorflow
$ source ~/tensorflow/bin/activate  # If using bash

tensorflowをpipでインストール。Macで動かしたので、今回はCPUベースで。
# Mac OS X, CPU only:
(tensorflow)$ pip install --upgrade https://storage.googleapis.com/tensorflow/mac/tensorflow-0.6.0-py2-none-any.whl

tensorflowの環境が出来たので、画像単体のカラー復元処理は実行できます。モデルと一緒にforward.pyがあるので、
(tensorflow)$ python forward.py

これだけでOK。

動画のカラー化

動画の場合は、前後にffmpegでの処理を行って、Automatic Colorizationをフォルダ内の画像に対してループするだけです。

動画を画像として切り出します。予め、framesフォルダの中に対象となる動画を入れておきます。1秒あたりの枚数が、あまり多くても処理時間がかかるのでとりあえず30フレームで。
(tensorflow)$ mkdir frames
(tensorflow)$ mkdir -p out/frames
(tensorflow)$ cd ./frames
(tensorflow)$ ffmpeg -i f1.mp4 -r 30 -f image2 %06d.jpg

Automatic Colorizationは画像サイズが224となっているので、リサイズしておきます。
(tensorflow)$ mogrify -resize 224x224 *.jpg
(tensorflow)$ mogrify -gravity center -background black -extent 224x224 *.jpg
(tensorflow)$ mogrify -colorspace sRGB -type TrueColor *.jpg
(tensorflow)$ cd ..

forward.pyをフォルダの中の画像に対して行うようにしたスクリプト(forward_movie.py)を実行します。
import tensorflow as tf
import skimage.transform
from skimage.io import imsave, imread
import os
from os import listdir, path
from os.path import isfile, join

def get_directory(folder):
    foundfile = []

    for path, subdirs, files in os.walk(folder):
        for name in files:
            found = os.path.join(path, name)
            if name.endswith('.jpg'):
                foundfile.append(found)
        break

    foundfile.sort()
    return foundfile

def load_image(path):
    img = imread(path)
    # crop image from center
    short_edge = min(img.shape[:2])
    yy = int((img.shape[0] - short_edge) / 2)
    xx = int((img.shape[1] - short_edge) / 2)
    crop_img = img[yy : yy + short_edge, xx : xx + short_edge]
    # resize to 224, 224
    img = skimage.transform.resize(crop_img, (224, 224))
    # desaturate image
    return (img[:,:,0] + img[:,:,1] + img[:,:,2]) / 3.0

with open("colorize.tfmodel", mode='rb') as f:
    fileContent = f.read()

    graph_def = tf.GraphDef()
    graph_def.ParseFromString(fileContent)
    grayscale = tf.placeholder("float", [1, 224, 224, 1])
    tf.import_graph_def(graph_def, input_map={ "grayscale": grayscale }, name='')

images = get_directory("frames")

for image in images:
    print image
    shark_gray = load_image(image).reshape(1, 224, 224, 1)

    with tf.Session() as sess:
        inferred_rgb = sess.graph.get_tensor_by_name("inferred_rgb:0")
        inferred_batch = sess.run(inferred_rgb, feed_dict={ grayscale: shark_gray })
        filename = "out/"+image
        imsave(filename, inferred_batch[0])
        print "saved " + filename
(tensorflow)$ python forward_movie.py

最後に、できあがった画像を結合して動画を作成します。
(tensorflow)$ cd out/frames
(tensorflow)$ cat *.jpg | ffmpeg -f image2pipe -r 30 -vcodec mjpeg -i - -vcodec libx264 out.mp4