【リアルタイム物体検出】YOLORまとめ｜第2回：マスク着用をリアルタイムに判定〜カスタムデータの学習〜

このシリーズではリアルタイム物体検出のYOLORについて、環境構築から学習の方法までまとめます。

YOLORは2021年に公開された最新バージョンであり、高速度な推論が可能です。

第2回目である今回は、マスク着用判定のデータセットを用いてモデルの学習方法を紹介します。

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

今回の目標

・データセットの準備

・カスタムデータによる学習

・作成したモデルの評価

1. YOLORとは
2. YOLOR の導入
3. 「Mask Wearing Dataset」データセット
3.1. データセットのダウンロード
3.2. yamlファイルの作成
3.3. namesファイルの作成
4. 学習
5. テスト
6. まとめ

YOLORとは

YOLORは最新のコンピュータービジョンで最速のオブジェクト検出アルゴリズムの1つになります。

詳細は前回の記事をご覧ください。

YOLOR の導入

ここからはGoogle colabを使用して実装していきます。

まずはGPUを使用できるように設定をします。

「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」をGPUに変更

今回紹介するコードは以下のボタンからコピーして使用していただくことも可能です。

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

まずはPytorchをバージョン指定してインストールします。

!pip install torch==1.10.0+cu111 torchvision==0.11.0+cu111 torchaudio==0.10.0 -f https://download.pytorch.org/whl/torch_stable.html

公式よりcloneしてきます。

!git clone https://github.com/WongKinYiu/yolor
%cd yolor

必要なライブラリをインストールします。

!pip install -qr requirements.txt

今回使用するYOLOR-P6をダウンロードします。

!gdown 'https://drive.google.com/uc?id=1Tdn3yqpZ79X7R1Ql0zNlNScB1Dv9Fp76'

以上で導入は終わりです。

「Mask Wearing Dataset」データセット

データセットのダウンロード

アノテーション済みのマスク画像データを提供しているMask Wearing Datasetからデータをダウンロードします。

ダウンロードボタンを押して、以下の画像のように「YOLO v5 PyTorch」を選択します。

ダウンロードが終わったら、フォルダ名は「mask」としておきましょう。

このフォルダごとGoogleドライブの「yolor/data」にアップします。

yolor　
  ┗ data
     ┗ mask    ←ここにダウンロードデータを格納（フォルダ名をmaskに変更）

yamlファイルの作成

学習するにあたり必要な情報を記述したyamlファイルを作成します。

以下にある「data.yaml」を書き換えます。

yolor　
  ┗ data
     ┗ mask
　　　 　 ┗ data.yaml ←このファイルをダウンロードして書き換える

ダウンロードしたら、以下の通りに書き換えて「mask_yolor」として保存します。

train: ./data/mask/train/images
val: ./data/mask/valid/images

nc: 2
names: ['mask', 'no-mask']

このファイルをyolorの直下にアップします。

yolor　
  ┠ data
  ┃  ┗ mask
  ┃　　  ┗ data.yaml
  ┗ mask_yolor.yaml

namesファイルの作成

物体検出のクラス（カテゴリー）をまとめたnamesファイルを作成します。

先程と同様に、以下にある「coco.names」を書き換えます。

yolor　
  ┗ data
　　　┗ coco.names ←このファイルをダウンロードして書き換える

ダウンロードしたら、以下の通りに書き換えて「mask.names」として保存します。

mask
no-mask

このファイルを先程ダウンロードした場所にアップします。

yolor　
  ┗ data
     ┠ coco.names 
　　　┗ mask.names ←ここにファイルをアップする

以上で学習に必要なデータの準備が終わりました。

学習

早速、学習を実行しましょう。

画像の入力サイズはデータセットに合わせて416としています。

!python train.py --batch-size 8 --img 416 416 --data mask_yolor.yaml --cfg cfg/yolor_p6.cfg --weights 'yolor_p6.pt' --device 0 --name yolor_p6 --hyp 'data/hyp.scratch.1280.yaml' --epochs 100

学習結果は「yolor/runs/train/yolor_p6/」に保存されます。

テスト

学習した結果を用いてテストを行ってみましょう。

!python detect.py --source inference/images/school.mp4 --cfg cfg/yolor_p6.cfg --weights runs/train/yolor_p63/weights/best.pt --names data/mask.names --conf 0.40 --img-size 1280 --device 0