MediaPipeでEdge TPU

目的

ここ最近、MediaPipeのサンプルをEdge TPUで動作させてみて、オリジナルから変えてみたこと、ハマったこと、感じたことをまとめてみる。

動機

前のブログでも書いたとおり、MediaPipeは気になっていた。

そこにGoogle Coral DevBoard向けのサンプルがでてきた。
DevBoardを持っておらず購入しようか迷ったが、USB Acceleratorでもできるのは？と思いやって見た結果、Jetson Nano + Edge TPU、Raspberry Pi + Edge TPUで動かすことができた。

できた！
MediaPipe Object Detection Edge TPU！！
JetsonNano + Edge TPU！！！
Edge TPU modelはSSDLite MobileNetEdgeTPUにしてみた！！！！👍

（Coral Dev Board Setupを参考にやってみた） pic.twitter.com/oja8VUIgyM

— nb.o (@Nextremer_nb_o) November 27, 2019

MediaPipe
Raspberry Pi 4 + Edge TPU
Object detection pic.twitter.com/rDWsuoo0KD

— nb.o (@Nextremer_nb_o) December 6, 2019

MediaPipe Object Detection and Tracking.
Raspberry Pi 4 + Edge TPU.

モデルはSSDLite MobileNetEdgeTPU。

今週悩んだけどようやく動いた👍。

ビルド手順・コードはここ。https://t.co/9CYDpOA8Qa pic.twitter.com/jsrYvvO2R8

— nb.o (@Nextremer_nb_o) December 13, 2019

メリットとデメリット

思いつくままにメリットとデメリットを記載したい。ただし、所感であり誤りも含んでいると思うので注意。

論文、ドキュメントは以下を参照。

• MediaPipe: A Framework for Building Perception Pipelines
• https://mediapipe.readthedocs.io/en/latest/scheduling_sync.html

メリット

モジュールの再利用性と非同期処理のサポートが特徴と感じる。

これはMediaPipeが目的でもあるMLアプリケーションのプロトタイプモデルをすばやく構築するために再利用性を高めることと、MLアプリケーションには必要になる非同期処理（重い推論と他の処理を分離）を意識していることからだと思われる。

• モジュールの再利用
• 非同期処理のサポート
• フロー制御のサポート

モジュールの再利用がしやすい

Calculatorとしてコンポーネントを分割することで、再利用しやすい作りとなっている。このコンポーネントをCalculatorと呼んでいる。特に推論を行う前処理、後処理は同じものを再利用したい。AnnotationOverlayCalculator（bonding boxの描画）やPacketResamplerCalculator（フレームの間引き）などがそれに当たる。あと、CPU⇔GPU間のメモリ割当もサポートしているため、AndroidのGPU delegateも簡単に構築できる。

非同期処理のサポート

一般に、推論部分は一番重い（処理時間がかかる）。

このため、GPU（もしくはAccelerator、CPUなど）で推論している間に他の処理をCPUで行うことで、トータルの処理時間を短縮させたい。一般的には、カメラ入力、推論、描画は非同期に行って描画を高速にしたい（見せたい）ということもある。しかし、一般的に非同期処理を実装するのは難易度が高い。

MediaPipeにはtimestampと呼ばれる概念で同期を取ることができる。timestampはデータ（packetと呼ばれる）が生成されるたびにインクリメントされていく。Calculatorは分散で処理していて、データ（packetと呼ばれる）に付属するtimestampによって、データの到達順が管理され、Calculator間のデータの同期が取られる。

フロー制御のサポート

推論など重い処理を担当するCalculatorは、データのボトルネックとなってしまう。適切な制御を行わないとオーバーフローやデッドロックの原因になってしまう。このため、MediaPipeは2つのフロー制御を持っている。

• The back-pressure mechanism
  ストリームでバッファリングされたPacketが制限に達すると、上位ノードの実行を抑制する。
• 専用ノードの挿入
  PacketResamplerCalculatorを使ったノードの場合は、Packetを間引きする（たとえば、フレームレートを1/10に間引くなど）。
  FlowLimiterCalculatorを使ったノードの場合は、処理が終わるまでの間のPacketを受け捨てる。

デメリット

まだ、ドキュメントやツール類が追いついていないところにデメリットを感じてしまう。このあたりは今後改善されるかもしれない。

• Bazelによるビルド
• どんなCalculatorがあるのか？
• デバッグツールがない？
• 実装がC++

Bazelによるビルド

Bazelを多少は理解していないといけないところは辛い。また、graphとBUILDの対応は手動で行う必要がある（graphからビルドするCalculatorのパス・ファイル名を割り出さないといけない）。もう少しスマートな方法がないのだろうか？

どんなCalculatorがあるのか？

これは、ドキュメントが整備されていないのが原因。今は、ソースから何があるのかを拾い出さないといけない。ドキュメントが整備されれば解消されるとは思う。

デバッグツールがない？

Calculatorが非同期で動作するということは、デバッグが大変である。論文を見る限りツールがありそうだがまだすべてが公開されていないようである。ブレークポイントを設定してのデバッグは大変かもしれない...

実装がC++

現状はアプリケーション、CalculatorともC++での実装が必要。C++は敷居が高いと思われる。少なくともアプリケーションはPythonで書きたいのでは？

Edge TPU向けのサンプルを作成したポイント

ここでは、Coral Dev Board Setup (experimental)からEdge TPU USB Acceleratorを使ったサンプルを作成したのポイントを残す。

作成したサンプルは２つ。

• Jetson Nano + Edge TPU（Object detection、Face detection）
• Raspberry Pi + Edge TPU（Object detection、Face detection、Object detection and Tracking）

Jetson Nano + Edge TPU

こちらは、ほとんどCoral Dev Board Setup (experimental)から変更はない。

Jetson Nano側にEdge TPUのRuntimeをインストールするぐらいである。OpenCVもJetpackにインストール済みのライブラリを使用すればよかった。

Raspberry Pi + Edge TPU

Raspbianは32bit OSであるため、クロスコンパイル時の指定を"--cpu=armv7a"にするぐらいである。ほぼ変更なく動作できるのは、MediaPipeの強みの一つだと思う。

Object Detection and Tracking

AndroidのデモをRaspberry Pi + Edge TPUで動かそうと思ったが、ハマってしまった...　最初の検出は動作するのだが、2回目以降の検出がされない...　原因は、Raspberry Pi（もとはCoarl）のメイン処理とPacketResamplerCalculatorにあった。

もともと、Object Detection and TrackingのGraphは、入力ストリーム（カメラ or ビデオ）を間引いて推論を行っている（トラッキングや表示は間引かない）。これは、時間のかかる・リソース消費が大きい推論を全フレームで行わず、トラッキングで補完することで、リアルタイム・リソース消費を軽くしている。このフレームの間引きをPacketResamplerCalculatorが担当し、デフォルトでは0.5フレームにしている（フロー制御の一つ）。

PacketResamplerCalculatorはデータのtimestampをチェックしてtimestampが指定したフレームレートの時間に達したら、後続のCalculatorに流すようになっている。このため、timestampは経過時間（実時間）を期待している。

timestampはデータの生成の際（メイン処理）に付与され、増加している。Raspberry Pi（もとはCoarl）のメイン処理では、timestampが+1しかされていなかった。このため、timestampの時間の流れが非常にゆっくりになってしまったのが原因だった...（timestampはμs単位なので、非常に長ーく待っていれば、2回目の推論がくる...）

  while (grab_frames) {
    // Capture opencv camera or video frame.
    cv::Mat camera_frame_raw;
    capture >> camera_frame_raw;
    if (camera_frame_raw.empty()) break;  // End of video.
    cv::Mat camera_frame;
    cv::cvtColor(camera_frame_raw, camera_frame, cv::COLOR_BGR2RGB);
    if (!load_video) {
      cv::flip(camera_frame, camera_frame, /*flipcode=HORIZONTAL*/ 1);
    }

    // Wrap Mat into an ImageFrame.
    auto input_frame = absl::make_unique(
        mediapipe::ImageFormat::SRGB, camera_frame.cols, camera_frame.rows,
        mediapipe::ImageFrame::kDefaultAlignmentBoundary);
    cv::Mat input_frame_mat = mediapipe::formats::MatView(input_frame.get());
    camera_frame.copyTo(input_frame_mat);

    // Send image packet into the graph.
    MP_RETURN_IF_ERROR(graph.AddPacketToInputStream(
        kInputStream, mediapipe::Adopt(input_frame.release())
                          .At(mediapipe::Timestamp(frame_timestamp++))));
    // ★↑ここでframe_timestampが+1されていることが原因...

    // Get the graph result packet, or stop if that fails.
    mediapipe::Packet packet;
    if (!poller.Next(&packet)) break;
    auto& output_frame = packet.Get();

    // Convert back to opencv for display or saving.
    cv::Mat output_frame_mat = mediapipe::formats::MatView(&output_frame);
    cv::cvtColor(output_frame_mat, output_frame_mat, cv::COLOR_RGB2BGR);
    if (save_video) {
      writer.write(output_frame_mat);
    } else {
      cv::imshow(kWindowName, output_frame_mat);
      // Press any key to exit.
      const int pressed_key = cv::waitKey(5);
      if (pressed_key >= 0 && pressed_key != 255) grab_frames = false;
    }
  }

このため、実時間（前回データ生成時との時間の差）を与えることで、問題が解消できたのである。修正のコードは下記。githubにもアップしている。std::chronoで前回からの時間差をとっている（ただ、この修正だと、ビデオファイル入力ではまずい...　ビデオファイル入力の場合は、フレームレートから算出しないといけない）。

  std::chrono::system_clock::time_point start;
  start = std::chrono::system_clock::now();
  while (grab_frames) {
    // Capture opencv camera or video frame.
    cv::Mat camera_frame_raw;
    capture >> camera_frame_raw;
    if (camera_frame_raw.empty()) break;  // End of video.
    cv::Mat camera_frame;
    cv::cvtColor(camera_frame_raw, camera_frame, cv::COLOR_BGR2RGB);
    if (!load_video) {
      cv::flip(camera_frame, camera_frame, /*flipcode=HORIZONTAL*/ 1);
    }

    // Wrap Mat into an ImageFrame.
    auto input_frame = absl::make_unique(
        mediapipe::ImageFormat::SRGB, camera_frame.cols, camera_frame.rows,
        mediapipe::ImageFrame::kDefaultAlignmentBoundary);
    cv::Mat input_frame_mat = mediapipe::formats::MatView(input_frame.get());
    camera_frame.copyTo(input_frame_mat);

    // Get timestamp.
    // ★ここで時間差(μs)を取得する。
    auto end = std::chrono::system_clock::now();
    frame_timestamp = std::chrono::duration_cast(end - start).count();
    
    // Send image packet into the graph.
    MP_RETURN_IF_ERROR(graph.AddPacketToInputStream(
        kInputStream, mediapipe::Adopt(input_frame.release())
                          .At(mediapipe::Timestamp(frame_timestamp))));

    // Get the graph result packet, or stop if that fails.
    mediapipe::Packet packet;
    if (!poller.Next(&packet)) break;
    auto& output_frame = packet.Get();

    // Convert back to opencv for display or saving.
    cv::Mat output_frame_mat = mediapipe::formats::MatView(&output_frame);
    cv::cvtColor(output_frame_mat, output_frame_mat, cv::COLOR_RGB2BGR);
    if (save_video) {
      writer.write(output_frame_mat);
    } else {
      cv::imshow(kWindowName, output_frame_mat);
      // Press any key to exit.
      const int pressed_key = cv::waitKey(5);
      if (pressed_key >= 0 && pressed_key != 255) grab_frames = false;
    }
  }

最後に

MediaPipeを使うとMLアプリケーションの開発がとても簡単になるように思える。

ただ、Calculator（コンポーネント）の実装やデバッグまで考えると難易度が高いようにも思える。

もう少し使われるようになるのか？2020年もウォッチしてみたいと思う。