UnityLaparoscopicSceneSimulator/Assets/Scripts/DepthRenderPassFeature.cs

using System;
using System.Collections.Generic;
using System.Threading.Tasks;
using UnityEditor;
using UnityEngine;
using UnityEngine.Rendering;
using UnityEngine.Rendering.Universal;
using UnityEngine.XR;

public class DepthRenderPassFeature : ScriptableRendererFeature
{
    private static Dictionary<int, Queue<TaskCompletionSource<float[]>>> m_PendingRequests = new();

    internal class DepthRenderPass : ScriptableRenderPass
    {
        public static Task<float[]> RequestDepthDataAsync(int id)
        {
            if(m_PendingRequests == null)
                m_PendingRequests = new Dictionary<int, Queue<TaskCompletionSource<float[]>>>();

            var tcs = new TaskCompletionSource<float[]>();
            int key = id;
            if (!m_PendingRequests.TryGetValue(key, out var queue))
            {
                queue = new Queue<TaskCompletionSource<float[]>>();
                m_PendingRequests.Add(key, queue);
            }
            queue.Enqueue(tcs);
            return tcs.Task;
        }

        ProfilingSampler m_ProfilingSampler = new ProfilingSampler("ColorBlit");
        private Material material;
        private RTHandle rtHandle;
        private RTHandle copyRtHandle;
        private bool depth = false;
        private ComputeBuffer computeBuffer;

        public DepthRenderPass(Material material)
        {
            this.material = material;
        }

        public void SetTarget(RTHandle handle, bool depth)
        {
            rtHandle = handle;
            this.depth = depth;

            if (computeBuffer != null) DestroyComputeBuffer();
                
            RenderTexture rt = rtHandle.rt;
            computeBuffer = new ComputeBuffer(rt.width * rt.height, 1 * sizeof(float));
            if(material != null)
            {
                material.SetBuffer("distBuffer", computeBuffer);
                material.SetFloat("res", rt.height);
                material.SetInt("_RenderDepth", this.depth ? 1 : 0);
            }
        }

        public override void OnCameraSetup(CommandBuffer cmd, ref RenderingData renderingData)
        {
            ConfigureTarget(rtHandle);
            var desc = renderingData.cameraData.cameraTargetDescriptor;
            RenderingUtils.ReAllocateIfNeeded(ref copyRtHandle, desc, FilterMode.Bilinear, TextureWrapMode.Clamp);
        }

        // Here you can implement the rendering logic.
        // Use <c>ScriptableRenderContext</c> to issue drawing commands or execute command buffers
        // https://docs.unity3d.com/ScriptReference/Rendering.ScriptableRenderContext.html
        // You don't have to call ScriptableRenderContext.submit, the render pipeline will call it at specific points in the pipeline.
        public override void Execute(ScriptableRenderContext context, ref RenderingData renderingData)
        {
            CommandBuffer cmd = CommandBufferPool.Get();

            using (new ProfilingScope(cmd, m_ProfilingSampler))
            {
                Blitter.BlitCameraTexture(cmd, rtHandle, rtHandle, material, 0);
            }

            context.ExecuteCommandBuffer(cmd);
            CommandBufferPool.Release(cmd);

            lock(m_PendingRequests)
                {
                    if (m_PendingRequests.Count > 0)
                    {
                        int key = renderingData.cameraData.camera.GetInstanceID();
                        if (m_PendingRequests.TryGetValue(key, out var queue) && queue.Count > 0)
                        {
                            float[] result = new float[computeBuffer.count];
                            material.SetBuffer("distBuffer", computeBuffer);
                            Graphics.SetRandomWriteTarget(1, computeBuffer,false);
                            computeBuffer.GetData(result);
                            while (queue.TryDequeue(out var tcs))
                                tcs.TrySetResult(result);
                        }
                    }
                }
        }

        // Cleanup any allocated resources that were created during the execution of this render pass.
        public override void OnCameraCleanup(CommandBuffer cmd)
        {
        }

        public void Dispose()
        {
#if UNITY_EDITOR
            if (EditorApplication.isPlaying)
            {
                Destroy(material);
            }
            else
            {
                DestroyImmediate(material);
            }
#else
            Destroy(material);
#endif
            if (rtHandle != null) rtHandle.Release();

            if (computeBuffer != null)
            {
                DestroyComputeBuffer();
            }
        }

        public void DestroyComputeBuffer()
        {
            computeBuffer.Release();
            computeBuffer = null;
        }
    }

    [SerializeField] Shader shader;
    [SerializeField] bool Depth = false;
    DepthRenderPass m_ScriptablePass;
    Material material;

    /// <inheritdoc/>
    public override void Create()
    {
        if(shader == null) shader = TexturesSettings.Instance.DepthShader;
        material = CoreUtils.CreateEngineMaterial(shader);
        m_ScriptablePass = new DepthRenderPass(material);

        // Configures where the render pass should be injected.
        m_ScriptablePass.renderPassEvent = RenderPassEvent.AfterRenderingPostProcessing;
    }

    // Here you can inject one or multiple render passes in the renderer.
    // This method is called when setting up the renderer once per-camera.
    public override void AddRenderPasses(ScriptableRenderer renderer, ref RenderingData renderingData)
    {
        if (renderingData.cameraData.camera.tag == "Eye")
        {
            renderer.EnqueuePass(m_ScriptablePass);
        }
    }

    public override void SetupRenderPasses(ScriptableRenderer renderer, in RenderingData renderingData)
    {
        m_ScriptablePass.ConfigureInput(ScriptableRenderPassInput.Color);
        m_ScriptablePass.SetTarget(renderer.cameraColorTargetHandle, Depth);
    }

    protected override void Dispose(bool disposing)
    {
        base.Dispose(disposing);
#if UNITY_EDITOR
        if (EditorApplication.isPlaying)
        {
            Destroy(material);
        }
        else
        {
            DestroyImmediate(material);
        }
#else
            Destroy(material);
#endif
        m_ScriptablePass.Dispose(); 

        if (m_PendingRequests != null)
        {
            foreach (var queue in m_PendingRequests.Values)
            {
                while (queue.TryDequeue(out var tcs))
                    tcs.TrySetCanceled();
            }
            m_PendingRequests.Clear();
        }
    }

}