TensorFlowのMagentaで高速に画風変換してみた

Magentaで画像の高速スタイル変換ができるようになったというニュースがあったので早速試してみました。

ニュースソースは、Google BrainチームのMagentaプロジェクトのこちらの記事（https://magenta.tensorflow.org/2016/11/01/multistyle-pastiche-generator/）とGoogle Research Blogのこちらの記事（https://research.googleblog.com/2016/10/supercharging-style-transfer.html）です。

なお、Magentaは、同じくGoogleが開発しているTensorFlow上で動作します。

環境構築

今回は、Docker上で動作させるため、ここ（https://hub.docker.com/r/tensorflow/magenta/）から公式Dockerイメージをダウンロードします。

以下のコマンドで、Magenta (0.1.6)のDockerイメージをダウンロードします。

$ docker pull tensorflow/magenta:0.1.6

GitHub（https://github.com/tensorflow/magenta/tree/v0.1.6/magenta/models/image_stylization）のREADME.mdでは、以下のコマンドでコンテナを生成するようになっていますが、

$ docker run -it -p 6006:6006 \
      -v /tmp/magenta:/magenta-data tensorflow/magenta:0.1.6

このコンテナは、セッションが終了すると、/magenta-dataディレクトリのファイルが消えてしまうので、以下のコマンドでホスト側に共有用ディレクトリを作成しておき、

$ mkdir -p ~/docker-data/magenta

Jupyter用のポートも設定して、以下のコマンドでコンテナを生成します。

$ docker run -it -p 6006:6006 -p 8888:8888 \
      -v ~/docker-data/magenta:/magenta-data tensorflow/magenta:0.1.6

コンテナに入って、以下のコマンドでバージョン等を確認してみました。

# python -c "import tensorflow; print(tensorflow.__version__)"
0.11.0rc2
# python -c "import magenta; print(magenta.__version__)"
0.1.6
# ls /usr/local/
bin etc games include lib man sbin share src

TensorFlowのバージョンは0.11.0rc2、Magentaのバージョンは0.1.6、残念ながら、CUDAライブラリがインストールされていないCPU版です。

なお、TensorFlow 0.10.0とMagenta 0.1.4の組み合わせで、GPU版のDockerイメージを作成できていたのですが、TensorFlow 0.11.0とMagenta 0.1.6の組み合わせでは、途中でエラーが発生しています。

高速スタイル変換

README.mdでは、画風変換コマンド（image_stylization_transform）の使い方は、

# image_stylization_transform \
      --num_styles= \
      --checkpoint=/path/to/model.ckpt \
      --input_image=/path/to/image.jpg \
      --which_styles="[0,1,2,5,14]" \
      --output_dir=/tmp/image_stylization/output \
      --output_basename="stylized"

その具体例は、以下のようになっています。

# image_stylization_transform \
      --num_styles=10 \
      --checkpoint=multistyle-pastiche-generator-monet.ckpt \
      --input_image=photo.jpg \
      --which_styles="{0:0.1,1:0.1,2:0.1,3:0.1,4:0.1,5:0.1,6:0.1,7:0.1,8:0.1,9:0.1}" \
      --output_dir=/tmp/image_stylization/output \
      --output_basename="all_monet_styles"

--checkpointオプションのmultistyle-pastiche-generator-monet.ckptは、トレーニング済のモデルで、公式Dockerイメージの場合、/magenta-modelsディレクトリに印象派クロード・モネのMonetモデル（multistyle-pastiche-generator-monet.ckpt）とvariedモデル（multistyle-pastiche-generator-varied.ckpt）の2つがダウンロード済みです。

以下のコマンドで、ホスト側の共有ディレクトリにコンテンツ画像をコピーします。

$ cp content.jpg ~/docker-data/magenta

上記の具体例から、環境に合わせて--checkpointオプションのパス、--input_imageオプションの画像ファイル名、--output_dirオプションのパスを変更した、以下のコマンドを実行します。

# image_stylization_transform \
      --num_styles=10 \
      --checkpoint=/magenta-models/multistyle-pastiche-generator-monet.ckpt \
      --input_image=content.jpg \
      --which_styles="{0:0.1,1:0.1,2:0.1,3:0.1,4:0.1,5:0.1,6:0.1,7:0.1,8:0.1,9:0.1}" \
      --output_dir=. \
      --output_basename="all_monet_styles"

すると、all_monet_styles_0_100_（中略）_9_100.pngというファイル名の画像が生成されます。

docker cpで、ホストにファイルをコピーすることもできますが、コンテナ起動時のオプションで-v ~/docker-data/magenta:/magenta-dataを設定してあるので、コンテナ側の/magenta-dataディレクトリのファイルを、ホスト側の~/docker-data/magentaディレクトリで確認することができます。

以下の図が変換された画像ファイルです。

README.mdのMonetモデルは10、variedモデルは32という--num_stylesオプションの意味がよく分からず、トレーニング済のモデルが2つでは少ないと思っていたのですが、Monetモデルに10種類、variedモデルに32種類のスタイルが含まれているようです。

上記のコマンドでは、--which_stylesオプションで"{0:0.1,...,9:0.1}"を指定しているため、上の図はMonetモデルの10種類のスタイル0~9が同じ重み0.1で合成された図となります。

それぞれのモデルのスタイルを確認するために、以下のシェル・スクリプト（monet.shとvaried.sh）を作成し、実行してみました。

# monet.sh
for i in $(seq 0 1 9);
do
  image_stylization_transform \
      --num_styles=10 \
      --checkpoint=/magenta-models/multistyle-pastiche-generator-monet.ckpt \
      --input_image=content.jpg \
      --which_styles="{$i:1}" \
      --output_dir=. \
      --output_basename="monet_styles"
done

# varied.sh
for i in $(seq 0 1 31);
do
  image_stylization_transform \
      --num_styles=32 \
      --checkpoint=/magenta-models/multistyle-pastiche-generator-varied.ckpt \
      --input_image=content.jpg \
      --which_styles="{$i:1}" \
      --output_dir=. \
      --output_basename="varied_styles"
done