Keypoint occlusion (#219)

* keypoint occlusion fixes 

does not report a labeled asset that is completely out of the frame, also individual keypoints are now labeled as 0 for off screen or not in model, 1 for in model and on screen but occluded by something else, and 2 everything works and is visible

* Keypoint templates now save their color info to json

* Updates for PR comments

com.unity.perception/CHANGELOG.md

 
 Added new ScenarioBase virtual lifecycle hooks: OnAwake, OnStart, OnComplete, and OnIdle.
 
+Keypoint occlusion has been added, no keypoint information will be recorded for a labeled asset completely out of the camera's frustum. 
+
+New keypoint tests have been added to test keypoint state.
+
+The color of keypoints and connections are now reported in the annotation definition json file for keypoint templates.
+
 ### Changed
 
 Renamed all appearances of the term `KeyPoint` within types and names to `Keypoint`.

com.unity.perception/Documentation~/Schema/Synthetic_Dataset_Schema.md

 
 Keypoint data, commonly used for human pose estimation. A keypoint capture is associated to a template that defines the keypoints (see annotation.definition file).
 Each keypoint record maps a tuple of (instance, label) to template, pose, and an array of keypoints. A keypoint will exist in this record for each keypoint defined in the template file.
-If a given keypoint doesn't exist in the labeled gameobject, then that keypoint will have a state value of 0; if it does exist then it will have a keypoint value of 2.
+If a given keypoint doesn't exist in the labeled gameobject, then that keypoint will have a state value of 0; if it exists but is not visible, it will have a state value of 1, 
+if it exists and is visible it will have a state value of 2.
 ```
 keypoints {
   label_id:      <int>   -- Integer identifier of the label
       index:     <int>   -- Index of keypoint in template
       x:         <float> -- X pixel coordinate of keypoint
       y:         <float> -- Y pixel coordinate of keypoint
-      state:     <int>   -- 0: keypoint does not exist, 2 keypoint exists
+      state:     <int>   -- 0: keypoint does not exist, 1 keypoint exists but is not visible, 2 keypoint exists and is visible
     }, ...
   ]
 }
     {
       label:         <str>           -- The label of the joint
       index:         <int>           -- The index of the joint
+      color {                        -- [Optional] The color to use for the visualization of the keypoint
+        r:           <float>         -- Value from 0 to 1 for the red channel
+        g:           <float>         -- Value from 0 to 1 for the green channel
+        b:           <float>         -- Value from 0 to 1 for the blue channel
+        a:           <float>         -- Value from 0 to 1 for the alpha channel
+      }
     }, ...
   ]
   skeleton [                         -- Array of skeletal connections (which joints have connections between one another) defined in this template
+      color {                        -- [Optional] The color to use for the visualization of the bone
+        r:           <float>         -- Value from 0 to 1 for the red channel
+        g:           <float>         -- Value from 0 to 1 for the green channel
+        b:           <float>         -- Value from 0 to 1 for the blue channel
+        a:           <float>         -- Value from 0 to 1 for the alpha channel
+      }
     }, ...
   ]
 }

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

         EntityQuery m_EntityQuery;
         Texture2D m_MissingTexture;
 
+        Dictionary<int, (AsyncAnnotation annotation, Dictionary<uint, KeypointEntry> keypoints)> m_AsyncAnnotations;
+        List<KeypointEntry> m_ToReport;
+
+        int m_CurrentFrame;
+
-        public event Action<List<KeypointEntry>> KeypointsComputed;
+        public event Action<int, List<KeypointEntry>> KeypointsComputed;
 
         /// <summary>
         /// Creates a new key point labeler. This constructor creates a labeler that
 
             m_EntityQuery = World.DefaultGameObjectInjectionWorld.EntityManager.CreateEntityQuery(typeof(Labeling), typeof(GroundTruthInfo));
 
-            m_KeypointEntries = new List<KeypointEntry>();
-
+
+            m_AsyncAnnotations = new Dictionary<int, (AsyncAnnotation, Dictionary<uint, KeypointEntry>)>();
+            m_ToReport = new List<KeypointEntry>();
+            m_CurrentFrame = 0;
+
+            perceptionCamera.InstanceSegmentationImageReadback += OnInstanceSegmentationImageReadback;
+        }
+
+        bool AreEqual(Color32 lhs, Color32 rhs)
+        {
+            return lhs.r == rhs.r && lhs.g == rhs.g && lhs.b == rhs.b && lhs.a == rhs.a;
+        }
+
+        void OnInstanceSegmentationImageReadback(int frameCount, NativeArray<Color32> data, RenderTexture renderTexture)
+        {
+            if (!m_AsyncAnnotations.TryGetValue(frameCount, out var asyncAnnotation))
+                return;
+
+            m_AsyncAnnotations.Remove(frameCount);
+
+            var width = renderTexture.width;
+
+            m_ToReport.Clear();
+
+            foreach (var keypointSet in asyncAnnotation.keypoints)
+            {
+                if (InstanceIdToColorMapping.TryGetColorFromInstanceId(keypointSet.Key, out var idColor))
+                {
+                    var shouldReport = false;
+
+                    foreach (var keypoint in keypointSet.Value.keypoints)
+                    {
+                        // If the keypoint isn't mapped to a body part keep it at 0
+                        if (keypoint.state == 0) continue;
+
+                        if (keypoint.x < 0 || keypoint.x > width || keypoint.y < 0 || keypoint.y > renderTexture.height)
+                        {
+                            keypoint.state = 0;
+                            keypoint.x = 0;
+                            keypoint.y = 0;
+                        }
+                        else
+                        {
+                            // Get the pixel color at the keypoints location
+                            var height = renderTexture.height - (int)keypoint.y;
+                            var pixel = data[height * width + (int)keypoint.x];
+
+                            keypoint.state = AreEqual(pixel, idColor) ? 2 : 1;
+                            shouldReport = true;
+                        }
+                    }
+
+                    if (shouldReport)
+                        m_ToReport.Add(keypointSet.Value);
+                }
+            }
+
+            KeypointsComputed?.Invoke(frameCount, m_ToReport);
+            asyncAnnotation.annotation.ReportValues(m_ToReport);
-            var reporter = perceptionCamera.SensorHandle.ReportAnnotationAsync(m_AnnotationDefinition);
+            m_CurrentFrame = Time.frameCount;
+
+            var annotation = perceptionCamera.SensorHandle.ReportAnnotationAsync(m_AnnotationDefinition);
+            var keypoints = new Dictionary<uint, KeypointEntry>();
+
+            m_AsyncAnnotations[m_CurrentFrame] = (annotation, keypoints);
-
-            m_KeypointEntries.Clear();
 
             foreach (var entity in entities)
             {
             entities.Dispose();
-
-            KeypointsComputed?.Invoke(m_KeypointEntries);
-            reporter.ReportValues(m_KeypointEntries);
         }
 
         // ReSharper disable InconsistentNaming
         // ReSharper restore NotAccessedField.Global
         // ReSharper restore NotAccessedField.Local
 
+        float GetCaptureHeight()
+        {
+            return perceptionCamera.attachedCamera.targetTexture != null ?
+                perceptionCamera.attachedCamera.targetTexture.height : Screen.height;
+        }
+
-            pt.y = Screen.height - pt.y;
+            pt.y = GetCaptureHeight() - pt.y;
-
-        List<KeypointEntry> m_KeypointEntries;
 
         struct CachedData
         {
                     cachedData.keypoints.pose = GetPose(cachedData.animator);
                 }
 
-                m_KeypointEntries.Add(cachedData.keypoints);
+                m_AsyncAnnotations[m_CurrentFrame].keypoints[labeledEntity.instanceId] = cachedData.keypoints;
             }
         }
 
         /// <inheritdoc/>
         protected override void OnVisualize()
         {
+            if (m_ToReport == null) return;
+
             var jointTexture = activeTemplate.jointTexture;
             if (jointTexture == null) jointTexture = m_MissingTexture;
 
-            foreach (var entry in m_KeypointEntries)
+            foreach (var entry in m_ToReport)
             {
                 foreach (var bone in activeTemplate.skeleton)
                 {
-                    if (joint1.state != 0 && joint2.state != 0)
+                    if (joint1.state == 2 && joint2.state == 2)
                     {
                         VisualizationHelper.DrawLine(joint1.x, joint1.y, joint2.x, joint2.y, bone.color, 8, skeletonTexture);
                     }
                 {
-                    if (keypoint.state != 0)
+                    if (keypoint.state == 2)
                         VisualizationHelper.DrawPoint(keypoint.x, keypoint.y, activeTemplate.keypoints[keypoint.index].color, 8, jointTexture);
                 }
             }
         {
             public string label;
             public int index;
+            public Color color;
         }
 
         [Serializable]
             public int joint2;
+            public Color color;
         }
 
         [Serializable]
                 json.key_points[i] = new JointJson
                 {
                     label = input.keypoints[i].label,
-                    index = i
+                    index = i,
+                    color = input.keypoints[i].color
                 };
             }
 
                 {
                     joint1 = input.skeleton[i].joint1,
-                    joint2 = input.skeleton[i].joint2
+                    joint2 = input.skeleton[i].joint2,
+                    color = input.skeleton[i].color
                 };
             }
 

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

     [TestFixture]
     public class KeypointGroundTruthTests : GroundTruthTestBase
     {
-        static GameObject SetupCamera(IdLabelConfig config, KeypointTemplate template, Action<List<KeypointLabeler.KeypointEntry>> computeListener)
+        static GameObject SetupCamera(IdLabelConfig config, KeypointTemplate template, Action<int, List<KeypointLabeler.KeypointEntry>> computeListener, RenderTexture renderTexture = null)
         {
             var cameraObject = new GameObject();
             cameraObject.SetActive(false);
             camera.farClipPlane = 1000;
 
             camera.transform.position = new Vector3(0, 0, -10);
+
+            if (renderTexture)
+            {
+                camera.targetTexture = renderTexture;
+            }
 
             var perceptionCamera = cameraObject.AddComponent<PerceptionCamera>();
             perceptionCamera.captureRgbImages = false;
 
         static void SetupCubeJoints(GameObject cube, KeypointTemplate template)
         {
-            SetupCubeJoint(cube, template, "FrontLowerLeft", -0.5f, -0.5f, -0.5f);
-            SetupCubeJoint(cube, template, "FrontUpperLeft", -0.5f, 0.5f, -0.5f);
-            SetupCubeJoint(cube, template, "FrontUpperRight", 0.5f, 0.5f, -0.5f);
-            SetupCubeJoint(cube, template, "FrontLowerRight", 0.5f, -0.5f, -0.5f);
-            SetupCubeJoint(cube, template, "BackLowerLeft", -0.5f, -0.5f, 0.5f);
-            SetupCubeJoint(cube, template, "BackUpperLeft", -0.5f, 0.5f, 0.5f);
-            SetupCubeJoint(cube, template, "BackUpperRight", 0.5f, 0.5f, 0.5f);
-            SetupCubeJoint(cube, template, "BackLowerRight", 0.5f, -0.5f, 0.5f);
+            SetupCubeJoint(cube, template, "FrontLowerLeft", -0.495f, -0.495f, -0.495f);
+            SetupCubeJoint(cube, template, "FrontUpperLeft", -0.495f, 0.495f, -0.495f);
+            SetupCubeJoint(cube, template, "FrontUpperRight", 0.495f, 0.495f, -0.495f);
+            SetupCubeJoint(cube, template, "FrontLowerRight", 0.495f, -0.495f, -0.495f);
+            SetupCubeJoint(cube, template, "BackLowerLeft", -0.495f, -0.495f, 0.495f);
+            SetupCubeJoint(cube, template, "BackUpperLeft", -0.495f, 0.495f, 0.495f);
+            SetupCubeJoint(cube, template, "BackUpperRight", 0.495f, 0.495f, 0.495f);
+            SetupCubeJoint(cube, template, "BackLowerRight", 0.495f, -0.495f, 0.495f);
         }
 
         [UnityTest]
             var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate");
+            var texture = new RenderTexture(1024, 768, 16);
+            texture.Create();
+            var cam = SetupCamera(SetUpLabelConfig(), template, (frame, data) =>
+            {
+                incoming.Add(data);
+            }, texture);
-            var cam = SetupCamera(SetUpLabelConfig(), template, (data) =>
+            var cube = TestHelper.CreateLabeledCube(scale: 6, z: 8);
+            SetupCubeJoints(cube, template);
+
+            cube.SetActive(true);
+            cam.SetActive(true);
+
+            AddTestObjectForCleanup(cam);
+            AddTestObjectForCleanup(cube);
+
+            yield return null;
+            yield return null;
+
+            texture.Release();
+
+            var testCase = incoming.Last();
+            Assert.AreEqual(1, testCase.Count);
+            var t = testCase.First();
+            Assert.NotNull(t);
+            Assert.AreEqual(1, t.instance_id);
+            Assert.AreEqual(1, t.label_id);
+            Assert.AreEqual(template.templateID.ToString(), t.template_guid);
+            Assert.AreEqual(9, t.keypoints.Length);
+
+            Assert.AreEqual(t.keypoints[0].x, t.keypoints[1].x);
+            Assert.AreEqual(t.keypoints[2].x, t.keypoints[3].x);
+            Assert.AreEqual(t.keypoints[4].x, t.keypoints[5].x);
+            Assert.AreEqual(t.keypoints[6].x, t.keypoints[7].x);
+
+            Assert.AreEqual(t.keypoints[0].y, t.keypoints[3].y);
+            Assert.AreEqual(t.keypoints[1].y, t.keypoints[2].y);
+            Assert.AreEqual(t.keypoints[4].y, t.keypoints[7].y);
+            Assert.AreEqual(t.keypoints[5].y, t.keypoints[6].y);
+
+            for (var i = 0; i < 9; i++) Assert.AreEqual(i, t.keypoints[i].index);
+            for (var i = 0; i < 8; i++) Assert.AreEqual(2, t.keypoints[i].state);
+            Assert.Zero(t.keypoints[8].state);
+            Assert.Zero(t.keypoints[8].x);
+            Assert.Zero(t.keypoints[8].y);
+        }
+
+        [UnityTest]
+        public IEnumerator Keypoint_TestAllOffScreen()
+        {
+            var incoming = new List<List<KeypointLabeler.KeypointEntry>>();
+            var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate");
+
+            var cam = SetupCamera(SetUpLabelConfig(), template, (frame, data) =>
             {
                 incoming.Add(data);
             });
 
+            cube.transform.position = new Vector3(-1000, -1000, 0);
+
             cube.SetActive(true);
             cam.SetActive(true);
 
             yield return null;
             yield return null;
 
+            foreach (var i in incoming)
+            {
+                Assert.Zero(i.Count);
+            }
+        }
+
+        [UnityTest]
+        public IEnumerator Keypoint_TestPartialOffScreen()
+        {
+            var incoming = new List<List<KeypointLabeler.KeypointEntry>>();
+            var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate");
+            var texture = new RenderTexture(1024, 768, 16);
+            texture.Create();
+            var cam = SetupCamera(SetUpLabelConfig(), template, (frame, data) =>
+            {
+                incoming.Add(data);
+            }, texture);
+
+            var cube = TestHelper.CreateLabeledCube(scale: 6, z: 8);
+            SetupCubeJoints(cube, template);
+            cube.transform.position += Vector3.right * 13.5f;
+
+            cube.SetActive(true);
+            cam.SetActive(true);
+
+            AddTestObjectForCleanup(cam);
+            AddTestObjectForCleanup(cube);
+
+            for (var i = 0; i < 2; i++)
+                yield return null;
+
+            if (texture != null) texture.Release();
+
+            var testCase = incoming.Last();
+            Assert.AreEqual(1, testCase.Count);
+            var t = testCase.First();
+            Assert.NotNull(t);
+            Assert.AreEqual(1, t.instance_id);
+            Assert.AreEqual(1, t.label_id);
+            Assert.AreEqual(template.templateID.ToString(), t.template_guid);
+            Assert.AreEqual(9, t.keypoints.Length);
+
+            Assert.NotZero(t.keypoints[0].state);
+            Assert.NotZero(t.keypoints[1].state);
+            Assert.NotZero(t.keypoints[4].state);
+            Assert.NotZero(t.keypoints[5].state);
+
+            Assert.Zero(t.keypoints[2].state);
+            Assert.Zero(t.keypoints[3].state);
+            Assert.Zero(t.keypoints[6].state);
+            Assert.Zero(t.keypoints[7].state);
+
+            for (var i = 0; i < 9; i++) Assert.AreEqual(i, t.keypoints[i].index);
+            Assert.Zero(t.keypoints[8].state);
+            Assert.Zero(t.keypoints[8].x);
+            Assert.Zero(t.keypoints[8].y);
+        }
+
+        [UnityTest]
+        public IEnumerator Keypoint_TestAllOnScreen()
+        {
+            var incoming = new List<List<KeypointLabeler.KeypointEntry>>();
+            var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate");
+            var texture = new RenderTexture(1024, 768, 16);
+            texture.Create();
+            var cam = SetupCamera(SetUpLabelConfig(), template, (frame, data) =>
+            {
+                incoming.Add(data);
+            }, texture);
+
+            var cube = TestHelper.CreateLabeledCube(scale: 6, z: 8);
+            SetupCubeJoints(cube, template);
+
+            cube.SetActive(true);
+            cam.SetActive(true);
+
+            AddTestObjectForCleanup(cam);
+            AddTestObjectForCleanup(cube);
+
+            for (var i = 0; i < 2; i++)
+                yield return null;
+
+            if (texture != null) texture.Release();
+
             var testCase = incoming.Last();
             Assert.AreEqual(1, testCase.Count);
             var t = testCase.First();
             Assert.Zero(t.keypoints[8].state);
             Assert.Zero(t.keypoints[8].x);
             Assert.Zero(t.keypoints[8].y);
+
+
+
+        [UnityTest]
+        public IEnumerator Keypoint_TestAllBlockedByOther()
+        {
+            var incoming = new List<List<KeypointLabeler.KeypointEntry>>();
+            var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate");
+            var texture = new RenderTexture(1024, 768, 16);
+            texture.Create();
+            var cam = SetupCamera(SetUpLabelConfig(), template, (frame, data) =>
+            {
+                incoming.Add(data);
+            }, texture);
+
+            var cube = TestHelper.CreateLabeledCube(scale: 6, z: 8);
+            SetupCubeJoints(cube, template);
+
+            var blocker = GameObject.CreatePrimitive(PrimitiveType.Cube);
+            blocker.transform.position = new Vector3(0, 0, 5);
+            blocker.transform.localScale = new Vector3(7, 7, 7);
+
+            cube.SetActive(true);
+            cam.SetActive(true);
+
+            AddTestObjectForCleanup(cam);
+            AddTestObjectForCleanup(cube);
+            AddTestObjectForCleanup(blocker);
+
+            for (var i = 0; i < 2; i++)
+                yield return null;
+
+            if (texture != null) texture.Release();
+
+            var testCase = incoming.Last();
+            Assert.AreEqual(1, testCase.Count);
+            var t = testCase.First();
+            Assert.NotNull(t);
+            Assert.AreEqual(1, t.instance_id);
+            Assert.AreEqual(1, t.label_id);
+            Assert.AreEqual(template.templateID.ToString(), t.template_guid);
+            Assert.AreEqual(9, t.keypoints.Length);
+
+            for (var i = 0; i < 8; i++)
+                Assert.AreEqual(t.keypoints[i].state, 1);
+
+            for (var i = 0; i < 9; i++) Assert.AreEqual(i, t.keypoints[i].index);
+            Assert.Zero(t.keypoints[8].state);
+            Assert.Zero(t.keypoints[8].x);
+            Assert.Zero(t.keypoints[8].y);
+        }
+
+        [UnityTest]
+        public IEnumerator Keypoint_TestPartiallyBlockedByOther()
+        {
+            var incoming = new List<List<KeypointLabeler.KeypointEntry>>();
+            var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate");
+            var texture = new RenderTexture(1024, 768, 16);
+            texture.Create();
+            var cam = SetupCamera(SetUpLabelConfig(), template, (frame, data) =>
+            {
+                incoming.Add(data);
+            }, texture);
+
+            var cube = TestHelper.CreateLabeledCube(scale: 6, z: 8);
+            SetupCubeJoints(cube, template);
+
+            var blocker = GameObject.CreatePrimitive(PrimitiveType.Cube);
+            blocker.transform.position = new Vector3(3, 0, 5);
+            blocker.transform.localScale = new Vector3(7, 7, 7);
+
+            cube.SetActive(true);
+            cam.SetActive(true);
+
+            AddTestObjectForCleanup(cam);
+            AddTestObjectForCleanup(cube);
+            AddTestObjectForCleanup(blocker);
+
+            for (var i = 0; i < 2; i++)
+                yield return null;
+
+            if (texture != null) texture.Release();
+
+            var testCase = incoming.Last();
+            Assert.AreEqual(1, testCase.Count);
+            var t = testCase.First();
+            Assert.NotNull(t);
+            Assert.AreEqual(1, t.instance_id);
+            Assert.AreEqual(1, t.label_id);
+            Assert.AreEqual(template.templateID.ToString(), t.template_guid);
+            Assert.AreEqual(9, t.keypoints.Length);
+
+            Assert.AreEqual(t.keypoints[0].state, 2);
+            Assert.AreEqual(t.keypoints[1].state, 2);
+            Assert.AreEqual(t.keypoints[4].state, 2);
+            Assert.AreEqual(t.keypoints[5].state, 2);
+
+            Assert.AreEqual(t.keypoints[2].state, 1);
+            Assert.AreEqual(t.keypoints[3].state, 1);
+            Assert.AreEqual(t.keypoints[6].state, 1);
+            Assert.AreEqual(t.keypoints[7].state, 1);
+
+            for (var i = 0; i < 9; i++) Assert.AreEqual(i, t.keypoints[i].index);
+            Assert.Zero(t.keypoints[8].state);
+            Assert.Zero(t.keypoints[8].x);
+            Assert.Zero(t.keypoints[8].y);
+        }
+
     }
 }