c++实现基于kd树的3d点云半径滤波 不使用pcl库

c++实现基于kd树的3d点云半径滤波 不使用pcl库

可优化的地方挺多的 是不使用kd树的运算时间的大约四分之一

#include <iostream>
#include <fstream>
#include <vector>
#include <cmath>
#include <algorithm>
#include <chrono>
 
struct Point3D
{
    float x, y, z;
};
 
struct TreeNode
{
    Point3D point;
    int split_dim;
    TreeNode* left_child;
    TreeNode* right_child;
    TreeNode(const Point3D& p) : point(p), split_dim(0), left_child(nullptr), right_child(nullptr) {}
};
 
bool compare_x(const Point3D& p1, const Point3D& p2)
{
    return p1.x < p2.x;
}
 
bool compare_y(const Point3D& p1, const Point3D& p2)
{
    return p1.y < p2.y;
}
 
bool compare_z(const Point3D& p1, const Point3D& p2)
{
    return p1.z < p2.z;
}
 
float distance(const Point3D& p1, const Point3D& p2)
{
    return std::sqrt(std::pow(p1.x - p2.x, 2) + std::pow(p1.y - p2.y, 2) + std::pow(p1.z - p2.z, 2));
}
 
void build_kd_tree(TreeNode*& node, std::vector<Point3D>& pointclouds, int start_index, int end_index, int depth)
{
    if (start_index > end_index) {
        node = nullptr;
        return;
    }
    int middle_index = start_index + (end_index - start_index) / 2;
    if (depth % 3 == 0) {
        std::sort(pointclouds.begin() + start_index, pointclouds.begin() + end_index + 1, compare_x);
    }
    else if (depth % 3 == 1) {
        std::sort(pointclouds.begin() + start_index, pointclouds.begin() + end_index + 1, compare_y);
    }
    else {
        std::sort(pointclouds.begin() + start_index, pointclouds.begin() + end_index + 1, compare_z);
    }
    node = new TreeNode(pointclouds[middle_index]);
    node->split_dim = depth % 3;
    build_kd_tree(node->left_child, pointclouds, start_index, middle_index - 1, depth + 1);
    build_kd_tree(node->right_child, pointclouds, middle_index + 1, end_index, depth + 1);
}
 
void search_knn(TreeNode* node, const Point3D& p, float radius, int k, std::vector<int>& indices, std::vector<float>& distances)
{
    if (node == nullptr) {
        return;
    }
    float dist = distance(node->point, p);
    if (dist <= radius) {
        indices.push_back(indices.size());
        distances.push_back(dist);
    }
    if (indices.size() < k || dist < distances[indices.size() - 1]) {
        if (node->split_dim == 0) {
            if (p.x < node->point.x) {
                search_knn(node->left_child, p, radius, k, indices, distances);
                search_knn(node->right_child, p, radius, k, indices, distances);
            }
            else {
                search_knn(node->right_child, p, radius, k, indices, distances);
                search_knn(node->left_child, p, radius, k, indices, distances);
            }
        }
        else if (node->split_dim == 1) {
            if (p.y < node->point.y) {
                search_knn(node->left_child, p, radius, k, indices, distances);
                search_knn(node->right_child, p, radius, k, indices, distances);
            }
            else {
                search_knn(node->right_child, p, radius, k, indices, distances);
                search_knn(node->left_child, p, radius, k, indices, distances);
            }
        }
        else {
            if (p.z < node->point.z) {
                search_knn(node->left_child, p, radius, k, indices, distances);
                search_knn(node->right_child, p, radius, k, indices, distances);
            }
            else {
                search_knn(node->right_child, p, radius, k, indices, distances);
                search_knn(node->left_child, p, radius, k, indices, distances);
            }
        }
 
    }
}
 
int main()
{
    // 获取当前系统时间点
    auto start_time = std::chrono::high_resolution_clock::now();
 
    float radius = 0.25f;
    int k = 10;
 
    // 读取点云数据
    std::vector<Point3D> pointclouds;
    std::ifstream infile(".txt");
    if (!infile) {
        std::cerr << "Failed to open file: filtered_points5.txt" << std::endl;
        return -1;
    }
    Point3D p;
    while (infile >> p.x >> p.y >> p.z) {
        pointclouds.emplace_back(p);
    }
 
    // 构建kd树
    TreeNode* root = nullptr;
    build_kd_tree(root, pointclouds, 0, pointclouds.size() - 1, 0);
 
    // 搜索邻域点
    std::vector<Point3D> filtered_cloud;
    for (int i = 0; i < pointclouds.size(); i++) {
        std::vector<int> indices;
        std::vector<float> distances;
        search_knn(root, pointclouds[i], radius, k, indices, distances);
        if (indices.size() >= k) {
            filtered_cloud.emplace_back(pointclouds[i]);
        }
    }
 
    std::ofstream outFile(".txt");
    for (int i = 0; i < filtered_cloud.size(); ++i)
    {
        outFile << filtered_cloud[i].x << " " << filtered_cloud[i].y << " " << filtered_cloud[i].z << std::endl;
    }
 
    outFile.close();
 
    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
    std::cout << "Execution time: " << duration.count() << " microseconds" << std::endl;
 
    return 0;
}