点云目标检测 -- 目标框内点的个数 -- python代码

此程序是计算目标框内点的个数，此段代码摘自second目标检测算法，把数据格式对应好便可直接运行。

输入数据为：
pcd格式的点云数据（x,y,z,i）
csv格式的标签数据（包括x,y,z,l,w,h,angle等）

注：部分代码写法繁琐，可自行简化，且速度不是太快，可自行加速。

import numpy as np
import os
import csv
import math
import numba
import pdb
PI_rads = math.pi/180

def create_data_infos(labelfile, lidar_path ):
    infos = {}
    gt_boxes = []
    points = []
    label_fpath = os.path.join(label_path,labelfile)
    with open(lidar_path, 'r') as f:
        line = f.readline().strip()
        while line:
            linestr = line.split(" ")
            if len(linestr) == 4:
                linestr_convert = list(map(float, linestr))
                points.append(linestr_convert)
            line = f.readline().strip()
    infos['points'] = np.array(points)
    with open(label_fpath,'r') as f:
        reader=csv.reader(f)
        for i in reader:
            #print(i)
            if int(i[1])==1:
                continue
            gt_boxes.append(
                [float(i[0])/100,float(i[1])/100,float(i[2])/100,float(i[3])/100,float(i[4])/100,float(i[5])/100, float(i[6])*PI_rads])   ##依次为x,y,z,l,w,h,angle，单位为厘米和角度，请对照自己的标签格式自行修改
    gt_boxes = np.array(gt_boxes)
    infos['gt_boxes'] = gt_boxes
    return infos

def corners_nd(dims, origin):
    ndim = int(dims.shape[1])
    corners_norm = np.stack(
        np.unravel_index(np.arange(2**ndim), [2] * ndim),
        axis=1).astype(dims.dtype)
    if ndim == 2:
        # generate clockwise box corners
        corners_norm = corners_norm[[0, 1, 3, 2]]
    elif ndim == 3:
        corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]]  
    corners_norm = corners_norm - np.array(origin, dtype=dims.dtype)
    corners = dims.reshape([-1, 1, ndim]) * corners_norm.reshape(
        [1, 2**ndim, ndim])
    return corners

def rotation_3d_in_axis(points, angles, axis=0):
    #print("points is :", points)
    #("angles is :", angles)
    rot_sin = np.sin(angles)
    rot_cos = np.cos(angles)
    ones = np.ones_like(rot_cos)
    zeros = np.zeros_like(rot_cos)
    if axis == 1:
        rot_mat_T = np.stack([[rot_cos, zeros, -rot_sin], [zeros, ones, zeros],
                              [rot_sin, zeros, rot_cos]])
    elif axis == 2 or axis == -1:
        rot_mat_T = np.stack([[rot_cos, rot_sin, zeros],
                              [-rot_sin, rot_cos, zeros], [zeros, zeros, ones]])
    elif axis == 0:
        rot_mat_T = np.stack([[zeros, rot_cos, -rot_sin],
                              [zeros, rot_sin, rot_cos], [ones, zeros, zeros]])
    else:
        raise ValueError('axis should in range')
    new_points = np.einsum('aij,jka->aik', points, rot_mat_T)
    return new_points

def center_to_corner_box3d(centers,
                           dims,
                           angles=None,
                           origin = None,
                           axis=1):
    corners = corners_nd(dims, origin=origin)
    # corners: [N, 8, 3]
    if angles is not None:
        corners = rotation_3d_in_axis(corners, angles, axis=axis)
    corners += centers.reshape([-1, 1, 3])
    for i in range(len(corners)):
        z_h = max(corners[i, :, -1])
        z_0 = min(corners[i, :, -1])
        z_max =z_h-(z_h-z_0)/2
        z_min = z_0 -(z_h-z_0) / 2
        corners[i, [0, 3, 4, 7], -1] = z_min
        corners[i, [1, 2, 5, 6], -1] = z_max
    return corners

def corner_to_surfaces_3d(corners):
    surfaces1 = np.array([
        [corners[:, 0], corners[:, 1], corners[:, 2], corners[:, 3]],
        [corners[:, 7], corners[:, 6], corners[:, 5], corners[:, 4]],
        [corners[:, 0], corners[:, 3], corners[:, 7], corners[:, 4]],
        [corners[:, 1], corners[:, 5], corners[:, 6], corners[:, 2]],
        [corners[:, 0], corners[:, 4], corners[:, 5], corners[:, 1]],
        [corners[:, 3], corners[:, 2], corners[:, 6], corners[:, 7]],
    ])
    surfaces = surfaces1.transpose([2, 0, 1, 3])
    return surfaces

def surface_equ_3d(polygon_surfaces):
    surface_vec = polygon_surfaces[:, :, :2, :] - \
        polygon_surfaces[:, :, 1:3, :]
    normal_vec = np.cross(surface_vec[:, :, 0, :], surface_vec[:, :, 1, :])
    d = np.einsum('aij, aij->ai', normal_vec, polygon_surfaces[:, :, 0, :])
    return normal_vec, -d


def points_in_convex_polygon_3d_jit(points,
                                    polygon_surfaces,
                                    num_surfaces=None):
    num_polygons = polygon_surfaces.shape[0]            ##目标框个数
    if num_surfaces is None:
        num_surfaces = np.full((num_polygons, ), 9999999, dtype=np.int64)
    normal_vec, d = surface_equ_3d(polygon_surfaces[:, :, :3, :])                  ##求面的法向量
    return _points_in_convex_polygon_3d_jit(points, polygon_surfaces,
                                            normal_vec, d, num_surfaces)          ##根据长方体内部点与面的法向量之间的关系求得在目标框内部的点

@numba.njit
def _points_in_convex_polygon_3d_jit(points, polygon_surfaces, normal_vec, d,
                                     num_surfaces):
    ##polygon_surfaces：N个目标框的面，每个框6个面
    ## polygon_surfaces.shape[1:3] 一个长方体6个面 每个面4个点
    max_num_surfaces, max_num_points_of_surface = polygon_surfaces.shape[1:3]
    num_points = points.shape[0]
    num_polygons = polygon_surfaces.shape[0]
    ret = np.ones((num_points, num_polygons), dtype=np.bool_)
    sign = 0.0
    for i in range(num_points):
        for j in range(num_polygons):
            for k in range(max_num_surfaces):
                if k > num_surfaces[j]:
                    break
                sign = (
                    points[i, 0] * normal_vec[j, k, 0] +
                    points[i, 1] * normal_vec[j, k, 1] +
                    points[i, 2] * normal_vec[j, k, 2] + d[j, k])
                if sign >= 0:
                    ret[i, j] = False
                    break
    return ret

def points_in_rbbox(points, gt_boxes):
    axis = 2             
    origin = (0.5, 0.5, 0)
    rbbox_corners = center_to_corner_box3d(
        gt_boxes[:, :3], gt_boxes[:, 3:6], gt_boxes[:, 6], origin=origin, axis=axis)    ##由中心点左边及长宽高计算出长方体角点
    surfaces = corner_to_surfaces_3d(rbbox_corners)       
    indices = points_in_convex_polygon_3d_jit(points[:, :3], surfaces)    ##计算长方体面内点的个数
    return indices

if __name__=="__main__":
    root_path=r'path/to/your/folder'
    label_path=os.path.join(root_path,'label')       ##csv格式标签所在文件夹
    data_path=os.path.join(root_path,'lidar')        ##pcd格式的激光雷达数据所在文件夹

    labelfiles = os.listdir(label_path)
    lidarfiles = os.listdir(data_path)
    labelfiles.sort()
    lidarfiles.sort()
    for labelfile in labelfiles:
        lidar_path= data_path+'/'+labelfile.split('.')[0]+'.pcd'
        infos = create_data_infos(labelfile, lidar_path)
        points = infos['points']
        gt_boxes = infos['gt_boxes']
        num_obj = gt_boxes.shape[0]
        point_indices = points_in_rbbox(points, gt_boxes)
        obj_indices = []
        for i in range(num_obj):
            gt_points = points[point_indices[:, i]]
            gt_points[:, :3] -= gt_boxes[i, :3]
            num_points_in_gt = gt_points.shape[0]           ##每个目标框内点的个数
            if num_points_in_gt >= 6:                       ##保留6个点以上的目标框
                obj_indices.append(i)
        gt_boxes_filterd = gt_boxes[obj_indices]            ##过滤后每帧点云的目标框，可自行保存到指定路径