今から始めるYOLOX

皆さんこんにちは、マツムラです。猛暑日が続き、PC排熱で部屋がより暑くなる季節になりました。

本記事ではGoogle ColabratoryでYOLOXを動かした後、一部改変を行い出力結果の表示を変更してみます。

これからYOLOXを触ってみたいという方も、サービスにYOLOXを利用したいという方も是非ご覧ください。

YOLOXって？

2021年に論文が発表された、物体検出アルゴリズムであるYou Only Look Once (YOLO)の派生の一つです。

ライセンスはApache-2.0 licenseとなっており、他のYOLO系の多くとは違い、商用利用がしやすくなっています。また、CVPR2021でおこなわれた Streaming Perception Challengeでは1位になっています。

YOLOXを動かしてみる

実際にYOLOXを動かしてみましょう。

今回はGoogle Colabratoryの環境を使ってYOLOXを試してみます。Google Colabratoryは、Googleアカウントを持っていれば、誰でも利用が可能です。ローカル環境でPythonを動かす際に比べて、難しい環境設定無しで手軽に動作させることができます。是非利用してみましょう。

ノートブック準備

ノートブックを新規作成を押します。

上のメニューから 編集＞ノートブックの設定 を選択します。

ハードウェア アクセラレータをGPUに変更し、保存を押します。

以上で設定は終わりです。ここからはコードを書いていきます。

基本的な操作ですが、＋コードを押しコードセルを作成します。実行するコードを記載後、左の▶を押し実行します。

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

公式プロジェクトをclone

公式プロジェクトをcloneしてきます。

!git clone https://github.com/Megvii-BaseDetection/YOLOX

ライブラリ等をインポート

%cd YOLOX
!pip install -U pip && pip install -r requirements.txt
!pip install -v -e .
!pip install cython
!pip install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

モデルをダウンロード

次にモデルをダウンロードします。

!wget https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_x.pth

これで動かす準備は完了です。

早速サンプル画像を判定していきます。

!python tools/demo.py image -n yolox-x -c yolox_x.pth --path assets/dog.jpg --save_result --device gpu

マイドライブ/YOLOX/YOLOX_output/yolox_x/vis_res/実行日時のフォルダ　に結果画像が出力されているので確認します。

無事にYOLOXで判定することができました。

YOLOXを少し改変してみる

動作確認ができたところで、出力される情報に少し制限を入れてみようと思います。

マイドライブ/YOLOX/tools/demo.py を開きます。

168行目にあるdef visualという関数で描画情報の設定が行われています。

今回は指定したカテゴリ(dog)のみが結果表示されるように改変しようと思います。

demo.pyの1行目に

import numpy as np

174行目に下記を追記してください。

# 指定したカテゴリのみ抽出
# 現在インデックス6(cls)の値が16(dog)を抽出
output = output[np.where(output[:, 6] == 16)[0]]

以上です。demo.pyを保存したのち、もう一度判定実行してみてください。

無事指定したカテゴリのみが結果表示されるようになりました。

今追記したoutput = output[np.where(output[:, 6] == 16)[0]]は、結果のカテゴリが16となっているものを抽出する処理でした。

カテゴリ一覧の確認方法ですが、マイドライブ/YOLOX/yolox/data/datasets/coco_classes.pyというファイルに記載されています。

COCO_CLASSES = (
    "person",
    "bicycle",
    "car",
    "motorcycle",
    "airplane",
    "bus",
    "train",
    "truck",
    "boat",
    "traffic light",
    "fire hydrant",
    "stop sign",
    "parking meter",
    "bench",
    "bird",
    "cat",
    "dog",
    "horse",
    "sheep",
    "cow",
    "elephant",
    .
    .
    .
    省略

dogは上から17番目にかかれている為、16を指定することで抽出が可能になります。

次に一番面積の大きいバウンディングボックスのみを表示する改変を行いたいと思います。

先ほど174行目に追記した分を削除します。

179行目(bboxes /= ratioの次の行)に下記を追記します。

# バウンディングボックスの面積が一番大きいものを探し抽出
area = (bboxes[:, 2]-bboxes[:, 0]) * (bboxes[:, 3]-bboxes[:, 1])
bboxes = bboxes[torch.argmax(area)]
bboxes = bboxes.reshape((-1, 4))
output = output[torch.argmax(area)]
output = output.reshape((-1, 7))

追記出来たら実行をします。

無事一番大きい面積のバウンディングボックスのみが表示されるようになりました。

まとめ

今回は
・指定したカテゴリのみを表示
・バウンディングボックスの一番大きいもののみを表示させる
ような改変を行いました。

他にも、指定した座標にあるバウンディングボックスだけを表示させることや、バウンディングボックスに付属されている確率が一番大きいものだけを表示させるような改変も可能です。

是非demo.pyの中を読んでみて、いろいろ試してみてください。