Argoverse数据集入门指南：从数据加载到轨迹预测实战

1次阅读

共计 3360 个字符，预计需要花费 9 分钟才能阅读完成。

自动驾驶研究需要大量真实场景数据来训练和验证算法。Argoverse 数据集由 Argo AI 发布，包含超过 1000 小时的真实驾驶数据，是目前最全面的自动驾驶开源数据集之一。对于刚接触这个领域的研究者来说，这个数据集既宝贵又充满挑战。

新手常遇到的困难包括：

数据规模庞大，单是运动预测子集就包含 30,000 多个场景，每个场景包含 5 秒的轨迹数据
数据结构复杂，包含高清地图、3D 物体跟踪和运动预测三个相互关联的子集
坐标系转换繁琐，需要在城市坐标系、车辆坐标系和像素坐标系之间转换
时间戳对齐困难，特别是当需要同步多个传感器数据时

Argoverse 数据集由三个主要子集组成，每个子集服务于不同的研究目标：

高清地图(HD Maps)
包含 6 个城市 (匹兹堡、迈阿密等) 的详细地图信息
提供车道几何形状、交通信号灯位置等语义信息
使用 JSON 格式存储，包含车道中心线、边界等信息
3D 跟踪(3D Tracking)
包含 113 个场景的连续帧点云和图像数据
每个场景约 15-30 秒，包含 LiDAR 点云和 7 个摄像头的图像
提供 3D 物体 (车辆、行人等) 的标注框和轨迹
运动预测(Motion Forecasting)
包含 324,557 个场景的轨迹数据
每个场景 5 秒长度，包含主车和周围物体的历史轨迹
挑战是预测物体在未来 3 秒内的运动

首先安装官方 Python API：

pip install argoverse-api

from argoverse.data_loading.argoverse_forecasting_loader import ArgoverseForecastingLoader

# 初始化数据加载器
root_dir = "/path/to/argoverse-forecasting/train"
avl = ArgoverseForecastingLoader(root_dir)

# 加载第一个场景
track_id_list = avl.track_id_list
seq_path = avl.seq_list[0]  # 第一个场景路径
data = avl.get(seq_path)

# 打印场景信息
print(f"场景 ID: {data.seq_id}")
print(f"城市: {data.city}")
print(f"轨迹数量: {len(data.tracks)}")

使用 Matplotlib 绘制场景中所有物体的轨迹：

import matplotlib.pyplot as plt

def plot_scene(data):
    fig, ax = plt.subplots(figsize=(10, 10))

    # 绘制每个物体的轨迹
    for track in data.tracks:
        x_coords = [obs.position[0] for obs in track.observations]
        y_coords = [obs.position[1] for obs in track.observations]
        ax.plot(x_coords, y_coords, marker='o', markersize=3, linestyle='-')

    # 设置图表
    ax.set_title(f"场景 {data.seq_id}")
    ax.set_xlabel("X 坐标(米)")
    ax.set_ylabel("Y 坐标(米)")
    ax.grid(True)
    ax.axis('equal')
    plt.show()

plot_scene(data)

使用 PyTorch 构建简单的 LSTM 轨迹预测模型：

import torch
import torch.nn as nn

class TrajectoryPredictor(nn.Module):
    def __init__(self, input_size=2, hidden_size=64, num_layers=2, output_size=2):
        super(TrajectoryPredictor, self).__init__()
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True
        )
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        # x 形状: (batch_size, seq_len, 2)
        lstm_out, _ = self.lstm(x)
        # 只取最后一个时间步的输出
        out = self.fc(lstm_out[:, -1, :])
        return out

# 示例使用
model = TrajectoryPredictor()
# 假设输入是 20 个时间步的 (x,y) 坐标
input_traj = torch.randn(32, 20, 2)  # batch_size=32, seq_len=20
prediction = model(input_traj)  # 预测下一个位置

处理大规模自动驾驶数据集时，内存和计算效率至关重要：

批处理加载：不要一次性加载所有数据，使用 PyTorch 的 DataLoader 实现批处理

from torch.utils.data import Dataset, DataLoader

class ArgoverseDataset(Dataset):
    def __init__(self, root_dir):
        self.avl = ArgoverseForecastingLoader(root_dir)
        self.seq_list = self.avl.seq_list

    def __len__(self):
        return len(self.seq_list)

    def __getitem__(self, idx):
        data = self.avl.get(self.seq_list[idx])
        # 转换为模型需要的格式
        # 这里简化为返回第一个物体的轨迹
        track = data.tracks[0]
        traj = torch.tensor([[obs.position[0], obs.position[1]] 
                            for obs in track.observations], dtype=torch.float32)
        return traj

# 创建数据加载器
dataset = ArgoverseDataset("/path/to/argoverse-forecasting/train")
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)