Waymo数据集实战:从TFRecord到激光雷达点云可视化的完整流程(附Python代码)

自动驾驶技术的快速发展离不开高质量的数据集支持。作为行业领先的开放数据集,Waymo Open Dataset以其丰富的传感器数据和精确标注成为研究人员的首选。本文将带您深入探索如何从原始TFRecord文件中提取激光雷达点云数据,并实现多种可视化方案。

1. 环境准备与数据获取

在开始处理Waymo数据集前,我们需要配置合适的开发环境。推荐使用Python 3.7+和TensorFlow 2.x环境,以下是必需的依赖包:

pip install waymo-open-dataset-tf-2-6-0
pip install open3d matplotlib numpy

Waymo数据集可通过官网申请下载,包含训练集、验证集和测试集。每个TFRecord文件约200MB-2GB不等,完整数据集大小超过1TB。下载后建议按以下结构组织文件:

waymo_dataset/
├── training/
│   ├── segment-1005085.tfrecord
│   ├── segment-1007239.tfrecord
│   └── ...
└── validation/
    ├── segment-1020363.tfrecord
    └── ...

提示:Waymo数据集使用Apache 2.0许可证,允许学术和商业用途,但需遵守其使用条款。

2. TFRecord文件结构解析

Waymo的TFRecord文件采用Protocol Buffers格式存储,每个文件包含约20秒的连续驾驶场景数据。主要数据结构包括:

  • Frame: 基本时间单元,包含所有传感器在该时刻的数据
  • Laser: 激光雷达点云数据(含强度、反射率等信息)
  • Camera: 多视角相机图像
  • Label: 3D边界框标注信息

关键字段说明:

字段名类型描述
contextContext场景元数据
timestamp_microsint64时间戳(微秒)
poseTransform自车位姿
lasersLaser激光雷达数据数组
imagesCamera相机图像数组

3. 点云数据提取实战

以下Python代码展示了如何从TFRecord中提取激光雷达点云:

import tensorflow as tf
from waymo_open_dataset import dataset_pb2
from waymo_open_dataset.utils import frame_utils

def extract_lidar_points(tfrecord_path, frame_index=0):
    dataset = tf.data.TFRecordDataset(tfrecord_path, compression_type='')
    frames = [dataset_pb2.Frame.FromString(data) for data in dataset]
    frame = frames[frame_index]
    
    # 解析激光雷达数据
    (range_images, camera_projections, _, range_image_top_pose) = (
        frame_utils.parse_range_image_and_camera_projection(frame))
    
    # 转换点云坐标
    points, cp_points = frame_utils.convert_range_image_to_point_cloud(
        frame, range_images, camera_projections, range_image_top_pose)
    
    # 合并五个激光雷达的点云
    all_points = np.concatenate(points, axis=0)
    return all_points

该函数返回的all_points是一个N×4的numpy数组,每行包含(x,y,z,intensity)四个值。Waymo车辆配备五个激光雷达,分别位于车顶、前、后、左、右位置。

4. 可视化方案对比

4.1 Matplotlib基础可视化

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def plot_with_matplotlib(points):
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 只绘制每10个点以提高性能
    ax.scatter(points[::10, 0], points[::10, 1], points[::10, 2], 
               c=points[::10, 3], s=0.1, cmap='viridis')
    
    ax.set_xlabel('X (m)')
    ax.set_ylabel('Y (m)')
    ax.set_zlabel('Z (m)')
    plt.title('Waymo LiDAR Point Cloud')
    plt.show()

Matplotlib适合快速查看点云,但处理大规模数据时性能有限。对于超过100万个点的场景,建议:

  1. 降采样显示(如代码中的points[::10]
  2. 使用更强的硬件
  3. 切换到专业可视化工具

4.2 Open3D高级可视化

Open3D提供了更专业的点云处理能力:

import open3d as o3d

def visualize_with_open3d(points):
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points[:, :3])
    pcd.colors = o3d.utility.Vector3dVector(
        np.tile(points[:, 3:], (1, 3)) / 255.0)  # 强度值映射为颜色
    
    # 可视化设置
    vis = o3d.visualization.Visualizer()
    vis.create_window(width=800, height=600)
    vis.add_geometry(pcd)
    
    # 设置视角
    ctr = vis.get_view_control()
    ctr.set_front([0, -1, 0.5])
    ctr.set_up([0, 0, 1])
    ctr.set_zoom(0.1)
    
    vis.run()
    vis.destroy_window()

Open3D的优势包括:

  • 支持实时旋转、缩放等交互
  • 提供多种点云处理算法(滤波、配准等)
  • 可保存可视化结果为图像或视频

5. 常见问题排查

5.1 数据读取问题

错误现象:`tfrecord文件无法解析**

解决方案:

  1. 确认使用的waymo-open-dataset库版本与TFRecord版本匹配
  2. 检查文件是否完整下载
  3. 尝试官方提供的验证脚本

5.2 可视化性能问题

卡顿处理方案

# 在Open3D中使用体素下采样
down_pcd = pcd.voxel_down_sample(voxel_size=0.1)

5.3 坐标转换问题

Waymo使用右手坐标系:

  • X: 车辆前进方向
  • Y: 左侧方向
  • Z: 上方方向

转换到其他坐标系时需注意变换矩阵:

# 示例:Waymo到KITTI坐标系转换
def waymo_to_kitti(points):
    # KITTI: X→前,Y→左,Z→上
    transform = np.array([[0, -1, 0],
                         [0, 0, -1],
                         [1, 0, 0]])
    return points @ transform.T

6. 进阶应用:点云分析与标注可视化

结合Waymo提供的3D标注框,我们可以实现更丰富的可视化:

def visualize_with_boxes(frame, points):
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points[:, :3])
    
    geometries = [pcd]
    for label in frame.laser_labels:
        box = label.box
        center = [box.center_x, box.center_y, box.center_z]
        size = [box.length, box.width, box.height]
        rotation = box.heading
        
        # 创建3D边界框
        bbox = o3d.geometry.OrientedBoundingBox(center, 
            [[np.cos(rotation), -np.sin(rotation), 0],
             [np.sin(rotation), np.cos(rotation), 0],
             [0, 0, 1]], size)
        bbox.color = [1, 0, 0]  # 红色
        geometries.append(bbox)
    
    o3d.visualization.draw_geometries(geometries)

这种可视化对于算法调试和数据分析非常有价值,可以直观地检查点云与标注的匹配情况。

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐