Box-Aware Tracker (BAT)

Pytorch-Lightning implementation of the Box-Aware Tracker.

Box-Aware Feature Enhancement for Single Object Tracking on Point Clouds. ICCV 2021

Chaoda Zheng, Xu Yan, Jiaotao Gao, Weibing Zhao, Wei Zhang, Zhen Li*, Shuguang Cui


  title={Box-Aware Feature Enhancement for Single Object Tracking on Point Clouds},
  author={Chaoda Zheng, Xu Yan, Jiaotao Gao, Weibing Zhao, Wei Zhang, Zhen Li, Shuguang Cui},



  • Modular design. It is easy to config the model and trainng/testing behaviors through just a .yaml file.
  • DDP support for both training and testing.
  • Provide a 3rd party implementation of P2B.



  • create the environment

    git clone
    cd BAT
    conda create -n bat  python=3.6
    conda activate bat
  • Install pytorch

    conda install pytorch==1.4.0 torchvision==0.5.0 cudatoolkit=10.1 -c pytorch

    Our code is well tested with pytorch 1.4.0 and CUDA 10.1. But other platforms may also work. Follow this to install another version of pytorch.

  • Install other dependencies

    pip install -r requirement.txt

KITTI dataset

  • Download the data for velodyne, calib and label_02 from KITTI Tracking.
  • Unzip the downloaded files.
  • Put the unzipped files under the same folder as following.
    [Parent Folder]
    --> [calib]
        --> {0000-0020}.txt
    --> [label_02]
        --> {0000-0020}.txt
    --> [velodyne]
        --> [0000-0020] folders with velodynes .bin files

Quick Start


To train a model, you must specify the .yaml file with --cfg argument. The .yaml file contains all the configurations of the dataset and the model. Currently, we provide three .yaml files under the cfgs directory. Note: Before running the code, you will need to edit the .yaml file by setting the path argument as the correct root of the dataset.

python --gpu 0 1 --cfg cfgs/BAT_Car.yaml  --batch_size 50 --epoch 60

After you start training, you can start Tensorboard to monitor the training process:

tensorboard --logdir=./ --port=6006

By default, the trainer runs a full evaluation on the full test split after training every epoch. You can set --check_val_every_n_epoch to a larger number to speed up the training.


To test a trained model, specify the checkpoint location with --checkpoint argument and send the --test flag to the command.

python --gpu 0 1 --cfg cfgs/BAT_Car.yaml  --checkpoint /path/to/checkpoint/xxx.ckpt --test


This codebase produces better results than those we report in our original paper.

Model Category Success Precision Checkpoint
BAT Car 65.37 78.88 pretrained_models/bat_kitti_car.ckpt
BAT Pedestrian 45.74 74.53 pretrained_models/bat_kitti_pedestrian.ckpt

Two Trained BAT models for KITTI dataset are provided in the pretrained_models directory. To reproduce the results, simply run the code with the corresponding .yaml file and checkpoint. For example, to reproduce the tracking results on Car, just run:

python --gpu 0 1 --cfg cfgs/BAT_Car.yaml  --checkpoint ./pretrained_models/bat_kitti_car.ckpt --test


  • [x] DDP support
  • [x] Multi-gpus testing
  • [ ] Add NuScenes dataset
  • [ ] Add codes for visualization
  • [ ] Add support for more methods


  • This repo is built upon P2B and SC3D.
  • Thank Erik Wijmans for his pytorch implementation of PointNet++