必要的包

from nuscenes.nuscenes import NuScenes
from pyquaternion import Quaternion
from nuscenes.utils.data_classes import LidarPointCloud
import numpy as np
from open3d import *
from nuscenes.utils.data_io import load_bin_file
from nuscenes.utils.geometry_utils import points_in_box
import os.path as osp
from functools import partial
from utils.points_process import *
from sklearn.neighbors import KDTree
import open3d as o3d
import argparse

初始化全局字典，用于存储中间静态点、姿态和标签

INTER_STATIC_POINTS = {
   }
INTER_STATIC_POSE = {
   }
INTER_STATIC_LABEL = {
   }

参数设置

dataroot: 数据集的根路径，类型为字符串，默认值为 ‘./project/data/nuscenes/’。

• save_path: 保存路径，类型为字符串，默认值为 ‘./project/data/nuscenes//occupancy2/’，该参数是可选的。
• num_sweeps: 每个示例的激光雷达扫描次数，类型为整数，默认值为 10，该参数是可选的。

def parse_args():
    parser = argparse.ArgumentParser(description='Data converter arg parser')
    parser.add_argument(
        '--dataroot',
        type=str,
        default='./project/data/nuscenes/',
        help='specify the root path of dataset')
    parser.add_argument(
        '--save_path',
        type=str,
        default='./project/data/nuscenes//occupancy2/',
        required=False,
        help='specify sweeps of lidar per example')
    parser.add_argument(
        '--num_sweeps',
        type=int,
        default=10,
        required=False,
        help='specify sweeps of lidar per example')
    args = parser.parse_args()
    return args

align_dynamic_thing：为了将动态物体的点云数据从上一帧对齐到当前帧

def align_dynamic_thing(box, prev_instance_token, nusc, prev_points, ego_frame_info):
        if prev_instance_token not in ego_frame_info['instance_tokens']:
            box_mask = points_in_box(box,
                                    prev_points[:3, :])
            return np.zeros((prev_points.shape[0], 0)), np.zeros((0, )), box_mask
        
        box_mask = points_in_box(box,
                                    prev_points[:3, :])
        box_points = prev_points[:, box_mask].copy()
        prev_bbox_center = box.center
        prev_rotate_matrix = box.rotation_matrix

        box_points = rotate(box_points, np.linalg.inv(prev_rotate_matrix), center=prev_bbox_center)
        target = ego_frame_info['instance_tokens'].index(prev_instance_token)
        ego_boxes_center = ego_frame_info['boxes'][target].center
        box_points = translate(box_points, ego_boxes_center-prev_bbox_center)
        box_points = rotate(box_points, ego_frame_info['boxes'][target].rotation_matrix, center=ego_boxes_center)
        box_points_mask = filter_points_in_ego(box_points, ego_frame_info, prev_instance_token)
        box_points = box_points[:, box_points_mask]
        box_label = np.full_like(box_points[0], nusc.lidarseg_name2idx_mapping[box.name]).copy()
        return box_points, box_label, box_mask

流程

1. 检查实例标识:
   if prev_instance_token ∉ ego_frame_info[‘instance_tokens’]:
   box_mask = points_in_box(box, prev_points[:3, :])
   return (0, 0, box_mask)

2. 计算边界框内的点:
   box_mask = points_in_box(box, prev_points[:3, :])
   box_points = prev_points[:, box_mask]

3. 获取上一帧边界框的中心和旋转矩阵:
   C_prev = box.center
   R_prev = box.rotation_matrix

4. 将点旋转到原点并平移到当前帧的中心:
   box_points = R_prev^-1 * (box_points - C_prev)

5. 获取目标边界框的中心和旋转矩阵:
   target = ego_frame_info[‘instance_tokens’].index(prev_instance_token)
   C_ego = ego_frame_info[‘boxes’][target].center
   R_ego = ego_frame_info[‘boxes’][target].rotation_matrix

6. 平移到当前帧的中心并再次旋转:
   box_points = box_points + (C_ego - C_prev)
   box_points = R_ego * box_points

7. 过滤当前帧边界框内的点:
   box_points_mask = points_in_box(ego_frame_info[‘boxes’][target], box_points[:3, :])
   box_points = box_points[:, box_points_mask]

8. 生成点云数据的标签:
   box_label = full_like(box_points[0], nusc.lidarseg_name2idx_mapping[box.name])

9. 返回结果:
   return (box_points, box_label, box_mask)

numpy.full_like()是根据现有数组的形状和数据类型来创建新数组，而numpy.full()则需要手动指定形状和数据类型。

旋转函数

def rotate(points, rot_matrix: np.ndarray, center=None) -> np.array:
    """
    Applies a rotation.
    :param rot_matrix: <np.float: 3, 3>. Rotation matrix.
    """
    if center is not None:
        points[:3, :] = np.dot(rot_matrix, points[:3, :]-center[:, None]) + center[:, None]
    else:

        points[:3, :] = np.dot(rot_matrix, points[:3, :])
    return points

平移公式

def translate(points, x: np.ndarray) -> np.array:
    """
    Applies a translation to the point cloud.
    :param x: <np.float: 3, 1>. Translation in x, y, z.
    """
    for i in range