M5Stack＋TensorFlow Liteでジェスチャ認識

TensorFlow Lite for Microcontrollersのサンプルには、前回試したHello worldとその前に試した音声認識の他に、「Magic Wand」があります。

これは、加速度センサの入力から3種類のジェスチャを認識するものです。
ジェスチャは「WING（ｗを描く）」「RING（丸を描く）」「SLOPE(「L」を描くような感じで右斜め上から左下へ移動、そのあと右へ水平移動）」となっています。

以下のビデオに、杖の先にマイコンを付けた「魔法の杖」で、ジェスチャに応じた色のLEDを点灯させるデモがあります。

新しいジェスチャのモデルを学習するためのスクリプトも、最近公開されました。

tensorflow/tensorflow/lite/micro/examples/magic_wand/train at master · tensorflow/tensorflow

M5Stack Fireの内蔵の加速度センサを使って、このMagic Wandを動作させてみました。
元のサンプルではArduino Nano 33 BLE Sense、およびSparkFun Edgeがターゲットになっていますが、それ以外のデバイスへの移植としては、Adafruit EdgeBadgeやParticle（STM32ベースのモジュール）の事例がネット上で公開されています。

Overview | TensorFlow Lite for EdgeBadge Quickstart | Adafruit Learning System

Use Tensor Flow Lite and a Particle Xenon to build the ML-gesture wand of your dreams – Particle Blog

これらを参考に、M5Stack FireにMagic Wandデモを移植していきます。

【2019/12/28追記：ソースコード一式を以下にアップロードしました。】

boochow/TFLite_Micro_MagicWand_M5Stack: M5Stack (ESP32) port of TensorFlow Lite for Microcontrollers demo "Magic Wand"

TensorFlowを使った推論処理

まず、そもそもTensorFlowは何をしてくれるのかを簡単に説明しておきます。
TensorFlowそのものについては、以下のドキュメントが分かりやすいのでご覧ください。

TensorFlowの使い方 DevFest 2019 – Google スライド

TensorFlow Lite for Microcontrollers(以下、TFLMと略記）では、学習は行わず、推論のみを行います。
推論とは、モデルに入力を与えて出力を得ることで、モデルとは事前にデータで訓練されたニューラルネットワークです。

モデルは、ニューロン間をつなぐ「重み」と、ニューロン内で行われる「演算」のセットです。
モデルを用いた推論処理は、TFLMではライブラリがインタープリタを提供しているので、それを使って行います。

モデルへの入力と出力はそれぞれテンソル（＝多次元の行列（＝多次元ベクトルを並べたもの））です。
入力データをモデルに合うように整形すること、出力データを解釈することはプログラム側のタスクです。

Magic Wandデモでは訓練済みのモデルが用意されていますので、モデルの中身はブラックボックスとして扱えます。
とはいえ、処理自体はソフトウェアで行うので、メモリ使用量や処理の内容を無視できるわけでもありません。
このデモのモデルは、前述の新しいジェスチャを学習するためのスクリプトの中で以下のように記述されています。

model.summary()の出力は以下のとおりです。

Model: "sequential_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d_2 (Conv2D)            (None, 128, 3, 8)         104       
_________________________________________________________________
max_pooling2d_2 (MaxPooling2 (None, 42, 1, 8)          0         
_________________________________________________________________
dropout_3 (Dropout)          (None, 42, 1, 8)          0         
_________________________________________________________________
conv2d_3 (Conv2D)            (None, 42, 1, 16)         528       
_________________________________________________________________
max_pooling2d_3 (MaxPooling2 (None, 14, 1, 16)         0         
_________________________________________________________________
dropout_4 (Dropout)          (None, 14, 1, 16)         0         
_________________________________________________________________
flatten_1 (Flatten)          (None, 224)               0         
_________________________________________________________________
dense_2 (Dense)              (None, 16)                3600      
_________________________________________________________________
dropout_5 (Dropout)          (None, 16)                0         
_________________________________________________________________
dense_3 (Dense)              (None, 4)                 68        
=================================================================
Total params: 4,300
Trainable params: 4,300
Non-trainable params: 0
_________________________________________________________________

コードの準備

まず、ローカルにクローンしたTensorflowのリポジトリから、ESP32用Magic Wandの元となるコードを生成します。
前回はr2.1のブランチを使用しましたが、最近TFLMのディレクトリがtensorflow/lite/experimental配下からtensorflow/lite直下に移動したこともあり、今回はmainブランチを使います。

以下のように、makeでソースコード一式を生成し、M5Stackの開発環境へ持ってきやすいようにzipでまとめておきます。

$ make -f tensorflow/lite/micro/tools/make/Makefile TARGET=esp generate_magic_wand_esp_project
$ cd tensorflow/lite/micro/tools/make/gen/esp_xtensa-esp32/prj/magic_wand/esp-idf/
$ ls -CF
CMakeLists.txt  LICENSE  README_ESP.md  components/  main/
$ zip -r esp32mw.zip components main

M5Stackの開発環境は、PlatformIOを使います。プラットホームはArduinoです。
PlatformIOでM5Stackを使用するための手順は以前書きましたので省略します。

さきほどzipファイルにまとめたコードを展開し、「components/tfmicro」フォルダを「lib」配下へ、「main」フォルダの中身を「src」配下へ移動します。

このままだとsetup()およびloop()がかぶってしまうので、PlatformIOが自動生成した「main.cpp」、TFLM側の「main.cc」「main_functions.h」は削除し、main_functions.ccをmain.cppへリネームしてください。
また、main.cppの中の#include "main_functions.h"は削除します。

ファイル構成は以下のようになります。

tfliteフォルダ配下にインクルードパスを通すため、「platformio.ini」を編集して、以下のようにインクルード設定を追加します。

ビルドするといくつかエラーが出ますので、直していきます。

tfmicro/third_party/flatbuffers/include/flatbuffers/base.h

35行目～39行目の

#if defined(ARDUINO) && !defined(ARDUINOSTL_M_H)
  #include <utility.h>
#else
  #include <utility>
#endif

でエラーが出ますが、#else節のほうを使って欲しいので、platformio.iniを編集して-DARDUINOSTL_M_Hをコンパイルオプションに追加します。
ついでにシリアルポートの設定も115200bpsにしておきます。

tfmicro/tensorflow/lite/kernels/internal/reference/concatenation.h

125行目～126行目の

static_cast(std::round(input_ptr[j] * scale + bias)) + output_zeropoint;

でstd::round()が未定義エラーになります。
これarduino-esp32のエラーっぽいのですが、とりあえずArduinoのround()マクロを使うように変更（std::round()→round()）します。

M5Stack Fireへの移植

以上でビルドが通るようになったので、M5Stack用の入出力のコードを追加していきます。

このデモのコードは、加速度センサの読み取りをaccelerometer_handler.cc、検出結果の処理をoutput_handler.ccへ実装するような構成になっています。初期状態では、前者はダミーコード、後者は結果をテキストでシリアルポートへ出力するコードになっています。

accelerometer_handler.cc

以下のAPIを実装します。

SetupAccelerometer()は初期化を行います。

ReadAccelerometer()はX、Y、Z方向の加速度を25Hzの頻度で計測し、結果を、X、Y、Zそれぞれmg単位（計測結果の1000倍）で時系列で内部のリングバッファに記録していきます。リングバッファの位置はbegin_indexが保持しています。
そして、リングバッファから配列inputへデータを最新のものからlength個コピーします。

TFLMのサンプルにはArduino用の実装例が含まれていますので、これを参考にしながら実装していきます。

注意しなければならないのは座標系で、これはモデルを訓練するときの座標系と一致していないとジェスチャを正しく認識できません。
これに関して、Adafruitのソースコードに以下のようなコメントがありました。

/* this is a little annoying to figure out, as a tip - when
 * holding the board straight, output should be (0, 0, 1)
 * tiling the board 90* left, output should be (0, 1, 0)
 * tilting the board 90* forward, output should be (1, 0, 0);
 */

つまり、M5Stackのディスプレイを手前に向けて立てたときx, y, zが(0, 0, 1）、左へ90度傾けたとき(0, 1, 0）、ディスプレイを上に向けて机に置いた状態のとき(1, 0, 0)となるようにする必要があります。
M5Stackの場合は、加速度センサのデータを(z, -x, -y)の順に並べることでこの条件を満たすことができます。

また、時間軸については、加速度計測のサンプリングレートがconstants.hの中で

const float kTargetHz = 25;

と定義されていますので、25Hzでのサンプリングが望ましいようです。

モデルの入力は連続した128個×座標軸3つ分の加速度データで、合計384個です。
サンプリングレートは25Hzなので、計測開始から5秒くらいはデータが不足することになります。
また、少なくとも128×3個分のデータを内部バッファに保存できることが必要です。

以上を踏まえ、accelerometer_handler.ccは以下のように実装しました。

M5Stackの加速度センサは、PlatformIOのM5Stack用ライブラリを利用して読み取ることが出来ます。

他のプラットホームの実装では、加速度センサにLIS3DHを使っていますが、M5StackはMPU9250を使っています。LIS3DHはFIFOを内蔵していて、自動でデータを周期的にFIFOへ貯めていくようです。

M5StackのライブラリはFIFOはサポートしていないっぽいので、毎回データを１つだけ計測しています。
また、サンプリングレートは25Hzが望ましいらしいので、前回計測したタイミングから40msec経過していないときは計測しないようにしました。

内部バッファを初期化した直後は、推論に必要なデータが十分蓄積されていません。この状態はpending_initial_dataフラグで表し、trueのときは加速度計測自体は行いますが推論は始まりません。
内部バッファ初期化は、起動時の他、何らかのジェスチャを認識した直後にも行います。（でないと同じジェスチャを何度も認識してしまうため）

後で書くように、認識したジェスチャをM5StackのLCDに表示するようにしたので、その後初期化したバッファに十分データがたまった段階（pending_initial_dataがfalseになるタイミング）で、LCDの表示を消去するようにしました。

main.cpp

M5StackおよびシリアルポートとI2Cバスの初期化を行います。I2Cバスは400KHzで動作するようにしています。

setup()の冒頭に、以下の3行を追加します。

  M5.begin();
  Serial.begin(115200);
  Wire.begin();
  Wire.setClock(400000);

Tensorflow関係のインクルードファイルの後に、以下のインクルードファイルを追加します。

#include <Arduino.h>
#include <M5Stack.h>

Arduino.hをインクルードするのは、TensorFlow関係のインクルードファイルよりも後にしてください。Arduino.hにはmin、maxなど様々なマクロが定義されているため、TensorFlow関係のインクルードファイルよりも前にインクルードしてしまうとコンパイルエラーを引き起こします。

メインループは以下のようにしました。

元のコードに追加したのは、16行目からのデバッグ用の出力です。
model_input->data.f[381..383]は、最新の加速度データが入っています。
interpreter->output(0)->data.f[0..2]は、3つのジェスチャが行われた可能性を示す推論結果（0.0～1.0、３つの合計が1.0）が入っています。

output_handler.cc

認識したジェスチャに応じた処理を行います。

デフォルトでは結果をシリアルポートへ出力するようになっていますが、M5StackのLCDに結果を表示するコードを追加しました。
以下がコードの全体ですが、特に説明を要するところは無いと思います。

テストと調整

以上でM5StackでMagic Wandデモが動くようになりましたが、ジェスチャ認識させるには結構コツが必要な感じです。
シリアルポートには加速度の値とジェスチャ検出結果をリアルタイムに出力していますので、それを見ながら少し設定を調整してみました。

ジェスチャ認識の閾値の設定

ジェスチャは時系列の信号なので、ジェスチャらしきものが一定時間続いてはじめてジェスチャとして認識されます。
イメージとしては、以下のような画像認識を行うような感じです。

・3軸の加速度が、それぞれ1ピクセルの明暗で表現されている。3つのピクセルが横に1行で並んでいる。
・40msecごと（25Hz）に、新しい行が追加されていく。縦に細長い画像ができていくイメージ。
・縦の長さが128行を超えると、スクロールして一番古い行が消える。
・この3×128ピクセルの画像の中から、ジェスチャを行ったときのパターンを探す。

ジェスチャの誤検出は起こりえますので、ジェスチャが実際にジェスチャとして認識されるのは、ジェスチャが連続して数回検出され、閾値を超えた後です。
この閾値はconstants.hに記載されていますが、値が大きいとジェスチャがより認識されにくくなります。また、異なるジェスチャには異なる閾値が設定できます。

元のファイルでは、Wing/Ring/Slopeのそれぞれの閾値が {15, 12, 10}となっていました。
Arduino用のconstants.hではもう少し敷居が下がって{8, 5, 4}となっています。
今回はカット＆トライの結果、閾値を{9, 7, 6}としました。

モデルの交換

動かしてみて、どうもあまり検出精度が思わしくなかった（誤検出が多いことも含め）のですが、調べてみるとモデルのデータ（magic_wand_model_data.cc）は最近更新されていることが分かりました。
冒頭に書いた、新しいジェスチャの訓練スクリプトの公開時に、新しいデータになったようです。

試しに、古いデータをr2.1のブランチから持ってきて動かしてみると、どうもこちらのほうが検出精度が高いように思えました。
このあたりは、個人の癖やデバイスとの相性もあるのかもしれません。