このシリーズでは物体検出でお馴染みのYOLOシリーズの最新版「YOLOX」について、環境構築から学習の方法までまとめます。

YOLOXは2021年8月に公開された最新バージョンであり、速度と精度の面で限界を押し広げています。

YOLOv5を超える性能で、使いやすいApache Licenseとなっています。

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

今回の目標

第3回目である今回は、YOLOXで必要となるデータセットの作成方法を中心に紹介します。

オリジナルのデータセットでアノテーションを行い、データの変換をした上で独自の物体検出モデルを作成していきましょう。

今回の内容

・YOLOXで必要となるデータセットの作成方法

・YOLOv5で作成したデータセットをYOLOX用のデータセットに変換

・独自の物体検出モデルを作成

データセットのための画像収集

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

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

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

%cd ./drive/MyDrive
#!git clone https://github.com/Megvii-BaseDetection/YOLOX

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

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

%cd YOLOX
!pip install icrawler

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

from icrawler.builtin import BingImageCrawler
crawler = BingImageCrawler(storage={"root_dir": "dogs"})
crawler.crawl(keyword="犬", max_num=10)
#keyword="犬":集めたい画像のキーワードを指定する
#max_num=10:収集する画像の最大枚数(あまり大きくしすぎない)

これを実行すると、「YOLOX/dogs」に犬の画像が10枚ほどアップされています。

画像を収集する際は「犬 ブルドッグ」や「犬 ペット」などキーワードを変えて収集してみましょう。

アノテーション

画像の中で物体検出したい対象の座標を明示する作業をアノテーションといいます。

例としてYOLOv5でのデータセットの場合には、学習のために画像(.jpg)とアノテーションファイル(.txt)のペアが必要となります。

画像「dog1.jpg」に対して、アノテーションファイル「dog1.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:枠線中心のx座標
#y_center:枠線中心のy座標
#width:枠線x方向長さ
#height:枠線y方向長さ

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

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

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

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

COCO形式へ変更

次に先ほど作成したデータセットをYOLOXで使用するため、「COCO形式」に変換していきます。

サードパーティにCOCO形式に変換するツールが用意されています。

詳細はhttps://github.com/RapidAI/YOLO2COCO をご確認ください。

ここからはYOLOv5のデータセットから変換する方法を紹介します。

!git clone https://github.com/RapidAI/YOLO2COCO

次に「YOLOX/YOLO2COCO/dataset/YOLOV5」に先ほど作成したデータセットを格納していきます。

格納場所は以下の通りです。

YOLOV5
├── classes.txt ←クラス一覧を記述
├── traincar ←ここにデータセットを格納する(フォルダ名は任意)
│   ├── images ←画像はここに格納する
│   └── labels ←テキストファイルはここに格納する
├── train.txt 
└── val.txt 

「classes.txt」・「train.txt」・「val.txt」は以下の例に従って書き換えてください。

#クラス一覧を記述(今回は1種類なので1行のみ)
traincar

#複数ある場合は1行ずつ書いていく
dataset/YOLOV5/traincar/images/00001.jpg
dataset/YOLOV5/traincar/images/00002.jpg
dataset/YOLOV5/traincar/images/00003.jpg
〜以下省略〜
dataset/YOLOV5/traincar/images/01001.jpg
dataset/YOLOV5/traincar/images/01002.jpg
dataset/YOLOV5/traincar/images/01003.jpg
〜以下省略〜

ここまで準備がきでたら、以下のコードを実行しましょう。

!python yolov5_2_coco.py --dir_path dataset/YOLOV5

結果は「YOLOX/YOLO2COCO/dataset/YOLOV5_COCO_format」に出力されます。

これでデータの変換は完了です。

前回の記事でも紹介した形式になっていることがわかるかと思います。

これらのフォルダを「YOLOX/datasets/traincar」※に移動します。

※traincarの部分は変更しても構いません。

データセットの格納

学習に必要な情報を記載した「expsファイル」を用意します。

まずは書き換えるための元データを用意します。

YOLOX/exps/example/customにある「yolox_x.py」をダウンロードしましょう。

これを以下のように書き換えます。

import os
from yolox.exp import Exp as MyExp

class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.depth = 1.33
        self.width = 1.50
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

        # Define yourself dataset path
        self.data_dir = "datasets/traincar" #※先ほど指定した「YOLOX/datasets/traincar」のフォルダ名
        self.train_ann = "instances_train2017.json"
        self.val_ann = "instances_val2017.json"

        self.num_classes = 1
        self.max_epoch = 300
        self.data_num_workers = 4
        self.eval_interval = 1

書き換えが終わったら、「traincar.py」とファイル名を変更して「YOLOX/exps/example/custom」にアップしましょう。

また、前回の記事でも紹介したように、ラベルの一覧表である「coco_classes.py」ファイルの中身を書き換えます。

これを行わないと、テスト時に正しくラベル付けが行われないので注意しましよう。

「YOLOX/yolox/data/datasets/coco_classes.py」を開いて、中身を編集します。

今回は以下のようにしました。

COCO_CLASSES = (
    "traincar",
)

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

学習

学習はコード1行で始まります。

ここでは先ほどアップしたtraincar.pyを指定しています。

なお、入力画像のサイズや学習の実行回数を変更したい場合は、先ほど作成した「expファイル」から変更します。

!python tools/train.py -f exps/example/custom/traincar.py -d 1 -b 4 --fp16 -o -c yolox_x.pth

学習が終わると、結果が「YOLOX/YOLOX_outputs/traincar/」に「best_ckpt.pth」として保存されます。

結果

実際にテスト画像を使って推論をしてみます。

使用するモデルは先ほど学習したものを指定しています。

TEST_IMAGE_PATH = "test.jpg"
MODEL_PATH = "YOLOX_outputs/traincar/best_ckpt.pth" 
!python tools/demo.py image -f exps/example/custom/traincar.py -c {MODEL_PATH} --path {TEST_IMAGE_PATH} --conf 0.7 --nms 0.45 --tsize 640 --save_result --device gpu

実行すると、結果が「YOLOX/YOLOX_outputs/traincar/vis_res/」に保存されます。

電車の車両の検出ができています。

YOLOv5のデータセットを変換して、YOLOXでモデルを作成することができました。

まとめ

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

今回はYOLOv5のデータセットを変換して、YOLOXでモデルを作成する方法を紹介しました。

以前のYOLOシリーズで使用したデータセットがあれば、簡単に高性能なオリジナルモデルが作れそうですね。

参考

学習時に下記のようなエラーが発生する場合がある。

TypeError: 'numpy.float64' object cannot be interpreted as an integer

この場合は「cocoeval.py」の506〜507行目を以下のように書き換える。

self.iouThrs = np.linspace(.5, 0.95, 10, endpoint=True)
self.recThrs = np.linspace(.0, 1.00, 101, endpoint=True)
よろしければおすすめ書籍もご覧ください。 ※これ以降は広告です

コメントを残す