0. 概要

サーキット走行中のゼッケン１の車両とそれ以外の車両をマークした画像から物体検出の機械学習モデル yolox-xを用いてAIを学習させます。
ゼッケンを文字認識しているのではなく、車体の識別をしていることになります。
結果は上の動画の通りで、青い枠がゼッケン1と判定したバイクです。実際のゼッケン1は先頭から3番目を走行しています。動画でみたときに、ちょこちょこと間違えています。
AIの学習に用いたデータは、coco-annotatorを使って作成しました。
リアルタイムでの判別を試みたかったのですが、より小さいサイズのモデルであるyolox-nanoでは精度が悪く、とりあえず判定を最大精度が期待できるyolox-xを非同期で実施したという結果です。
開発環境はすべてGoogle colaboratory上に作成しました。

1. 手順

スマートフォンで物体検出のinstall.ipynbを使ってcolaboratoryから接続されたgoogle drive上にyoloxをインストールします。
coco-annotatorによってCOCO(Common Objects in Context)形式のカスタムデータを作成します。
colaboratory上で、yolox-xのモデルを学習させます。
学習結果のモデルファイルを用いて、動画からゼッケン1のバイク検出を実行します。
検出結果はGoogle Drive上に保存されるので、ダウンロードしてPC上で確認します。

2. yoloxのインストール

上記サイトの3. yoloxのGoogle Driveへのインストール (install.ipynb)を実行してください。基本的にinstall.ipynbファイルをgithubからダウンロードしてcolaboratoryにアップして実行だけすれば良いです。

3. カスタムデータの作成

COCO(common objects in context)形式のカスタムデータを作成します。カスタムデータとは、自分自身で作成する機械学習モデルのトレーニングデータです。画像からオブジェクトを検出するタスクの場合には、画像に対してどこに、何が映っているのかを指定したデータのことを指します。

カスタムデータを作成するための様々なツールが提案されていますが、今回は上記のcoco-annotatorを使用します。

coco annotatorのドキュメントにあるように、docker上にインストールします。

↑ の画面は、インストール後のDocker Desktopの画面です。Containersに追加されたcoco-annotatorの右側のActionsで再生の三角ボタンを押してcoco-annotatorを起動してください。

起動したcoco-annotatorへは、ブラウザで localhost:5000 で接続します。クラス値になるcategoriesを追加して、datasetsを作成し、画像をアップしてアノテーションを付けます。アノテーションは、物体をより詳細に指定する多角形で指定するものもありますが、yoloxはbboxと呼ばれる長方形を出力するので、上記のように長方形で指定すれば良いです。

↑ の例では、左側の先頭車両から3台目がゼッケン１の車両で、そこだけアノテーションのクラスが異なります。それ以外の車両は同じクラスとしています。

アノテーションを付けた後は、exportによってアノテーションデータを保存、ダウンロードします。このアノテーションデータと画像ファイルをあとでgoogle driveにアップします。

4. カスタムデータのアップロード

アノテーションのjsonファイルは、以下の2つの箇所にアップロードしてください。フォルダがない場合は適宜作成してください。

YOLOX/datasets/COCO/annotations/instances_train2017.json
YOLOX/datasets/COCO/annotations/instances_val2017.json

今回はとりあえず、トレーニングデータと検証データを同じデータで行います。本来は分けるものなのですけど、とりあえず実行して結果を見るためです。またその場合、検証用データを用意しないということが考えられるわけですが、yoloxの学習をまだしっかりと理解していないので、検証用データなしで動作させる方法が分からず、とりあえず同じデータを割り当てるということをしています。

画像ファイルは次のフォルダにそれぞれアップロードします。画像ファイルのファイル名は、coco-annotatorで作成したときから変更しないでください。

YOLOX/datasets/COCO/val2017/
YOLOX/datasets/COCO/train2017/

このフォルダ名は、次に示す学習をコントロールするexpファイルのデフォルトの場所です。expファイルでフォルダやファイル名を変更することもできるのですが、ここではその方法は省略いたします。

5. 学習をコントロールするexpファイルの作成

install.ipynbと同じ階層(/content/drive/MyDrive/Colab Notebooks/yolox/)に以下のような myexp_yolox_x.py ファイルを作成します。このファイルは、yoloxのモデルの学習や推論時に、ハイパーパラメータ等を制御するためのもので、YOLOX/exps/default/yolox_x.pyからコピーして最後の2行を追加したものです。

myexp_yolox_x.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 = 1.33
        self.width = 1.25
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

        # created next 2 line for this projects
        self.num_classes = 2
        self.max_epoch = 100

追加した個所は、クラス値が2つ、そして学習回数のepochをデフォルトの300から100に変更しています。これは、colaboratoryの無料枠の範囲で学習が終了するようにするためです。

6. モデルの学習(custom_dataset.ipynb)

ここで記述するプログラムはすべて上記サイトのcustom_dataset.ipynbに記載されています。またすでに上記のinstall.ipynbまたはそれ相当によってgoogle drive上にyoloxがインストールされていることを前提とします。

まずcolaboratoryのノートブックの設定で、GPUを選択します。ノートブックの設定は、colaboratoryの編集メニュー - ノートブックの設定から開くことができます。

Google Driveへの接続

# google drive のマウント
current_folder = "/content/drive/MyDrive/Colab Notebooks/yolox/YOLOX"
from google.colab import drive
drive.mount('/content/drive')
%cd $current_folder

最初に、google driveに接続します。上記の実行すると許可を求められるので許可してください。さらにセキュリティ確認のメールが送られてくるので自分で操作したことを確認してください。

loguru, thopのインストール

# yoloxで使用する loguru, thop のインストール
!pip install loguru
!pip install thop

機械学習モデルの再学習に必要なモジュールをインストールします。

各種のバージョンの確認

import numpy as np
import torch
import torch.onnx

# 各種のバージョンを確認する。
!python --version
print(f"numpy {np.__version__}")
print(f"pytorch {torch.__version__}")

Python 3.10.6
numpy 1.22.4
pytorch 2.0.1+cu118

ここで記事作成時点(2023.08.01)の各モジュールのバージョンを確認しておきます。上記のようになっています。

公式の学習済み重みファイルのダウンロード

# 学習済み重みファイルのダウンロード
!mkdir ../checkpoints
%cd ../checkpoints
!wget https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_x.pth
%cd ../YOLOX

yoloxの学習を行う時に使用する重みの初期値を、公式の学習済み重みファイルを用いてセットします。↑ のプログラムを実行すると、checkpointsフォルダを作成してそこにyolox_x.pthファイルをダウンロードします。

pythonモジュール検索パスに現階層を含める

# 現在のフォルダ (current_folder変数に入っている) をpythonのモジュール検索パスに含める。
import sys
if current_folder not in sys.path:
  sys.path.append(current_folder)
print("\n".join(sys.path))

pythonがモジュールを検索するときに使用するパスに、この時点での現階層(/content/drive/MyDrive/Colab Notebooks/yolox/YOLOX)を含めます。現階層は、current_folder変数に入っています。

モデルの学習

# スマートフォンでもなんとか動く nano のモデルを使用する場合
# %run tools/train.py -f ../myexp_yolox_nano.py -d 1 -b 8 --fp16 -c ../checkpoints/yolox_nano.pth

# リアルタイムをやめて非同期で動画ファイルから検出する場合に、最高精度の出る yolox-x のモデルを使用する場合
%run tools/train.py -f ../myexp_yolox_x.py -d 1 -b 8 --fp16 -c ../checkpoints/yolox_x.pth

yolox-xのモデルを学習します。コメントアウトしている上の部分は、yolox-nanoモデルの場合の例です。重みの初期値には、公式からダウンロードできる学習済みファイル(yolox_x.pth)を使用します。

学習時間には、約30分がかかります。

7. 検出の実行と結果の確認

動画ファイルから検出する

# トレーニングデータのうち1枚を使用して推論の実行をしてみる。
# 実行結果は、YOLOX/YOLOX_outputs/[expのファイル名]/vis_res/[日付のフォルダ(UTC)]/ に、画像でも動画ファイルでも保存される。

# トレーニングデータの画像から1つを選んで、試しに推論を実行してみる場合
# %run tools/demo.py image -f ../myexp_yolox_nano.py -c YOLOX_outputs/myexp_yolox_nano/best_ckpt.pth --path datasets/COCO/train2017/suzuka_chicane_camera_fixed_04.png --conf 0.25 --nms 0.25 --tsize 640 --save_result --device gpu

# 動画ファイルを指定してすべてのフレームに推論を実行してみる場合
%run tools/demo.py video -f ../myexp_yolox_x.py -c YOLOX_outputs/myexp_yolox_x/best_ckpt.pth --path ../suzuka_chicane_camera_moved_part.mp4 --conf 0.25 --nms 0.45 --save_result

上記のコメントアウトの部分は、動画ではなく画像ファイルで推論を実行する場合の例です。プログラムを実行すると、出力は YOLOX_outputs/モデル名/vis_resにUTC日付のフォルダが作成されてそこに動画ファイルが作られます。

推論は1つのフレームで約5秒程度かかります。よって例えば30fpsで15秒の動画から物体検出しようとすると、40分近い時間がかかります。

↑ のプログラムを実行して作成される動画が、この記事の最初で示した動画です。検証結果をしっかりと検証していませんが、目視で7割程度？ぐらいの認識率だと思います。ただこの認識率は高く見積もられています。その理由は、モデルの学習に使用したデータは、この動画から8枚の画像を抜き出して作成したからです。よって違う状況の動画では認識率は大きく下がると思われます。

また、ある程度はリアルタイムで物体検出できるyolox-nanoというサイズの小さいモデルを用いたところ、画像ベースで以下のような結果となりました。

上記の3番目のバイクがゼッケン１のバイクです。しかし他のバイクと同じと判定されています。この画像はモデルの学習に用いたものなので、この画像で判別ができないということは、他の画像でもほぼ判別できないということを意味します。

8. まとめ

物体検出のAIであるyoloxを用いて、レース中の特定の車両を見つけ出すAIを開発しました。結果は、yolox-xを用いて非同期で検出するとある程度できているようにも見えます。ただまだいろいろな工夫が必要そうです。一方リアルタイムで、例えば観戦中にスマートフォンをコースにかざしてお目当ての車両が映った時に知らせてくれる、などの使い方はまだちょっと難しそうです。

参考

yoloxの公式サイト(github)
COCO(common objects in context)について
 カスタムデータ作成ツール coco-annotator

ゼッケン１のバイクを捜せ！