【物体検出】YOLOシリーズ最新版のYOLOXを試してみる② マスク着用判定〜カスタムデータセットの学習〜

#はじめに

前回に引き続き、物体検出でお馴染みのYOLOシリーズの最新版「YOLOX」を試していきます。今回はマスク着用検出モデル作成を通じて、YOLOXによる学習の実装を紹介します。Google colab環境を使用しますので、誰でも簡単に試すことが可能です。

（こちらではより詳細に紹介しております。よかったらご覧ください。）

#YOLOXとは
YOLOXは2021年8月に公開された最新バージョンであり、速度と精度の面で限界を押し広げています。YOLOv5を超える性能で、使いやすいApache Licenseとなっています。YOLOXのベンチマーク結果は前回の記事をご参照ください。

#環境構築

早速YOLOXを使用していきましょう。以下、Google colab環境で進めていきます。
まずは本家よりcloneしてきます。

python

from google.colab import drive
drive.mount('/content/drive')

%cd ./drive/MyDrive
!git clone https://github.com/Megvii-BaseDetection/YOLOX

次に必要なライブラリをインポートします。

python

%cd YOLOX
!pip install -U pip && pip install -r requirements.txt
!pip install -v -e . 

こちらも合わせて準備します。

python

!pip install cython
!pip install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

以上で準備完了です。
YOLOXを使えるようになりました。

#データセットの準備
今回は、アノテーション済みのマスク画像データを提供している「Mask Wearing Dataset」からデータをダウンロードします。

リンク先から「COCO」を選択してダウンロードしてきます。

#データセットの格納
次に、先ほどダウンロードしてきたデータセットを以下のように変更していきます。
順番は前後しても問題ありません。

①先ほどダウンロードしたデータのフォルダ名を「face_mask」に変更
②「train」フォルダと「valid」フォルダ以外のファイルとフォルダを消去
③「train」フォルダと「valid」フォルダの中にあるjsonファイルをそれぞれ「train.json」「val.json」にファイル名を変更
④「annotations」フォルダを作成して、先ほどファイル名変更した「train.json」「val.json」を「annotations」に移動
⑤「train」フォルダと「valid」フォルダを「train2017」「val2017」にそれぞれ変更

以上のように変更します。
変更が終わったら、YOLOXの中の「datasets」に「face_mask」のフォルダごとアップします。

#expファイルの準備

学習に必要な情報をまとめた「expファイル」を作成します。
フォルダの格納場所や学習回数などの条件はここで指定します。

face_mask.py

import os
from yolox.exp import Exp as MyExp

class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.depth = 0.33
        self.width = 0.50
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

        # Define yourself dataset path
        self.data_dir = "datasets/face_mask"
        self.train_ann = "train.json"
        self.val_ann = "val.json"

        self.num_classes = 3
        self.max_epoch = 300
        self.data_num_workers = 4
        self.eval_interval = 1

作成が終わったら、YOLOXの中の「exps/example/custom」にアップします。
以上で準備が完了しました。

#いざ学習
学習はコード1行で始まります。ここでは先ほどアップしたface_mask.pyを指定しています。
なお、入力画像のサイズや学習の実行回数を変更したい場合は、先ほど作成した「expファイル」から変更します。

python

!python tools/train.py -f exps/example/custom/face_mask.py -d 1 -b 8 --fp16 -o -c yolox_s.pth

学習が終わると、結果がYOLOXの中の「outputs/face_mask/」に「best_ckpt.pth」として保存されます。

#学習結果
実際にテスト画像で、先ほどのモデルのテスト結果を表示させてみます。

python

TEST_IMAGE_PATH = "test.jpg"
MODEL_PATH = "YOLOX_outputs/face_mask/best_ckpt.pth" 
!python tools/demo.py image -f exps/example/custom/face_mask.py -c {MODEL_PATH} --path {TEST_IMAGE_PATH} --conf 0.3 --nms 0.45 --tsize 640 --save_result --device gpu

※classesファイルを変更しないと正しくラベル付けができないので注意

マスクを着けている人の検出ができています。
しっかりとしたデータセットを用意すれば、簡単に高精度な検出器ができますね。

#まとめ
最後までご覧いただきありがとうございました。
今回は「Mask Wearing Dataset」データセットを用いたYOLOXによる学習の実装を紹介しました。
データセットが用意できれば、簡単に高性能なオリジナルモデルが作れそうですね。

こちらではオリジナルのデータセットを用いたYOLOXの学習方法を紹介します。
よかったらご覧ください。