ChenJunX1an / SDNet

Code of SD-Net: Spatially-Disentangled Point Cloud Completion Network

SD-Net: Spatially-Disentangled Point Cloud Completion Network

[Paper]

This repository contains PyTorch implementation for SD-Net: Spatially-Disentangled Point Cloud Completion Network (ACM MM 2023 Oral Presentation).

pip install -r requirements.txt

NOTE: PyTorch >= 1.7 and GCC >= 4.9 are required.

# Chamfer Distance
bash install.sh

The solution for a common bug in chamfer distance installation can be found in Issue #6

# PointNet++
pip install "git+git://github.com/erikwijmans/Pointnet2_PyTorch.git#egg=pointnet2_ops&subdirectory=pointnet2_ops_lib"
# GPU kNN
pip install --upgrade https://github.com/unlimblue/KNN_CUDA/releases/download/0.2/KNN_CUDA-0.2-py3-none-any.whl

Dataset

The details of our new ShapeNet-55/34 datasets and other existing datasets can be found in DATASET.md.

Evaluation

To evaluate a pre-trained SDNet model on the Dataset with single GPU, run:

bash ./scripts/test.sh <GPU_IDS>  \
    --ckpts <path> \
    --config <config> \
    --exp_name <name> \
    [--mode <easy/median/hard>]

Some examples:

Test the SDNet pretrained model on ShapeNet55 benchmark (easy mode):

bash ./scripts/test.sh 0 \
    --ckpts ./pretrained/SDNet_ShapeNet55.pth \
    --config ./cfgs/ShapeNet55_models/PoinTr.yaml \
    --mode easy \
    --exp_name example

Training

To train a point cloud completion model from scratch, run:

# Use DistributedDataParallel (DDP)
bash ./scripts/dist_train.sh <NUM_GPU> <port> \
    --config <config> \
    --exp_name <name> \
    [--resume] \
    [--start_ckpts <path>] \
    [--val_freq <int>]
# or just use DataParallel (DP)
bash ./scripts/train.sh <GPUIDS> \
    --config <config> \
    --exp_name <name> \
    [--resume] \
    [--start_ckpts <path>] \
    [--val_freq <int>]

Some examples:

Train a SDNet model on ShapeNet55 benchmark with 2 gpus:

CUDA_VISIBLE_DEVICES=0,1 bash ./scripts/dist_train.sh 2 13232 \
    --config ./cfgs/ShapeNet55_models/SDNet.yaml \
    --exp_name example

Resume a checkpoint:

CUDA_VISIBLE_DEVICES=0,1 bash ./scripts/dist_train.sh 2 13232 \
    --config ./cfgs/ShapeNet55_models/SDNet.yaml \
    --exp_name example --resume

Train a SDNet model with a single GPU:

bash ./scripts/train.sh 0 \
    --config ./cfgs/ShapeNet55_models/SDNet.yaml \
    --exp_name example

Acknowledgements

Our code is inspired by PoinTr, SnowflakeNet and SeedFormer.