Raspberry PiでTensorFlow Lite Support Task Libraryをやってみる。

目的

TensorFlow Lite Support Task Library (C++)をつかってRaspberry Pi4で動くカメラキャプチャのサンプルを作ってみる。

TensorFlow LiteのAPIと比べて簡単に実装できるのかを確認してみる。

動機

2020年9月ぐらいから tensorflow / tflite-support のリポジトリに気がついたり、公式ブログでもアナウンスがあったりしたので気になっていた。

ようやく年末年始の休みにやってみることにした。

TensorFlow / exampleに追加されたサンプル、task libraryってあるけど、tflite-supportとの関係はあるの？

うーん。ちゃんと見ようかな...https://t.co/D466nvVvWv

— nb.o (@Nextremer_nb_o) September 30, 2020

⬇️でやってた、TensorFlow Lite SupportのTFLite Task library - C++を使ってRaspberry Piで動くサンプル(Object detection)をアップした。

TFLite Model MetadataとVision Task Librariesで少ないコードで実現できるのが魅力。input shapeや型、labelを意識しなくていい。https://t.co/vKqUObGp8J https://t.co/0Ao7LuLzwG

— nb.o (@Nextremer_nb_o) January 3, 2021

TensorFlow Lite Support Task Libraryとは？

公式のドキュメントは日本語にも翻訳されている。

自分の理解は以下。

• TensorFlow Lite APIよりも簡単に扱うことができるAPIを用意。
• タスク（画像分類、物体検出、自然言語処理、...etc）ごとにAPIを用意。
• モデルのInput / Outputの形式（Float, INT, Shape...）を気にしなくてよい。

TensorFlow Lite APIと比べて扱いやすいAPIを用意することが目的と推測。

また、TensorFlow Lite Support Task Library は、他のライブラリを含めてTensorFlow Lite Supportと称している模様。

TensorFlow Lite Supportに含まれるライブラリは

• TensorFlow Lite Support Library
• TensorFlow Lite Model Metadata
• TensorFlow Lite Support Codegen Tool
• TensorFlow Lite Support Task Library ← 今回はこれ

があり、これらを使うことで特にモバイルのアプリの作成やデプロイを簡単にしようとする目的だと推測。

以下、簡単な特徴を記載。

これらは今後変更になる可能性があるので注意。

サポート言語

Native、Android、iOSの各プラットフォームで開発できる言語をサポートしている。

• Java
• C++ (WIP)
• Swift (WIP)

用意されているタスク

2021.01.17時点では以下がサポートされている。

画像系

• ImageClassifier（画像分類タスク）
• ObjectDetector（物体検出タスク、SSD系のモデルが動作）
• ImageSegmenter（Image Segmantationタスク、DeepLabが動作）

自然言語系

• NLClassifier
  （どんなモデルか不明。直ダウンロードリンクで詳細がよくわからなかった）
• BertNLCLassifier
  （こちらもBertのどのようなモデルまでかは詳細がTODOとなっている。）
• BertQuestionAnswerer
  （Mobile BERT Q&A modelが動作）

タスクを自作するには？

上記のタスク以外で、独自のタスクを実装することも可能。

モデル

通常のTensorFlow Lite モデルも可能だが、metadataを追加したTensorFlow Lite モデルを扱うことができる。TensorFlow Lite モデルにmetadataを追加することで

• Task LibraryがInput, Outputの違い（型、サイズ）を吸収してくれる
  （アプリがリサイズ、型変換を意識しなくてよい）
• Labelファイルが不要となる（モデルに埋め込める）。
  （モデルとラベルを一元管理でき、Task Libraryが結果からラベルを返してくれる）

といった利点が出てくる。

また、metadataにはモデルの説明や著作権などのライセンス情報も埋め込むことができる。これはOn-deviceな実行の場合、モデルが端末に配布（ダウンロード）される。ユーザーにはモデルが見えるので、ライセンスが明示できることは非常にありがたいと思う。Labelもモデルに埋め込めれば管理も容易になる。

TensorFlow Hubではmetadataが組み込まれたTensorFlow Liteモデルがある。例えば、ssd_mobilenet_v1を確認するとmetadetaにはモデルの説明、ライセンス、Input、Outputの詳細が確認できる。

metadataを追加したTensorFlow Lite モデルについては、次回以降にもう少し詳細化してみたい。

ラズパイ4でカメラキャプチャのサンプルを作る

今回は、画像のタスクのサンプル（CLI Demos for C++ Vision Task APIs）を参考に、ラズパイ4（64bit）とPiCameraでカメラキャプチャでタスクを実行するサンプル（C++）を作った。

Task Library（C++）の使い勝手を確認してみる。

元のリポジトリからForkしたリポジトリを作成。

用意したサンプル

画像で用意されているタスクのサンプルを作成してみた。

• Image Classifier
• Object Detector
• Image Segmenter

環境

自分の環境は以下。

• Raspberry Pi 4 4GB
• Raspberry Pi OS 64bit
• Raspberry Pi Camera Module V2.1（UVCカメラでもOKなはず）

ビルド環境の準備と必要なモジュールのインストール

サンプルは元リポジトリと同様Bazelをつかってビルドしている。また、カメラキャプチャするためOpenCVを利用。ビルドはクロスコンパイルしたかったが、今回はHostのラズパイでビルドする。

クロスコンパイルするにはMediaPipeのサンプルと同じようにDokcerで行なえばいいと思う。

# Install required library
$ sudo apt install git libopencv-dev

# Install build tool.
$ wget https://github.com/bazelbuild/bazel/releases/download/3.7.2/bazel-3.7.2-linux-arm64
$ chmod +x bazel-3.7.2-linux-arm64
$ sudo mv bazel-3.7.2-linux-arm64 /usr/local/bin/bazel
$ sudo apt install openjdk-11-jdk

ビルド

リポジトリをClone後、それぞれをビルド。ビルド時間はおおよそ１５～２０分程度。

# Clone repository
$ git clone https://github.com/NobuoTsukamoto/tflite-support.git
$ cd tflite-support

Image Classifier

# Build Image Classifier
$ bazel build \
    --verbose_failures \
    tensorflow_lite_support/examples/task/vision/pi/image_classifier_capture

Object Detector

# Build Image Classifier
$ bazel build \
    --verbose_failures \
    tensorflow_lite_support/examples/task/vision/pi/object_detector_capture

Image Segmenter

# Build Image Classifier
$ bazel build \
    --verbose_failures \
    tensorflow_lite_support/examples/task/vision/pi/image_segmenter_capture

実行

TensorFlow Hubからmetadataが組み込まれたTensorFlow Liteモデルをダウンロードして実行する。

Image Classifier

# Download the model
$ curl \
   -L 'https://tfhub.dev/google/lite-model/aiy/vision/classifier/birds_V1/3?lite-format=tflite' \
   -o ./aiy_vision_classifier_birds_V1_3.tflite

# Run the classification tool.
$ ./bazel-bin/tensorflow_lite_support/examples/task/vision/pi/image_classifier_capture \
    --model_path=./aiy_vision_classifier_birds_V1_3.tflite \
    --num_thread=4

Object Detector

# Download the model.
$ curl \
   -L 'https://tfhub.dev/tensorflow/lite-model/ssd_mobilenet_v1/1/metadata/2?lite-format=tflite' \
   -o ./ssd_mobilenet_v1_1_metadata_2.tflite

# Run the detection tool.
$ ./bazel-bin/tensorflow_lite_support/examples/task/vision/pi/object_detector_capture \
    --model_path=./ssd_mobilenet_v1_1_metadata_2.tflite \
    --score_threshold=0.5 \
    --num_thread=4

Image Segmenter

# Download the model.
$ curl \
    -L 'https://tfhub.dev/tensorflow/lite-model/deeplabv3/1/metadata/1?lite-format=tflite'  \
    -o ./deeplabv3_1_metadata_1.tflite

# Run the segmantation tool.
$ ./bazel-bin/tensorflow_lite_support/examples/task/vision/pi/image_segmenter_capture \
    --model_path=./deeplabv3_1_metadata_1.tflite \
    --num_thread=4

ハマったこと

WORKSPACEにOpenCVのBUILDを追加したところビルドエラーが発生。

/usr/include/c++/8/cstdlib:75:15: fatal error: stdlib.h: No such file or directory
#include_next <stdlib.h>

詳細はこのIssueにある通りで、ビルドのパラメータに「build --spawn_strategy=standalone」が指定されていると、/usr/includeのインクルードパスが追加されなくなるためである模様。

ビルドパラメータを削除して対応した。

TensorFlow Lite Support Task Libraryの使いやすさ

ある程度、サンプルを実装してみた感想。画像系のタスクでの結果なので、自然言語系のタスクは触れていないので注意。

モデルのロード（タスクの生成）

モデルのロードはそのタスクのクラスを生成することで実現する。モデルパス、閾値やthread数などのパラメータもオプションとして指定する。

物体検出タスクだとこんな感じでObjectDetectorのインスタンスを生成。

  // Build ObjectDetector.
  const ObjectDetectorOptions& options = BuildOptions();
  ASSIGN_OR_RETURN(std::unique_ptr<ObjectDetector> object_detector,
                   ObjectDetector::CreateFromOptions(options));

オプションの指定（BuildOptions関数）はこんな感じ。

ObjectDetectorOptions BuildOptions() {
  ObjectDetectorOptions options;
  // モデルパスを指定
  options.mutable_model_file_with_metadata()->set_file_name(
      absl::GetFlag(FLAGS_model_path));
  // 出力の最大数
  options.set_max_results(absl::GetFlag(FLAGS_max_results));
  // 推論でのスレッドの並列数
  options.set_num_threads(absl::GetFlag(FLAGS_num_thread));
  // スコアの閾値
  if (absl::GetFlag(FLAGS_score_threshold) >
      std::numeric_limits<float>::lowest()) {
    options.set_score_threshold(absl::GetFlag(FLAGS_score_threshold));
  }
  // 出力クラスのホワイトリスト
  for (const std::string& class_name :
       absl::GetFlag(FLAGS_class_name_whitelist)) {
    options.add_class_name_whitelist(class_name);
  }
  // 出力クラスのブラックリスト
  for (const std::string& class_name :
       absl::GetFlag(FLAGS_class_name_blacklist)) {
    options.add_class_name_blacklist(class_name);
  }
  return options;
}                 

スコアの閾値や出力数の個数はアプリ側で制御すると煩雑になりがちなのでオプションで指定できるのは便利。また、画像分類、物体検出タスクでは出力のホワイトリスト、ブラックリストが指定できる。TensorFlow Hubなどで用意されたPre-trainedモデルを使う場合で、必要なクラスを制御したいときは便利だと思う。

各クラスのオプションはprotoとして記述されている。以下を参照。

画像系

• ImageClassifierOptions
  
• ObjectDetectorOptions
  
• ImageSegmenterOptions

自然言語系にはオプションはないように見える？

入力（画像）

TF-Liteで面倒なのは画像の入力だと思う。

モデルにあわせてリサイズや型（Float、INT）の変換、標準化が必要。必要な情報はモデルから取得できるのだが、すべての形式にあわせようとするとかなり冗長なコードとなってしまう。自分はFloat、INTモデルに限らず、モデルの入力はINT8で統一してしまうことでリサイズだけ意識するようにしている。ただ、Floatモデルの場合はINT⇒Floatへのdequantizedを挟むのでほんの少しだけもったいない気がする。

Task Libraryではモデルの入力を意識する必要がない（Libraryが吸収してくれる）。

アプリは画像データをFlatBuffer形式するだけ。各入力データの形式（Gray, RGB, RGA, YUV, Raw）からFlatBufferに変換するIFが用意されている。

OpenCVでカメラキャプチャしたデータの場合、BGR⇒RGBに変換、CreateFromRgbRawBufferを使ってFlatBufferを生成すればよい。CreateFromRgbRawBufferの引数にはcv::MatのサイズとRawデータを与える。

    cap >> frame; // capture frame.
    cv::cvtColor(frame, input_im, cv::COLOR_BGR2RGB); // BGR to RGB

    // Frame in a FrameBuffer.
    std::unique_ptr<FrameBuffer> frame_buffer;
    frame_buffer = CreateFromRgbRawBuffer(input_im.data, {input_im.cols, input_im.rows});                

２，３行で入力データが生成できるのはとてもありがたい。

また、CreateFromRgbRawBufferの引数には画像の向きも指定できる（引数のFrameBuffer::Orientation orientation）。

モバイルの場合、９軸センサーの値などから、スマホの向きを考慮して画像の向きを考える必要がある。Task Libraryではアプリで画像を回転する必要がなく、ライブラリ内部で回転してくれる。

画像のデータだけでなくて、サイズや向き、フォーマットを指定するのでFlatBufferで入力データを指定するということで理解した。

推論

推論自身は各タスクのメソッド（Classify、Detect、Segment、etc...）を呼び出し、入力データのFlatBufferを指定してあげる。

さほどTensorFlow Lite APIのinvokeと変わらない。

    // Run object detection and draw results on input image.
    ASSIGN_OR_RETURN(DetectionResult result,
                     object_detector->Detect(*frame_buffer));          

入力したデータはどのように変換されるかはここに記載がある。

• RGBAやYUVなどの形式の場合はRGBに変換される。
• アスペクトを維持せず、モデルの入力サイズにリサイズ。
  （アスペクト比を維持しないので要注意）
• Orientationのパラメータによって、画像を回転して推論。

出力

画像分類、物体検出タスクの場合、それぞれ結果はClassificationResult、DetectionResult

として取得することができる。生のOutputTensorを意識する必要はない。

物体検出タスクの場合は下記のようにBoundingBoxの位置やサイズ、クラスのラベル名が取得できる。モデルにラベルが組み込まれていれば、出力クラスのindexから該当のラベルの文字列をとってくる処理も不要になる。これはほんとに便利。

（DrawCaptionはOpenCVの文字列描画を行う独自の関数）

absl::Status EncodeResultToMat(const DetectionResult& result,
                               cv::Mat& image) {
  for (int index = 0; index < result.detections_size(); ++index) {
    // Get bounding box as left, top, right, bottom.
    const BoundingBox& box = result.detections(index).bounding_box();
    const Detection& detection = result.detections(index);
    const int x = box.origin_x();
    const int y = box.origin_y();
    const int width = box.width();
    const int height = box.height();

    // Draw. Boxes might have coordinates outside of [0, w( x [0, h( so clamping
    // is applied.
    cv::rectangle(image, cv::Rect(x, y, width, height), kBuleColor, kLineThickness);

    // Draw. Caption.
    std::ostringstream caption;

    if (detection.classes_size() == 0) {
      caption << "  No top-1 class available";
    } else {
      const Class& classification = detection.classes(0);

      if (classification.has_class_name()) {
        caption << classification.class_name();
      } else {
        caption << classification.index();
      }
      caption << " (" << std::fixed << std::setprecision(2) << classification.score() << ")";
      DrawCaption(image, cv::Point(x-3, y), caption.str());
    }
  }

  return absl::OkStatus();
}           

また、Image Segmenterの場合、ループ処理でOutputとのマスクをとる必要がある。OpenCVを利用している場合はcv::Mat::forEachでループしてあげれば並列化も期待できると思う。なお、このサンプルではcolored_labelsでLabelごとのcolormap（いわゆるPASCAL VOCのcolormap）で色付けしている。

std::unique_ptr<cv::Mat> EncodeMaskToMat(const SegmentationResult& result) {
  if (result.segmentation_size() != 1) {
    std::cout << "Image segmentation models with multiple output segmentations are not "
        "supported by this tool." << std::endl;
    return nullptr;
  }
  const Segmentation& segmentation = result.segmentation(0);
  // Extract raw mask data as a uint8 pointer.
  const uint8* raw_mask =
      reinterpret_cast<const uint8*>(segmentation.category_mask().data());

  // Create RgbImageData for the output mask.
  auto seg_im = std::make_unique<cv::Mat>(cv::Size(segmentation.width(), segmentation.height()), CV_8UC3);
  auto wdith = seg_im->cols;
  seg_im->forEach<cv::Vec3b>([&](cv::Vec3b &src, const int position[2]) -> void {
    size_t index = position[0] * wdith + position[1];
    Segmentation::ColoredLabel colored_label =
        segmentation.colored_labels(raw_mask[index]);
        src[0] = colored_label.b();
        src[1] = colored_label.g();
        src[2] = colored_label.r();
    });
  
  return seg_im;
}        

感想

使ってみた感想

Native C++しか使っていないが、

使いやすい点

• TensorFlow Liteモデルの入出力の型、サイズを意識しなくてよいので実装がとても楽。
• TensorFlow Lite APIを使た場合と比べて1/2～1/3の実装で済む。
• OpenCVを使っても実装が楽。とくにInputをFlatBufferへの変換。
  （たぶんこれはAndroid、iOSの場合もそうかも？）

使いにくかった点

• Bazelを使ったビルド（これは自分が慣れていないせいもある）。
  とくにOpenCVを追加したらなぜかビルドエラー。。。
• APIのリファレンスがまだ整備されていない。
  まだ正式リリースでもない状態なので仕方がない。今後に期待。

MediaPipeとは何が違うの？

MediaPipeもTensorFlow Lite モデルを扱うことができるクロスプラットフォームなライブラリである。

MediaPipeとTask Libraryを比べた場合、

• MediaPipeは画像系のタスクに特化、Task Libraryは画像以外のタスクも可能。
• MediaPipeは推論部分だけでなくて、前処理、後処理も含めてのフレームワーク。
  作成したコンポーネントを再利用可能として、開発を容易とする。
  Task Libraryは推論部分の実装を容易とする。

で、それぞれ目的が異なると思う。

TF-Liteモデルを使ってお手軽にアプリを実装したい場合は、Task Libraryを使うほうが良いと思う。MediaPipeは簡単にという訳にはいかない。ある程度、MediaPipeのFrameworkとしての内容を理解していないと難しいと思う。

Coralとは何が違うの

Coral EdgeTPUのPyCoral API（Python）やlibcoral API（C++）もTF-Liteモデルを扱うことができる（EdgeTPU delegateだけでなくて）。どちらもTensorFlow Lite モデルをより簡単に扱うためのAPIを提供しているようにも見える。

• Task LibraryはEdgeTPU delegateができない。
• Pythonで扱うことができるAPIはPyCoralのみ。

で、EdgeTPUを扱う場合はPyCoral、libcoarl APIを扱う以外はないのが現状。２つのライブラリがわかれているのがもったいない気もするけど、、、

次は？

今回はTensorFlow Lite Support Task LibraryのNative C++を扱ってみた。楽に実装ができるのはいいねと思うけど、あとはBazelとかドキュメントが充実してくるといいなぁと思う。

つぎはTask Libraryとセットで必要になるTensorFlow Liteモデルへのmetadataの組み込みをやってみよう。