Waymo Motion Open Dataset 介绍与使用

编程入门行业动态更新时间:2024-10-27 10:25:41

Waymo <a href=https://www.elefans.com/category/jswz/34/1711525.html style= Motion Open Dataset 介绍与使用"/>

Waymo Motion Open Dataset 介绍与使用

Waymo Motion Open Dataset 的使用

自动驾驶数据集使用交流 771845071
1. 下载
2. 安装
3. 使用
- (1) 解析数据
- (2) scenario 信息
3. 可视化
- 1. 地图可视化
- 2. 场景可视化可以使用metadrive仿真器将场景导入。

自动驾驶数据集使用交流 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  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 .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_pb2raw_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文本文件中查看# with open('./one_scenario.txt', 'w+') as f:#     f.write(str(scenario))# f.close()# print("--------------------- save successfully! ---------------------")# quit()print(type(scenario))info['scenario_id'] = scenario.scenario_idinfo['timestamps_seconds'] = list(scenario.timestamps_seconds)  # list of int of shape (91)info['current_time_index'] = scenario.current_time_index  # int, 10info['sdc_track_index'] = scenario.sdc_track_index  # intinfo['objects_of_interest'] = list(scenario.objects_of_interest)  # list, could be empty listinfo['tracks_to_predict'] = {'track_index': [cur_pred.track_index for cur_pred in scenario.tracks_to_predict],'difficulty': [cur_pred.difficulty for cur_pred in scenario.tracks_to_predict]}  # for training: suggestion of objects to train on, for val/test: need to be predictedinfo['tracks'] = list(scenario.tracks)info['dynamic_map_states'] = list(scenario.dynamic_map_states)output_file = os.path.join(process_data_path, f'sample_{scenario.scenario_id}.pkl')with open(output_file, 'wb') as f:pickle.dump(info, f)

每个场景包含哪些内容？我们这里将每一份数据（一个完整的20s数据）姑且称为一个 scenario，官方好像称为 segment 。

(2) scenario 信息

scenario.scenario_id：场景唯一id

scenario.timestamps_seconds：从0开始的每个场景的时间步

scenario.current_time_index：当前帧的时间索引（前面为历史，后面为预测）

scenario.tracks：包含每个对象的轨迹

id：每个对象的唯一id
object_type：每个对象的类型
states：当前 id 对应对象在每一帧时的状态信息
center_x: 几何中心x坐标
center_y: 几何中心y坐标
center_z: 几何中心z坐标
length: 几何长度
width: 几何宽度
height: 几何高度
heading: 车头角度
velocity_x: 横向速度
velocity_y: 纵向速度
valid: 该帧是否可以观测到对象

scenario.dynamic_map_states：跨时间步长的交通信息状态

lane_states：包含给定时间步长的交通信号状态集及其控制的车道 id

scenario.map_features：地图数据

lane centers：车道中心线
lane boundaries：车道边界
road boundaries：道路边界
crosswalks：人行横道
speed bumps：减速带位置
stop signs：停车点位置

scenario.sdc_track_index：自动驾驶车辆（主车）在场景中的索引

scenario.objects_of_interest：可能对研究训练有用的行为的对象

scenario.tracks_to_predict：指示必须预测哪些对象，仅在训练和验证集中提供

3. 可视化

1. 地图可视化

import os
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from waymo_open_dataset.protos import scenario_pb2plt.figure(figsize=(30,30))
plt.rcParams['axes.facecolor']='grey'# 选择交互车辆
data_path = '/data1/zyp/drl/PiH/MTR_PiH/data/training_20s/validation'
file_list = os.listdir(data_path)for cnt_file, file in enumerate(file_list):file_path = os.path.join(data_path, file)dataset = tf.data.TFRecordDataset(file_path, compression_type='')for cnt_scenar, data in enumerate(dataset):scenario = scenario_pb2.Scenario()scenario.ParseFromString(bytearray(data.numpy()))print(scenario.scenario_id)# 画地图线for i in range(len(scenario.map_features)):# 车道线if str(scenario.map_features[i].lane) != '':line_x = [z.x for z in scenario.map_features[i].lane.polyline]line_y = [z.y for z in scenario.map_features[i].lane.polyline]plt.scatter(line_x, line_y, c='g', s=5)# plt.text(line_x[0], line_y[0], str(scenario.map_features[i].id), fontdict={'family': 'serif', 'size': 20, 'color': 'black'})# 边界线if str(scenario.map_features[i].road_edge) != '':road_edge_x = [polyline.x for polyline in scenario.map_features[i].road_edge.polyline]road_edge_y = [polyline.y for polyline in scenario.map_features[i].road_edge.polyline]plt.scatter(road_edge_x, road_edge_y)# plt.text(road_edge_x[0], road_edge_y[0], scenario.map_features[i].road_edge.type, fontdict={'family': 'serif', 'size': 20, 'color': 'black'})if scenario.map_features[i].road_edge.type == 2:plt.scatter(road_edge_x, road_edge_y, c='k')elif scenario.map_features[i].road_edge.type == 3:plt.scatter(road_edge_x, road_edge_y, c='purple')print(scenario.map_features[i].road_edge)else:plt.scatter(road_edge_x, road_edge_y, c='k')# 道路边界线if str(scenario.map_features[i].road_line) != '':road_line_x = [j.x for j in scenario.map_features[i].road_line.polyline]road_line_y = [j.y for j in scenario.map_features[i].road_line.polyline]if scenario.map_features[i].road_line.type == 7:  # 双实黄线plt.plot(road_line_x, road_line_y, c='y')elif scenario.map_features[i].road_line.type == 8:  # 双虚实黄线plt.plot(road_line_x, road_line_y, c='y') elif scenario.map_features[i].road_line.type == 6:  # 单实黄线plt.plot(road_line_x, road_line_y, c='y')elif scenario.map_features[i].road_line.type == 1:  # 单虚白线for i in range(int(len(road_line_x)/7)):plt.plot(road_line_x[i*7:5+i*7], road_line_y[i*7:5+i*7], color='w')elif scenario.map_features[i].road_line.type == 2:  # 单实白线plt.plot(road_line_x, road_line_y, c='w')else:plt.plot(road_line_x, road_line_y, c='w')# 画车及轨迹for i in range(len(scenario.tracks)):if i==scenario.sdc_track_index:traj_x = [center.center_x for center in scenario.tracks[i].states if center.center_x != 0.0]traj_y = [center.center_y for center in scenario.tracks[i].states if center.center_y != 0.0]head = [center.heading for center in scenario.tracks[i].states if center.center_y != 0.0]plt.scatter(traj_x[0], traj_y[0], s=140, c='r', marker='s')# plt.imshow(img1,extent=[traj_x[0]-3, traj_x[0]+3,traj_y[0]-1.5, traj_y[0]+1.5])plt.scatter(traj_x, traj_y, s=14, c='r')else:traj_x = [center.center_x for center in scenario.tracks[i].states if center.center_x != 0.0]traj_y = [center.center_y for center in scenario.tracks[i].states if center.center_y != 0.0]head = [center.heading for center in scenario.tracks[i].states if center.center_y != 0.0]plt.scatter(traj_x[0], traj_y[0], s=140, c='k', marker='s')# plt.imshow(img1,extent=[traj_x[0]-3, traj_x[0]+3,traj_y[0]-1.5, traj_y[0]+1.5])plt.scatter(traj_x, traj_y, s=14, c='b')breakbreak