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

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

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

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

第3回目である今回は物体検出のためのデータセット作成の方法と学習の実装を紹介します。

実際にデータセットを用意して、物体検出のモデルを作成してみましょう。

この記事では電車を検出するモデルを作成してみることにします。

yolov5におけるモデル作成に必要なデータセットとは

物体検出の学習で必要となる情報は「画像」、「クラス（人、自動車など対象物）」、「クラスの位置（座標）」となります。

これらをまとめて機械が学習できるように整理したものをデータセットといいます。

難しそうに聞こえるかもしれませんが、画像に対して何がどこに写っているのかを教えてあげればいいわけです。

以下の流れに沿って、実際にデータセットを作成していきましょう。

STEP1.画像収集

まずは数百枚程度の画像を収集しましょう。

手動で集めることも可能ですが、例としてスクレイピングで画像収集する方法を紹介します。

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

スクレイピングに必要なライブラリをインストールします。

今回は「icrawler」を使用します。

!pip install -qr requirements.txt

ここでは「犬」の画像を集める例を示します。

from icrawler.builtin import BingImageCrawler
crawler = BingImageCrawler(storage={"root_dir": "dogs"})
crawler.crawl(keyword="犬", max_num=10)

#keyword="犬":集めたい画像のキーワードを指定する
#max_num=10:収集する画像の最大枚数（あまり大きくしすぎない）

これを実行すると、「dogs」フォルダに犬の画像が10枚ほど保存されています。

画像を収集する際は「犬 ブルドッグ」や「犬 ペット」などキーワードを変えることで、多くの画像を取得することができます。

STEP2.アノテーション

画像に対して「クラス（人、自動車など対象物）」と「クラスの位置（座標）」を明示する作業をアノテーションといいます。

YOLOv5では、学習のために画像(.jpg)とアノテーションファイル（.txt）のペアが必要です。

画像「zidane.jpg」に対して、アノテーションファイル「zidane.txt」を作成していきます。

実際に画像と共にファイルの中身を見てましょう。

# [oject-class] [x_center] [y_center] [width] [height]
0 0.480000 0.630000 0.690000 0.710000
0 0.740000 0.520000 0.310000 0.930000
27  0.360000 0.790000 0.070000 0.400000

#oject-class:クラスの番号。yamlファイルの内容と合わせる
#x_center:bounding box 中心のx座標
#y_center:bounding box 中心のy座標
#width:bounding box x方向長さ
#height:bounding box y方向長さ

各数値と画像を見比べると、意味がわかるかと思います。

いずれの数値も画像の各方向の長さに対する比率であることに注意してください。

手作業で１枚ずつ数値を入力することもできますが、効率よく行うためにアノテーションツールを使うことをお勧めします。

物体検出用のアノテーションツールとして、VoTT, labelImgなどいくつか種類があるので使いやすいものを使用してください。

STEP3.yamlファイルの作成

データセット画像とアノテーションファイルの格納先、クラスの種類などの情報をひとまとめにしたyamlファイルを作成します。

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

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

nc: 1
names: ['traincar',]

以上でデータセットの作成は完了となります。

これらのデータを前回と同様にそれぞれ格納していきます。

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

なお、データセットの完成形は前回の記事でも紹介しております。

学習がうまくできない場合には確認をしてみましょう。

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

学習

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

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

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

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

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

学習結果

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

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

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

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

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

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

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

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

まとめ

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

今回はデータセットの作成方法と学習の実装を紹介しました。

データセットが用意できれば、簡単に高性能なオリジナルモデルが作れそうですね。

次回は作成したモデルとtorch hub用いたテストの方法を紹介します。

ぜひご覧ください。