This repo is an unofficial TensorFlow implementation of Semantic Query Network (SQN). Yet, it achieves comparable or even better performance on S3DIS as the SQN paper (w/o any additional training strategies, e.g., re-training w. pseudo labels), check Results section for details.

New (Nov. 4, 2021): this SQN implementation achieves further performance boost using imbalanced learning techniques. Under 0.01%, it achieves 48.98% mIoU, outperforming the official SQN by 3.68%..

New update(Oct. 28, 2021): our repo also outperforms recent SOTA One Thing One Click on S3DIS under 0.02% weak labels, achieving 50.72 mIoU. check Results section for details.

Our repo achieves better performance (47.56% mIoU) on S3DIS than the original SQN paper (45.30% mIoU) under w. only 0.01% weak labels. Check its checkpoint for details.

The latest codes are tested on two Ubuntu settings:

• Ubuntu 18.04, Nvidia 1080, CUDA 10.1, TensorFlow 1.13 and Python 3.6
• Ubuntu 18.04, Nvidia 3090, CUDA 11.3, TensorFlow 1.13 and Python 3.6

git clone https://github.com/PointCloudYC/SQN_tensorflow && cd SQN_tensorflow 

create a conda environment

# DL is the conda env name conda create -n DL python=3.5 source activate DL pip install -r helper_requirements.txt # compile the sub-sampling and knn op sh compile_op.sh 

For more details to set up the development environment, check the official RandLA-Net repo.

You can download the S3DIS dataset from here (4.8 GB). You only need to download the file named Stanford3dDataset_v1.2.zip, unzip and move (or link) it to a folder. (same as the RandLA-Net repo setting.)

# assume S3DIS dataset is downloaded at /media/yinchao/dataset/S3DIS ln -s /media/yinchao/dataset/S3DIS ./data/S3DIS/Stanford3dDataset_v1.2_Aligned_Version 

You can use the s3dis-prepare-sqn.sh script to prepare the S3DIS dataset with weak labels.

# prepare the dataset, each room (Note: each area is preprocessed in the CLoserLook3D code) will result in four files 1 file in the original_ply folder for raw_pc.ply, and 3 files in the input_0.040 for sub_pc.py, sub_pc.kdtree, and project_indices file for each raw point, check data_prepare_s3dis_sqn.py for details. python utils/data_prepare_s3dis_sqn.py # check #rooms in npy format, should be 272 rooms find *.npy | wc -l 

The data file structure should look like:

<root> ├── ... ├── data │ └── S3DIS │ └── Stanford3dDataset_v1.2_Aligned_Version │ ├── Area_1 │ ├── Area_2 │ ├── Area_3 │ ├── Area_4 │ ├── Area_5 │ └── Area_6 │ └── input_0.040 │ └── original_ply │ └── weak_label_0.01 └── ... 

Only tf_ops/3d_interpolation CUDA ops need to be compiled, which will used for three trilinear interpolation.

check the tf_interpolate_compile.sh; You may need to tailor the CUDA_ROOT and TF_ROOT path according to your own system.

#/bin/bash CUDA_ROOT="/usr/local/cuda-10.1" TF_ROOT="/home/yinchao/miniconda3/envs/DL/lib/python3.6/site-packages/tensorflow" # TF1.4 (Note: -L ${TF_ROOT} should have a space in between) g++ -std=c++11 tf_interpolate.cpp -o tf_interpolate_so.so -shared -fPIC -I ${TF_ROOT}/include -I ${CUDA_ROOT}/include -I ${TF_ROOT}/include/external/nsync/public -lcudart -L ${CUDA_ROOT}/lib64/ -L ${TF_ROOT} -ltensorflow_framework -O2 # -D_GLIBCXX_USE_CXX11_ABI=0 

For more details, check Charles' PointNet2

To train the SQN, run this command:

python main_S3DIS_Sqn.py \ --gpu 0 \ --mode train \ --test_area 5 

For more arguments, see main_S3DIS_Sqn.py or use python main_S3DIS_Sqn.py --help to see documentation.

P.S.: you can use run-s3dis-Sqn.sh bash script to train multiple settings or do ablation study.

To evaluate our model on S3DIS, run:

python main_S3DIS_Sqn.py \ --gpu 0 \ --mode test \ --test_area 5 

For more arguments, see main_S3DIS_Sqn.py or use python main_S3DIS_Sqn.py --help to see documentation.

Our SQN achieves the following performance on S3DIS:

We use Nvidia 1080 GPU to train the replicated SQN with a small batch size; The performance might be improved if a powerful GPU w. larger memory could be used. We will update this performance table with a Nvidia 3090 in the future.

Note: experiments are still in progress due to my slow GPU. Stay in tuned.

Those numbers surpassing the official SQN is highlighted in bold in the table.

You can download pre-trained models and training log here:

• weak ratio 10% (todo)
• weak ratio 1% (todo)
• weak ratio 0.1%
• weak ratio 0.01%
• weak ratio 0.0067%

Each check point folder has the following files:

Use the below script to run checkpoint model on S3DIS:

python -B main_S3DIS_Sqn.py \ --gpu 0 \ --mode test \ --test_area 5 \ --model_path [your_checkpoint_path, e.g., /path/xx/snap-27001; no need add the file extension] 

Our pytorch codes borrowed a lot from official RandLA-Net and the custom trilinear interoplation CUDA ops are modified from official Pointnet2.

• re-train w. pseudo labels to further improve performance
• apply self-supervised learning techniques, e.g., Contrastive Learning.
• implement the training strategy mentioned in the Appendix of the paper.
• ablation study
• benchmark weak supervision
• dataset preparation implementation
• SQN architecture implementation consisting of encoder and query network
• training an evaluation loop modification

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

@code{SQN_tensorflow_yc, Author = {YIN, Chao}, Title = {SQN TensorFlow implementation}, Journal = {https://github.com/PointCloudYC/SQN_tensorflow}, Year = {2021} } @article{hu2021sqn, title={SQN: Weakly-Supervised Semantic Segmentation of Large-Scale 3D Point Clouds with 1000x Fewer Labels}, author={Hu, Qingyong and Yang, Bo and Fang, Guangchi and Guo, Yulan and Leonardis, Ales and Trigoni, Niki and Markham, Andrew}, journal={arXiv preprint arXiv:2104.04891}, year={2021} }