python: lidar点云转BEV投影及pillar/voxel

bev投影

点云bev投影代码

pillar/voxel

#!/usr/bin/env python
# -*- encoding: utf-8 -*-



import os
import numpy as np
import math
import pickle


n_max_points_per_pillar = 32
n_max_pillars = 20736  # 144*144 = 20736
n_in_features = 7  # x y z i xc yc zc

x_step = 0.32  # 32cm 32*144 = 4608 = 46.08m
y_step = 0.16  # 16cm 16*144 =  = 15.36m

max_intensity = 255

def read_pcd(pcd_file, readline=11):
    # np.fromfile read pcd 文本/二进制
    # cloud = np.fromfile(str(pcd_file), dtype=np.float32, count=-1)
    data = []
    with open(pcd_file, 'r') as f:
        lines = f.readlines()[readline:]
        for line in lines:
            line = np.array(line.strip('\n').split(' '), dtype=np.float32)
            # print(line.shape, line)
            x,y,z,i = line[0], line[1], line[2], line[3]
            # print(x,y,z,i)
            # r = math.sqrt(x**2 + y**2 + z**2) * i  # for rslidar
            # r = float(line[4]) # for ouster lidar
            data.append([x,y,z,i])
        # points = list(map(lambda line: list(map(lambda x: float(x), line.split(' '))), lines))
        points = np.array(data)

    return points

def points_to_pillars(points, list_roi_xyz = [0.0, 69.12, -11.52, 11.52, -2.5, 0.5],
                          list_pillar_xy = [0.48, 0.16],
                          list_grid_xy = [144, 144],  # (69.12 - 0) / 0.48
                          list_pillar_szie = [20736, 32, 7]
                          ):
    x_min, x_max, y_min, y_max, z_min, z_max = list_roi_xyz
    x_grid, y_grid = list_grid_xy
    x_pillar, y_pillar = list_pillar_xy
    n_max_pillars, n_max_points_per_pillar, n_in_features = list_pillar_szie

    # ROI
    idx = np.where((points[:, 0] > x_min) & (points[:, 0] < x_max) &
                   (points[:, 1] > y_min) & (points[:, 1] < y_max) &
                   (points[:, 2] > z_min) & (points[:, 2] < z_max))
    points_roi = points[idx[0]]
    # intensity normlize
    points_roi[:,3] /= 255

    # 计算每个点所在的voxel索引
    # Compute the indices of the pillars that each point belongs to.
    x_img = (points_roi[:,0] - x_min) // x_pillar
    y_img = (points_roi[:,1] - y_min) // y_pillar
    x_img = x_img.reshape(-1, 1).astype(int)
    y_img = y_img.reshape(-1, 1).astype(int)
    # pillar_coords = np.concatenate((x_img, y_img), axis=1)

    # 计算pillar的索引
    pillar_idx = x_img * y_grid + y_img

    # 找到每个pillar的点数
    unique_pillars, num_points_per_pillars = np.unique(pillar_idx, return_counts=True)
    # print("num pillars: ",unique_pillars, num_points_per_pillars)

    num_points_per_pillars = np.minimum(num_points_per_pillars, n_max_points_per_pillar)

    # Initialize the pillar representation with zeros.
    pillars = np.zeros([n_max_pillars, n_max_points_per_pillar, n_in_features]).astype(np.float32)
    pillar_indices = np.zeros((n_max_pillars, 3), dtype=int)
    for i, pillar_id in enumerate(unique_pillars):
        # 找出pillar对应的点
        indices = np.where(pillar_idx == pillar_id)[0]  # 每个pillar内所有点的索引
        
        # (可选)当pillar内点云大于n, 则对z进行排序，取前n个最小z的点
        if len(indices) > num_points_per_pillars[i]:
            points_z_sort = np.argsort(points_roi[indices, 2], axis=0)  # 返回z从小到大排序的索引值，axis=0 按列排序
            indices = indices[points_z_sort]
        indices = indices[:num_points_per_pillars[i]]  # 每个pillar内取前n个点的索引

        # 计算pillar的中心坐标
        pillar_c = np.mean(points_roi[indices, :3], axis=0)

        # 计算pillar的中心坐标所在pillar的index
        x_pillar_idx = ((pillar_c[0] - x_min) // x_pillar).astype(int)
        y_pillar_idx = ((pillar_c[1] - y_min) // y_pillar).astype(int)

        # 将每个pillar坐标位置存入
        pillar_indices[pillar_id, 1:] = [x_pillar_idx, y_pillar_idx]

        # 计算pillar的中心坐标和每个点的偏移
        offsets = points_roi[indices, :3] - pillar_c

        # 将点和偏移添加到pillar中
        pillars[pillar_id, :num_points_per_pillars[i], :4] = points_roi[indices]
        pillars[pillar_id, :num_points_per_pillars[i], 4:7] = offsets
    # print(pillars[:5, :5, :], np.unique(pillars[:, 0, 0]), pillar_indices[:, :])

    return pillars, pillar_indices

    # # initialize the number of points in each pillar
    # num_points_in_pillar = np.zeros((n_max_pillars,))
    # # Fill the pillar representation with the pointcloud data.
    # for i, p in enumerate(points_roi):
    #     pillar_index = pillar_idx[i]
    #     point_index = i
    #     # check if the pillar is full
    #     if num_points_in_pillar[pillar_index] < n_max_points_per_pillar:
    #         # add the point to the pillar
    #         pillars[pillar_index, point_index, :4] = p.reshape(1, -1)

    #     # calculate the centroid of the points in the pillar
    #     if num_points_in_pillar[pillar_index] > 0:
    #         points_num = num_points_in_pillar[pillar_index][0].astype(int)
    #         pillar_points = pillars[pillar_index, 0:points_num, 0:3]
    #         centroid = np.mean(pillar_points, axis=0)
    #         # print(p[:3] - centroid, p, centroid.shape)
    #         # print(pillars[pillar_index, 0:points_num, 4:8])
    #         pillars[pillar_index, 0:points_num, 4:8] = p[:3] - centroid

    #     # increment the number of points in the pillar
    #     num_points_in_pillar[pillar_index] += 1
    # print(pillars[pillar_index, i, ...].shape)


if __name__=="__main__":
    root_dir = "/dataset"
    # for i, f in enumerate(os.listdir(root_dir)):
    #     file = os.path.join(root_dir, f)
    #     print(i,f)
    for root, dirs, files in os.walk(root_dir):
        for file in files:
            file_name, suffix = os.path.splitext(file)
            if suffix == '.pcd':
                file = os.path.join(root, file)
                pillar_path = os.path.dirname(os.path.dirname(file)) + "/pillar_tensor_sort"
                pillar_file = os.path.join(pillar_path, file_name+'.pkl')
                if not os.path.exists(pillar_path):
                    print("pillar path is not exist.")
                    os.makedirs(pillar_path)
                # print(pillar_file)
                # print(pillar_path, suffix , file_name)
                if not os.path.exists(pillar_file):
                    points = read_pcd(file)
                    pillars, pillar_indices = points_to_pillars(points)
                    with open(pillar_file, 'wb') as f:
                        pickle.dump([pillars, pillar_indices], f, pickle.HIGHEST_PROTOCOL)
                        print(pillar_file)
                    print(pillars.shape, pillar_indices.shape)