ディープラーニングで軽量のスタイル変換を手軽に楽しむ方法

ディープラーニングでスタイル変換

 以下の記事に「17kBのスタイル変換のモデルの作り方」が書かれていたので、ちょっと試して見ました。

Creating a 17 KB style transfer model with layer pruning and quantization

 私は手抜き簡略化して、上記の記事のpruning(枝借り)やquantization(量子化)のテクニックは使っていないのですが、それでも556kBという1Mを切る軽量なスタイル変換モデルが簡単にできました。

 基本上記の記事の通りにやるだけなのですが、軽量モデルということなのでJetson Nanoでリアルタイムで変換できるソフトを作って動かして見ました。せっかくなので、動かし方を簡単に日本語で紹介します。

ディープラーニングで軽量のスタイル変換を行う方法

モデルの学習

 Google Colaboratoryのノートブックで、PCのブラウザ上で学習できます。学習済みのモデルも用意してあるので、とりあえず試したい人はここは飛ばしてもOKです。

 また、Google Colaboratory自体に関しては、以下記事参照下さい。

 以下のノートブックを上から順に実行していきます。

colab.research.google.com

 「Stylize an image!」のfiles.download('/content/fritz-models/style_transfer/data/my_style_025.h5')までを実行して、my_style_025.h5というモデルファイルをダウンロードできればOKです。

 それ以降は、さらなるモデル軽量化の処理なので、今回は実施しません(興味ある人は実施してみて下さい)。

 途中のimage_url = 'https://farm3.staticflickr.com/2907/14746369554_b783ba8d13_o_d.png' #@param {type:"string"}という箇所の画像を、好きな画像に入れ替えることで、自分の好きなスタイルへ変換するモデルを学習可能です。

Jetson Nanoの環境構築

 Jetson Nanoは、以下記事を参考にTensorFlowと単体のKeras、keras-contribをインストールして下さい。

 既にJetson Nanoを使っていて、とにかく最短で環境作りたいという方は、何も考えず以下のコマンドを順に実行してみて下さい。

$ git clone https://github.com/karaage0703/jetson-nano-tools
$ cd jetson-nano-tools
$ ./install-tensorflow.sh
$ sudo apt install libatlas-base-dev gfortran
$ pip3 install -U cython
$ pip3 install keras
$ pip3 install git+https://www.github.com/keras-team/keras-contrib.git

スタイル変換を実施

 以下のリポジトリをforkしてJetson Nano用でリアルタイムに変換できるスクリプトを用意しました。

 Jetson Nano上で以下実行すればダウンロード可能です。

$ git clone https://github.com/karaage0703/fritz-models
$ cd fritz-models
$ git checkout jetson_nano
$ cd style_transfer

 以下実行すれば、本記事の最初にあるようにリアルタイムの変換ができます。

$ python3 stylize_movie.py --model-checkpoint='example/umbrella_style.h5' -d='jetson_nano_raspi_cam'

 自分の好きなスタイルに変換したい場合は、先ほど、学習したモデルをexampleディレクトリの下に置いて、上記コマンドの'example/umbrella_style.h5'をモデル名に合わせて修正して下さい。

Macで動かす方法

 おまけ(というか自分へのメモ)です。Macでも簡単に動かせます。

 以下記事を参考に、TensorFlow、Kerasをセットアップしておきます。

 あとは、以下でkeras-contribをインストールします。

$ pip install git+https://www.github.com/keras-team/keras-contrib.git

 以下実行して、セットアップ

$ git clone https://github.com/karaage0703/fritz-models
$ cd fritz-models
$ git checkout jetson_nano
$ cd style_transfer

 以下を実行するとMacでも同様にリアルタイム変換できます。

$ python stylize_movie.py --model-checkpoint='example/umbrella_style.h5'

まとめ

 Jetson Nano(とついでにMac)で、ディープラーニングを使ったリアルタイムのスタイル変換をしてみました。

 スタイル変換自体は、3年ほど前にPC上でChainer上でやったことがあります。

 このときは、訳が分からないまま、大変な思いをして時間をかけて静止画を変換していた記憶があるのですが、今だとサクッと20分程度で軽量なモデルをJetson NanoのようなAIエッジデバイスでリアルタイムで動かせるようになったので、技術の進歩は素晴らしいですね。自分も少しくらい進歩していたら嬉しいのですが。

 多分、頑張ればスマホとかRaspberry Piとかでも動かせると思います。興味ある方は是非。そしてできたら教えて下さい(笑)

 こちらは、趣味TECH祭におまけとして展示しようかなーと考えていたのですが、スペース的にもインパクト的にもイマイチなので、多分出さないと思います。お蔵入りも勿体無いので、Webで供養いたしました(笑)一応持ってはいくつもりですので、興味ある方は直接声かけて下さい。

関連記事