PiMAE: Point cloud and Image Interactive Masked Autoencoders for 3D Object Detecion

Accepted to CVPR2023. 🔥

Paper | BibTeX | Webpage

Pytorch pretraining & downstream-training code for PiMAE. We propose a MAE-based self-supervised pre-training framework that promotes 3D and 2D interaction to improve model performance on downstream object detection tasks.

📣 check out another lattest work from us, I2P-MAE, which obtains superior 3D representations from 2D pre-trained models via Image-to-Point Masked Autoencoders. 📣

Model Zoo

We provide our pretrained weights (on SUNRGBD) and finetuned transformer-based baseline models (on SUNRGBD, ScanNetV2, KITTI), including 3DETR, DETR, MonoDETR.

	task	name	Dataset	AP(gain)	download
0	Pretrain	PiMAE	SUNRGBD	-	model |
1	3D Object Detection	3DETR	SUNRGBD	[AP25] 59.4(+1.4)	model | logs
2	3D Object Detection	3DETR	ScanNetV2	[AP25] 62.6(+0.5)	model | logs
5	Monocular 3D Object Detection	MonoDETR	KITTI	[Easy] 26.6(+3.5)	model | logs
6	2D Object Detection	DETR	ScanNetV2	[AP50] 46.5(+6.7)	model | logs

Easy-Usage

We have provided a easy tutorial to use PiMAE's pre-trained 3D extractor. You can easily modify the code to fit in your model.

Get our pretrained models from here and place it as ./Pretrain/pimae.pth.

Install minimum required dependencies then simply run the tutorial code by:

pip install torch torchvision

python Pretrain/tutorial_load.py

Pretraining

Install

First, clone this repository into your local machine.

git clone https://github.com/BLVLab/PiMAE.git

Next, install required dependencies.

cd Pretrain

sh install.sh

Data Preparation

We follow the VoteNet to preprocess our data. The instructions for preparing SUN RGB-D are here. Remember to Edit the dataset paths in Pretrain/datasets/sunrgbd.py.

Training

python main.py --config cfgs/pretrain_JD_pc2img.yaml --exp_name pimae

Visualization

To get reconstruction visualization like this.

python main_vis.py \
    --test \
    --ckpts ./experiments/pretrain/cfgs/pimae/ckpt-last.pth \
    --config ./experiments/pretrain/cfgs/pimae/config.yaml \
    --exp_name vis_pimae \

Downstream

3DETR

Follow 3DETR codebase to prepare the training data (SUNRGBD & ScanNetV2).

Install required dependencies by

cd Downstream/3detr

sh install.sh

Run the training code (you can specify training configure in the script)

sh run.sh

MonoDETR

Follow MonoDETR codebase to prepare the training data (KITTI). Install required dependencies by

cd Downstream/MonoDETR

sh install.sh

Run the code for training and testing (remember to check monodetr.yaml where we specify path to pimae weights).

bash train.sh configs/monodetr.yaml > logs/monodetr.log # training

bash test.sh configs/monodetr.yaml # testing

DETR

Follow DETR to prepare data and required dependencies. Then train it by

cd Downstream/detr/d2

python train_net.py --config configs/detr_256_6_6_torchvision.yaml --num-gpus 8

Acknowledgement

This repository is based on 3DETR, MonoDETR, DETR, timm, MAE repositories, we thank them for their great work.

Citation

If you find this repository helpful, please consider citing our work:

@inproceedings{chen2023pimae，
  title={PiMAE: Point Cloud and Image Interactive Masked Autoencoders for 3D Object Detection},
  author = {Chen, Anthony and Zhang, Kevin and Zhang, Renrui and Wang, Zihan and Lu, Yuheng and Guo, Yandong and Zhang, Shanghang},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  year={2023}
}