OpenCV使用 GenericIndex 进行 KNN 搜索

以前不了解 OpenCV 的 flann 模块，做 K-nearest-neighbour 搜索一直用的是 Github 上的一个很简单的实现：kd-tree。这个 KD tree 实现很简洁，仅有一个头文件，方便加入自己的工程，而且使用也特别方便。不过据说 cv::flann 实现的 KNN 速度更快，而且功能更加强大，所以大概看了一下怎么用。

鉴于目前网上相关资料很少，而且基本上能找到的都是关于已经被弃用的 cv::flann::Index，OpenCV 官方文档也缺乏相关说明，我在这里记录一下自己所了解的 cv::flann::GenericIndex 的用法（其实还有很多东西没搞明白，欢迎评论区交流）。

示例代码

// generate points
constexpr size_t num = 100;
constexpr int width = 640, height = 480;
vector<Point2f> pts;
Mat image(height, width, CV_8UC3, cv::Scalar::all(255));
{
    
    mt19937 gen(0);
    uniform_real_distribution<float> dis_x(0.f, float(width));
    uniform_real_distribution<float> dis_y(0.f, float(height));
    pts.reserve(num);
    for (size_t i = 0; i < num; ++i) {
    
        pts.emplace_back(dis_x(gen), dis_y(gen));
        circle(image, pts.back(), 2, cv::Scalar::all(0), cv::FILLED);
    }
}
cout << "num of points: " << pts.size() << '\n';
Point2f query(width / 2.f, height / 2.f);
cout << "query: " << query.x << " " << query.y << '\n';
circle(image, query, 2, {
     0, 0, 255 }, cv::FILLED);

// build kd-tree
auto kdtree = flann::GenericIndex(Mat(pts).reshape(1), cvflann::KDTreeIndexParams {
     2 }, flann::L2<float> {
    });

// knn search
constexpr int K = 4;
vector<int> indices(K);
vector<float> dists(K);
kdtree.knnSearch({
     query.x, query.y }, indices, dists, K, cvflann::SearchParams {
    });
cout << " nearest " << K << ": " << Mat(indices).t() << endl;
{
    
    Mat result = image.clone();
    for (int i : indices) {
    
        line(result, query, pts[i], {
     0, 255, 0 });
        circle(result, pts[i], 2, {
     255, 0, 0 }, cv::FILLED);
    }
    imwrite("knn_result.png", result);
}

// radius search
constexpr float radius = 50.f;
constexpr size_t max_num = 10;
indices.resize(max_num, -1);
dists.resize(max_num);
kdtree.radiusSearch({
     query.x, query.y }, indices, dists, radius * radius, cvflann::SearchParams {
    });
cout << "in radius " << radius << ": " << Mat(indices).t() << endl;
{
    
    Mat result = image.clone();
    for (int i : indices) {
    
        if (i < 0)
            break;
        line(result, query, pts[i], {
     0, 255, 0 });
        circle(result, pts[i], 2, {
     255, 0, 0 }, cv::FILLED);
    }
    circle(result, query, radius, {
     0, 0, 255 });
    imwrite("radius_result.png", result);
}

这个例子是在大小 640×480 图像上随机生成 100 个点，然后利用 KD tree 查找距离图像中心最近的 4 个点（KNN search），以及与图像中心距离小于 50 像素的点（radius search）：

flann::GenericIndex 支持暴力搜索 (LinearIndexParams)、KD tree (KDTreeIndexParams)、层级聚类 (HierarchicalClusteringIndexParams) 等搜索策略，同时支持 L1、L2、Hamming 等距离计算方法，功能确实很强大，参见 OpenCV Docs。

需要注意的是，在使用 L2 距离的情况下， radiusSearch 中的 radius 参数其实是半径的平方（所以代码中写的是 radius*radius），参见 issue#12683。