🦊

ブラウザ上で顔のパーツを検出する。

2020/11/03に公開

JavaScript

はじめに

世にあるいろいろな機械学習モデルをブラウザ上で簡単に動かせるようにしたいぞ！ということで、少し前からtensorflowjs + webworkerで動くnpmパッケージを作成する作業をコツコツとしています[link]。前回はWhite-box-Cartoonizationのモデルを使ってブラウザ上で画像をアニメ化するお話を書きました。

今回は、ブラウザ上で顔のパーツを検出してみようと思います。こんな感じのものです。なお、左の絵は、https://thispersondoesnotexist.com/ で作成したものです。

今回やりたいこと（と、やらないこと）

顔のパーツを検出するAIモデルはいくつか種類があります。
ランドマークを検出するモデルは以前の記事で簡単にご紹介しました。なので、今回はこちらは扱いません。

今回は、上記の"はじめに"で示したように顔のパーツをセグメンテーションしてみるモデルを扱います。使用するモデルはBiseNetV2です。この表を見ると超高速で動くようです。

Specifically, for a 2,048×1,024 input, we achieve 72.6% Mean IoU on the Cityscapes test set with a speed of 156 FPS on one NVIDIA GeForce GTX 1080 Ti card, which is significantly faster than existing methods, yet we achieve better segmentation accuracy.

GPUを使っているとはいえ、2048x1024で156FPSってほんまでっか？とは思ってしまいますが。。。さて、ブラウザ上ではどのくらいのスピードが出るのでしょうか。今回は非公式実装ですがtensorflowjsへのコンバートの簡単さから、tensorflowで実装されているBiseNetV2-Tensorflowを使いたいと思います。

やること

基本的には次のことをします。

リポジトリをclone
チェックポイントをダウンロード
チェックポイントをfreeze
tensorflowjsにコンバート

ただ、上記リポジトリのfreeze用のスクリプトはなぜかoptimizationがコメントアウトされている^[1]ので、それを戻してあげる必要があります。また、inputの解像度(shape)が448x448に固定されているので、これを可変にしたいと思います。

やったこと

(1) まずは、githubからBiseNetV2-Tensorflowをcloneしてきます。

git clone https://github.com/MaybeShewill-CV/bisenetv2-tensorflow.git

(2) 次にReadmeに記載のトレーニング済みのチェックポイントをダウンロードしてきます。今回はcheckpointフォルダを作成し、そこにファイルを入れることにします。

$ ls checkpoint/ -lah
-rw-r--r--  1   33M 11月  3 06:24 celebamaskhq.ckpt.data-00000-of-00001
-rw-r--r--  1   36K 11月  3 06:24 celebamaskhq.ckpt.index
-rw-r--r--  1   11M 11月  3 06:24 celebamaskhq.ckpt.meta
-rw-r--r--  1    43 11月  3 06:24 checkpoint

(3) 次に、チェックポイントをfreezeします。
なお、以下では、いろいろとコードを変更するので面倒な方は、こちらのリポジトリをcloneして使ってください。Forkして変更し終わったファイルをpushしてあります。(3-1)まで読み飛ばしてください。

freezeはtools/celebamask_hq/freeze_celebamaskhq_bisenetv2_model.pyを使って行うのですが、2020/11/03時点ではoptimizationがコメントアウトされています。これを解除してあげます。下記の部分です。

    # optimize_inference_model(
    #     frozen_pb_file_path=args.frozen_pb_file_path,
    #     output_pb_file_path=args.optimized_pb_file_path
    # )

また、このスクリプトでは、inputの解像度(shape)が448x448に固定されているので、任意の解像度を受け入れるように変更します。

    input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 448, 448, 3], name='input_tensor')
    -> input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, None, None, 3], name='input_tensor')

この変更により、未知のShapeのTensorのままネットワークを構築することになるので、一部ネットワークの定義の処理が動かなくなります。その部分も修正してあげましょう。具体的にはbisenet_model/bisenet_v2.pyを修正します。name='semantic_upsample_features'とname='guided_upsample_features'のところが未知のshapeに対応できていないので、tf.shapeで実行時のshapeを参照できるようにしてあげます。

                x_shape = tf.shape(detail_input_tensor)                  # <------ here
                semantic_branch_upsample = tf.image.resize_bilinear(
                    semantic_branch_upsample,
                    x_shape[1:3],                                        # <------ here
                    name='semantic_upsample_features'
                )

あとは、upsamplingのところも対応していないようなので、そこも直します。

        input_tensor_size = input_tensor.get_shape().as_list()[1:3]
        -> input_tensor_size = tf.shape(input_tensor)
        output_tensor_size = [int(tmp * ratio) for tmp in input_tensor_size]
        -> output_tensor_size = [tf.multiply(input_tensor_size[1],ratio), tf.multiply(input_tensor_size[2],ratio)]

(3-1) さて、この変更後のファイルを使ってfreezeします。

$ python3 tools/celebamask_hq/freeze_celebamaskhq_bisenetv2_model.py \
    --weights_path checkpoint/celebamaskhq.ckpt \
    --frozen_pb_file_path ./checkpoint/bisenetv2_celebamask_frozen.pb \
    --optimized_pb_file_path ./checkpoint/bisenetv2_celebamask_optimized.pb

(4) 次に、tensorflowjsにコンバートします。

$ tensorflowjs_converter --input_format tf_frozen_model \
    --output_node_names final_output \
    checkpoint/bisenetv2_celebamask_optimized.pb \
    webmodel

以上で、tensorflowjsのモデルの作成が完了しました。

なお、(3)でも言及しましたが、ソースコードを触るのがそこそこ面倒だったりするので、難しく感じられる方はこのリポジトリをcloneしてReadmeに従って作業して頂くといいと思います。

動作確認

それでは、これを使って顔のパーツを検出してみましょう。

インプットされる画像のサイズによって処理時間と品質が変わりますので、いくつかサイズを変えてやってみたいと思います。なお、ここでは、GeForce GTX 1660を搭載しているLinux PC&Chromeで実験しています。

まずは448x448の画像をインプットにした場合です。約12FPSくらいです。あれ！？思ったよりだいぶ遅い。
精度はそこそこありそうですね。

次に1028x1028の画像をインプットしてみました。うわ、、遅い。0.5FPSになってしまいました。
精度も448x448と大きく変わりませんでした。

次に288x288で試してみます。20FPSくらい出ているようです。ようやく現実的な速度になってきました。
しかし、やはり少し精度が落ちますね。残念。

ということで、今回の動作確認では、速度の面で大幅に期待を下回る結果となってしまいました。やはりデフォルトの448x448で使うのが良さそうです。
デモは下記のページにあるので、是非お試しください。

今回使ったデモのソースは下記のリポジトリに置いてあります。今回のモデルをwebworkerとして利用するnpmモジュールも同リポジトリで公開していますので、ご活用ください。

なお別の実装(pytorch)のリポジトリに性能評価が出てましたが、やはりFP32だと16FPSくらいしか出ていないようです。inputが1024x1024ではありますが。いずれにせよ、156FPSには程遠いですね。私が、なにか勘違いしているのかもしれません。。。

動作確認(追加)

せっかくなので、前回と同じく、ThinkpadとMacBookでも確認してみます。
Mac Bookのスペックはcorei5 2.4GHzのものです。Safariで動かします。上記デモでお試しになる場合は、SafariはWebworker上でWebGLが使えないので、右上のコントローラからprocessOnLocalをONにして、reload modelボタンを押してください。
また、ThinkPadのスペックはcorei5 1.7GHzのものです。

inputを448x448にしたときMacBookでは0.8程度でした。Thinkipadは0.4程度でした。リアルタイムで使うには厳しいですね。

まとめ

ブラウザ上で顔のパーツを検出するチャレンジをしてみました。GPU搭載のPC上であれば、そこそこの性能でリアルタイムに検出できそうです。GPU非搭載だと、かなり厳しそうです。今後のPCやスマホの性能向上に期待したいです。

リファレンス

画像は下記のページのものを利用させていただきました。
https://www.pakutaso.com/

脚注

2020/11/03時点(commit 710db8646ceb505999b9283c0837c0b5cf67876d ) ↩︎