您最多选择25个主题
主题必须以中文或者字母或数字开头，可以包含连字符 (-)，并且长度不得超过35个字符
884 行
36 KiB
884 行
36 KiB
using System;
|
|
using System.Collections.Generic;
|
|
using System.Linq;
|
|
using Unity.Collections;
|
|
using Unity.Collections.LowLevel.Unsafe;
|
|
using Unity.Mathematics;
|
|
using Unity.Profiling;
|
|
using Unity.Simulation;
|
|
using UnityEngine.Experimental.Rendering;
|
|
using UnityEngine.Rendering;
|
|
using UnityEngine.Serialization;
|
|
|
|
namespace UnityEngine.Perception.GroundTruth
|
|
{
|
|
/// <summary>
|
|
/// Produces keypoint annotations for a humanoid model. This labeler supports generic
|
|
/// <see cref="KeypointTemplate"/>. Template values are mapped to rigged
|
|
/// <see cref="Animator"/> <seealso cref="Avatar"/>. Custom joints can be
|
|
/// created by applying <see cref="JointLabel"/> to empty game objects at a body
|
|
/// part's location.
|
|
/// </summary>
|
|
[Serializable]
|
|
public sealed class KeypointLabeler : CameraLabeler
|
|
{
|
|
// Smaller texture sizes produce assertion failures in the engine
|
|
const int k_MinTextureWidth = 8;
|
|
|
|
static ProfilerMarker k_OnEndRenderingMarker = new ProfilerMarker($"KeypointLabeler OnEndRendering");
|
|
static ProfilerMarker k_OnVisualizeMarker = new ProfilerMarker($"KeypointLabeler OnVisualize");
|
|
|
|
/// <summary>
|
|
/// The active keypoint template. Required to annotate keypoint data.
|
|
/// </summary>
|
|
public KeypointTemplate activeTemplate;
|
|
|
|
/// <inheritdoc/>
|
|
public override string description
|
|
{
|
|
get => "Produces keypoint annotations for all visible labeled objects that have a humanoid animation avatar component.";
|
|
protected set { }
|
|
}
|
|
|
|
///<inheritdoc/>
|
|
protected override bool supportsVisualization => true;
|
|
|
|
// ReSharper disable MemberCanBePrivate.Global
|
|
/// <summary>
|
|
/// The GUID id to associate with the annotations produced by this labeler.
|
|
/// </summary>
|
|
public string annotationId = "8b3ef246-daa7-4dd5-a0e8-a943f6e7f8c2";
|
|
/// <summary>
|
|
/// The <see cref="IdLabelConfig"/> which associates objects with labels.
|
|
/// </summary>
|
|
public IdLabelConfig idLabelConfig;
|
|
|
|
public float visThickness = 6;
|
|
|
|
public bool drawBones = true;
|
|
|
|
/// <summary>
|
|
/// Controls which objects will have keypoints recorded in the dataset.
|
|
/// <see cref="KeypointObjectFilter"/>
|
|
/// </summary>
|
|
public KeypointObjectFilter objectFilter;
|
|
// ReSharper restore MemberCanBePrivate.Global
|
|
|
|
AnnotationDefinition m_AnnotationDefinition;
|
|
Texture2D m_MissingTexture;
|
|
Material m_MaterialDepthCheck;
|
|
Texture2D m_KeypointPositionsTexture;
|
|
Texture2D m_KeypointCheckDepthTexture;
|
|
|
|
struct FrameKeypointData
|
|
{
|
|
public AsyncAnnotation annotation;
|
|
public int pointsPerEntry;
|
|
public List<KeypointEntry> keypoints;
|
|
public bool isDepthCheckComplete;
|
|
public bool isInstanceSegmentationCheckComplete;
|
|
public NativeArray<RenderedObjectInfo> objectInfos;
|
|
}
|
|
|
|
Dictionary<int, FrameKeypointData> m_FrameKeypointData;
|
|
|
|
List<KeypointEntry> m_KeypointEntriesToReport;
|
|
|
|
int m_CurrentFrame;
|
|
|
|
/// <summary>
|
|
/// Action that gets triggered when a new frame of key points are computed.
|
|
/// </summary>
|
|
public event Action<int, List<KeypointEntry>> KeypointsComputed;
|
|
|
|
/// <summary>
|
|
/// Creates a new key point labeler. This constructor creates a labeler that
|
|
/// is not valid until a <see cref="IdLabelConfig"/> and <see cref="KeypointTemplate"/>
|
|
/// are assigned.
|
|
/// </summary>
|
|
public KeypointLabeler() { }
|
|
|
|
/// <summary>
|
|
/// Creates a new key point labeler.
|
|
/// </summary>
|
|
/// <param name="config">The Id label config for the labeler</param>
|
|
/// <param name="template">The active keypoint template</param>
|
|
public KeypointLabeler(IdLabelConfig config, KeypointTemplate template)
|
|
{
|
|
this.idLabelConfig = config;
|
|
this.activeTemplate = template;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Array of animation pose labels which map animation clip times to ground truth pose labels.
|
|
/// </summary>
|
|
public List<AnimationPoseConfig> animationPoseConfigs;
|
|
|
|
private ComputeShader m_KeypointDepthTestShader;
|
|
private RenderTexture m_ResultsBuffer;
|
|
private RenderTextureReader<Color32> m_DepthCheckReader;
|
|
|
|
|
|
/// <inheritdoc/>
|
|
protected override void Setup()
|
|
{
|
|
if (idLabelConfig == null)
|
|
throw new InvalidOperationException($"{nameof(KeypointLabeler)}'s idLabelConfig field must be assigned");
|
|
|
|
m_AnnotationDefinition = DatasetCapture.RegisterAnnotationDefinition("keypoints", new []{TemplateToJson(activeTemplate)},
|
|
"pixel coordinates of keypoints in a model, along with skeletal connectivity data", id: new Guid(annotationId));
|
|
|
|
// Texture to use in case the template does not contain a texture for the joints or the skeletal connections
|
|
m_MissingTexture = new Texture2D(1, 1);
|
|
|
|
m_KnownStatus = new Dictionary<uint, CachedData>();
|
|
|
|
m_FrameKeypointData = new Dictionary<int, FrameKeypointData>();
|
|
m_KeypointEntriesToReport = new List<KeypointEntry>();
|
|
m_CurrentFrame = 0;
|
|
m_KeypointDepthTestShader = (ComputeShader) Resources.Load("KeypointDepthTest");
|
|
|
|
var depthCheckShader = Shader.Find("Perception/KeypointDepthCheck");
|
|
|
|
var shaderVariantCollection = new ShaderVariantCollection();
|
|
m_MaterialDepthCheck = new Material(depthCheckShader);
|
|
string keyword;
|
|
if (SRPSupport.GetCurrentPipelineRenderingType() == RenderingPipelineType.HDRP)
|
|
keyword = "HDRP_ENABLED";
|
|
else
|
|
keyword = "HDRP_DISABLED";
|
|
|
|
m_MaterialDepthCheck.EnableKeyword(keyword);
|
|
|
|
shaderVariantCollection.Add(
|
|
new ShaderVariantCollection.ShaderVariant(depthCheckShader, PassType.ScriptableRenderPipeline, keyword));
|
|
shaderVariantCollection.WarmUp();
|
|
|
|
perceptionCamera.attachedCamera.depthTextureMode = DepthTextureMode.Depth;
|
|
#if URP_PRESENT
|
|
var cameraData = UnityEngine.Rendering.Universal.CameraExtensions.GetUniversalAdditionalCameraData(perceptionCamera.attachedCamera);
|
|
cameraData.requiresDepthOption = UnityEngine.Rendering.Universal.CameraOverrideOption.On;
|
|
cameraData.requiresDepthTexture = true;
|
|
#endif
|
|
|
|
perceptionCamera.InstanceSegmentationImageReadback += OnInstanceSegmentationImageReadback;
|
|
perceptionCamera.RenderedObjectInfosCalculated += OnRenderedObjectInfoReadback;
|
|
}
|
|
|
|
private void SetupDepthCheckBuffers(int size)
|
|
{
|
|
var textureDimensions = TextureDimensions(size);
|
|
if (m_ResultsBuffer != null &&
|
|
textureDimensions.x == m_ResultsBuffer.width &&
|
|
textureDimensions.y == m_ResultsBuffer.height)
|
|
return;
|
|
|
|
|
|
if (m_ResultsBuffer != null)
|
|
{
|
|
m_ResultsBuffer.Release();
|
|
m_DepthCheckReader.Dispose(false);
|
|
// Object.Destroy(m_KeypointPositionsTexture);
|
|
// Object.Destroy(m_KeypointCheckDepthTexture);
|
|
}
|
|
|
|
m_KeypointPositionsTexture = new Texture2D(textureDimensions.x, textureDimensions.y, GraphicsFormat.R16G16_SFloat, TextureCreationFlags.None);
|
|
m_KeypointCheckDepthTexture = new Texture2D(textureDimensions.x, textureDimensions.y, GraphicsFormat.R16_SFloat, TextureCreationFlags.None);
|
|
|
|
m_ResultsBuffer = new RenderTexture(textureDimensions.x, textureDimensions.y, 0, GraphicsFormat.R8G8B8A8_UNorm);
|
|
m_DepthCheckReader = new RenderTextureReader<Color32>(m_ResultsBuffer);
|
|
}
|
|
|
|
bool AreEqual(Color32 lhs, Color32 rhs)
|
|
{
|
|
return lhs.r == rhs.r && lhs.g == rhs.g && lhs.b == rhs.b && lhs.a == rhs.a;
|
|
}
|
|
|
|
bool PixelOnScreen(int2 pixelLocation, (int x, int y) dimensions)
|
|
{
|
|
return pixelLocation.x >= 0 && pixelLocation.x < dimensions.x && pixelLocation.y >= 0 && pixelLocation.y < dimensions.y;
|
|
}
|
|
|
|
bool PixelsMatch(int x, int y, Color32 idColor, (int x, int y) dimensions, NativeArray<Color32> data)
|
|
{
|
|
var h = dimensions.y - 1 - y;
|
|
var pixelColor = data[h * dimensions.x + x];
|
|
return AreEqual(pixelColor, idColor);
|
|
}
|
|
|
|
static int s_PixelTolerance = 1;
|
|
|
|
// Determine the state of a keypoint. A keypoint is considered visible (state = 2) if it is on screen and not occluded
|
|
// by itself or another object. Self-occlusion has already been checked, so the input keypoint may be state 2, 1, or 0.
|
|
// We determine if a point is occluded by other objects is by checking the pixel location of the keypoint
|
|
// against the instance segmentation mask for the frame. The instance segmentation mask provides the instance id of the
|
|
// visible object at a pixel location. Which means, if the keypoint does not match the visible pixel, then another
|
|
// object is in front of the keypoint occluding it from view. An important note here is that the keypoint is an infintely small
|
|
// point in space, which can lead to false negatives due to rounding issues if the keypoint is on the edge of an object or very
|
|
// close to the edge of the screen. Because of this we will test not only the keypoint pixel, but also the immediate surrounding
|
|
// pixels to determine if the pixel is really visible. This method returns 1 if the pixel is not visible but on screen, and 0
|
|
// if the pixel is off of the screen (taken the tolerance into account).
|
|
int DetermineKeypointState(Keypoint keypoint, Color32 instanceIdColor, (int x, int y) dimensions, NativeArray<Color32> data)
|
|
{
|
|
if (keypoint.state == 0) return 0;
|
|
|
|
var pixelLocation = PixelLocationFromScreenPoint(keypoint);
|
|
|
|
if (!PixelOnScreen(pixelLocation, dimensions))
|
|
return 0;
|
|
|
|
var pixelMatched = false;
|
|
|
|
for (var y = pixelLocation.y - s_PixelTolerance; y <= pixelLocation.y + s_PixelTolerance; y++)
|
|
{
|
|
for (var x = pixelLocation.x - s_PixelTolerance; x <= pixelLocation.x + s_PixelTolerance; x++)
|
|
{
|
|
if (!PixelOnScreen(new int2(x, y), dimensions)) continue;
|
|
|
|
pixelMatched = true;
|
|
if (PixelsMatch(x, y, instanceIdColor, dimensions, data))
|
|
{
|
|
return keypoint.state;
|
|
}
|
|
}
|
|
}
|
|
|
|
return pixelMatched ? 1 : 0;
|
|
}
|
|
|
|
private static int2 PixelLocationFromScreenPoint(Keypoint keypoint)
|
|
{
|
|
var centerX = Mathf.FloorToInt(keypoint.x);
|
|
var centerY = Mathf.FloorToInt(keypoint.y);
|
|
int2 pixelLocation = new int2(centerX, centerY);
|
|
return pixelLocation;
|
|
}
|
|
|
|
void OnInstanceSegmentationImageReadback(int frameCount, NativeArray<Color32> data, RenderTexture renderTexture)
|
|
{
|
|
if (!m_FrameKeypointData.TryGetValue(frameCount, out var frameKeypointData))
|
|
return;
|
|
|
|
var dimensions = (renderTexture.width, renderTexture.height);
|
|
|
|
foreach (var keypointEntry in frameKeypointData.keypoints)
|
|
{
|
|
if (InstanceIdToColorMapping.TryGetColorFromInstanceId(keypointEntry.instance_id, out var idColor))
|
|
{
|
|
for (var i = 0; i < keypointEntry.keypoints.Length; i++)
|
|
{
|
|
var keypoint = keypointEntry.keypoints[i];
|
|
keypoint.state = DetermineKeypointState(keypoint, idColor, dimensions, data);
|
|
|
|
if (keypoint.state == 0)
|
|
{
|
|
keypoint.x = 0;
|
|
keypoint.y = 0;
|
|
}
|
|
else
|
|
{
|
|
keypoint.x = math.clamp(keypoint.x, 0, dimensions.width - .001f);
|
|
keypoint.y = math.clamp(keypoint.y, 0, dimensions.height - .001f);
|
|
}
|
|
|
|
keypointEntry.keypoints[i] = keypoint;
|
|
}
|
|
}
|
|
}
|
|
|
|
frameKeypointData.isInstanceSegmentationCheckComplete = true;
|
|
m_FrameKeypointData[frameCount] = frameKeypointData;
|
|
ReportIfComplete(frameCount, frameKeypointData);
|
|
}
|
|
|
|
private void OnRenderedObjectInfoReadback(int frameCount, NativeArray<RenderedObjectInfo> objectInfos)
|
|
{
|
|
if (!m_FrameKeypointData.TryGetValue(frameCount, out var frameKeypointData))
|
|
return;
|
|
|
|
frameKeypointData.objectInfos = new NativeArray<RenderedObjectInfo>(objectInfos, Allocator.Persistent);
|
|
m_FrameKeypointData[frameCount] = frameKeypointData;
|
|
ReportIfComplete(frameCount, frameKeypointData);
|
|
}
|
|
|
|
private void ReportIfComplete(int frameCount, FrameKeypointData frameKeypointData)
|
|
{
|
|
if (!frameKeypointData.isInstanceSegmentationCheckComplete || !frameKeypointData.isDepthCheckComplete || !frameKeypointData.objectInfos.IsCreated)
|
|
return;
|
|
|
|
m_KeypointEntriesToReport.Clear();
|
|
|
|
//filter out objects that are not visible
|
|
foreach (var entry in frameKeypointData.keypoints)
|
|
{
|
|
var include = false;
|
|
if (objectFilter == KeypointObjectFilter.All)
|
|
include = true;
|
|
else
|
|
{
|
|
foreach (var objectInfo in frameKeypointData.objectInfos)
|
|
{
|
|
if (entry.instance_id == objectInfo.instanceId)
|
|
{
|
|
include = true;
|
|
break;
|
|
}
|
|
}
|
|
|
|
if (!include && objectFilter == KeypointObjectFilter.VisibleAndOccluded)
|
|
include = entry.keypoints.Any(k => k.state == 1);
|
|
}
|
|
|
|
if (include)
|
|
m_KeypointEntriesToReport.Add(entry);
|
|
}
|
|
m_FrameKeypointData.Remove(frameCount);
|
|
KeypointsComputed?.Invoke(frameCount, m_KeypointEntriesToReport);
|
|
frameKeypointData.annotation.ReportValues(m_KeypointEntriesToReport);
|
|
frameKeypointData.objectInfos.Dispose();
|
|
}
|
|
|
|
/// <param name="scriptableRenderContext"></param>
|
|
/// <inheritdoc/>
|
|
protected override void OnEndRendering(ScriptableRenderContext scriptableRenderContext)
|
|
{
|
|
using (k_OnEndRenderingMarker.Auto())
|
|
{
|
|
m_CurrentFrame = Time.frameCount;
|
|
|
|
var annotation = perceptionCamera.SensorHandle.ReportAnnotationAsync(m_AnnotationDefinition);
|
|
var keypointEntries = new List<KeypointEntry>();
|
|
var checkLocations = new NativeList<float3>(512, Allocator.Persistent);
|
|
|
|
foreach (var label in LabelManager.singleton.registeredLabels)
|
|
ProcessLabel(label, keypointEntries, checkLocations);
|
|
|
|
m_FrameKeypointData[m_CurrentFrame] = new FrameKeypointData
|
|
{
|
|
annotation = annotation,
|
|
keypoints = keypointEntries,
|
|
pointsPerEntry = activeTemplate.keypoints.Length
|
|
};
|
|
|
|
if (keypointEntries.Count != 0)
|
|
DoDepthCheck(scriptableRenderContext, keypointEntries, checkLocations);
|
|
else
|
|
{
|
|
var frameKeypointData = m_FrameKeypointData[m_CurrentFrame];
|
|
frameKeypointData.isDepthCheckComplete = true;
|
|
m_FrameKeypointData[m_CurrentFrame] = frameKeypointData;
|
|
}
|
|
|
|
checkLocations.Dispose();
|
|
}
|
|
}
|
|
|
|
private void DoDepthCheck(ScriptableRenderContext scriptableRenderContext, List<KeypointEntry> keypointEntries, NativeList<float3> checkLocations)
|
|
{
|
|
var keypointCount = keypointEntries.Count * activeTemplate.keypoints.Length;
|
|
|
|
var commandBuffer = CommandBufferPool.Get("KeypointDepthCheck");
|
|
|
|
var textureDimensions = TextureDimensions(keypointCount);
|
|
|
|
|
|
SetupDepthCheckBuffers(checkLocations.Length);
|
|
|
|
var positionsPixeldata = new NativeArray<half>(textureDimensions.x * textureDimensions.y * 2, Allocator.Temp);
|
|
var depthPixeldata = new NativeArray<half>(textureDimensions.x * textureDimensions.y, Allocator.Temp);
|
|
var depthTexture = Shader.GetGlobalTexture("_CameraDepthTexture");
|
|
|
|
for (int i = 0; i < checkLocations.Length; i++)
|
|
{
|
|
var pos = checkLocations[i];
|
|
positionsPixeldata[i * 2] = new half(pos.x);
|
|
positionsPixeldata[i * 2 + 1] = new half(pos.y);
|
|
depthPixeldata[i] = new half(pos.z);
|
|
}
|
|
|
|
m_KeypointPositionsTexture.SetPixelData(positionsPixeldata, 0);
|
|
m_KeypointPositionsTexture.Apply();
|
|
m_KeypointCheckDepthTexture.SetPixelData(depthPixeldata, 0);
|
|
m_KeypointCheckDepthTexture.Apply();
|
|
|
|
positionsPixeldata.Dispose();
|
|
depthPixeldata.Dispose();
|
|
|
|
m_MaterialDepthCheck.SetTexture("_Positions", m_KeypointPositionsTexture);
|
|
m_MaterialDepthCheck.SetTexture("_KeypointCheckDepth", m_KeypointCheckDepthTexture);
|
|
m_MaterialDepthCheck.SetTexture("_CameraDepthTexture", depthTexture);
|
|
commandBuffer.Blit(null, m_ResultsBuffer, m_MaterialDepthCheck);
|
|
|
|
scriptableRenderContext.ExecuteCommandBuffer(commandBuffer);
|
|
scriptableRenderContext.Submit();
|
|
CommandBufferPool.Release(commandBuffer);
|
|
m_DepthCheckReader.Capture(scriptableRenderContext, OnDepthCheckReadback);
|
|
}
|
|
|
|
private static Vector2Int TextureDimensions(int keypointCount)
|
|
{
|
|
var width = Math.Max(k_MinTextureWidth, Mathf.NextPowerOfTwo((int)Math.Ceiling(Math.Sqrt(keypointCount))));
|
|
var height = width;
|
|
|
|
var textureDimensions = new Vector2Int(width, height);
|
|
return textureDimensions;
|
|
}
|
|
|
|
private void OnDepthCheckReadback(int frame, NativeArray<Color32> data, RenderTexture renderTexture)
|
|
{
|
|
DoDepthCheckReadback(frame, data);
|
|
}
|
|
|
|
private void OnDepthCheckReadback(int frameCount, AsyncGPUReadbackRequest obj)
|
|
{
|
|
var data = obj.GetData<uint>();
|
|
//DoDepthCheckReadback(frameCount, data);
|
|
}
|
|
|
|
private void DoDepthCheckReadback(int frameCount, NativeArray<Color32> data)
|
|
{
|
|
var frameKeypointData = m_FrameKeypointData[frameCount];
|
|
var totalLength = frameKeypointData.pointsPerEntry * frameKeypointData.keypoints.Count;
|
|
Debug.Assert(totalLength < data.Length);
|
|
for (var i = 0; i < totalLength; i++)
|
|
{
|
|
var value = data[i];
|
|
if (value.r == 0)
|
|
{
|
|
var keypoints = frameKeypointData.keypoints[i / frameKeypointData.pointsPerEntry];
|
|
var indexInObject = i % frameKeypointData.pointsPerEntry;
|
|
var keypoint = keypoints.keypoints[indexInObject];
|
|
keypoint.state = 1;
|
|
keypoints.keypoints[indexInObject] = keypoint;
|
|
}
|
|
}
|
|
|
|
frameKeypointData.isDepthCheckComplete = true;
|
|
m_FrameKeypointData[frameCount] = frameKeypointData;
|
|
ReportIfComplete(frameCount, frameKeypointData);
|
|
}
|
|
|
|
// ReSharper disable InconsistentNaming
|
|
// ReSharper disable NotAccessedField.Global
|
|
// ReSharper disable NotAccessedField.Local
|
|
/// <summary>
|
|
/// Record storing all of the keypoint data of a labeled gameobject.
|
|
/// </summary>
|
|
[Serializable]
|
|
public class KeypointEntry
|
|
{
|
|
/// <summary>
|
|
/// The label id of the entity
|
|
/// </summary>
|
|
public int label_id;
|
|
/// <summary>
|
|
/// The instance id of the entity
|
|
/// </summary>
|
|
public uint instance_id;
|
|
/// <summary>
|
|
/// The template that the points are based on
|
|
/// </summary>
|
|
public string template_guid;
|
|
/// <summary>
|
|
/// Pose ground truth for the current set of keypoints
|
|
/// </summary>
|
|
public string pose = "unset";
|
|
/// <summary>
|
|
/// Array of all of the keypoints
|
|
/// </summary>
|
|
public Keypoint[] keypoints;
|
|
}
|
|
|
|
/// <summary>
|
|
/// The values of a specific keypoint
|
|
/// </summary>
|
|
[Serializable]
|
|
public struct Keypoint
|
|
{
|
|
/// <summary>
|
|
/// The index of the keypoint in the template file
|
|
/// </summary>
|
|
public int index;
|
|
/// <summary>
|
|
/// The keypoint's x-coordinate pixel location
|
|
/// </summary>
|
|
public float x;
|
|
/// <summary>
|
|
/// The keypoint's y-coordinate pixel location
|
|
/// </summary>
|
|
public float y;
|
|
/// <summary>
|
|
/// The state of the point,
|
|
/// 0 = not present,
|
|
/// 1 = keypoint is present but not visible,
|
|
/// 2 = keypoint is present and visible
|
|
/// </summary>
|
|
public int state;
|
|
}
|
|
// ReSharper restore InconsistentNaming
|
|
// ReSharper restore NotAccessedField.Global
|
|
// ReSharper restore NotAccessedField.Local
|
|
|
|
float GetCaptureHeight()
|
|
{
|
|
var targetTexture = perceptionCamera.attachedCamera.targetTexture;
|
|
return targetTexture != null ?
|
|
targetTexture.height : Screen.height;
|
|
}
|
|
Vector3 ConvertToScreenSpace(Vector3 worldLocation)
|
|
{
|
|
var pt = perceptionCamera.attachedCamera.WorldToScreenPoint(worldLocation);
|
|
pt.y = GetCaptureHeight() - pt.y;
|
|
if (Mathf.Approximately(pt.y, perceptionCamera.attachedCamera.pixelHeight))
|
|
pt.y -= .0001f;
|
|
if (Mathf.Approximately(pt.x, perceptionCamera.attachedCamera.pixelWidth))
|
|
pt.x -= .0001f;
|
|
|
|
return pt;
|
|
}
|
|
|
|
struct CachedData
|
|
{
|
|
public bool status;
|
|
public Animator animator;
|
|
public KeypointEntry keypoints;
|
|
public List<(JointLabel, int)> overrides;
|
|
public float occlusionScalar;
|
|
}
|
|
|
|
Dictionary<uint, CachedData> m_KnownStatus;
|
|
|
|
bool TryToGetTemplateIndexForJoint(KeypointTemplate template, JointLabel joint, out int index)
|
|
{
|
|
index = -1;
|
|
|
|
foreach (var label in joint.labels)
|
|
{
|
|
for (var i = 0; i < template.keypoints.Length; i++)
|
|
{
|
|
if (template.keypoints[i].label == label)
|
|
{
|
|
index = i;
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
|
|
return false;
|
|
}
|
|
|
|
bool DoesTemplateContainJoint(JointLabel jointLabel)
|
|
{
|
|
return TryToGetTemplateIndexForJoint(activeTemplate, jointLabel, out _);
|
|
}
|
|
|
|
void ProcessLabel(Labeling labeledEntity, List<KeypointEntry> keypointEntries, NativeList<float3> checkLocations)
|
|
{
|
|
if (!idLabelConfig.TryGetLabelEntryFromInstanceId(labeledEntity.instanceId, out var labelEntry))
|
|
return;
|
|
|
|
// Cache out the data of a labeled game object the first time we see it, this will
|
|
// save performance each frame. Also checks to see if a labeled game object can be annotated.
|
|
if (!m_KnownStatus.ContainsKey(labeledEntity.instanceId))
|
|
{
|
|
var cached = new CachedData()
|
|
{
|
|
status = false,
|
|
animator = null,
|
|
keypoints = new KeypointEntry(),
|
|
overrides = new List<(JointLabel, int)>(),
|
|
occlusionScalar = 1.0f
|
|
};
|
|
|
|
var entityGameObject = labeledEntity.gameObject;
|
|
|
|
cached.keypoints.instance_id = labeledEntity.instanceId;
|
|
cached.keypoints.label_id = labelEntry.id;
|
|
cached.keypoints.template_guid = activeTemplate.templateID;
|
|
|
|
cached.keypoints.keypoints = new Keypoint[activeTemplate.keypoints.Length];
|
|
for (var i = 0; i < cached.keypoints.keypoints.Length; i++)
|
|
{
|
|
cached.keypoints.keypoints[i] = new Keypoint { index = i, state = 0 };
|
|
}
|
|
|
|
var animator = entityGameObject.transform.GetComponentInChildren<Animator>();
|
|
if (animator != null)
|
|
{
|
|
cached.animator = animator;
|
|
cached.status = true;
|
|
}
|
|
|
|
foreach (var joint in entityGameObject.transform.GetComponentsInChildren<JointLabel>())
|
|
{
|
|
if (TryToGetTemplateIndexForJoint(activeTemplate, joint, out var idx))
|
|
{
|
|
cached.overrides.Add((joint, idx));
|
|
cached.status = true;
|
|
}
|
|
}
|
|
|
|
var occlusionOverrider = labeledEntity.GetComponent<KeypointOcclusionOverrides>();
|
|
if (occlusionOverrider != null)
|
|
{
|
|
cached.occlusionScalar = occlusionOverrider.overrideDistanceScale;
|
|
}
|
|
|
|
m_KnownStatus[labeledEntity.instanceId] = cached;
|
|
}
|
|
|
|
var cachedData = m_KnownStatus[labeledEntity.instanceId];
|
|
|
|
if (cachedData.status)
|
|
{
|
|
var animator = cachedData.animator;
|
|
|
|
var listStart = checkLocations.Length;
|
|
checkLocations.Resize(checkLocations.Length + activeTemplate.keypoints.Length, NativeArrayOptions.ClearMemory);
|
|
//grab the slice of the list for the current object to assign positions in
|
|
var checkLocationsSlice = new NativeSlice<float3>(checkLocations, listStart);
|
|
|
|
var cameraPosition = perceptionCamera.transform.position;
|
|
var cameraforward = perceptionCamera.transform.forward;
|
|
|
|
// Go through all of the rig keypoints and get their location
|
|
for (var i = 0; i < activeTemplate.keypoints.Length; i++)
|
|
{
|
|
var pt = activeTemplate.keypoints[i];
|
|
if (pt.associateToRig)
|
|
{
|
|
var bone = animator.GetBoneTransform(pt.rigLabel);
|
|
if (bone != null)
|
|
{
|
|
var bonePosition = bone.position;
|
|
|
|
// Check to see if
|
|
|
|
var occlusionDistance = pt.selfOcclusionDistance * cachedData.occlusionScalar;
|
|
|
|
var jointSelfOcclusionDistance = JointSelfOcclusionDistance(bone, bonePosition, cameraPosition, cameraforward, occlusionDistance);
|
|
InitKeypoint(bonePosition, cachedData, checkLocationsSlice, i, jointSelfOcclusionDistance);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Go through all of the additional or override points defined by joint labels and get
|
|
// their locations
|
|
foreach (var (joint, templateIdx) in cachedData.overrides)
|
|
{
|
|
var jointTransform = joint.transform;
|
|
var jointPosition = jointTransform.position;
|
|
float resolvedSelfOcclusionDistance;
|
|
if (joint.overrideSelfOcclusionDistance)
|
|
resolvedSelfOcclusionDistance = joint.selfOcclusionDistance;
|
|
else
|
|
resolvedSelfOcclusionDistance = activeTemplate.keypoints[templateIdx].selfOcclusionDistance;
|
|
|
|
var jointSelfOcclusionDistance = JointSelfOcclusionDistance(joint.transform, jointPosition, cameraPosition, cameraforward, resolvedSelfOcclusionDistance);
|
|
|
|
InitKeypoint(jointPosition, cachedData, checkLocationsSlice, templateIdx, jointSelfOcclusionDistance);
|
|
}
|
|
|
|
cachedData.keypoints.pose = "unset";
|
|
|
|
if (cachedData.animator != null)
|
|
{
|
|
cachedData.keypoints.pose = GetPose(cachedData.animator);
|
|
}
|
|
|
|
var cachedKeypointEntry = cachedData.keypoints;
|
|
var keypointEntry = new KeypointEntry()
|
|
{
|
|
instance_id = cachedKeypointEntry.instance_id,
|
|
keypoints = cachedKeypointEntry.keypoints.ToArray(),
|
|
label_id = cachedKeypointEntry.label_id,
|
|
pose = cachedKeypointEntry.pose,
|
|
template_guid = cachedKeypointEntry.template_guid
|
|
};
|
|
keypointEntries.Add(keypointEntry);
|
|
}
|
|
}
|
|
|
|
private float JointSelfOcclusionDistance(Transform transform, Vector3 jointPosition, Vector3 cameraPosition,
|
|
Vector3 cameraforward, float configuredSelfOcclusionDistance)
|
|
{
|
|
var depthOfJoint = Vector3.Dot(jointPosition - cameraPosition, cameraforward);
|
|
var cameraEffectivePosition = jointPosition - cameraforward * depthOfJoint;
|
|
|
|
var jointRelativeCameraPosition = transform.InverseTransformPoint(cameraEffectivePosition);
|
|
var jointRelativeCheckPosition = jointRelativeCameraPosition.normalized * configuredSelfOcclusionDistance;
|
|
var worldSpaceCheckVector = transform.TransformVector(jointRelativeCheckPosition);
|
|
return worldSpaceCheckVector.magnitude;
|
|
}
|
|
|
|
private void InitKeypoint(Vector3 position, CachedData cachedData, NativeSlice<float3> checkLocations, int idx,
|
|
float occlusionDistance)
|
|
{
|
|
var loc = ConvertToScreenSpace(position);
|
|
|
|
var keypoints = cachedData.keypoints.keypoints;
|
|
keypoints[idx].index = idx;
|
|
if (loc.z < 0)
|
|
{
|
|
keypoints[idx].x = 0;
|
|
keypoints[idx].y = 0;
|
|
keypoints[idx].state = 0;
|
|
}
|
|
else
|
|
{
|
|
keypoints[idx].x = loc.x;
|
|
keypoints[idx].y = loc.y;
|
|
keypoints[idx].state = 2;
|
|
}
|
|
|
|
//TODO: move this code
|
|
var pixelLocation = PixelLocationFromScreenPoint(keypoints[idx]);
|
|
if (pixelLocation.x < 0 || pixelLocation.y < 0)
|
|
{
|
|
pixelLocation = new int2(int.MaxValue, int.MaxValue);
|
|
}
|
|
|
|
checkLocations[idx] = new float3(pixelLocation.x + .5f, pixelLocation.y + .5f, loc.z - occlusionDistance);
|
|
}
|
|
|
|
string GetPose(Animator animator)
|
|
{
|
|
var info = animator.GetCurrentAnimatorClipInfo(0);
|
|
|
|
if (info != null && info.Length > 0)
|
|
{
|
|
var clip = info[0].clip;
|
|
var timeOffset = animator.GetCurrentAnimatorStateInfo(0).normalizedTime;
|
|
|
|
if (animationPoseConfigs != null)
|
|
{
|
|
foreach (var p in animationPoseConfigs)
|
|
{
|
|
if (p != null && p.animationClip == clip)
|
|
{
|
|
var time = timeOffset;
|
|
var label = p.GetPoseAtTime(time);
|
|
return label;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return "unset";
|
|
}
|
|
|
|
Keypoint? GetKeypointForJoint(KeypointEntry entry, int joint)
|
|
{
|
|
if (joint < 0 || joint >= entry.keypoints.Length) return null;
|
|
return entry.keypoints[joint];
|
|
}
|
|
|
|
/// <inheritdoc/>
|
|
protected override void OnVisualize()
|
|
{
|
|
// TODO - remove this, it is just for debugging
|
|
hudPanel.UpdateEntry(this, "frame", m_CurrentFrame.ToString());
|
|
// END OF TODO
|
|
|
|
if (m_KeypointEntriesToReport == null) return;
|
|
using (k_OnVisualizeMarker.Auto())
|
|
{
|
|
var jointTexture = activeTemplate.jointTexture;
|
|
if (jointTexture == null) jointTexture = m_MissingTexture;
|
|
|
|
var skeletonTexture = activeTemplate.skeletonTexture;
|
|
if (skeletonTexture == null) skeletonTexture = m_MissingTexture;
|
|
|
|
foreach (var entry in m_KeypointEntriesToReport)
|
|
{
|
|
if (drawBones)
|
|
{
|
|
foreach (var bone in activeTemplate.skeleton)
|
|
{
|
|
var joint1 = GetKeypointForJoint(entry, bone.joint1);
|
|
var joint2 = GetKeypointForJoint(entry, bone.joint2);
|
|
|
|
if (joint1 != null && joint1.Value.state == 2 && joint2 != null && joint2.Value.state == 2)
|
|
{
|
|
VisualizationHelper.DrawLine(joint1.Value.x, joint1.Value.y, joint2.Value.x, joint2.Value.y, bone.color, visThickness, skeletonTexture);
|
|
}
|
|
}
|
|
}
|
|
|
|
foreach (var keypoint in entry.keypoints)
|
|
{
|
|
#if false
|
|
var color = keypoint.state == 1 ? Color.black : activeTemplate.keypoints[keypoint.index].color;
|
|
|
|
if (keypoint.state > 0)
|
|
VisualizationHelper.DrawPoint(keypoint.x, keypoint.y, color, visualizationBaseSize * activeTemplate.keypoints[keypoint.index].selfOcclusionDistance, jointTexture);
|
|
#else
|
|
if (keypoint.state == 2)
|
|
VisualizationHelper.DrawPoint(keypoint.x, keypoint.y, activeTemplate.keypoints[keypoint.index].color, visThickness, jointTexture);
|
|
#endif
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ReSharper disable InconsistentNaming
|
|
// ReSharper disable NotAccessedField.Local
|
|
[Serializable]
|
|
struct JointJson
|
|
{
|
|
public string label;
|
|
public int index;
|
|
public Color color;
|
|
}
|
|
|
|
[Serializable]
|
|
struct SkeletonJson
|
|
{
|
|
public int joint1;
|
|
public int joint2;
|
|
public Color color;
|
|
}
|
|
|
|
[Serializable]
|
|
struct KeypointJson
|
|
{
|
|
public string template_id;
|
|
public string template_name;
|
|
public JointJson[] key_points;
|
|
public SkeletonJson[] skeleton;
|
|
}
|
|
// ReSharper restore InconsistentNaming
|
|
// ReSharper restore NotAccessedField.Local
|
|
|
|
KeypointJson TemplateToJson(KeypointTemplate input)
|
|
{
|
|
var json = new KeypointJson();
|
|
json.template_id = input.templateID;
|
|
json.template_name = input.templateName;
|
|
json.key_points = new JointJson[input.keypoints.Length];
|
|
json.skeleton = new SkeletonJson[input.skeleton.Length];
|
|
|
|
for (var i = 0; i < input.keypoints.Length; i++)
|
|
{
|
|
json.key_points[i] = new JointJson
|
|
{
|
|
label = input.keypoints[i].label,
|
|
index = i,
|
|
color = input.keypoints[i].color
|
|
};
|
|
}
|
|
|
|
for (var i = 0; i < input.skeleton.Length; i++)
|
|
{
|
|
json.skeleton[i] = new SkeletonJson()
|
|
{
|
|
joint1 = input.skeleton[i].joint1,
|
|
joint2 = input.skeleton[i].joint2,
|
|
color = input.skeleton[i].color
|
|
};
|
|
}
|
|
|
|
return json;
|
|
}
|
|
}
|
|
}
|