YOLOv5 でマスク検出のトレーニングさせてみる

ラズパイでの YOLOv5 は（性能的に）結構キツイというのは分かりましたが、この YOLOv5 はモデルの学習が簡単にできるっていうおもしろそうな機能がついているので、これをちょっと掘り下げてみます

まず今回、モデルの学習（トレーニング）を回すにあたって分かったことがあります（知ってる人にとっては当たり前）
学習まわすなら Google colab

【目次】

まずは学習用のデータセットを入手

こちらのサイトで YOLOv5に対応したいろいろな画像データセットが入手できます
Object Detection Datasets

2020年8月現在、やはり人気があるのか一番上にマスクのデータセットがあります

ダウンロードしちゃいましょう
（初回はメール登録が必要です）

YOLOv5 PyTorch を選びます

ダウンロードしたら、一旦そのまま放置

YOLOv5での学習（トレーニング）の準備

Github から YOLOv5 をダウンロードすると、train.py という学習用の pythonコードがあらかじめ用意されています
自分の環境でも学習させることができますが（ラズパイ上ではムリ！）、今回は Google Colab 上でGPUを使って学習させてみます

Google Colabの練習

YOLOv5 の Github サイトにチュートリアルなどありまして
GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

Google Colab上でとりあえず実行してみたい場合は、Environments の 「Google in Colab」 をクリックして、ブラウザ内で Goolge Colab を起動させて、どんなもんかを試してみることもできます

自動的に Google Colab のサイトに接続されます

Google Colab 上に YOLOv5 をセット

では、オレ用学習環境を Google Colab に作ってみましょう。ノートブックを新規作成します

まずはGPUが搭載されているか確認します

新規作成すると、None になっていると思われるのでGPUを選択します

次に自分のGoolgeドライブをマウントして、学習データを保存できるようにします

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

Goolgeドライブへの認証を求められるので、リンクをクリックして進めます

作業用のフォルダを作ります

%cd /content/drive/My\ Drive/Colab\ Notebooks
%mkdir yolo
%cd yolo

Github から YOLOv5 をダウンロード

!git clone https://github.com/ultralytics/yolov5 

YOLOv5 に必要なモジュールをインストールし、インポートします

%cd yolov5
!pip install -qr requirements.txt  # install dependencies (ignore errors)

import torch
from IPython.display import Image, clear_output  # to display images
from utils.google_utils import gdrive_download  # to download models/datasets

clear_output()
print('Setup complete. Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))

PyTorch のバージョンや GPU が使える状態を確認できます

まずは、チュートリアルにもある、物体検出をデフォルトのまま実行してみます

!python detect.py --weights yolov5s.pt --img 416 --conf 0.4 --source inference/images/
Image(filename='inference/output/zidane.jpg', width=600)

YOLOv5 でお馴染みの画像に物体検出かましたものが、出力されました

YOLOv5の学習（トレーニング）

学習用のデータセットをGoogleドライブにアップロードし、学習用のプログラムを回していきます
学習の可視化には TensorBoard を使います

マスク画像のデータセットをアップロード

ダウンロードしておいたマスク画像データセット zipファイル、
"Mask Wearing.v1-416x416-black-padding.yolov5pytorch.zip"
を解凍し、Goolgeドライブにアップロードします
・trainフォルダ、valdファルダ
　　-> yolo/

・data.yaml
　　-> yolo/yolov5/data/

こんな感じです

TensorBoardを起動

学習といえば可視化です、可視化と言えばTensorBoard
Google Colab でもTensorBoardが簡単に使えます

学習経過のログが保存される runsフォルダを指定 TensorBoard を起動
始めはデータが無いので、なんにも表示されないかもしれません

# TensorBoard を起動する
%load_ext tensorboard
%tensorboard --logdir "/content/drive/My Drive/Colab Notebooks/yolo/yolov5/runs"

学習スタート

YOLOv5 には学習済みモデルが、サイズ順に4つ用意されています。
　・yolov5s
　・yolov5m
　・yolov5l
　・yolov5x

モデルのパラメータが多いほど、複雑な表現を解析できますが、学習に時間がかかったり、処理重たくなったりします

Google Colab上のマシンは高スペックなので、迷わず yolov5x で学習行います
とりあえず、まずテストで学習を実行してみましょう
epoch 10回で2-3分で終わります

# Train YOLOv5s on mask_dataset for 3 epochs
!python train.py --img 416 --batch 16 --epochs 10 --data data.yaml --cfg yolov5x.yaml --weights yolov5x.pt --nosave --cache --name test_yolov5x

（2022/3/7追記 エポック数がっつり回すときは nosave外しておきましょう！）

学習開始すると、TensorBoard をリロードすることで学習過程をグラフで確認できます

実際に使える学習データには、epoch数を 1000くらいは回した方がよさそうです

学習結果の確認

Epoch 1000回分学習をまわした学習済みモデルで、マスク検出してみましょう

学習結果のweightsフォルダに学習パラメータが保存されています
”runs/exp2_maskdetect_yolov5x/weights/maskdetect_yolov5x.pt"
yolo5フォルダへコピーして
weights の引数に指定します

!python detect.py --weights maskdetect_yolov5x.pt --img 416 --conf 0.4 --source inference/images/

結果、人の顔自体は検出できていて "mask" or "no-mask" どちらかを表示してくれています
マスクをしていれば、"mask"検出はできますが、
マスクをしていなければ、"no-mask" と判定してほしいですが、結構な確率で "mask" 判定がでちゃってます

今回利用したマスク画像のデータセットが、149枚だったので学習するには少なかったですが、学習自体はむちゃくちゃ簡単にできます。すげー
もうちょっと（倍くらい？）データセット増やしてやれば、精度が変わってくれそうな気がしなくもないです
逆に考えるとこの枚数くらいで、こんなにも簡単に学習できるのであれば、アノテーションツールに手を出してみるのもアリかもしれません