AIで室内の人間をカウントする研究記録⑥

概要

今回はYOLOv3で独自データを用いて学習を行う方法をご紹介します。
YOLOv3とは何か、YOLOv3の導入方法についてはAIで店舗内の人間をカウントする研究記録を参考にしてみてください。
※学習用の画像データは用意してあるていで進めていきます。

アノテーション

まずはアノテーション作業を行います。
機械学習でのアノテーションとは、画像中の物体が何なのか、それはどこにあるのかという情報ファイルを作成することを言います。
アノテーションソフトはいろいろありますが、YOLO形式のアノテーションファイルをエクスポートすることができるVoTTというものをおすすめします。（ソフトの使い方はこちらのサイトに詳しくまとめられています）
VoTTのメニュー上部にあるObject Detection -> Export Tags でYOLOを選択するとYOLO形式のアノテーションファイルをエクスポートすることができます。

アノテーションファイルはそのままkeras版YOLOv3には使えない

VoTTでエクスポートしたYOLO用アノテーションファイルはkeras-YOLOv3の形式と少し異なり、そのまま使うことが出来ません。
なので、エクスポートしたファイルからkeras-YOLOv3用のアノテーションファイルを作成する必要があります。

フォーマット

keras-YOLOv3のアノテーションファイルのフォーマットは以下のような形式です。

画像のパス x1,y1,x2,y2,クラスの番号

１行が1つの画像を表しており、行頭にファイルパス、スペース区切りでバウンディングボックスの情報を記載します(これはいくつあってもok)。
クラスの番号はクラス情報を記載したクラスファイルを作り、それを参照して記載します。番号は0スタートです。

クラスファイル

今回は検出対象が人だけなので、クラスファイルは以下のような形式になります。

アノテーションファイルを生成する処理

import os
from glob import glob
from PIL import Image

def make_annotation_file(dir_path, images_dir):
    annotation_files_dir = glob(os.path.join(dir_path, '*'))
    new_annotation = ""
    for annotation in annotation_files_dir:
        root, ext = os.path.splitext(annotation)
        filename = os.path.basename(annotation)
        if ext=='.txt' and filename!='classes.txt':
            new_annotation += './train_images'+'/'+filename.rstrip('text')+'jpg'

            try:
                image = Image.open(images_dir+'\\'+filename.rstrip('text')+'jpg')
            except FileNotFoundError:
                continue

            with open(annotation, 'r') as annotation_data:
                line = annotation_data.readline()
                while line:
                    new_annotation += " "
                    list = line.strip().split()
                    list.pop(0)

                    x_center = float(list[0]) * image.width
                    y_center = float(list[1])  * image.height
                    width = float(list[2]) * image.width
                    height = float(list[3]) * image.height

                    x_min = x_center - width / 2
                    y_min = y_center - height / 2
                    x_max = x_min + width
                    y_max = y_min + height
                    val = [x_min, y_min, x_max, y_max]
                    for i in range(4):
                        new_annotation += str(round(val[i]))+','
                    new_annotation += '0'  #クラスはpersonだけなのですべて0
                    line = annotation_data.readline()

            new_annotation += '\n'
    new_annotation = new_annotation.rstrip()
    f = open('./train.txt', 'w')
    f.write(new_annotation)
    f.close()

if __name__ == '__main__':
    make_annotation_file("YOLOアノテーションファイルのディレクトリ", '画像ファイルのディレクトリ')

これを実行すると、"train.txt"という名前のアノテーションフィルが生成されます。

学習の実行

ではいよいよ学習を行います。

まずは先程作成した"train.txt"と画像ファイルのディレクトリを同じ階層のディレクトリに配置します。
次にtrain.pyの"annotation_path"、"classes_path"、"input_shape"をそれぞれ書き換えます。
"annotation_path"と"classes_path"にはそれぞれ自分で作成した"train.txt"のパスとクラスファイルのパスを指定、"input_shape"には入力サイズ（必ず32の倍数）を指定します。
ちなみに、"input_shape"は必ずしも入力する画像のサイズと同じサイズでなければならない、ということはありません。
入力画像のサイズよりも小さいサイズを指定することも可能です。

def _main():
    annotation_path = 'train.txt'
    log_dir = 'logs/000/'
    classes_path = 'model_data/voc_classes.txt'
    anchors_path = 'model_data/yolo_anchors.txt'
    class_names = get_classes(classes_path)
    num_classes = len(class_names)
    anchors = get_anchors(anchors_path)

    input_shape = (416,416) # multiple of 32, hw

ここまできたら、"train.py"を実行します。
学習が終わると、"trained_weights_finel.h5"という名前でモデルが保存されます。
このモデルは"log_dir"に指定された場所に保存されるので、今回の場合logs/000/に保存されます。

ResourceExhaustedErrorが出力されてしまう時

このエラーはマシンのメモリ不足によって引き起こされるエラーなので、GPUのメモリサイズを大きくすれば解決します。
また、多くの場合入力画像またはバッチサイズが大きすぎることが原因です。
なので、このエラーが出力される際は入力サイズかバッチサイズ（または両方）を下げて実行してみてください。
YOLOv3は学習が2つのステージに分かれているため、"train.py"の57行目と76行目にある"batch_size"をそれぞれ変えます。

ログファイルを使用して続きから学習を行う方法

"train.py"を実行すると、"trained_weights_stage_1.h5"と"trained_weights_finel.h5"の2つのモデルが生成されます。
前述した通り、YOLOv3の学習は２つのステージにわかれており、ステージ1が終了した時点でのモデルが"trained_weights_stage_1.h5"、ステージ２まで
終了したモデルが"trained_weights_finel.h5"として保存されます。

"train.py"ではそれぞれのステージをif文のTrue/False指定で制御することが出来ます。
このif文は"train.py"の52行目と70行目にあります。
ステージ1終了時点でのモデルを使用してステージ２から学習をし直す場合は52行目のif文をFlaseにし、33行目の"weights_path"に"trained_weights_stage_1.h5"のパスを指定することができます。

    # Train with frozen layers first, to get a stable loss.
    # Adjust num epochs to your dataset. This step is enough to obtain a not bad model.
    if Flase:  #TrueからFalseに変更
        model.compile(optimizer=Adam(lr=1e-3), loss={
            # use custom yolo_loss Lambda layer.
            'yolo_loss': lambda y_true, y_pred: y_pred})

        batch_size = 32

   else:
        model = create_model(input_shape, anchors, num_classes,
            freeze_body=2, weights_path='logs/000/trained_weights_stage_1.h5') # weights_pathを変更

次回

学習済みモデルのファインチューニングの方法についてご紹介します。

---

次の記事へ

前の記事へ 目次へ戻る