YOLOv5を使って物体検出する

概要

YOLOv5を使って物体検出にチャレンジした記録のまとめです。基本的な流れをはじめに書きますが、ちょっとずつ書き足す予定です。

YOLOv5を使ってサンプルデータではなく自前のデータに対して学習を行う場合、最低限行わなければならないのは以下の3点です。

1. 学習のためのフォルダ構成を作る
2. アノテーションの情報が入ったテキストファイルを作成する
3. フォルダ構成やラベル設定を書いたyamlファイルを作成する

というわけで上から順番にやっていきます。

学習のためのフォルダ構成を作る

まずはフォルダ構成です。今回は以下のような構成を作成しました。

file/
├─ yolov5/  <-- githubからcloneしてくる。cloneした状態から特にいじってないです。
│    ├─ utils/
│    ├─ runs/ <-- この配下に結果が入る
│    ├─ data/ <-- あとで3番でやるyamlファイルを入れる
│    ├─ models/ <--この配下にあるyamlファイルからモデルを選ぶ
│    └─ train.py, val.py, LICENSEなどなど
├─ train/  <-- 訓練データを入れる
│    ├─ images/ 
│    └─ labels/
├─ val/  <-- バリデーションデータを入れる
│    ├─ images/
│    └─ labels/
└─ test/ <-- 検出をしたいデータを入れる

yolov5のフォルダは以下のファイルからcloneしてきます。

チュートリアルを見ながらターミナル上で以下を実行しました。

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

フォルダ構成ができたので画像をそれぞれtrain/images/, vel/images/, test/フォルダに入れていきます。

アノテーションの情報が入ったテキストファイルを作成する

次にアノテーションとラベルの情報を用意していきます。YOLOv5の学習のためには以下のようなテキスト形式のアノテーションデータが各画像に対して１つ必要になります。

ラベル	x	y	w	h
画像のラベル番号	領域の中心x軸	領域の中心y軸	領域の幅	領域の高さ

なお、x, y, w, hはすべて画像の幅と高さで正規化されます。
例えば以下のようなテキストファイルになります。

0 1.48 0.9245833333333333 2.7575 1.6825
3 0.81 1.05375 1.1675 1.0908333333333333
1 2.5796875 0.9845833333333334 0.105625 1.3275
1 2.43 0.9754166666666667 0.0875 1.3208333333333333
1 2.2865625 0.9770833333333333 0.101875 1.3125

この例の場合、この画像には5つの物体(ラベル0が1つ、ラベル2が3つ、ラベル3が1つ)があることになります。このファイルには画像と同じ名前をつけておき(train_01.jpgに対するアノテーションデータならtrain_01.txt)、train/labelsとval/labelsに保存していきます。

フォルダ構成やラベル設定を書いたyamlファイルを作成する

次にデータの場所やラベルの情報を記述したyamlファイルを作成します。cloneしてきたyolov5/dataフォルダ内に
あるyamlファイルを参考にしながら以下のように作りました。

train: file/train  #trainのパス
val: file/val  #valのパス

nc: 4  #ラベルの数
names: ['label_0', 'label_1', 'label_2', 'label_3']  #ラベルの名前

これをyolov5/dataフォルダの中に保存しておきます。

学習

ここまでで事前準備が終わったので学習に移ります。ターミナルを開いてyolov5のフォルダに移動し、以下を実行します。

python train.py --data data.yaml --cfg yolov5s.yaml --weights '' --batch-size 8 --epochs 200

それぞれのコマンドは以下のような意味になっています。

変数名	意味
--data	yolov5/data内にある設定が書いてあるyamlファイルの指定
--cfg	yolov5/models内にあるモデル構造の指定(pretrainモデルを指定しない場合に指定する)
--weights	重みを指定(推奨はyolov5s.ptなどのpretrainされたものを使う。これで指定するとダウンロードされる。''で指定すると1からの学習になり、--cbfで構造の指定が必要になる)
--batch-size	バッチサイズ
--epochs	エポック数

これを実行すると学習が始まります。(えらい時間かかりました。mAPなどのスコアを見ながら学習の様子を観察し、適当なところで止めました。)

うまくいっていればyolov5/runsフォルダの中にtrain/expフォルダが作成され、その中に色々な情報(各エポックのスコア一覧や訓練データのラベルの統計情報、またweights内に学習された重みのptファイルなど)が入っています。

検出

ターミナル上で以下を実行します。

python detect.py --source ../test/ --weights ./yolov5/runs/train/exp/best.pt --conf 0.5

それぞれのコマンドは以下のような意味になっています。

変数名	意味
--sorce	物体検出したい画像が入ったフォルダの指定
--weights	検出につかう重み
--conf	検出の閾値の設定

これを実行するとyolov5/runsフォルダ内にdetect/expフォルダが作成され、その中に検出結果が描画されて表示されます。(検出領域とラベル名、confidenceの値が描画されます)

数値的な結果はjupyter notebook上でpytouch を使ってみることができます。

!pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

import cv2
import touch

image = cv2.imread('./test/test_01.jpg', 1)

#検出器の準備
#自作した検出器を使用する場合
model = touch.hub.load('./yolov5', 'custom', path='./yolov5/runs/train/exp/best.pt', source='local')
#学習済みモデルを使用する場合
#model = touch.hub.load('ultralytics/yolov5', 'yolov5s')

#閾値の設定
model.conf = 0.5

#検出
results = model(image)
#複数画像を渡すこともファイルパスで渡すこともできます
#results = model(glob.glob('./test/*.jpg'))

#結果の確認
reults.print()
#--->
#image 1/1: 1200x1600 1 label_0, 2 label_1, 4 label_3
#Speed: 22.4ms pre-process, 952.5ms inference, 21.2ms NMS per image at shape (1, 3, 480, 640)

最後のreults.print()で簡単なサマリーをみることができます。

一行目には画像サイズと、それぞれの画像がいくつ入っていたか(例ではlabel_0が1つ、label_1が2つ、label_3が4つ)がわかります。また、二行目では処理時間が表示されています。

以下で各検出の座標を得ることができます。

#tensor形式での表示(results.xyxy自体はリストで返ってくる)
results.xyxy[0]
#--->
#tensor([[100.00, 100.00, 200.00, 200.00, 0.9534, 0.0000],
#        [30.123, 45.678, 70.345, 98.765, 0.8642, 1.0000], 
#                             (省略)

#DataFrame形式での表示１
results.pandas().xyxy[0]

#DataFrame形式での表示2
results.pandas().xywh[0]

xyxyとxywhの出力は以下の差があります。

形式	0列目	1列目	2列目	3列目	4列目	5列目	6列目
xyxy	左上x座標	左上y座標	右下x座標	右下y座標	confidence	class	ラベル名
xywh	中心x座標	中心y座標	横幅	高さ	confidence	class	ラベル名

なおラベル名のカラムはテンソルの出力だと存在しません。