この記事はUMITRON Advent Calendar 2022 8日目の記事です

まえがき

ウミトロンでは物体検出を複数のプロダクトで活用しています。
アプリケーションの要件によっては、WebGLを使いブラウザで動作させることがあります。
今回は、比較的最近の物体検出アルゴリズムであるNanoDetをTensorFlow.jsを使いWebGLで動作させたので、その際に得た知見をまとめます。

NanoDetについて

NanoDetは高速な物体検出アルゴリズムであり、いくつかのモデルが提供されています。
CPUで推論する場合、ShuffleNetV2をバックボーンとするNanoDet-Plusが精度及び速度の観点から良さそうです。
今回はGPUを使用したいので、Custom CSP NetをバックボーンとするNanoDet-gを使います。

実際に、NVIDIA Quadro P2000を使い処理時間を比較してみました。
40フレームの映像を使い、1フレームの処理にかかった平均時間を計測しました。
ネットワーク部分の推論時間だけでなく前処理や後処理も含めて時間を計測しています。
また、参考までに高速な物体検出手法の一つであるMobileNet SSDとも比較しています。
結果は以下の通りです。

手法	処理時間[msec.]
MobileNet SSD	106
NanoDet-Puls	203
NanoDet-g	79

結果から、GPUで推論する場合、NanoDet-PulsよりもNanoDet-gの方が倍以上高速なことがわかります。
また、MobileNet SSDと比べても1.3倍程度高速なことがわかります。

モデルの学習

READMEのHow to Trainに従って学習します。
学習データのアノテーションはVOC XMLフォーマットとMS COCOフォーマットに対応しています。
NanoDet-gのデフォルトの設定ファイルはこちらです。
設定ファイルを編集した後、以下のようにtrain.pyを実行することで学習します。

python tools/train.py [config file path]

モデルの変換

CNNをWebGLで動作させる方法はいくつかありますが、NanoDetを組み込みたかったアプリケーションとの親和性の観点から、今回はTensorFlow.jsを採用しました。
NanoDetはPyTrochで実装されているため、学習済みモデルをTensorFlow.jsに変換する必要があります。
手順は以下のとおりです。

1. PyTrochからONNXに変換

READMEに記載されている通り、export_onnx.pyを使って変換できるのですが、ONNXのoperator setのバージョンを10に下げないと後のTensorFlowへの変換時にエラーがでるので、export_onnx.pyのopset_versionを11から10に変更した後、以下のように変換します。

python tools/export_onnx.py --cfg_path config/legacy_v0.x_configs/nanodet-g.yml --model_path [network ckpt path]

2. ONNXからTensorFlowに変換

onnx-tfを使って変換します。

onnx-tf convert -i [onnx file path] -o [output TensorFlow model directory]

3. TensorFlowからTensorFlow.jsに変換

tensorflowjsに含まれるtensorflowjs_converterを使って変換します。
このとき、TensorFlow.jsがint64に対応していないことによるオーバーフローを防ぐために、int32にキャストする必要があります（参考）。

tensorflowjs_converter --input_format tf_saved_model --output_format tfjs_graph_model --strip_debug_ops=* --weight_shard_size_bytes 18388608 [input TensorFlow model directory] [output TensorFlow.js model directory]

TensofFlow.jsでの推論

上記手順で変換されたモデルは、リサイズ及び正規化された画像を入力とし、クラス尤度と矩形位置が混ざったテンソルを出力します。
そのため、前処理と後処理は別途実装する必要があります。
参考までに、著者らによる前処理と後処理の実装箇所を記載しておきます。

前処理: 画像の正規化

BGRの画像を平均と分散で正規化しています（該当箇所）。
画像サイズはNanoDet-gであれば416x416です。

後処理: 矩形の位置とクラス尤度を計算

NanoDetはYOLOの様に画像をグリッドセルに分割し推定を行います。
そのため、出力は各グリッドセルにおけるクラス尤度と矩形位置になり、テンソルの形状は
[グリッドセル数, クラス数 + 出力矩形数 * 4]
となります。

こちらの関数でこのテンソルを各クラスの矩形に変換しているのですが、テンソルの中身がもう少し分かっていた方が実装を追いやすいと思うので、以下で簡単に説明します。

NanoDetはマルチスケールでグリッドセルに分割しており、グリッドセル数は学習時に指定された入力画像の解像度と各スケールのグリッドセルの分割数に依存します。
NanoDet-gの場合、入力画像の解像度はcofigのinput_size、グリッドセルの分割数はconfigのstridesで指定できます。
デフォルトの設定では、各スケールのグリッドセル数は(416 / 8)^2 = 2704、(416 / 16)^2 = 676、(416 / 32)^2 = 169、総数は3549となります。
また、クラス数はconfigのnum_classes、矩形出力数はconfigのreg_maxで指定できます。reg_max + 1個の矩形が出力されるので注意してください。
矩形の位置は、各グリッドセルの中心から矩形の左端、上端、右端、下端までの距離が尤度としてそれぞれ格納されています。
NanoDet-Plusの場合、中心位置ではなくグリッドセルの左上からの距離になっているので注意してください。
そのため、後処理で使われているget_single_level_center_point()はNanoDet-PlusのものではなくNanoDetのものを参考にする必要があります。

余談

NanoDetと同様に高速な物体検出アルゴリズムであるPicoDetも候補に挙がりました。
ONNXにエクスポートする際、後処理もモデルに含めてくれるオプションがあるので実装コストは下がるのですが、ONNX Runtime Webで動作させる際に一部のオペレーターがサポートされておらず断念しました(ONNXからTensoFlowに変換する場合は後処理を含めることはできませんでした)。
ONNX Runtime Webでのオペレーターのサポート状況はこちらで確認できます。
現在は少なくともHardSigmoidが未対応ですが、将来的に利用できるようになるかもしれません。

また、PicoDetのバックボーンはEnhance ShuffleNetというShuffleNteV2を改良したものなので、GPUで推論する場合、CPUでの推論時間とあまり差が無いかもしれません。

ウミトロンでは一緒に働く仲間を募集しております。持続可能な水産養殖を地球に実装するというミッションの元で、私たちと一緒に水産養殖xテクノロジーに取り組みませんか？

NanoDetをWebGLで動かす