您最多选择25个主题 主题必须以中文或者字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符

106 行
4.7 KiB

using System.Collections.Generic;
using System.Threading;
using Unity.Entities;
namespace UnityEngine.Perception.GroundTruth
{
struct IdAssignmentParameters : IComponentData
{
public uint idStart;
public uint idStep;
}
/// <summary>
/// System which notifies the registered <see cref="IGroundTruthGenerator"/> about <see cref="Labeling"/> additions.
/// </summary>
public class GroundTruthLabelSetupSystem : ComponentSystem
{
List<IGroundTruthGenerator> m_ActiveGenerators = new List<IGroundTruthGenerator>();
ThreadLocal<MaterialPropertyBlock> m_MaterialPropertyBlocks = new ThreadLocal<MaterialPropertyBlock>();
int m_CurrentObjectIndex = -1;
/// <inheritdoc/>
protected override void OnCreate()
{
//These are here to inform the system runner the queries we are interested in. Without these calls, OnUpdate() might not be called
GetEntityQuery(ComponentType.Exclude<GroundTruthInfo>(), ComponentType.ReadOnly<Labeling>());
GetEntityQuery(ComponentType.ReadOnly<GroundTruthInfo>(), ComponentType.ReadOnly<Labeling>());
}
/// <inheritdoc/>
protected override void OnUpdate()
{
var entityQuery = Entities.WithAll<IdAssignmentParameters>().ToEntityQuery();
IdAssignmentParameters idAssignmentParameters;
if (entityQuery.CalculateEntityCount() == 1)
idAssignmentParameters = entityQuery.GetSingleton<IdAssignmentParameters>();
else
idAssignmentParameters = new IdAssignmentParameters {idStart = 1, idStep = 1};
var entityCount = Entities.WithAll<Labeling, GroundTruthInfo>().ToEntityQuery().CalculateEntityCount();
if (entityCount == 0)
m_CurrentObjectIndex = -1;
Entities.WithNone<GroundTruthInfo>().ForEach((Entity e, Labeling labeling) =>
{
var objectIndex = (uint)Interlocked.Increment(ref m_CurrentObjectIndex);
var instanceId = idAssignmentParameters.idStart + objectIndex * idAssignmentParameters.idStep;
var gameObject = labeling.gameObject;
if (!m_MaterialPropertyBlocks.IsValueCreated)
m_MaterialPropertyBlocks.Value = new MaterialPropertyBlock();
InitGameObjectRecursive(gameObject, m_MaterialPropertyBlocks.Value, labeling, instanceId);
EntityManager.AddComponentData(e, new GroundTruthInfo
{
instanceId = instanceId
});
});
}
void InitGameObjectRecursive(GameObject gameObject, MaterialPropertyBlock mpb, Labeling labeling, uint instanceId)
{
var meshRenderer = gameObject.GetComponent<MeshRenderer>();
if (meshRenderer != null)
{
meshRenderer.GetPropertyBlock(mpb);
foreach (var pass in m_ActiveGenerators)
pass.SetupMaterialProperties(mpb, meshRenderer, labeling, instanceId);
meshRenderer.SetPropertyBlock(mpb);
}
for (var i = 0; i < gameObject.transform.childCount; i++)
{
var child = gameObject.transform.GetChild(i).gameObject;
if (child.GetComponent<Labeling>() != null)
continue;
InitGameObjectRecursive(child, mpb, labeling, instanceId);
}
}
/// <summary>
/// Activates the given <see cref="IGroundTruthGenerator"/>. <see cref="IGroundTruthGenerator.SetupMaterialProperties"/>
/// will be called for all <see cref="MeshRenderer"/> instances under each object containing a <see cref="Labeling"/> component.
/// </summary>
/// <param name="generator">The generator to register</param>
public void Activate(IGroundTruthGenerator generator)
{
m_ActiveGenerators.Add(generator);
Entities.ForEach((Labeling labeling, ref GroundTruthInfo info) =>
{
var gameObject = labeling.gameObject;
InitGameObjectRecursive(gameObject, m_MaterialPropertyBlocks.Value, labeling, info.instanceId);
});
}
/// <summary>
/// Deactivates the given <see cref="IGroundTruthGenerator"/>. It will no longer receive calls when <see cref="Labeling"/> instances are created.
/// </summary>
/// <param name="generator">The generator to deactivate</param>
/// <returns>True if the <see cref="generator"/> was successfully removed. False if <see cref="generator"/> was not active.</returns>
public bool Deactivate(IGroundTruthGenerator generator)
{
return m_ActiveGenerators.Remove(generator);
}
}
}