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


#はじめに

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

![mask.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1559455/75068066-f8da-a2ed-b4ab-a9fc2fbf556a.jpeg)

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

https://tt-tsukumochi.com/archives/tag/yolox

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

https://qiita.com/hkwsdgea_ttt2/items/970f34e1aa59059f7c69

#環境構築

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

```ruby:python
from google.colab import drive
drive.mount('/content/drive')

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

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

```ruby:python
%cd YOLOX
!pip install -U pip && pip install -r requirements.txt
!pip install -v -e . 
```

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

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

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

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

https://public.roboflow.com/object-detection/mask-wearing

リンク先から「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ファイル」を作成します。
フォルダの格納場所や学習回数などの条件はここで指定します。

```ruby: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ファイル」から変更します。


```ruby: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」として保存されます。

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

```ruby: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.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1559455/75068066-f8da-a2ed-b4ab-a9fc2fbf556a.jpeg)


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

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

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

https://tt-tsukumochi.com/archives/tag/yolox