Edited at

Jetson Nanoでリアルタイムに物体検出をする方法(TensorFlow Object Detection API/NVIDIA TensorRT)

jetson_detection.gif


Jetson Nanoでの物体検出

 Jetson Nanoでディープラーニングでの画像認識を試したので、次は物体検出にチャレンジしてみました。そこで、本記事では、TensorFlowの「Object Detection API」と「Object Detection API」を簡単に使うための自作ツール「Object Detection Tools」を活用します。

 これらに関して詳細は、以下ブログ記事を参照下さい。Jetson Nanoでの物体検出自体は、本記事で完結するのでブログ記事は読まなくても大丈夫です。

TensorFlowの物体検出用ライブラリ「Object Detection API」を手軽に使えるソフト「Object Detection Tools」を作ってみた

 これ以降の本記事の作業は、全部Jetson Nanoで行います(Jetson Nano自体のセットアップ作業を除く)。


Jetson Nanoで最も簡単に物体検出する方法

 まずは、以下記事をもとにJetson NanoのTensorFlowのセットアップまでを実施します。

Jetson Nanoをセットアップしてディープラーニングで画像認識試してみた

 具体的には、Jetson Nanoが起動するところまでできたら、以下4行のスクリプト実行するだけです。TensorFlowやOpenCV、Matplotlibなどの必要なライブラリがインストールされます。

$ git clone https://github.com/karaage0703/jetson-nano-tools

$ cd jetson-nano-tools
$ ./install-tensorflow.sh
$ ./install-tools.sh

 また、以下記事のSwapファイルの設定も実施しておいて下さい。これを実施しないと思い処理(モデルの変換等)を実施したとき、プログラムが落ちる可能性があります。

Jetson Nano関係のTIPSまとめ

 あとは、SSD Lite Mobilenet V2という軽いモデルを使えば、そのまま物体検出のデモを動かすことができます。具体的には以下5行のコマンドを実行するだけです。

 使用しているプログラムは「Object Detection API」のサンプルプログラムをベースに、私が作成したものです。

$ cd && git clone https://github.com/karaage0703/object_detection_tools

$ cd ~/object_detection_tools/models
$ ./get_ssdlite_mobilenet_v2_coco_model.sh
$ cd ~/object_detection_tools
$ python3 scripts/object_detection.py -l='models/coco-labels-paper.txt' -m='models/ssdlite_mobilenet_v2_coco_2018_05_09/frozen_inference_graph.pb' -d='jetson_nano_raspi_cam'

 -lでラベルファイルの指定。-mでモデルファイルの指定。-dで使用するカメラの指定をしています。今回は、Jetson NanoにRaspberry Piカメラモジュールv2を繋いで使用していますが、-djetson_nano_web_camを指定すれば、USBのWebカメラでも使用可能です。

 実行結果は、以下のようになります(ソフトのバージョンによって、若干表示は異なります)

jetson_detection.gif


NVIDIA TensorRTを使ったモデルの最適化

 物体検出を更に高速化したい場合や、もう少し大きいSSD Mobilenet V1を使いたい場合は、NVIDIA TensorRTという仕組みを使って、モデルを最適化することで高速化することが可能です。この場合、Jetson NanoでObject Detection APIの環境構築とモデル変換が必要になるため、少し手順が複雑になります。

 順に説明していきます。


Jetson Nanoへの「Object Detection API」環境構築

 環境構築は、「Object Detection API」の公式のセットアップ情報を参考に、Jetson Nano用に修正したものとなります(Jetson NanoはArmコアなので公式情報そのままではセットアップできません)。あくまで非公式情報ですのでご了承下さい。

 次に、TensorFlowの「Object Detection API」を含むリポジトリをホームディレクトリにクローンします。ファイルサイズ大きいので結構時間かかります。

$ cd && git clone https://github.com/tensorflow/models

 Protocol Bufferをインストールしてセットアップします。

$ cd ~/models/research

$ wget -O protobuf.zip https://github.com/protocolbuffers/protobuf/releases/download/v3.7.0/protoc-3.7.0-linux-aarch_64.zip
$ unzip protobuf.zip
$ ./bin/protoc object_detection/protos/*.proto --python_out=.

 以下でObject Detection APIの動作テストをします。

$ cd ~/models/research

$ export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
$ python3 object_detection/builders/model_builder_test.py

 OKという表示が出たら準備は完了です。


NVIDIA TensorRTを使ってモデルファイルを最適化する

 次に、学習モデルを最適化します。先ほど普通のPCではサクサク動くような学習モデルも、Jetson Nanoでそのまま動かそうとするともの凄い遅く、使い物になりません。Jetson Nanoでリアルタイムに動かすには、モデルをNVIDIA TensorRTというNVIDIAの環境を用いて、モデルファイルを最適化する必要があります。NVIDIA TensorRTに関して詳しくは公式サイトを参照下さい(私も詳細理解できていません…)

 まずは、最適化する学習モデルをダウンロードしましょう。ちなみにJetson Nanoで最適化できるモデルは、私の環境ではmobilenet等の小さいモデルのみでした(ssd_inception_v2等のモデルで試したら、GPUがnvinfer1::OutofMemoryエラーになりました)。

 あらかじめ、「Object Detection API」のパスを通しておきます。

$ cd ~/models/research

$ export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

 続いて、Object Detection Tools(自作のツール)と学習モデルをダウンロードします。以下3行でOKです。最初の物体検出のところで、既にObject Detection Toolsをダウンロード済みの場合は、最初の1行は省略して下さい。

$ cd && git clone https://github.com/karaage0703/object_detection_tools

$ cd ~/object_detection_tools/models
$ ./get_ssd_mobilenet_v1_coco_model.sh

 続いて「NVIDIA TensorRT」をダウンロードします。

$ cd && git clone https://github.com/NVIDIA-AI-IOT/tf_trt_models

$ cd tf_trt_models

 自作の変換ツールを用いて、学習モデルをJetson Nanoに最適化したモデルに変換します。

$ cp ../object_detection_tools/scripts/convert_rt_model.py .

$ python3 convert_rt_model.py -c='../models/research/object_detection/samples/configs/ssd_mobilenet_v1_coco.config' -m='../object_detection_tools/models/ssd_mobilenet_v1_coco_2018_01_28/model.ckpt' -o='./frozen_inference_graph_trt.pb'

 -cオプションが使用するコンフィグファイル、-mが最適化する対象の学習モデル -oが最適化した後の学習モデルのファイル名となります。


Jetson Nanoでの最適化したモデルでのリアルタイム物体検出

 最適化したモデルで物体検出を行います。

$ cd ~/object_detection_tools

$ python3 scripts/object_detection.py -l='models/coco-labels-paper.txt' -m='../tf_trt_models/frozen_inference_graph_trt.pb' -d='jetson_nano_raspi_cam'

 もし、ここで以下のようなエラーがでたら

NotFoundError: Op type not registered 'TRTEngineOp' in binary running on your-jetson. ...

 以下のように、object_detection.pyに1行追記をして下さい。

import tensorflow as tf

import tensorflow.contrib.tensorrt as trt # <-- add

 このエラー対処は @brownbro さんにコメントで指摘いただいたものです。感謝です。


まとめ

 Jetson Nanoでリアルタイムの物体検出をしました。簡単にできるかと思ったのですが、TensorFlowの「Object Detection API」が思ったより難しく大変で、結局ツールを自作することになってしまいました。

 PCでもJetson Nanoのようなモバイルデバイスでも、モデルファイルを変換するだけで、同じプログラムで学習から推論までできるというのは素晴らしいですね。これって、意外に今まではできそうで出来ていなかったことと思います(大体、デバイスに特化したプログラムを学習でも推論でも用意する必要があるので)。NVIDIAのデファクトの強さを感じますね。

 次は、自前データの学習方法に関しても、まとめる予定です。


参考リンク

 今回も何度も絶望したのですが、その度に多くの先人の知恵に助けられました。感謝です。参考にした情報にリンクいたします。

 特に @PINTO さんと @tsutof さんの情報なければ、実現できなかったと思います感謝です。


関連記事

Jetson Nano関係のTIPSまとめ

Jetson Nano関係の情報はこちらにまとめています