Home

Awesome

Diffusion-SS3D: Diffusion Model for Semi-supervised 3D Object Detection

Introduction

This is the code release of our NeurIPS 2023 paper "Diffusion-SS3D: Diffusion Model for Semi-supervised 3D Object Detection".

In this repository, we provide an implementation (with Pytorch) based on VoteNet, 3DIoUMatch and DiffusionDet with some modification, as well as the training and evaluation scripts on ScanNet.

Installation

This repo is tested under the following environment:

You can follow the steps below to install the above dependencies:

# Create and activate virtualenv
conda create -n myenv python=3.7.16
conda activate myenv

Install NumPy

pip install numpy==1.21.5

Install PyTorch according to your CUDA version.

conda install pytorch==1.10.1 torchvision==0.11.2 torchaudio==0.10.1 cudatoolkit=11.3 -c pytorch -c conda-forge

Install TensorFlow (for TensorBoard) (We tested this repo with TensorFlow 2.9.1.)

pip install tensorflow

Install and register for wandb

Wandb QuickStart When you have a wandb account, please substitute "entity" in wandb.init() with your own account name. (for both train.py and pretrain.py)

pip install wandb

Compile the CUDA code for PointNet++, which is used in the backbone network. If you have any probelm about this part, you can refer to Pointnet2/Pointnet++ PyTorch

cd pointnet2
python setup.py install

Compile the CUDA code for general 3D IoU calculation in OpenPCDet:

cd OpenPCDet
python setup.py develop

Install dependencies:

pip install -r requirements.txt

Dataset

ScanNet

Please download the ScanNet data following the README in scannet folder.

Download Pre-trained and Trained Models

We provide the pre-trained models of ScanNet 5%: Link

We provide the trained model of ScanNet 5%: Link

Pre-training

sh run_pretrain.sh <GPU_ID> <LOG_DIR> <DATA_RATIO> <LABELED_LIST>

For example:

sh run_pretrain.sh 0 results/pretrain 0.05 scannetv2_train_0.05.txt

Training

sh run_train.sh <GPU_ID> <LOG_DIR> <DATA_RATIO> <LABELED_LIST> <PRETRAINED_CKPT>

For example:

sh run_train.sh 0 results/train 0.05 scannetv2_train_0.05.txt results/pretrain/best_checkpoint_sum.tar

Evaluation

sh run_eval.sh <GPU_ID> <LOG_DIR> <DATA_RATIO> <LABELED_LIST> <CKPT>

For example:

sh run_eval.sh 0 results/eval 0.05 scannetv2_train_0.05.txt results/train/best_checkpoint_sum.tar

Evaluation with IoU Optimization

sh run_eval_opt.sh <GPU_ID> <LOG_DIR> <DATA_RATIO> <LABELED_LIST> <CKPT> <OPT_RATE>

The number of steps (of optimization) is by default 10.

Acknowledgements

Our implementation uses code from the following repositories: