はじめに

今回はObject Detection APIを使用し、弊社社員を元データに物体検出を試しました。
以前に pytorch で SSD300のネットワークモデルを構築して、物体検出をやったことありますが、 DataLoader の作成やネットワークモデルの構築が思った以上に面倒で、 DataAugment はかなり大変でした。
その面倒で大変な作業をライブラリとして提供してくれたのが Object Detection API です。
当たり前ですが tensorflow で作られているため、簡単に tflite に変換し、スマホで動かすことが出来ます。
今回のゴールは Android で動かし、実際に遊ぶところまでとします。

開発環境

  • tensorflow-gpu
  • python3
  • GPU(Geforce 980M)
  • docker

データ

スマホで弊社社員を撮って集めました。
今回は転移学習になるのでデータ数は少なめで実験してみました。
学習データは8枚!!テストは3枚でやっています。
Object Detection API は今回初めて使うのでターゲット社員は一人でやってみました(つまり、ラベルは1つになります)。
なおプライバシー保護のため、黒線を入れていますが学習データは加工前のものを使用しています。

弊社社員の画像

アノテーション付け

アノテーション付けは labelImg を使いました。
顔だけにするか迷いましたが、思い切って全体図にしてみました。

弊社社員のアノテーション図

Object Detection API 周りの設定

使用するモデル

最終的にスマホで動かしたいため、 ssdmobilenetv1_coco を使いました。

また、今回使用する学習マシンは簡単な実験用のものでグラボの性能があまり良くないので
OOM を避けるためにこのモデルを選びました。

環境構築

Object Detection API は SSD を簡単に構築できる素晴らしいものですが、環境構築は難しい方だと思います。
様々なライブラリに依存していて、開発環境が汚れしまうため Docker で組むことをオススメします。

FROM nvidia/cuda:10.0-cudnn7-runtime

ENV DEBIAN_FRONTEND=noninteractive

RUN apt update && apt install -y python3-pip

RUN apt-get update
RUN apt-get install gcc make -y
RUN apt-get install python-lxml -y
RUN apt-get install python-dev libxml2 libxml2-dev libxslt-dev curl -y
RUN pip3 install --upgrade pip setuptools
RUN pip3 install Cython
RUN pip3 install docopt
RUN pip3 install clint
RUN pip3 install crontab
RUN pip3 install tablib
RUN pip3 install coco
RUN pip3 install matplotlib
RUN pip3 install Pillow
RUN pip3 install pycocotools
RUN pip3 install tensorflow-gpu
RUN pip3 install lxml

RUN export PYTHONPATH=$PYTHONPATH:/tensorflow/models/research:/tensorflow/

WORKDIR /tmp
RUN apt-get install -y unzip
RUN curl -OL https://github.com/google/protobuf/releases/download/v3.9.0/protoc-3.9.0-linux-x86_64.zip
RUN unzip -d protoc3 protoc-3.9.0-linux-x86_64.zip
RUN mv protoc3/bin/* /usr/local/bin/
RUN mv protoc3/include/* /usr/local/include/
RUN rm -rf protoc-3.9.0-linux-x86_64.zip protoc3

RUN apt-get install git -y
RUN git clone https://github.com/cocodataset/cocoapi.git

WORKDIR /tmp/cocoapi/PythonAPI
RUN python3 setup.py build_ext --inplace
RUN rm -rf build
RUN python3 setup.py build_ext install
RUN rm -rf build

RUN mkdir /tensorflow
# tensorflow は事前にホストの方で git clone しておきます
COPY tensorflow /tensorflow
WORKDIR /tensorflow/models/research
RUN protoc object_detection/protos/*.proto --python_out=.

ラベル map の作成

item {
  id: 1
  name: 'worker1'
}
item {
  id: 2
  name: 'worker2'
}

このように作成します。

tfRecord の作成

labelImg でアノテーションをつけたデータセットを tfrecord に変換する必要があります。
変換する際はObject Detection API(dataset_tools)を参考にしてください。

pipeline.config の修正

ラベル map の PATH や tfRecord の PATH などを編集します。
batch_size もここで修正できるため、もし VRAM が足らずに OOM が起こるならここで batch 数を調整してみてください。

学習

今回はデータ数も少なく、そこまで回す必要がないのですが loss の動きとかをみたいので10000step で回してみました。
スマホで撮った画像をそのまま使っているのでサイズは大きく、100step でも200秒とかなりかかりました。
夜中に回し、翌朝出社した時には学習は完了していました。

Android アプリで動かすための準備

tflite 用モデルの作成

tensorflow/models/research/object_detection/export_tflite_ssd_graph.py \
--pipeline_config_path=pipeline.config \
--trained_checkpoint_prefix=model.ckpt-10000 \
--output_directory=DIR \
--add_postprocessing_op=true

出来上がったものに対して、 bazel を使うと tflite 用モデルを作ることができます。

bazel run -c opt tensorflow/lite/toco:toco -- \
--input_file=tflite_graph.pb \
--output_file=detect.tflite \
--input_shapes=1,300,300,3 \
--input_arrays=normalized_input_image_tensor \
--output_arrays='TFLite_Detection_PostProcess','TFLite_Detection_PostProcess:1','TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3'  \
--inference_type=FLOAT \
--allow_custom_ops

android で動かす

TensorFlow(Object Detection)からサンプルコードを取ってきて、公式ドキュメントに従っていくつか修正すると android で簡単に動かすことができます。
今回は実験的に社員一人だけを入れて試してみましたが、一人だけだと人間としての判定をされるようで、他の社員でも高精度で判定されてしまいました。次回は二人以上で実験してみたいと思います。

実際に試した画像

おわりに

0から作るとかなり大変な SSD ですが、 ObjectDetectionAPI を使えば、機械学習初心者の方でも簡単に物体検出を楽しむことができます。また、 android や iOS などで動かすことも簡単に出来るようになっているため、興味があったら是非試してみてください。