はじめに

任意のデータセットで物体検出のモデルを作成する手順について備忘録としてまとめました。
今回はGoogle ColabでYOLOv5を使用して物体検出モデルを作成しますので、Googleアカウントがあれば誰でも簡単に試すことができます。なお、ローカル環境でのモデル作成は以前の記事を参照してください。

（参考記事）ローカル環境でのYOLOv5モデル作成

成果物

以下のようにYolov5の物体検出により動画または画像から電車を抽出して、車両形式を判定します。
今回は一般型(E231系,E233系など)、特急(E257系,E353系など)、機関車(EH200系,EH500系など)など14種類の車両形式を対象とします。

物体検出とは

物体検出（object detection）とは、画像内の「どこに」「何が」写っているかを検出する技術のことです。
物体検出としては、SSDやYOLOといったものがよく使用されます。

YOLOv5とは

YOLOv5とは、物体検出をするアルゴリズムですがYOLOv3の後継にあたり、2020年に公開された最新のモデルです。YOLOv4という高精度化したYOLOv3の後継モデルもありますが、YOLOv5は推論処理時間がより速くなっているのが特徴です。

環境構築

まずはGoogleドライブをマウントします。

cmd

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

クローンしてきます。

%cd ./drive/MyDrive
!git clone https://github.com/ultralytics/yolov5

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

%cd yolov5
!pip install -qr requirements.txt

python

import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True)

以上で準備完了です。

準備

YOLOv5でモデル作成の際に必要となるのは、教師データ・アノテーションデータ・yamlファイルです。これらをGoogleドライブにアップします。

教師画像とアノテーションファイルを作成したうえで、それぞれyolov5内の場所に格納していきます。
①画像(.jpg)
②アノテーションファイル(.jpg)
③yamlファイル(.yaml)

yolov5
  ┠ data
  ┃  ┠ train
  ┃  ┃   ┠ images
  ┃  ┃   ┃   ┗ *.jpg
  ┃  ┃   ┗ labels
  ┃  ┃       ┗ *.txt
  ┃  ┠ valid
  ┃      ┗ images
  ┃          ┗ *.jpg
  ┗ traincar.yaml

coco.yamlをコピーしてtraincar.yamlと名前を付けて、ファイル内を以下のように書き換えます。

traincar.yaml

train: .data/train/images # 学習の画像のパス
val: .data/valid/images # 検証用画像のパス
#number of classes
nc: 14
names: [ 'E231','E233',...]

number of classes：物体検出対象のクラス数です。今回は14種類なので14とします。
class names：物体検出対象の名前です。任意の名前にすることが可能です。

学習

学習は以下のコマンドを実行すれば開始します。
初めから学習するときは以下のコマンドを入力を実行して学習を進めます。

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

上記の引数は以下の通りです。
batch:バッチサイズを指定します。
epochs:学習回数を指定します。
data.yaml:先程作成したyamlファイルを指定します。
weights:ここではyolov5x.ptを指定します。

#前回の続きから学習
YOLOv5では学習するごとに結果がGoogleドライブに自動で保存されます。
--resumeを付け加えることで、前回までの結果の途中から学習することができます。

!python train.py --resume --img 640 --batch 4 --epochs 1000 --data traincar.yaml --weights yolov5x.pt --name traincar

結果

学習が終わるごとに/runs/train/traincar/weights/にbest.ptと保存されます。

まとめ

実際にbest.ptでテストしてみます。
電車の車両形式判定モデルを作ることができました。

おわりに

最後までご覧いただきありがとうございました。
訂正要望・ご意見等がございましたら、ご連絡頂けますと幸いです。