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

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

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

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

今回の目標

・データセットの準備

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

・作成したモデルの評価

YOLORとは

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

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

【リアルタイム物体検出】YOLORまとめ|第1回:導入からチュートリアルまで

このシリーズではリアルタイム物体検出のYOLORについて、環境構築から学習の方法までまとめます。 YOLORは2021年に公開された最新バージョンであり、高速度な推論が可能で…

YOLOR の導入

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

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

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

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

Open In Colab

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

まとめ

最後までご覧いただきありがとうございました。

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

なお、YOLORは推論速度重視のため、精度についてはやや物足りなさがあります。精度向上の手法については別な回でまとめますので、ぜひ合わせて御覧ください。