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 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文本文件中查看
        # 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_id
        info['timestamps_seconds'] = list(scenario.timestamps_seconds)  # list of int of shape (91)
        info['current_time_index'] = scenario.current_time_index  # int, 10
        info['sdc_track_index'] = scenario.sdc_track_index  # int
        info['objects_of_interest'] = list(scenario.objects_of_interest)  # list, could be empty list

        info['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 predicted
        info['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_pb2


plt.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')
        break
    break

Waymo Motion Open Dataset 介绍与使用_第1张图片

2. 场景可视化可以使用metadrive仿真器将场景导入。

你可能感兴趣的:(自动驾驶,自动驾驶,深度学习,人工智能)