sample.cpp

#include <opencv2/opencv.hpp>
#include <iostream>
#include <array>
#include <vector>

#include "kdtree.h"

// user-defined point type
// inherits std::array in order to use operator[]
class MyPoint : public std::array<double, 2>
{
public:

	// dimension of space (or "k" of k-d tree)
	// KDTree class accesses this member
	static const int DIM = 2;

	// the constructors
	MyPoint() {}
	MyPoint(double x, double y)
	{ 
		(*this)[0] = x;
		(*this)[1] = y;
	}

	// conversion to OpenCV Point2d
	operator cv::Point2d() const { return cv::Point2d((*this)[0], (*this)[1]); }
};

void kdtree_knn_dump(struct kdtree *tree, double *candidates, int k) 
{
        int i, j = 0;
        struct knn_list *p = tree->knn_list_head.next;
        const int dim = tree->dim;
        printf("The nearest %d samples are as follows:\n", k);
        while (p != &tree->knn_list_head) {
                putchar('(');
                for (i = 0; i < tree->dim; i++) {
                        if (i == tree->dim - 1) {
                                printf("%.2lf) Distance:%lf\n", p->node->coord[i], sqrt(p->distance));
                        } else {
                                printf("%.2lf, ", p->node->coord[i]);
                        }
                        candidates[j*dim+i] = p->node->coord[i];
                }
                p = p->next;
                j++;
        }
}

int main(int argc, char **argv)
{
	const int seed = argc > 1 ? std::stoi(argv[1]) : 0;
	srand(seed);

	// generate space
	const int width = 500;
	const int height = 500;
	cv::Mat img = cv::Mat::zeros(cv::Size(width, height), CV_8UC3);

	// generate points
	const int npoints = 1000;
	std::vector<MyPoint> points(npoints);
	for (int i = 0; i < npoints; i++)
	{
		const int x = rand() % width;
		const int y = rand() % height;
		points[i] = MyPoint(x, y);
	}

	for (const auto& pt : points)
		cv::circle(img, cv::Point2d(pt), 1, cv::Scalar(0, 255, 255), -1);

	// build k-d tree
	int dim = 2;
	struct kdtree *tree = kdtree_init(dim);
	for (int i = 0; i < npoints; i++)
	{
		// 1. Declare an array to hold the coordinates
		double my_coord[] = {points[i][0], points[i][1]};

		// 2. Pass the array (which acts as a pointer) to the function
		kdtree_insert(tree, my_coord);	
	}
	kdtree_rebuild(tree);

	// generate query (center of the space)
	const int k = 100;
	const MyPoint query(0.5 * width, 0.5 * height);
	cv::circle(img, cv::Point2d(query), 1, cv::Scalar(0, 0, 255), -1);

	// nearest neigbor search
	const cv::Mat I0 = img.clone();
	double target1[] = {query[0], query[1]};
	kdtree_knn_search(tree, target1, 1);
	// kdtree_knn_dump(tree, nullptr, k);
	// const int idx = kdtree.nnSearch(query);
	// cv::circle(I0, cv::Point2d(points[idx]), 1, cv::Scalar(255, 255, 0), -1);
	// cv::line(I0, cv::Point2d(query), cv::Point2d(points[idx]), cv::Scalar(0, 0, 255));

	// k-nearest neigbors search
	const cv::Mat I1 = img.clone();
	double candidates[k*dim];
	kdtree_knn_search(tree, target1, k);
	kdtree_knn_dump(tree, candidates, k);
	for (int i = 0; i < k; i++){
		cv::circle(I1, cv::Point2d(candidates[i*dim], candidates[i*dim + 1]), 1, cv::Scalar(255, 255, 0), -1);	
		cv::line(I1, cv::Point2d(query), cv::Point2d(candidates[i*dim], candidates[i*dim + 1]), cv::Scalar(0, 0, 255));
	}
		
	// const std::vector<int> knnIndices = kdtree.knnSearch(query, k);
	// for (int i : knnIndices)
	// {
	// 	cv::circle(I1, cv::Point2d(points[i]), 1, cv::Scalar(255, 255, 0), -1);
	// 	cv::line(I1, cv::Point2d(query), cv::Point2d(points[i]), cv::Scalar(0, 0, 255));
	// }
	
	// radius search
	const cv::Mat I2 = img.clone();
	// const double radius = 50;
	// const std::vector<int> radIndices = kdtree.radiusSearch(query, radius);
	// for (int i : radIndices)
	// 	cv::circle(I2, cv::Point2d(points[i]), 1, cv::Scalar(255, 255, 0), -1);
	// cv::circle(I2, cv::Point2d(query), cvRound(radius), cv::Scalar(0, 0, 255));

	// show results
	cv::imshow("Nearest neigbor search", I0);
	// cv::imshow("K-nearest neigbors search (k = 10)", I1);
	cv::imshow("K-nearest neighbors search (k = " + std::to_string(k) + ")", I1);
	cv::imshow("Radius search (radius = 50)", I2);

	cv::waitKey(0);

	return 0;
}