YOLOv5 with Oriented Bounding Box

こちらの記事で、Oriented Bounding Boxのアノテーションについて説明しました。この記事では作成されたOriented Bounding Boxのアノテーションデータを使ってYOLOv5の学習・推論方法について解説します。

オリジナルのYOLOv5はOriented Bounding Box(OBB)を扱うことができません。YOLOv5にOriented Bounding Box拡張が行われたYolov5 for Oriented Object Detection(yolov5_obb)を使用します。

アノテーションデータの変換

こちらの記事でOriented Bounding Boxのアノテーションのフォーマットを解説しましたが、<robndbox>...</robndbox>の中に<cx>、<cy>、<w>、<h>、<angle>があり、それぞれ、Orineted Bounding Boxの中央のx座標、中央のy座標、幅、高さ、角度となっていました。

https://github.com/hukaixuan19970627/yolov5_obb/blob/master/docs/GetStart.mdによると、yolov5_obbのOriented Bounding Boxは(x1, y1, x2, y2, x3, y3, x4, y4)で表され、(xi, yi)はi番目のOriented Bounding Boxを表し、座標は時計回りに移動するようです。また座標の後に名前そしてdifficultyが続くようです。difficultyは0で構いません。yolov5_obbではスペース区切りで座標と名前そしてdifficultyを以下のように記述しています。

x1 y1 x2 y2 x3 y3 x4 y4 name 0
Code language: plaintext (plaintext)

XMLアノテーションデータをyolov5_obbアノテーションデータに変換するスクリプトをこちらからダウンロードしアノテーションデータの変換を行います。なお以下で変換するXMLアノテーションデータはこちらの記事で作成したアノテーションデータとなります。

git clone https://github.com/otamajakusi/yolov5-utils.git
python3 yolov5-utils/voc2yolo5_obb.py --path data/car-parking.xml --class-file data/classes.txt
Code language: Bash (bash)

変換されたアノテーションデータdata/car-parking.txtが作成されていることが分かります。

384.90906545146345 769.5000055959972 274.67376491516154 714.9137141146996 301.27973454853657 661.1837944040028 411.5150350848385 715.7700858853004 car 0
416.90906545146345 713.5000055959972 306.67376491516154 658.9137141146996 333.27973454853657 605.1837944040028 443.5150350848385 659.7700858853004 car 0
443.90906545146345 656.5000055959972 333.67376491516154 601.9137141146996 360.27973454853657 548.1837944040028 470.5150350848385 602.7700858853004 car 0
484.90906545146345 601.5000055959972 374.67376491516154 546.9137141146996 401.27973454853657 493.1837944040028 511.5150350848385 547.7700858853004 car 0
...
Code language: plaintext (plaintext)

アノテーションデータが正しく作成されていることを以下のスクリプトで確認することができます。

python3 yolov5-utils/show_obb.py --image data/car-parking.jpg --anno data/car-parking.txt
Code language: Bash (bash)

なおこのスクリプトはyolov5_obbアノテーションデータだけでなく、XMLアノテーションデータに対しても確認することができます。

学習データとテストデータの分割

自身のGPUを使用して学習を進めることもできますが、セットアップの手軽さからここでは、Google Colaboratoryで作業を進めます。Google Colaboratoryにアクセスし、ノートブックを新規作成、を選びます。

GPUインスタンスを利用するためハードウェアアクセラレータを設定します。

ランタイム > ランタイムのタイプを変更、を選択します。

ハードウェアアクセラレータ > GPU、を選択します。

次に、これから学習を始めるにあたって、Colaboratoryを閉じてもデータを残しておけるように、GoogleログインアカウントのGoogleドライブにColaboratoryを接続し、そこにデータを保存します。

フォルダアイコンをクリックし、ドライブをマウントをクリックします。

Google ドライブに接続、をクリックします。

Googleドライブディレクトリに移動します。

%cd drive/MyDrive
Code language: Bash (bash)

Colaboratoryでは実行するコマンドを以下のテキストボックスに入力し、再生ボタンを押しコマンドを実行します。また先頭の % や ! も入力する必要があります。

yolov5_obbをクローンします。

!git clone https://github.com/hukaixuan19970627/yolov5_obb
%cd yolov5_obb
Code language: Bash (bash)

さて学習するためにはたくさんの学習データ(画像とそのアノテーションデータ)が必要になります。今回練習のためこちらで用意したアナログメータの学習データを使用します。

!gdown 1XwnMKQKSyyvMTfcHNDUu7sqnJpzvgChY
!unzip analog-meter.zip
Code language: Bash (bash)

画像とXMLアノテーションデータは1というディレクトリにあります。またクラスファイルはclasses.txtにあります。アノテーションデータを変換するスクリプトをダウンロードしアノテーションデータを変換します。なおアノテーションデータを変換するスクリプトは、--path にディレクトリを指定するとディレクトリ以下のすべての*.xmlを変換します。

!git clone https://github.com/otamajakusi/yolov5-utils
!python3 yolov5-utils/voc2yolo5_obb.py --path 1 --class-file classes.txt
Code language: Bash (bash)

次にデータをトレーニングデータとテストデータに分割します。 yolov5_obbは以下のディレクトリ構成を必要とします。

├── images
│   ├── train
│   │   ├── 0123.jpg
│   │   ├── ...
│   │   └── 4567.jpg
│   └── valid
│       ├── 89ab.jpg
│       ├── ...
│       └── cdef.jpg
└── labelTxt
    ├── train
    │   ├── 0123.txt
    │   ├── ...
    │   └── 4567.txt
    └── valid
        ├── 89ab.txt
        ├── ...
        └── cdef.txt
Code language: plaintext (plaintext)

yolov5はlabelsディレクトリを必要としますが、yolov5_obbはlabelTxtディレクトリを必要とする点が異なります。

データ分割スクリプトでデータを分割します。データ分割スクリプトはimagesとlabelsディレクトリに分割するため、yolov5_obbに対応するためlabelsディレクトリにlabelTxtのシンボリックリンクを作成します。

!python3 yolov5-utils/data_split.py --datapath 1 --outpath analog-meter
!(cd meter;ln -s labels labelTxt)
Code language: Bash (bash)

yolov5_obbはtorch version 1.12に対応していないため予めyolov5_obbが対応するtorchライブラリをインストールし、その後pipでyolov5_obbのrequirements.txtをインストールします。

!pip3 install torch==1.10.1+cu113 torchvision==0.11.2+cu113 torchaudio==0.10.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
Code language: Bash (bash)

ライブラリのビルドとパッケージのインストールを行います。

!(cd utils/nms_rotated; python3 setup.py develop)
!python3 -m pip install -r requirements.txt
Code language: Bash (bash)

今回学習するデータの設定データを作成します。2つのクラスneedle(針)とscale(目盛り)です。

%%bash
cat << EOS > data/analog-meter.yaml

train: analog-meter/images/train/    # train images (relative to 'path') 
val: analog-meter/images/valid/   # val images (relative to 'path') 

# Classes
nc: 2  # number of classes
names: [
  'needle',
  'scale',
]
EOS
Code language: Bash (bash)

学習

学習を開始します。

!python3 train.py --data data/analog-meter.yaml --cfg models/yolov5m.yaml --batch-size 4 --device 0 --epochs 10000
Code language: Bash (bash)

手元の環境では500回程度で、学習が打ち切られるようです。

tensorboardを表示してみます。

%load_ext tensorboard
%tensorboard --logdir runs/train/exp
Code language: Bash (bash)

自分のPCで学習する際はリアルタイムにグラフが見られるのですが、Colaboratoryでは学習が終わった後にしか見ることができません。

推論

テストデータの推論をしてみます。

!python detect.py --weights runs/train/exp/weights/best.pt --source analog-meter/images/valid --device 0 --hide-conf --hide-labels
Code language: Bash (bash)

作成された画像を見てみます。

import glob
from IPython.display import Image,display_jpeg

imgs = glob.glob("runs/detect/exp/*")
for img in imgs:
  display_jpeg(Image(img))
Code language: Bash (bash)

結果あまり良くありませんが、学習データを増やすことで改善されるかもしれません。

以上です。

参照

• https://github.com/hukaixuan19970627/yolov5_obb
• https://github.com/otamajakusi/yolov5-utils.git
• https://captain-whu.github.io/DOTA/dataset.html