環境

・GoogleColaboratory

流れ

・学習データをアップロード
・YOLOXの学習
・学習したモデルを使って物体検出を試す

学習データをアップロード

GoogleColaboratory上で学習を行うので、学習データをGoogleDriveにアップロードする必要があります。

YOLOXで使う学習データは実際の画像ファイルと、画像のどこに物体が写っているかといった情報が記載されたjsonファイルの2種類になります。
それぞれ訓練用と検証用に分かれています。

今回使う学習データには第一回で作成した自作のデータセットを使用します．このデータはすでに訓練データと検証データに分けてあります。分ける割合におそらく正解はないと思いますが8:2くらいの割合が一般的なのではないかと思います。

第一回で作成したデータはそれぞれ
画像データがcoco-annotator/datasets/my_dog_trainとcoco-annotator/datasets/my_dog_validの中に
jsonデータがmy_dog_train.jsonとmy_dog_valid.jsonといった名前で保存してあるはずです．
これらのデータをGoogleDrive上に正しく配置していきます．

まずはこのjsonファイルの名前を下記のように変更します．

     変更前               変更後
my_dog_train.json  =>  train.json
my_dog_valid.json  =>  val.json

次にGoogleDrive上のYOLOX/datasets/下にmydog/ディレクトリを作って下記のようにこれらのファイルを配置していきます。
画像ファイルの名前はこの例の通りに変更する必要はありません．

YOLOX
└── datasets
    └── mydog
        ├── annotations
        │   ├── train.json      <-- 訓練用jsonファイル
        │   └── val.json        <-- 検証用jsonファイル
        │
        ├── train2017
        │   ├── image0.jpg      <-- 訓練用画像ファイル
        │   ├── image1.jpg      <-- 訓練用画像ファイル
        │   ├── image2.jpg      <-- 訓練用画像ファイル
        │        ・
        │        ・
        │        ・
        │
        ├── val2017
        │   ├── image0.jpg      <-- 検証用画像ファイル
        │   ├── image1.jpg      <-- 検証用画像ファイル
        │   ├── image2.jpg      <-- 検証用画像ファイル
        │        ・
        │        ・
        │        ・

YOLOXの学習

学習に使う設定ファイルを作成します。
今回はデフォルトで用意されているものをコピー編集して使います。
YOLOX/exps/default/yolox_s.pyをYOLOX/exps/example/customにコピーしてmydog.pyにリネームします。
そうしたらmydog.pyを編集していきます。

編集前

#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

import os

from yolox.exp import Exp as MyExp


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

編集後

#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

import os

from yolox.exp import Exp as MyExp


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


        # Define yourself dataset path
        self.data_dir = "datasets/mydog"
        self.train_ann = "train.json"
        self.val_ann = "val.json"

        self.num_classes = 1

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

編集するのは以下の部分になります。

        # Define yourself dataset path
        self.data_dir = "datasets/mydog"
        self.train_ann = "train.json"
        self.val_ann = "val.json"

        self.num_classes = 1

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

学習用のデータが入っているディレクトリの場所とアノテーションファイルを指定しています。
また、今回は学習させるクラスの数は一つだけなのでself.num_classes = 1としています。
その他学習パラメータの設定をしています．

次に、YOLOX/yolox/data/datasets/coco_classes.pyを編集します。このファイルには検出時に表示するクラスの名前が記載されています。デフォルトだといろんな種類のクラスが列挙されていますが、今回は1つのクラスのみ検出するので下記のように編集します。(検出時の表示に使うだけなので間違ったクラスを書いてしまっても実は問題ないです。そもそも編集しなくてもちゃんと動きます．)

#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

COCO_CLASSES = (
    "mydog",
)

ここまできたらあとは学習させるだけです．
前回同様GoogleColaboratoryでDriveをマウントして環境構築します．ランタイムの種類をGPUにするのを忘れずに．

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

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

%cd /content/drive/MyDrive/YOLOX
!pip install -U pip && pip install -r requirements.txt
!pip install -v -e .

環境が構築できたら学習させます．
学習はコマンド一発です．

!python tools/train.py -f exps/example/custom/mydog.py -d 1 -b 8 --fp16 -o -c yolox_s.pth

学習には時間がかかります．
学習が終わるとYOLOX/YOLOX_outputs/mydogの中に各エポックごとの学習済みデータと，中でも一番精度の良い学習済みデータとしてbest_ckpt.pthが保存されています．

学習したモデルを使って物体検出を試す

学習済みデータbest_ckpt.pthを使って物体検出を試してみます．
試す方法は前回とほとんど一緒です．

まずは検出試したい画像をYOLOX/の中に保存します．ファイル名はmydog.jpgとします．
そして下記のコマンドで検出を行います．

!python tools/demo.py image -f exps/example/custom/mydog.py -c ./YOLOX_outputs/mydog/best_ckpt.pth --path ./mydog.jpg --conf 0.3 --nms 0.45 --tsize 640 --save_result --device gpu

結果はYOLOX/YOLOX_outputs/mydog/vis_resの中に日付ごとに保存されています．
98.4%の確率で私の超可愛い愛犬が検出できていますね．
みなさんはうまく検出できましたでしょうか．