manolas
/
vcglib
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
vcglib/vcg/space/index/kdtree/kdtree.h

533 lines
18 KiB
C++
Executable File
Raw Blame History

/****************************************************************************
* VCGLib o o *
* Visual and Computer Graphics Library o o *
* _ O _ *
* Copyright(C) 2004-2016 \/)\/ *
* Visual Computing Lab /\/| *
* ISTI - Italian National Research Council | *
* \ *
* All rights reserved. *
* *
* This program is free software; you can redistribute it and/or modify *
* it under the terms of the GNU General Public License as published by *
* the Free Software Foundation; either version 2 of the License, or *
* (at your option) any later version. *
* *
* This program is distributed in the hope that it will be useful, *
* but WITHOUT ANY WARRANTY; without even the implied warranty of *
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the *
* GNU General Public License (http://www.gnu.org/licenses/gpl.txt) *
* for more details. *
* *
****************************************************************************/
#ifndef KDTREE_VCG_H
#define KDTREE_VCG_H
#include <vcg/space/point3.h>
#include <vcg/space/box3.h>
#include <vcg/space/index/kdtree/priorityqueue.h>
#include <vector>
#include <limits>
#include <iostream>
#include <cstdint>
namespace vcg {
template<typename _DataType>
class ConstDataWrapper
{
public:
typedef _DataType DataType;
inline ConstDataWrapper()
: mpData(0), mStride(0), mSize(0)
{}
inline ConstDataWrapper(const DataType* pData, int size, int64_t stride = sizeof(DataType))
: mpData(reinterpret_cast<const unsigned char*>(pData)), mStride(stride), mSize(size)
{}
inline const DataType& operator[] (int i) const
{
return *reinterpret_cast<const DataType*>(mpData + i*mStride);
}
inline size_t size() const { return mSize; }
protected:
const unsigned char* mpData;
int64_t mStride;
size_t mSize;
};
template<class StdVectorType>
class VectorConstDataWrapper :public ConstDataWrapper<typename StdVectorType::value_type>
{
public:
inline VectorConstDataWrapper(StdVectorType &vec) :
ConstDataWrapper<typename StdVectorType::value_type>(&(vec[0]), vec.size(), sizeof(typename StdVectorType::value_type))
{}
};
template<class MeshType>
class VertexConstDataWrapper :public ConstDataWrapper<typename MeshType::CoordType>
{
public:
inline VertexConstDataWrapper(MeshType &m) :
ConstDataWrapper<typename MeshType::CoordType>(&(m.vert[0].P()), m.vert.size(), sizeof(typename MeshType::VertexType))
{}
};
/**
* This class allows to create a Kd-Tree thought to perform the neighbour query (radius search, knn-nearest serach and closest search).
* The class implemetantion is thread-safe.
*/
template<typename _Scalar>
class KdTree
{
public:
typedef _Scalar Scalar;
typedef vcg::Point3<Scalar> VectorType;
typedef vcg::Box3<Scalar> AxisAlignedBoxType;
typedef HeapMaxPriorityQueue<int, Scalar> PriorityQueue;
struct Node
{
union {
//standard node
struct {
Scalar splitValue;
unsigned int firstChildId : 24;
unsigned int dim : 2;
unsigned int leaf : 1;
};
//leaf
struct {
unsigned int start;
unsigned short size;
};
};
};
typedef std::vector<Node> NodeList;
// return the protected members which store the nodes and the points list
inline const NodeList& _getNodes(void) { return mNodes; }
inline const std::vector<VectorType>& _getPoints(void) { return mPoints; }
inline unsigned int _getNumLevel(void) { return numLevel; }
inline const AxisAlignedBoxType& _getAABBox(void) { return mAABB; }
public:
KdTree(const ConstDataWrapper<VectorType>& points, unsigned int nofPointsPerCell = 16, unsigned int maxDepth = 64, bool balanced = false);
~KdTree();
void doQueryK(const VectorType& queryPoint, int k, PriorityQueue& mNeighborQueue);
void doQueryDist(const VectorType& queryPoint, float dist, std::vector<unsigned int>& points, std::vector<Scalar>& sqrareDists);
void doQueryClosest(const VectorType& queryPoint, unsigned int& index, Scalar& dist);
protected:
// element of the stack
struct QueryNode
{
QueryNode() {}
QueryNode(unsigned int id) : nodeId(id) {}
unsigned int nodeId; // id of the next node
Scalar sq; // squared distance to the next node
};
// used to build the tree: split the subset [start..end[ according to dim and splitValue,
// and returns the index of the first element of the second subset
unsigned int split(int start, int end, unsigned int dim, float splitValue);
int createTree(unsigned int nodeId, unsigned int start, unsigned int end, unsigned int level);
protected:
AxisAlignedBoxType mAABB; //BoundingBox
NodeList mNodes; //kd-tree nodes
std::vector<VectorType> mPoints; //points read from the input DataWrapper
std::vector<unsigned int> mIndices; //points indices
unsigned int targetCellSize; //min number of point in a leaf
unsigned int targetMaxDepth; //max tree depth
unsigned int numLevel; //actual tree depth
bool isBalanced; //true if the tree is balanced
};
template<typename Scalar>
KdTree<Scalar>::KdTree(const ConstDataWrapper<VectorType>& points, unsigned int nofPointsPerCell, unsigned int maxDepth, bool balanced)
: mPoints(points.size()), mIndices(points.size())
{
// compute the AABB of the input
mPoints[0] = points[0];
mAABB.Set(mPoints[0]);
for (unsigned int i = 1; i < mPoints.size(); ++i)
{
mPoints[i] = points[i];
mIndices[i] = i;
mAABB.Add(mPoints[i]);
}
targetMaxDepth = maxDepth;
targetCellSize = nofPointsPerCell;
isBalanced = balanced;
//mNodes.reserve(4 * mPoints.size() / nofPointsPerCell);
//first node inserted (no leaf). The others are made by the createTree function (recursively)
mNodes.resize(1);
mNodes.back().leaf = 0;
numLevel = createTree(0, 0, mPoints.size(), 1);
}
template<typename Scalar>
KdTree<Scalar>::~KdTree()
{
}
/** Performs the kNN query.
*
* This algorithm uses the simple distance to the split plane to prune nodes.
* A more elaborated approach consists to track the closest corner of the cell
* relatively to the current query point. This strategy allows to save about 5%
* of the leaves. However, in practice the slight overhead due to this tracking
* reduces the overall performance.
*
* This algorithm also use a simple stack while a priority queue using the squared
* distances to the cells as a priority values allows to save about 10% of the leaves.
* But, again, priority queue insertions and deletions are quite involved, and therefore
* a simple stack is by far much faster.
*
* The result of the query, the k-nearest neighbors, are stored into the stack mNeighborQueue, where the
* topmost element [0] is NOT the nearest but the farthest!! (they are not sorted but arranged into a heap).
*/
template<typename Scalar>
void KdTree<Scalar>::doQueryK(const VectorType& queryPoint, int k, PriorityQueue& mNeighborQueue)
{
mNeighborQueue.setMaxSize(k);
mNeighborQueue.init();
std::vector<QueryNode> mNodeStack(numLevel + 1);
mNodeStack[0].nodeId = 0;
mNodeStack[0].sq = 0.f;
unsigned int count = 1;
while (count)
{
//we select the last node (AABB) inserted in the stack
QueryNode& qnode = mNodeStack[count - 1];
//while going down the tree qnode.nodeId is the nearest sub-tree, otherwise,
//in backtracking, qnode.nodeId is the other sub-tree that will be visited iff
//the actual nearest node is further than the split distance.
Node& node = mNodes[qnode.nodeId];
//if the distance is less than the top of the max-heap, it could be one of the k-nearest neighbours
if (mNeighborQueue.getNofElements() < k || qnode.sq < mNeighborQueue.getTopWeight())
{
//when we arrive to a leaf
if (node.leaf)
{
--count; //pop of the leaf
//end is the index of the last element of the leaf in mPoints
unsigned int end = node.start + node.size;
//adding the element of the leaf to the heap
for (unsigned int i = node.start; i < end; ++i)
mNeighborQueue.insert(mIndices[i], vcg::SquaredNorm(queryPoint - mPoints[i]));
}
//otherwise, if we're not on a leaf
else
{
// the new offset is the distance between the searched point and the actual split coordinate
float new_off = queryPoint[node.dim] - node.splitValue;
//left sub-tree
if (new_off < 0.)
{
mNodeStack[count].nodeId = node.firstChildId;
//in the father's nodeId we save the index of the other sub-tree (for backtracking)
qnode.nodeId = node.firstChildId + 1;
}
//right sub-tree (same as above)
else
{
mNodeStack[count].nodeId = node.firstChildId + 1;
qnode.nodeId = node.firstChildId;
}
//distance is inherited from the father (while descending the tree it's equal to 0)
mNodeStack[count].sq = qnode.sq;
//distance of the father is the squared distance from the split plane
qnode.sq = new_off*new_off;
++count;
}
}
else
{
// pop
--count;
}
}
}
/** Performs the distance query.
*
* The result of the query, all the points within the distance dist form the query point, is the vector of the indeces
* and the vector of the squared distances from the query point.
*/
template<typename Scalar>
void KdTree<Scalar>::doQueryDist(const VectorType& queryPoint, float dist, std::vector<unsigned int>& points, std::vector<Scalar>& sqrareDists)
{
std::vector<QueryNode> mNodeStack(numLevel + 1);
mNodeStack[0].nodeId = 0;
mNodeStack[0].sq = 0.f;
unsigned int count = 1;
float sqrareDist = dist*dist;
while (count)
{
QueryNode& qnode = mNodeStack[count - 1];
Node & node = mNodes[qnode.nodeId];
if (qnode.sq < sqrareDist)
{
if (node.leaf)
{
--count; // pop
unsigned int end = node.start + node.size;
for (unsigned int i = node.start; i < end; ++i)
{
float pointSquareDist = vcg::SquaredNorm(queryPoint - mPoints[i]);
if (pointSquareDist < sqrareDist)
{
points.push_back(mIndices[i]);
sqrareDists.push_back(pointSquareDist);
}
}
}
else
{
// replace the stack top by the farthest and push the closest
float new_off = queryPoint[node.dim] - node.splitValue;
if (new_off < 0.)
{
mNodeStack[count].nodeId = node.firstChildId;
qnode.nodeId = node.firstChildId + 1;
}
else
{
mNodeStack[count].nodeId = node.firstChildId + 1;
qnode.nodeId = node.firstChildId;
}
mNodeStack[count].sq = qnode.sq;
qnode.sq = new_off*new_off;
++count;
}
}
else
{
// pop
--count;
}
}
}
/** Searchs the closest point.
*
* The result of the query, the closest point to the query point, is the index of the point and
* and the squared distance from the query point.
*/
template<typename Scalar>
void KdTree<Scalar>::doQueryClosest(const VectorType& queryPoint, unsigned int& index, Scalar& dist)
{
std::vector<QueryNode> mNodeStack(numLevel + 1);
mNodeStack[0].nodeId = 0;
mNodeStack[0].sq = 0.f;
unsigned int count = 1;
int minIndex = mIndices.size() / 2;
Scalar minDist = vcg::SquaredNorm(queryPoint - mPoints[minIndex]);
minIndex = mIndices[minIndex];
while (count)
{
QueryNode& qnode = mNodeStack[count - 1];
Node & node = mNodes[qnode.nodeId];
if (qnode.sq < minDist)
{
if (node.leaf)
{
--count; // pop
unsigned int end = node.start + node.size;
for (unsigned int i = node.start; i < end; ++i)
{
float pointSquareDist = vcg::SquaredNorm(queryPoint - mPoints[i]);
if (pointSquareDist < minDist)
{
minDist = pointSquareDist;
minIndex = mIndices[i];
}
}
}
else
{
// replace the stack top by the farthest and push the closest
float new_off = queryPoint[node.dim] - node.splitValue;
if (new_off < 0.)
{
mNodeStack[count].nodeId = node.firstChildId;
qnode.nodeId = node.firstChildId + 1;
}
else
{
mNodeStack[count].nodeId = node.firstChildId + 1;
qnode.nodeId = node.firstChildId;
}
mNodeStack[count].sq = qnode.sq;
qnode.sq = new_off*new_off;
++count;
}
}
else
{
// pop
--count;
}
}
index = minIndex;
dist = minDist;
}
/**
* Split the subarray between start and end in two part, one with the elements less than splitValue,
* the other with the elements greater or equal than splitValue. The elements are compared
* using the "dim" coordinate [0 = x, 1 = y, 2 = z].
*/
template<typename Scalar>
unsigned int KdTree<Scalar>::split(int start, int end, unsigned int dim, float splitValue)
{
int l(start), r(end - 1);
for (; l < r; ++l, --r)
{
while (l < end && mPoints[l][dim] < splitValue)
l++;
while (r >= start && mPoints[r][dim] >= splitValue)
r--;
if (l > r)
break;
std::swap(mPoints[l], mPoints[r]);
std::swap(mIndices[l], mIndices[r]);
}
//returns the index of the first element on the second part
return (mPoints[l][dim] < splitValue ? l + 1 : l);
}
/** recursively builds the kdtree
*
* The heuristic is the following:
* - if the number of points in the node is lower than targetCellsize then make a leaf
* - else compute the AABB of the points of the node and split it at the middle of
* the largest AABB dimension.
*
* This strategy might look not optimal because it does not explicitly prune empty space,
* unlike more advanced SAH-like techniques used for RT. On the other hand it leads to a shorter tree,
* faster to traverse and our experience shown that in the special case of kNN queries,
* this strategy is indeed more efficient (and much faster to build). Moreover, for volume data
* (e.g., fluid simulation) pruning the empty space is useless.
*
* Actually, storing at each node the exact AABB (we therefore have a binary BVH) allows
* to prune only about 10% of the leaves, but the overhead of this pruning (ball/ABBB intersection)
* is more expensive than the gain it provides and the memory consumption is x4 higher !
*/
template<typename Scalar>
int KdTree<Scalar>::createTree(unsigned int nodeId, unsigned int start, unsigned int end, unsigned int level)
{
//select the first node
Node& node = mNodes[nodeId];
AxisAlignedBoxType aabb;
//putting all the points in the bounding box
aabb.Set(mPoints[start]);
for (unsigned int i = start + 1; i < end; ++i)
aabb.Add(mPoints[i]);
//bounding box diagonal
VectorType diag = aabb.max - aabb.min;
//the split "dim" is the dimension of the box with the biggest value
unsigned int dim;
if (diag.X() > diag.Y())
dim = diag.X() > diag.Z() ? 0 : 2;
else
dim = diag.Y() > diag.Z() ? 1 : 2;
node.dim = dim;
if (isBalanced) //we divide the points using the median value along the "dim" dimension
{
std::vector<Scalar> tempVector;
for (unsigned int i = start + 1; i < end; ++i)
tempVector.push_back(mPoints[i][dim]);
std::sort(tempVector.begin(), tempVector.end());
node.splitValue = (tempVector[tempVector.size() / 2.0] + tempVector[tempVector.size() / 2.0 + 1]) / 2.0;
}
else //we divide the bounding box in 2 partitions, considering the average of the "dim" dimension
node.splitValue = Scalar(0.5*(aabb.max[dim] + aabb.min[dim]));
//midId is the index of the first element in the second partition
unsigned int midId = split(start, end, dim, node.splitValue);
node.firstChildId = mNodes.size();
mNodes.resize(mNodes.size() + 2);
bool flag = (midId == start) || (midId == end);
int leftLevel, rightLevel;
{
// left child
unsigned int childId = mNodes[nodeId].firstChildId;
Node& child = mNodes[childId];
if (flag || (midId - start) <= targetCellSize || level >= targetMaxDepth)
{
child.leaf = 1;
child.start = start;
child.size = midId - start;
leftLevel = level;
}
else
{
child.leaf = 0;
leftLevel = createTree(childId, start, midId, level + 1);
}
}
{
// right child
unsigned int childId = mNodes[nodeId].firstChildId + 1;
Node& child = mNodes[childId];
if (flag || (end - midId) <= targetCellSize || level >= targetMaxDepth)
{
child.leaf = 1;
child.start = midId;
child.size = end - midId;
rightLevel = level;
}
else
{
child.leaf = 0;
rightLevel = createTree(childId, midId, end, level + 1);
}
}
if (leftLevel > rightLevel)
return leftLevel;
return rightLevel;
}
}
#endif
Reference in New Issue View Git Blame Copy Permalink