Moving self occlusion distance to the KeypointTemplate

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

         [SerializeField]
         public List<string> labels = new List<string>();
 
-        public SelfOcclusionDistanceSource selfOcclusionDistanceSource = SelfOcclusionDistanceSource.JointLabel;
+        /// <summary>
+        /// Whether <see cref="selfOcclusionDistance"/> should be used instead of the one specified in the <see cref="KeypointTemplate"/>.
+        /// </summary>
+        public bool overrideSelfOcclusionDistance = false;
+        /// <summary>
+        /// Whether <see cref="selfOcclusionDistance"/> should be used instead of the one specified in the <see cref="KeypointTemplate"/>.
+        /// </summary>
-        public void OnBeforeSerialize()
+        /// <summary>
+        /// Internal method for serialization.
+        /// </summary>
+        void ISerializationCallbackReceiver.OnBeforeSerialize()
-        public void OnAfterDeserialize()
+        /// <summary>
+        /// Internal method for serialization.
+        /// </summary>
+        void ISerializationCallbackReceiver.OnAfterDeserialize()
         {
             if (templateInformation != null)
             {
         {
             Gizmos.DrawIcon(transform.position, "Packages/com.unity.perception/Editor/Icons/Keypoint.png", false);
         }
-        
+
         private void OnDrawGizmosSelected()
         {
             if (singlePerceptionCamera == null)
 
 #endif
             float occlusionDistance;
-            switch (selfOcclusionDistanceSource)
+            if (this.overrideSelfOcclusionDistance)
-                case SelfOcclusionDistanceSource.JointLabel:
-                    occlusionDistance = selfOcclusionDistance;
-                    break;
-                case SelfOcclusionDistanceSource.KeypointLabeler:
-                    if (singlePerceptionCamera == null)
-                    {
-                        occlusionDistance = KeypointLabeler.defaultSelfOcclusionDistance;
-                    }
+                occlusionDistance = selfOcclusionDistance;
+            }
+            else
+            {
+                if (singlePerceptionCamera == null)
+                {
+                    occlusionDistance = KeypointDefinition.defaultSelfOcclusionDistance;
+                }
+                else
+                {
+                    var keypointLabeler = (KeypointLabeler) singlePerceptionCamera.labelers.FirstOrDefault(l => l is KeypointLabeler);
+                    var template = keypointLabeler?.activeTemplate;
+                    if (template == null)
+                        occlusionDistance = KeypointDefinition.defaultSelfOcclusionDistance;
-                        var keypointLabeler = (KeypointLabeler) singlePerceptionCamera.labelers.FirstOrDefault(l => l is KeypointLabeler);
-                        if (keypointLabeler == null)
-                            occlusionDistance = KeypointLabeler.defaultSelfOcclusionDistance;
+                        KeypointDefinition matchingKeypoint = null;
+                        foreach (var k in template.keypoints)
+                        {
+                            if (this.labels.Contains(k.label))
+                            {
+                                matchingKeypoint = k;
+                                break;
+                            }
+                        }
+
+                        if (matchingKeypoint == null)
+                            occlusionDistance = KeypointDefinition.defaultSelfOcclusionDistance;
-                            occlusionDistance = keypointLabeler.selfOcclusionDistance;
+                            occlusionDistance = matchingKeypoint.selfOcclusionDistance;
-                    break;
-                default:
-                    throw new InvalidOperationException("Invalid SelfOcclusionDistanceSource");
+                }
             }
 
             var occlusionDistanceScale = transform.lossyScale * occlusionDistance;

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

     [Serializable]
     public sealed class KeypointLabeler : CameraLabeler
     {
-        internal const float defaultSelfOcclusionDistance = 0.15f;
         // Smaller texture sizes produce assertion failures in the engine
         const int k_MinTextureWidth = 8;
 
         /// <see cref="KeypointObjectFilter"/>
         /// </summary>
         public KeypointObjectFilter objectFilter;
-        /// <summary>
-        /// The max distance a keypoint can be from the front of an object before it is considered occluded
-        /// </summary>
-        public float selfOcclusionDistance = defaultSelfOcclusionDistance;
         // ReSharper restore MemberCanBePrivate.Global
 
         AnnotationDefinition m_AnnotationDefinition;
                         if (bone != null)
                         {
                             var bonePosition = bone.position;
-                            var jointSelfOcclusionDistance = JointSelfOcclusionDistance(bone, bonePosition, cameraPosition, cameraforward, this.selfOcclusionDistance);
+                            var jointSelfOcclusionDistance = JointSelfOcclusionDistance(bone, bonePosition, cameraPosition, cameraforward, pt.selfOcclusionDistance);
                             InitKeypoint(bonePosition, cachedData, checkLocationsSlice, i, jointSelfOcclusionDistance);
                         }
                     }
                 // their locations
-                foreach (var (joint, idx) in cachedData.overrides)
+                foreach (var (joint, templateIdx) in cachedData.overrides)
-                    switch (joint.selfOcclusionDistanceSource)
-                    {
-                        case SelfOcclusionDistanceSource.JointLabel:
-                            resolvedSelfOcclusionDistance = joint.selfOcclusionDistance;
-                            break;
-                        case SelfOcclusionDistanceSource.KeypointLabeler:
-                            resolvedSelfOcclusionDistance = selfOcclusionDistance;
-                            break;
-                        default:
-                            throw new NotImplementedException("Invalid SelfOcclusionDistanceSource");
-                    }
+                    if (joint.overrideSelfOcclusionDistance)
+                        resolvedSelfOcclusionDistance = joint.selfOcclusionDistance;
+                    else
+                        resolvedSelfOcclusionDistance = activeTemplate.keypoints[templateIdx].selfOcclusionDistance;
+
-                    InitKeypoint(jointPosition, cachedData, checkLocationsSlice, idx, jointSelfOcclusionDistance);
+                    InitKeypoint(jointPosition, cachedData, checkLocationsSlice, templateIdx, jointSelfOcclusionDistance);
                 }
 
                 cachedData.keypoints.pose = "unset";

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

     [Serializable]
     public class KeypointDefinition
     {
+        internal const float defaultSelfOcclusionDistance = 0.15f;
         /// <summary>
         /// The name of the keypoint
         /// </summary>
         /// The color of the keypoint in the visualization
         /// </summary>
         public Color color = Color.blue;
+        /// <summary>
+        /// The max distance the keypoint can be behind an object before it is considered occluded. Distance is computed
+        /// in the local space of the joint, so object scale is computed
+        /// </summary>
+        public float selfOcclusionDistance = defaultSelfOcclusionDistance;
     }
 
     /// <summary>

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

 #endif
         }
 
-        static GameObject SetupCamera(IdLabelConfig config, KeypointTemplate template, Action<int, List<KeypointLabeler.KeypointEntry>> computeListener, RenderTexture renderTexture = null, KeypointObjectFilter keypointObjectFilter = KeypointObjectFilter.Visible, float defaultSelfOcclusionDistance = 0.15f)
+        static GameObject SetupCamera(IdLabelConfig config, KeypointTemplate template, Action<int, List<KeypointLabeler.KeypointEntry>> computeListener, RenderTexture renderTexture = null, KeypointObjectFilter keypointObjectFilter = KeypointObjectFilter.Visible)
         {
             var cameraObject = new GameObject();
             cameraObject.SetActive(false);
             perceptionCamera.captureRgbImages = false;
             var keyPointLabeler = new KeypointLabeler(config, template);
             keyPointLabeler.objectFilter = keypointObjectFilter;
-            keyPointLabeler.selfOcclusionDistance = defaultSelfOcclusionDistance;
             if (computeListener != null)
                 keyPointLabeler.KeypointsComputed += computeListener;
 
         }
 
-        static KeypointTemplate CreateTestTemplate(Guid guid, string label)
+        static KeypointTemplate CreateTestTemplate(Guid guid, string label, float selfOcclusionDistance = 0.15f)
         {
             var keypoints = new[]
             {
                     associateToRig = false,
-                    color = Color.black
+                    color = Color.black,
+                    selfOcclusionDistance = selfOcclusionDistance
-                    color = Color.black
+                    color = Color.black,
+                    selfOcclusionDistance = selfOcclusionDistance
-                    color = Color.black
+                    color = Color.black,
+                    selfOcclusionDistance = selfOcclusionDistance
-                    color = Color.black
+                    color = Color.black,
+                    selfOcclusionDistance = selfOcclusionDistance
-                    color = Color.black
+                    color = Color.black,
+                    selfOcclusionDistance = selfOcclusionDistance
-                    color = Color.black
+                    color = Color.black,
+                    selfOcclusionDistance = selfOcclusionDistance
-                    color = Color.black
+                    color = Color.black,
+                    selfOcclusionDistance = selfOcclusionDistance
-                    color = Color.black
+                    color = Color.black,
+                    selfOcclusionDistance = selfOcclusionDistance
-                    color = Color.black
+                    color = Color.black,
+                    selfOcclusionDistance = selfOcclusionDistance
                 }
             };
 
             jointLabel.labels.Add(label);
             if (selfOcclusionDistance.HasValue)
             {
-                jointLabel.selfOcclusionDistanceSource = SelfOcclusionDistanceSource.JointLabel;
+                jointLabel.overrideSelfOcclusionDistance = true;
-                jointLabel.selfOcclusionDistanceSource = SelfOcclusionDistanceSource.KeypointLabeler;
+                jointLabel.overrideSelfOcclusionDistance = false;
         }
 
         static void SetupCubeJoints(GameObject cube, KeypointTemplate template, float? selfOcclusionDistance = null)
             [Values(ProjectionKind.Orthographic, ProjectionKind.Projection)] ProjectionKind projectionKind)
         {
             var incoming = new List<List<KeypointLabeler.KeypointEntry>>();
-            var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate");
-            var frameSize = 1024;
-            var texture = new RenderTexture(frameSize, frameSize, 16);
+            var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate", selfOcclusionDistance: labelerSelfOcclusionDistance);
+            var frameSize = 1024;
+            var texture = new RenderTexture(frameSize, frameSize, 16);
-            }, texture, defaultSelfOcclusionDistance: labelerSelfOcclusionDistance);
+            }, texture);
             var camComponent = cam.GetComponent<Camera>();
             camComponent.fieldOfView = args.cameraFieldOfView;
             camComponent.farClipPlane = 100f;
             (Vector3 objectScale, Quaternion rotation, float checkDistance, Vector3 pointLocalPosition, bool expectOccluded) args)
         {
             var incoming = new List<List<KeypointLabeler.KeypointEntry>>();
-            var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate");
+            var labelerSelfOcclusionDistance = 0.5f;
+            var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate", labelerSelfOcclusionDistance);
-            var labelerSelfOcclusionDistance = 0.5f;
-            }, texture, defaultSelfOcclusionDistance: labelerSelfOcclusionDistance);
+            }, texture);
             var camComponent = cam.GetComponent<Camera>();
             camComponent.orthographic = true;
             camComponent.orthographicSize = 100f;
         public IEnumerator ManyObjects_LabelsCorrectly()
         {
             var incoming = new List<List<KeypointLabeler.KeypointEntry>>();
-            var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate");
+            var labelerSelfOcclusionDistance = 0.5f;
+            var template = CreateTestTemplate(Guid.NewGuid(), "TestTemplate", selfOcclusionDistance: labelerSelfOcclusionDistance);
-            var labelerSelfOcclusionDistance = 0.5f;
-            }, texture, defaultSelfOcclusionDistance: labelerSelfOcclusionDistance);
+            }, texture);
 
             void PlaceObjects(Rect rect, float z, Vector2Int count)
             {