您最多选择25个主题
主题必须以中文或者字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符
387 行
11 KiB
387 行
11 KiB
using System;
|
|
using UnityEngine;
|
|
using UnityEngine.Profiling;
|
|
using Unity.Collections;
|
|
using Unity.Collections.LowLevel.Unsafe;
|
|
using Unity.Jobs;
|
|
using Unity.Burst;
|
|
|
|
[Serializable]
|
|
public class KdTree3
|
|
{
|
|
public struct Node
|
|
{
|
|
public int point;
|
|
public int stepL;
|
|
public int stepR;
|
|
}
|
|
|
|
public struct Point3
|
|
{
|
|
public float x;// note: 'x' MUST be first field
|
|
public float y;
|
|
public float z;
|
|
public int index;
|
|
}
|
|
|
|
public int size;
|
|
public Node[] nodes;
|
|
public Point3[] points;
|
|
|
|
public KdTree3(Vector3[] pointCloud, int pointCount)
|
|
{
|
|
BuildFrom(pointCloud, pointCount);
|
|
}
|
|
|
|
//----------------------
|
|
// kd-tree construction
|
|
|
|
public void BuildFrom(Vector3[] pointBuffer, int pointCount)
|
|
{
|
|
ArrayUtils.ResizeCheckedIfLessThan(ref this.nodes, pointCount);
|
|
ArrayUtils.ResizeCheckedIfLessThan(ref this.points, pointCount);
|
|
|
|
Profiler.BeginSample("kd-prep");
|
|
for (int i = 0; i != pointCount; i++)
|
|
{
|
|
this.points[i].x = pointBuffer[i].x;
|
|
this.points[i].y = pointBuffer[i].y;
|
|
this.points[i].z = pointBuffer[i].z;
|
|
this.points[i].index = i;
|
|
}
|
|
Profiler.EndSample();
|
|
|
|
Profiler.BeginSample("kd-build");
|
|
if (pointCount > 0)
|
|
{
|
|
unsafe
|
|
{
|
|
int numThreads = SystemInfo.processorCount;
|
|
//Debug.Log("numThreads = " + numThreads);
|
|
|
|
int schedLeaves = Mathf.NextPowerOfTwo(numThreads) * 2;
|
|
int schedNodes = BalancedBinaryTreeInfo.NodeCountFromLeafCount(schedLeaves);
|
|
int schedDepth = BalancedBinaryTreeInfo.DepthFromLeafCount(schedLeaves) - 1;
|
|
|
|
//Debug.Log("numThreads=" + numThreads + ", schedLeaves=" + schedLeaves + ", schedNodes=" + schedNodes + ", schedDepth=" + schedDepth);
|
|
|
|
var jobArray = (JobHandle*)UnsafeUtility.Malloc(sizeof(JobHandle) * schedNodes, 1, Allocator.Temp);
|
|
var jobSched = jobArray;
|
|
|
|
fixed (Node* __node = this.nodes)
|
|
fixed (Point3* __points = this.points)
|
|
{
|
|
ScheduleNode(ref jobSched, null, __node, __points, 0, pointCount, 0, schedDepth);
|
|
|
|
JobHandle.ScheduleBatchedJobs();
|
|
|
|
long jobCount = jobSched - jobArray;
|
|
for (long i = 0; i != jobCount; i++)
|
|
{
|
|
jobArray[i].Complete();
|
|
}
|
|
}
|
|
|
|
UnsafeUtility.Free(jobArray, Allocator.Temp);
|
|
}
|
|
}
|
|
Profiler.EndSample();
|
|
|
|
size = pointCount;
|
|
}
|
|
|
|
unsafe static void ScheduleNode(ref JobHandle* jobSched, JobHandle* jobDepends, Node* node, Point3* points, int offset, int length, int depth, int depthSingleThreaded)
|
|
{
|
|
// pick median
|
|
int median = length >> 1;
|
|
|
|
// calc offsets
|
|
var offsetL = offset;
|
|
var offsetR = offset + median + 1;
|
|
var lengthL = median;
|
|
var lengthR = length - median - 1;
|
|
var stepL = lengthL > 0 ? 1 : 0;
|
|
var stepR = lengthR > 0 ? 1 + lengthL : 0;
|
|
|
|
// make job
|
|
var job = new BuildNodeJob()
|
|
{
|
|
node = node,
|
|
points = points,
|
|
offset = offset,
|
|
length = length,
|
|
depth = depth,
|
|
leaf = (depth >= depthSingleThreaded),
|
|
};
|
|
|
|
var jobHandle = jobSched++;
|
|
if (jobDepends != null)
|
|
*jobHandle = job.Schedule(*jobDepends);
|
|
else
|
|
*jobHandle = job.Schedule();
|
|
|
|
if (depth == 0)
|
|
JobHandle.ScheduleBatchedJobs();
|
|
|
|
if (job.leaf)
|
|
return;
|
|
|
|
// schedule subtrees
|
|
if (lengthL > 0) ScheduleNode(ref jobSched, jobHandle, node + stepL, points, offsetL, lengthL, depth + 1, depthSingleThreaded);
|
|
if (lengthR > 0) ScheduleNode(ref jobSched, jobHandle, node + stepR, points, offsetR, lengthR, depth + 1, depthSingleThreaded);
|
|
}
|
|
|
|
[BurstCompile]
|
|
unsafe struct BuildNodeJob : IJob
|
|
{
|
|
[NativeDisableUnsafePtrRestriction] public Node* node;
|
|
[NativeDisableUnsafePtrRestriction] public Point3* points;// shared
|
|
|
|
public int offset;
|
|
public int length;
|
|
public int depth;
|
|
public bool leaf;
|
|
|
|
public void Execute()
|
|
{
|
|
if (leaf)
|
|
BuildNode(node, points, offset, length, depth);
|
|
else
|
|
BuildNodeDeferSubtrees(node, points, offset, length, depth);
|
|
}
|
|
}
|
|
|
|
unsafe static void BuildNode(Node* node, Point3* points, int offset, int length, int depth)
|
|
{
|
|
// pick median
|
|
int median = length >> 1;
|
|
|
|
// pick splitting axis
|
|
int axis = depth % 3;
|
|
|
|
// split points by median
|
|
SelectByAxis(median, points + offset, length, axis);
|
|
|
|
// calc offsets
|
|
var offsetL = offset;
|
|
var offsetR = offset + median + 1;
|
|
var lengthL = median;
|
|
var lengthR = length - median - 1;
|
|
var stepL = lengthL > 0 ? 1 : 0;
|
|
var stepR = lengthR > 0 ? 1 + lengthL : 0;
|
|
|
|
// make node
|
|
node->point = offset + median;
|
|
node->stepL = stepL;
|
|
node->stepR = stepR;
|
|
|
|
// build subtrees
|
|
if (lengthL > 0) BuildNode(node + stepL, points, offsetL, lengthL, depth + 1);
|
|
if (lengthR > 0) BuildNode(node + stepR, points, offsetR, lengthR, depth + 1);
|
|
}
|
|
|
|
unsafe static void BuildNodeDeferSubtrees(Node* node, Point3* points, int offset, int length, int depth)
|
|
{
|
|
// pick median
|
|
int median = length >> 1;
|
|
|
|
// pick splitting axis
|
|
int axis = depth % 3;
|
|
|
|
// split points by median
|
|
SelectByAxis(median, points + offset, length, axis);
|
|
|
|
// calc offsets
|
|
var lengthL = median;
|
|
var lengthR = length - median - 1;
|
|
var stepL = lengthL > 0 ? 1 : 0;
|
|
var stepR = lengthR > 0 ? 1 + lengthL : 0;
|
|
|
|
// make node
|
|
node->point = offset + median;
|
|
node->stepL = stepL;
|
|
node->stepR = stepR;
|
|
}
|
|
|
|
unsafe static void SelectByAxis(int nth, Point3* points, int length, int axis)
|
|
{
|
|
// note: this function adapted from public domain variants listed in KdTreeUtils
|
|
const int strideLsh = 2;
|
|
|
|
var i = 0;
|
|
var j = length - 1;
|
|
var v = &points->x + axis;
|
|
|
|
while (i < j)
|
|
{
|
|
var r = i;
|
|
var w = j;
|
|
|
|
float k = v[((r + w) >> 1) << strideLsh];
|
|
|
|
while (r < w)
|
|
{
|
|
if (v[r << strideLsh] >= k)
|
|
{
|
|
Point3 pw = points[w];
|
|
points[w] = points[r];
|
|
points[r] = pw;
|
|
w--;
|
|
}
|
|
else
|
|
{
|
|
r++;
|
|
}
|
|
}
|
|
|
|
if (v[r << strideLsh] > k)
|
|
r--;
|
|
|
|
if (nth <= r)
|
|
j = r;
|
|
else
|
|
i = r + 1;
|
|
}
|
|
}
|
|
|
|
//-----------------
|
|
// kd-tree queries
|
|
|
|
public int FindNearest(ref Vector3 target)
|
|
{
|
|
var bestDist = float.PositiveInfinity;
|
|
var bestNode = -1;
|
|
|
|
Profiler.BeginSample("kd-nearest");
|
|
FindNearest(ref bestDist, ref bestNode, 0, 0, ref target);
|
|
Profiler.EndSample();
|
|
|
|
if (bestNode != -1)
|
|
return points[nodes[bestNode].point].index;
|
|
else
|
|
return -1;
|
|
}
|
|
|
|
public bool FindNearest(ref float bestDist, ref int bestNode, ref Vector3 target)
|
|
{
|
|
var __bestDist = float.PositiveInfinity;
|
|
var __bestNode = -1;
|
|
|
|
Profiler.BeginSample("kd-nearest");
|
|
FindNearest(ref __bestDist, ref __bestNode, 0, 0, ref target);
|
|
Profiler.EndSample();
|
|
|
|
if (__bestNode != -1)
|
|
{
|
|
bestDist = __bestDist;
|
|
bestNode = points[nodes[__bestNode].point].index;
|
|
return true;
|
|
}
|
|
else
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
|
|
unsafe void FindNearest(ref float bestDist, ref int bestNode, int node, int depth, ref Vector3 target)
|
|
{
|
|
// update best index
|
|
int point = nodes[node].point;
|
|
Vector3 r;
|
|
r.x = target.x - points[point].x;
|
|
r.y = target.y - points[point].y;
|
|
r.z = target.z - points[point].z;
|
|
|
|
var dist = r.x * r.x + r.y * r.y + r.z * r.z;
|
|
if (dist < bestDist)
|
|
{
|
|
bestDist = dist;
|
|
bestNode = node;
|
|
}
|
|
|
|
// pick search axis
|
|
int axis = depth % 3;
|
|
|
|
// pick near, far
|
|
var delta = *(&r.x + axis);// avoid calling operator[]
|
|
int stepN = delta < 0.0f ? nodes[node].stepL : nodes[node].stepR;
|
|
int stepF = delta < 0.0f ? nodes[node].stepR : nodes[node].stepL;
|
|
|
|
// search near
|
|
if (stepN != 0)
|
|
{
|
|
FindNearest(ref bestDist, ref bestNode, node + stepN, depth + 1, ref target);
|
|
}
|
|
|
|
// search far
|
|
if (stepF != 0 && delta * delta < bestDist)
|
|
{
|
|
FindNearest(ref bestDist, ref bestNode, node + stepF, depth + 1, ref target);
|
|
}
|
|
}
|
|
|
|
public int RaycastApprox(ref Vector3 origin, ref Vector3 direction, int maxIterations = 100)
|
|
{
|
|
var position = origin;
|
|
var numIterations = 0;
|
|
|
|
var bestDist = float.PositiveInfinity;
|
|
var bestNode = -1;
|
|
|
|
while (numIterations++ < maxIterations)
|
|
{
|
|
var stepDist = float.PositiveInfinity;
|
|
var stepNode = -1;
|
|
|
|
FindNearest(ref stepDist, ref stepNode, ref position);
|
|
|
|
if (stepDist < bestDist)
|
|
{
|
|
bestDist = stepDist;
|
|
bestNode = stepNode;
|
|
|
|
if (bestDist < float.Epsilon)
|
|
{
|
|
break;// crude termination criteria
|
|
}
|
|
}
|
|
|
|
position = position + Mathf.Sqrt(stepDist) * direction;
|
|
}
|
|
|
|
return bestNode;
|
|
}
|
|
}
|
|
|
|
public static class BalancedBinaryTreeInfo
|
|
{
|
|
public static int LeafCountFromDepth(int depth)
|
|
{
|
|
int leafCount = (depth > 0) ? 1 : 0;
|
|
while (--depth > 0) leafCount = (leafCount << 1);
|
|
return leafCount;
|
|
}
|
|
|
|
public static int NodeCountFromDepth(int depth)
|
|
{
|
|
int nodeCount = 0;
|
|
while (--depth >= 0) nodeCount = (nodeCount << 1) | 1;
|
|
return nodeCount;
|
|
}
|
|
|
|
public static int NodeCountFromLeafCount(int leafCount)
|
|
{
|
|
Debug.Assert(Mathf.IsPowerOfTwo(leafCount));
|
|
int nodeCount = leafCount;
|
|
while ((leafCount >>= 1) > 0) nodeCount += leafCount;
|
|
return nodeCount;
|
|
}
|
|
|
|
public static int DepthFromLeafCount(int leafCount)
|
|
{
|
|
Debug.Assert(Mathf.IsPowerOfTwo(leafCount));
|
|
int depth = 0;
|
|
while (leafCount > 0) { leafCount >>= 1; depth++; }
|
|
return depth;
|
|
}
|
|
}
|