SemAffiNet: Semantic-Affine Transformation for Point Cloud Segmentation

Created by Ziyi Wang, Yongming Rao, Xumin Yu, Jie Zhou, Jiwen Lu

This repository is an official implementation of SemAffiNet (CVPR 2022).

[arxiv]

Installation

Prerequisites

• Python 3.8
• PyTorch 1.8.1
• MinkowskiEngine 0.5.4
• timm
• open3d
• cv2, tensorboardX, imageio, SharedArray, scipy, tqdm, h5py

conda create -n semaffinet python=3.8
conda activate semaffinet
conda install pytorch==1.8.1 torchvision==0.9.1 cudatoolkit=10.2

git clone https://github.com/NVIDIA/MinkowskiEngine.git
cd MinkowskiEngine
export CXX=g++-7
conda install openblas
python setup.py install --blas_include_dirs=${CONDA_PREFIX}/include --blas=openblas

pip install timm
pip install open3d
pip install opencv-python
conda install tensorboardX, imageio, sharedarray, plyfile, tqdm

Usage

Data Preparation

• Download the official ScanNetV2 dataset.

• Prepare ScanNetV2 2D data: Please follow instructions in 3DMV repo.

  python prepare_2d_data.py --scannet_path SCANNET_INPUT_PATH --output_path SCANNET_OUTPUT_PATH --export_label_images
  

• Prepare ScanNetV2 3D data:

  python dataset/preprocess_3d_scannet.py
  

• Group ScanNetV2 2D views: preprocess 2D data and group multiple views of one scene into several groups. You will need to install pointnet2_ops from PointNet++ PyTorch repo to run the following command:

  python dataset/pregroup_2d_scannet.py
  

  You can also download our processed group results here.

• The data is expected to be in the following file structure:

  SemAffiNet/
  |-- data/
      |-- 2D/
          |-- scene0000_00/
              |-- color/
                  |-- 0.jpg
              |-- depth/
                  |-- 0.png
              |-- label/
                  |-- 0.png
              |-- pose/
                  |-- 0.txt
      |-- 3D/
          |-- train/
              |-- scene0000_00_vh_clean_2.pth
          |-- val/
              |-- scene0011_00_vh_clean_2.pth
          |-- test/
              |-- scene0707_00_vh_clean_2.pth
      |-- view_groups/
          |-- view_groups_train.pth
          |-- view_groups_val.pth
          |-- view_groups_test.pth
  

Init model preparation

Download the pre-trained resnet34d weights and place it in the initmodel folder. The pre-trained weight is from the timm repository.

Train

• ScanNetV2 5cm voxelization setting:

  bash tool/train.sh SemAffiNet_5cm config/scannet/semaffinet_5cm.yaml scannet 2
  

• ScanNetV2 2cm voxelization setting:

  bash tool/train.sh SemAffiNet_2cm config/scannet/semaffinet_2cm.yaml scannet 2
  

Test

• ScanNetV2 5cm voxelization setting:

  bash tool/test.sh SemAffiNet_5cm config/scannet/semaffinet_5cm.yaml scannet 2
  

• ScanNetV2 2cm voxelization setting:

  bash tool/test.sh SemAffiNet_2cm config/scannet/semaffinet_2cm.yaml scannet 2
  

Results

We provide pre-trained SemAffiNet models:

Dataset	URL	3D mIoU	2D mIoU
ScanNetV2 5cm	Google Drive	72.1	68.2
ScanNetV2 2cm	Google Drive	74.5	74.2

Please rename the checkpoints as model_best.pth.tar and organize the directory as the following structure:

    SemAffiNet/
    |-- initmodel/
        |-- resnet34d_ra2-f8dcfcaf.pth
    |-- Exp/
        |-- scannet/
            |-- SemAffiNet_2cm/
                |-- model/
                    |-- model_best.pth.tar
            |-- SemAffiNet_5cm/
                |-- model/
                    |-- model_best.pth.tar

Citation

If you find our work useful in your research, please consider citing:

@inproceedings{wang2022semaff,
  title={SemAffiNet: Semantic-Affine Transformation for Point Cloud Segmentation},
  author={Wang, Ziyi and Rao, Yongming and Yu, Xumin and Zhou, Jie and Lu, Jiwen},
  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2022}
}

Acknowledgements

Our code is inspired by BPNet. Some of the data preprocessing codes for ScanNetV2 are inspired by 3DMV.