Finishing fix for out of sync data

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

 
         [SuppressMessage("ReSharper", "InconsistentNaming")]
         [SuppressMessage("ReSharper", "NotAccessedField.Local")]
-        struct BoundingBoxValue
+        internal struct BoundingBoxValue
         {
             public int label_id;
             public string label_name;
         [FormerlySerializedAs("labelingConfiguration")]
         public IdLabelConfig idLabelConfig;
 
-        Dictionary<int, AsyncAnnotation> m_AsyncAnnotations;
+        Dictionary<int, (AsyncAnnotation annotation, LabelEntryMatchCache labelEntryMatchCache)> m_AsyncData;
         AnnotationDefinition m_BoundingBoxAnnotationDefinition;
         List<BoundingBoxValue> m_BoundingBoxValues;
 
         /// <inheritdoc/>
         protected override bool supportsVisualization => true;
 
+
+
+        /// <summary>
+        /// Event information for <see cref="BoundingBox2DLabeler.BoundingBoxesCalculated"/>
+        /// </summary>
+        internal struct BoundingBoxesCalculatedEventArgs
+        {
+            /// <summary>
+            /// The <see cref="Time.frameCount"/> on which the data was derived. This may be multiple frames in the past.
+            /// </summary>
+            public int frameCount;
+            /// <summary>
+            /// Bounding boxes.
+            /// </summary>
+            public IEnumerable<BoundingBoxValue> data;
+        }
+
+        /// <summary>
+        /// Event which is called each frame a semantic segmentation image is read back from the GPU.
+        /// </summary>
+        internal event Action<BoundingBoxesCalculatedEventArgs> boundingBoxesCalculated;
+
         /// <inheritdoc/>
         protected override void Setup()
         {
-            m_AsyncAnnotations = new Dictionary<int, AsyncAnnotation>();
+            m_AsyncData = new Dictionary<int, (AsyncAnnotation annotation, LabelEntryMatchCache labelEntryMatchCache)>();
             m_BoundingBoxValues = new List<BoundingBoxValue>();
 
             m_BoundingBoxAnnotationDefinition = DatasetCapture.RegisterAnnotationDefinition("bounding box", idLabelConfig.GetAnnotationSpecification(),
         /// <inheritdoc/>
         protected override void OnBeginRendering(ScriptableRenderContext scriptableRenderContext)
         {
-            m_AsyncAnnotations[Time.frameCount] = perceptionCamera.SensorHandle.ReportAnnotationAsync(m_BoundingBoxAnnotationDefinition);
+            m_AsyncData[Time.frameCount] =
+                (perceptionCamera.SensorHandle.ReportAnnotationAsync(m_BoundingBoxAnnotationDefinition),
+                 idLabelConfig.CreateLabelEntryMatchCache(Allocator.TempJob));
-            if (!m_AsyncAnnotations.TryGetValue(frameCount, out var asyncAnnotation))
+            if (!m_AsyncData.TryGetValue(frameCount, out var asyncData))
-            m_AsyncAnnotations.Remove(frameCount);
+            m_AsyncData.Remove(frameCount);
             using (s_BoundingBoxCallback.Auto())
             {
                 m_BoundingBoxValues.Clear();
-                    if (!idLabelConfig.TryGetLabelEntryFromInstanceId(objectInfo.instanceId, out var labelEntry))
+                    if (!asyncData.labelEntryMatchCache.TryGetLabelEntryFromInstanceId(objectInfo.instanceId, out var labelEntry, out _))
                         continue;
 
                     m_BoundingBoxValues.Add(new BoundingBoxValue
                 if (!CaptureOptions.useAsyncReadbackIfSupported && frameCount != Time.frameCount)
                     Debug.LogWarning("Not on current frame: " + frameCount + "(" + Time.frameCount + ")");
 
-                asyncAnnotation.ReportValues(m_BoundingBoxValues);
+                boundingBoxesCalculated?.Invoke(new BoundingBoxesCalculatedEventArgs()
+                {
+                    data = m_BoundingBoxValues,
+                    frameCount = frameCount
+                });
+                asyncData.annotation.ReportValues(m_BoundingBoxValues);
+                asyncData.labelEntryMatchCache.Dispose();
             }
         }
 

com.unity.perception/Runtime/GroundTruth/Labeling/IdLabelConfig.cs

         /// <returns>True if a labelId is found for the given instanceId.</returns>
         public bool TryGetLabelEntryFromInstanceId(uint instanceId, out IdLabelEntry labelEntry, out int index)
         {
-            if (m_LabelEntryMatchCache == null)
-                m_LabelEntryMatchCache = new LabelEntryMatchCache(this);
+            EnsureInitLabelEntryMatchCache();
+        }
+
+        private void EnsureInitLabelEntryMatchCache()
+        {
+            if (m_LabelEntryMatchCache == null)
+                m_LabelEntryMatchCache = new LabelEntryMatchCache(this, Allocator.Persistent);
         }
 
         /// <inheritdoc/>
             }).ToArray();
         }
 
-        public static IdLabelMap GetIdLabelCache(Allocator allocator = Allocator.Temp)
-        {
-            return new IdLabelMap();
-        }
-    }
-
-    public struct IdLabelMap : IDisposable
-    {
-        public void Dispose()
+        /// <summary>
+        /// Creates a LabelEntryMatchCache from the currently registered labeled objects, which can be used to look up
+        /// labeling information in future frames, even after the objects have been destroyed. Due to timing of labeled
+        /// object registration, if this is called during or before LateUpdate, this cache may become invalid.
+        ///
+        /// It is recommended to only use this method in rendering, as the cache is guaranteed to be in its final state
+        /// for ground truth generation.
+        /// </summary>
+        /// <param name="allocator">The allocator for creating the cache.</param>
+        /// <returns>The created cache.</returns>
+        public LabelEntryMatchCache CreateLabelEntryMatchCache(Allocator allocator)
+            EnsureInitLabelEntryMatchCache();
+            return m_LabelEntryMatchCache.CloneCurrentState(allocator);
         }
     }
 }

com.unity.perception/Runtime/GroundTruth/Labeling/LabelEntryMatchCache.cs

 {
     /// <summary>
     /// Cache of instance id -> label entry index for a LabelConfig. This is not well optimized and is the source of
-    /// a known memory leak for apps that create new instances frequently
+    /// a known memory leak for apps that create new instances frequently.
-    class LabelEntryMatchCache : IGroundTruthGenerator, IDisposable
+    public class LabelEntryMatchCache : IGroundTruthGenerator, IDisposable
+        private bool m_ReceiveUpdates;
-        public LabelEntryMatchCache(IdLabelConfig idLabelConfig)
+        internal LabelEntryMatchCache(IdLabelConfig idLabelConfig, Allocator allocator = Allocator.Persistent, bool receiveUpdates = true)
-            m_InstanceIdToLabelEntryIndexLookup = new NativeList<ushort>(k_StartingObjectCount, Allocator.Persistent);
-            LabelManager.singleton.Activate(this);
+            m_InstanceIdToLabelEntryIndexLookup = new NativeList<ushort>(k_StartingObjectCount, allocator);
+            m_ReceiveUpdates = receiveUpdates;
+            if (receiveUpdates)
+                LabelManager.singleton.Activate(this);
+        }
+
+        private LabelEntryMatchCache(LabelEntryMatchCache labelEntryMatchCache, Allocator allocator)
+        {
+            m_IdLabelConfig = labelEntryMatchCache.m_IdLabelConfig;
+            m_InstanceIdToLabelEntryIndexLookup = new NativeList<ushort>(labelEntryMatchCache.m_InstanceIdToLabelEntryIndexLookup.Length, allocator);
+            m_InstanceIdToLabelEntryIndexLookup.AddRange(labelEntryMatchCache.m_InstanceIdToLabelEntryIndexLookup.AsArray());
+            m_ReceiveUpdates = false;
+        /// <summary>
+        /// Retrieves the label entry for the given instance id.
+        /// </summary>
+        /// <param name="instanceId">The instance id to look up</param>
+        /// <param name="labelEntry">The <see cref="IdLabelEntry"/> of the match if found. Otherwise returns <code>default(IdlabelEntry)</code>.</param>
+        /// <param name="index">The index of the matched <see cref="IdLabelEntry"/> in the <see cref="IdLabelConfig"/> if found. Otherwise returns -1.</param>
+        /// <returns>True if a the instance id was found in the cache. </returns>
         public bool TryGetLabelEntryFromInstanceId(uint instanceId, out IdLabelEntry labelEntry, out int index)
         {
             labelEntry = default;
             }
         }
 
+        /// <inheritdoc/>
-            LabelManager.singleton.Deactivate(this);
+            if (m_ReceiveUpdates)
+                LabelManager.singleton.Deactivate(this);
+
+        }
+
+        internal LabelEntryMatchCache CloneCurrentState(Allocator allocator)
+        {
+            var clone = new LabelEntryMatchCache(this, Allocator.Persistent);
+            return clone;
         }
     }
 }

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

         void CaptureInstanceSegmentation(ScriptableRenderContext scriptableRenderContext)
         {
             //var myCache = new List<IdLabelEntry>();
-            var cache = IdLabelConfig.GetIdLabelCache();
             var width = m_InstanceSegmentationTexture.width;
 
             m_InstanceSegmentationReader.Capture(scriptableRenderContext,  (frameCount, data, renderTexture) =>
                 if(RenderedObjectInfosCalculated != null)
                 {
                     m_RenderedObjectInfoGenerator.Compute(data, width,
-                        BoundingBoxOrigin.TopLeft, cache, out var renderedObjectInfos, Allocator.Temp);
+                        BoundingBoxOrigin.TopLeft, out var renderedObjectInfos, Allocator.Temp);
-
-                cache.Dispose();
             });
         }
 

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

         /// <param name="boundingBoxOrigin">Whether bounding boxes should be top-left or bottom-right-based.</param>
         /// <param name="renderedObjectInfos">When this method returns, filled with RenderedObjectInfo entries for each object visible in the frame.</param>
         /// <param name="allocator">The allocator to use for allocating renderedObjectInfos and perLabelEntryObjectCount.</param>
-        public void Compute(NativeArray<Color32> instanceSegmentationRawData, int stride, BoundingBoxOrigin boundingBoxOrigin, IdLabelMap idLabelMap, out NativeArray<RenderedObjectInfo> renderedObjectInfos, Allocator allocator)
+        public void Compute(NativeArray<Color32> instanceSegmentationRawData, int stride, BoundingBoxOrigin boundingBoxOrigin, out NativeArray<RenderedObjectInfo> renderedObjectInfos, Allocator allocator)
         {
             const int jobCount = 24;
             var height = instanceSegmentationRawData.Length / stride;

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

             Object.DestroyImmediate(@object);
             m_ObjectsToDestroy.Remove(@object);
         }
+
+        public GameObject SetupCamera(Action<PerceptionCamera> initPerceptionCamera, bool activate = true)
+        {
+            var cameraObject = new GameObject();
+            cameraObject.SetActive(false);
+            var camera = cameraObject.AddComponent<Camera>();
+            camera.orthographic = true;
+            camera.orthographicSize = 1;
+
+            var perceptionCamera = cameraObject.AddComponent<PerceptionCamera>();
+            perceptionCamera.captureRgbImages = false;
+            initPerceptionCamera?.Invoke(perceptionCamera);
+
+            if (activate)
+                cameraObject.SetActive(true);
+
+            AddTestObjectForCleanup(cameraObject);
+            return cameraObject;
+        }
     }
 }

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

 using System.Collections;
 using NUnit.Framework;
+using Unity.Collections;
 using UnityEngine;
 using UnityEngine.Perception.GroundTruth;
 using UnityEngine.TestTools;
         public void TryGet_ReturnsFalse_ForInvalidInstanceId()
         {
             var config = ScriptableObject.CreateInstance<IdLabelConfig>();
-            using (var cache = new LabelEntryMatchCache(config))
+            using (var cache = new LabelEntryMatchCache(config, Allocator.Persistent))
             {
                 Assert.IsFalse(cache.TryGetLabelEntryFromInstanceId(100, out var labelEntry, out var index));
                 Assert.AreEqual(-1, index);
                     label = label
                 },
             });
-            using (var cache = new LabelEntryMatchCache(config))
+            using (var cache = new LabelEntryMatchCache(config, Allocator.Persistent))
             {
                 //allow label to be registered
                 yield return null;
             var labeledPlane = TestHelper.CreateLabeledPlane(label: label);
             AddTestObjectForCleanup(labeledPlane);
             var config = ScriptableObject.CreateInstance<IdLabelConfig>();
-            using (var cache = new LabelEntryMatchCache(config))
+            using (var cache = new LabelEntryMatchCache(config, Allocator.Persistent))
             {
                 //allow label to be registered
                 yield return null;
                     label = label
                 },
             });
-            using (var cache = new LabelEntryMatchCache(config))
+            using (var cache = new LabelEntryMatchCache(config, Allocator.Persistent))
             {
                 //allow label to be registered
                 yield return null;
                 },
             });
 
-            using (var cache = new LabelEntryMatchCache(config))
+            using (var cache = new LabelEntryMatchCache(config, Allocator.Persistent))
             {
                 //allow label to be registered
                 yield return null;

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

             });
             return labelingConfiguration;
         }
-
-        void SetupCamera(Action<PerceptionCamera> initPerceptionCamera, bool activate = true)
-        {
-            var cameraObject = new GameObject();
-            cameraObject.SetActive(false);
-            var camera = cameraObject.AddComponent<Camera>();
-            camera.orthographic = true;
-            camera.orthographicSize = 1;
-
-            var perceptionCamera = cameraObject.AddComponent<PerceptionCamera>();
-            perceptionCamera.captureRgbImages = false;
-            initPerceptionCamera?.Invoke(perceptionCamera);
-
-            if (activate)
-                cameraObject.SetActive(true);
-
-            AddTestObjectForCleanup(cameraObject);
-        }
     }
 }

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

 using System;
 using System.Collections;
 using System.Collections.Generic;
+using System.Linq;
+using Object = UnityEngine.Object;
 
 namespace GroundTruthTests
 {
 
             var dataNativeArray = new NativeArray<Color32>(producesCorrectObjectInfoData.data, Allocator.Persistent);
 
-            var idLabelMap = IdLabelConfig.GetIdLabelCache();
-            renderedObjectInfoGenerator.Compute(dataNativeArray, producesCorrectObjectInfoData.stride, producesCorrectObjectInfoData.boundingBoxOrigin, idLabelMap, out var boundingBoxes, Allocator.Temp);
+            var cache = labelingConfiguration.CreateLabelEntryMatchCache(Allocator.Persistent);
+            renderedObjectInfoGenerator.Compute(dataNativeArray, producesCorrectObjectInfoData.stride, producesCorrectObjectInfoData.boundingBoxOrigin, out var boundingBoxes, Allocator.Temp);
-            idLabelMap.Dispose();
+            cache.Dispose();
+        }
+
+
+        [UnityTest]
+        public IEnumerator LabelsCorrectWhenIdsReset()
+        {
+            int timesInfoReceived = 0;
+            Dictionary<int, int> expectedLabelIdAtFrame = null;
+
+            //TestHelper.LoadAndStartRenderDocCapture(out var gameView);
+
+            void OnBoundingBoxesReceived(BoundingBox2DLabeler.BoundingBoxesCalculatedEventArgs eventArgs)
+            {
+                if (expectedLabelIdAtFrame == null || !expectedLabelIdAtFrame.ContainsKey(eventArgs.frameCount)) return;
+
+                timesInfoReceived++;
+
+                Debug.Log($"Bounding boxes received. FrameCount: {eventArgs.frameCount}");
+
+                Assert.AreEqual(1, eventArgs.data.Count());
+                Assert.AreEqual(expectedLabelIdAtFrame[eventArgs.frameCount], eventArgs.data.First().label_id);
+            }
+
+            var idLabelConfig = ScriptableObject.CreateInstance<IdLabelConfig>();
+            idLabelConfig.Init(new []
+            {
+                new IdLabelEntry()
+                {
+                    id = 1,
+                    label = "label1"
+                },
+                new IdLabelEntry()
+                {
+                    id = 2,
+                    label = "label2"
+                },
+                new IdLabelEntry()
+                {
+                    id = 3,
+                    label = "label3"
+                },
+            });
+            AddTestObjectForCleanup(idLabelConfig);
+
+            var cameraObject = SetupCameraBoundingBox2D(OnBoundingBoxesReceived, idLabelConfig);
+
+            expectedLabelIdAtFrame = new Dictionary<int, int>
+            {
+                {Time.frameCount    , 1},
+                {Time.frameCount + 1, 2},
+                {Time.frameCount + 2, 3}
+            };
+            GameObject planeObject;
+
+            //Put a plane in front of the camera
+            planeObject = TestHelper.CreateLabeledPlane(label: "label1");
+            yield return null;
+
+            //UnityEditorInternal.RenderDoc.EndCaptureRenderDoc(gameView);
+            Object.DestroyImmediate(planeObject);
+            planeObject = TestHelper.CreateLabeledPlane(label: "label2");
+
+            //TestHelper.LoadAndStartRenderDocCapture(out gameView);
+            yield return null;
+
+            //UnityEditorInternal.RenderDoc.EndCaptureRenderDoc(gameView);
+            Object.DestroyImmediate(planeObject);
+            planeObject = TestHelper.CreateLabeledPlane(label: "label3");
+            yield return null;
+            Object.DestroyImmediate(planeObject);
+            yield return null;
+
+            //destroy the object to force all pending segmented image readbacks to finish and events to be fired.
+            DestroyTestObject(cameraObject);
+
+            Assert.AreEqual(3, timesInfoReceived);
+        }
+
+        private GameObject SetupCameraBoundingBox2D(Action<BoundingBox2DLabeler.BoundingBoxesCalculatedEventArgs> onBoundingBoxesCalculated, IdLabelConfig idLabelConfig)
+        {
+            var cameraObject = SetupCamera(camera =>
+            {
+                camera.showVisualizations = false;
+                var boundingBox2DLabeler = new BoundingBox2DLabeler(idLabelConfig);
+                boundingBox2DLabeler.boundingBoxesCalculated += onBoundingBoxesCalculated;
+                camera.AddLabeler(boundingBox2DLabeler);
+            });
+            return cameraObject;
         }
     }
 }

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

                 if (frames < 10)
                     return;
 
-                var idLabelMap = IdLabelConfig.GetIdLabelCache();
-
                 // Calculate the bounding box
                 if (fLensDistortionEnabled == false)
                 {
-                    renderedObjectInfoGenerator.Compute(data, tex.width, BoundingBoxOrigin.TopLeft, idLabelMap, out var boundingBoxes, Allocator.Temp);
+                    renderedObjectInfoGenerator.Compute(data, tex.width, BoundingBoxOrigin.TopLeft, out var boundingBoxes, Allocator.Temp);
 
                     boundingBoxWithoutLensDistortion = boundingBoxes[0].boundingBox;
 
                 {
                     var renderedObjectInfoGenerator = new RenderedObjectInfoGenerator();
 
-                    renderedObjectInfoGenerator.Compute(data, tex.width, BoundingBoxOrigin.TopLeft, idLabelMap, out var boundingBoxes, Allocator.Temp);
+                    renderedObjectInfoGenerator.Compute(data, tex.width, BoundingBoxOrigin.TopLeft, out var boundingBoxes, Allocator.Temp);
 
                     boundingBoxWithLensDistortion = boundingBoxes[0].boundingBox;
 
                     fDone = true;
                 }
-                idLabelMap.Dispose();
             }
 
             cameraObject = SetupCamera(out perceptionCamera, false);