Fixing various bugs in keypoint self occlusion, including (hopefully) properly supporting player builds

com.unity.perception/Runtime/GroundTruth/Labelers/KeypointLabeler.cs

                 new ShaderVariantCollection.ShaderVariant(depthCheckShader, PassType.ScriptableRenderPipeline, keyword));
             shaderVariantCollection.WarmUp();
 
-            //TODO: Proper resizing
-            m_ResultsBuffer = new RenderTexture(1024, 1, 0, GraphicsFormat.R8G8B8A8_UNorm);
-            m_DepthCheckReader = new RenderTextureReader<Color32>(m_ResultsBuffer);
+            SetupResultsBuffer(1024);
 
             perceptionCamera.attachedCamera.depthTextureMode = DepthTextureMode.Depth;
 #if URP_PRESENT
 
             perceptionCamera.InstanceSegmentationImageReadback += OnInstanceSegmentationImageReadback;
             perceptionCamera.RenderedObjectInfosCalculated += OnRenderedObjectInfoReadback;
+        }
+
+        private void SetupResultsBuffer(int size)
+        {
+            if (m_ResultsBuffer != null && m_ResultsBuffer.width >= size)
+                return;
+
+            if (m_ResultsBuffer != null)
+            {
+                m_ResultsBuffer.Release();
+                m_DepthCheckReader.Dispose(false);
+            }
+
+            m_ResultsBuffer = new RenderTexture(size, 1, 0, GraphicsFormat.R8G8B8A8_UNorm);
+            m_DepthCheckReader = new RenderTextureReader<Color32>(m_ResultsBuffer);
         }
 
         bool AreEqual(Color32 lhs, Color32 rhs)
             m_MaterialDepthCheck.SetTexture("_KeypointCheckDepth", keypointCheckDepthTexture);
             m_MaterialDepthCheck.SetTexture("_CameraDepthTexture", depthTexture);
 
+            SetupResultsBuffer(checkLocations.Length);
-
-
-            // var resultsComputeBuffer =
-            //     new ComputeBuffer(keypointCount, 4, ComputeBufferType.Default, ComputeBufferMode.Dynamic);
-            // var depthTexture = Shader.GetGlobalTexture("_CameraDepthTexture");
-            // var positionSize = UnsafeUtility.SizeOf<float3>();
-            // commandBuffer.SetComputeTextureParam(m_KeypointDepthTestShader, 0, "DepthBuffer", depthTexture);
-            //
-            // var positionSize = UnsafeUtility.SizeOf<float3>();
-            // var keypointPositionsBuffer = new ComputeBuffer(keypointCount * positionSize, positionSize,
-            //     ComputeBufferType.Default, ComputeBufferMode.Dynamic);
-            // keypointPositionsBuffer.SetData(positions.AsArray());
-            //
-            // var resultsComputeBuffer =
-            //     new ComputeBuffer(keypointCount, 4, ComputeBufferType.Default, ComputeBufferMode.Dynamic);
-            //
-            // commandBuffer.SetComputeBufferParam(m_KeypointDepthTestShader, 0, "CheckPositions", keypointPositionsBuffer);
-            // commandBuffer.SetComputeBufferParam(m_KeypointDepthTestShader, 0, "CheckResults", resultsComputeBuffer);
-            // commandBuffer.DispatchCompute(m_KeypointDepthTestShader, 0, keypointCount, 1, 1);
-            //
-            // commandBuffer.RequestAsyncReadback(resultsComputeBuffer, request => OnDepthCheckReadback(currentFrame, request));
 
             scriptableRenderContext.ExecuteCommandBuffer(commandBuffer);
             scriptableRenderContext.Submit();

com.unity.perception/Runtime/GroundTruth/RenderTextureReader.cs

             {
                 RenderTexture.active = m_Source;
 
-                if (m_CpuTexture == null)
+                if (m_CpuTexture == null || m_CpuTexture.width != m_Source.width || m_CpuTexture.height != m_Source.height)
+                {
+                    if (m_CpuTexture != null)
+                        m_CpuTexture.Resize(m_Source.width, m_Source.height);
+
+                }
 
                 m_CpuTexture.ReadPixels(new Rect(
                     Vector2.zero,
         /// <summary>
         /// Shut down the reader, waiting for all requests to return.
         /// </summary>
-        public void Dispose()
+        public void Dispose() => Dispose(true);
+
+        /// <summary>
+        /// Shut down the reader, optionally waiting for all requests to return.
+        /// </summary>
+        /// <param name="waitForAllImages">Whether this should block on waiting for asynchronous readbacks</param>
+        public void Dispose(bool waitForAllImages)
         {
             WaitForAllImages();
             if (m_CpuTexture != null)

com.unity.perception/Tests/Runtime/GroundTruthTests/KeypointGroundTruthTests.cs

             AddTestObjectForCleanup(cam);
             AddTestObjectForCleanup(cube);
 
-            for (int i = 0; i < 10000; i++)
+            //for (int i = 0; i < 10000; i++)
              yield return null;
 
             //force all async readbacks to complete
             Assert.AreEqual(1, testCase.Count);
             var t = testCase.First();
             Assert.AreEqual(args.expectOccluded ? 1 : 2, t.keypoints[8].state);
+        }
+
+        [UnityTest]
+        public IEnumerator ManyObjects_LabelsCorrectly()
+        {
+            var incoming = new List<List<KeypointLabeler.KeypointEntry>>();
+            var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate");
+            var frameSize = 1024;
+            var texture = new RenderTexture(frameSize, frameSize, 16);
+            var labelerSelfOcclusionDistance = 0.5f;
+            var cam = SetupCamera(SetUpLabelConfig(), template, (frame, data) =>
+            {
+                incoming.Add(data);
+            }, texture, defaultSelfOcclusionDistance: labelerSelfOcclusionDistance);
+            var count = new Vector2Int(100, 100);
+
+            for (int x = 0; x < count.x; x++)
+            {
+                for (int y = 0; y < count.y; y++)
+                {
+                    var cube = TestHelper.CreateLabeledCube(scale: 1f / count.x - .01f, x: 1f / count.x * x * 2 - 1, y: 1f / count.y * y * 2 - 1);
+                    SetupCubeJoints(cube, template);
+                    cube.SetActive(true);
+                    AddTestObjectForCleanup(cube);
+                }
+            }
+            cam.SetActive(true);
+
+            AddTestObjectForCleanup(cam);
+
+            //for (int i = 0; i < 10000; i++)
+            yield return null;
+
+            //force all async readbacks to complete
+            DestroyTestObject(cam);
+            texture.Release();
+
+            var testCase = incoming.Last();
+            Assert.AreEqual(count.x * count.y, testCase.Count);
+            var idx = 0;
+            foreach (var entry in testCase)
+            {
+                Assert.AreEqual(9, entry.keypoints.Length);
+
+                CollectionAssert.AreEqual(Enumerable.Repeat(2, 4),
+                    entry.keypoints.Take(4).Select(k => k.state),
+                    $"State mismatch on front in entry {idx}");
+                CollectionAssert.AreEqual(Enumerable.Repeat(1, 4),
+                    entry.keypoints.Skip(4).Take(4).Select(k => k.state),
+                    $"State mismatch on back in entry {idx}");
+                idx++;
+            }
         }
     }
 }

com.unity.perception/Editor/GroundTruth/ShaderPreprocessor.cs

+using System.Collections.Generic;
+using System.Linq;
+using UnityEditor.Build;
+using UnityEditor.Rendering;
+using UnityEngine;
+using UnityEngine.Rendering;
+
+namespace UnityEditor.Perception.GroundTruth
+{
+    public class ShaderPreprocessor : IPreprocessShaders
+    {
+        private string[] shadersToPreprocess = new[]
+        {
+            "Perception/KeypointDepthCheck"
+        };
+        public int callbackOrder => 0;
+        public void OnProcessShader(Shader shader, ShaderSnippetData snippet, IList<ShaderCompilerData> data)
+        {
+            if (!shadersToPreprocess.Contains(shader.name))
+                return;
+
+#if HDRP_PRESENT || URP_PRESENT
+#if HDRP_PRESENT
+            ShaderKeyword keyword = new ShaderKeyword(shader, "HDRP_ENABLED");
+#else
+            ShaderKeyword keyword = new ShaderKeyword(shader, "HDRP_DISABLED");
+#endif
+            for (int i = data.Count - 1; i >= 0; --i)
+            {
+                if (!data[i].shaderKeywordSet.IsEnabled(keyword))
+                    continue;
+
+                data.RemoveAt(i);
+            }
+#endif
+        }
+    }
+}

com.unity.perception/Editor/GroundTruth/ShaderPreprocessor.cs.meta

+fileFormatVersion: 2
+guid: 2156803809034108bf01e9ab96302a9f
+timeCreated: 1619623320

com.unity.perception/Runtime/GroundTruth/Resources/KeypointDepthCheck.shader

+Shader "Perception/KeypointDepthCheck"
+{
+    Properties
+    {
+        _Positions("Positions", 2D) = "defaultTexture" {}
+        _KeypointCheckDepth("KeypointCheckDepth", 2D) = "defaultTexture" {}
+        _DepthTexture("Depth", 2DArray) = "defaultTexture" {}
+    }
+
+    HLSLINCLUDE
+
+    #pragma target 4.5
+    #pragma only_renderers d3d11 ps4 xboxone vulkan metal switch
+
+    //enable GPU instancing support
+    #pragma multi_compile HDRP_DISABLED HDRP_ENABLED
+    #pragma multi_compile_instancing
+
+    ENDHLSL
+
+    SubShader
+    {
+        Pass
+        {
+            Tags { "LightMode" = "SRP" }
+
+            Name "KeypointDepthCheck"
+            ZWrite Off
+            ZTest Always
+            Blend SrcAlpha OneMinusSrcAlpha
+            Cull Off
+
+            HLSLPROGRAM
+            #pragma only_renderers d3d11 vulkan metal
+            #pragma target 4.5
+            #pragma vertex Vert
+            #pragma fragment Frag
+
+#if HDRP_ENABLED
+
+
+            #pragma enable_d3d11_debug_symbols
+
+            Texture2D _Positions;
+            SamplerState my_point_clamp_sampler;
+            Texture2D _KeypointCheckDepth;
+
+            #include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/RenderPass/CustomPass/CustomPassCommon.hlsl"
+
+            float4 Frag(Varyings varyings) : SV_Target
+            {
+                UNITY_SETUP_STEREO_EYE_INDEX_POST_VERTEX(varyings);
+
+                float4 checkPosition = _Positions.Load(float3(varyings.positionCS.xy, 0));
+                float checkDepth = _KeypointCheckDepth.Load(float3(varyings.positionCS.xy, 0)).r;
+
+                float depth = LoadCameraDepth(float2(checkPosition.x, _ScreenSize.y - checkPosition.y));
+                PositionInputs positionInputs = GetPositionInput(checkPosition, _ScreenSize.zw, depth, UNITY_MATRIX_I_VP, UNITY_MATRIX_V);
+                depth = positionInputs.linearDepth;
+
+                uint result = depth >= checkDepth ? 1 : 0;
+                return float4(result, result, result, 1);
+            }
+#else
+            #include "UnityCG.cginc"
+
+            Texture2D _Positions;
+            SamplerState my_point_clamp_sampler;
+            Texture2D _KeypointCheckDepth;
+
+            //copied from UnityInput.hlsl
+            float4x4 _InvViewProjMatrix;
+            float4x4 _InvProjMatrix;
+
+            //#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
+
+            struct appdata
+            {
+                float4 vertex : POSITION;
+                float2 uv : TEXCOORD0;
+            };
+
+            struct v2f
+            {
+                float2 uv : TEXCOORD0;
+                float4 vertex : SV_POSITION;
+            };
+
+            v2f Vert (appdata v)
+            {
+                v2f o;
+                o.vertex = UnityObjectToClipPos(v.vertex);
+                o.uv = v.uv;
+                return o;
+            }
+
+            bool IsPerspectiveProjection()
+            {
+                return unity_OrthoParams.w == 0;
+            }
+            float ViewSpaceDepth(float depth)
+            {
+                if (IsPerspectiveProjection())
+                    return LinearEyeDepth(depth);
+                else
+                    return _ProjectionParams.y + (_ProjectionParams.z - _ProjectionParams.y) * (1 - depth);
+            }
+
+            float EncodeAndDecodeDepth(float vsDepth)
+            {
+                if (IsPerspectiveProjection())
+                {
+                    //derived from vsDepth = 1.0 / (_ZBufferParams.z * dtDepth + _ZBufferParams.w);
+                    float dtDepth = (1.0 / vsDepth - _ZBufferParams.w) / _ZBufferParams.z;
+                    dtDepth = dtDepth;
+                    return LinearEyeDepth(dtDepth);
+                }
+                else //in orthographic projections depth is linear so there is no loss of precision.
+                    return vsDepth;
+            }
+
+            Texture2D _CameraDepthTexture;
+
+            float LoadSceneDepth(uint2 uv)
+            {
+                return _CameraDepthTexture.Load(float3(uv, 0)).r;
+            }
+
+            float4 ComputeClipSpacePosition(float2 positionNDC, float deviceDepth)
+            {
+                float4 positionCS = float4(positionNDC * 2.0 - 1.0, deviceDepth, 1.0);
+
+            #if UNITY_UV_STARTS_AT_TOP
+                // Our world space, view space, screen space and NDC space are Y-up.
+                // Our clip space is flipped upside-down due to poor legacy Unity design.
+                // The flip is baked into the projection matrix, so we only have to flip
+                // manually when going from CS to NDC and back.
+                positionCS.y = -positionCS.y;
+            #endif
+
+                return positionCS;
+            }
+
+            float3 ComputeWorldSpacePosition(float2 positionNDC, float deviceDepth, float4x4 invViewProjMatrix)
+            {
+                float4 positionCS  = ComputeClipSpacePosition(positionNDC, deviceDepth);
+                float4 hpositionWS = mul(invViewProjMatrix, positionCS);
+                return hpositionWS.xyz / hpositionWS.w;
+            }
+            fixed4 Frag (v2f i) : SV_Target
+            {
+                float4 checkPosition = _Positions.Load(float3(i.vertex.xy, 0));
+                float depthVSToCheck = _KeypointCheckDepth.Load(float3(i.vertex.xy, 0)).r;
+
+                float depth = LoadSceneDepth(float2(checkPosition.x, _ScreenParams.y - checkPosition.y));
+                float depthVSActual = ViewSpaceDepth(depth);
+
+                uint result = depthVSActual >= depthVSToCheck  ? 1 : 0;
+                return float4(result, result, result, 1);
+            }
+#endif
+            ENDHLSL
+        }
+    }
+    Fallback Off
+}

com.unity.perception/Runtime/GroundTruth/Resources/KeypointDepthCheckHDRP.shader

-Shader "Perception/KeypointDepthCheck"
-{
-    Properties
-    {
-        _Positions("Positions", 2D) = "defaultTexture" {}
-        _KeypointCheckDepth("KeypointCheckDepth", 2D) = "defaultTexture" {}
-        _DepthTexture("Depth", 2DArray) = "defaultTexture" {}
-    }
-
-    HLSLINCLUDE
-
-    #pragma target 4.5
-    #pragma only_renderers d3d11 ps4 xboxone vulkan metal switch
-
-    //enable GPU instancing support
-    #pragma multi_compile_instancing
-
-    ENDHLSL
-
-    SubShader
-    {
-        Pass
-        {
-            Tags { "LightMode" = "SRP" }
-
-            Name "KeypointDepthCheck"
-            ZWrite Off
-            ZTest Always
-            Blend SrcAlpha OneMinusSrcAlpha
-            Cull Off
-
-            HLSLPROGRAM
-
-            #pragma multi_compile HDRP_DISABLED HDRP_ENABLED
-            #pragma enable_d3d11_debug_symbols
-            #pragma only_renderers d3d11 vulkan metal
-            #pragma target 4.5
-            #pragma vertex Vert
-            #pragma fragment FullScreenPass
-
-
-            //UNITY_DECLARE_TEX2DARRAY(_DepthTexture);
-            //sampler2D _CameraDepthTexture;
-
-            Texture2D _Positions;
-            SamplerState my_point_clamp_sampler;
-            Texture2D _KeypointCheckDepth;
-
-#if HDRP_ENABLED
-
-            #include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/RenderPass/CustomPass/CustomPassCommon.hlsl"
-
-            float4 FullScreenPass(Varyings varyings) : SV_Target
-            {
-                UNITY_SETUP_STEREO_EYE_INDEX_POST_VERTEX(varyings);
-
-                float4 checkPosition = _Positions.Load(float3(varyings.positionCS.xy, 0));
-                float4 checkDepth = _KeypointCheckDepth.Load(float3(varyings.positionCS.xy, 0));
-
-                float depth = LoadCameraDepth(float2(checkPosition.x, _ScreenSize.y - checkPosition.y));
-                PositionInputs positionInputs = GetPositionInput(checkPosition, _ScreenSize.zw, depth, UNITY_MATRIX_I_VP, UNITY_MATRIX_V);
-                depth = positionInputs.linearDepth;
-
-                //encode and decode checkDepth to account for loss of precision with depth values close to far plane
-                float4 viewPos = mul(UNITY_MATRIX_I_VP, float4(positionInputs.positionWS.x, positionInputs.positionWS.y, checkDepth.x, 1.0));
-                float4 positionCheckWS = mul(UNITY_MATRIX_VP, viewPos);
-
-                //depth = LinearEyeDepth(depth, _ZBufferParams);
-                // float depth = UNITY_SAMPLE_TEX2DARRAY(_DepthTexture, float3(checkPosition.xy, 0)).r; //SAMPLE_DEPTH_TEXTURE(_DepthTexture, checkPosition.xy);
-                //float depth_decoded = LinearEyeDepth(depth);
-                // float depth_decoded = Linear01Depth(depth);
-                uint result = depth >= positionCheckWS.z ? 1 : 0;
-                return float4(result, result, result, 1);
-            }
-#else
-
-            #include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
-            #include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Filtering.hlsl"
-            #include "Packages/com.unity.render-pipelines.universal/ShaderLibrary/DeclareDepthTexture.hlsl"
-            #include "Packages/com.unity.render-pipelines.universal/Shaders/PostProcessing/Common.hlsl"
-
-            bool IsPerspectiveProjection()
-            {
-                return unity_OrthoParams.w == 0;
-            }
-            float ViewSpaceDepth(float depth)
-            {
-                if (IsPerspectiveProjection())
-                    return LinearEyeDepth(depth, _ZBufferParams);
-                else
-                    return _ProjectionParams.y + (_ProjectionParams.z - _ProjectionParams.y) * (1 - depth);
-            }
-
-            float4 FullScreenPass(Varyings varyings) : SV_Target
-            {
-                UNITY_SETUP_STEREO_EYE_INDEX_POST_VERTEX(varyings);
-
-                float4 checkPosition = _Positions.Load(float3(varyings.positionCS.xy, 0));
-                float checkDepth = _KeypointCheckDepth.Load(float3(varyings.positionCS.xy, 0)).r;
-
-                float depth = LoadSceneDepth(float2(checkPosition.x, _ScreenParams.y - checkPosition.y));
-
-                depth = ViewSpaceDepth(depth);
-
-                //encode and decode checkDepth to account for loss of precision with depth values close to far plane
-                PositionInputs positionInputs = GetPositionInput(checkPosition, _ScreenParams.zw - float2(1, 1), depth, UNITY_MATRIX_I_VP, UNITY_MATRIX_V);
-                float4 viewPos = mul(UNITY_MATRIX_V, float4(positionInputs.positionWS.x, positionInputs.positionWS.y, checkDepth, 1.0));
-                float4 positionCheckWS = mul(UNITY_MATRIX_I_V, viewPos);
-                float depthCompare = positionCheckWS.z;
-
-                //float depthCompare = checkDepth;
-
-                //depth = LinearEyeDepth(depth, _ZBufferParams);
-                // float depth = UNITY_SAMPLE_TEX2DARRAY(_DepthTexture, float3(checkPosition.xy, 0)).r; //SAMPLE_DEPTH_TEXTURE(_DepthTexture, checkPosition.xy);
-                //float depth_decoded = LinearEyeDepth(depth);
-                // float depth_decoded = Linear01Depth(depth);
-                uint result = depth >= depthCompare ? 1 : 0;
-                return float4(result, result, result, 1);
-            }
-#endif
-            ENDHLSL
-        }
-    }
-    Fallback Off
-}