在工作时需要用到dbscan算法来对点云进行聚类,在网上搜索到的版本都是在cpu上进行处理,速度很慢,因此自己写了一个cuda版本,供大家参考。经实测,3060显卡上5000个点云处理速度为20ms以内。

代码已经封装好,可以直接调用。

dbscan.h

#ifndef DBSCAN_H
#define DBSCAN_H

#include <iostream>
#include <vector>
#include <cmath>
#include <algorithm>
#include <thrust/device_vector.h>

using namespace std;

class DBScan
    {
        public:
        DBScan();
        ~DBScan();
        float calculateDistance(float p1_x, float p1_y, float p1_z, float p2_x, float p2_y, float p2_z);
        void getNeighborhood(const float* points_cpu, int points_num, int *points_id_gpu, int *points_id_num,
         int pIdx, float epsilon);
        void getAllNeighborhood(const float* points, int points_num, int *points_id_all, int *points_id_num_all, float epsilon);
        void expandAllCluster(const int* points_id_all, int *points_id_num_all, int clusterId, int pIdx, float epsilon,
            int minPts, int* clusterLabels_cpu);
        void dbscanCluster(const float* points_cpu, int points_num, float epsilon, int minPts, std::vector<int>& clusterLabels);
        void expandAllCluster(const int* points_id_all, int *points_id_num_all, int clusterId, int pIdx, float epsilon,
            int minPts, std::vector<int>& clusterLabels); 

        public:
        int *outsnum_gpu_;
        int *points_id_gpu_;
        int *points_id_cpu_;
        int *clusterLabels_cpu_;
        int *clusterLabels_gpu_;
        int *points_id_all_gpu_;
        int *points_id_all_cpu_;
        int *points_id_num_all_gpu_;
        int *points_id_num_all_cpu_;
        int max_points_num_ = 5000; //这里修改最大点云数量


    };


#endif // !DBSCAN_H

dbscan.cu

#include "dbscan.h"

    DBScan::DBScan()
    {
        cudaMallocHost((void **)&clusterLabels_cpu_, max_points_num_ * sizeof(int));
        cudaMalloc((void **)&points_id_all_gpu_, max_points_num_ * max_points_num_ * sizeof(int));
        cudaMalloc((void **)&points_id_num_all_gpu_, max_points_num_ * sizeof(int));
        cudaMallocHost((void **)&points_id_all_cpu_, max_points_num_ * max_points_num_ * sizeof(int));
        cudaMallocHost((void **)&points_id_num_all_cpu_, max_points_num_ * sizeof(int));
        

    }
    DBScan::~DBScan()
    {
        cudaFree(points_id_all_gpu_);
        cudaFree(points_id_num_all_gpu_);
        cudaFreeHost(points_id_all_cpu_);
        cudaFreeHost(points_id_num_all_cpu_);
        cudaFreeHost(clusterLabels_cpu_);
    }

    __global__ void cuda_getAllNeighborhood(const float* points, int points_num, 
     float epsilon, int *points_id_all, int *points_id_num_all, int max_points_num)
    {
        int id_1 = blockDim.x * blockIdx.x + threadIdx.x;
        int id_2 = blockDim.y * blockIdx.y + threadIdx.y;
        
		if (id_1 >= points_num||id_2 >= points_num||id_1==id_2)
		{
			return;
		}
        

        float point1_x = points[4*id_1];
        float point1_y = points[4*id_1+1];
        float point1_z = points[4*id_1+2];

        float point2_x = points[4*id_2];
        float point2_y = points[4*id_2+1];
        float point2_z = points[4*id_2+2];

        float dist = sqrt((point1_x-point2_x)*(point1_x-point2_x)+(point1_y-point2_y)*(point1_y-point2_y)+
         (point1_z-point2_z)*(point1_z-point2_z));

        if (dist > epsilon)
            return;
        int cur_idx_1 = atomicAdd(&points_id_num_all[id_1], (int)1);       
        points_id_all[id_1*max_points_num+cur_idx_1] = id_2;
            
    }

    // 获取给定点的邻域点
    void DBScan::getAllNeighborhood(const float* points, int points_num, int *points_id_all, int *points_id_num_all, float epsilon) 
    {

        dim3 thread1d(32, 32, 1);
		dim3 block1d(0, 0, 1);
		block1d.x = (int)((points_num + thread1d.x - 1) / thread1d.x);
        block1d.y = (int)((points_num + thread1d.y - 1) / thread1d.y);
        
		cuda_getAllNeighborhood<<<block1d, thread1d>>>(points, points_num, epsilon, points_id_all, points_id_num_all, max_points_num_);
        
		cudaDeviceSynchronize();
        
    }

    // 执行DBScan聚类
    void DBScan::dbscanCluster(const float* points, int points_num, float epsilon, int minPts, std::vector<int>& clusterLabels) 
    {
        int clusterId = 0;
        cudaMemset(points_id_num_all_gpu_, 0, sizeof(int)*points_num);
        getAllNeighborhood(points, points_num, points_id_all_gpu_, points_id_num_all_gpu_, epsilon);
        cudaMemcpy(points_id_num_all_cpu_, points_id_num_all_gpu_, sizeof(int)*points_num, cudaMemcpyDeviceToHost);
        cudaMemcpy(points_id_all_cpu_, points_id_all_gpu_, sizeof(int)*points_num*max_points_num_, cudaMemcpyDeviceToHost);
        
        for (int i = 0; i < points_num; ++i) 
        {
            if (clusterLabels[i] != 0) 
            {  // 已经被标记为噪声点或已被聚类
                continue;
            }

            if (points_id_num_all_cpu_[i] < minPts)
            {
                clusterLabels[i] = -1;
            }
            else
            {
                ++clusterId;
                expandAllCluster(points_id_all_cpu_, points_id_num_all_cpu_, clusterId, i, epsilon, minPts, clusterLabels);

            }
        }
    }

    void DBScan::expandAllCluster(const int* points_id_all, int *points_id_num_all, int clusterId, int pIdx, float epsilon,
     int minPts, std::vector<int>& clusterLabels) 
    {
        clusterLabels[pIdx] = clusterId;  // 将当前点标记为聚类
        for (int j = 0; j < points_id_num_all[pIdx]; j++)
        {
            int neighborIdx = points_id_all[pIdx*max_points_num_+j];
            if (clusterLabels[neighborIdx] == -1 || clusterLabels[neighborIdx] == 0) 
            {  
                // 如果邻域点是噪声点则重新分类到当前聚类
                clusterLabels[neighborIdx] = clusterId;
                if (points_id_num_all[neighborIdx] > 0)
                {
                    expandAllCluster(points_id_all, points_id_num_all, clusterId, neighborIdx, epsilon, minPts, clusterLabels);
                }
            }
        }

    }

代码解读:先使用cuda核函数计算所有点云相互之间的距离值,当距离值小于设定的阈值时,则两个点云之间互为邻域点。将所有点云的领域点下标保存下来,放在一个max_points_num*max_points_num的矩阵中,矩阵的第一行保存的为下标为1的点云的所有领域点的下标,第二行保存的为下标为2的点云的所有领域点的下标,以此类推。

后面就是正常的dbscan聚类流程了,在此不做赘述。

Logo

欢迎来到FlagOS开发社区,这里是一个汇聚了AI开发者、数据科学家、机器学习爱好者以及业界专家的活力平台。我们致力于成为业内领先的Triton技术交流与应用分享的殿堂,为推动人工智能技术的普及与深化应用贡献力量。

更多推荐