【物体検出】YOLOXまとめ｜第4回：ONNXRuntimeとテスト結果（スコア、座標）の出力方法

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

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

第4回目である今回は、YOLOXでpytorchモデルをONNX変換する方法と、ONNXRuntimeを実行する方法を紹介します。

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

今回の目標

・YOLOXでpytorchモデルをOnnxに変換する方法

・Onnxruntimeデモの方法

・テスト結果の出力方法

1. ONNXとは
2. YOLOXpytorchモデルをONNXに変換
2.1. YOLOX の導入
2.2. 標準のYOLOXモデルをONNXに変換
3. YOLOXでONNXRuntime
3.1. 物体検出結果（スコアや座標など）を表形式で出力する
4. オリジナルモデルでONNXRuntime
4.1. 作成したYOLOXモデルをONNXに変換
4.2. 作成したYOLOXモデルでテスト
5. まとめ

ONNXとは

ONNX（Open Neural Network Exchange）とは、Microsoftがオープンソースで公開している学習モデルを表現するフォーマット形式のことです。

PyTorch、Tensorflow、 Scikit-learnなどの各種フレームワークで学習したモデルを、別のフレームワークで読み込めるようになります。

CPUやCUDA、TensorRTなど様々なデバイスに対応しており、学習済みモデルの管理や運用が容易になります。

ONNX Runtime はサーバーサイドの推論を高速化するだけでなく、ONNX Runtime for Mobile も提供されています。

さらに2021年9月にはORT Web が新た供されたことで、ブラウザー内での推論が可能になりました。

YOLOXpytorchモデルをONNXに変換

早速、YOLOXのPytorchモデルをONNXに変換する方法を紹介します。

ここからはGoogle colabを使用して実装していきます。

まずはGPUを使用できるように設定をします。

「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」をGPUに変更

今回紹介するコードは以下のボタンからコピーして使用していただくことも可能です。

YOLOX の導入

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

すでにYOLO Xを使用している方はクローンは不要です。

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

公式よりcloneしてきます。

!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

今回初めてYOLOXを使用する方は、環境構築確認のためテストしてみましょう。

!python tools/demo.py image -n yolox-x -c yolox_x.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu

標準のYOLOXモデルをONNXに変換

ここからはYOLOXのモデルをONNXに変換してみましょう。

まずは標準のYOLOXモデルを変換します。

以下の例では先ほどダウンロードした「yolox_x.pth」を「yolox_x.onnx」に変換します。

変換する際はサイズを統一する必要がありますので-n をyolox-xとします。

!python tools/export_onnx.py --output-name yolox_x.onnx -n yolox-x -c yolox_x.pth

# -n：ベースとなるモデル名。[yolox-s、m、l、xおよびyolox-nane、yolox-tiny、yolov3]から選択。
# -c：変換前のモデル名
# --output-name：変換後のモデル名

変換が終了しました。

YOLOXでONNXRuntime

先ほど変換したモデルを使ってテストをしてみましょう。

まずは必要なライブラリをインポートします。

import argparse
import os
import cv2
import numpy as np
import pandas as pd
import onnxruntime
from google.colab.patches import cv2_imshow
from yolox.data.data_augment import preproc as preprocess
from yolox.data.datasets import COCO_CLASSES
from yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis

今回使用する画像とモデルを指定します。

「image_path 」にはテストする画像、「model」には先ほど変換したモデル名を入力します。

output_dir ='onnx_out'
image_path = 'yolox_test1.jpeg'
model = 'yolox_x.onnx'

テストを実行します。

input_shape = (640,640)
origin_img = cv2.imread(image_path)
img, ratio = preprocess(origin_img, input_shape)
session = onnxruntime.InferenceSession(model)
ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
output = session.run(None, ort_inputs)
predictions = demo_postprocess(output[0], input_shape)[0]
boxes = predictions[:, :4]
scores = predictions[:, 4:5] * predictions[:, 5:]
boxes_xyxy = np.ones_like(boxes)
boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
boxes_xyxy /= ratio
dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.5)
if dets is not None:
    final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
    origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
                      0.3, class_names=COCO_CLASSES)

mkdir(output_dir)
output_path = os.path.join(output_dir, os.path.basename(image_path))

cv2.imwrite(output_path, origin_img)
cv2_imshow(origin_img)

ONNXでテストすることができました。

物体検出結果（スコアや座標など）を表形式で出力する

先ほどの結果を表形式で表示してましょう。

result = []
[result.extend((final_cls_inds[x],COCO_CLASSES[int(final_cls_inds[x])],final_scores[x],final_boxes[x][0],final_boxes[x][1],final_boxes[x][2],final_boxes[x][3]) for x in range(len(final_scores)))]
df = pd.DataFrame(result, columns = ['class-id','class','score','x-min','y-min','x-max','y-max'])
df

以下の表を得ることができました。

オリジナルモデルでONNXRuntime

次に自分で作成したモデルのONNX変換とテストを試してみましょう。

今回は、以前作成したマスク着用の判定モデルを使用します。

詳細は以下の記事よりご確認ください。

作成したYOLOXモデルをONNXに変換

先ほどと同様に変換をしていきましょう。

変換前のモデルには作成したマスク着用判定モデル、「-f」には学習時に設定したexpsファイルを指定します。

!python tools/export_onnx.py --output-name face_mask.onnx -f exps/example/custom/face_mask.py -c YOLOX_outputs/face_mask/last_epoch_ckpt.pth

# -n：ベースとなるモデル名。モデル名は、[yolox-s、m、l、xおよびyolox-nane、yolox-tiny、yolov3]から選択。
# -c：変換前のモデル名
# --output-name：変換後のモデル名
# -f：学習時に使用したexpsファイル

作成したYOLOXモデルでテスト

変換したモデルで、テストをしてみましょう。

output_dir ='onnx_out'
image_path = 'OK.jpg'
model = 'face_mask.onnx'
    
input_shape = (640,640)
origin_img = cv2.imread(image_path)
img, ratio = preprocess(origin_img, input_shape)
session = onnxruntime.InferenceSession(model)
ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
output = session.run(None, ort_inputs)
predictions = demo_postprocess(output[0], input_shape)[0]
boxes = predictions[:, :4]
scores = predictions[:, 4:5] * predictions[:, 5:]
boxes_xyxy = np.ones_like(boxes)
boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
boxes_xyxy /= ratio
dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.5)
if dets is not None:
    final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
    origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
                      0.3, class_names=COCO_CLASSES)
mkdir(output_dir)
output_path = os.path.join(output_dir, os.path.basename(image_path))
cv2.imwrite(output_path, origin_img)