YOLOv8のclassificationモデルをJetson nanoで動かしてみた

はじめに

この記事では、JetsonにYOLOv8のclassificationモデルを組込んでDeepStream上でリアルタイムに分類を行う方法を紹介します。参考になれば幸いです。

目的

YOLOv8のclassificationモデルをJetsonのDeepStream上で動かせるようにする。サンプルアプリdeepstream-appを使って推論を行う。

環境

• Jetson Nano 4GB
• Jetpack 4.6
• DeepStream SDK 6.0

組込み

以前作成した花の2クラス分類モデル（daisy, dandelion）を使用する。

組込み手順

1. モデルのTensorRT化: モデルをDeepStreamで使える形式に変換する。デバイスに合わせてモデルを高速化する。
2. コンフィグファイルの作成: 前処理の設定などを行う。
3. パーサー作成: 後処理を行うファイルを作成する。

1. TensorRT化

.pt -> .onnx -> .planの順に変換する。　　
.onnxファイルまではgoogle colabで変換し、.planファイルへの変換はJetsonで行う。

.ptファイルの変換

下記から重みファイルyolov8s_cls_flower.ptをダウンロードし、GoogleDriveの作業ディレクトリに保存する。
https://drive.google.com/drive/folders/1DN1pciFgE-U_ZmoCjcvAQV-1XUB4aAau?usp=sharing

Google colabを開き、Google driveをマウントする。
作業ディレクトリに移動し、以下を実行する。

!pip3 install ultralytics
!pip3 install onnx

from ultralytics import YOLO
import torch

重みファイルのパスを指定してモデルを読み込む。

file_path = "yolov8s_cls_flower.pt"
model = YOLO(file_path)

以下を実行し、.onnxファイルに変換する。
今回はinput size = 416 x 416で学習したため、input_w, input_hに416を設定する。

input_w = 416
input_h = 416

onnx_file_name = "yolov8s_cls_flower.onnx"

dummy_input = torch.randn(1, 3, input_h, input_w, device="cpu")
input_names = ["input1"]
output_names = ["output1"]

torch.onnx.export(model.model.eval(), 
                  dummy_input, 
                  onnx_file_name, 
                  verbose=True, 
                  input_names=input_names, 
                  output_names=output_names, 
                  opset_version=11)

作業ディレクトリに.onnxファイルが生成する。

.onnxファイルの変換

Jetsonで行う。
.onnxがインストールされていない場合は先にインストールする。
numpyのバージョンが1.19.5以上だと変換ができないことがあるため、numpy==1.19.4をインストールする。

pip3 install onnx==1.11.0
pip3 install numpy==1.19.4

以下のリポジトリをダウンロードし、yolov8_cls_tensorrtフォルダをJetsonのホームディレクトリにコピーする。
https://github.com/mihara-shoko/yolov8_cls

yolov8_cls_tensorrtフォルダに先ほど変換した.onnxファイルをコピーする。

yolov8_cls_tensorrtフォルダに移動して以下を実行し、.planファイルに変換する。

python3 engine.py yolov8s_cls_flower.onnx . yolov8s_cls_flower.plan

yolov8_cls_tensorrtフォルダにyolov8s_cls_flower.planが生成する。

2. コンフィグファイルの作成

前処理、モデルファイル、信頼度の閾値、クラス数の設定などを行う。
先ほどダウンロードした、yolov8_clsフォルダの中のyolov8_cls_parserフォルダをJetsonの/opt/nvidia/deepstream/deepstream-6.0/sources/ディレクトリにコピーする。
「1. モデルのTensorRT化」で作成したyolov8s_cls_flower.planファイルを/opt/nvidia/deepstream/deepstream-6.0/sources/yolov8_cls_parser下にコピーする。
/opt/nvidia/deepstream/deepstream-6.0/sources/yolov8_cls_parser下にあるconfig_infer_primary.txtを必要に応じて修正する。

config_infer_primary.txt

[property]
gpu-id=0
net-scale-factor=0.0039215697906911373
model-color-format=0

# モデルファイルのパス
model-engine-file=yolov8s_cls_flower.plan　

batch-size=1
network-mode=0
interval=0
gie-unique-id=1
process-mode=1
network-type=1

maintain-aspect-ratio=1
symmetric-padding=1

前処理に関係する項目

項目	説明	yolov8の場合の設定値
net-scale-factor	ピクセル値にかける値 モデルにはyのピクセル値が入力される。 y = net scale factor*(x-offset)	0.0039215697906911373 0-255を0-1のrangeにしている
offset	ピクセル値から引く値 RGBそれぞれ別に設定できる。 モデルにはyのピクセル値が入力される。 y = net scale factor*(x-offset)	0:0:0 設定しなくて良い
model-color-format	0:RGB 1:BGR 2:GRAY	0
maintain-aspect-ratio	リサイズするときにアスペクト比を維持するかどうか	1

詳しくはNVIDIAのページを参照
https://docs.nvidia.com/metropolis/deepstream/dev-guide/text/DS_plugin_gst-nvinfer.html

3. パーサー作成

後処理を行うファイルを作成する。今回は最も信頼度の高いクラスを出力する処理を記載している。
Jetsonの/opt/nvidia/deepstream/deepstream-6.0/sources/yolov8_cls_parserに移動する。
nvdsinfer_customparse_yolov8_cls.cppのクラス数、クラスラベルのリストを必要に応じて修正する。

nvdsinfer_customparse_yolov8_cls.cpp

#include <cstring>
#include <iostream>
#include <string.h>
#include "nvdsinfer_custom_impl.h"

// クラス数
int num_class = 2;

// クラスラベルのリスト
std::vector<std::string> class_list = { "daisy", "dandelion" };

コンパイルする。

cd /opt/nvidia/deepstream/deepstream-6.0/sources
make -C yolov8_cls_parser

yolov8_cls_parserフォルダ下に生成した.soファイルがパーサーとなる。

「2. コンフィグファイルの作成」で作成したコンフィグファイルにパーサーの情報を追記する。

config_infer_primary.txt

parse-classifier-func-name=NvDsInferParseCustomYolov8Cls
custom-lib-path=libnvds_infercustomparser_yolov8_cls.so

アプリ実行

アプリ設定

/opt/nvidia/deepstream/deepstream-6.0/sources/yolov8_cls_parserに移動する。
deepstream_app_config.txtのuriに読み込ませたい動画のパスを設定する。
今回は下記からflower.mp4をダウンロードして使用する。
https://drive.google.com/drive/folders/1DN1pciFgE-U_ZmoCjcvAQV-1XUB4aAau?usp=share_link

flower.mp4ファイルをyolov8_cls_parserフォルダにコピーする。

[source0]
enable=1
type=3
#uri=file:///opt/nvidia/deepstream/deepstream/samples/streams/sample_1080p_h264.mp4
///opt/nvidia/deepstream/deepstream-6.0/sources/yolov8_cls_parser/flower.mp4
num-sources=1
gpu-id=0
cudadec-memtype=0

以下のコマンドを実行し、アプリを実行する。

deepstream-app -c deepstream_app_config.txt

実行結果

推論画面の左上のテキストボックスに推論されたクラス名が表示されています。正しく分類できている様子が見られました！FPSは22程度でした。

おわりに

この記事では、JetsonにYOLOv8のclassificationモデルを組込んでリアルタイムに分類を行う方法を紹介しました。新しいモデルを組込むときはonnx化してからTensorRT化することをお勧めします。Jetsonでいろいろインストールするのが大変なので。。。

参考ページ

• https://docs.ultralytics.com/tasks/classify/#predict