OpenCVを使って動画からpix2pixの学習モデルを作ってみる

こんにちは、お久しぶりです。mediba広告システム開発部の原です。

前回はpython+TensorFlowで画像から顔認識と分類をする簡単なモデルについて書きました。

機械学習で芸能人の顔を分類してみよう!

で、今回ですが、やっぱり流行りのアレ。

流行ってますよね、pix2pix!

ということで、pix2pixを使うのに必要な学習素材を動画から簡単に作れますよ、今すぐ始められますよ、という内容です。

開発環境

最初に環境の話です。 本記事の作成・検証環境は以下のとおりです。

  • Mac OS X 10.12.3 (10.10.5とかでも動作すること確認済み)
  • Python 2.7.10
  • OpenCV2 2.4.12

概要

pix2pix

pix2pixとは、簡単には画像と画像の間の関係性/2画像間の変換パターンの特徴を学習、DNNで表現してしまおう、というプログラムです。

例えば地図から航空写真に変換するための学習を行い、架空の町並みを作ってみたり、線画からカラー画像を作ったりすることができます。

GAN (Generative Adversarial Network)という仕組みを基にしているわけですが、画像から画像への変換という仕組みは数あれど、pix2pixはとにかく簡単に始められて、しかも精度がすごい、と最近話題になってます。

参考:pix2pix

作る学習モデルについて

今回はシンプルに、モノクロ画像を彩色するというモデルを作ってみようと思います。

pix2pixでの学習に用いるのは変換前後それぞれ256×256ピクセルの画像を左右に結合した、512×256の画像(便宜上、これを素材タイルと呼びます)になります。

今回は、弊社メディアであるZ TOKYOのプロモーション動画から、学習に必要な素材タイルを作ってみたいと思います。

素材タイル作成

元動画について

今回、素材に用いる動画のスペックは下記のとおりです。

  • MP4フォーマット
  • 1920 × 1080
  • 30FPS
  • 60sec

動画を読み込み1フレームずつ画像として処理

さて、上記の動画を1フレームずつ処理すれば、1920 × 1080の画像が約1800枚出来ることになります。

※実際には60秒を少し越えていたので、1875フレームありました

これでは多すぎるので、あとで適当な枚数に減らすとして、まずはこの操作を書いていきます。 また、今回は彩色モデルを作るので、併せて画像をグレースケールに変換します。

#coding=utf-8

import cv2

# 入力する動画パスを指定
cap = cv2.VideoCapture("sample.mp4")

counter = 0

while(cap.isOpened()):
    ret, frame = cap.read()

    if ret == True :
        counter = counter + 1
        gray2 = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)

        # 変換するとカラー情報が落ちるので、一度ファイルに保存して開き直す
        cv2.imwrite('./gray.png', gray2)
        gray = cv2.imread('./gray.png')

        print(frame.shape)
        print(gray2.shape)
        print(gray.shape)

    if ret == False and counter != 0 :
        break

cap.release()
cv2.destroyAllWindows()

これだけです。

デバッグも兼ねて各フレームのshape(高さ、幅、チャネル数)を表示しています。

これを実行すると、

(1080, 1920, 3)
(1080, 1920)
(1080, 1920, 3)

という表示がフレーム数分表示されます。 上から元のフレーム/グレースケール変換/変換後のファイルをファイルから読み直したものの、それぞれのshapeになります。

単純なグレースケール変換ではチャネル数のデータが落ちていることがわかると思います。

あとで素材タイルに結合する際に、データ形式が違うと実行エラーが発生してしまいますので、ここで一旦ファイル保存を経由してデータ形式を揃えています(もっと効率がいい方法があるかもしれませんので、ご存じの方はこっそりおしえてください)

各画像から一部を切り出してタイルに結合

さて、上記処理で作ったframeおよびgrayは、どちらも幅1920×高さ1080の画像になっています。

素材タイルは256pix × 256pixの正方形なので次の図のように位置を指定して切り出していくことにします。

frameおよびgrayは色を除けば同じ画像ですので、それぞれから同じ位置の画像を切り抜けば、彩色有無だけが違った正方形の画像データを取り出すことが出来ます。

このやり方であれば、1フレームごとに28枚の画像を切り出すことが出来ます。

pix2pixの標準トレーニングデータは全体で600枚くらいですので、そのくらいのタイルセットが取れればとりあえずは十分だと思います。

今回は特にデータのスクリーニングを行いませんが、例えばほとんど一色のタイルなどは学習時にゴミになりえます。 そこで、この時点では少し多めにデータを用意することにして、70フレームに一度、この処理を行うことにしました。

(1800 / 70 ) * 28 = 720枚のタイルセットが出来ることになります。

pythonで書くとこうなります。

            # 切り出す画像の縦横幅定義
            height = 256
            width = 256

            # 切り出す位置の初期定義
            defX = 28 # 縦位置(pixel)初期値
            defY = 64 # 横位置(pixel)初期値

            # 縦方向ループ
            for numX in range(4):
                # 横方向ループ
                for numY in range(7):
                    cutX = defX + (numX * height)
                    cutY = defY + (numY * width)

                    cutImg = frame[cutX:cutX+256, cutY:cutY+256]
                    cutGray = gray[cutX:cutX+256, cutY:cutY+256]

train/test/valに画像振り分け

さて、720枚のデータセットができたところで、最後に素材タイルへと結合していきます。

また、pix2pix標準のデータセット600枚は以下のような内訳になっています。

  • train(学習用データ)400枚程度
  • test(テスト用データ)100枚程度
  • val(モデルの検証用データ)100枚程度

これに併せてtrain:test:valを4:1:1になるようにランダムで振り分ける仕組みもついでに実装すれば、動画から素材タイルを作り出すプログラムの完成です。

で、これまで書いたものとあわせて出来上がったものがこちら。

#coding=utf-8

import cv2
from numpy.random import *

# 入力する動画パスを指定
cap = cv2.VideoCapture("sample.mp4")

counter = 0
dataset_counter = 0

while(cap.isOpened()):
    ret, frame = cap.read()

    if ret == True :
        counter = counter + 1

        if counter % 70 == 0 : #70フレームに一度画像処理を行う

            # グレースケールに変換
            gray2 = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)

            # 変換するとカラー情報が落ちるので、一度ファイルに保存して開き直す
            cv2.imwrite('./gray.png', gray2)
            gray = cv2.imread('./gray.png')

            # 切り出す位置の初期定義
            defX = 28 # 縦位置(pixel)初期値
            defY = 64 # 横位置(pixel)初期値

            # 切り出す画像の縦横幅定義
            height = 256
            width = 256

            # 縦方向ループ
            for numX in range(4):
                # 横方向ループ
                for numY in range(7):
                    dataset_counter = dataset_counter+1
                    cutX = defX + (numX * height)
                    cutY = defY + (numY * width)

                    cutImg = frame[cutX:cutX+256, cutY:cutY+256]
                    cutGray = gray[cutX:cutX+256, cutY:cutY+256]

                    # 画像の結合
                    imgAdd = cv2.hconcat([cutImg, cutGray])

                    # train:test:val = 4:1:1になるように保存ディレクトリ振り分け
                    key =  = rand() * 6

                if key < 4:
                    saveDir = 'train'
                elif key < 5:
                    saveDir = 'test'
                else:
                    saveDir = 'val'

                cv2.imwrite('./datasets/sample/%s/%d.png' % (saveDir, dataset_counter), imgAdd)


    if ret == False and counter != 0 :
        break

cap.release()
cv2.destroyAllWindows() 

これで720枚の素材を振り分けることが出来ました。

学習・結果評価

データ移動

作成したトレーニングデータをpix2pixのdatasets以下に移動します。

$ cp -a datasets/sample ~/pix2pix/datasets/movie

トレーニング

移動したデータを下にpix2pixのトレーニングを実施します。

$ DATA_ROOT=./datasets/movie \
  name=movieDump2  which_direction=BtoA \
  gpu=0 cudnn=0 batchSize=20 \
  save_epoch_freq=5 save_latest_freq=10 \
  continue_train=0 niter=10 th train.lua

パラメータについて以下の設定をしています。 gpu=0 cudnn=0:GPUおよびCUDNNを利用しない save_epoch_freq=5:5回学習毎にモデルを保存する save_latest_freq=10:10回回学習毎に「最新学習モデル」を保存する continue_train=0:前回の「最新学習モデル」を利用して継続学習する ※「最新学習モデル」がないと利用できない niter=10:学習ループ回数定義(デフォルトは200回)

pix2pixによる彩色の実施

最後に、ここで作成したモデルで彩色を実施します。

$ DATA_ROOT=./datasets/movie \
  name=movieDump2   which_direction=BtoA \
  phase=val   gpu=0 cudnn=0 th test.lua

ここでもGPUおよびCUDNNは利用しません。

評価

実行した結果がこちらになります。

左:グレースケール変換した画像 中:pix2pixが機械的に彩色した画像 右:元動画から切り出した(正解の)画像

今回は10回ループでしたが、これでも結構学習してくれているのがわかると思います。

一方で、

こんな状況になっている画像もありました。

幾つか原因はあると思うのですが、

  • 元動画には上下に黒帯が入っているのに除去しなかったこと
  • 機械的に学習タイルを作成したので、元のタイルが全体的にマットで均一色彩の場合はほぼ効果が出ない(むしろじゃまになりかねない)

ということなどが考えられます。

これらは切り出す位置を工夫するとか、ヒストグラムを作って均一な色合いのタイルを予め除去しておくなどの対応が取れると思います。 せっかく多めの素材タイル用意したんだからやっておけよ、という話

このように課題も見つかりましたが、簡単にpix2pixの学習素材が作れることはわかっていただけたかと思います。 それにしても、GANすごいなあ……

まとめ

  • GANすごい
  • GPUない環境でやるもんじゃない
  • Z TOKYO見に来てね!