Skip to content

A Unified Framework for scalable Vehicle Trajectory Prediction

License

Notifications You must be signed in to change notification settings

ChengkaiYang/UniTraj

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

21 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

UniTraj

Code | Paper

A Unified Framework for Scalable Vehicle Trajectory Prediction

💡UniTraj allows users to train and evaluate trajectory prediction models from real-world datasets like Waymo, nuPlan, nuScenes and Argoverse2 in a unified pipeline.

system

🔥Powered by Hydra, Pytorch-lightinig, and WandB, the framework is easy to configure, train and log.

system

🛠 Quick Start

  1. Create a new conda environment
conda create -n unitraj python=3.9
conda activate unitraj
  1. Install ScenarioNet: https://scenarionet.readthedocs.io/en/latest/install.html

  2. Install Unitraj:

git clone git@github.com:metadriverse/scenarionet.git
cd unitraj
pip install -r requirements.txt
pip setup.py develop

You can verify the installation of UniTraj via running the training script:

python train.py method=autobot

The model will be trained on several sample data.

Code Structure

There are three main components in UniTraj: dataset, model and config. The structure of the code is as follows:

unitraj
├── configs
│   ├── config.yaml
│   ├── method
│   │   ├── autobot.yaml
│   │   ├── MTR.yaml
│   │   ├── wayformer.yaml
├── datasets
│   ├── base_dataset.py
│   ├── autobot_dataset.py
│   ├── wayformer_dataset.py
│   ├── MTR_dataset.py
├── models
│   ├── autobot
│   ├── mtr
│   ├── wayformer
│   ├── base_model
├── utils

There is a base config, dataset and model class, and each model has its own config, dataset and model class that inherit from the base class.

Pipeline

1. Data Preparation

UniTraj takes data from ScenarioNet as input. Process the data with ScenarioNet in advance.

2. Configuration

UniTraj uses Hydra to manage configuration files.

Universal configuration file is located in unitraj/config/config.yaml. Each model has its own configuration file in unitraj/config/method/, for example, unitraj/config/method/autobot.yaml.

The configuration file is organized in a hierarchical structure, and the configuration of the model is inherited from the universal configuration file.

Configuration Example

Please refer to config.yaml and method/autobot.yaml for more details.

2. Train

python train.py

3. Evaluation

  1. In config.yaml, set the ckpt_path to the path of the trained model and val_data_path to the validation data path.
  2. (Optional) In config.yaml, set eval_waymo or eval_nuscenes to True if you want to evaluate the model with Waymo or nuScenes official evaluation tool. (Install waymo-open-dataset and nuscenes-devkit first)
  3. Runpython evaluation.py

4. Dataset Analysis

python data_analysis.py

Contribute to UniTraj

Implement a new model

  1. Create a new config file in unitraj/config/ folder, for example, unitraj/config/new_model.yaml
  2. (Optional) Create a new dataset class in unitraj/datasets/ folder, for example, unitraj/datasets/new_dataset.py, and inherit unitraj/dataset/base_dataset.py
  3. Create a new model class in unitraj/model/ folder, for example, unitraj/model/lanegcn.py, and inherit from pl.LightningModule

Batch Data Format

The default batch data format is:

For trajectory input: obj_trajs [num_centered_obj, num_surrounding_objs, past_time_steps, num_attribute]

[0:3] position (x, y, z) [3:6] size (l, w, h) [6:11] type_onehot [11:23] time_onehot [23:35] heading_encoding [25:37] vx,vy [27:39] ax,ay

For map input: map_polylines [num_centered_obj, num_surrounding_lines, max_points_per_lane, num_attribute=9]

[0:3] position (x, y, z) [3:6] direction (x, y, z) [6] type [7:9] previous_point_xy

About

A Unified Framework for scalable Vehicle Trajectory Prediction

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 82.7%
  • Cuda 10.7%
  • C++ 6.3%
  • C 0.3%