TFLiteを使ったAndroid上でのセマンティックセグメンテーション Part 2

この記事は弊社のAbirがMediumに投稿した以下の記事を日本語訳したものです。

このシリーズのパート1では、Pasal-vocデータセットを使ってDeepLab-v3モデルを学習し、そのモデルを257x257の入力サイズで frozen_inference_graph.pbファイルとしてエクスポートする方法を学びました。今回は、この凍結グラフを画像セグメンテーション用のandroidアプリで使えるTfLiteモデルに変換する方法を紹介します。

1. インストール前提条件

TensorFlow v2.2.0
Numpy

$ pip install tensorflow==2.2.0
$ pip install numpy

ここでは、パート1で学習したpbモデルを使用します。または、このbash スクリプトを使って同じモデルをダウンロードすることもできます。スクリプトをダウンロードして、プロジェクトのrootで実行してください。

$ bash download.sh

モデルを含むzipファイルをダウンロードし、プロジェクトルートに展開します。

2. モデル変換のための読み込み操作

では、convert.py という名前の python スクリプトを作成し、必要なパッケージをインポートしてみましょう。

gist.github.com

次に、最後のステップでダウンロードしたモデルのパスを指定します。

MODEL_FILE = "frozen_inference_graph_257.pb"

TensorFlow TFLite コンバータを使用して我々のTensorFlowモデルのグラフ定義 frozen_inference_graph_257.pbをmodel.tfliteに変換しましょう。次の関数を使ってTensorflowモデルをコンバーターに読み込ませましょう。

gist.github.com

input_arraysは，学習したモデルの入力レイヤー名なので，sub_2に設定されています。変換が正しく動作するためには、入力レイヤー名と出力レイヤー名（ここではResizeBilinear_2）を指定しなければなりません。また、Netronモデルビューアを使って、モデルの入力レイヤー名と出力レイヤー名を調べることもできます。(なぜ指定するのかを知りたい場合は、TensorFlowのレポジトリのIssueで詳しく説明しています)

このモデルの input_shapes は [1,257,257,3] に設定されており、1 は入力画像のバッチサイズ、257 は入力画像の形状、3 は入力画像のカラーチャンネル数を表しています。257x257の入力画像サイズを選択したのは、現在のところ、推論のための入力画像としてより大きな画像が与えられると、tfliteのandroid インタプリタがバッファオーバーフローエラーを投げるためです。

3 モデルの変換と保存

モデルが読み込まれたら、モデルを変換して deeplabv3_mnv2_custom_257.tflite として保存します。

gist.github.com

変換中はトレーニング後の最適化と量子化のためのデフォルトのパラメータを選択しました。もっと詳しく知りたい方は、以下のドキュメントを参照してください。

www.tensorflow.org

変換されたモデルを保存した後、Netronを使用して、入力と出力のレイヤー名とモデルのinput_shapesが正しいかどうかを検査します。

これで、推論のためのAndroidアプリとして展開できるようになりました。[変換済みモデルへのリンク]

次のノートブックには、モデルをテストするためのPythonの推論と可視化コードとともに、全体の変換コードが含まれています。

colab.research.google.com

4. Androidアプリのビルド

推論用のモデルを簡単に展開するために、TensorFlowチームが既に提供している既存のコードベースを使用します。

github.com

4.1 TensorFlowのサンプルソースコードをcloneする

TensorFlow の例 GitHub リポジトリをコンピュータにクローンして、デモアプリケーションを取得します。

git clone https://github.com/tensorflow/examples

4.2 Android Studioにサンプルアプリをインポート

Android StudioでTensorFlowのソースコードを開きます。そのためには、Android Studioを開き、Import Projects（Gradle、Eclipse ADTなど）を選択し、フォルダをexamples/lite/examples/image_segmentation/androidに設定します。

4.3 サンプルアプリの実行Androidデバイスをコンピュータに接続し、携帯電話に表示されるADBの許可プロンプトを承認することを確認してください。[実行] -> [アプリの実行] を選択します。アプリがインストールされるデバイスに接続されているデバイスで展開先を選択します。これでデバイスにアプリがインストールされます。アプリをテストするには、お使いのデバイスでTFL Image Segmentationというアプリを開きます。アプリを再インストールするには、以前にインストールしたものをアンインストールする必要があります。アプリに提供されているデフォルトモデルでアプリが正常に動作することを確認できるようになったので、このアプリにカスタムトレーニングされ変換されたモデルを追加していきます。

5. 変換済みの独自モデルのデプロイ

変換した deeplabv3_mnv2_custom_257.tflite モデルファイルを[app_root]/app/src/main/assets に配置してください。アプリ内のデフォルトモデルはdeeplabv3_mnv2_dm10_257.tfliteという名前になっていますが、変換されたモデルには異なるファイル名になっているため、アプリ内の名前と一致するように変更する必要があります。モデルに学習オブジェクトの異なるクラスがある場合は、クラスのラベルのテキストとクラスの数を変更します。その後、アプリケーションをビルドしてテスト実行します。