🔰YOLOv5で実装する物体検出入門｜第2回：マスク着用判定〜カスタムデータの学習〜

このシリーズでは物体検出でお馴染みの「YOLOv5」を用いて、物体検出の実装を基礎から学ぶことができます。

環境構築から学習の方法、さらに活用方法までをまとめています。

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

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

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

YOLOv5の導入

Google colabを立ち上げて、早速実装していきましょう。

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

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

まずはGoogle ドライブのファイルにアクセスするため、マウントします。

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

今回初めてYOLOv5を使用する方はクローンしましょう。

すでに前回クローンしている方は実行不要です。

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

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

%cd yolov5/
!pip install -qr requirements.txt

以上で準備は終了です。

学習に必要なデータの準備

「Mask Wearing Dataset」データセット

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

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

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

yolov5　
  ┠ data
     ┠ images
     ┃  ┗ zidane.jpg  ←前回使用したサンプル画像
     ┠ mask    ←ここにダウンロードデータを格納（フォルダ名をmaskに変更）
     ┠ coco128.yaml

yamlファイルの作成

今回はデータセットをア画像とアノテーションファイルがどこにあるのか、クラスの種類はいくつあるのかなどの情報をひとまとめにしたファイルを作成します。

「yolov5/data」の中に「coco128.yaml」ファイルがありますので、コピーして書き換えましょう。

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

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

「mask.yaml」として、yolov5の直下に保存しましょう。

yolov5　
 ┠ data
 ┃  ┠ images
 ┃  ┃  ┗ zidane.jpg  ←前回使用したサンプル画像
 ┃  ┠ mask    ←ここにダウンロードデータを格納（フォルダ名をmaskに変更）
 ┃  ┠ coco128.yaml
 ┗ mask.yaml ←書き換えたyamlファイルを格納

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

学習

以下の通り学習を実行しましょう。

各引数は記載の通りです。

!python train.py --img 640 --batch 8 --epochs 1000 --data mask.yaml --weights yolov5x.pt --name mask

# --img 640　入力サイズ
# --batch 8 　バッチサイズ
# --epochs 1000 学習回数
# --data mask.yaml　先程作成したファイルを指定
# --weights yolov5x.pt　使用するモデル

学習が終わると「yolov5/runs/train/mask/weights」にbest.pt保存されます。

学習結果

作成したモデルを使用してテストをしてみましょう。

まずは先程作成したモデルを読み込みます。

import torch
model = torch.hub.load('', 'custom', path='runs/train/mask/weights/best.pt', source='local') 

次にテスト画像を用意しましょう。

今回はこちらから画像を拝借しました。

準備ができたら、画像を以下の場所にアップします。

yolov5　
  ┠ data
     ┠ images
        ┠ zidane.jpg  ←前回使用したサンプル画像
        ┠ test.jpg.  ←ここにテスト画像を格納

画像を指定してテストを実行します。

imgs = './data/images/test.jpg'
results = model(imgs)
results.save()

マスク着用を正しく判定するモデルを作ることができました。

まとめ

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

今回は物体検出のためのデータセット作成方法とYOLOv5学習の実装方法をまとめました。

次回の記事ではオリジナルのデータセットを作成方法について紹介します。

オリジナルデータで自分だけの学習モデルを作成してみましょう。

ぜひ、ご覧ください。

🔰YOLOv5で実装する物体検出入門｜第3回：データセット作成と学習

このシリーズでは物体検出でお馴染みの「YOLOv5」を用いて、物体検出の実装を基礎から学ぶことができます。 環境構築から学習の方法、さらに活用方法までをまとめています…

つくもちブログ　〜Python&AIまとめ〜