【Android】カスタムモデルで物体検出をする (モデル作成編)

はじめに

　TensorFlowを使用して独自のモデルを作成し、Androidで10円玉と100円玉を検出するアプリを作成しました。チュートリアルを見ている限りだと簡単そうに見えましたが、かなり苦戦したのでメモを残しておきます。

※ ニューラルネットワーク周りに関する知識がほぼないのでおかしなことを書いている可能性があります。自己責任でお願いします。

この記事でやること

　独自の画像データからAndroidのObjectDetectorクラスで物体検出できるtfliteモデルを作成する。

Android用のTensorFlow Liteモデルについて

　AndroidのObjectDetectorクラスでは入出力の形式がMobileNetと同じであれば、処理することができるようです。RetinaNetなども試してみましたが、出力の形式を合わせる方法がわからず断念しました。入出力の形式についてはNetronで確認できます。

TensorFlowのバージョンについて

　TensorFlowは2.10から2.11の変更で、contribというアドオンパッケージが使えなくなりました。チュートリアルの物体検出ではcontribに依存しているため2.11以上のTensorFlowだとエラーになります。
　チュートリアル「Object detection with Model Garden」だと2.15で tfliteモデルを作成できますが、AndroidのObjectDetectorクラスで使える形式のtfliteを作成することができませんでした。原因は、AndroidのObjectDetectorクラスに合う出力形式のモデルがないことです。
　モデルの出力形式を最適化できれば最新の2.15でもできそうですが、知識がないため断念し2.10で出力することにしました。

手順

1.画像データのアノテーションを行う

2.アノテーションデータをTensorFlowで扱えるデータ型（.tfrecord）に変換する

3.ラベルマップファイルを作成する

4..tfrecordを使って物体検出の学習を行う

5.モデルを.tflite形式で出力する

6.モデルにメタデータを付与する

環境

・Windows 11

・docker

・Label Studio(アノテーションツール)

・TensorFlow (ver. 2.10.0 学習用)

・TensorFlow (ver. 2.15.0 データ変換用)

・Netron(モデルのメタデータ確認)

　Label Studio, TensorFlow 2.10.0はdockerで環境を構築します。TensorFlow 2.15.0はdockerでもwslでもお好きな環境で。

画像データのアノテーションを行う

　独自の画像データにアノテーションを行っていきます。Label Imgを使っているサイトが多いですが、後継のLabel Studioが出ているのでそちらで行います。
　trainとvalの二つのtfrecordが必要になるので、「coin_train」と「coin_val」の二つのプロジェクトを作りました。tfrecordがすでにある場合は物体検出の学習まで飛ばしてください。

Label Studio Install and upgrade

使い方等については他のサイトを参照。

出力形式について

　今回はcocoファイル形式で出力します。最終的にtfrecordが作れればなんでもいいです。

tfrecordに変換する

TensorFlow 2.15で行う

　tf-models-officialを使うのにTensorFlowの2.11より上のバージョンが必要。wslでもDockerで準備してもOK。Jupyter上で行っていますがshellでもできると思います

!pip install -U -q "tf-models-official"

　アノテーションした画像のフォルダとファイルを指定して実行するとoutputに出力される。train と val でtfrecordを作成する。

TRAIN_DATA_DIR='train/images'
TRAIN_ANNOTATION_FILE_DIR='train/result.json'
OUTPUT_TFRECORD_TRAIN='output/train'

# Need to provide
  # 1. image_dir: where images are present
  # 2. object_annotations_file: where annotations are listed in json format
  # 3. output_file_prefix: where to write output convered TFRecords files
!python3 -m official.vision.data.create_coco_tf_record --logtostderr \
  --image_dir='train' \
  --object_annotations_file='train/result.json' \
  --output_file_prefix='output/train' \
  --num_shards=1

出力した train.tfrecord と val.tfrecord で学習を行う

Labelmapを作成する

　ラベルマップファイルはラベルのIdと名前をもったテキストファイル以下のようにアノテーションしたラベルの数だけ記入する。ファイル名は拡張子に .pbtxt を指定する。

item {
  id: 1
  name: 'ten'
}

item {
  id: 2
  name: 'hundred'
}

(Option) DockerでTensorFlowのGPU環境を準備する

　WindowsではTensorFlowのGPU環境は提供されなくっているようですので、Dockerで環境を作ります。物体検出に対応するためバージョンを 2.10.0 にしています。
　Jupyter Labを入れたかったのでjupyterなしのバージョンを指定しています。最初からjupyterが欲しい場合は、タグを 2.10.0-gpu-jupyter にしてください

docker pull tensorflow/tensorflow:2.10.0-gpu

　初回のコンテナがないときのみ以下のコマンドで初期化します。引数は以下です。

• --name コンテナの名前
• -p マッピングするポート
• -v E:/docker_mount:/tf/mount/ ローカルのフォルダーをマウントする

  この例だと Eドライブの docker_mount フォルダをコンテナの /tf/mount にマウントしている。

docker run --name tensorflow-gpu-2.10.0 --gpus all -it -p 8888:8888 -v E:/docker_mout:/tf/mount/ tensorflow/tensorflow:2.10.0-gpu

Jupyter Labを入れる

pip install jupyterlab
python3 -m pip install --upgrade pip
// 起動コマンド
jupyter lab --allow-root --ip 0.0.0.0

物体検出の学習を行う

ここからはTensorFlowのバージョンは2.10.0を使う

setup

必要なアプリを入れる

!apt update
!apt install -y git wget protobuf-compiler libgl1

モデルを取得する。 modelsフォルダーができる

!git clone --depth 1 https://github.com/tensorflow/models

setupファイルを作る

%%bash
cd models/research
protoc object_detection/protos/*.proto --python_out=.
cp object_detection/packages/tf2/setup.py .

pipで必要なパッケージをインストールをする

!pip install models/research
!pip install tensorflow==2.10.0

opencvのバージョンを指定する

# opencv-python-headlessとversionを合わせる必要がある？よくわからん

!pip list | grep opencv
!pip install opencv-python==4.5.2.52

modelの設定

転移学習用のモデルをダウンロードしてきます

# モデルのダウンロード
import tarfile
checkpoint = "ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.tar.gz"
url_tar = "http://download.tensorflow.org/models/object_detection/tf2/20200711/" + checkpoint
!wget {url_tar}

ファイルを解凍してtarファイルを削除します。

# tar failの解凍
tar = tarfile.open(checkpoint) 
tar.extractall()
tar.close()
!rm {checkpoint}

/tf/mount/にpipeline.configという設定ファイルをコピーします

# 設定ファイルをコピー
!cp ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8/pipeline.config pipeline.config
# 名前が長いので短縮
!mv ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8 ssd_mobilenet_v2

pipeline.configの設定

pipeline.configを編集します。設定する部分は以下のところです。

fine_tune_checkpoint は 「/tf/mount/ssd_mobilenet_v2/checkpoint/ckpt-0」のような形式で指定する

転移学習を行う

引数の設定

pipeline_path = 'pipeline_file.config'
ckpt_dir = 'ckpt'
num_steps = 500

学習する。指定したステップ数の学習を行う。

!python3 models/research/object_detection/model_main_tf2.py \
    --pipeline_config_path={pipeline_path} \
    --model_dir={ckpt_dir} \
    --alsologtostderr \
    --num_train_steps={num_steps} 

転移学習が終わるとckpt_dirに出力される。

モデルを出力する

モデルの変換

初めにチェックポイントから.pb 形式に変換する

# 出力ディレクトリの作成
!mkdir custom_model_lite

!python3 models/research/object_detection/export_tflite_graph_tf2.py \
    --trained_checkpoint_dir "/tf/mount/ckpt" \
    --output_directory "/tf/mount/cutom_model_lite" \
    --pipeline_config_path "/tf/mount/pipeline_file.config" \

次に.pb形式から.tflite形式に変換する

import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_saved_model("cutom_model_lite/saved_model/")
tflite_model = converter.convert()

.tflite形式の変数をファイルに出力する

!mkdir export
with open("export/output.tflite", "wb") as f: f.write(tflite_model)

export ディレクトリにoutput.tflite が作成されます。

モデルの確認

Netronで出力したモデルの確認ができます。以下のように入力が一つ、出力が4つあればOKです。

このoutput.tfliteをAndroidのObjectDetectorクラスに渡すとMetadataがないとエラーになります。

メタデータを付与する

ほぼ「TensorFlow Lite Metadata Writer API 」のままファイルパスのみ変更している。

Label Fileにはlabelmap.pbtxtではなくlabelmap.txtを指定する。内容は検出したいラベル名を羅列したテキストファイル。

ten
hundred

必要なライブラリのインストール

!apt install -y libusb-1.0-0-dev

ライブラリのインポート

from tflite_support.metadata_writers import writer_utils
from tflite_support.metadata_writers import object_detector

メタデータをつけて.tfliteを出力。
outputはoutput_with_metadata_v4.tfliteにしている。

ObjectDetectorWriter = object_detector.MetadataWriter
_MODEL_PATH = "/tf/mount/export/output.tflite"
# Task Library expects label files that are in the same format as the one below.
_LABEL_FILE = "/tf/mount/labelmap.txt"
_SAVE_TO_PATH = "/tf/mount/export/output_with_metadata.tflite"
# Normalization parameters is required when reprocessing the image. It is
# optional if the image pixel values are in range of [0, 255] and the input
# tensor is quantized to uint8. See the introduction for normalization and
# quantization parameters below for more details.
# https://www.tensorflow.org/lite/models/convert/metadata#normalization_and_quantization_parameters)
_INPUT_NORM_MEAN = 127.5
_INPUT_NORM_STD = 127.5

# Create the metadata writer.
writer = ObjectDetectorWriter.create_for_inference(
    writer_utils.load_file(_MODEL_PATH), [_INPUT_NORM_MEAN], [_INPUT_NORM_STD],
    [_LABEL_FILE])

# Verify the metadata generated by metadata writer.
print(writer.get_metadata_json())

# Populate the metadata into the model.
writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)

この出力したoutput_with_metadata_v4.tflite がAndroidのObjectDetectorクラスで使用できました。

まとめ

TensorFlow 2.11より上で何とか作りたかったけど、断念して2.10で作りました。何かしら方法があって知らないだけな気がします。