概要
HoloLab DNN PackagesというUnity Sentisをベースにしたディープラーニングによる画像認識を実装したパッケージを公開しました。
この記事ではHoloLab DNN Packagesのデフォルトではサポートされていないモデルの1つとしてMobileSAMによるセグメンテーションの推論クラスを実装したUnityアプリのサンプルを紹介します。
学習済みモデルの準備
MobileSAMはMetaの公開したSegment Anythin Model(SAM)を高速化したモデルです。 SAMはエンコーダーとデコーダーで構成されており、MobileSAMはエンコーダーをViT-HからTiny-ViTに置き換えることで大幅に高速化しています。
MobileSAMをONNXにエクスポートするには、ONNX Runtimeで実装された以下のリポジトリを参考にエンコーダーとデコーダーをそれぞれ分割してエクスポートしました。
公式のエクスポートスクリプトでスクリプトすることができますが、エンコーダーとデコーダーが一緒になったエンドツーエンドでエクスポートされます。
サンプルプログラム
サンプルのUnityアプリは以下のリポジトリで公開しています。 今回は簡易的な実装のためアノテーションはクリック位置の1つのみでセグメンテーションを行うアプリになっています。
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);
まとめ
Unityで簡単にディープラーニングによる画像認識を実装できるHoloLab DNN Packagesを利用してMobileSAMによるセグメンテーションの実装方法を紹介しました。Unityアプリにディープラーニングによる画像認識を使った機能を組み込んで遊んでみてください。