概要

HoloLab DNN PackagesというUnity Sentisをベースにしたディープラーニングによる画像認識を実装したパッケージを公開しました。
この記事ではHoloLab DNN Packagesのデフォルトではサポートされていないモデルの推論クラスを実装してUnityアプリに組み込む方法を紹介します。

blog.hololab.co.jp

学習済みモデルの準備

YOLOXのモデルでは出力テンソルから検出されたオブジェクトのクラス番号、信頼度、位置と大きさを解釈する後処理を実装する必要があります。 HoloLab DNN Packagesのオブジェクト検出パッケージではこの後処理が実装済みの推論クラスが提供されています。

sugiura-lab.hatenablog.com

このように通常はユーザーコードで後処理を実装しますが、モデルそのものに後処理を組み込むこともできます。 PINTO Model Zooでは後処理が組み込まれ、オブジェクトのクラス番号、信頼度、位置と大きさが直接出力されるモデルがいくつか提供されています。これにより、ユーザーコードで後処理を実装することなく簡単に結果を取り出すことができます。ただし、HoloLab DNN Packagesのオブジェクト検出パッケージではこの形式の出力をデフォルトではサポートしていません。

このようなHoloLab DNN Packagesがサポートしていないモデルはどうしたらいいのでしょうか？ここでは後処理が組み込まれているYOLOXを題材として、HoloLab DNN Packagesのデフォルトではサポートされていないモデルの推論クラスを新たに実装する方法を紹介します。題材となるモデルは以下のページで公開されいてます。

426_YOLOX-Body-Head-Hand | PINTO Model Zoo

パッケージのインポート

Unityのメニューから[Window]>[Package Manager]を開きます。 Package Managerの左上にある[+]>[Add package from git URL]から以下のURLを入力して[Add]ボタンを押します。ここではHoloLab DNN Packagesのうち「基本パッケージ」と「オブジェクト検出パッケージ」の2つをインポートしていきます。
（推論クラスを新たに実装するだけなら「基本パッケージ」のみ導入するだけでいいのですが、推論結果の可視化機能を利用するため「オブジェクト検出パッケージ」も導入しています。）

基本パッケージ

まずは基本パッケージをインポートします。
```
 https://github.com/HoloLabInc/HoloLabDnnPackages.git?path=packages/jp.co.hololab.dnn.base
```
基本パッケージのインポート
オブジェクト検出パッケージ

次にオブジェクト検出パッケージをインポートします。
```
 https://github.com/HoloLabInc/HoloLabDnnPackages.git?path=packages/jp.co.hololab.dnn.objectdetection
```
オブジェクト検出パッケージのインポート

これで準備は完了です。

推論クラスを実装

HoloLab DNN Packagesのデフォルトではサポートされていないモデルの推論クラスを新たに実装していきます。ここで実装している推論クラスの全体はこちらを参照してください。

名前空間を参照する

Unity SentisとHoloLab DNN Packagesの名前空間を参照します。ここでは基本モデルを利用するのでHoloLab.DNN.Baseを参照しています。
```
 using Unity.Sentis;
 using HoloLab.DNN.Base;
```

基本モデルを継承した推論クラスを定義する

基本的な推論機能が実装されているBaseModelを継承して推論クラスを実装します。
ここではObjectDetectionModel_YOLOXWithPostProcessクラスを作成しています。

コンストラクタはファイルパスを受け取るタイプとアセットを受け取るタイプを定義します。モデルと実行タイプはBaseModelのコンストラクタにそのまま渡します。

もしメモリの解放など終了処理を記述する必要がある場合は、Dispose()のメソッドを定義します。ここでは特に終了処理は無いため、基本クラスのBaseModel.Dispose()を呼び出すだけです。（終了処理が無い場合は書かなくてもいいです。）

 public class ObjectDetectionModel_YOLOXWithPostProcess : BaseModel, IDisposable
 {
     private int input_width = 0;
     private int input_height = 0;

     public ObjectDetectionModel_YOLOXWithPostProcess(string file_path, BackendType backend_type = BackendType.GPUCompute)
         : base(file_path, backend_type)
     {
         Initialize();
     }

     public ObjectDetectionModel_YOLOXWithPostProcess(ModelAsset model_asset, BackendType backend_type = BackendType.GPUCompute)
         : base(model_asset, backend_type)
     {
         Initialize();
     }

     public new void Dispose()
     {
         base.Dispose();
     }

     /* ... */
 }

クラスにメソッドを定義していきます。

まずは基本モデルを設定するなどの初期化をするメソッドの定義です。

初期化処理として基本モデルが担当している入力テンソルの正規化のパラメータを設定しています。 YOLOXでは0.0~1.0ではなく0~255の値を持つ画像を入力するため、基本モデルのBaseModel.SetInputMax()を呼び出してスケールを設定しています。

 private void Initialize()
 {
     SetInputMax(255.0f);

     var input_shape = GetInputShapes().First().Value;
     input_width = input_shape[3];
     input_height = input_shape[2];
 }

次に画像を受け取り検出結果を返すメソッドを定義です。

入力された画像を基本モデルのBaseModel.Predict()に渡して出力テンソルを取得します。 Unity SentisではTensor.MakeReadable()を呼び出してテンソルの値を取得する準備をします。準備が出来た後、TensorFloat.ToReadOnlySpan()を呼びだしてReadOnlySpan<float>型で値を取得しています。

あとは値を取り出して解釈を行います。とはいえ、前述のとおりモデルそのものに後処理が組み込まれているため、ユーザーコードでやることは値を取り出すだけです。
YOLOX-Body-Head-HandのモデルではクラスID、信頼度、位置と大きさ（Xmin、Ymin、Xmax、Ymax）の順番で値が格納されています。これを取り出し、Objectクラスを生成しています。

 public List<HoloLab.DNN.ObjectDetection.Object> Detect(Texture2D image, float score_threshold = 0.8f)
 {
     var output_tensors = Predict(image);

     var output_name = runtime_model.outputs[0];
     var output_tensor = output_tensors[output_name] as TensorFloat;

     output_tensor.MakeReadable();
     var output_span = output_tensor.ToReadOnlySpan();

     var objects = new List<HoloLab.DNN.ObjectDetection.Object>();
     for (var i = 0; i < output_tensor.shape[0]; i++)
     {
         // class_id, score, x_min, y_min, x_max, y_max
         var span = output_span.Slice(i * output_tensor.shape[1], output_tensor.shape[1]);

         var class_id = (int)span[1];

         var score = span[2];
         if (score < score_threshold)
         {
             continue;
         }

         var x_min = (int)(Math.Max(0, span[3]) * image.width / input_width);
         var y_min = (int)(Math.Max(0, span[4]) * image.height / input_height);
         var x_max = (int)(Math.Min(span[5], input_width) * image.width / input_width);
         var y_max = (int)(Math.Min(span[6], input_height) * image.height / input_height);
         var rect = new Rect(x_min, y_min, x_max - x_min, y_max - y_min);

         objects.Add(new Object(rect, class_id, score));
     }

     output_tensors.AllDispose();

     return objects;
 }

サンプルの実行

ここまででHoloLab DNN Packagesのデフォルトではサポートされていないモデルの推論クラスを新たに実装できました。あとは推論クラスのインスタンスを生成して推論すれば利用できます。

サンプルプロジェクトはこちらのリポジトリで公開しています。

github.com

まとめ

HoloLab DNN Packagesのデフォルトではサポートされていないモデルの推論クラスを新たに実装する方法を紹介しました。
どのようなモデルでもこのように実装できるというわけではありませんが、HoloLab DNN Packagesはある程度の柔軟性を備えた作りになっているつもりです。（もし、BaseModelなどHoloLab DNN Packagesそのものに変更を加えたいという場合はフォークする、またはプルリクエストを送ってください。歓迎します。）

HoloLab DNN Packagesを使ってUnityでディープラーニングによる画像認識を遊んでみてください！