This repository contains a C++ implementation of a k-dimensional tree (k-d tree) specifically designed for handling 256-dimensional embeddings. The implementation is tailored for applications requiring efficient nearest neighbor searches in high-dimensional spaces, such as image processing and machine learning tasks in the Nvidia DeepStream framework.
- K-D Tree Implementation: A basic, yet efficient k-d tree for handling 256-dimensional data points.
- Nearest Neighbor Search: Quickly find the closest point in the tree to a given 256-dimensional point.
- Embedding to ID Mapping: Efficiently map a 256-dimensional embedding to a corresponding ID.
Clone the repository to your local machine:
git clone https://github.com/KamoliddinS/EmbeddingKDTree.git
cd EmbeddingKDTreeInclude the kdtree.hpp in your C++ project. Here's a simple example to demonstrate usage:
#include "kdtree.hpp"
int main() {
// Create a k-d tree instance
KDTree tree;
// Insert embeddings with corresponding IDs
tree.insert({/* 256-dimensional embedding here */}, /* ID here */);
// Search for the nearest embedding
int nearestId = tree.findNearest({/* 256-dimensional embedding here */});
return 0;
}Note: This project created in mind for reid applications, https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tao/models/reidentificationnet
To integrate this k-d tree into the Nvidia DeepStream framework, ensure thread safety and consider GPU optimization based on your specific use case.
Contributions to enhance this implementation, including performance optimizations, are welcome. Please follow the standard pull request process for this repository.