概要

HoloLab DNN PackagesというUnity Sentisをベースにしたディープラーニングによる画像認識を実装したパッケージを公開しました。

blog.hololab.co.jp

この記事ではHoloLab DNN Packagesのデフォルトではサポートされていないモデルの1つとしてMobileSAMによるセグメンテーションの推論クラスを実装したUnityアプリのサンプルを紹介します。

学習済みモデルの準備

MobileSAMはMetaの公開したSegment Anythin Model（SAM）を高速化したモデルです。 SAMはエンコーダーとデコーダーで構成されており、MobileSAMはエンコーダーをViT-HからTiny-ViTに置き換えることで大幅に高速化しています。

github.com

MobileSAMをONNXにエクスポートするには、ONNX Runtimeで実装された以下のリポジトリを参考にエンコーダーとデコーダーをそれぞれ分割してエクスポートしました。
公式のエクスポートスクリプトでスクリプトすることができますが、エンコーダーとデコーダーが一緒になったエンドツーエンドでエクスポートされます。

github.com

サンプルプログラム

サンプルのUnityアプリは以下のリポジトリで公開しています。今回は簡易的な実装のためアノテーションはクリック位置の1つのみでセグメンテーションを行うアプリになっています。

github.com

MobileSAMによるセグメンテーションの推論を実装

サンプルアプリに含まれるMobileSAMの推論クラスSegmentationModel_MobileSAMの使い方を簡単に紹介します。推論クラスはSegmentationModel_MobileSAM.csに実装されているので、実装の中身が気になる方はそちらを参照してください。

名前空間を参照する

Unity SentisとHoloLab DNN Packagesの名前空間を参照します。
ここではセグメンテーションを利用するのでHoloLab.DNN.Segmentationを参照しています。
```
 using Unity.Sentis;
 using HoloLab.DNN.Segmentation;
```
クラスを生成する

MobileSAMによるセグメンテーションのSegmentationModel_MobileSAMクラスを生成します。
コンストラクタにエンコーダーとデコーダーのモデルを設定します。
```
 [SerializeField] private ModelAsset encoder_asset;
 [SerializeField] private ModelAsset decoder_asset;
 var model = new SegmentationModel_MobileSAM(encoder_asset, decoder_asset);
```
画像を準備する

Texture2Dクラスで入力画像を準備します。
WebCamTextureやTexture2D.LoadImage()などお好みの方法で取得します。
```
 // e.g. Load Image
 var data = File.ReadAllBytes("./image.jpg");
 var texure = new Texture2D(1, 1);
 texture.LoadImage(data);
```
アノテーションのポイントを準備する

MobileSAMに与えるアノテーションのポイントを準備します。
サンプルアプリでは一例としてクリックした位置をVector2クラスで取得しています。

クリックイベントを取得
セグメンテーションする

モデルと画像、ポイントが用意できたらセグメンテーションします。
サンプルアプリでは1点のみ利用していますが、複数の点を与えることもできます。そのさいは、点に対応するラベルも一緒に与えてください。
ラベルはポイント（領域外）が0、ポイント（領域内）が1、バウンディングボックス（左上、右下）が2と3になります。
セグメンテーションの結果はColor.rに領域内が1、領域外が0のインデックスが振られた入力画像と同じ大きさのマスク画像です。
```
 // segmentation with one annotation point
 var indices_texture = model.Segment(input_texture, click_point);

 // segmentation with multiple annotation points
 var points = new List<Vector2>() { new Vector2(x, y), ... };
 var labels = new List<float>() { 1.0f, ... };
 var indices_texture = model.Segment(input_texture, points, labels);
```
セグメンテーション結果を可視化する

セグメンテーション結果のマスク画像は目的に合わせて利用できます。
サンプルアプリではではHoloLab DNN Packagesのセグメンテーションパッケージに含まれる可視化クラスを利用してマスク画像に色を付けて可視化しています。
ここではインデックス1（領域内）を赤、インデックス0（領域外）を透明として着色しています。
```
 var colors = new List<Color>() { Color.clear, new Color(1.0f, 0.0f, 0.0f, 0.5f) };
 var colorized_texture = Visualizer.ColorizeArea(indices_texture, colors);
```