Zenn
🦊

ブラウザ上ででポートレートを生成する。U^2-Net

2020/12/14に公開
JavaScript
Node.js
TensorFlow
TypeScript
機械学習
tech

はじめに

気軽にブラウザで機械学習のモデルを動かしたい、という欲求から、tensorflowjsでいろんなモデルを動かすライブラリを作っています(github)。今回は、写真から肖像画を生成する機能をブラウザで動かしてみようと思います。

たとえばこんな感じに、左の写真から右の絵を生成することが出来るようになります。なお、左の画像はいつものとおりぱくたそからお借りしました。河村さんです。いつ見てもキュート。
image

U^2-Netについて少しだけ説明

今回はU^2-Netを使用して肖像画を生成します。少し前にバズってたやつです。

機械学習で肖像画を生成する手法は、U^2-Netが参照しているAPDrawingGANをはじめとして、いくつか提案されているようです。肖像画の生成も画像生成の一つなのでなんとなくGANの延長線上でやることを考えてしまいます。つまり、APDrawingGANが行っているように、どこに目や鼻があるかなど顔全体の特徴から再構成した画像と、目や鼻の個別の特徴から再構成した画像をフュージョンして肖像画を作るという感じです。しかし、U^2-Netは全く違ったアプローチをしています。U^2-Netは、実はSalient Object Detection​という、画像内の注目している対象と背景を分離する分野の研究から出てきたものです。そのため、論文本体には全く肖像画の話は出てきていません。

U^2-Netのリポジトリを参照していただければわかりますが、後から、「あれ？これってこういうふうにも使えるんじゃない？」と思いついたと言っています。写真をまねて描いた肖像画を教師データとして使えば、セグメンテーションのモデルが作れるでしょ、と。「マジか！？」という思いと、「たしかにそうかも、、」という思いが共存してしまいます。しかし、実際動いているので、「すごい！」と言うしかありません。とはいえ、マルチスケールとマルチレイヤの特徴をうまく汲み取れるU^2-Netのネットワークだからできたのであって、他のバックエンドを使っているセグメンテーションモデルではうまく行かないのではないかと思います。（もう少しアーキテクチャなどU^2-Netの詳細が知りたい方は論文を読んでみてください。）

image
リポジトリからリンクがはられているページの和訳。

今回は、こういう意外な背景がある面白いモデルを使って肖像画を生成する機能をブラウザで動かしてみたいと思います。

やること

U^2-Netは今の所PyTorchでの実装しかありません。Tensorflowjsで動かすためには、モデルを変換する必要があります。PyTorchのモデルをTensorflowjsに変換する方法はPINTOさんの記事が詳しいです。基本的には、このPINTOさんの記事に従って変換しているためここでは多くは説明しません。PINTOさんの記事の記事をご参照ください。なお、dockerで変換処理を動かす方法は前回の記事で説明しているので、そちらも参考にしていただければと思います。

変換後は、Tensorflowjsに組み込んで完了です。

やったこと

ここでは、前回の記事を参考にdockerの環境を構築していることを前提とします。以下、作業の流れを示します。

(1) githubからU^2-Netのリポジトリをcloneして、トレーニング済みのモデルをダウンロードします。

$ git clone https://github.com/NathanUA/U-2-Net.git
$ cd U-2-Net/
$ mkdir ./saved_models/u2net_portrait/
$ curl -sc /tmp/cookie "https://drive.google.com/uc?export=download&id=1IG3HdpcRiDoWNookbncQjeaPN28t90yW" > /dev/null
$ CODE="$(awk '/_warning_/ {print $NF}' /tmp/cookie)"  
$ curl -Lb /tmp/cookie "https://drive.google.com/uc?export=download&confirm=${CODE}&id=1IG3HdpcRiDoWNookbncQjeaPN28t90yW" -o saved_models/u2net_portrait/u2net_portrait.pth

(2) 次に変換処理を行います。モデル名--model-name U2NETやトレーニング済みのモデルのファイル名--weights saved_models/u2net_portrait/u2net_portrait.pthが違うだけで、基本は前回の記事で紹介したU^2-Netでの作業と同様になります。pytorch_to_onnx.pyの出力ファイル名を前回記事と同じにしたので、以降の処理は前回記事をご参考ください。ただし、openvino2tensorflow実行時にエラーが出てSavedModelが作成できないと思います。これは、SavedModelの2Gの壁に引っかかってしまっているからと思われます（issue）。FrozenModelは生成されますので、このモデルからtensorflowjsにコンバートしてください。なお、全部で150MB以上になります。

$ export PYTHONPATH=/U-2-Net
$ SIZE=512
$ python3 /opt/intel/openvino_2021/deployment_tools/tools/model_downloader/pytorch_to_onnx.py  \
	--import-module model.u2net   \
	--model-name U2NET   \
	--input-shape 1,3,${SIZE},${SIZE}   \
	--weights saved_models/u2net_portrait/u2net_portrait.pth   \
	--output-file u2netp_${SIZE}x${SIZE}.onnx   --input-names "x"   \
	--output-names "a/F.sigmoid(d0)"

(3) 変換が完了したらウェブサイトに配置してTensorflowjsでロードして使用します。基本的な使い方は他のモデルと同じです。精度の高い画像生成を行うためには、入力や出力のノーマライズや閾値設定が必要となりますのでお気をつけください。私が使っている閾値等は下記のリポジトリで公開しているソースコードからご確認ください。（まだ試行錯誤中なので。）

動作確認

それでは、これを使って肖像画を生成してみましょう。

いつものようにインプットする画像のサイズを変えて処理時間と品質を見てみます。なお、ここでは、GeForce GTX 1660を搭載しているLinux PC&Chromeで実験しています。

320x320だとこのような感じになります。 結構しっかりなポートレートになってますね。処理時間はGPUで1フレームあたり1.5秒くらいです。重い。
image

256x256にしてみました。結構崩れちゃいますね。処理時間は1.1秒くらい。
image

192x192だと少し実用に耐えない感じですね。処理時間は１フレームあたり0.4秒程度でした。
image

512x512にします。もともとのトレーニングが512x512でやられているようなので、さすがとてもよくできていますね。目元の感じがとても魅力的に描けています。残念ながら、私の環境だとGPUのメモリが足りなくて処理できませんでした。CPUで処理させて１フレームあたり4分くらいかかります。うーん、つらい。
image

ためしに1024x1024も試してみました。これもとても良くできています。ただ、512x512も良かったので、人によってどっちがいいか評価が別れるかも知れません。ちなみに、こちらもCPUで実施したのですが、実施して昼寝して起きたらできてました。おそらく、１時間位かかってると思います。
image

以上です。肖像画のクオリティはかなり高いけど処理時間がかかりすぎてるかな、と思います。

リポジトリとデモ

今回作成したデモは下記のリポジトリに格納してあります。また、下記のサイトでデモが動かせます。

  1. とても遅いです。最初の画像が出るまで１分位は待ってください。
  2. 192x192のCPUで動く設定にしています。右上のコントローラで解像度とGPUへの切り替えができます。
  3. useTFWasmBackendをOffにしてreload modelボタンを押すとGPUモードになります。
  4. modelPathのプルダウンメニューから解像度を選んでreload modelボタンを押すとその解像度で処理されるようになります。
  5. inputから入力する画像を変えられます。カメラも使えます。
  6. 解像度を512x512以上に上げるときはGPUモードにしないほうが良いと思います。メモリ溢れの危険性があります。

まとめ

ブラウザ上で肖像画を作成する機能を動かしてみました。GPU搭載マシンでもリアルタイムではちょっと厳しそうです。このくらいの負荷がかかるものだと、さすがにブラウザ上で動かさずサーバで処理をするほうが良いかもしれません。

追記

解像度順に並べ変えた。自分的には512が一番いい感じかな。
(192->256->320->512->1024)
image
image
image
image
image

Discussion