560 lines
14 KiB
C++
560 lines
14 KiB
C++
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
#include "PoseSearch/KDTree.h"
|
|
#include "Stats/Stats.h"
|
|
|
|
#ifndef UE_POSE_SEARCH_USE_NANOFLANN
|
|
#define UE_POSE_SEARCH_USE_NANOFLANN 1
|
|
#endif
|
|
|
|
// @third party code - BEGIN nanoflann
|
|
#if UE_POSE_SEARCH_USE_NANOFLANN
|
|
THIRD_PARTY_INCLUDES_START
|
|
#include "nanoflann/nanoflann.hpp"
|
|
THIRD_PARTY_INCLUDES_END
|
|
#endif
|
|
// @third party code - END nanoflann
|
|
|
|
namespace UE::PoseSearch
|
|
{
|
|
|
|
#if UE_POSE_SEARCH_USE_NANOFLANN
|
|
using FKDTreeImplementationBase = nanoflann::KDTreeSingleIndexAdaptor<nanoflann::L2_Simple_Adaptor<float, FKDTree::FDataSource>, FKDTree::FDataSource, -1, AccessorType>;
|
|
struct FKDTreeImplementation : FKDTreeImplementationBase
|
|
{
|
|
using FKDTreeImplementationBase::FKDTreeImplementationBase;
|
|
|
|
bool operator==(const FKDTreeImplementation& Other) const
|
|
{
|
|
if (m_size != Other.m_size)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
if (dim != Other.dim)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
const AccessorType RootBBoxSize = root_bbox.size();
|
|
if (RootBBoxSize != Other.root_bbox.size())
|
|
{
|
|
return false;
|
|
}
|
|
|
|
for (AccessorType Index = 0; Index < RootBBoxSize; ++Index)
|
|
{
|
|
const Interval& ThisInterval = root_bbox[Index];
|
|
const Interval& OtherInterval = Other.root_bbox[Index];
|
|
|
|
if (ThisInterval.high != OtherInterval.high)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
if (ThisInterval.low != OtherInterval.low)
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
|
|
if (m_leaf_max_size != Other.m_leaf_max_size)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
if (vAcc != Other.vAcc)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
if (!CompareNodes(root_node, Other.root_node))
|
|
{
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
private:
|
|
static bool CompareNodes(const NodePtr& NodeA, const NodePtr& NodeB)
|
|
{
|
|
const bool bAnyNodeAChild1 = NodeA->child1 != nullptr;
|
|
const bool bAnyNodeBChild1 = NodeA->child1 != nullptr;
|
|
if (bAnyNodeAChild1 != bAnyNodeBChild1)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
const bool bAnyNodeAChild2 = NodeA->child2 != nullptr;
|
|
const bool bAnyNodeBChild2 = NodeA->child2 != nullptr;
|
|
if (bAnyNodeAChild2 != bAnyNodeBChild2)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
const bool bIsLeafNode = !bAnyNodeAChild1 && !bAnyNodeAChild2;
|
|
if (bIsLeafNode)
|
|
{
|
|
if (NodeA->node_type.lr.left != NodeB->node_type.lr.left)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
if (NodeA->node_type.lr.right != NodeB->node_type.lr.right)
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
else
|
|
{
|
|
if (NodeA->node_type.sub.divfeat != NodeB->node_type.sub.divfeat)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
if (NodeA->node_type.sub.divhigh != NodeB->node_type.sub.divhigh)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
if (NodeA->node_type.sub.divlow != NodeB->node_type.sub.divlow)
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
|
|
if (bAnyNodeAChild1)
|
|
{
|
|
if (!CompareNodes(NodeA->child1, NodeB->child1))
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
|
|
if (bAnyNodeAChild2)
|
|
{
|
|
if (!CompareNodes(NodeA->child2, NodeB->child2))
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
};
|
|
#endif
|
|
|
|
FKDTree::FKDTree(AccessorType Count, AccessorType Dim, const float* Data, AccessorType MaxLeafSize)
|
|
: DataSource(Count, Dim, Data)
|
|
, KDTreeImplementation(nullptr)
|
|
{
|
|
#if UE_POSE_SEARCH_USE_NANOFLANN
|
|
if (Count > 0 && Dim > 0 && Data)
|
|
{
|
|
KDTreeImplementation = new FKDTreeImplementation(Dim, DataSource, nanoflann::KDTreeSingleIndexAdaptorParams(MaxLeafSize));
|
|
}
|
|
#endif // UE_POSE_SEARCH_USE_NANOFLANN
|
|
}
|
|
|
|
FKDTree::FKDTree()
|
|
: DataSource(0, 0, nullptr)
|
|
, KDTreeImplementation(nullptr)
|
|
{
|
|
}
|
|
|
|
FKDTree::~FKDTree()
|
|
{
|
|
Reset();
|
|
}
|
|
|
|
#if UE_POSE_SEARCH_USE_NANOFLANN
|
|
void CopySubTree(FKDTree& KDTree, FKDTreeImplementation::NodePtr& ThisNode, const FKDTreeImplementation::NodePtr& OtherNode)
|
|
{
|
|
check(KDTree.KDTreeImplementation);
|
|
|
|
ThisNode = KDTree.KDTreeImplementation->pool.template allocate<FKDTreeImplementation::Node>();
|
|
|
|
ThisNode->node_type = OtherNode->node_type;
|
|
|
|
if (OtherNode->child1 != nullptr)
|
|
{
|
|
CopySubTree(KDTree, ThisNode->child1, OtherNode->child1);
|
|
}
|
|
else
|
|
{
|
|
ThisNode->child1 = nullptr;
|
|
}
|
|
|
|
if (OtherNode->child2 != nullptr)
|
|
{
|
|
CopySubTree(KDTree, ThisNode->child2, OtherNode->child2);
|
|
}
|
|
else
|
|
{
|
|
ThisNode->child2 = nullptr;
|
|
}
|
|
}
|
|
|
|
#endif // UE_POSE_SEARCH_USE_NANOFLANN
|
|
|
|
FKDTree::FKDTree(const FKDTree& Other)
|
|
: DataSource()
|
|
, KDTreeImplementation(nullptr)
|
|
{
|
|
#if UE_POSE_SEARCH_USE_NANOFLANN
|
|
if (this != &Other && Other.KDTreeImplementation)
|
|
{
|
|
KDTreeImplementation = new FKDTreeImplementation(0, DataSource, nanoflann::KDTreeSingleIndexAdaptorParams(0));
|
|
|
|
DataSource = Other.DataSource;
|
|
|
|
check(Other.KDTreeImplementation->m_size <= AccessorTypeMax);
|
|
KDTreeImplementation->m_size = Other.KDTreeImplementation->m_size;
|
|
|
|
if (KDTreeImplementation->m_size > 0)
|
|
{
|
|
KDTreeImplementation->dim = Other.KDTreeImplementation->dim;
|
|
|
|
check(Other.KDTreeImplementation->root_bbox.size() <= AccessorTypeMax);
|
|
const AccessorType root_bbox_size = Other.KDTreeImplementation->root_bbox.size();
|
|
KDTreeImplementation->root_bbox.resize(root_bbox_size);
|
|
|
|
for (AccessorType i = 0; i < root_bbox_size; ++i)
|
|
{
|
|
KDTreeImplementation->root_bbox[i] = Other.KDTreeImplementation->root_bbox[i];
|
|
}
|
|
|
|
check(Other.KDTreeImplementation->m_leaf_max_size <= AccessorTypeMax);
|
|
const AccessorType KDTreeLeafMaxSize = Other.KDTreeImplementation->m_leaf_max_size;
|
|
KDTreeImplementation->m_leaf_max_size = KDTreeLeafMaxSize;
|
|
|
|
check(Other.KDTreeImplementation->vAcc.size() <= AccessorTypeMax);
|
|
const AccessorType VAccSize = Other.KDTreeImplementation->vAcc.size();
|
|
KDTreeImplementation->vAcc.resize(VAccSize);
|
|
|
|
for (AccessorType i = 0; i < VAccSize; ++i)
|
|
{
|
|
KDTreeImplementation->vAcc[i] = Other.KDTreeImplementation->vAcc[i];
|
|
}
|
|
|
|
CopySubTree(*this, KDTreeImplementation->root_node, Other.KDTreeImplementation->root_node);
|
|
}
|
|
}
|
|
#endif // UE_POSE_SEARCH_USE_NANOFLANN
|
|
}
|
|
|
|
FKDTree::FKDTree(FKDTree&& Other)
|
|
: DataSource()
|
|
, KDTreeImplementation(nullptr)
|
|
{
|
|
#if UE_POSE_SEARCH_USE_NANOFLANN
|
|
if (this != &Other)
|
|
{
|
|
if (Other.KDTreeImplementation)
|
|
{
|
|
KDTreeImplementation = new FKDTreeImplementation(0, DataSource, nanoflann::KDTreeSingleIndexAdaptorParams(0));
|
|
|
|
DataSource = MoveTemp(Other.DataSource);
|
|
|
|
check(Other.KDTreeImplementation->m_size <= AccessorTypeMax);
|
|
KDTreeImplementation->m_size = Other.KDTreeImplementation->m_size;
|
|
|
|
if (KDTreeImplementation->m_size > 0)
|
|
{
|
|
KDTreeImplementation->dim = Other.KDTreeImplementation->dim;
|
|
|
|
check(Other.KDTreeImplementation->root_bbox.size() <= AccessorTypeMax);
|
|
const AccessorType root_bbox_size = Other.KDTreeImplementation->root_bbox.size();
|
|
KDTreeImplementation->root_bbox.resize(root_bbox_size);
|
|
|
|
for (AccessorType i = 0; i < root_bbox_size; ++i)
|
|
{
|
|
KDTreeImplementation->root_bbox[i] = Other.KDTreeImplementation->root_bbox[i];
|
|
}
|
|
|
|
check(Other.KDTreeImplementation->m_leaf_max_size <= AccessorTypeMax);
|
|
const AccessorType KDTreeLeafMaxSize = Other.KDTreeImplementation->m_leaf_max_size;
|
|
KDTreeImplementation->m_leaf_max_size = KDTreeLeafMaxSize;
|
|
|
|
check(Other.KDTreeImplementation->vAcc.size() <= AccessorTypeMax);
|
|
const AccessorType VAccSize = Other.KDTreeImplementation->vAcc.size();
|
|
KDTreeImplementation->vAcc.resize(VAccSize);
|
|
|
|
for (AccessorType i = 0; i < VAccSize; ++i)
|
|
{
|
|
KDTreeImplementation->vAcc[i] = Other.KDTreeImplementation->vAcc[i];
|
|
}
|
|
|
|
CopySubTree(*this, KDTreeImplementation->root_node, Other.KDTreeImplementation->root_node);
|
|
}
|
|
}
|
|
|
|
Other.Reset();
|
|
}
|
|
#endif // UE_POSE_SEARCH_USE_NANOFLANN
|
|
}
|
|
|
|
FKDTree& FKDTree::operator=(const FKDTree& Other)
|
|
{
|
|
if (this != &Other)
|
|
{
|
|
this->~FKDTree();
|
|
new(this) FKDTree(Other);
|
|
}
|
|
return *this;
|
|
}
|
|
|
|
FKDTree& FKDTree::operator=(FKDTree&& Other)
|
|
{
|
|
if (this != &Other)
|
|
{
|
|
this->~FKDTree();
|
|
new(this) FKDTree(MoveTemp(Other));
|
|
}
|
|
return *this;
|
|
}
|
|
|
|
bool FKDTree::operator==(const FKDTree& Other) const
|
|
{
|
|
const bool bAnyImpl = KDTreeImplementation != nullptr;
|
|
const bool bAnyOtherImpl = Other.KDTreeImplementation != nullptr;
|
|
if (bAnyImpl != bAnyOtherImpl)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
#if UE_POSE_SEARCH_USE_NANOFLANN
|
|
if (bAnyImpl && bAnyOtherImpl)
|
|
{
|
|
if (*KDTreeImplementation != *Other.KDTreeImplementation)
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
#endif // UE_POSE_SEARCH_USE_NANOFLANN
|
|
|
|
if (DataSource != Other.DataSource)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
void FKDTree::Reset()
|
|
{
|
|
#if UE_POSE_SEARCH_USE_NANOFLANN
|
|
delete KDTreeImplementation;
|
|
KDTreeImplementation = nullptr;
|
|
#endif // UE_POSE_SEARCH_USE_NANOFLANN
|
|
DataSource = FDataSource();
|
|
}
|
|
|
|
void FKDTree::Construct(AccessorType Count, AccessorType Dim, const float* Data, AccessorType MaxLeafSize)
|
|
{
|
|
this->~FKDTree();
|
|
new(this) FKDTree(Count, Dim, Data, MaxLeafSize);
|
|
}
|
|
|
|
template <typename RESULTSET>
|
|
inline int32 FindNeighborsInternal(FKDTreeImplementation* KDTreeImplementation, RESULTSET& Result, TConstArrayView<float> Query)
|
|
{
|
|
#if UE_POSE_SEARCH_USE_NANOFLANN
|
|
|
|
QUICK_SCOPE_CYCLE_COUNTER(STAT_FKDTree_FindNeighbors);
|
|
|
|
if (KDTreeImplementation)
|
|
{
|
|
check(Query.GetData() && Query.Num() == KDTreeImplementation->dim && KDTreeImplementation->root_node);
|
|
|
|
const nanoflann::SearchParams SearchParams(
|
|
32, // Ignored parameter (Kept for compatibility with the FLANN interface).
|
|
0.f, // search for eps-approximate neighbours (default: 0)
|
|
false); // only for radius search, require neighbours sorted by
|
|
KDTreeImplementation->findNeighbors(Result, Query.GetData(), SearchParams);
|
|
return Result.Num();
|
|
}
|
|
return 0;
|
|
|
|
#else // UE_POSE_SEARCH_USE_NANOFLANN
|
|
|
|
checkNoEntry(); // unimplemented
|
|
return 0;
|
|
|
|
#endif // UE_POSE_SEARCH_USE_NANOFLANN
|
|
}
|
|
|
|
int32 FKDTree::FindNeighbors(FKNNResultSet& Result, TConstArrayView<float> Query) const
|
|
{
|
|
return FindNeighborsInternal(KDTreeImplementation, Result, Query);
|
|
}
|
|
|
|
int32 FKDTree::FindNeighbors(FFilteredKNNResultSet& Result, TConstArrayView<float> Query) const
|
|
{
|
|
return FindNeighborsInternal(KDTreeImplementation, Result, Query);
|
|
}
|
|
|
|
int32 FKDTree::FindNeighbors(FRadiusResultSet& Result, TConstArrayView<float> Query) const
|
|
{
|
|
return FindNeighborsInternal(KDTreeImplementation, Result, Query);
|
|
}
|
|
|
|
int32 FKDTree::FindNeighbors(FKNNMaxHeapResultSet& Result, TConstArrayView<float> Query) const
|
|
{
|
|
return FindNeighborsInternal(KDTreeImplementation, Result, Query);
|
|
}
|
|
|
|
int32 FKDTree::FindNeighbors(FFilteredKNNMaxHeapResultSet& Result, TConstArrayView<float> Query) const
|
|
{
|
|
return FindNeighborsInternal(KDTreeImplementation, Result, Query);
|
|
}
|
|
|
|
int32 FKDTree::FindNeighbors(FRadiusMaxHeapResultSet& Result, TConstArrayView<float> Query) const
|
|
{
|
|
return FindNeighborsInternal(KDTreeImplementation, Result, Query);
|
|
}
|
|
|
|
SIZE_T FKDTree::GetAllocatedSize() const
|
|
{
|
|
SIZE_T AllocatedSize = sizeof(FKDTree);
|
|
|
|
#if UE_POSE_SEARCH_USE_NANOFLANN
|
|
if (KDTreeImplementation)
|
|
{
|
|
AllocatedSize += sizeof(FKDTreeImplementation);
|
|
AllocatedSize += KDTreeImplementation->usedMemory(*KDTreeImplementation);
|
|
}
|
|
#endif // UE_POSE_SEARCH_USE_NANOFLANN
|
|
|
|
return AllocatedSize;
|
|
}
|
|
|
|
#if UE_POSE_SEARCH_USE_NANOFLANN
|
|
FArchive& SerializeSubTree(FArchive& Ar, FKDTree& KDTree, FKDTreeImplementation::NodePtr& KDTreeNode)
|
|
{
|
|
check(KDTree.KDTreeImplementation);
|
|
|
|
if (Ar.IsLoading())
|
|
{
|
|
KDTreeNode = KDTree.KDTreeImplementation->pool.template allocate<FKDTreeImplementation::Node>();
|
|
// zeroing FKDTreeImplementation::Node memory since it contains a union and doesn't have a constructor
|
|
FMemory::Memzero(KDTreeNode, sizeof(FKDTreeImplementation::Node));
|
|
}
|
|
|
|
bool bAnyNodeChild1 = KDTreeNode->child1 != nullptr;
|
|
bool bAnyNodeChild2 = KDTreeNode->child2 != nullptr;
|
|
Ar << bAnyNodeChild1;
|
|
Ar << bAnyNodeChild2;
|
|
|
|
const bool bIsLeafNode = !bAnyNodeChild1 && !bAnyNodeChild2;
|
|
if (bIsLeafNode)
|
|
{
|
|
check(KDTreeNode->node_type.lr.left <= AccessorTypeMax);
|
|
check(KDTreeNode->node_type.lr.right <= AccessorTypeMax);
|
|
|
|
AccessorType OffsetLeft = KDTreeNode->node_type.lr.left;
|
|
AccessorType OffsetRight = KDTreeNode->node_type.lr.right;
|
|
|
|
Ar << OffsetLeft;
|
|
Ar << OffsetRight;
|
|
|
|
KDTreeNode->node_type.lr.left = OffsetLeft;
|
|
KDTreeNode->node_type.lr.right = OffsetRight;
|
|
}
|
|
else
|
|
{
|
|
Ar << KDTreeNode->node_type.sub.divfeat;
|
|
Ar << KDTreeNode->node_type.sub.divhigh;
|
|
Ar << KDTreeNode->node_type.sub.divlow;
|
|
}
|
|
|
|
if (bAnyNodeChild1)
|
|
{
|
|
SerializeSubTree(Ar, KDTree, KDTreeNode->child1);
|
|
}
|
|
else if (Ar.IsLoading())
|
|
{
|
|
KDTreeNode->child1 = nullptr;
|
|
}
|
|
|
|
if (bAnyNodeChild2)
|
|
{
|
|
SerializeSubTree(Ar, KDTree, KDTreeNode->child2);
|
|
}
|
|
else if (Ar.IsLoading())
|
|
{
|
|
KDTreeNode->child2 = nullptr;
|
|
}
|
|
return Ar;
|
|
}
|
|
|
|
#endif // UE_POSE_SEARCH_USE_NANOFLANN
|
|
|
|
FArchive& Serialize(FArchive& Ar, FKDTree& KDTree, const float* KDTreeData)
|
|
{
|
|
#if UE_POSE_SEARCH_USE_NANOFLANN
|
|
check(!KDTree.KDTreeImplementation || KDTree.KDTreeImplementation->m_size <= AccessorTypeMax);
|
|
|
|
AccessorType KDTreeSize = KDTree.KDTreeImplementation ? KDTree.KDTreeImplementation->m_size : 0;
|
|
|
|
Ar << KDTreeSize;
|
|
|
|
if (KDTreeSize > 0)
|
|
{
|
|
if (Ar.IsLoading() && !KDTree.KDTreeImplementation)
|
|
{
|
|
KDTree.KDTreeImplementation = new FKDTreeImplementation(0, KDTree.DataSource, nanoflann::KDTreeSingleIndexAdaptorParams(0));
|
|
}
|
|
|
|
KDTree.KDTreeImplementation->m_size = KDTreeSize;
|
|
|
|
Ar << KDTree.KDTreeImplementation->dim;
|
|
|
|
AccessorType root_bbox_size = KDTree.KDTreeImplementation->root_bbox.size();
|
|
check(KDTree.KDTreeImplementation->root_bbox.size() <= AccessorTypeMax);
|
|
Ar << root_bbox_size;
|
|
|
|
if (Ar.IsLoading())
|
|
{
|
|
KDTree.DataSource.Data = KDTreeData;
|
|
KDTree.DataSource.PointDim = KDTree.KDTreeImplementation->dim;
|
|
KDTree.DataSource.PointCount = KDTree.KDTreeImplementation->m_size;
|
|
|
|
KDTree.KDTreeImplementation->root_bbox.resize(root_bbox_size);
|
|
}
|
|
|
|
for (FKDTreeImplementation::Interval& el : KDTree.KDTreeImplementation->root_bbox)
|
|
{
|
|
Ar.Serialize(&el, sizeof(FKDTreeImplementation::Interval));
|
|
}
|
|
|
|
check(KDTree.KDTreeImplementation->m_leaf_max_size <= AccessorTypeMax);
|
|
AccessorType KDTreeLeafMaxSize = KDTree.KDTreeImplementation->m_leaf_max_size;
|
|
Ar << KDTreeLeafMaxSize;
|
|
KDTree.KDTreeImplementation->m_leaf_max_size = KDTreeLeafMaxSize;
|
|
|
|
check(KDTree.KDTreeImplementation->vAcc.size() <= AccessorTypeMax);
|
|
AccessorType VAccSize = KDTree.KDTreeImplementation->vAcc.size();
|
|
Ar << VAccSize;
|
|
if (Ar.IsLoading())
|
|
{
|
|
KDTree.KDTreeImplementation->vAcc.resize(VAccSize);
|
|
}
|
|
for (AccessorType& el : KDTree.KDTreeImplementation->vAcc)
|
|
{
|
|
Ar << el;
|
|
}
|
|
SerializeSubTree(Ar, KDTree, KDTree.KDTreeImplementation->root_node);
|
|
}
|
|
else if (Ar.IsLoading())
|
|
{
|
|
KDTree.Reset();
|
|
}
|
|
#endif // UE_POSE_SEARCH_USE_NANOFLANN
|
|
|
|
return Ar;
|
|
}
|
|
|
|
} // namespace UE::PoseSearch
|