この記事は2023年8月9日に弊社ブログにて公開されたものを移植したものです。
古い情報が含まれている可能性がありますのでご注意ください。

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

今回は下記の改変を行いました。

指定したカテゴリのみを表示
バウンディングボックスの一番大きいもののみを表示させる

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

YOLOXって？

2021年に論文が発表された、物体検出アルゴリズムであるYou Only Look Once (YOLO)の派生の一つです。
CVPR2021でおこなわれた Streaming Perception Challengeでは1位になっています。

YOLOXを動かしてみる

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

ノートブック準備

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

上のメニューから編集＞ノートブックの設定を選択します。
ハードウェアアクセラレータをGPUに変更し、保存を押します。

以上で設定は終わりです。ここからはコードを書いていきます。
基本的な操作ですが、＋コードを押しコードセルを作成します。実行するコードを記載後、左の▶を押し実行します。

自身のGoogleドライブに接続します。

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の中を読んでみて、いろいろ試してみてください。