基于KDTree的点云外包盒提取

解决问题

使用 RandLA-Net 框架得到结果，从中提取目标外包盒，并去除误判的点云块。

参考方法

链接: link

代码

Ubuntu环境，代码是基于RandLA-Net 代码片.

from helper_ply import read_ply   				# 函数来自RandLA-Net
from collections import Counter
from sklearn.neighbors import KDTree
from helper_tool import DataProcessing as DP	# 函数来自RandLA-Net
from sklearn.cluster import KMeans
import cv2
import open3d
import numpy as np
from open3d import read_point_cloud


def tow_array_sorted(label, xyz_coordinate):
	# 根据label值排序
    sorted_indices = label.argsort()  # 排序索引
    xyz_sorted = xyz_coordinate[sorted_indices]
    list_numbers = Counter(label)       # {0.0: 40942368, 1.0: 60380, 4.0: 46125, 3.0: 26614, 2.0: 7047}
    return xyz_sorted, list_numbers


def KDtreeRecursion(xyz, P_vector=[], Q_vector=None, kdtree_k=80,  kdtree_distance_max=15.0):
	'''
    基于KDTree的递归聚类点云
    xyz: 输入点云
    P_vector: 已经用KDTree搜索过的点的索引组成的向量
    Q_vector:已经聚类到一起的点索引
    kdtree_k: KDTree搜索点个个数
    kdtree_distance_max: 距离阈值
    :return: Q_vector
    '''
    global Q_vector1
    tree = KDTree(xyz)
    if Q_vector is None:	# 第一次Q_vector为空，需要赋值
        xyz_distance_indices = tree.query(xyz[:1], k=kdtree_k)
        Q_vector = xyz_distance_indices[1][np.where(xyz_distance_indices[0] < kdtree_distance_max)]
    in_Q_out_P = np.array(list(set(Q_vector) - set(P_vector)))	# 得到在Q集合，不在P集合的点索引
    Q_vector_add = Q_vector
    for indx in in_Q_out_P:        # 遍历所有点                                         # ##################
        P_vector = np.append(P_vector, indx)
        xyz_distance_indices = tree.query(xyz[indx:indx+1], k=kdtree_k)
        Q_vector_new = xyz_distance_indices[1][np.where( xyz_distance_indices[0] < kdtree_distance_max)]
        Q_vector_add = np.unique(np.append(Q_vector_add, Q_vector_new))
    if len(Q_vector_add) != len(Q_vector):	# 判断是否递归
        Q_vector = Q_vector_add
        KDtreeRecursion(xyz, P_vector, Q_vector)
    else:
        Q_vector1 = Q_vector
        # print('长度：', xyz.shape, len(np.unique(Q_vector)))
    return Q_vector1


def rectrangle_pcd(xyz):
    '''
    根据坐标值绘制长方体
    :return:
    '''
    rectrangle_points = np.empty([8,3], dtype=int)
    ind = 0
    for i in range(2):
        for j in range(2):
            for k in range(2):
                point = [xyz[:,0][i],xyz[:,1][j],xyz[:,2][k]]
                rectrangle_points[ind,:] = point
                ind = ind + 1
    lines = [[0, 1], [1, 3], [3, 2],[2, 0], [4, 5], [5, 7],[7, 6], [6, 4], [0, 4],[1, 5], [2, 6], [3, 7]]  # Right leg
    colors = [[0, 0, 0] for i in range(len(lines))]  # Default blue

    # 定义三角形三条连接线
    line_pcd = open3d.LineSet()
    line_pcd.lines = open3d.Vector2iVector(lines)
    line_pcd.colors = open3d.Vector3dVector(colors)
    line_pcd.points = open3d.Vector3dVector(rectrangle_points)

    return line_pcd


# 读取ply点云
ply_file_path = "/root/results/Area_4_*_1.ply"
ply_points = read_ply(ply_file_path)		
xyz = np.vstack((ply_points['x'], ply_points['y'], ply_points['z'])).T
predict_class = ply_points['class'].astype(int)
xyz_sorted, list_numbers = tow_array_sorted(predict_class, xyz)
unique_class = np.unique(predict_class)
pcd = []
Q_vector1 = None
for i in range(2):
    print(i, '\t', list_numbers[i])
    if i == 0:
        temp_number = list_numbers[i]
        continue
    else:
        temp_xyz = xyz_sorted[temp_number:, :]
        temp_xyz = DP.grid_sub_sampling(points=temp_xyz, grid_size=0.16)
        print(len(temp_xyz))
        # KDTree
        tree = KDTree(temp_xyz)
        while len(temp_xyz)>100:
            xyz_indx = KDtreeRecursion(temp_xyz)
            xyz = temp_xyz[xyz_indx]
            temp_xyz = np.delete(temp_xyz, xyz_indx, axis=0)
            if len(xyz_indx)>800:
                xyz_max = np.max(xyz, axis=0)
                xyz_min = np.min(xyz, axis=0)
                min_max_xyz = np.array([xyz_min, xyz_max])
                line_pcd = rectrangle_pcd(min_max_xyz)
                pcd.append(line_pcd)
# 方法1（非阻塞显示）
ply_file_path1 = "/root/results/*.ply"
point_pcd = read_point_cloud(ply_file_path1)
vis = open3d.Visualizer()
vis.create_window(window_name='Open3D_1')
vis.get_render_option().point_size = 10  # 设置点的大小
# 先把点云对象添加给Visualizer
vis.add_geometry(point_pcd)
for i in range(len(pcd)):
    vis.add_geometry(pcd[i])

while True:
    vis.poll_events()
    vis.update_renderer()
    cv2.waitKey(100)

效果图