diff --git a/SynthDet/Assets/Scripts/Runtime/SynthDetScenario.cs b/SynthDet/Assets/Scripts/Runtime/SynthDetScenario.cs index 90fa1285..d4cf26f9 100644 --- a/SynthDet/Assets/Scripts/Runtime/SynthDetScenario.cs +++ b/SynthDet/Assets/Scripts/Runtime/SynthDetScenario.cs @@ -1,4 +1,5 @@ -using System.Collections.Generic; +using System; +using System.Collections.Generic; using System.Linq; using SynthDet.Randomizers; using UnityEngine; @@ -8,7 +9,7 @@ using UnityEngine.Perception.GroundTruth; using UnityEngine.Perception.Randomization.Parameters; using UnityEngine.Perception.Randomization.Scenarios; using UnityEngine.ResourceManagement.AsyncOperations; -using UnityEngine.ResourceManagement.ResourceLocations; +using Debug = UnityEngine.Debug; namespace SynthDet.Scenarios { @@ -19,142 +20,102 @@ namespace SynthDet.Scenarios [AddComponentMenu("SynthDet/SynthDet Scenario")] public class SynthDetScenario : FixedLengthScenario { - int m_NumPrefabsToLoad; - List m_Prefabs = new List(); - AssetLoadingStatus m_LoadingStatus = AssetLoadingStatus.LoadingCatalog; - AsyncOperationHandle[] m_CatalogHandles; - + static HashSet s_ColorsAlreadyUsed = new HashSet(); + static ColorRgbParameter s_SemanticColorParameter = new ColorRgbParameter(); + List m_LabelStringsForAutoLabelConfig = new List(); - - /// - protected override bool isScenarioReadyToStart - { - get - { - switch (m_LoadingStatus) - { - case AssetLoadingStatus.Complete: - return true; - case AssetLoadingStatus.LoadingCatalog: - { - foreach (var handle in m_CatalogHandles) - { - if (!handle.IsDone) - return false; - } - - foreach (var handle in m_CatalogHandles) - { - if (handle.Status != AsyncOperationStatus.Succeeded) - { - Debug.LogError("Catalog download failed"); - m_LoadingStatus = AssetLoadingStatus.Failed; - return false; - } - } - - LoadAndLabelPrefabs(); - m_LoadingStatus = AssetLoadingStatus.LoadingPrefabs; - return false; - } - case AssetLoadingStatus.LoadingPrefabs: - { - lock (m_Prefabs) - { - if (m_Prefabs.Count < m_NumPrefabsToLoad) - return false; - SetupLabelConfigs(); - var randomizer = GetRandomizer(); - m_Prefabs.Sort((prefab1, prefab2) => prefab1.name.CompareTo(prefab2.name)); - randomizer.prefabs = m_Prefabs.ToArray(); - m_LoadingStatus = AssetLoadingStatus.Complete; - return true; - } - } - case AssetLoadingStatus.Failed: - { - Quit(); - return false; - } - default: - return false; - } - } - } + + /// + /// Skip the first frame since the simulation capture package cannot capture it. + /// + protected override bool isScenarioReadyToStart => Time.frameCount != 1; protected override void OnAwake() { base.OnAwake(); - - Caching.ClearCache(); - - string InternalIdTransformFunc(IResourceLocation location) + try { - var internalId = m_BundleToUrlMap.ContainsKey(location.PrimaryKey) - ? m_BundleToUrlMap[location.PrimaryKey] - : location.InternalId; - - Debug.Log($"InternalIdTransformFunc. Old: {location.InternalId} New: {internalId}"); - return internalId; + LoadAssets(); + } + catch (Exception) + { + Quit(); + throw; + } + } + + void LoadAssets() + { + if (m_CatalogUrls.Count == 0) + throw new Exception("No content catalogs specified to load"); + + // Map bundle urls from app-param to addressables resource locations. + // This mapping enables dynamically assigned urls to be used for remote bundle locations. + Addressables.InternalIdTransformFunc = location => m_BundleToUrlMap.ContainsKey(location.PrimaryKey) + ? m_BundleToUrlMap[location.PrimaryKey] + : location.InternalId; + + var catalogs = LoadCatalogs(); + LoadAndLabelPrefabs(catalogs); + } + + List LoadCatalogs() + { + // Clear addressables cache to ensure that we are loading asset bundles from their remote locations + Caching.ClearCache(); + + // Begin loading the remote content catalogs included in the app-param + var catalogHandles = new List>(); + foreach (var url in m_CatalogUrls) + { + var catalogHandle = Addressables.LoadContentCatalogAsync(url); + catalogHandles.Add(catalogHandle); + } + + // Wait for catalogs to load + for (var i = 0; i < catalogHandles.Count; i++) + { + var handle = catalogHandles[i]; + handle.WaitForCompletion(); + if (handle.Status != AsyncOperationStatus.Succeeded) + throw new Exception($"Catalog failed to load from URL {m_CatalogUrls[i]}"); } - Addressables.InternalIdTransformFunc = InternalIdTransformFunc; - - // Addressables.InternalIdTransformFunc = TransformFunc; - - m_CatalogHandles = new AsyncOperationHandle[m_CatalogUrls.Count]; - for(var i = 0 ; i < m_CatalogUrls.Count; i++) - m_CatalogHandles[i] = Addressables.LoadContentCatalogAsync(m_CatalogUrls[i], false); + return catalogHandles.Select(handle => handle.Result).ToList(); } - string TransformFunc(IResourceLocation location) + void LoadAndLabelPrefabs(List catalogs) { - Debug.Log(location.InternalId); - if (location.InternalId.Contains("catalog") && location.InternalId.Contains(".json")) - return location.InternalId; - if (location.InternalId.Contains("__placeholder__") && location.InternalId.Contains("_assets_all_")) //"doorgroup6_assets_all_ef13245136c2b5bec4ea18ac92071dca.bundle")) - return location.InternalId.Replace("__placeholder__/foreground_group_assets_all_0c7829f95130438cde35283652f1fd3f.bundle", "https://storage.googleapis.com/addressables-synthdet/e2e/newe2e/foreground_group_assets_all_0c7829f95130438cde35283652f1fd3f.bundle"); - if (location.InternalId.Contains("__placeholder__") && location.InternalId.Contains("_unitybuiltinshaders_")) - //"doorgroup6_unitybuiltinshaders_fb3cc65dc055f6c5ef84d30de128788a.bundle")) - return location.InternalId.Replace("__placeholder__/foreground_group_unitybuiltinshaders_51bd82438b9ab720fa52c2dbd163e111.bundle","https://storage.googleapis.com/addressables-synthdet/e2e/newe2e/foreground_group_unitybuiltinshaders_51bd82438b9ab720fa52c2dbd163e111.bundle"); - return location.InternalId; - } - - void LoadAndLabelPrefabs() - { - foreach (var handle in m_CatalogHandles) + // Gather only the prefab keys from the remote catalogs + var prefabKeys = new List(); + foreach (var catalog in catalogs) { - foreach (var key in handle.Result.Keys) + foreach (var key in catalog.Keys) { - if (!key.ToString().Contains(".prefab")) - continue; - - m_NumPrefabsToLoad++; - Addressables.LoadAssetAsync(key).Completed += prefabHandle => - { - if (prefabHandle.Status == AsyncOperationStatus.Failed) - { - Debug.LogError($"Failed to load prefab from key '{key}'"); - m_LoadingStatus = AssetLoadingStatus.Failed; - return; - } - lock (m_Prefabs) - { - ConfigureLabeling(prefabHandle.Result); - prefabHandle.Result.layer = LayerMask.NameToLayer("Foreground"); - m_Prefabs.Add(prefabHandle.Result); - } - }; + var keyPath = key.ToString(); + if (keyPath.Contains(".prefab")) + prefabKeys.Add(keyPath); } } - } + + // Load all the prefabs from the remote catalogs, label them, and finally assign + // them to the ForegroundObjectPlacementRandomizer. + var handle = Addressables.LoadAssetsAsync( + prefabKeys, null, Addressables.MergeMode.Union); + handle.WaitForCompletion(); + if (handle.Status != AsyncOperationStatus.Succeeded) + throw new Exception("Prefabs failed to load from content catalogs"); - enum AssetLoadingStatus - { - Complete, - LoadingCatalog, - LoadingPrefabs, - Failed + // Sort prefabs and configure labeling + SetupLabelConfigs(); + var prefabsList = new List(handle.Result); + prefabsList.Sort((prefab1, prefab2) => prefab1.name.CompareTo(prefab2.name)); + foreach (var prefab in prefabsList) + ConfigureLabeling(prefab); + + // Inject the loaded prefabs into the ForegroundObjectPlacementRandomizer + var randomizer = GetRandomizer(); + randomizer.prefabs = prefabsList.ToArray(); } void ConfigureLabeling(GameObject gObj) @@ -225,13 +186,11 @@ namespace SynthDet.Scenarios labeler.Init(perceptionCamera); } } - - static HashSet s_ColorsAlreadyUsed = new HashSet(); - static ColorRgbParameter s_SemanticColorParameter = new ColorRgbParameter(); + static Color GetUniqueSemanticSegmentationColor() { var sampledColor = s_SemanticColorParameter.Sample(); - var maxTries = 1000; + const int maxTries = 1000; var count = 0; while (s_ColorsAlreadyUsed.Contains(sampledColor) && count <= maxTries)