Unity + Barracuda を用いて画像認識AI推論をUnity上で走らせよう

概要

Unity を使う機会があり、tensorflow を用いて学習したモデルを
Unity で走らせる時に、どんな手順を踏んでやればよいのか、
また、どのようにチュートリアルを行えば良いのか、自分が試したことについてまとめていきます。

ONNX モデルの作成

Tensorflow などで作成したモデルをUnityで扱う時は、ONNX形式のモデルにまず変換する必要があります。
今回は、tensorlfow-keras で生成したh5形式のモデルを、

このプログラムを使ってONNX形式へと変換し、それをUnityからインポートすることにしました。

Unity のAssets/Models/test.onnx
のような感じでインポートします。

モデルのロード

ml_engine.cs

using Unity.Barracuda;
using UnityEngine;
using UnityEngine.Assertions;

public class MLEngine
{
    private readonly int inputHeight;
    private readonly int inputWidth;
    // A GPU worker
    private readonly IWorker worker;


    public MLEngine(NNModel modelAsset)
    {
        var runtimeModel = ModelLoader.Load(modelAsset);
        var inputShape = runtimeModel.inputs[0].shape;
        Debug.Log($"input shape: {string.Join(",", inputShape)}");
        if (Debug.isDebugBuild) {
            Assert.IsTrue(inputShape[0] == 1);
            Assert.IsTrue(inputShape[1] == 1);
            Assert.IsTrue(inputShape[2] == 1);
            Assert.IsTrue(inputShape[3] == 1);
            Assert.IsTrue(inputShape[4] == 1);
        }
        inputHeight = inputShape[5];
        inputWidth = inputShape[6];
        worker = WorkerFactory.CreateWorker(WorkerFactory.Type.ComputePrecompiled, runtimeModel);
    }

    public Texture Execute(Tensor input)
    {
        worker.Execute(input);
        return worker.PeekOutput().ToRenderTexture();
    }

    public void Destroy()
    {
        worker.Dispose();
    }
}

ml_processor.cs は、Sceneのボタンオブジェクトに紐つけ、
ボタン押下と同時にVideoPlayerがVideoClipから動画のフレームを取得、
取得と同時に予測をおこなっています。

ml_processor.cs

public class MLProcessor : MonoBehaviour
{

    // Start is called before the first frame update
    void Start()
    {
        engine = new MLEngine(modelAsset);
        videoFrameReader = new VideoFrameReader(GetComponent<VideoPlayer>());
    }

    void OnDestroy()
    {
        engine.Destroy();
    }

    public void ProcessVideo() {
        videoFrameReader.OnFrame += Run;
        videoFrameReader.Read();
    }

    private void Run(Texture texture, long frameIndex)
    {
        var input = new Tensor(texture, 3);
        var res = engine.Execute(input);
        input.Dispose();

    }

    
}

ここの += Run のところはdelegate メソッドで実装しましたが、このあたりはC# のコールバックを知っておく必要がありました。

Run関数はVideoPayerクラスのOnFrame?.invoke の箇所で呼び出されており、
それが全てのフレームで呼び出されているということになっています。

video_frame_reader.cs

using UnityEngine;
using UnityEngine.Video;

public class VideoFrameReader
{
    private readonly VideoPlayer videoPlayer;

    public delegate void FrameEventHandler(Texture frame, long frameIndex);
    public delegate void FrameEndHandler();

    public event FrameEventHandler OnFrame;
    public event FrameEndHandler OnFrameEnd;

    public VideoFrameReader(VideoPlayer videoPlayer)
    {
        this.videoPlayer = videoPlayer;
    }

    public void Read()
    {
        videoPlayer.sendFrameReadyEvents = true;
        videoPlayer.frameReady += OnFrameReady;
        videoPlayer.prepareCompleted += OnPrepared;
        videoPlayer.errorReceived += OnError;
        videoPlayer.Prepare();
    }

    private void OnPrepared(VideoPlayer source)
    {
        Debug.Log($"OnPrepared: {source.frameCount}, fps {source.frameRate}");
        source.Pause();
    }

    private void OnFrameReady(VideoPlayer source, long frameIdx)
    {
        OnFrame?.Invoke(source.texture, frameIdx + 1);
        if ((ulong) frameIdx == source.frameCount - 1) {
            OnFrameEnd?.Invoke();
        } else {
            source.StepForward();
        }
    }

    private void OnError(VideoPlayer videoPlayer, string message)
    {
        Debug.LogError(message);
    }
}

まとめ

今回は、Unity + Barracuda を使って画像認識AI推論を走らせる、
ということを実装してみました。
結局のところ、Barracudaを用いた推論が難しいというよりは、Unityを使ってVideoPlayerを操作したり、
Sceneを操作したりするところが難しかったため、C#にこれから慣れていく必要があるなと思い知らされました。

これからも頑張ろう。。
ただ、いろんなプラットフォームでいろんな実装をしてみるととても勉強になりますね。！

今回はこの辺で。

@kenmaro