調和技研 技術ブログ

調和技研で取り組んでいる技術を公開していきます

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

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

medium.com

物体検出は自動運転車やロボットなど様々なアプリケーションに広く利用されるようになってきました。その中でも、物体内に埋め込まれた文字情報を抽出することが急務となっています。現実世界のオブジェクトはあらゆる形や大きさのものがあるため、画像内の古典的なオブジェクト検出のバウンディングボックスを特定するだけでは、抽出パイプラインの初期部分として使用するには十分な精度が得られません。

そこで、画像上で検出されたオブジェクトの位置を特定するために、バウンディングボックスを使用する代わりに、ポリゴンをオブジェクトマスクとして使用することが考え出されました。これらのオブジェクトマスクは、各オブジェクトの形状や大きさを正確に識別することができます。このタスクを達成するために、画像のセグメンテーションには、セマンティックセグメンテーションインスタンスセグメンテーションなどのいくつかの方法があります。セマンティックセグメンテーションは、画像の各ピクセルにオブジェクトマスクを割り当て、それらにラベルクラスを割り当てます。ここでは理論的な側面の説明は割愛します。

 

DeepLab v3モデルアーキテクチャでは、この手法を用いてピクセルごとにマスクを予測し、分類を行います。オブジェクトセグメンテーションを実行するためにこのモデルを必要とするアプリケーションは多岐にわたるため、APIゲートウェイを介してGPUサーバ上でこのモデルを実行することが可能です。しかし、このモデルをAndroid上で画像を推論することができれば、GPUサーバを利用する処理コストが削減され、アプリはインターネット接続を使わずにセグメンテーションモデルを利用することができるようになります。また、オンデバイスでのモデル推論は、入力画像と出力画像をネットワークを経由してサーバーに送る必要がないため、処理時間の短縮にもつながります。

このプロジェクトの目的は、カメラ画像のセマンティックセグメンテーション推論のためにDeepLabモデルをAndroidで実行することです。しかし、DeepLab TensorFlowモデル(saved_model.pb)は、一度学習したモデルそのままではAndroid上で推論することはできません。そのため、TfLiteインタプリタを使ってAndroid推論用のTensorFlow Lite(Tf Lite)モデルに変換する必要があります。モデルを訓練するためのドキュメントはたくさんある一方で、訓練後に凍結されたモデルをエクスポートしてTfLiteに変換してAndroidアプリに展開するためのドキュメントが不足しているようです。実際に変換を試みると、入力サイズの不一致や、TfLiteインタプリタが257x257以上の画像サイズを扱えないなどの問題に直面しました。DeepLabのAndroidアプリのサンプルコードは、入力サイズ257x257で変換されたモデルを提供しています。しかし、どのようにして訓練されたモデルグラフからモデルを変換したのかについては言及されていません。

このチュートリアルシリーズの目的は、DeepLab v3を使用してセマンティックセグメンテーションモデルを訓練し、モデルを凍結グラフとしてエクスポートし、TensorFlow liteに変換し、変換したモデルをアンドロイドアプリに展開してカメラ画像からの推論に使用することです。

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

TensorFlowチームはこのためのレポジトリを用意しており、我々はそれを使ってmobilenet v3バックボーンを使ってpascal-vocデータセットを使ってモデルを訓練すします。必要条件は以下の通りです。

python 3.6.8 環境では、pip でパッケージをインストールします。

$ pip install tensorflow-gpu==1.15.3
$ pip install tf_slim==1.0.0
$ pip install numpy
$ pip install Pillow==2.2.1
$ pip install matplotlib
$ pip install PrettyTable

2. DeepLabのレポジトリのclone

$ cd ~
$ mkdir deeplab;cd deeplab
$ git clone --depth 1 https://github.com/tensorflow/models.git

もしくは以下のリリースからコードとモデルを入手します。

github.com

(任意) - もし独自のデータセットを使いたい場合は、この記事のステップ4,5,6,7に従ってください。以下のステップ3と4は無視してください。

3. Pascal-VOC datasetのダウンロード

$ cd models/research/deeplab/datasets
$ mkdir pascal_voc_seg
$
curl -sc /tmp/cookie \
 "https://drive.google.com/uc?export=download&id=1rATNHizJdVHnaJtt-hW9MOgjxoaajzdh" > /dev/null
$ CODE="$(awk '/_warning_/ {print $NF}' /tmp/cookie)"
$ curl -Lb /tmp/cookie \
  "https://drive.google.com/uc?export=download&confirm=${CODE}&id=1rATNHizJdVHnaJtt-hW9MOgjxoaajzdh" \
  -o pascal_voc_seg/VOCtrainval_11-May-2012.tar

4. 学習用の画像を取り出しTFRecordへ変換

以下のスクリプトは、データセットを抽出し、モデル学習の入力形状で必要とされる形状に変換し、TfRecordとして保存します。 

$ sed -i -e "s/python .\/remove_gt_colormap.py/python3 .\/remove_gt_colormap.py/g" \
      -i -e "s/python .\/build_voc2012_data.py/python3 .\/build_voc2012_data.py/g" \
      download_and_convert_voc2012.sh
$
sh download_and_convert_voc2012.sh

5. ログ用のフォルダの作成

$ cd ../..
$ mkdir -p deeplab/datasets/pascal_voc_seg/exp/train_on_train_set/train
$ mkdir -p deeplab/datasets/pascal_voc_seg/exp/train_on_train_set/eval
$ mkdir -p deeplab/datasets/pascal_voc_seg/exp/train_on_train_set/vis 

以下のフォルダを環境変数に設定します。

$ export PATH_TO_TRAIN_DIR=${HOME}/deeplab/models/research/deeplab/datasets/pascal_voc_seg/exp/train_on_train_set/train
$ export PATH_TO_DATASET=${HOME}/deeplab/models/research/deeplab/datasets/pascal_voc_seg/tfrecord
$ export PYTHONPATH=${HOME}/deeplab/models/research:${HOME}/deeplab/models/research/deeplab:${HOME}/deeplab/models/research/slim:${PYTHONPATH}

6. 学習を開始

train.pyスクリプトはモデルを訓練するために使われます。このスクリプトには、モデルの学習パラメータを変更するためのいくつかのコマンドラインフラグがあります。

model_variantフラグはどのオブジェクト検出バックボーンを使用するかを選択するために使用されます。選択肢の一覧はこちらにあります。我々はより高速な学習のために mobilenet_v3_small_seg を使用します。

training_number_of_steps フラグは、トレーニングのステップを変更するために使用します。gpuの計算能力が限られている場合は、少ないステップ数から始めてください。モデルはsave_interval_secsの値ごとにチェックポイントファイルに保存されます。そのため,学習を再開するときは最後のチェックポイントから開始されます.

$ python3 deeplab/train.py \
    --logtostderr \
    --training_number_of_steps=500000 \
    --train_split="train" \
    --model_variant="mobilenet_v3_small_seg" \
    --decoder_output_stride=16 \
    --train_crop_size="513,513" \
    --train_batch_size=8 \
    --dataset="pascal_voc_seg" \
    --save_interval_secs=300 \
    --save_summaries_secs=300 \
    --save_summaries_images=True \
    --log_steps=100 \
    --train_logdir=${PATH_TO_TRAIN_DIR} \
    --dataset_dir=${PATH_TO_DATASET}

7. 学習済みモデルを凍結グラフとしてエクスポートする

export_model.pyスクリプトは、保存された学習チェックポイントをfrozen_inference_graph.pbモデルファイルにエクスポートするために使用されます。これはPart-2でさらにmodel.tflite形式に変換するために使用されます。このスクリプトにはmodel_variantフラグがあり、これは以前のtrain.pyスクリプトと同じでなければなりません。

crop_sizeフラグは推論ステップ(Androidアプリ、Part-3)で前処理する画像のサイズを指定します。推論の入力に大きな画像を与えると、tfliteのandroidインタプリタバッファオーバーフローエラーを出すので、crop_sizeを257x257にしてエクスポートしています。

$ python3 deeplab/export_model.py \
--model_variant="mobilenet_v3_small_seg" \
--crop_size=257 \
--crop_size=257 \
--checkpoint_path=${PATH_TO_TRAIN_DIR}/model.ckpt-30000 \
--export_path=${HOME}/Downloads/deeplabv3_mnv2_pascal_trainval_2018_01_29/deeplabv3_mnv2_pascal_trainval/frozen_inference_graph.pb

以下のcolab notebookには、参考のために、すべてのトレーニングとエクスポート用のスクリプトが記載されています。

colab.research.google.com

Part2に続きます