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

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

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

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

※2023年11月13日「COCO形式へ変更」の部分を修正

今回の目標

第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形式へ変更

※Label Convertのアップデートに伴い記事の一部を修正しました (2023年11月13日)

まず、先に作成したデータセットをYOLOXで使用するためには、「COCO形式」に変換する必要があります。

この変換には専用のツールがあります。その詳細については、こちらのリンクでご確認いただけます。

本記事では、YOLOv5のデータセットをCOCO形式に変換する手順を解説します。

変換するデータセットは、すべて「YOLOV5」という名前のフォルダに収めておきましょう。

始めに、変換ツールをインストールする必要があります。

!pip install label_convert

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

以下の通りの構成となります。

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

「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
〜以下省略〜

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

!yolov5_to_coco --data_dir dataset/YOLOV5/traincar --mode_list train,val

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

出力後は以下のようなファイル構成になっています。

YOLOX
 ├── datasets
     ├── YOLOV5   
        ├── traincar ←ここにデータセットを格納する(フォルダ名は任意)
            |    ├── classes.txt
            |   ├── train.txt 
            |    ├── val.txt 
            |    ├── images ←画像はここに格納する
            |    └── labels ←テキストファイルはここに格納する
            ├── traincar_COCO_format←変換されたデータセット

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

「traincar_COCO_format」のフォルダの中身を見てみると、前回の記事でも紹介した形式になっていることがわかります。

この「traincar_COCO_format」を「traincar」にフォルダ名を変更して、「YOLOX/datasets/traincar」に移動します。

最終的に以下のような構成にします。

YOLOX
├── datasets
     ├── traincar←先ほど出力した「traincar_COCO_format」を「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)

コメントを残す