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.

This repo is tested under the following environment:

• Python 3.7.16
• NumPy 1.21.5
• pytorch 1.10.1, cuda 11.3, torchvision 0.11.2
• Pointnet2 from Pointnet2/Pointnet++ PyTorch
• OpenPCDet from OpenPCDet

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

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

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

We provide the trained model of ScanNet 5%: Link

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

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

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

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.

Our implementation uses code from the following repositories:

• Deep Hough Voting for 3D Object Detection in Point Clouds
• 3DIoUMatch
• Pointnet2/Pointnet++ PyTorch
• OpenPCDet
• DiffusionDet