CasA: A Cascade Attention Network for 3D Object Detection from LiDAR point clouds

CasA is a simple multi-stage 3D object detection framework based on a Cascade Attention design. CasA can be integrated into many SoTA 3D detectors and greatly improve their detection performance. The paper of "CasA: A Cascade Attention Network for 3D Object Detection from LiDAR point clouds" can be found here.
This code is mostly built upon OpenPCDet. Note that, the CasA++ is based on a transfer learning framework: pre-training on Waymo and fine-tuning on KITTI. Since additional data has been included, we did not release the CasA++ codes.

Overview

• Cascade Attention Design
• Model Zoo
• Getting Started
• Citation

Cascade Attention Design

Cascade frameworks have been widely studied in 2D object detection but less investigated in 3D space. Conventional cascade structures use multiple separate sub-networks to sequentially refine region proposals. Such methods, however, have limited ability to measure proposal quality in all stages, and hard to achieve a desirable detection performance improvement in 3D space. We propose a new cascade framework, termed CasA, for 3D object detection from point clouds. CasA consists of a Region Proposal Network (RPN) and a Cascade Refinement Network (CRN). In this CRN, we designed a new Cascade Attention Module that uses multiple sub-networks and attention modules to aggregate the object features from different stages and progressively refine region proposals. CasA can be integrated into various two-stage 3D detectors and greatly improve their detection performance. Extensive experimental results on KITTI and Waymo datasets with various baseline detectors demonstrate the universality and superiority of our CasA. In particular, based on one variant of Voxel-RCNN, we achieve state-of-the-art results on KITTI 3D object detection benchmark.

Update Log

• 2022/10/15 Update a 3D multi-object tracker CasTrack based on the CasA detections, currently rank first on the KITTI tracking leader-board :fire:!

• 2022/9/30 Update details of installation. Update environment we tested. Update Spconv2.X support :rocket:!

• 2022/3/3 Initial update, achieve SOTA performance on the KITTI 3D detection leader-board

Model Zoo

KITTI 3D Object Detection Results

The results are the 3D detection performance of moderate difficulty on the val set of KITTI dataset. Currently, this repo supports CasA-PV, CasA-V, CasA-T and CasA-PV2. The base detectors are PV-RCNN, Voxel-RCNN, CT3D and PV-RCNN++, respectively.

• All released models are trained with 2 3090 GPUs and are available for download.
• These models are not suitable to directly report results on KITTI test set, please use slightly lower score threshold and train the models on all or 80% training data to achieve a desirable performance on KITTI test set.

PV-RCNN VS. CasA-PV

Voxel-RCNN VS. CasA-V

CT3D VS. CasA-T

PV-RCNN++ VS. CasA-PV2

Where * denodes reproduced results of a simplified version using their open-source codes.

Waymo Open Dataset Results

Here we provided two models on WOD, where CasA-V-center denotes that the center-based RPN are used. All models are trained with a single frame on 8 V100 GPUs, and the results of each cell here are mAP/mAPH calculated by the official Waymo evaluation metrics on the whole validation set (version 1.2).

Where * denodes reproduced results using their open-source codes.

We could not provide the above pretrained models due to Waymo Dataset License Agreement, but you could easily achieve similar performance by training with the default configs.

Getting Started

conda create -n spconv2 python=3.9
conda activate spconv2
pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
pip install numpy==1.19.5 protobuf==3.19.4 scikit-image==0.19.2 waymo-open-dataset-tf-2-5-0 nuscenes-devkit==1.0.5 spconv-cu111 numba scipy pyyaml easydict fire tqdm shapely matplotlib opencv-python addict pyquaternion awscli open3d pandas future pybind11 tensorboardX tensorboard Cython prefetch-generator

Environment we tested

Our released implementation is tested on.

• Ubuntu 18.04
• Python 3.6.9
• PyTorch 1.8.1
• Numba 0.53.1
• Spconv 1.2.1
• NVIDIA CUDA 11.1
• 8x Tesla V100 GPUs

We also tested on.

• Ubuntu 18.04
• Python 3.9.13
• PyTorch 1.8.1
• Numba 0.53.1
• Spconv 2.1.22 # pip install spconv-cu111
• NVIDIA CUDA 11.1
• 2x 3090 GPUs

Prepare Dataset

KITTI Dataset

• Please download the official KITTI 3D object detection dataset and organize the downloaded files as follows (the road planes could be downloaded from [road plane], which are optional for data augmentation in the training):

CasA
├── data
│   ├── kitti
│   │   │── ImageSets
│   │   │── training
│   │   │   ├──calib & velodyne & label_2 & image_2 & (optional: planes)
│   │   │── testing
│   │   │   ├──calib & velodyne & image_2
├── pcdet
├── tools

Run following command to creat dataset infos:

python3 -m pcdet.datasets.kitti.kitti_dataset create_kitti_infos tools/cfgs/dataset_configs/kitti_dataset.yaml

Waymo Dataset

CasA
├── data
│   ├── waymo
│   │   │── ImageSets
│   │   │── raw_data
│   │   │   │── segment-xxxxxxxx.tfrecord
|   |   |   |── ...
|   |   |── waymo_processed_data_train_val_test
│   │   │   │── segment-xxxxxxxx/
|   |   |   |── ...
│   │   │── pcdet_waymo_track_dbinfos_train_cp.pkl
│   │   │── waymo_infos_test.pkl
│   │   │── waymo_infos_train.pkl
│   │   │── waymo_infos_val.pkl
├── pcdet
├── tools

Run following command to creat dataset infos:

python3 -m pcdet.datasets.waymo.waymo_tracking_dataset --cfg_file tools/cfgs/dataset_configs/waymo_tracking_dataset.yaml 

Installation

git clone https://github.com/hailanyi/CasA.git
cd CasA
python3 setup.py develop

Training and Evaluation

Evaluation

cd tools
python3 test.py --cfg_file ${CONFIG_FILE} --batch_size ${BATCH_SIZE} --ckpt ${CKPT}

For example, if you test the CasA-V model:

cd tools
python3 test.py --cfg_file cfgs/kitti_models/CasA-V.yaml --ckpt CasA-V.pth

Multiple GPU test: you need modify the gpu number in the dist_test.sh and run

sh dist_test.sh 

The log infos are saved into log-test.txt You can run cat log-test.txt to view the test results.

Training

cd tools
python3 train.py --cfg_file ${CONFIG_FILE}

For example, if you train the CasA-V model:

cd tools
python3 train.py --cfg_file cfgs/kitti_models/CasA-V.yaml

Multiple GPU train: you can modify the gpu number in the dist_train.sh and run

sh dist_train.sh

The log infos are saved into log.txt You can run cat log.txt to view the training process.

Acknowledgement

This repo is developed from OpenPCDet 0.3, we thank shaoshuai shi for his implementation of OpenPCDet.

Citation

If you find this project useful in your research, please consider cite:

@article{casa2022,
    title={CasA: A Cascade Attention Network for 3D Object Detection from LiDAR point clouds},
    author={Wu, Hai and Deng, Jinhao and Wen, Chenglu and Li, Xin and Wang, Cheng},
    journal={IEEE Transactions on Geoscience and Remote Sensing},
    year={2022}
}