自作データセットで物体検出モデル yolov9 を訓練する

mim

2024年3月3日 12:35

先月末に、物体検出モデル「YOLO」のバージョン 9 が発表されたので、手元の PC (Ubuntu 22.04) で動かしてみました。
論文はこちら
Github はこちら

(2024.3.24 追記) 論文紹介を書きました。

※この投稿は 2024 年 3 月 3 日時点 (v0.1) の情報に基づいています。コードは日々更新されており、時間が経てばもう少しユーザーに親切な構成になるとは思いますが、「今すぐ動かしたいんだ」という方向けに共有します。

動作の確認

リポジトリとデータのコピー

デスクトップに yolov9 のレポジトリをコピーします

cd Desktop
git clone https://github.com/WongKinYiu/yolov9
cd yolov9

pretrained というフォルダを作り、事前訓練済みのモデルもダウンロードします。現時点では v0.1 ですが、変わるかもしれません。

wget -P pretrained -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-c.pt
wget -P pretrained -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/yolov9-e.pt
wget -P pretrained -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-c.pt
wget -P pretrained -q https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-e.pt

同様に、お試しで検出する画像を sample_images フォルダにダウンロードします。

wget -P sample_images https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/bus.jpg
wget -P sample_images https://raw.githubusercontent.com/Megvii-BaseDetection/YOLOX/main/assets/dog.jpg

ライブラリのインストール

Anaconda の仮想環境を構築し、必要なライブラリをインストールします。python のバージョンは適当です (3.9 で動きました)。

conda create -n yolov9 python=3.9
conda activate yolov9
pip install -r requirements.txt

事前学習モデルによる検出

python detect.py --weights pretrained/yolov9-c.pt --source sample_images \n
  --data data/coco.yaml --device 0

AttributeError: 'list' object has no attribute 'device'

output

自分が動かしたときには、上のようなエラーが出ました。GitHub の issue を見ると、yolov9-c, yolov9-e で検出する場合には、utils/general.py の 903 行目を以下のように修正する必要があるそうです。

prediction = prediction[0][1]　# prediction[0]

この修正をしたうえでもう一度検出したところ、正常に動きました。runs/detect/exp フォルダに以下のファイルが出力されます。

自作データセットでの訓練

データセットの準備

データセットの構成はこれまでの yolo と同じです。画像名は任意ですが、同じ名前のテキストファイルを /labels に作る必要があります。

└─{データセット名}
    ├─train
    │  ├─images
    │  │   0001.jpg
    │  │   0002.jpg
    │  │   ...
    │  └─labels
    │      0001.txt
    │      0002.txt
    │      ...
    ├─val
    │  ├─images
    │  │   ...
    │  └─labels
    │      ...
    │
    └─test
       ├─images
       │   ...
       └─labels
           ...

データファイルの作成

data フォルダにある coco.yaml を参考に、データセットのパスとクラス名を指定します。自分は他のバージョンの yolo も使っているので、データセットは yolo9 フォルダの外に置いています。custom_dataset.yaml という名前で同じフォルダに保存します。

# Train/val/test sets
path: ../yolo_datasets/dataset1  # dataset root dir
train: train  # train images (relative to 'path')
val: val  # val images (relative to 'path')
test: test # test images (optional)

# number of classes
nc: 3

# Classes
names:
  0: tooth
  1: denticle
  2: saw-toothed

0から訓練

GitHub の記載に従って、とりあえず 3 epoch で訓練を実施してみます。

python train_dual.py --workers 8 --device 0 --batch 16 \n
  --data data/custom_dataset.yaml --cfg models/detect/yolov9-c.yaml \n
  --img 640 --weights '' --hyp hyp.scratch-high.yaml --min-items 0 --epochs 3

自分はプログラムは回ったものの、しつこく AttributeError が出ました。GitHub の Issue を参考に、Pillow を再インストールしたら直りました。

訓練中はログが出力されますが、かなり低い精度になっていると思います。引数 --weights を指定しない場合、ランダムに初期化された重みからの訓練になるようです。

事前学習モデルのファインチューン

そこで、先ほどダウンロードした事前学習モデルを使ってファインチューニングをしてみます。先ほどよりは精度が改善すると思います。

python train_dual.py --workers 8 --device 0 --batch 16 \n
  --data data/custom_dataset.yaml --cfg models/detect/yolov9-c.yaml \n
  --img 640 --weights pretrained/yolov9-c.pt --hyp hyp.scratch-high.yaml \n
   --min-items 0 --epochs 3

テスト

テストデータを使ってテストをします。--weights には先ほどの訓練で保存されたモデルのパスを指定します（exp を exp2, exp3 などと書き替えてください）。

python val.py --data data/custom_dataset.yaml --img 640 --batch 32 \n
  --conf 0.001 --iou 0.7 --device 0 \n
  --weights runs/train/exp/weights/best.pt --task test

結果は runs/val 以下に出力されます。PR曲線を比べてみると、事前学習モデルを使った方が精度が高いことが分かります。