【機械学習】YOLOv5を使って独自データから物体検出モデルを作ってみよう

今回はYOLOv5を使って機械学習をし独自データから物体検出モデルを作り、PyTorch上で使用する方法を解説します。
開発環境はGoogle Colab上で開発をし、独自データセットは準備できてるものとして進めていきます。

前提

今回物体検出モデルを作るにあたって独自データセット(画像データ・ラベルデータ)が必要です。
images(画像データ),labels(ラベルデータ)フォルダを用意し、Google Drive上にアップロードできているとして進めます。

モデルを作成する

では実際にGoogle Colabを利用して独自物体検出モデルを作成していきます。
Google Colabでは無料でGPUを使うことができ、モデルを作成する機械学習を行うにあたってCPUよりも早く処理ができるため、ランタイム設定からランタイムのタイプを変更し、ハードウェアアクセラレーターをGPUに設定します。

YOLOv5をインストールする

今回使うYOLOv5をインストールします。
コードセルを作成し、以下コードを入力・実行してください。

!git clone https://github.com/ultralytics/yolov5
!pip install -r yolov5/requirements.txt

上記コードを実行すると以下の様なログが出力され、インストールが完了します。

Cloning into 'yolov5'...
remote: Enumerating objects: 10488, done.
~~~省略~~~
Installing collected packages: thop, PyYAML
  Attempting uninstall: PyYAML
    Found existing installation: PyYAML 3.13
    Uninstalling PyYAML-3.13:
      Successfully uninstalled PyYAML-3.13
Successfully installed PyYAML-6.0 thop-0.0.31.post2005241907

Google Driveを使えるようにする

実際に自作データセットをGoogle Driveから読み込むためにColab上にGoogle Driveをマウントします。
コードセルを作成し以下コードを実行します。

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

上記コードを実行すると/content/driveディレクトリとしてGoogle Driveを扱うことができるようになります。

データセット設定ファイルを作成する

YOLOv5のデータセットの設定ファイルを作成します。
データセットの設定ファイルはdataset.yamlとし、以下の様なディレクトリ構造に配置します。
また、ここではクラス数は10個とし全て連番としてカウントします。

content/drive/MyDrive/
 └ data/
     ├ dataset.yaml # 今回作成するファイル
     ├ val/ # 確認用データ
     │   ├ images/
     │   │   ├ 0.png
     │   │   ├ 1.png
     │   │   └ ...
     │   └ labels/
     │       ├ 0.txt
     │       ├ 1.txt
     │       └ ...
     └ train/ # 学習用データ
         ├ images/
         │   ├ 0.png
         │   ├ 1.png
         │   └ ...
         └ labels/
             ├ 0.txt
             ├ 1.txt
             └ ...

実際に書き込むにdataset.yamlなります。
以下内容を書き込むことでデータセットを設定できます。

# train val data
path: /content/drive/MyDrive/data/
train: /content/drive/MyDrive/data/train/
val: /content/drive/MyDrive/data/val/

# クラス数
nc: 10

# クラス名
names: ['0','1','2','3','4','5','6','7','8','9']

機械学習をする

では実際に機械学習をし、モデルデータを作成します。
機械学習をするにはYOLOv5に同梱されているtrain.pyファイルを実行します。
実際に以下のコードで機械学習を開始します。

!python train.py --batch 16 --epochs 35 --data /content/data/dataset.yaml --weights yolov5s.pt

それではここで実際に使用している引数を説明します。
--batchオプションはバッチサイズを指定することができます。
--epochsオプションはエポックサイズ(学習回数)を指定することができます。
--dataオプションは独自データの設定ファイルまでのパスを指定します。
--weightsオプションは学習済みの重みを指定でき、変更することでより高品質なモデルを作成することができます。
実際に実行すると機械学習が開始され、以下の様なログが出力されます。

WARNING: --img-size 1080 must be multiple of max stride 32, updating to 1088
Overriding model.yaml nc=80 with nc=17
                 from n params  module             arguments
0                -1   1 3520    models.common.Conv [3, 32, 6, 2, 2]
1                -1   1 18560   models.common.Conv [32, 64, 3, 2]
2                -1   1 18816   models.common.C3   [64, 64, 1]
3                -1   1 73984   models.common.Conv [64, 128, 3, 2]
4                -1   2 115712  models.common.C3   [128, 128, 2]
5                -1   1 295424  models.common.Conv [128, 256, 3, 2]
6                -1   3 625152  models.common.C3   [256, 256, 3]
7                -1   1 1180672 models.common.Conv [256, 512, 3, 2]
~~~省略~~~

実際に機械学習が完了すると/runs/train/exp/weightsディレクトリに学習済みモデルbest.ptが生成されています。
次はこれを使用して画像から物体認識をしていきたいと思います。

モデルをPyTorchで使用する

では実際に機械学習したモデルをPyTorchで使用して物体認識をしてみましょう。
YOLOv5で作成したモデルをPyTorchで読み込むためにはtorch.hub.loadメソッドを使います。
実際に以下コードで画像認識をしてみましょう。

import torch

# モデルを読み込む
# best.ptまでのパスは自身の物に変更してください
model = torch.hub.load('ultralytics/yolov5', '/content/yolov5/runs/train/exp/weights/best.pt')

# モデルで画像を処理する
# 画像へのパスは自身の物に変更してください
# ここではダミーを選択しています。
result = model('/content/yolov5/data/images/zidane.jpg')

# モデルで画像を処理した後に保存する。
# 引数無しの場合 runs/detect/exp に保存される
model.save()

実際に実行すると以下の様な画像から物体認識を行いハイライトされた画像が出力されます。

ここではダミーとしていますが実際には学習した物体が認識できているかと思います。

まとめ

今回は自作データを元に機械学習をし、モデルを作成しました。
YOLOv5を用いることで誰でも簡単に学習をしモデルを作成することができます。
YOLOv5自体、比較的新しい技術なのでこれから広がって行くと思います。
Colab上では無料で誰でもGPUを使うことができるのでどんどん機械学習をしモデルを作成してみましょう。