【物体検出】YOLOXまとめ｜第2回：マスク着用判定〜カスタムデータの学習〜

このシリーズでは物体検出でお馴染みのYOLOシリーズの最新版「YOLOX」について、環境構築から学習の方法までまとめます。

YOLOXは2021年8月に公開された最新バージョンであり、速度と精度の面で限界を押し広げています。YOLOv5を超える性能で、使いやすいApache Licenseとなっています。

Google colabを使用して簡単に物体検出のモデルを実装することができますので、ぜひ最後までご覧ください。

前回の記事はこちら

今回の目標

第2回目である今回は、マスクの着用有無を判定するための「Mask Wearing Dataset」データセットによる学習の実装を紹介します。

YOLOXで必要となるデータを準備と学習の実装について学んでいきましょう。

今回の内容

・YOLOXで必要となるデータを準備

・「Mask Wearing Dataset」データセットによるYOLOX学習の実装

1. 今回の目標
2. 「Mask Wearing Dataset」データセット
3. 環境構築
4. データ準備
4.1. 「Mask Wearing Dataset」データセット
4.2. expsファイル
4.3. coco_classesファイル
5. 学習
6. 結果
7. まとめ

URL：https://arxiv.org/abs/2107.08430
出典：Zheng Ge, Songtao Liu, Feng Wang, Zeming Li, Jian Sun; arXiv:2107.08430 [cs.CV]
コード：https://github.com/Megvii-BaseDetection/YOLOX

「Mask Wearing Dataset」データセット

アノテーション済みのマスク画像データを提供しているMask Wearing Datasetからデータをダウンロードします。
ダウンロードボタンを押して、以下の画像のように「COCO」を選択します。

環境構築

前回と同様にGoogle colab環境で進めていきます。

まずは公式よりcloneしてきます。

冒頭で紹介しているYOLOXのGitHubとなります。

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

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

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

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

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

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

データ準備

ここからは必要なデータとその格納場所について紹介していきます。

「Mask Wearing Dataset」データセット

先ほどダウンロードしたデータの中身を編集します。

ダウンロード後のデータを解凍すると「変更前」のようになっています。

Mask Mask ~ .coco
  ┠ 
  ┠ train
  ┃   ┠ *.json
  ┃   ┗ *.jpg
  ┠ test
  ┃   ┠ *.json
  ┃   ┗ *.jpg 
  ┠ valid
  ┃   ┠ *.json
  ┃   ┗ *.jpg
  ┗*.txt

ここから以下の通り編集していきます。

下に示す階層と比較しながら進めていきましょう。

①フォルダ名を「face_mask」に変更

②「train」フォルダと「valid」フォルダ以外のファイルとフォルダを削除

③「train」フォルダと「valid」フォルダの中にあるjsonファイルをそれぞれ「train.json」「val.json」にファイル名を変更

④「annotations」フォルダを作成して、先ほどファイル名変更した「train.json」「val.json」を「annotations」に移動

⑤「train」フォルダと「valid」フォルダを「train2017」「val2017」にそれぞれ変更

Mask Mask ~ .coco　←①「face_mask」に変更
  ┠ 
  ┠ train　←⑤　「train2017」に変更
  ┃   ┠ *.json　←③「train.json」に変更
  ┃   ┗ *.jpg
  ┠ test　←②消去
  ┃   ┠ *.json
  ┃   ┗ *.jpg 
  ┠ valid　←⑤「val2017」に変更
  ┃   ┠ *.json　←③「val.json」に変更　
  ┃   ┗ *.jpg
  ┠*.txt　←②消去
  ┃
  ┗ annotations　←④「annotations」を作成して③の「train.json」「val.json」を移動

face_mask　　
  ┠ 
  ┠ train2017　　
  ┃   ┗ *.jpg
  ┠ val2017
  ┃   ┗ *.jpg
  ┗ 　annotations　
      ┗ train.json
      ┗ 　val.json

ここまで無事に終わると上図のようになっているはずです。

この「face_mask」フォルダをGoogleドライブから「YOLOX/datasets 」にアップロードします。

expsファイル

学習に必要な情報を記載した「expsファイル」を用意します。まずは書き換えるための元データを用意します。

「YOLOX/exps/example/custom」にある「yolox_s.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_s.py」ファイルがあった「YOLOX/exps/example/custom」にアップしましょう。

なお、ファイル名は「face_mask.py」としました。

coco_classesファイル

ラベルの一覧表である「coco_classes.py」ファイルの中身を書き変えます。

これを行わないと、テスト時に正しくラベル付けが行われないので注意しましよう。

「YOLOX/yolox/data/datasets/coco_classes.py」を開いて、中身を編集します。

今回は以下のようにしました。

COCO_CLASSES = (
    "person",
    "mask",
    "no-mask",
)

以上で準備が完了しました。

学習

学習はコード1行で始まります。

ここでは先ほどアップしたface_mask.pyを指定しています。

なお、入力画像のサイズや学習の実行回数を変更したい場合は、先ほど作成した「expファイル」から変更します。

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

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

結果

最後にテスト画像を使って推論をしてみます。

使用するモデルは先ほど学習したものを指定しています。

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