【物体検出】YOLOv5で電車抽出モデルを自作する

はじめに

物体検出のモデルを作成する手順について備忘録としてまとめました。
YOLOv5を使用して、物体検出モデルをローカル環境で作成します。

（参考記事）Google ColabでのYOLOv5モデル作成

成果物

以下のように動画または画像から電車を抽出するモデルを作成します。
(YOLOv5で列車の抽出は可能ですが、車両正面を正確に囲ってくれる電車に特化したモデルを作るため自作していきます)

物体検出とは

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

YOLOv5とは

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

環境構築

①まずはPyTorchを導入します。

pip3 install torch==1.8.1+cu102 torchvision==0.9.1+cu102 torchaudio===0.8.1 -f YOLOv5

②次に以下のURLからYOLOV5をダウンロードしてきます。

https://github.com/ultralytics/yolov5からクローンします。

③ダウンロードした場所に移動して必要なものをインストール

!pip install -qr requirements.txt

以上で準備完了です。

モデル作成に必要なもの

画像収集とアノテーションファイルを作成したうえで、それぞれ指定の場所に格納していきます。

①画像(.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: 1

# class names
names: [ 'traincar']

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

学習

以下のコマンドを入力

python train.py ---batch 8 --epochs 500 --data traincar.yaml --weights yolov5x.pt --name traincar

上記の引数は以下の通りです。

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

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

電車抽出モデルを作ることができました。

訂正要望がありましたら、ご連絡頂けますと幸いです。

おまけ
結果を動画にもしてみました。
よかったらご覧ください。

続き