はじめに

movie_001

単眼深度推定モデルのMiDaSv2をUnityで実行します。
Unity BarracudaとCompute Shaderを用いて前処理、推論、後処理まで全てGPU上で処理します。
※Ryzen7 2700 & RTX3060tiでエディタ上で80FPSくらい出ています

実装内容

リポジトリ

実装を公開しています。

推論処理の流れ

初期化フェーズ（Visualizer.Start）

RunInferenceクラスのインスタンスが生成され、必要なリソースや設定を初期化します。

実行フェーズ（Visualizer.Update）

フレームごとにテクスチャを更新します。
RunInference.ProcessImageが呼び出され、以下の処理を行います
- 前処理: 画像データを推論に適した形に変換。
- 推論: 相対的な逆深度を推論。
- 後処理: 推論結果を最終的な出力形式に変換。

終了フェーズ（Visualizer.OnDisable）

アプリケーションが終了または非アクティブになると、RunInference.Disposeが呼び出され、使用したリソースを解放します。

推論クラス（推論実行におけるエントリポイント）

ProcessImageメソッドで前処理、推論、後処理が実行されます

using Unity.Barracuda;
using UnityEngine;
using System;

namespace MiDaSv2
{
    public class RunInference : IDisposable
    {
        private IWorker worker;
        private ComputeBuffer inputBuffer;
        private ComputeBuffer outputBuffer;
        private ResourceSet resources;
        public RenderTexture DepthTexture { get; private set; }
        ThreadSize _threadSize;

        private const int inputWidth = 256;
        private const int inputHeight = 256;
        private const int inputChannels = 3;

        public RunInference(ResourceSet resourceSet)
        {
            Initialize(resourceSet);
        }

        private void Initialize(ResourceSet resourceSet)
        {
            resources = resourceSet;

            var model = ModelLoader.Load(resources.model);
            worker = WorkerFactory.CreateWorker(WorkerFactory.Type.ComputePrecompiled, model);

            inputBuffer = new ComputeBuffer(1 * inputWidth * inputHeight * inputChannels, sizeof(float));
            outputBuffer = new ComputeBuffer(1 * inputWidth * inputHeight, sizeof(float));

            DepthTexture = Utils.CreateRenderTexture(inputWidth, inputHeight);
        }

        public void ProcessImage(RenderTexture sourceTexture)
        {
            // Preprocessing
            var preprocess = resources.preprocess;
            preprocess.GetKernelThreadGroupSizes(0, out _threadSize.x, out _threadSize.y, out _threadSize.z);
            preprocess.SetTexture(0, "InputTexture", sourceTexture);
            preprocess.SetBuffer(0, "OutputTensor", inputBuffer);
            preprocess.SetInt("Size", inputWidth);
            preprocess.Dispatch(0, inputWidth / (int)_threadSize.x, inputHeight / (int)_threadSize.y, (int)_threadSize.z);

            // Inference
            using (var tensor = new Tensor(1, inputWidth, inputHeight, inputChannels, inputBuffer))
            {
                worker.Execute(tensor);
            }

            var output = worker.PeekOutput();
            outputBuffer.SetData(output.AsFloats());

            // Postprocessing
            var postprocess = resources.postprocess;
            postprocess.GetKernelThreadGroupSizes(0, out _threadSize.x, out _threadSize.y, out _threadSize.z);
            postprocess.SetBuffer(0, "InputTensor", outputBuffer);
            postprocess.SetTexture(0, "OutputTexture", DepthTexture);
            postprocess.Dispatch(0, inputWidth / (int)_threadSize.x, inputHeight / (int)_threadSize.y, (int)_threadSize.z);
        }

        public void Dispose()
        {
            worker.Dispose();
            inputBuffer.Dispose();
            outputBuffer.Dispose();

            if (DepthTexture != null)
            {
                DepthTexture.Release();
                DepthTexture = null;
            }
        }
    }
}

前処理のシェーダ

// Preprocess
#pragma kernel Preprocess_Texture2D

Texture2D<float4> InputTexture;
RWStructuredBuffer<float> OutputTensor;
uint Size;

[numthreads(16, 16, 1)]
void Preprocess_Texture2D(uint3 id : SV_DispatchThreadID)
{
    uint textureWidth, textureHeight;
    InputTexture.GetDimensions(textureWidth, textureHeight);
    
    int tx = id.x * textureWidth / Size;
    int ty = id.y * textureHeight / Size;
    float4 color = InputTexture[float2(ty, tx)];

    // インデックス: (バッチ番号 * 高さ * 幅 * チャネル数) + (y座標 * 幅 * チャネル数) + (x座標 * チャネル数) + チャネル
    uint tensorIndex = id.y * Size * 3 + id.x * 3;
    OutputTensor[tensorIndex] = color.r;      // チャネル = 0
    OutputTensor[tensorIndex + 1] = color.g;  // チャネル = 1
    OutputTensor[tensorIndex + 2] = color.b;  // チャネル = 2
}

後処理のシェーダ

// Postprocess
#pragma kernel Postprocess

StructuredBuffer<float> InputTensor;
RWTexture2D<float4> OutputTexture;

[numthreads(16, 16, 1)]
void Postprocess (uint3 id : SV_DispatchThreadID)
{
    float value = InputTensor[id.x + id.y * 256] / 1000.0f;
    OutputTexture[id.xy] = float4(value, value, value, 1.0f);
}

Discussion