Waymo Motion Open Dataset 介绍与使用

自动驾驶数据集使用交流 771845071

Waymo数据集分为两部分：motion 和 perception，其中motion数据集的主要用途是Sim Agents，Motion Prediction，Interaction Prediction，和 Occupancy and Flow Prediction。官方主页：Waymo Motion Dataset

1. 下载

Waymo_motion 原始数据集保存在谷歌云盘中waymo_open_dataset_motion_v_1_1_0 ，需要自行下载。

其目录结构如下：

• uncompressed/
  • occupancy_flow_challenge/
  • tf_example/
  • scenario/
    • training/ 训练集
    • validation/ 验证集
    • testing/ 测试集
    • validation_interactive/ 交互预测的验证集
    • testing_interactive/ 交互预测的测试集
    • training_20s/ 20s长时数据集（总共86.6GB）

其中 scenario 文件夹下的数据与 tf_example 文件夹下数据基本相同（数据保存格式不同，tf_example 文件夹下为 tensor 格式），occupancy_flow_challenge文件夹下数据没有使用过，后面使用到了再做补充。

scenario 下普通数据集每个场景共9秒91帧（历史10帧，当前1帧，未来80帧），在轨迹预测或其他任务中一般使用1s历史数据预测8s未来数据。这9s的数据是由下面20s的数据切分得到的，测试集只包含1s的历史数据（防作弊），后续1.2.0版本的数据集中加入了1s的历史雷达点云数据。

training_20s 数据集共包含70506个场景，每个场景20秒199帧数据，通常根据实验用途自己进行切分。

2. 安装

# 新建 conda 环境
conda create -n waymo python=3.8
conda activate waymo

# 安装 tensorflow （需要注意tensorflow与cuda版本对应，我这里使用cuda 11.0版本）
pip install tensorflow==2.4.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

#安装 pytorch (pytorch-lightning 是否安装自行选择）
pip install pytorch-lightning==1.4.0
pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html

# 测试gpu是否可用（返回 gpu 信息和 True 则表示安装正确）
python
>>> import tensorflow as tf
>>> tf.config.list_physical_devices('GPU')
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
>>> import torch
>>> torch.cuda.is_available()
True
>>> quit()

 # 安装 waymo-open-dataset 需要与tensorflow版本对应
 pip install waymo-open-dataset-tf-2-4-0

# 其他库缺啥安装啥

3. 使用

(1) 解析数据

import os, glob, pickle
import tensorflow as tf
from waymo_open_dataset.protos import scenario_pb2

raw_data_path = 'training_20s/'
process_data_path = 'training_20s_process/'

raw_data = glob.glob(os.path.join(raw_data_path, '*.tfrecord*'))
raw_data.sort()

for data_file in raw_data:
    dataset = tf.data.TFRecordDataset(data_file, compression_type='')
    for cnt, data in enumerate(dataset):
        info = {
   
   }
        scenario = scenario_pb2.Scenario()
        scenario.ParseFromString(bytearray(data.numpy()))
		
		# 一下几行代码可以将每个scenario的所有信息保存到txt文本文件中查看