CCCMKホールディングス TECH Labの Tech Blog

TECH Labスタッフによる格闘記録やマーケティング界隈についての記事など

ブログタイトル

一般物体検知アルゴリズム Single Shot MultiBox Detector(SSD)を動かす

こんにちは、技術開発の三浦です。前回一般物体検知アルゴリズム Single Shot MultiBox Detector(SSD)の仕組みと実装について紹介しました。今回は実装したプログラムを使って自分だけのオリジナル物体検知器を作る手順について紹介します。

画像を集める

検出したい物体が写っている画像を用意します。今回は100枚くらい用意しました。

f:id:miu4930:20200510101715p:plain

撮影する角度を変えたり撮影場所を変えたりしました。

f:id:miu4930:20200510101739p:plain

学習データを作る(アノテーション)

次に収集した画像のどこに何が写っているのかというデータを作成します。PASCAL VOCのデータフォーマットで作ります。以下のようなxml形式です。

<annotation>
    <folder>img</folder>
    <filename>IMG20200505124111.jpg</filename>
    <path>/home/xxx/img/IMG20200505124111.jpg</path>
    <source>
        <database>Unknown</database>
    </source>
    <size>
        <width>1280</width>
        <height>960</height>
        <depth>3</depth>
    </size>
    <segmented>0</segmented>
    <object>
        <name>Drakee</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>3</xmin>
            <ymin>253</ymin>
            <xmax>274</xmax>
            <ymax>428</ymax>
        </bndbox>
    </object>
    <object>
        <name>Golem</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>189</xmin>
            <ymin>37</ymin>
            <xmax>404</xmax>
            <ymax>246</ymax>
        </bndbox>
    </object>
    <object>
        <name>Slime</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>221</xmin>
            <ymin>459</ymin>
            <xmax>378</xmax>
            <ymax>608</ymax>
        </bndbox>
    </object>
    <object>
        <name>Metal Babble</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>477</xmin>
            <ymin>233</ymin>
            <xmax>646</xmax>
            <ymax>331</ymax>
        </bndbox>
    </object>
</annotation>

ではどうやってこのデータを作るのかというと、アノテーション用のツールは色々と公開されているのでそれを利用します。LabelImgというツールがDockerで簡単に導入できそうだったのでこちらを利用しました。

github.com

以下、コマンドを記載します。

git clone https://github.com/tzutalin/labelImg.git
cd labelImg

docker run -it \
--user $(id -u) \
-e DISPLAY=unix$DISPLAY \
--workdir=$(pwd) \
--volume="/home/$USER:/home/$USER" \
--volume="/etc/group:/etc/group:ro" \
--volume="/etc/passwd:/etc/passwd:ro" \
--volume="/etc/shadow:/etc/shadow:ro" \
--volume="/etc/sudoers.d:/etc/sudoers.d:ro" \
-v /tmp/.X11-unix:/tmp/.X11-unix \
tzutalin/py2qt4

make qt4py2;./labelImg.py

ツールが立ち上がります。

f:id:miu4930:20200510101922p:plain

あとはアノテーション作業をするだけです。かなり疲れる作業ですが、根性で進めます。

f:id:miu4930:20200510102936p:plain

こういう作業をしていると、何故か昔のことを思い出します。中学校の頃のことを思い出しました

f:id:miu4930:20200510103012p:plain

1つの画像ファイルに対し、1つのxmlファルが生成されます。次にこのxmlファイルを加工してPythonのプログラムで読み込むpklファイルに変換します。次のようなプログラムになります。

import numpy as np
import os
from xml.etree import ElementTree

class XML_preprocessor(object):

    def __init__(self, data_path, classes):
        self.path_prefix= data_path
        self.num_classes = len(classes)
        self.classes = classes
        self.data = dict()
        self._preprocess_XML()

    def _preprocess_XML(self):
        filenames = os.listdir(self.path_prefix)
        for filename in filenames:
            tree = ElementTree.parse(self.path_prefix + filename)
            root = tree.getroot()
            bounding_boxes = []
            one_hot_classes = []
            size_tree = root.find('size')
            width = float(size_tree.find('width').text)
            height = float(size_tree.find('height').text)
            for object_tree in root.findall('object'):
                for bounding_box in object_tree.iter('bndbox'):
                    xmin = float(bounding_box.find('xmin').text)/width
                    ymin = float(bounding_box.find('ymin').text)/height
                    xmax = float(bounding_box.find('xmax').text)/width
                    ymax = float(bounding_box.find('ymax').text)/height
                bounding_box = [xmin,ymin,xmax,ymax]
                bounding_boxes.append(bounding_box)
                class_name = object_tree.find('name').text
                one_hot_class = self._to_one_hot(class_name)
                one_hot_classes.append(one_hot_class)
            image_name = object_tree.find('name').text.split('/')[-1]
            bounding_box = np.asarray(bounding_box)
            one_hot_classes = np.asarray(one_hot_classes)
            image_data = np.hstack((bounding_boxes, one_hot_classes))
            self.data[image_name] = image_data

    def _to_one_hot(self, name):
        one_hot_vector = [0] * self.num_classes
        for i, target in enumerate(self.classes):
            if target == name:
                one_hot_vector[i] = 1
        return one_hot_vector

if __name__ == '__main__':
    import pickle
    classes = ['object0',
               'object1',
               'object2',
               'object3',
               'object4',
               'object5',
               'object6',
               'object7',
               'object8']
    data = XML_preprocessor('xmls/',classes).data
    pickle.dump(data, open('Object.pkl','wb'))

prior boxと学習済みの重みファイルを取得する

前回作成したプログラムの中のBBoxUtilityにセットするprior boxとVGG16の学習済みの重みファイルを用意します。これらはKerasの実装ページからダウンロードすることができます。

github.com

必要なファイルを配置する

次のようにファイルを配置します。

.
├── SSD_training.ipynb
├── ssd.py
├── ssd_layers.py
├── ssd_training.py
├── ssd_utils.py
├── Object.pkl
├── prior_boxes_ssd300.pkl
├── weights_SSD300.hdf5
├── traindata
│   ├── images
│       ├── xxx.jpg
        ├── xxx.jpg
        ...

prior_boxes_ssd300.pklとweights_SSD300.pklがダインロードしてきたprior box及びVGG16の学習済みの重み、Object.pklがアノテーションツールで生成したxmlファイルを変換したファイル、traindata/imagesに元の画像を入れます。その他は前回の記事で作成したKerasで作られたプログラムをTensorflow v2で動くよう編集したファイルになります。

Notebookの変更箇所

Notebookの変更箇所をまとめます。Tensorflow v2対応とScipyの対応については前回の記事を参考にしてください。ここではそれ以外の変更箇所を記載します。

まず検出したい物体の種類を変更します。NUM_CLASSESで指定しています。

# some constants
NUM_CLASSES = 4
input_shape = (300, 300, 3)

次にアノテーションデータを読み込むところです。作成したpklファイル名に変更します。

gt = pickle.load(open('gt_pascal.pkl', 'rb'))
keys = sorted(gt.keys())
num_train = int(round(0.8 * len(keys)))
train_keys = keys[:num_train]
val_keys = keys[num_train:]
num_val = len(val_keys)

画像ファイルの格納パスも変更します。path_prefixです。

path_prefix = '../../frames/'
gen = Generator(gt, bbox_util, 16, '../../frames/',
                train_keys, val_keys,
                (input_shape[0], input_shape[1]), do_crop=False)

最後に推計した結果を画像に表示する部分です。枠のカラーを生成するところが元のnotebookだと4クラスで固定されています。これだと4を超えるクラスの検出をしようとするとエラーになります。

colors = plt.cm.hsv(np.linspace(0, 1, 4)).tolist()

colors = plt.cm.hsv(np.linspace(0, 1, NUM_CLASSES)).tolist()

に変更します。

やってみる

やってみました。ちなみに検出する物体は大好きなドラクエのモンスターたちです。昔飲み物のオマケで集めました。

学習開始してちょっとしてテストした結果。

f:id:miu4930:20200510103434p:plain

もう少し学習させてテスト。シルバーデビルを認識しはじめた!

f:id:miu4930:20200510103450p:plain

ちなみに検証データでなく学習データで試すとこのような感じ。学習データにはかなりフィットしているようです。

f:id:miu4930:20200510103514p:plain
シルバーデビルとモーモン、たしかに似てるかも

せっかくドラクエのモンスターなので検出結果もそれっぽくしてみました。本当は画像上に表示したいのですが、日本語表示で文字化けしてしまい、まだ出来ていません。

f:id:miu4930:20200510085104p:plain

最後に

SSDについて前回で実装、今回で実際に動かすところまで紹介しました!実際に動かすと面白いですね。今度はリアルタイムで物体検出したり、学習させたモデルをJetson Nanoで動かしてみたりしようと思います!