Unity AIのドキュメントを読むその６１（出力データの使用）

本日はUnityの技術調査枠です。
Unity AIのドキュメントを読みながら実際に操作を試して記事に残します。

Unity AI

以下のUnity AIのドキュメントを試しながら実行時のキャプチャをしていきます。
docs.unity3d.com

出力データの使用

モデルからの出力をテンソルとして取得したらそのデータを後処理してプロジェクトで使用できます。

CPU テンソルへのダウンロード

ReadbackAndCloneまたはReadbackAndCloneAsyncを使用してグラフィックスプロセッシングユニット (GPU) 上のテンソルを中央処理装置 (CPU) に移動して読み取ります。

var outputTensor = worker.Schedule(inputTensor).PeekOutput() as Tensor<float>;
var cpuTensor = outputTensor.ReadbackAndClone();

返されるテンソルは出力テンソルのCPUベースの読み書き可能なコピーです。

レンダーテクスチャへの変換

テンソルをレンダーテクスチャに変換するにはTextureConverter.RenderToTextureを使用してテンソルデータを既存のレンダーテクスチャに書き込みます。
TextureConverter.RenderToTextureを使用する際、テンソルの次元がレンダーテクスチャの次元と一致しない場合、推論エンジンは以下の調整を行います。

次元が一致しない場合はテンソルを線形サンプリングします。
レンダーテクスチャのチャンネル数がテンソルのチャンネル数より少ない場合は末尾からチャンネルを削除します。
レンダーテクスチャのチャンネル数がテンソルのチャンネル数より多い場合はRGBチャンネルの値を0に、アルファチャンネルの値を1に設定します。

// レンダリングテクスチャをインスタンス化する
public RenderTexture rt = new RenderTexture(24, 32, 0, RenderTextureFormat.ARGB32);
Worker worker;
Tensor inputTensor;

void Start()
{
    // モデルの出力をテンソルとして取得する
    Tensor<float> outputTensor = worker.PeekOutput() as Tensor<float>;

    // テンソルをテクスチャに変換し、レンダリングテクスチャに保存します
    TextureConverter.RenderToTexture(outputTensor, rt);
}

画面へのコピー

出力テンソルを画面にコピーするには次の手順に従います。

Camera.mainのCamera.targetTextureプロパティをnullに設定します。
スクリプトを作成し、カメラにアタッチします。
スクリプト内でOnRenderImageなどのイベント関数でTextureConverter.RenderToScreenを使用します。

画像が明るすぎる場合、出力テンソルが0～1ではなく0～255の値を使用している可能性があります。
RenderToScreenを呼び出す前に「モデルの編集」を行って出力テンソルの値を再マッピングできます。

次のサンプルスクリプトはモデルを使用してテクスチャを変更し、その結果を画面にコピーします。
modelAssetをONNXのスタイル転送モデルの1つに設定し、inputImageをテクスチャに設定します。
テクスチャのインポート設定を確認し、テクスチャがモデルに必要な形状とレイアウトと一致していることを確認してください。

using UnityEngine;
using Unity.InferenceEngine;

public class StyleTransfer : MonoBehaviour
{
    public ModelAsset modelAsset;
    public Model runtimeModel;
    public Texture2D inputImage;
    public RenderTexture outputTexture;

    Worker worker;
    Tensor<float> inputTensor;

    void Start()
    {
        var sourceModel = ModelLoader.Load(modelAsset);
        var graph = new FunctionalGraph();
        var input = graph.AddInput(sourceModel, 0);
        var output = Functional.Forward(sourceModel, input)[0];
        // ソースモデルの出力を再スケールする
        output /= 255f;
        graph.AddOutput(output);
        var runtimeModel = graph.Compile();

        worker = new Worker(runtimeModel, BackendType.GPUCompute);
        inputTensor = new Tensor<float>(new TensorShape(1, 3, 256, 256));
    }

    void OnRenderImage(RenderTexture source, RenderTexture destination)
    {
        // テクスチャから入力テンソルを作成する
        TextureConverter.ToTensor(inputImage, inputTensor, new TextureTransform());

        // モデルを実行し、出力をテンソルとして取得する
        worker.Schedule(inputTensor);
        Tensor<float> outputTensor = worker.PeekOutput() as Tensor<float>;

        // 再スケールされたテンソルをテクスチャとして画面にコピーする
        TextureConverter.RenderToScreen(outputTensor);
    }

    void OnDisable()
    {
        worker.Dispose();
    }
}

ユニバーサルレンダーパイプライン（URP）または高解像度レンダーパイプライン（HDRP）を使用する場合はRenderPipelineManager.endFrameRenderingまたはRenderPipelineManager.endContextRenderingコールバックでRenderToScreenを呼び出します。