Introduction

This repository is the code for our article Raising context awareness in motion forecasting, accepted at CVPR Workshop WAD 2022. It was built based on the Trajectron ++ repository.

Setup

Start by downloading the nuScenes dataset (v1.0) along with the map expansion pack (v1.1). Then, install the requirements

pip install -r requirements.txt

Pre-processing

The preprocessing code is similar to the original one, excepts that it splits the data according to the Challenge split (see this issue). To pre-process the data, please run

cd CAB/trajectron && 
python process_data_challenge.py --data /path/to/nuscenes/ --version "v1.0-trainval" --output_path /desired/preprocessed/location

This script will create the following .pkl files:

nuScenes_train_full.pkl  
nuScenes_train_val_full.pkl  
nuScenes_val_full.pkl

Train models

Here are the commands you can use to train the models presented in our article. Before training, please specify in the data_dir field of the config files the location of the pre-processed dataset (the value of --output_path in the process_data_challenge.py script).

CAB

cd CAB/trajectron &&
python train.py -o ../experiments/nuScenes/models/cab/config.yaml --log_dir /desired/log/dir

RubiZ

cd CAB/trajectron &&
python train.py -o ../experiments/nuScenes/models/rubiz/config.yaml --log_dir /desired/log/dir

Reweight

cd CAB/trajectron &&
python train.py -o ../experiments/nuScenes/models/reweight/config.yaml --log_dir /desired/log/dir

Trajectron

cd CAB/trajectron &&
python train.py -o ../experiments/nuScenes/models/trajectron/config.yaml --log_dir /desired/log/dir

Trajectron (no-context)

cd CAB/trajectron &&
python train.py -o ../experiments/nuScenes/models/trajectron_nc/config.yaml --log_dir /desired/log/dir

HalentNet

cd CAB/trajectron &&
python train_halentnet.py -o ../experiments/nuScenes/models/halentnet/config.yaml --log_dir /desired/log/dir

Evaluate models

Model evaluation is done using two scripts. The first one computes all the predictions on the given dataset.

MODEL_DIR="/directory/of/saved/experiment" && CKPT=epoch_you_want_to_evaluate
python evaluate_new.py -o $MODEL_DIR/config.yaml --resume $CKPT  --eval_data_dict nuScenes_val_full.pkl

This will create a folder in your logs that contains all the predictions. Then, you can use these predictions to computes the metrics values:

python compute_val_metrics.py --model_dir $MODEL_DIR --checkpoint $CKPT

Metrics values are prompted and saved in $MODEL_DIR/metrics/epoch,$CKPT/val/metrics.pkl

Citation

If you use our code and/or our article, you can cite us using:

@InProceedings{cab2022,
    author    = {Ben-Younes, H\'edi and Zablocki, \'Eloi and Chen, Micka\"el and P\'erez, Patrick and Cord, Matthieu},
    title     = {Raising Context Awareness in Motion Forecasting},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops},
    month     = {June},
    year      = {2022},
    pages     = {4409-4418}
}

valeoai / CAB