├── .gitignore ├── LICENSE ├── README.md ├── RandLANet.py ├── SqnNet.py ├── compile_op.sh ├── compute-point-distribution.sh ├── eval-s3dis-sqn.sh ├── helper_ply.py ├── helper_requirements.txt ├── helper_tf_util.py ├── helper_tool.py ├── imgs ├── checkpoint_log.jpg └── replication-roadmap.jpg ├── jobs_6_fold_cv_s3dis.sh ├── main_S3DIS.py ├── main_S3DIS_Sqn.py ├── main_S3DIS_weak_benchmark.py ├── prepare-s3dis-sqn.sh ├── run-s3dis-Sqn.sh ├── tester_S3DIS.py ├── tester_S3DIS_Sqn.py ├── tf_ops ├── 3d_interpolation │ ├── interpolate.cpp │ ├── tf_interpolate.cpp │ ├── tf_interpolate.py │ ├── tf_interpolate_compile.sh │ ├── tf_interpolate_op_test.py │ └── visu_interpolation.py ├── readme.md └── test.py └── utils ├── 6_fold_cv.py ├── compute_weak_point_distribution.py ├── cpp_wrappers ├── compile_wrappers.sh ├── cpp_subsampling │ ├── grid_subsampling │ │ ├── grid_subsampling.cpp │ │ └── grid_subsampling.h │ ├── setup.py │ └── wrapper.cpp └── cpp_utils │ ├── cloud │ ├── cloud.cpp │ └── cloud.h │ └── nanoflann │ └── nanoflann.hpp ├── data_prepare_s3dis.py ├── data_prepare_s3dis_sqn.py ├── knn_interpolation.py ├── meta ├── anno_paths.txt └── class_names.txt └── nearest_neighbors ├── KDTreeTableAdaptor.h ├── knn.pyx ├── knn_.cxx ├── knn_.h ├── nanoflann.hpp ├── setup.py └── test.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | # data 132 | data/* 133 | 134 | # vscode 135 | .vscode/* 136 | 137 | # outputs of compiled ops 138 | utils/*/build/* 139 | utils/cpp_wrappers/cpp_subsampling/build/* 140 | utils/cpp_wrappers/cpp_subsampling/*.so 141 | utils/nearest_neighbors/lib/python/* 142 | utils/nearest_neighbors/knn.cpp 143 | __pycache__/* 144 | */__pycache__/* 145 | 146 | # logs 147 | train_log/* 148 | log_*.txt 149 | 150 | # tf_ops 151 | tf_ops/*/*.so 152 | tf_ops/*/__pycache__/* 153 | tf_ops/*/*.py.bak 154 | 155 | # checkpoints 156 | results-Sqn/*/snapshots/* 157 | results_Sqn/Log_*/snapshots/* 158 | 159 | train_log_Sqn/* 160 | 161 | # bugs 162 | bugs/* 163 | imgs/records/* 164 | yc.md 165 | results_Sqn/Log_*/* 166 | test/* 167 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 PointCloudYC 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # SQN_tensorflow 2 | 3 | This repo is an unofficial TensorFlow implementation of **[Semantic Query Network (SQN)](https://arxiv.org/abs/2104.04891)**. 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](#results) section** for details. 4 | 5 | >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%.**. 6 | 7 | >New update(Oct. 28, 2021): our repo also outperforms recent SOTA [One Thing One Click](https://arxiv.org/abs/2104.02246) on S3DIS under **0.02%** weak labels, achieving 50.72 mIoU. **check [Results](#results) section** for details. 8 | 9 | >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](https://hkustconnect-my.sharepoint.com/:f:/g/personal/cyinac_connect_ust_hk/EvFhPeuZKp5Hq8IR-XpsLJkB9EObsJHHFhOevRTzJ36mXg?e=TQbCJ3) for details. 10 | 11 | 12 | ## Requirements 13 | 14 | The latest codes are tested on two Ubuntu settings: 15 | 16 | - [x] Ubuntu 18.04, Nvidia 1080, CUDA 10.1, TensorFlow 1.13 and Python 3.6 17 | - [ ] Ubuntu 18.04, Nvidia 3090, CUDA 11.3, TensorFlow 1.13 and Python 3.6 18 | 19 | ### Clone the repository 20 | 21 | ``` 22 | git clone https://github.com/PointCloudYC/SQN_tensorflow && cd SQN_tensorflow 23 | ``` 24 | 25 | ### Setup python environment 26 | 27 | create a conda environment 28 | 29 | ``` 30 | # DL is the conda env name 31 | conda create -n DL python=3.5 32 | source activate DL 33 | pip install -r helper_requirements.txt 34 | # compile the sub-sampling and knn op 35 | sh compile_op.sh 36 | ``` 37 | 38 | For more details to set up the development environment, check [the official RandLA-Net repo](https://github.com/QingyongHu/RandLA-Net). 39 | 40 | ### Download S3DIS (and make a symlink) 41 | 42 | You can download the S3DIS dataset from [here](https://goo.gl/forms/4SoGp4KtH1jfRqEj2") (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.) 43 | 44 | ``` 45 | # assume S3DIS dataset is downloaded at /media/yinchao/dataset/S3DIS 46 | ln -s /media/yinchao/dataset/S3DIS ./data/S3DIS/Stanford3dDataset_v1.2_Aligned_Version 47 | ``` 48 | 49 | ### Preprocess S3DIS dataset 50 | 51 | You can use the `s3dis-prepare-sqn.sh` script to prepare the S3DIS dataset with weak labels. 52 | 53 | ``` 54 | # 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. 55 | 56 | python utils/data_prepare_s3dis_sqn.py 57 | 58 | # check #rooms in npy format, should be 272 rooms 59 | find *.npy | wc -l 60 | ``` 61 | 62 | The data file structure should look like: 63 | 64 | ``` 65 | 66 | ├── ... 67 | ├── data 68 | │ └── S3DIS 69 | │ └── Stanford3dDataset_v1.2_Aligned_Version 70 | │ ├── Area_1 71 | │ ├── Area_2 72 | │ ├── Area_3 73 | │ ├── Area_4 74 | │ ├── Area_5 75 | │ └── Area_6 76 | │ └── input_0.040 77 | │ └── original_ply 78 | │ └── weak_label_0.01 79 | └── ... 80 | ``` 81 | 82 | ### Compile custom CUDA tf_ops 83 | 84 | Only `tf_ops/3d_interpolation` CUDA ops need to be compiled, which will used for three trilinear interpolation. 85 | 86 | check the `tf_interpolate_compile.sh`; You may need to tailor the `CUDA_ROOT` and `TF_ROOT` path according to your own system. 87 | 88 | ``` 89 | #/bin/bash 90 | CUDA_ROOT="/usr/local/cuda-10.1" 91 | TF_ROOT="/home/yinchao/miniconda3/envs/DL/lib/python3.6/site-packages/tensorflow" 92 | 93 | # TF1.4 (Note: -L ${TF_ROOT} should have a space in between) 94 | 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 95 | ``` 96 | 97 | For more details, check [Charles' PointNet2](https://github.com/charlesq34/pointnet2) 98 | 99 | ## Training 100 | 101 | To train the SQN, run this command: 102 | 103 | ``` 104 | python main_S3DIS_Sqn.py \ 105 | --gpu 0 \ 106 | --mode train \ 107 | --test_area 5 108 | ``` 109 | 110 | >For more arguments, see `main_S3DIS_Sqn.py` or use `python main_S3DIS_Sqn.py --help` to see documentation. 111 | 112 | P.S.: you can use `run-s3dis-Sqn.sh` bash script to train multiple settings or do ablation study. 113 | 114 | ## Evaluation 115 | 116 | To evaluate our model on S3DIS, run: 117 | 118 | ``` 119 | python main_S3DIS_Sqn.py \ 120 | --gpu 0 \ 121 | --mode test \ 122 | --test_area 5 123 | ``` 124 | 125 | >For more arguments, see `main_S3DIS_Sqn.py` or use `python main_S3DIS_Sqn.py --help` to see documentation. 126 | 127 | 128 | ## Results 129 | 130 | Our SQN achieves the following performance on S3DIS: 131 | >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. 132 | 133 | | Model | Weak ratio | mIoU(%) | Description| 134 | |-------|------------|-----------------------|------------| 135 | | SQN(Official)|100%| 63.73| trained with full labels| 136 | | SQN(Official)|10%| 64.67| Note: add **retrain w. pseudo labels**| 137 | | SQN(this repo)|10%| in progress| no retraining w. pseudo labels| 138 | | SQN(Official)|1%| 63.65| Note: add **retrain w. pseudo labels**| 139 | | SQN(this repo)|1%|in progress| no retraining w. pseudo labels| 140 | | SQN(Official)|0.1%| 61.41| Note: add **retrain w. pseudo labels**| 141 | | SQN(this repo)|0.1%| 55.25 | no retraining w. pseudo labels| 142 | | One-Thing-One-Click |0.02%| 50.1| CVPR2021 [paper](https://arxiv.org/abs/2104.02246) | 143 | | SQN(this repo)|0.02%| **50.72** | no retraining w. pseudo labels| 144 | | SQN(Official)|0.01%| 45.30| Note: add **retrain w. pseudo labels**| 145 | | SQN(this repo)|0.01%| **48.98** | no retraining w. pseudo labels| 146 | | SQN(this repo)|0.0067%| **46.81** | no retraining w. pseudo labels| 147 | | SQN(this repo)|0.005%| **45.27** | no retraining w. pseudo labels| 148 | | SQN(this repo)|0.002%| **39.54** | no retraining w. pseudo labels| 149 | | SQN(this repo)|0.001%| **35.65** | no retraining w. pseudo labels| 150 | 151 | Note: experiments are still in progress due to my slow GPU. Stay in tuned. 152 | >Those numbers surpassing the official SQN is highlighted in bold in the table. 153 | 154 | 155 | ## Pre-trained Models 156 | 157 | You can download pre-trained models and training log here: 158 | 159 | - [weak ratio 10% (todo)](TODO) 160 | - [weak ratio 1% (todo)](TODO) 161 | - [weak ratio 0.1%](https://hkustconnect-my.sharepoint.com/:f:/g/personal/cyinac_connect_ust_hk/EoS3AYSkrxZLovTfRvyV6xABGsDtOZxiu6bgBcAe0-S-dw?e=7JcwSr) 162 | - [weak ratio 0.01% ](https://hkustconnect-my.sharepoint.com/:f:/g/personal/cyinac_connect_ust_hk/EvFhPeuZKp5Hq8IR-XpsLJkB9EObsJHHFhOevRTzJ36mXg?e=TQbCJ3) 163 | - [weak ratio 0.0067%](https://hkustconnect-my.sharepoint.com/:f:/g/personal/cyinac_connect_ust_hk/EsMOq5fqavpOv2ayemdK9boB8-u22eFTcCDWQSPJSN7SbA?e=ivlXAA) 164 | 165 | 166 | Each check point folder has the following files: 167 | 168 | checkpoint folder 169 | 170 | Use the below script to run checkpoint model on S3DIS: 171 | 172 | ``` 173 | python -B main_S3DIS_Sqn.py \ 174 | --gpu 0 \ 175 | --mode test \ 176 | --test_area 5 \ 177 | --model_path [your_checkpoint_path, e.g., /path/xx/snap-27001; no need add the file extension] 178 | ``` 179 | 180 | ## Acknowledgements 181 | 182 | Our pytorch codes borrowed a lot from [official RandLA-Net](https://github.com/QingyongHu/RandLA-Net) and the custom trilinear interoplation CUDA ops are modified from [official Pointnet2](https://github.com/charlesq34/pointnet2). 183 | 184 | ## TODOs 185 | 186 | - [ ] **re-train w. pseudo labels to further improve performance** 187 | - [ ] **apply self-supervised learning techniques, e.g., Contrastive Learning**. 188 | - [ ] implement the training strategy mentioned in the Appendix of the paper. 189 | - [ ] ablation study 190 | - [ ] benchmark weak supervision 191 | - [x] dataset preparation implementation 192 | - [x] SQN architecture implementation consisting of encoder and query network 193 | - [x] training an evaluation loop modification 194 | 195 | roadmap 196 | 197 | 198 | 199 | ## Citation 200 | 201 | If you find our work useful in your research, please consider citing: 202 | 203 | ``` 204 | @code{SQN_tensorflow_yc, 205 | Author = {YIN, Chao}, 206 | Title = {SQN TensorFlow implementation}, 207 | Journal = {https://github.com/PointCloudYC/SQN_tensorflow}, 208 | Year = {2021} 209 | } 210 | 211 | @article{hu2021sqn, 212 | title={SQN: Weakly-Supervised Semantic Segmentation of Large-Scale 3D Point Clouds with 1000x Fewer Labels}, 213 | author={Hu, Qingyong and Yang, Bo and Fang, Guangchi and Guo, Yulan and Leonardis, Ales and Trigoni, Niki and Markham, Andrew}, 214 | journal={arXiv preprint arXiv:2104.04891}, 215 | year={2021} 216 | } 217 | ``` 218 | -------------------------------------------------------------------------------- /compile_op.sh: -------------------------------------------------------------------------------- 1 | cd utils/nearest_neighbors 2 | python setup.py install --home="." 3 | cd ../../ 4 | 5 | cd utils/cpp_wrappers 6 | sh compile_wrappers.sh 7 | cd ../../../ -------------------------------------------------------------------------------- /compute-point-distribution.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | SUB_GRID_SIZES=(0.04) 4 | WEAK_LABEL_RATIOS=(0.001 0.0001) 5 | 6 | for sub_grid_size in "${SUB_GRID_SIZES[@]}"; do 7 | for weak_label_ratio in "${WEAK_LABEL_RATIOS[@]}"; do 8 | 9 | echo "sub_grid_size is ${sub_grid_size}" 10 | echo "weak_label_ratio is ${weak_label_ratio}" 11 | 12 | time python utils/compute_weak_point_distribution.py --sub_grid_size ${sub_grid_size} \ 13 | --weak_label_ratio ${weak_label_ratio} 14 | 15 | echo "sub_grid_size is ${sub_grid_size}" 16 | echo "weak_label_ratio is ${weak_label_ratio}" 17 | 18 | done 19 | done 20 | 21 | echo "finish computing!" -------------------------------------------------------------------------------- /eval-s3dis-sqn.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | mode='test' 4 | gpu=0 5 | # 4 costs about 8GB GPU, others can be used on grp server GPUs 6 | batch_size=(3) # 3 just 8GB, 4--> OOM on 1080 GPU 7 | # 8 works on my nvidia 1080 gpu, others can be tried on grp server GPUs 8 | val_batch_size=4 # 6-->system collapse 9 | num_points=40960 10 | max_epoch=100 # 400 11 | 12 | LOG_NAMES=( 13 | 'Log_weak_0.001_2021-11-01_16-25-48_WCE' 14 | 'Log_weak_0.0001_2021-11-03_08-03-05_WCE' 15 | 'Log_weak_2e-05_2021-11-04_15-31-17_WCE' 16 | 'Log_weak_5e-05_2021-11-04_03-36-32_WCE' 17 | ) 18 | 19 | WEAK_LABEL_RATIOS=( 20 | 0.001 21 | 0.0001 22 | 0.00005 23 | 0.00002 24 | ) 25 | 26 | # check above each item from LOG_NAMES's snapishots folder 27 | EPOCHS=( 28 | 36001 29 | 35501 30 | 27001 31 | 23501 32 | ) 33 | 34 | for i in "${!LOG_NAMES[@]}"; do 35 | 36 | echo "weak_label_ratio: ${WEAK_LABEL_RATIOS[$i]}" 37 | echo "model_path: results_Sqn/${LOG_NAMES[$i]}/snapshots/snap-${EPOCHS[$i]}" 38 | 39 | time python -B main_S3DIS_Sqn.py \ 40 | --gpu ${gpu} \ 41 | --mode ${mode} \ 42 | --test_area 5 \ 43 | --batch_size ${batch_size} \ 44 | --val_batch_size ${val_batch_size} \ 45 | --num_points ${num_points} \ 46 | --max_epoch ${max_epoch} \ 47 | --weak_label_ratio ${WEAK_LABEL_RATIOS[$i]} \ 48 | --model_path results_Sqn/${LOG_NAMES[$i]}/snapshots/snap-${EPOCHS[$i]} 49 | 50 | echo "weak_label_ratio: ${WEAK_LABEL_RATIOS[$i]}" 51 | echo "model_path: results_Sqn/${LOG_NAMES[$i]}/snapshots/snap-${EPOCHS[$i]}" 52 | done 53 | 54 | echo "finish testing." -------------------------------------------------------------------------------- /helper_ply.py: -------------------------------------------------------------------------------- 1 | # 2 | # 3 | # 0===============================0 4 | # | PLY files reader/writer | 5 | # 0===============================0 6 | # 7 | # 8 | # ---------------------------------------------------------------------------------------------------------------------- 9 | # 10 | # function to read/write .ply files 11 | # 12 | # ---------------------------------------------------------------------------------------------------------------------- 13 | # 14 | # Hugues THOMAS - 10/02/2017 15 | # 16 | 17 | 18 | # ---------------------------------------------------------------------------------------------------------------------- 19 | # 20 | # Imports and global variables 21 | # \**********************************/ 22 | # 23 | 24 | 25 | # Basic libs 26 | import numpy as np 27 | import sys 28 | 29 | 30 | # Define PLY types 31 | ply_dtypes = dict([ 32 | (b'int8', 'i1'), 33 | (b'char', 'i1'), 34 | (b'uint8', 'u1'), 35 | (b'uchar', 'u1'), 36 | (b'int16', 'i2'), 37 | (b'short', 'i2'), 38 | (b'uint16', 'u2'), 39 | (b'ushort', 'u2'), 40 | (b'int32', 'i4'), 41 | (b'int', 'i4'), 42 | (b'uint32', 'u4'), 43 | (b'uint', 'u4'), 44 | (b'float32', 'f4'), 45 | (b'float', 'f4'), 46 | (b'float64', 'f8'), 47 | (b'double', 'f8') 48 | ]) 49 | 50 | # Numpy reader format 51 | valid_formats = {'ascii': '', 'binary_big_endian': '>', 52 | 'binary_little_endian': '<'} 53 | 54 | 55 | # ---------------------------------------------------------------------------------------------------------------------- 56 | # 57 | # Functions 58 | # \***************/ 59 | # 60 | 61 | 62 | def parse_header(plyfile, ext): 63 | # Variables 64 | line = [] 65 | properties = [] 66 | num_points = None 67 | 68 | while b'end_header' not in line and line != b'': 69 | line = plyfile.readline() 70 | 71 | if b'element' in line: 72 | line = line.split() 73 | num_points = int(line[2]) 74 | 75 | elif b'property' in line: 76 | line = line.split() 77 | properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) 78 | 79 | return num_points, properties 80 | 81 | 82 | def parse_mesh_header(plyfile, ext): 83 | # Variables 84 | line = [] 85 | vertex_properties = [] 86 | num_points = None 87 | num_faces = None 88 | current_element = None 89 | 90 | 91 | while b'end_header' not in line and line != b'': 92 | line = plyfile.readline() 93 | 94 | # Find point element 95 | if b'element vertex' in line: 96 | current_element = 'vertex' 97 | line = line.split() 98 | num_points = int(line[2]) 99 | 100 | elif b'element face' in line: 101 | current_element = 'face' 102 | line = line.split() 103 | num_faces = int(line[2]) 104 | 105 | elif b'property' in line: 106 | if current_element == 'vertex': 107 | line = line.split() 108 | vertex_properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) 109 | elif current_element == 'vertex': 110 | if not line.startswith('property list uchar int'): 111 | raise ValueError('Unsupported faces property : ' + line) 112 | 113 | return num_points, num_faces, vertex_properties 114 | 115 | 116 | def read_ply(filename, triangular_mesh=False): 117 | """ 118 | Read ".ply" files 119 | 120 | Parameters 121 | ---------- 122 | filename : string 123 | the name of the file to read. 124 | 125 | Returns 126 | ------- 127 | result : array 128 | data stored in the file 129 | 130 | Examples 131 | -------- 132 | Store data in file 133 | 134 | >>> points = np.random.rand(5, 3) 135 | >>> values = np.random.randint(2, size=10) 136 | >>> write_ply('example.ply', [points, values], ['x', 'y', 'z', 'values']) 137 | 138 | Read the file 139 | 140 | >>> data = read_ply('example.ply') 141 | >>> values = data['values'] 142 | array([0, 0, 1, 1, 0]) 143 | 144 | >>> points = np.vstack((data['x'], data['y'], data['z'])).T 145 | array([[ 0.466 0.595 0.324] 146 | [ 0.538 0.407 0.654] 147 | [ 0.850 0.018 0.988] 148 | [ 0.395 0.394 0.363] 149 | [ 0.873 0.996 0.092]]) 150 | 151 | """ 152 | 153 | with open(filename, 'rb') as plyfile: 154 | 155 | 156 | # Check if the file start with ply 157 | if b'ply' not in plyfile.readline(): 158 | raise ValueError('The file does not start whith the word ply') 159 | 160 | # get binary_little/big or ascii 161 | fmt = plyfile.readline().split()[1].decode() 162 | if fmt == "ascii": 163 | raise ValueError('The file is not binary') 164 | 165 | # get extension for building the numpy dtypes 166 | ext = valid_formats[fmt] 167 | 168 | # PointCloud reader vs mesh reader 169 | if triangular_mesh: 170 | 171 | # Parse header 172 | num_points, num_faces, properties = parse_mesh_header(plyfile, ext) 173 | 174 | # Get point data 175 | vertex_data = np.fromfile(plyfile, dtype=properties, count=num_points) 176 | 177 | # Get face data 178 | face_properties = [('k', ext + 'u1'), 179 | ('v1', ext + 'i4'), 180 | ('v2', ext + 'i4'), 181 | ('v3', ext + 'i4')] 182 | faces_data = np.fromfile(plyfile, dtype=face_properties, count=num_faces) 183 | 184 | # Return vertex data and concatenated faces 185 | faces = np.vstack((faces_data['v1'], faces_data['v2'], faces_data['v3'])).T 186 | data = [vertex_data, faces] 187 | 188 | else: 189 | 190 | # Parse header 191 | num_points, properties = parse_header(plyfile, ext) 192 | 193 | # Get data 194 | data = np.fromfile(plyfile, dtype=properties, count=num_points) 195 | 196 | return data 197 | 198 | 199 | def header_properties(field_list, field_names): 200 | 201 | # List of lines to write 202 | lines = [] 203 | 204 | # First line describing element vertex 205 | lines.append('element vertex %d' % field_list[0].shape[0]) 206 | 207 | # Properties lines 208 | i = 0 209 | for fields in field_list: 210 | for field in fields.T: 211 | lines.append('property %s %s' % (field.dtype.name, field_names[i])) 212 | i += 1 213 | 214 | return lines 215 | 216 | 217 | def write_ply(filename, field_list, field_names, triangular_faces=None): 218 | """ 219 | Write ".ply" files 220 | 221 | Parameters 222 | ---------- 223 | filename : string 224 | the name of the file to which the data is saved. A '.ply' extension will be appended to the 225 | file name if it does no already have one. 226 | 227 | field_list : list, tuple, numpy array 228 | the fields to be saved in the ply file. Either a numpy array, a list of numpy arrays or a 229 | tuple of numpy arrays. Each 1D numpy array and each column of 2D numpy arrays are considered 230 | as one field. 231 | 232 | field_names : list 233 | the name of each fields as a list of strings. Has to be the same length as the number of 234 | fields. 235 | 236 | Examples 237 | -------- 238 | >>> points = np.random.rand(10, 3) 239 | >>> write_ply('example1.ply', points, ['x', 'y', 'z']) 240 | 241 | >>> values = np.random.randint(2, size=10) 242 | >>> write_ply('example2.ply', [points, values], ['x', 'y', 'z', 'values']) 243 | 244 | >>> colors = np.random.randint(255, size=(10,3), dtype=np.uint8) 245 | >>> field_names = ['x', 'y', 'z', 'red', 'green', 'blue', values'] 246 | >>> write_ply('example3.ply', [points, colors, values], field_names) 247 | 248 | """ 249 | 250 | # Format list input to the right form 251 | field_list = list(field_list) if (type(field_list) == list or type(field_list) == tuple) else list((field_list,)) 252 | for i, field in enumerate(field_list): 253 | if field.ndim < 2: 254 | field_list[i] = field.reshape(-1, 1) 255 | if field.ndim > 2: 256 | print('fields have more than 2 dimensions') 257 | return False 258 | 259 | # check all fields have the same number of data 260 | n_points = [field.shape[0] for field in field_list] 261 | if not np.all(np.equal(n_points, n_points[0])): 262 | print('wrong field dimensions') 263 | return False 264 | 265 | # Check if field_names and field_list have same nb of column 266 | n_fields = np.sum([field.shape[1] for field in field_list]) 267 | if (n_fields != len(field_names)): 268 | print('wrong number of field names') 269 | return False 270 | 271 | # Add extension if not there 272 | if not filename.endswith('.ply'): 273 | filename += '.ply' 274 | 275 | # open in text mode to write the header 276 | with open(filename, 'w') as plyfile: 277 | 278 | # First magical word 279 | header = ['ply'] 280 | 281 | # Encoding format 282 | header.append('format binary_' + sys.byteorder + '_endian 1.0') 283 | 284 | # Points properties description 285 | header.extend(header_properties(field_list, field_names)) 286 | 287 | # Add faces if needded 288 | if triangular_faces is not None: 289 | header.append('element face {:d}'.format(triangular_faces.shape[0])) 290 | header.append('property list uchar int vertex_indices') 291 | 292 | # End of header 293 | header.append('end_header') 294 | 295 | # Write all lines 296 | for line in header: 297 | plyfile.write("%s\n" % line) 298 | 299 | # open in binary/append to use tofile 300 | with open(filename, 'ab') as plyfile: 301 | 302 | # Create a structured array 303 | i = 0 304 | type_list = [] 305 | for fields in field_list: 306 | for field in fields.T: 307 | type_list += [(field_names[i], field.dtype.str)] 308 | i += 1 309 | data = np.empty(field_list[0].shape[0], dtype=type_list) 310 | i = 0 311 | for fields in field_list: 312 | for field in fields.T: 313 | data[field_names[i]] = field 314 | i += 1 315 | 316 | data.tofile(plyfile) 317 | 318 | if triangular_faces is not None: 319 | triangular_faces = triangular_faces.astype(np.int32) 320 | type_list = [('k', 'uint8')] + [(str(ind), 'int32') for ind in range(3)] 321 | data = np.empty(triangular_faces.shape[0], dtype=type_list) 322 | data['k'] = np.full((triangular_faces.shape[0],), 3, dtype=np.uint8) 323 | data['0'] = triangular_faces[:, 0] 324 | data['1'] = triangular_faces[:, 1] 325 | data['2'] = triangular_faces[:, 2] 326 | data.tofile(plyfile) 327 | 328 | return True 329 | 330 | 331 | def describe_element(name, df): 332 | """ Takes the columns of the dataframe and builds a ply-like description 333 | 334 | Parameters 335 | ---------- 336 | name: str 337 | df: pandas DataFrame 338 | 339 | Returns 340 | ------- 341 | element: list[str] 342 | """ 343 | property_formats = {'f': 'float', 'u': 'uchar', 'i': 'int'} 344 | element = ['element ' + name + ' ' + str(len(df))] 345 | 346 | if name == 'face': 347 | element.append("property list uchar int points_indices") 348 | 349 | else: 350 | for i in range(len(df.columns)): 351 | # get first letter of dtype to infer format 352 | f = property_formats[str(df.dtypes[i])[0]] 353 | element.append('property ' + f + ' ' + df.columns.values[i]) 354 | 355 | return element 356 | 357 | -------------------------------------------------------------------------------- /helper_requirements.txt: -------------------------------------------------------------------------------- 1 | numpy==1.16.1 2 | h5py==2.10.0 3 | cython==0.29.15 4 | open3d-python==0.3.0 5 | pandas==0.25.3 6 | scikit-learn==0.21.3 7 | scipy==1.4.1 8 | PyYAML==5.4 9 | -------------------------------------------------------------------------------- /helper_tool.py: -------------------------------------------------------------------------------- 1 | from open3d import linux as open3d 2 | from os.path import join 3 | import numpy as np 4 | import colorsys, random, os, sys 5 | import pandas as pd 6 | 7 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 8 | 9 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 10 | 11 | sys.path.append(BASE_DIR) 12 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 13 | 14 | import utils.cpp_wrappers.cpp_subsampling.grid_subsampling as cpp_subsampling 15 | import utils.nearest_neighbors.lib.python.nearest_neighbors as nearest_neighbors 16 | 17 | 18 | class ConfigSemanticKITTI: 19 | k_n = 16 # KNN 20 | num_layers = 4 # Number of layers 21 | num_points = 4096 * 11 # Number of input points 22 | num_classes = 19 # Number of valid classes 23 | sub_grid_size = 0.06 # preprocess_parameter 24 | 25 | batch_size = 6 # batch_size during training 26 | val_batch_size = 20 # batch_size during validation and test 27 | train_steps = 500 # Number of steps per epochs 28 | val_steps = 100 # Number of validation steps per epoch 29 | 30 | sub_sampling_ratio = [4, 4, 4, 4] # sampling ratio of random sampling at each layer 31 | d_out = [16, 64, 128, 256] # feature dimension 32 | num_sub_points = [num_points // 4, num_points // 16, num_points // 64, num_points // 256] 33 | 34 | noise_init = 3.5 # noise initial parameter 35 | max_epoch = 100 # maximum epoch during training 36 | learning_rate = 1e-2 # initial learning rate 37 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 38 | 39 | train_sum_dir = 'train_log' 40 | saving = True 41 | saving_path = None 42 | 43 | class ConfigS3DIS: 44 | k_n = 16 # KNN 45 | num_layers = 5 # Number of layers 46 | num_points = 40960 # Number of input points 47 | num_classes = 13 # Number of valid classes 48 | sub_grid_size = 0.04 # preprocess_parameter 49 | 50 | batch_size = 2 # batch_size during training 51 | val_batch_size = 20 # batch_size during validation and test 52 | train_steps = 500 # Number of steps per epochs 53 | val_steps = 100 # Number of validation steps per epoch 54 | 55 | sub_sampling_ratio = [4, 4, 4, 4, 2] # sampling ratio of random sampling at each layer 56 | d_out = [16, 64, 128, 256, 512] # feature dimension 57 | 58 | noise_init = 3.5 # noise initial parameter 59 | max_epoch = 100 # maximum epoch during training 60 | learning_rate = 1e-2 # initial learning rate 61 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 62 | 63 | train_sum_dir = 'train_log' 64 | saving = True 65 | saving_path = None 66 | 67 | class ConfigS3DIS_Benchmark: 68 | k_n = 16 # KNN 69 | num_layers = 5 # Number of layers 70 | num_points = 40960 # Number of input points 71 | num_classes = 13 # Number of valid classes 72 | sub_grid_size = 0.04 # preprocess_parameter 73 | 74 | batch_size = 2 # batch_size during training 75 | val_batch_size = 20 # batch_size during validation and test 76 | train_steps = 500 # Number of steps per epochs 77 | val_steps = 100 # Number of validation steps per epoch 78 | 79 | sub_sampling_ratio = [4, 4, 4, 4, 2] # sampling ratio of random sampling at each layer 80 | d_out = [16, 64, 128, 256, 512] # feature dimension 81 | 82 | noise_init = 3.5 # noise initial parameter 83 | max_epoch = 100 # maximum epoch during training 84 | learning_rate = 1e-2 # initial learning rate 85 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 86 | 87 | train_sum_dir = 'train_log_weak_benchmark' 88 | saving = True 89 | saving_path = None 90 | results_dir = 'results_weak_benchmark' 91 | 92 | # weakly semantic segmentation 93 | weak_label_ratio = 0.01 # 0.1, etc 94 | 95 | # Faithfully follow the Sqn paper 96 | class ConfigS3DIS_Sqn: 97 | k_n = 16 # KNN 98 | num_layers = 4 # Number of layers 99 | num_points = 40960 # Number of input points 100 | num_classes = 13 # Number of valid classes 101 | sub_grid_size = 0.04 # preprocess_parameter 102 | 103 | # NOTE: for a Nvidia 1080 GPU (8GB memory), use batch_size=3 costing about 7.9 GB GPU memory. 104 | batch_size = 3 # batch_size during training, change to a large number if owning sufficient GPU memory 105 | # val_batch_size = 20 # batch_size during validation and test for the RandLA-Net 106 | val_batch_size = 4 # change to a very sm number, 1 or 2, as the SQN need lg memory for three_nearest_interpolation() 107 | train_steps = 500 # Number of steps per epochs 108 | val_steps = 100 # Number of validation steps per epoch 109 | 110 | sub_sampling_ratio = [4, 4, 4, 4] # sampling ratio of random sampling at each layer 111 | d_out = [16, 64, 128, 256] # feature dimension 112 | 113 | noise_init = 3.5 # noise initial parameter 114 | max_epoch = 100 # maximum epoch during training 115 | learning_rate = 1e-2 # initial learning rate 116 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 117 | 118 | train_sum_dir = 'train_log_Sqn' 119 | saving = True 120 | saving_path = None 121 | results_dir = 'results_Sqn' 122 | 123 | # weakly semantic segmentation 124 | weak_label_ratio = 0.01 # 0.1, etc 125 | # the number of neighbors for interpolating feature for the query network 126 | k_interpolation = 3 127 | # how to concat point query features, default '1234' denoting concating 1-4th point query features, other options are: '123', etc. 128 | concat_type = '1234' 129 | 130 | 131 | class ConfigSemantic3D: 132 | k_n = 16 # KNN 133 | num_layers = 5 # Number of layers 134 | num_points = 65536 # Number of input points 135 | num_classes = 8 # Number of valid classes 136 | sub_grid_size = 0.06 # preprocess_parameter 137 | 138 | batch_size = 4 # batch_size during training 139 | val_batch_size = 16 # batch_size during validation and test 140 | train_steps = 500 # Number of steps per epochs 141 | val_steps = 100 # Number of validation steps per epoch 142 | 143 | sub_sampling_ratio = [4, 4, 4, 4, 2] # sampling ratio of random sampling at each layer 144 | d_out = [16, 64, 128, 256, 512] # feature dimension 145 | 146 | noise_init = 3.5 # noise initial parameter 147 | max_epoch = 100 # maximum epoch during training 148 | learning_rate = 1e-2 # initial learning rate 149 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 150 | 151 | train_sum_dir = 'train_log' 152 | saving = True 153 | saving_path = None 154 | 155 | augment_scale_anisotropic = True 156 | augment_symmetries = [True, False, False] 157 | augment_rotation = 'vertical' 158 | augment_scale_min = 0.8 159 | augment_scale_max = 1.2 160 | augment_noise = 0.001 161 | augment_occlusion = 'none' 162 | augment_color = 0.8 163 | 164 | 165 | class DataProcessing: 166 | @staticmethod 167 | def load_pc_semantic3d(filename): 168 | pc_pd = pd.read_csv(filename, header=None, delim_whitespace=True, dtype=np.float16) 169 | pc = pc_pd.values 170 | return pc 171 | 172 | @staticmethod 173 | def load_label_semantic3d(filename): 174 | label_pd = pd.read_csv(filename, header=None, delim_whitespace=True, dtype=np.uint8) 175 | cloud_labels = label_pd.values 176 | return cloud_labels 177 | 178 | @staticmethod 179 | def load_pc_kitti(pc_path): 180 | scan = np.fromfile(pc_path, dtype=np.float32) 181 | scan = scan.reshape((-1, 4)) 182 | points = scan[:, 0:3] # get xyz 183 | return points 184 | 185 | @staticmethod 186 | def load_label_kitti(label_path, remap_lut): 187 | label = np.fromfile(label_path, dtype=np.uint32) 188 | label = label.reshape((-1)) 189 | sem_label = label & 0xFFFF # semantic label in lower half 190 | inst_label = label >> 16 # instance id in upper half 191 | assert ((sem_label + (inst_label << 16) == label).all()) 192 | sem_label = remap_lut[sem_label] 193 | return sem_label.astype(np.int32) 194 | 195 | @staticmethod 196 | def get_file_list(dataset_path, test_scan_num): 197 | seq_list = np.sort(os.listdir(dataset_path)) 198 | 199 | train_file_list = [] 200 | test_file_list = [] 201 | val_file_list = [] 202 | for seq_id in seq_list: 203 | seq_path = join(dataset_path, seq_id) 204 | pc_path = join(seq_path, 'velodyne') 205 | if seq_id == '08': 206 | val_file_list.append([join(pc_path, f) for f in np.sort(os.listdir(pc_path))]) 207 | if seq_id == test_scan_num: 208 | test_file_list.append([join(pc_path, f) for f in np.sort(os.listdir(pc_path))]) 209 | elif int(seq_id) >= 11 and seq_id == test_scan_num: 210 | test_file_list.append([join(pc_path, f) for f in np.sort(os.listdir(pc_path))]) 211 | elif seq_id in ['00', '01', '02', '03', '04', '05', '06', '07', '09', '10']: 212 | train_file_list.append([join(pc_path, f) for f in np.sort(os.listdir(pc_path))]) 213 | 214 | train_file_list = np.concatenate(train_file_list, axis=0) 215 | val_file_list = np.concatenate(val_file_list, axis=0) 216 | test_file_list = np.concatenate(test_file_list, axis=0) 217 | return train_file_list, val_file_list, test_file_list 218 | 219 | @staticmethod 220 | def knn_search(support_pts, query_pts, k): 221 | """ 222 | :param support_pts: points you have, B*N1*3 223 | :param query_pts: points you want to know the neighbour index, B*N2*3 224 | :param k: Number of neighbours in knn search 225 | :return: neighbor_idx: neighboring points indexes, B*N2*k 226 | """ 227 | 228 | neighbor_idx = nearest_neighbors.knn_batch(support_pts, query_pts, k, omp=True) 229 | return neighbor_idx.astype(np.int32) 230 | 231 | @staticmethod 232 | def data_aug(xyz, color, labels, idx, num_out): 233 | """add padded points for points less than desired number-yc 234 | 235 | Args: 236 | xyz ([type]): [description] (N', 3) 237 | color ([type]): [description] (N', 3) 238 | labels ([type]): [description] (N',) 239 | idx ([type]): [description] (N',) 240 | num_out ([type]): [description] scalar, desired number of points 241 | 242 | Returns: 243 | [type]: [description] 244 | """ 245 | num_in = len(xyz) 246 | dup = np.random.choice(num_in, num_out - num_in) 247 | xyz_dup = xyz[dup, ...] 248 | xyz_aug = np.concatenate([xyz, xyz_dup], 0) # 249 | color_dup = color[dup, ...] 250 | color_aug = np.concatenate([color, color_dup], 0) 251 | idx_dup = list(range(num_in)) + list(dup) 252 | idx_aug = idx[idx_dup] 253 | label_aug = labels[idx_dup] 254 | return xyz_aug, color_aug, idx_aug, label_aug 255 | 256 | @staticmethod 257 | def data_aug_Sqn(xyz, color, labels, masks, idx, num_out): 258 | """add padded points for points less than desired number-yc 259 | 260 | Args: 261 | xyz ([type]): [description] (N', 3) 262 | color ([type]): [description] (N', 3) 263 | labels ([type]): [description] (N',) 264 | idx ([type]): [description] (N',) 265 | num_out ([type]): [description] scalar, desired number of points 266 | 267 | Returns: 268 | [type]: [description] 269 | """ 270 | num_in = len(xyz) 271 | dup = np.random.choice(num_in, num_out - num_in) 272 | xyz_dup = xyz[dup, ...] 273 | xyz_aug = np.concatenate([xyz, xyz_dup], 0) # 274 | color_dup = color[dup, ...] 275 | color_aug = np.concatenate([color, color_dup], 0) 276 | idx_dup = list(range(num_in)) + list(dup) 277 | idx_aug = idx[idx_dup] 278 | label_aug = labels[idx_dup] 279 | mask_aug = masks[idx_dup] 280 | return xyz_aug, color_aug, idx_aug, label_aug, mask_aug 281 | 282 | @staticmethod 283 | def shuffle_idx(x): 284 | # random shuffle the index 285 | idx = np.arange(len(x)) 286 | np.random.shuffle(idx) 287 | return x[idx] 288 | 289 | @staticmethod 290 | def shuffle_list(data_list): 291 | indices = np.arange(np.shape(data_list)[0]) 292 | np.random.shuffle(indices) 293 | data_list = data_list[indices] 294 | return data_list 295 | 296 | @staticmethod 297 | def grid_sub_sampling(points, features=None, labels=None, grid_size=0.1, verbose=0): 298 | """ 299 | CPP wrapper for a grid sub_sampling (method = barycenter for points and features 300 | :param points: (N, 3) matrix of input points 301 | :param features: optional (N, d) matrix of features (floating number) 302 | :param labels: optional (N,) matrix of integer labels 303 | :param grid_size: parameter defining the size of grid voxels 304 | :param verbose: 1 to display 305 | :return: sub_sampled points, with features and/or labels depending of the input 306 | """ 307 | 308 | if (features is None) and (labels is None): 309 | return cpp_subsampling.compute(points, sampleDl=grid_size, verbose=verbose) 310 | elif labels is None: 311 | return cpp_subsampling.compute(points, features=features, sampleDl=grid_size, verbose=verbose) 312 | elif features is None: 313 | return cpp_subsampling.compute(points, classes=labels, sampleDl=grid_size, verbose=verbose) 314 | else: 315 | return cpp_subsampling.compute(points, features=features, classes=labels, sampleDl=grid_size, 316 | verbose=verbose) 317 | 318 | @staticmethod 319 | def IoU_from_confusions(confusions): 320 | """ 321 | Computes IoU from confusion matrices. 322 | :param confusions: ([..., n_c, n_c] np.int32). Can be any dimension, the confusion matrices should be described by 323 | the last axes. n_c = number of classes 324 | :return: ([..., n_c] np.float32) IoU score 325 | """ 326 | 327 | # Compute TP, FP, FN. This assume that the second to last axis counts the truths (like the first axis of a 328 | # confusion matrix), and that the last axis counts the predictions (like the second axis of a confusion matrix) 329 | TP = np.diagonal(confusions, axis1=-2, axis2=-1) 330 | TP_plus_FN = np.sum(confusions, axis=-1) 331 | TP_plus_FP = np.sum(confusions, axis=-2) 332 | 333 | # Compute IoU 334 | IoU = TP / (TP_plus_FP + TP_plus_FN - TP + 1e-6) 335 | 336 | # Compute mIoU with only the actual classes 337 | mask = TP_plus_FN < 1e-3 338 | counts = np.sum(1 - mask, axis=-1, keepdims=True) 339 | mIoU = np.sum(IoU, axis=-1, keepdims=True) / (counts + 1e-6) 340 | 341 | # If class is absent, place mIoU in place of 0 IoU to get the actual mean later 342 | IoU += mask * mIoU 343 | return IoU 344 | 345 | @staticmethod 346 | def get_class_weights(dataset_name): 347 | # pre-calculate the number of points in each category 348 | num_per_class = [] 349 | 350 | # the number of points for each sub-sampling category-the RandLA-Net author's 351 | # if 'S3DIS' in dataset_name: 352 | # num_per_class = np.array([3370714, 2856755, 4919229, 318158, 375640, 478001, 974733, 353 | # 650464, 791496, 88727, 1284130, 229758, 2272837], dtype=np.int32) 354 | # the number of points for each sub-sampling category-yc 355 | if 'S3DIS_SQN' in dataset_name: 356 | num_per_class = np.array([3440776 , 2924550 , 4983774 , 313721 , 370052 , 465707 , 943685 , 660455 , 773345 , 87439 , 1266527 , 221644 , 2158967], dtype=np.int32) 357 | elif dataset_name is 'Semantic3D': 358 | num_per_class = np.array([5181602, 5012952, 6830086, 1311528, 10476365, 946982, 334860, 269353], 359 | dtype=np.int32) 360 | elif dataset_name is 'SemanticKITTI': 361 | num_per_class = np.array([55437630, 320797, 541736, 2578735, 3274484, 552662, 184064, 78858, 362 | 240942562, 17294618, 170599734, 6369672, 230413074, 101130274, 476491114, 363 | 9833174, 129609852, 4506626, 1168181]) 364 | weight = num_per_class / float(sum(num_per_class)) 365 | ce_label_weight = 1 / (weight + 0.02) 366 | return np.expand_dims(ce_label_weight, axis=0) 367 | 368 | 369 | class Plot: 370 | @staticmethod 371 | def random_colors(N, bright=True, seed=0): 372 | brightness = 1.0 if bright else 0.7 373 | hsv = [(0.15 + i / float(N), 1, brightness) for i in range(N)] 374 | colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv)) 375 | random.seed(seed) 376 | random.shuffle(colors) 377 | return colors 378 | 379 | @staticmethod 380 | def draw_pc(pc_xyzrgb): 381 | pc = open3d.PointCloud() 382 | pc.points = open3d.Vector3dVector(pc_xyzrgb[:, 0:3]) 383 | if pc_xyzrgb.shape[1] == 3: 384 | open3d.draw_geometries([pc]) 385 | return 0 386 | if np.max(pc_xyzrgb[:, 3:6]) > 20: ## 0-255 387 | pc.colors = open3d.Vector3dVector(pc_xyzrgb[:, 3:6] / 255.) 388 | else: 389 | pc.colors = open3d.Vector3dVector(pc_xyzrgb[:, 3:6]) 390 | open3d.draw_geometries([pc]) 391 | return 0 392 | 393 | @staticmethod 394 | def draw_pc_sem_ins(pc_xyz, pc_sem_ins, plot_colors=None): 395 | """ 396 | pc_xyz: 3D coordinates of point clouds 397 | pc_sem_ins: semantic or instance labels 398 | plot_colors: custom color list 399 | """ 400 | if plot_colors is not None: 401 | ins_colors = plot_colors 402 | else: 403 | ins_colors = Plot.random_colors(len(np.unique(pc_sem_ins)) + 1, seed=2) 404 | 405 | ############################## 406 | sem_ins_labels = np.unique(pc_sem_ins) 407 | sem_ins_bbox = [] 408 | Y_colors = np.zeros((pc_sem_ins.shape[0], 3)) 409 | for id, semins in enumerate(sem_ins_labels): 410 | valid_ind = np.argwhere(pc_sem_ins == semins)[:, 0] 411 | if semins <= -1: 412 | tp = [0, 0, 0] 413 | else: 414 | if plot_colors is not None: 415 | tp = ins_colors[semins] 416 | else: 417 | tp = ins_colors[id] 418 | 419 | Y_colors[valid_ind] = tp 420 | 421 | ### bbox 422 | valid_xyz = pc_xyz[valid_ind] 423 | 424 | xmin = np.min(valid_xyz[:, 0]); 425 | xmax = np.max(valid_xyz[:, 0]) 426 | ymin = np.min(valid_xyz[:, 1]); 427 | ymax = np.max(valid_xyz[:, 1]) 428 | zmin = np.min(valid_xyz[:, 2]); 429 | zmax = np.max(valid_xyz[:, 2]) 430 | sem_ins_bbox.append( 431 | [[xmin, ymin, zmin], [xmax, ymax, zmax], [min(tp[0], 1.), min(tp[1], 1.), min(tp[2], 1.)]]) 432 | 433 | Y_semins = np.concatenate([pc_xyz[:, 0:3], Y_colors], axis=-1) 434 | Plot.draw_pc(Y_semins) 435 | return Y_semins 436 | 437 | def log_out(out_str, f_out): 438 | f_out.write(out_str + '\n') 439 | f_out.flush() 440 | print(out_str) -------------------------------------------------------------------------------- /imgs/checkpoint_log.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PointCloudYC/SQN_tensorflow/1fde5085541f051eab7119202152388f95370ed2/imgs/checkpoint_log.jpg -------------------------------------------------------------------------------- /imgs/replication-roadmap.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PointCloudYC/SQN_tensorflow/1fde5085541f051eab7119202152388f95370ed2/imgs/replication-roadmap.jpg -------------------------------------------------------------------------------- /jobs_6_fold_cv_s3dis.sh: -------------------------------------------------------------------------------- 1 | python -B main_S3DIS_sqn.py --gpu 0 --mode train --test_area 1 2 | python -B main_S3DIS_sqn.py --gpu 0 --mode test --test_area 1 3 | python -B main_S3DIS_sqn.py --gpu 0 --mode train --test_area 2 4 | python -B main_S3DIS_sqn.py --gpu 0 --mode test --test_area 2 5 | python -B main_S3DIS_sqn.py --gpu 0 --mode train --test_area 3 6 | python -B main_S3DIS_sqn.py --gpu 0 --mode test --test_area 3 7 | python -B main_S3DIS_sqn.py --gpu 0 --mode train --test_area 4 8 | python -B main_S3DIS_sqn.py --gpu 0 --mode test --test_area 4 9 | python -B main_S3DIS_sqn.py --gpu 0 --mode train --test_area 5 10 | python -B main_S3DIS_sqn.py --gpu 0 --mode test --test_area 5 11 | python -B main_S3DIS_sqn.py --gpu 0 --mode train --test_area 6 12 | python -B main_S3DIS_sqn.py --gpu 0 --mode test --test_area 6 13 | 14 | 15 | -------------------------------------------------------------------------------- /main_S3DIS.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | from RandLANet import Network 3 | from tester_S3DIS import ModelTester 4 | from helper_ply import read_ply 5 | # S3DIS's configs 6 | from helper_tool import ConfigS3DIS as cfg 7 | from helper_tool import DataProcessing as DP 8 | from helper_tool import Plot 9 | import tensorflow as tf 10 | import numpy as np 11 | import time, pickle, argparse, glob, os 12 | 13 | 14 | class S3DIS: 15 | """S3DIS dataset class w/o inheriting any TensorFlow built-in classes 16 | Despite not inheriting any TensorFlow built-in classes, this class follow the tensorflow's dataset pattern: input pipeline w. iterator and initilizer 17 | - __int__(): initialize the dataset basic settings, e.g., dataset path, test_area_idx=5, classes, categories, and physical file paths, etc.Then it will call load_sub_sampled_clouds(). 18 | - load_sub_sampled_clouds(): load S3DIS dataset physical sub-sampled files as training and test clouds; (note: these sub-sub-sampled files are prepared by the data_prepare_s3dis.py) 19 | - init_input_pipeline(): create tensorflow built-in dataset object using the `from_generator` method, then create its iterator and train,val_init_op operator for session running. 20 | - get_batch_gen(): use for the above init_input_pipeline() to obtain batch data 21 | - get_tf_mapping2(): use for the above init_input_pipeline() to organize each stage's points,neighbors,pools and up_samples into a list. 22 | """ 23 | def __init__(self, test_area_idx): 24 | self.name = 'S3DIS' 25 | # self.path = '/data/S3DIS' 26 | self.path = 'data/S3DIS' 27 | self.label_to_names = {0: 'ceiling', 28 | 1: 'floor', 29 | 2: 'wall', 30 | 3: 'beam', 31 | 4: 'column', 32 | 5: 'window', 33 | 6: 'door', 34 | 7: 'table', 35 | 8: 'chair', 36 | 9: 'sofa', 37 | 10: 'bookcase', 38 | 11: 'board', 39 | 12: 'clutter'} 40 | self.num_classes = len(self.label_to_names) 41 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 42 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 43 | self.ignored_labels = np.array([]) 44 | 45 | self.val_split = 'Area_' + str(test_area_idx) 46 | self.all_files = glob.glob(join(self.path, 'original_ply', '*.ply')) # scan the folder for all ply files 47 | 48 | # Initiate containers 49 | # validation projection indice list, each item represents projection nnest id over a sub_pc for each corresponding raw pc pt-yc 50 | self.val_proj = [] 51 | # validation labels list, each item represent a validation sub_pc's label-yc 52 | self.val_labels = [] 53 | # possibility for control to randomly choose a point in the sub_pc evenly-yc 54 | self.possibility = {} 55 | self.min_possibility = {} 56 | # {training,validation} sub_pc's kd_trees, colors, labels and names-yc 57 | self.input_trees = {'training': [], 'validation': []} 58 | self.input_colors = {'training': [], 'validation': []} 59 | self.input_labels = {'training': [], 'validation': []} 60 | self.input_names = {'training': [], 'validation': []} 61 | # fill the above containers by reading physical sub_pc files-yc 62 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 63 | 64 | def load_sub_sampled_clouds(self, sub_grid_size): 65 | """load sub_sampled physical files and fill all the containers, input_{trees, colors, labels, names} and val_{proj,labels}-yc 66 | 67 | Args: 68 | sub_grid_size ([type]): sub-sampling grid size, e.g., 0.040 69 | """ 70 | tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size)) 71 | for i, file_path in enumerate(self.all_files): 72 | t0 = time.time() 73 | cloud_name = file_path.split('/')[-1][:-4] 74 | if self.val_split in cloud_name: 75 | cloud_split = 'validation' 76 | else: 77 | cloud_split = 'training' 78 | 79 | # Name of the input files 80 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) # e.g., Area_1_conferenceRoom_1_KDTree.pkl 81 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) # e.g., Area_1_conferenceRoom_1.ply 82 | 83 | data = read_ply(sub_ply_file) # ply format: x,y,z,red,gree,blue,class 84 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T # (N',3), note the transpose symbol 85 | sub_labels = data['class'] 86 | 87 | # Read pkl with search tree 88 | with open(kd_tree_file, 'rb') as f: 89 | search_tree = pickle.load(f) 90 | 91 | # input_xx is a dict contain training or validation info for all sub_pc, each of them contain a list 92 | self.input_trees[cloud_split] += [search_tree] 93 | self.input_colors[cloud_split] += [sub_colors] 94 | self.input_labels[cloud_split] += [sub_labels] 95 | self.input_names[cloud_split] += [cloud_name] 96 | 97 | size = sub_colors.shape[0] * 4 * 7 98 | print('{:s} {:.1f} MB loaded in {:.1f}s'.format(kd_tree_file.split('/')[-1], size * 1e-6, time.time() - t0)) 99 | 100 | print('\nPreparing reprojected indices for testing') 101 | 102 | # Get validation and test reprojected indices and labels (this is useful for validating on all raw points) 103 | for i, file_path in enumerate(self.all_files): 104 | t0 = time.time() 105 | cloud_name = file_path.split('/')[-1][:-4] 106 | 107 | # Validation projection and labels 108 | if self.val_split in cloud_name: 109 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 110 | with open(proj_file, 'rb') as f: 111 | proj_idx, labels = pickle.load(f) 112 | self.val_proj += [proj_idx] 113 | self.val_labels += [labels] 114 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 115 | 116 | # Generate the input data flow 117 | # Intuitively, it prepare data training examples, each pc will generate numerous point cloud training/validation examples by selecting a center point in the pc evenly, then select center point's neighboring points within a radius but not more than a threshold (e.g.,10000)-yc 118 | def get_batch_gen(self, split): 119 | if split == 'training': 120 | num_per_epoch = cfg.train_steps * cfg.batch_size 121 | elif split == 'validation': 122 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 123 | # assign a possibility for all sub_pc and their containing points-yc 124 | self.possibility[split] = [] 125 | self.min_possibility[split] = [] 126 | # Random initialize 127 | for i, tree in enumerate(self.input_colors[split]): 128 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 129 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 130 | 131 | def spatially_regular_gen(): 132 | # Generator loop 133 | for i in range(num_per_epoch): 134 | 135 | # Choose the cloud with the lowest probability 136 | cloud_idx = int(np.argmin(self.min_possibility[split])) 137 | 138 | # choose the point with the minimum of possibility in the cloud as query point 139 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 140 | 141 | # Get all points within the cloud from tree structure 142 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 143 | 144 | # Center point of input region 145 | center_point = points[point_ind, :].reshape(1, -1) 146 | 147 | # Add noise to the center point 148 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 149 | pick_point = center_point + noise.astype(center_point.dtype) 150 | 151 | # Check if the number of points in the selected cloud is less than the predefined num_points 152 | if len(points) < cfg.num_points: 153 | # Query all points within the cloud 154 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=len(points))[1][0] 155 | else: 156 | # Query the predefined number of points 157 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 158 | 159 | # Shuffle index 160 | queried_idx = DP.shuffle_idx(queried_idx) 161 | # Get corresponding points and colors based on the index 162 | queried_pc_xyz = points[queried_idx] 163 | queried_pc_xyz = queried_pc_xyz - pick_point 164 | queried_pc_colors = self.input_colors[split][cloud_idx][queried_idx] 165 | queried_pc_labels = self.input_labels[split][cloud_idx][queried_idx] 166 | 167 | # Update the possibility of the selected points 168 | dists = np.sum(np.square((points[queried_idx] - pick_point).astype(np.float32)), axis=1) 169 | delta = np.square(1 - dists / np.max(dists)) 170 | self.possibility[split][cloud_idx][queried_idx] += delta 171 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 172 | 173 | # up_sampled with replacement 174 | if len(points) < cfg.num_points: 175 | queried_pc_xyz, queried_pc_colors, queried_idx, queried_pc_labels = \ 176 | DP.data_aug(queried_pc_xyz, queried_pc_colors, queried_pc_labels, queried_idx, cfg.num_points) 177 | 178 | if True: 179 | yield (queried_pc_xyz.astype(np.float32), 180 | queried_pc_colors.astype(np.float32), 181 | queried_pc_labels, 182 | queried_idx.astype(np.int32), 183 | np.array([cloud_idx], dtype=np.int32)) 184 | 185 | gen_func = spatially_regular_gen 186 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32) 187 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None]) 188 | return gen_func, gen_types, gen_shapes 189 | 190 | @staticmethod 191 | def get_tf_mapping2(): 192 | """mapping for tranlating dataset's tensor to another form-yc 193 | The params of tf_map() just corresponds to {xyz,features,labels,idx,cloud_idx}; 194 | Considering there are cfg.num_layers(e.g., 4) stages in the encoder, each stage will have a sub-sampling process, use a list (named flat_inputs) for managing all of them (i.e., sub_sampled point cloud info at these stages). For example, if we have 4 sub-sampling processes, the flat_inputs list will have 20 items, like: [input_points, input_neighbors, input_pools, input_up_samples, batch_features, batch_labels, batch_pc_idx(i.e. points idx in the cloud),batch_cloud_idx], so 4*(cfg.num_layers+1) items in total 195 | Returns: 196 | [type]: [description] 197 | """ 198 | def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 199 | batch_features = tf.concat([batch_xyz, batch_features], axis=-1) 200 | input_points = [] 201 | input_neighbors = [] 202 | input_pools = [] 203 | input_up_samples = [] 204 | 205 | for i in range(cfg.num_layers): 206 | neighbour_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) # (B,N,k) 207 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] # retrieve first N/sub_sampling_ratio pts 208 | pool_i = neighbour_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] # sub_sampled points' id 209 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) # (B,N,K) over the sub_points 210 | input_points.append(batch_xyz) 211 | input_neighbors.append(neighbour_idx) 212 | input_pools.append(pool_i) 213 | input_up_samples.append(up_i) 214 | batch_xyz = sub_points 215 | 216 | input_list = input_points + input_neighbors + input_pools + input_up_samples 217 | input_list += [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx] 218 | 219 | return input_list # contains: [input_points, input_neighbors, input_pools, input_up_samples, batch_features, batch_labels, batch_pc_idx(i.e. points idx in the cloud),batch_cloud_idx], so 4*(cfg.num_layers+1) items in total 220 | 221 | return tf_map 222 | 223 | def init_input_pipeline(self): 224 | """ 225 | obtain X,Y pair and {train,val}_init_op operator following tensorflow pipline pattern. 226 | """ 227 | print('Initiating input pipelines') 228 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 229 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 230 | gen_function_val, _, _ = self.get_batch_gen('validation') 231 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) # create the dataset from a generator 232 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 233 | 234 | self.batch_train_data = self.train_data.batch(cfg.batch_size) # batch the dataset object 235 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 236 | map_func = self.get_tf_mapping2() 237 | 238 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) # map to another form, each batch is a list containing 4*(num_layers+1) items corresponding to points, features, labels, cloud_idx, point_idx at different sub-sampled stages for each batch 239 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 240 | 241 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 242 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 243 | 244 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 245 | self.flat_inputs = iter.get_next() # iterator, each returns a flat_inputs list containing 20 items 246 | # prepare operator for session to run 247 | self.train_init_op = iter.make_initializer(self.batch_train_data) 248 | self.val_init_op = iter.make_initializer(self.batch_val_data) 249 | 250 | 251 | if __name__ == '__main__': 252 | parser = argparse.ArgumentParser() 253 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 254 | parser.add_argument('--test_area', type=int, default=5, help='Which area to use for test, option: 1-6 [default: 5]') 255 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 256 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 257 | FLAGS = parser.parse_args() 258 | 259 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 260 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 261 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 262 | Mode = FLAGS.mode 263 | 264 | test_area = FLAGS.test_area 265 | 266 | # create S3DIS dataset object using test_area as validation/test set, the rest as training set-yc 267 | dataset = S3DIS(test_area) 268 | dataset.init_input_pipeline() 269 | 270 | """provide 3 functionality: training, testing and visualization-yc 271 | - training; pass the dataset object and dataset config to create the Network object, then start training 272 | - testing; pass in the model checkpoint and create ModelTest object, then start testing 273 | - visualization; plot the raw pc and sub_pc 274 | """ 275 | if Mode == 'train': 276 | # NOTE: cfg is S3DIS object w. common configs, a global variable here. 277 | model = Network(dataset, cfg) 278 | model.train(dataset) 279 | elif Mode == 'test': 280 | cfg.saving = False 281 | model = Network(dataset, cfg) 282 | if FLAGS.model_path is not 'None': 283 | chosen_snap = FLAGS.model_path 284 | else: 285 | chosen_snapshot = -1 286 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 287 | chosen_folder = logs[-1] 288 | snap_path = join(chosen_folder, 'snapshots') 289 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 290 | chosen_step = np.sort(snap_steps)[-1] 291 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 292 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 293 | tester.test(model, dataset) 294 | else: 295 | ################## 296 | # Visualize data # 297 | ################## 298 | 299 | with tf.Session() as sess: 300 | sess.run(tf.global_variables_initializer()) 301 | sess.run(dataset.train_init_op) 302 | while True: 303 | flat_inputs = sess.run(dataset.flat_inputs) 304 | pc_xyz = flat_inputs[0] 305 | sub_pc_xyz = flat_inputs[1] 306 | labels = flat_inputs[21] 307 | Plot.draw_pc_sem_ins(pc_xyz[0, :, :], labels[0, :]) 308 | Plot.draw_pc_sem_ins(sub_pc_xyz[0, :, :], labels[0, 0:np.shape(sub_pc_xyz)[1]]) 309 | -------------------------------------------------------------------------------- /main_S3DIS_weak_benchmark.py: -------------------------------------------------------------------------------- 1 | """ 2 | main entry for benchmarking weak supervision for RandLA-Net, reproduced based on the SQN paper, check https://arxiv.org/abs/2104.04891 3 | Essentially, this weak supervision is almost similar to the official RandLA-Net except that a point cloud example does not have full labels, which are partially labeled. 4 | - config same as RandLA-Net's config (but need add weak_label_ratio attribute) 5 | - dataset preparation, add weak_labels for raw and sub_pc and same as preparing SQN dataset, check dataset_prepare_s3dis_sqn.py and main_S3DIS_Sqn.py 6 | - model same as the RandLA-Net 7 | - Loss: only compute losses of weakly points 8 | - train and evaluation are the same 9 | Author: Chao YIN 10 | Email: cyinac@connect.ust.hk 11 | Date: Oct. 23, 2021 12 | """ 13 | 14 | import numpy as np 15 | import time, pickle, argparse, glob, os, random 16 | from os.path import join 17 | # from main_S3DIS_Sqn import S3DIS_SQN 18 | from RandLANet import Network 19 | from tester_S3DIS import ModelTester 20 | from helper_ply import read_ply 21 | 22 | # S3DIS for Sqn config 23 | from helper_tool import ConfigS3DIS_Benchmark as cfg 24 | from helper_tool import DataProcessing as DP 25 | from helper_tool import Plot 26 | import tensorflow as tf 27 | # tf.enable_eager_execution() 28 | 29 | class S3DIS_Benchmark: 30 | """S3DIS dataset class w/o inheriting any TensorFlow built-in classes 31 | Despite not inheriting any TensorFlow built-in classes, this class follow the tensorflow's dataset pattern: input pipeline w. iterator and initilizer 32 | - __int__(): initialize the dataset basic settings, e.g., dataset path, test_area_idx=5, classes, categories, and physical file paths, etc.Then it will call load_sub_sampled_clouds(). 33 | - load_sub_sampled_clouds(): load S3DIS dataset physical sub-sampled files as training and test clouds; (note: these sub-sub-sampled files are prepared by the data_prepare_s3dis.py) 34 | - init_input_pipeline(): create tensorflow built-in dataset object using the `from_generator` method, then create its iterator and train,val_init_op operator for session running. 35 | - get_batch_gen(): use for the above init_input_pipeline() to obtain batch data 36 | - get_tf_mapping2(): use for the above init_input_pipeline() to organize each stage's points,neighbors,pools and up_samples into a list. 37 | """ 38 | def __init__(self, test_area_idx): 39 | self.name = 'S3DIS_Benchmark' 40 | # self.path = '/data/S3DIS' 41 | self.path = 'data/S3DIS' 42 | self.label_to_names = {0: 'ceiling', 43 | 1: 'floor', 44 | 2: 'wall', 45 | 3: 'beam', 46 | 4: 'column', 47 | 5: 'window', 48 | 6: 'door', 49 | 7: 'table', 50 | 8: 'chair', 51 | 9: 'sofa', 52 | 10: 'bookcase', 53 | 11: 'board', 54 | 12: 'clutter'} 55 | self.num_classes = len(self.label_to_names) 56 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 57 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 58 | self.ignored_labels = np.array([]) 59 | 60 | self.val_split = 'Area_' + str(test_area_idx) 61 | self.all_files = glob.glob(join(self.path, 'original_ply', '*.ply')) # scan the folder for all ply files 62 | 63 | # Initiate containers 64 | # validation projection indice list, each item represents projection nnest id over a sub_pc for each corresponding raw pc pt-yc 65 | self.val_proj = [] 66 | # validation labels list, each item represent a validation sub_pc's label-yc 67 | self.val_labels = [] 68 | # possibility for control to randomly choose a point in the sub_pc evenly-yc 69 | self.possibility = {} 70 | self.min_possibility = {} 71 | # {training,validation} sub_pc's kd_trees, colors, labels and names-yc 72 | self.input_trees = {'training': [], 'validation': []} 73 | self.input_colors = {'training': [], 'validation': []} 74 | self.input_labels = {'training': [], 'validation': []} 75 | self.input_weak_labels = {'training': [], 'validation': []} 76 | self.input_names = {'training': [], 'validation': []} 77 | # fill the above containers by reading physical sub_pc files-yc 78 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 79 | 80 | def load_sub_sampled_clouds(self, sub_grid_size): 81 | """load sub_sampled physical files and fill all the containers, input_{trees, colors, labels, names} and val_{proj,labels}-yc 82 | 83 | Args: 84 | sub_grid_size ([type]): sub-sampling grid size, e.g., 0.040 85 | """ 86 | tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size)) 87 | for i, file_path in enumerate(self.all_files): 88 | t0 = time.time() 89 | cloud_name = file_path.split('/')[-1][:-4] 90 | if self.val_split in cloud_name: 91 | cloud_split = 'validation' 92 | else: 93 | cloud_split = 'training' 94 | 95 | # Name of the input files 96 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) # e.g., Area_1_conferenceRoom_1_KDTree.pkl 97 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) # e.g., Area_1_conferenceRoom_1.ply 98 | 99 | data = read_ply(sub_ply_file) # ply format: x,y,z,red,gree,blue,class 100 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T # (N',3), note the transpose symbol 101 | sub_labels = data['class'] 102 | 103 | # read weak labels for sub_pc 104 | weak_label_folder = join(self.path, 'weak_label_{}'.format(self.weak_label_ratio)) 105 | weak_label_sub_file = join(weak_label_folder, file_path.split('/')[-1][:-4] + '_sub_weak_label.ply') 106 | if os.path.exists(weak_label_sub_file): 107 | weak_data = read_ply(weak_label_sub_file) # ply format: x,y,z,red,gree,blue,class 108 | weak_label_sub_mask = weak_data['weak_mask'] # (N',) to align same shape as sub_labels 109 | else: 110 | raise NotImplementedError("run the dataset_prepare_s3dis_sqn.py to generate weak labels for raw and sub PC") 111 | 112 | # Read pkl with search tree 113 | with open(kd_tree_file, 'rb') as f: 114 | search_tree = pickle.load(f) 115 | 116 | # input_xx is a dict contain training or validation info for all sub_pc, each of them contain a list 117 | self.input_trees[cloud_split] += [search_tree] 118 | self.input_colors[cloud_split] += [sub_colors] 119 | self.input_labels[cloud_split] += [sub_labels] 120 | # HACK: for validation set, all points should have labels meaning the weak_label_ratio is 1(i.e., all points have labels) 121 | if cloud_split == 'validation': 122 | self.input_weak_labels[cloud_split] += [np.ones_like(weak_label_sub_mask)] 123 | else: 124 | self.input_weak_labels[cloud_split] += [weak_label_sub_mask] 125 | self.input_names[cloud_split] += [cloud_name] 126 | 127 | size = sub_colors.shape[0] * 4 * 7 128 | print('{:s} {:.1f} MB loaded in {:.1f}s'.format(kd_tree_file.split('/')[-1], size * 1e-6, time.time() - t0)) 129 | 130 | print('\nPreparing reprojected indices for testing') 131 | 132 | # Get validation and test reprojected indices and labels (this is useful for validating on all raw points) 133 | for i, file_path in enumerate(self.all_files): 134 | t0 = time.time() 135 | cloud_name = file_path.split('/')[-1][:-4] 136 | 137 | # Validation projection and labels 138 | if self.val_split in cloud_name: 139 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 140 | with open(proj_file, 'rb') as f: 141 | proj_idx, labels = pickle.load(f) 142 | self.val_proj += [proj_idx] 143 | self.val_labels += [labels] 144 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 145 | 146 | # Generate the input data flow 147 | # Intuitively, it prepare data training examples, each pc will generate numerous point cloud training/validation examples by selecting a center point in the pc evenly, then select center point's neighboring points within a radius but not more than a threshold (e.g.,10000)-yc 148 | def get_batch_gen(self, split): 149 | if split == 'training': 150 | num_per_epoch = cfg.train_steps * cfg.batch_size 151 | elif split == 'validation': 152 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 153 | # assign a possibility for all sub_pc and their containing points-yc 154 | self.possibility[split] = [] 155 | self.min_possibility[split] = [] 156 | # Random initialize 157 | for i, tree in enumerate(self.input_colors[split]): 158 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 159 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 160 | 161 | def spatially_regular_gen(): 162 | # Generator loop 163 | for i in range(num_per_epoch): 164 | 165 | # Choose the cloud with the lowest probability 166 | cloud_idx = int(np.argmin(self.min_possibility[split])) 167 | 168 | # choose the point with the minimum of possibility in the cloud as query point 169 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 170 | 171 | # Get all points within the cloud from tree structure 172 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 173 | 174 | # Center point of input region 175 | center_point = points[point_ind, :].reshape(1, -1) 176 | 177 | # Add noise to the center point 178 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 179 | pick_point = center_point + noise.astype(center_point.dtype) 180 | 181 | # Check if the number of points in the selected cloud is less than the predefined num_points 182 | if len(points) < cfg.num_points: 183 | # Query all points within the cloud 184 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=len(points))[1][0] 185 | else: 186 | # Query the predefined number of points 187 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 188 | 189 | # Shuffle index 190 | queried_idx = DP.shuffle_idx(queried_idx) 191 | # Get corresponding points and colors based on the index 192 | queried_pc_xyz = points[queried_idx] 193 | queried_pc_xyz = queried_pc_xyz - pick_point 194 | queried_pc_colors = self.input_colors[split][cloud_idx][queried_idx] 195 | queried_pc_labels = self.input_labels[split][cloud_idx][queried_idx] 196 | queried_pc_weak_label_mask = self.input_weak_labels[split][cloud_idx][queried_idx] 197 | 198 | # Update the possibility of the selected points 199 | dists = np.sum(np.square((points[queried_idx] - pick_point).astype(np.float32)), axis=1) 200 | delta = np.square(1 - dists / np.max(dists)) 201 | self.possibility[split][cloud_idx][queried_idx] += delta 202 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 203 | 204 | # up_sampled with replacement 205 | if len(points) < cfg.num_points: 206 | queried_pc_xyz, queried_pc_colors, queried_idx, queried_pc_labels, queried_pc_weak_label_mask = \ 207 | DP.data_aug_Sqn(queried_pc_xyz, queried_pc_colors, queried_pc_labels, 208 | queried_pc_weak_label_mask, queried_idx, cfg.num_points) 209 | if True: 210 | yield (queried_pc_xyz.astype(np.float32), 211 | queried_pc_colors.astype(np.float32), 212 | queried_pc_labels, 213 | queried_pc_weak_label_mask, 214 | queried_idx.astype(np.int32), 215 | np.array([cloud_idx], dtype=np.int32)) 216 | 217 | gen_func = spatially_regular_gen 218 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32, tf.int32) 219 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None], [None]) 220 | return gen_func, gen_types, gen_shapes 221 | 222 | @staticmethod 223 | def get_tf_mapping2(): 224 | """mapping for tranlating dataset's tensor to another form-yc 225 | The params of tf_map() just corresponds to {xyz,features,labels,idx,cloud_idx}; 226 | Considering there are cfg.num_layers(e.g., 4) stages in the encoder, each stage will have a sub-sampling process, use a list (named flat_inputs) for managing all of them (i.e., sub_sampled point cloud info at these stages). For example, if we have 4 sub-sampling processes, the flat_inputs list will have 20 items, like: [input_points, input_neighbors, input_pools, input_up_samples, batch_features, batch_labels, batch_pc_idx(i.e. points idx in the cloud),batch_cloud_idx], so 4*(cfg.num_layers+1) items in total 227 | Returns: 228 | [type]: [description] 229 | """ 230 | def tf_map(batch_xyz, batch_features, batch_labels,batch_weak_label_mask, batch_pc_idx, batch_cloud_idx): 231 | batch_features = tf.concat([batch_xyz, batch_features], axis=-1) 232 | input_points = [] 233 | input_neighbors = [] 234 | input_pools = [] 235 | input_up_samples = [] 236 | 237 | for i in range(cfg.num_layers): 238 | neighbour_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) # (B,N,k) 239 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] # retrieve first N/sub_sampling_ratio pts 240 | pool_i = neighbour_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] # sub_sampled points' id 241 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) # (B,N,K) over the sub_points 242 | input_points.append(batch_xyz) 243 | input_neighbors.append(neighbour_idx) 244 | input_pools.append(pool_i) 245 | input_up_samples.append(up_i) 246 | batch_xyz = sub_points 247 | 248 | input_list = input_points + input_neighbors + input_pools + input_up_samples 249 | input_list += [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx, batch_weak_label_mask] 250 | 251 | return input_list # contains: [input_points, input_neighbors, input_pools, input_up_samples, batch_features, batch_labels, batch_pc_idx(i.e. points idx in the cloud),batch_cloud_idx], so 4*(cfg.num_layers+1) items in total 252 | 253 | return tf_map 254 | 255 | def init_input_pipeline(self): 256 | """ 257 | obtain X,Y pair and {train,val}_init_op operator following tensorflow pipline pattern. 258 | """ 259 | print('Initiating input pipelines') 260 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 261 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 262 | gen_function_val, _, _ = self.get_batch_gen('validation') 263 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) # create the dataset from a generator 264 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 265 | 266 | self.batch_train_data = self.train_data.batch(cfg.batch_size) # batch the dataset object 267 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 268 | map_func = self.get_tf_mapping2() 269 | 270 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) # map to another form, each batch is a list containing 4*(num_layers+1) items corresponding to points, features, labels, cloud_idx, point_idx at different sub-sampled stages for each batch 271 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 272 | 273 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 274 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 275 | 276 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 277 | self.flat_inputs = iter.get_next() # iterator, each returns a flat_inputs list containing 20 items 278 | # prepare operator for session to run 279 | self.train_init_op = iter.make_initializer(self.batch_train_data) 280 | self.val_init_op = iter.make_initializer(self.batch_val_data) 281 | 282 | 283 | if __name__ == '__main__': 284 | parser = argparse.ArgumentParser() 285 | parser.add_argument("--rng_seed", type=int, default=123, help='manual seed') 286 | parser.add_argument("--num_points", type=int, default=40960, help='the number of points for each PC example') 287 | parser.add_argument("--batch_size", type=int, default=4, help='batch size for training') 288 | parser.add_argument("--val_batch_size", type=int, default=1, help='batch size for validation') 289 | parser.add_argument("--max_epoch", type=int, default=400, help='max epoch for training') 290 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 291 | parser.add_argument('--test_area', type=int, default=5, help='Which area to use for test, option: 1-6 [default: 5]') 292 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 293 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 294 | parser.add_argument('--sub_grid_size', type=float, default=0.04, help='grid-sampling size') 295 | parser.add_argument('--weak_label_ratio', type=float, default=0.01, help='the weakly semantic segmentation ratio') 296 | FLAGS = parser.parse_args() 297 | 298 | # set fixed seeds for reproducible results 299 | random.seed(FLAGS.rng_seed) 300 | np.random.seed(FLAGS.rng_seed) 301 | # tf.random.set_seed(FLAGS.rng_seed) 302 | tf.random.set_random_seed(FLAGS.rng_seed) 303 | 304 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 305 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 306 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 307 | Mode = FLAGS.mode 308 | test_area = FLAGS.test_area 309 | 310 | # override the config with argparse's arguments 311 | cfg.num_points=FLAGS.num_points 312 | cfg.max_epoch=FLAGS.max_epoch 313 | cfg.batch_size=FLAGS.batch_size 314 | cfg.val_batch_size=FLAGS.val_batch_size 315 | cfg.sub_grid_size=FLAGS.sub_grid_size 316 | cfg.weak_label_ratio=FLAGS.weak_label_ratio 317 | 318 | # create S3DIS dataset object for weakly semseg using test_area as validation/test set, the rest as training set-yc 319 | dataset = S3DIS_Benchmark(test_area, cfg) 320 | dataset.init_input_pipeline() 321 | 322 | """provide 3 functionality: training, testing and visualization-yc 323 | - training; pass the dataset object and dataset config to create the Network object, then start training 324 | - testing; pass in the model checkpoint and create ModelTest object, then start testing 325 | - visualization; plot the raw pc and sub_pc 326 | """ 327 | if Mode == 'train': 328 | # NOTE: cfg is S3DIS object w. common configs, a global variable here. 329 | model = Network(dataset, cfg) 330 | model.train(dataset) 331 | elif Mode == 'test': 332 | cfg.saving = False 333 | model = Network(dataset, cfg) 334 | if FLAGS.model_path is not 'None': 335 | chosen_snap = FLAGS.model_path 336 | else: 337 | chosen_snapshot = -1 338 | logs = np.sort([os.path.join(cfg.results_dir, f) for f in os.listdir(cfg.results_dir) if f.startswith('Log')]) 339 | chosen_folder = logs[-1] 340 | snap_path = join(chosen_folder, 'snapshots') 341 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 342 | chosen_step = np.sort(snap_steps)[-1] 343 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 344 | # TODO: 345 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 346 | tester.test(model, dataset) 347 | else: 348 | ################## 349 | # Visualize data # 350 | ################## 351 | 352 | with tf.Session() as sess: 353 | sess.run(tf.global_variables_initializer()) 354 | # use session to start the dataset iterator 355 | sess.run(dataset.train_init_op) 356 | while True: 357 | 358 | # obtain the iterator's next element 359 | flat_inputs = sess.run(dataset.flat_inputs) 360 | pc_xyz = flat_inputs[0] 361 | sub_pc_xyz = flat_inputs[1] 362 | labels = flat_inputs[4*cfg.num_layers +1] 363 | Plot.draw_pc_sem_ins(pc_xyz[0, :, :], labels[0, :]) 364 | Plot.draw_pc_sem_ins(sub_pc_xyz[0, :, :], labels[0, 0:np.shape(sub_pc_xyz)[1]]) -------------------------------------------------------------------------------- /prepare-s3dis-sqn.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # SUB_GRID_SIZES=(0.02) 4 | SUB_GRID_SIZES=(0.04) 5 | # WEAK_LABEL_RATIOS=(0.00001 0.00002 0.00005) 6 | # WEAK_LABEL_RATIOS=(0.00001 0.00002 0.00005) 7 | # WEAK_LABEL_RATIOS=(0.005) 8 | # WEAK_LABEL_RATIOS=(0.0006) 9 | WEAK_LABEL_RATIOS=(0.001 0.0001 0.0005 0.0002) 10 | # WEAK_LABEL_RATIOS=(0.0005 0.0002 0.0001) 11 | # WEAK_LABEL_RATIOS=(0.1 0.01 0.0001) 12 | # WEAK_LABEL_RATIOS=(0.1) 13 | 14 | for sub_grid_size in "${SUB_GRID_SIZES[@]}"; do 15 | for weak_label_ratio in "${WEAK_LABEL_RATIOS[@]}"; do 16 | 17 | echo "sub_grid_size is ${sub_grid_size}" 18 | echo "weak_label_ratio is ${weak_label_ratio}" 19 | 20 | time python utils/data_prepare_s3dis_sqn.py --sub_grid_size ${sub_grid_size} \ 21 | --weak_label_ratio ${weak_label_ratio} 22 | 23 | echo "sub_grid_size is ${sub_grid_size}" 24 | echo "weak_label_ratio is ${weak_label_ratio}" 25 | done 26 | done 27 | 28 | echo "finish preparing the S3DIS dataset for weakly semantic segmentation!" -------------------------------------------------------------------------------- /run-s3dis-Sqn.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | MODES=('train') 4 | # MODES=('train' 'test') 5 | gpu=0 6 | # 4 costs about 8GB GPU, others can be used on grp server GPUs 7 | BATCH_SIZES=(3) # 3 just 8GB, 4--> OOM on 1080 GPU 8 | # VAL_BATCH_SIZES=(20) # 16 10 8 6 9 | # 8 works on my nvidia 1080 gpu, others can be tried on grp server GPUs 10 | VAL_BATCH_SIZES=(4) # 6-->system collapse 11 | num_points=40960 12 | max_epoch=100 # 400 13 | 14 | # KEY: weak label ratio is defined as the number of weak points over the raw poinits 15 | # WEAK_LABEL_RATIOS=(0.001) 16 | # WEAK_LABEL_RATIOS=(0.1 0.01 0.001 0.0001) 17 | # WEAK_LABEL_RATIOS=(0.001 0.0001) 18 | # WEAK_LABEL_RATIOS=(0.00001 0.00005) # 0.01 can not on 1080gpu, suffer OOM 19 | # WEAK_LABEL_RATIOS=(0.00001 0.00002 0.00005) 20 | WEAK_LABEL_RATIOS=(0.0001 0.00005 0.00002) 21 | 22 | # 0.005 batch_size=2, OOM 23 | # 0.001 batch_size 3, val_batch_size 4 24 | # WEAK_LABEL_RATIOS=(0.001 0.01 0.0001) 25 | # new weak_label_ratio refers to weak points/#raw_pc, has a 15 times relation 26 | # TODO: normalize the ratios 27 | # WEAK_LABEL_RATIOS=(0.015 0.15 0.0015) # corresponds to 0.001/0.01/0.0001 28 | 29 | # TODO: ablation study 30 | # num_k_query_pts 31 | # how to concat features 32 | # CONCAT_TYPES=('1234' '123' '234' '12' '1') 33 | CONCAT_TYPES=('1234') 34 | 35 | 36 | echo "training using WCE loss" 37 | for mode in "${MODES[@]}"; do 38 | for weak_label_ratio in "${WEAK_LABEL_RATIOS[@]}"; do 39 | for batch_size in "${BATCH_SIZES[@]}"; do 40 | for val_batch_size in "${VAL_BATCH_SIZES[@]}"; do 41 | for concat_type in "${CONCAT_TYPES[@]}"; do 42 | echo "batch_size: ${batch_size}" 43 | echo "val_batch_size: ${val_batch_size}" 44 | echo "num_points: ${num_points}" 45 | echo "max_epoch: ${max_epoch}" 46 | echo "weak_label_ratio: ${weak_label_ratio}" 47 | echo "concat_type: ${concat_type}" 48 | 49 | time python -B main_S3DIS_Sqn.py \ 50 | --gpu ${gpu} \ 51 | --mode ${mode} \ 52 | --test_area 5 \ 53 | --batch_size ${batch_size} \ 54 | --val_batch_size ${val_batch_size} \ 55 | --num_points ${num_points} \ 56 | --max_epoch ${max_epoch} \ 57 | --weak_label_ratio ${weak_label_ratio} \ 58 | --concat_type ${concat_type} 59 | 60 | echo "batch_size: ${batch_size}" 61 | echo "val_batch_size: ${val_batch_size}" 62 | echo "num_points: ${num_points}" 63 | echo "max_epoch: ${max_epoch}" 64 | echo "weak_label_ratio: ${weak_label_ratio}" 65 | echo "concat_type: ${concat_type}" 66 | done 67 | done 68 | done 69 | done 70 | done 71 | 72 | echo "training using WCE loss" 73 | echo "finish training." -------------------------------------------------------------------------------- /tester_S3DIS.py: -------------------------------------------------------------------------------- 1 | from os import makedirs 2 | from os.path import exists, join 3 | from helper_ply import write_ply 4 | from sklearn.metrics import confusion_matrix 5 | from helper_tool import DataProcessing as DP 6 | import tensorflow as tf 7 | import numpy as np 8 | import time 9 | 10 | 11 | def log_out(out_str, log_f_out): 12 | log_f_out.write(out_str + '\n') 13 | log_f_out.flush() 14 | print(out_str) 15 | 16 | 17 | class ModelTester: 18 | def __init__(self, model, dataset, restore_snap=None): 19 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 20 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 21 | self.Log_file = open('log_test_' + str(dataset.val_split) + '.txt', 'a') 22 | 23 | # Create a session for running Ops on the Graph. 24 | on_cpu = False 25 | if on_cpu: 26 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 27 | else: 28 | c_proto = tf.ConfigProto() 29 | c_proto.gpu_options.allow_growth = True 30 | self.sess = tf.Session(config=c_proto) 31 | self.sess.run(tf.global_variables_initializer()) 32 | 33 | # Load trained model 34 | if restore_snap is not None: 35 | self.saver.restore(self.sess, restore_snap) 36 | print("Model restored from " + restore_snap) 37 | 38 | self.prob_logits = tf.nn.softmax(model.logits) 39 | 40 | # Initiate global prediction over all test clouds 41 | self.test_probs = [np.zeros(shape=[l.shape[0], model.config.num_classes], dtype=np.float32) 42 | for l in dataset.input_labels['validation']] 43 | 44 | def test(self, model, dataset, num_votes=100): 45 | 46 | # Smoothing parameter for votes 47 | test_smooth = 0.95 48 | 49 | # Initialise iterator with validation/test data 50 | self.sess.run(dataset.val_init_op) 51 | 52 | # Number of points per class in validation set 53 | val_proportions = np.zeros(model.config.num_classes, dtype=np.float32) 54 | i = 0 55 | for label_val in dataset.label_values: 56 | if label_val not in dataset.ignored_labels: 57 | val_proportions[i] = np.sum([np.sum(labels == label_val) for labels in dataset.val_labels]) 58 | i += 1 59 | 60 | # Test saving path 61 | saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 62 | test_path = join('test', saving_path.split('/')[-1]) 63 | makedirs(test_path) if not exists(test_path) else None 64 | makedirs(join(test_path, 'val_preds')) if not exists(join(test_path, 'val_preds')) else None 65 | 66 | step_id = 0 67 | epoch_id = 0 68 | last_min = -0.5 69 | 70 | while last_min < num_votes: 71 | try: 72 | ops = (self.prob_logits, 73 | model.labels, 74 | model.inputs['input_inds'], 75 | model.inputs['cloud_inds'], 76 | ) 77 | 78 | stacked_probs, stacked_labels, point_idx, cloud_idx = self.sess.run(ops, {model.is_training: False}) 79 | correct = np.sum(np.argmax(stacked_probs, axis=1) == stacked_labels) 80 | acc = correct / float(np.prod(np.shape(stacked_labels))) 81 | print('step' + str(step_id) + ' acc:' + str(acc)) 82 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, model.config.num_points, 83 | model.config.num_classes]) 84 | 85 | for j in range(np.shape(stacked_probs)[0]): 86 | probs = stacked_probs[j, :, :] 87 | p_idx = point_idx[j, :] 88 | c_i = cloud_idx[j][0] 89 | self.test_probs[c_i][p_idx] = test_smooth * self.test_probs[c_i][p_idx] + (1 - test_smooth) * probs 90 | step_id += 1 91 | 92 | except tf.errors.OutOfRangeError: 93 | 94 | new_min = np.min(dataset.min_possibility['validation']) 95 | log_out('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_id, new_min), self.Log_file) 96 | 97 | if last_min + 1 < new_min: 98 | 99 | # Update last_min 100 | last_min += 1 101 | 102 | # Show vote results (On subcloud so it is not the good values here) 103 | log_out('\nConfusion on sub clouds', self.Log_file) 104 | confusion_list = [] 105 | 106 | num_val = len(dataset.input_labels['validation']) 107 | 108 | for i_test in range(num_val): 109 | probs = self.test_probs[i_test] 110 | preds = dataset.label_values[np.argmax(probs, axis=1)].astype(np.int32) 111 | labels = dataset.input_labels['validation'][i_test] 112 | 113 | # Confs 114 | confusion_list += [confusion_matrix(labels, preds, dataset.label_values)] 115 | 116 | # Regroup confusions 117 | C = np.sum(np.stack(confusion_list), axis=0).astype(np.float32) 118 | 119 | # Rescale with the right number of point per class 120 | C *= np.expand_dims(val_proportions / (np.sum(C, axis=1) + 1e-6), 1) 121 | 122 | # Compute IoUs 123 | IoUs = DP.IoU_from_confusions(C) 124 | m_IoU = np.mean(IoUs) 125 | s = '{:5.2f} | '.format(100 * m_IoU) 126 | for IoU in IoUs: 127 | s += '{:5.2f} '.format(100 * IoU) 128 | log_out(s + '\n', self.Log_file) 129 | 130 | if int(np.ceil(new_min)) % 1 == 0: 131 | 132 | # Project predictions 133 | log_out('\nReproject Vote #{:d}'.format(int(np.floor(new_min))), self.Log_file) 134 | proj_probs_list = [] 135 | 136 | for i_val in range(num_val): 137 | # Reproject probs back to the evaluations points 138 | proj_idx = dataset.val_proj[i_val] 139 | probs = self.test_probs[i_val][proj_idx, :] 140 | proj_probs_list += [probs] 141 | 142 | # Show vote results 143 | log_out('Confusion on full clouds', self.Log_file) 144 | confusion_list = [] 145 | for i_test in range(num_val): 146 | # Get the predicted labels 147 | preds = dataset.label_values[np.argmax(proj_probs_list[i_test], axis=1)].astype(np.uint8) 148 | 149 | # Confusion 150 | labels = dataset.val_labels[i_test] 151 | acc = np.sum(preds == labels) / len(labels) 152 | log_out(dataset.input_names['validation'][i_test] + ' Acc:' + str(acc), self.Log_file) 153 | 154 | confusion_list += [confusion_matrix(labels, preds, dataset.label_values)] 155 | name = dataset.input_names['validation'][i_test] + '.ply' 156 | write_ply(join(test_path, 'val_preds', name), [preds, labels], ['pred', 'label']) 157 | 158 | # Regroup confusions 159 | C = np.sum(np.stack(confusion_list), axis=0) 160 | 161 | IoUs = DP.IoU_from_confusions(C) 162 | m_IoU = np.mean(IoUs) 163 | s = '{:5.2f} | '.format(100 * m_IoU) 164 | for IoU in IoUs: 165 | s += '{:5.2f} '.format(100 * IoU) 166 | log_out('-' * len(s), self.Log_file) 167 | log_out(s, self.Log_file) 168 | log_out('-' * len(s) + '\n', self.Log_file) 169 | print('finished \n') 170 | self.sess.close() 171 | return 172 | 173 | self.sess.run(dataset.val_init_op) 174 | epoch_id += 1 175 | step_id = 0 176 | continue 177 | 178 | return 179 | -------------------------------------------------------------------------------- /tester_S3DIS_Sqn.py: -------------------------------------------------------------------------------- 1 | from os import makedirs 2 | from os.path import exists, join 3 | from helper_ply import write_ply 4 | from sklearn.metrics import confusion_matrix 5 | from helper_tool import DataProcessing as DP 6 | import tensorflow as tf 7 | import numpy as np 8 | import time 9 | 10 | 11 | def log_out(out_str, log_f_out): 12 | log_f_out.write(out_str + '\n') 13 | log_f_out.flush() 14 | print(out_str) 15 | 16 | 17 | class ModelTester: 18 | def __init__(self, model, dataset, restore_snap=None): 19 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 20 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 21 | self.Log_file = open('log_test_' + str(dataset.val_split) + '.txt', 'a') 22 | 23 | # Create a session for running Ops on the Graph. 24 | on_cpu = False 25 | if on_cpu: 26 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 27 | else: 28 | c_proto = tf.ConfigProto() 29 | c_proto.gpu_options.allow_growth = True 30 | self.sess = tf.Session(config=c_proto) 31 | self.sess.run(tf.global_variables_initializer()) 32 | 33 | # Load trained model 34 | if restore_snap is not None: 35 | self.saver.restore(self.sess, restore_snap) 36 | print("Model restored from " + restore_snap) 37 | 38 | self.prob_logits = tf.nn.softmax(model.logits) 39 | 40 | # Initiate global prediction over all test clouds 41 | self.test_probs = [np.zeros(shape=[l.shape[0], model.config.num_classes], dtype=np.float32) 42 | for l in dataset.input_labels['validation']] 43 | 44 | def test(self, model, dataset, num_votes=100): 45 | 46 | # Smoothing parameter for votes 47 | test_smooth = 0.95 48 | 49 | # Initialise iterator with validation/test data 50 | self.sess.run(dataset.val_init_op) 51 | 52 | # Number of points per class in validation set 53 | val_proportions = np.zeros(model.config.num_classes, dtype=np.float32) 54 | i = 0 55 | for label_val in dataset.label_values: 56 | if label_val not in dataset.ignored_labels: 57 | val_proportions[i] = np.sum([np.sum(labels == label_val) for labels in dataset.val_labels]) 58 | i += 1 59 | 60 | # Test saving path 61 | # saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 62 | saving_path = time.strftime('{}/Log_weak_{}_%Y-%m-%d_%H-%M-%S-{}'.format(model.config.results_dir,dataset.weak_label_ratio), time.gmtime()) 63 | test_path = join('test', saving_path.split('/')[-1]) 64 | makedirs(test_path) if not exists(test_path) else None 65 | makedirs(join(test_path, 'val_preds')) if not exists(join(test_path, 'val_preds')) else None 66 | 67 | step_id = 0 68 | epoch_id = 0 69 | last_min = -0.5 70 | 71 | while last_min < num_votes: 72 | try: 73 | ops = (self.prob_logits, 74 | model.weak_labels, 75 | model.inputs['input_inds'], 76 | model.inputs['cloud_inds'], 77 | ) 78 | 79 | stacked_probs, stacked_labels, point_idx, cloud_idx = self.sess.run(ops, {model.is_training: False}) 80 | correct = np.sum(np.argmax(stacked_probs, axis=1) == stacked_labels) 81 | acc = correct / float(np.prod(np.shape(stacked_labels))) 82 | print('step' + str(step_id) + ' acc:' + str(acc)) 83 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, model.config.num_points, 84 | model.config.num_classes]) 85 | 86 | for j in range(np.shape(stacked_probs)[0]): 87 | probs = stacked_probs[j, :, :] 88 | p_idx = point_idx[j, :] 89 | c_i = cloud_idx[j][0] 90 | self.test_probs[c_i][p_idx] = test_smooth * self.test_probs[c_i][p_idx] + (1 - test_smooth) * probs 91 | step_id += 1 92 | 93 | except tf.errors.OutOfRangeError: 94 | 95 | new_min = np.min(dataset.min_possibility['validation']) 96 | log_out('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_id, new_min), self.Log_file) 97 | 98 | if last_min + 1 < new_min: 99 | 100 | # Update last_min 101 | last_min += 1 102 | 103 | # Show vote results (On subcloud so it is not the good values here) 104 | log_out('\nConfusion on sub clouds', self.Log_file) 105 | confusion_list = [] 106 | 107 | num_val = len(dataset.input_labels['validation']) 108 | 109 | for i_test in range(num_val): 110 | probs = self.test_probs[i_test] 111 | preds = dataset.label_values[np.argmax(probs, axis=1)].astype(np.int32) 112 | labels = dataset.input_labels['validation'][i_test] 113 | 114 | # Confs 115 | confusion_list += [confusion_matrix(labels, preds, dataset.label_values)] 116 | 117 | # Regroup confusions 118 | C = np.sum(np.stack(confusion_list), axis=0).astype(np.float32) 119 | 120 | # Rescale with the right number of point per class 121 | C *= np.expand_dims(val_proportions / (np.sum(C, axis=1) + 1e-6), 1) 122 | 123 | # Compute IoUs 124 | IoUs = DP.IoU_from_confusions(C) 125 | m_IoU = np.mean(IoUs) 126 | s = '{:5.2f} | '.format(100 * m_IoU) 127 | for IoU in IoUs: 128 | s += '{:5.2f} '.format(100 * IoU) 129 | log_out(s + '\n', self.Log_file) 130 | 131 | if int(np.ceil(new_min)) % 1 == 0: 132 | 133 | # Project predictions 134 | log_out('\nReproject Vote #{:d}'.format(int(np.floor(new_min))), self.Log_file) 135 | proj_probs_list = [] 136 | 137 | for i_val in range(num_val): 138 | # Reproject probs back to the evaluations points 139 | proj_idx = dataset.val_proj[i_val] 140 | probs = self.test_probs[i_val][proj_idx, :] 141 | proj_probs_list += [probs] 142 | 143 | # Show vote results 144 | log_out('Confusion on full clouds', self.Log_file) 145 | confusion_list = [] 146 | for i_test in range(num_val): 147 | # Get the predicted labels 148 | preds = dataset.label_values[np.argmax(proj_probs_list[i_test], axis=1)].astype(np.uint8) 149 | 150 | # Confusion 151 | labels = dataset.val_labels[i_test] 152 | acc = np.sum(preds == labels) / len(labels) 153 | log_out(dataset.input_names['validation'][i_test] + ' Acc:' + str(acc), self.Log_file) 154 | 155 | confusion_list += [confusion_matrix(labels, preds, dataset.label_values)] 156 | name = dataset.input_names['validation'][i_test] + '.ply' 157 | write_ply(join(test_path, 'val_preds', name), [preds, labels], ['pred', 'label']) 158 | 159 | # Regroup confusions 160 | C = np.sum(np.stack(confusion_list), axis=0) 161 | 162 | IoUs = DP.IoU_from_confusions(C) 163 | m_IoU = np.mean(IoUs) 164 | s = '{:5.2f} | '.format(100 * m_IoU) 165 | for IoU in IoUs: 166 | s += '{:5.2f} '.format(100 * IoU) 167 | log_out('-' * len(s), self.Log_file) 168 | log_out(s, self.Log_file) 169 | log_out('-' * len(s) + '\n', self.Log_file) 170 | print('finished \n') 171 | self.sess.close() 172 | return 173 | 174 | self.sess.run(dataset.val_init_op) 175 | epoch_id += 1 176 | step_id = 0 177 | continue 178 | 179 | return 180 | -------------------------------------------------------------------------------- /tf_ops/3d_interpolation/interpolate.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // Find three nearest neigbors with square distance 19 | // input: xyz1 (b,n,3), xyz2(b,m,3) 20 | // output: dist (b,n,3), idx (b,n,3) 21 | void threenn_cpu(int b, int n, int m, const float *xyz1, const float *xyz2, float *dist, int *idx) { 22 | for (int i=0;i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include "tensorflow/core/framework/op.h" 7 | #include "tensorflow/core/framework/op_kernel.h" 8 | #include "tensorflow/core/framework/shape_inference.h" 9 | #include "tensorflow/core/framework/common_shape_fns.h" 10 | using namespace tensorflow; 11 | 12 | REGISTER_OP("ThreeNN") 13 | .Input("xyz1: float32") 14 | .Input("xyz2: float32") 15 | .Output("dist: float32") 16 | .Output("idx: int32") 17 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 18 | c->set_output(0, c->input(0)); 19 | c->set_output(1, c->input(0)); 20 | return Status::OK(); 21 | }); 22 | REGISTER_OP("ThreeInterpolate") 23 | .Input("points: float32") 24 | .Input("idx: int32") 25 | .Input("weight: float32") 26 | .Output("out: float32") 27 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 28 | ::tensorflow::shape_inference::ShapeHandle dims1; // (b,m,c) 29 | c->WithRank(c->input(0), 3, &dims1); 30 | ::tensorflow::shape_inference::ShapeHandle dims2; // (b,n,3) 31 | c->WithRank(c->input(1), 3, &dims2); 32 | // (b,n,c) 33 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), c->Dim(dims2, 1), c->Dim(dims1, 2)}); 34 | c->set_output(0, output); 35 | return Status::OK(); 36 | }); 37 | REGISTER_OP("ThreeInterpolateGrad") 38 | .Input("points: float32") 39 | .Input("idx: int32") 40 | .Input("weight: float32") 41 | .Input("grad_out: float32") 42 | .Output("grad_points: float32") 43 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 44 | c->set_output(0, c->input(0)); 45 | return Status::OK(); 46 | }); 47 | 48 | float randomf(){ 49 | return (rand()+0.5)/(RAND_MAX+1.0); 50 | } 51 | static double get_time(){ 52 | timespec tp; 53 | clock_gettime(CLOCK_MONOTONIC,&tp); 54 | return tp.tv_sec+tp.tv_nsec*1e-9; 55 | } 56 | 57 | // Find three nearest neigbors with square distance 58 | // input: xyz1 (b,n,3), xyz2(b,m,3) 59 | // output: dist (b,n,3), idx (b,n,3) 60 | void threenn_cpu(int b, int n, int m, const float *xyz1, const float *xyz2, float *dist, int *idx) { 61 | for (int i=0;iinput(0); 163 | OP_REQUIRES(context, xyz1_tensor.dims()==3 && xyz1_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeNN expects (b,n,3) xyz1 shape.")); 164 | int b = xyz1_tensor.shape().dim_size(0); 165 | int n = xyz1_tensor.shape().dim_size(1); 166 | 167 | const Tensor& xyz2_tensor = context->input(1); 168 | OP_REQUIRES(context, xyz2_tensor.dims()==3 && xyz2_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeNN expects (b,m,3) xyz2 shape.")); 169 | int m = xyz2_tensor.shape().dim_size(1); 170 | 171 | Tensor *dist_tensor = nullptr; 172 | OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape{b,n,3}, &dist_tensor)); 173 | Tensor *idx_tensor = nullptr; 174 | OP_REQUIRES_OK(context, context->allocate_output(1, TensorShape{b,n,3}, &idx_tensor)); 175 | 176 | auto xyz1_flat = xyz1_tensor.flat(); 177 | const float *xyz1 = &(xyz1_flat(0)); 178 | auto xyz2_flat = xyz2_tensor.flat(); 179 | const float *xyz2 = &(xyz2_flat(0)); 180 | auto dist_flat = dist_tensor->flat(); 181 | float *dist = &(dist_flat(0)); 182 | auto idx_flat = idx_tensor->flat(); 183 | int *idx = &(idx_flat(0)); 184 | threenn_cpu(b,n,m,xyz1,xyz2,dist,idx); 185 | } 186 | }; 187 | REGISTER_KERNEL_BUILDER(Name("ThreeNN").Device(DEVICE_CPU), ThreeNNOp); 188 | 189 | 190 | 191 | class ThreeInterpolateOp: public OpKernel{ 192 | public: 193 | explicit ThreeInterpolateOp(OpKernelConstruction * context):OpKernel(context){} 194 | 195 | void Compute(OpKernelContext * context) override { 196 | const Tensor& points_tensor=context->input(0); 197 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("ThreeInterpolate expects (b,m,c) points shape")); 198 | int b = points_tensor.shape().dim_size(0); 199 | int m = points_tensor.shape().dim_size(1); 200 | int c = points_tensor.shape().dim_size(2); 201 | 202 | const Tensor& idx_tensor=context->input(1); 203 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b && idx_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeInterpolate expects (b,n,3) idx shape")); 204 | int n = idx_tensor.shape().dim_size(1); 205 | const Tensor& weight_tensor=context->input(2); 206 | OP_REQUIRES(context,weight_tensor.dims()==3 && weight_tensor.shape().dim_size(0)==b && weight_tensor.shape().dim_size(1)==n && weight_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeInterpolate expects (b,n,3) weight shape")); 207 | 208 | Tensor * out_tensor = nullptr; 209 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,n,c}, &out_tensor)); 210 | 211 | auto points_flat = points_tensor.flat(); 212 | const float *points = &(points_flat(0)); 213 | auto idx_flat = idx_tensor.flat(); 214 | const int *idx = &(idx_flat(0)); 215 | auto weight_flat = weight_tensor.flat(); 216 | const float *weight = &(weight_flat(0)); 217 | auto out_flat = out_tensor->flat(); 218 | float *out = &(out_flat(0)); 219 | threeinterpolate_cpu(b,m,c,n,points,idx,weight,out); 220 | } 221 | }; 222 | REGISTER_KERNEL_BUILDER(Name("ThreeInterpolate").Device(DEVICE_CPU),ThreeInterpolateOp); 223 | 224 | 225 | class ThreeInterpolateGradOp: public OpKernel{ 226 | public: 227 | explicit ThreeInterpolateGradOp(OpKernelConstruction * context):OpKernel(context){} 228 | 229 | void Compute(OpKernelContext * context) override { 230 | const Tensor& points_tensor=context->input(0); 231 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("ThreeInterpolateGrad expects (b,m,c) points shape")); 232 | int b = points_tensor.shape().dim_size(0); 233 | int m = points_tensor.shape().dim_size(1); 234 | int c = points_tensor.shape().dim_size(2); 235 | 236 | const Tensor& idx_tensor=context->input(1); 237 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b, errors::InvalidArgument("ThreeInterpolateGrad expects (b,n,3) idx shape")); 238 | int n = idx_tensor.shape().dim_size(1); 239 | const Tensor& weight_tensor=context->input(2); 240 | OP_REQUIRES(context,weight_tensor.dims()==3 && weight_tensor.shape().dim_size(0)==b && weight_tensor.shape().dim_size(1)==n && weight_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeInterpolateGrad expects (b,n,3) weight shape")); 241 | 242 | const Tensor& grad_out_tensor=context->input(3); 243 | OP_REQUIRES(context,grad_out_tensor.dims()==3 && grad_out_tensor.shape().dim_size(0)==b && grad_out_tensor.shape().dim_size(1)==n && grad_out_tensor.shape().dim_size(2)==c, errors::InvalidArgument("ThreeInterpolateGrad expects (b,n,c) grad_out shape")); 244 | 245 | Tensor * grad_points_tensor = nullptr; 246 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,m,c}, &grad_points_tensor)); 247 | 248 | auto points_flat = points_tensor.flat(); 249 | const float *points = &(points_flat(0)); 250 | auto idx_flat = idx_tensor.flat(); 251 | const int *idx = &(idx_flat(0)); 252 | auto weight_flat = weight_tensor.flat(); 253 | const float *weight = &(weight_flat(0)); 254 | auto grad_out_flat = grad_out_tensor.flat(); 255 | const float *grad_out = &(grad_out_flat(0)); 256 | auto grad_points_flat = grad_points_tensor->flat(); 257 | float *grad_points = &(grad_points_flat(0)); 258 | memset(grad_points, 0, sizeof(float)*b*m*c); 259 | threeinterpolate_grad_cpu(b,n,c,m,grad_out,idx,weight,grad_points); 260 | } 261 | }; 262 | REGISTER_KERNEL_BUILDER(Name("ThreeInterpolateGrad").Device(DEVICE_CPU),ThreeInterpolateGradOp); 263 | 264 | 265 | -------------------------------------------------------------------------------- /tf_ops/3d_interpolation/tf_interpolate.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.python.framework import ops 3 | import sys 4 | import os 5 | BASE_DIR = os.path.dirname(__file__) 6 | sys.path.append(BASE_DIR) 7 | 8 | # load custom tf interpolate lib 9 | try: 10 | if os.path.exists(os.path.join(BASE_DIR, 'tf_interpolate_so.so')): 11 | interpolate_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_interpolate_so.so')) 12 | else: 13 | raise NotImplementedError("Your TensorFlow has suffered a problem") 14 | except: 15 | print(f'Attention! your tensorflow version should below 1.14! yours is {tf.__version__}\n') 16 | print('can not load tf_interpolate_so custom ops correctly! Check your tensorflow version!') 17 | 18 | def three_nn(xyz_query, xyz_support): 19 | '''find xyz_query's nearest 3 neighbors of xyz_support 20 | Input: 21 | xyz_query: (b,n,3) float32 array, unknown/query points 22 | xyz_support: (b,m,3) float32 array, known/support points 23 | Output: 24 | dist: (b,n,3) float32 array, distances to known points 25 | idx: (b,n,3) int32 array, indices to known points 26 | ''' 27 | return interpolate_module.three_nn(xyz_query, xyz_support) 28 | 29 | ops.NoGradient('ThreeNN') 30 | 31 | def three_interpolate(features_support, query_idx_over_support, weight): 32 | '''interpolate features for the xyz_query(determined by idx) 33 | Input: 34 | features_support: (b,m,c) float32 array, known/support features of the corresponding xyz_support 35 | query_idx_over_support: (b,n,3) int32 array, indices of nearest 3 neighbors in the known/support points for each query point 36 | weight: (b,n,3) float32 array, weights for query_idx_over_support 37 | 38 | Output: 39 | out: (b,n,c) float32 array, interpolated point features 40 | ''' 41 | return interpolate_module.three_interpolate(features_support, query_idx_over_support, weight) 42 | 43 | @tf.RegisterGradient('ThreeInterpolate') 44 | def _three_interpolate_grad(op, grad_out): 45 | points = op.inputs[0] 46 | idx = op.inputs[1] 47 | weight = op.inputs[2] 48 | return [interpolate_module.three_interpolate_grad(points, idx, weight, grad_out), None, None] 49 | 50 | if __name__=='__main__': 51 | 52 | import numpy as np 53 | import time 54 | np.random.seed(100) 55 | features = np.random.random((32,128,64)).astype('float32') 56 | xyz1 = np.random.random((32,512,3)).astype('float32') 57 | xyz2 = np.random.random((32,128,3)).astype('float32') 58 | 59 | with tf.device('/cpu:0'): 60 | points = tf.constant(features) 61 | xyz1 = tf.constant(xyz1) 62 | xyz2 = tf.constant(xyz2) 63 | dist, idx = three_nn(xyz1, xyz2) 64 | weight = tf.ones_like(dist)/3.0 65 | interpolated_points = three_interpolate(points, idx, weight) 66 | 67 | with tf.Session('') as sess: 68 | now = time.time() 69 | for _ in range(100): 70 | ret = sess.run(interpolated_points) 71 | print(time.time() - now) 72 | print(ret.shape, ret.dtype) -------------------------------------------------------------------------------- /tf_ops/3d_interpolation/tf_interpolate_compile.sh: -------------------------------------------------------------------------------- 1 | #/bin/bash 2 | CUDA_ROOT="/usr/local/cuda-10.1" 3 | TF_ROOT="/home/yinchao/miniconda3/envs/DL/lib/python3.6/site-packages/tensorflow" 4 | 5 | # TF1.4 (Note: -L ${TF_ROOT} should have a space in between) 6 | 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 7 | -------------------------------------------------------------------------------- /tf_ops/3d_interpolation/tf_interpolate_op_test.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | from tf_interpolate import three_nn, three_interpolate 4 | 5 | class GroupPointTest(tf.test.TestCase): 6 | def test(self): 7 | pass 8 | 9 | def test_grad(self): 10 | with self.test_session(): 11 | features = tf.constant(np.random.random((1,8,16)).astype('float32')) # features, (1,8,16) 12 | print(features) 13 | xyz1 = tf.constant(np.random.random((1,128,3)).astype('float32')) 14 | xyz2 = tf.constant(np.random.random((1,8,3)).astype('float32')) 15 | dist, idx = three_nn(xyz1, xyz2) # (1,128,3), (1,128,3) 16 | weight = tf.ones_like(dist)/3.0 # (1,128,3) 17 | interpolated_features = three_interpolate(features, idx, weight) # (1,128,16) 18 | print(interpolated_features) 19 | err = tf.test.compute_gradient_error(features, (1,8,16), interpolated_features, (1,128,16)) 20 | print(err) 21 | self.assertLess(err, 1e-4) 22 | 23 | if __name__=='__main__': 24 | tf.test.main() 25 | -------------------------------------------------------------------------------- /tf_ops/3d_interpolation/visu_interpolation.py: -------------------------------------------------------------------------------- 1 | ''' Visualize part segmentation ''' 2 | import os 3 | import sys 4 | ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 5 | sys.path.append('/home/rqi/Projects/toolkits/visualization') 6 | from show3d_balls import showpoints 7 | import numpy as np 8 | from tf_interpolate import three_nn, three_interpolate 9 | import tensorflow as tf 10 | 11 | 12 | pts2 = np.array([[0,0,1],[1,0,0],[0,1,0],[1,1,0]]).astype('float32') 13 | xyz1 = np.random.random((100,3)).astype('float32') 14 | xyz2 = np.array([[0,0,0],[1,0,0],[0,1,0],[1,1,1]]).astype('float32') 15 | 16 | def fun(xyz1,xyz2,pts2): 17 | with tf.device('/cpu:0'): 18 | points = tf.constant(np.expand_dims(pts2,0)) 19 | xyz1 = tf.constant(np.expand_dims(xyz1,0)) 20 | xyz2 = tf.constant(np.expand_dims(xyz2,0)) 21 | dist, idx = three_nn(xyz1, xyz2) 22 | #weight = tf.ones_like(dist)/3.0 23 | dist = tf.maximum(dist, 1e-10) 24 | norm = tf.reduce_sum((1.0/dist),axis=2,keep_dims=True) 25 | norm = tf.tile(norm, [1,1,3]) 26 | print(norm) 27 | weight = (1.0/dist) / norm 28 | interpolated_points = three_interpolate(points, idx, weight) 29 | with tf.Session('') as sess: 30 | tmp,pts1,d,w = sess.run([xyz1, interpolated_points, dist, weight]) 31 | #print w 32 | pts1 = pts1.squeeze() 33 | return pts1 34 | 35 | pts1 = fun(xyz1,xyz2,pts2) 36 | all_pts = np.zeros((104,3)) 37 | all_pts[0:100,:] = pts1 38 | all_pts[100:,:] = pts2 39 | all_xyz = np.zeros((104,3)) 40 | all_xyz[0:100,:]=xyz1 41 | all_xyz[100:,:]=xyz2 42 | showpoints(xyz2, pts2, ballradius=8) 43 | showpoints(xyz1, pts1, ballradius=8) 44 | showpoints(all_xyz, all_pts, ballradius=8) 45 | -------------------------------------------------------------------------------- /tf_ops/readme.md: -------------------------------------------------------------------------------- 1 | # understanding the custom CUDA ops 2 | 3 | 4 | # problems 5 | 6 | ## shell script reports $'\r': command not found 7 | 8 | - reason: these shell scripts are copied from windows, need remove those '\r\n' characters(windows OS) for each line in the file 9 | - solution: 10 | - use dos2unix; `dos2unix [file]` 11 | - manual; remove them or create a new file, copy these liens w/o '\r\n' 12 | 13 | refs: 14 | * [shell脚本执行错误 $'\r':command not found_liuxiangke0210的专栏-CSDN博客](https://blog.csdn.net/liuxiangke0210/article/details/80395707) 15 | * [How do I fix "$'\r': command not found" errors running Bash scripts in WSL? - Ask Ubuntu](https://askubuntu.com/questions/966488/how-do-i-fix-r-command-not-found-errors-running-bash-scripts-in-wsl) 16 | 17 | ## When compiling PointNet++ compile CUDA op, it report undefined symbol: _ZTIN10tensorflow8OpKernelE 18 | 19 | - solution: need comment `#-D_GLIBCXX_USE_CXX11_ABI=0` 20 | 21 | ``` 22 | CUDA_ROOT="/usr/local/cuda-10.1" 23 | TF_ROOT="/home/yinchao/miniconda3/envs/DL/lib/python3.6/site-packages/tensorflow" 24 | 25 | # TF1.13 26 | 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 27 | ``` 28 | 29 | refs: 30 | * [Undefined symbol: · Issue #48 · charlesq34/pointnet2](https://github.com/charlesq34/pointnet2/issues/48) 31 | 32 | 33 | ## When compiling PointNet++ compile CUDA op, it report /usr/bin/ld: cannot find -lcudart 34 | 35 | - solution: symlink libcudart.so, e.g., `sudo ln -s /usr/local/cuda/lib64/libcudart.so /usr/lib/libcudart.so` 36 | 37 | refs: 38 | * [compiling - "/usr/bin/ld: cannot find -lcudart" - Ask Ubuntu](https://askubuntu.com/questions/510176/usr-bin-ld-cannot-find-lcudart) 39 | 40 | 41 | ## PointNet++ some python files is python 2.x, e.g., print is version 2 42 | 43 | - solution: convert python 2.x to python 3.x source code by 2to3 script 44 | 45 | ``` 46 | 2to3 -2 file.py 47 | ``` 48 | 49 | refs: 50 | * [2to3 - Automated Python 2 to 3 code translation — Python 3.10.0 documentation](https://docs.python.org/3/library/2to3.html) 51 | 52 | ## always report /usr/bin/ld: cannot find -lcudart when compiling tf_ops/3d_interpolation 53 | 54 | - reason: some minor characters 55 | - solution: I remove all the files under 3d_interpolation folder, then re-download the files from the github repo, and then re-edit the file. 56 | 57 | ``` 58 | #/bin/bash 59 | CUDA_ROOT="/usr/local/cuda-10.1" 60 | TF_ROOT="/home/yinchao/miniconda3/envs/DL/lib/python3.6/site-packages/tensorflow" 61 | 62 | # TF1.4 (Note: -L ${TF_ROOT} should have a space in between) 63 | 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 64 | ``` -------------------------------------------------------------------------------- /tf_ops/test.py: -------------------------------------------------------------------------------- 1 | import os 2 | from os import makedirs 3 | from os.path import exists, join 4 | import sys 5 | import numpy as np 6 | import tensorflow as tf 7 | from sklearn.metrics import confusion_matrix 8 | 9 | 10 | # custom tf ops based on PointNet++(https://github.com/charlesq34/pointnet2) 11 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 12 | sys.path.append(os.path.join(BASE_DIR, '../tf_ops/3d_interpolation')) 13 | # sys.path.append(os.path.join(BASE_DIR, '../tf_ops/sampling')) 14 | # from tf_sampling import farthest_point_sample, gather_point 15 | from tf_interpolate import three_nn, three_interpolate 16 | 17 | # print(gather_point) 18 | 19 | print(three_nn) 20 | print(three_interpolate) -------------------------------------------------------------------------------- /utils/6_fold_cv.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import glob, os, sys 3 | 4 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 5 | ROOT_DIR = os.path.dirname(BASE_DIR) 6 | sys.path.append(ROOT_DIR) 7 | from helper_ply import read_ply 8 | from helper_tool import Plot 9 | 10 | if __name__ == '__main__': 11 | base_dir = '/data/S3DIS/results' 12 | original_data_dir = '/data/S3DIS/original_ply' 13 | data_path = glob.glob(os.path.join(base_dir, '*.ply')) 14 | data_path = np.sort(data_path) 15 | 16 | test_total_correct = 0 17 | test_total_seen = 0 18 | gt_classes = [0 for _ in range(13)] 19 | positive_classes = [0 for _ in range(13)] 20 | true_positive_classes = [0 for _ in range(13)] 21 | visualization = False 22 | 23 | for file_name in data_path: 24 | pred_data = read_ply(file_name) 25 | pred = pred_data['pred'] 26 | original_data = read_ply(os.path.join(original_data_dir, file_name.split('/')[-1][:-4] + '.ply')) 27 | labels = original_data['class'] 28 | points = np.vstack((original_data['x'], original_data['y'], original_data['z'])).T 29 | 30 | ################## 31 | # Visualize data # 32 | ################## 33 | if visualization: 34 | colors = np.vstack((original_data['red'], original_data['green'], original_data['blue'])).T 35 | xyzrgb = np.concatenate([points, colors], axis=-1) 36 | Plot.draw_pc(xyzrgb) # visualize raw point clouds 37 | Plot.draw_pc_sem_ins(points, labels) # visualize ground-truth 38 | Plot.draw_pc_sem_ins(points, pred) # visualize prediction 39 | 40 | correct = np.sum(pred == labels) 41 | print(str(file_name.split('/')[-1][:-4]) + '_acc:' + str(correct / float(len(labels)))) 42 | test_total_correct += correct 43 | test_total_seen += len(labels) 44 | 45 | for j in range(len(labels)): 46 | gt_l = int(labels[j]) 47 | pred_l = int(pred[j]) 48 | gt_classes[gt_l] += 1 49 | positive_classes[pred_l] += 1 50 | true_positive_classes[gt_l] += int(gt_l == pred_l) 51 | 52 | iou_list = [] 53 | for n in range(13): 54 | iou = true_positive_classes[n] / float(gt_classes[n] + positive_classes[n] - true_positive_classes[n]) 55 | iou_list.append(iou) 56 | mean_iou = sum(iou_list) / 13.0 57 | print('eval accuracy: {}'.format(test_total_correct / float(test_total_seen))) 58 | print('mean IOU:{}'.format(mean_iou)) 59 | print(iou_list) 60 | 61 | acc_list = [] 62 | for n in range(13): 63 | acc = true_positive_classes[n] / float(gt_classes[n]) 64 | acc_list.append(acc) 65 | mean_acc = sum(acc_list) / 13.0 66 | print('mAcc value is :{}'.format(mean_acc)) 67 | -------------------------------------------------------------------------------- /utils/compute_weak_point_distribution.py: -------------------------------------------------------------------------------- 1 | """ 2 | compute S3DIS's point distribution for raw, sub-sampled and weak settings, ending w. generating a xlxs file 3 | Author: Chao YIN 4 | Email: cyinac@connect.ust.hk 5 | 6 | base: data_prepare_s3dis_sqn.py 7 | history: 8 | - Nov. 1, 2021, init the file 9 | """ 10 | 11 | import os, sys, glob, pickle, argparse, random 12 | from os.path import join, exists, dirname, abspath 13 | import numpy as np 14 | import xlsxwriter # write info to xlsx file 15 | 16 | BASE_DIR = dirname(abspath(__file__)) 17 | ROOT_DIR = dirname(BASE_DIR) 18 | sys.path.append(BASE_DIR) 19 | sys.path.append(ROOT_DIR) 20 | from helper_ply import write_ply, read_ply 21 | from helper_tool import DataProcessing as DP 22 | 23 | 24 | def compute_num_points(anno_path, save_path, sub_pc_folder, 25 | weak_label_folder, weak_label_ratio, sub_grid_size, 26 | gt_class, gt_class2label): 27 | 28 | num_raw_points = 0 # number of raw points for current room 29 | num_sub_points = 0 # number of sub-sampled points for current room 30 | 31 | # save raw_cloud 32 | if not os.path.exists(save_path): 33 | raise NotImplementedError("run the dataset_prepare_s3dis_sqn.py to generate raw and sub PC") 34 | else: 35 | # if existed then read this ply file to fill the data/xyz/colors/labels 36 | data = read_ply(save_path) # ply format: x,y,z,red,gree,blue,class 37 | xyz = np.vstack((data['x'], data['y'], data['z'])).T # (N',3), note the transpose symbol 38 | num_raw_points = xyz.shape[0] 39 | colors = np.vstack((data['red'], data['green'], data['blue'])).T # (N',3), note the transpose symbol 40 | labels = data['class'] 41 | pc_label = np.concatenate((xyz, colors, np.expand_dims(labels, axis=1)),axis=1) # (N,7) 42 | 43 | 44 | # save sub_cloud 45 | sub_ply_file = join(sub_pc_folder, save_path.split('/')[-1][:-4] + '.ply') 46 | if not os.path.exists(sub_ply_file): 47 | raise NotImplementedError("run the dataset_prepare_s3dis_sqn.py to generate raw and sub PC") 48 | else: 49 | data = read_ply(sub_ply_file) # ply format: x,y,z,red,gree,blue,class 50 | sub_xyz = np.vstack((data['x'], data['y'], data['z'])).T # (N',3), note the transpose symbol 51 | num_sub_points = sub_xyz.shape[0] 52 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T # (N',3), note the transpose symbol 53 | sub_labels = data['class'] 54 | 55 | 56 | # weak labels 57 | weak_label_sub_file = join(weak_label_folder, save_path.split('/')[-1][:-4] + '_sub_weak_label.ply') 58 | if not os.path.exists(weak_label_sub_file): 59 | raise NotImplementedError("run the dataset_prepare_s3dis_sqn.py to generate weak labels") 60 | else: 61 | data = read_ply(weak_label_sub_file) 62 | weak_label_mask = data['weak_mask'] 63 | 64 | # compute number of points for each class in this room 65 | num_points_raw = np.zeros(len(gt_class), dtype=np.int32) 66 | num_points_sub = np.zeros(len(gt_class), dtype=np.int32) 67 | num_points_weak = np.zeros(len(gt_class), dtype=np.int32) 68 | # raw 69 | for i,item in enumerate(labels): 70 | num_points_raw[int(item)]+=1 71 | # sub-sampled points 72 | for i,item in enumerate(sub_labels): 73 | num_points_sub[int(item)]+=1 74 | # weak points 75 | for i,item in enumerate(sub_labels[weak_label_mask.astype(bool)]): 76 | num_points_weak[int(item)]+=1 77 | 78 | return num_points_raw, num_points_sub, num_points_weak 79 | 80 | 81 | if __name__ == '__main__': 82 | 83 | parser = argparse.ArgumentParser() 84 | parser.add_argument("--rng_seed", type=int, default=123, help='manual seed') 85 | parser.add_argument('--dataset_path', type=str, default='./data/S3DIS/Stanford3dDataset_v1.2_Aligned_Version', help='dataset path') 86 | parser.add_argument('--sub_grid_size', type=float, default=0.04, help='grid-sampling size') 87 | parser.add_argument('--weak_label_ratio', type=float, default=0.001, help='the weakly semantic segmentation ratio') 88 | FLAGS = parser.parse_args() 89 | 90 | # set fixed seeds for reproducible results 91 | random.seed(FLAGS.rng_seed) 92 | np.random.seed(FLAGS.rng_seed) 93 | 94 | dataset_path = FLAGS.dataset_path 95 | anno_paths = [line.rstrip() for line in open(join(BASE_DIR, 'meta/anno_paths.txt'))] 96 | anno_paths = [join(dataset_path, p) for p in anno_paths] # each room's path 97 | # object categories for the S3DIS dataset 98 | gt_class = [x.rstrip() for x in open(join(BASE_DIR, 'meta/class_names.txt'))] 99 | gt_class2label = {cls: i for i, cls in enumerate(gt_class)} 100 | sub_grid_size = FLAGS.sub_grid_size # grid_subsampling size 101 | 102 | """ 103 | create 3 folder 104 | - input_0.040, for sub_pc.py, kdtree for sub_pc and the projection indices 105 | - original_ply, raw_pc.ply 106 | - weak_label_0.01, weak labels for raw and sub_pc 107 | """ 108 | original_pc_folder = join(dirname(dataset_path), 'original_ply') 109 | sub_pc_folder = join(dirname(dataset_path), 'input_{:.3f}'.format(sub_grid_size)) 110 | weak_label_ratio = FLAGS.weak_label_ratio 111 | weak_label_folder = join(dirname(dataset_path), 'weak_label_{}'.format(weak_label_ratio)) 112 | os.mkdir(original_pc_folder) if not exists(original_pc_folder) else None 113 | os.mkdir(sub_pc_folder) if not exists(sub_pc_folder) else None 114 | os.mkdir(weak_label_folder) if not exists(weak_label_folder) else None 115 | out_format = '.ply' 116 | 117 | # Note: there is an extra character in the v1.2 data in Area_5/hallway_6. It's fixed manually. 118 | num_points_each_class_raw = np.zeros(len(gt_class), dtype=np.int32) 119 | num_points_each_class_sub = np.zeros(len(gt_class), dtype=np.int32) 120 | num_points_each_class_weak = np.zeros(len(gt_class), dtype=np.int32) 121 | 122 | for annotation_path in anno_paths: 123 | 124 | # e.g.: data/S3DIS/Stanford3dDataset_v1.2_Aligned_Version/Area_1/conferenceRoom_1/Annotations 125 | print(annotation_path) 126 | elements = str(annotation_path).split('/') 127 | # e.g.: Area_1_conferenceRoom_1.ply 128 | out_file_name = elements[-3] + '_' + elements[-2] + out_format 129 | 130 | # convert each room's pc to ply and more(kdtree and projection indices for raw points over its corresponding sub_pc) 131 | num_points_room_raw, num_points_room_sub, num_points_room_weak = compute_num_points(annotation_path, 132 | join(original_pc_folder, out_file_name), 133 | sub_pc_folder, weak_label_folder, 134 | weak_label_ratio, sub_grid_size, 135 | gt_class, gt_class2label) 136 | 137 | num_points_each_class_raw += num_points_room_raw 138 | num_points_each_class_sub += num_points_room_sub 139 | num_points_each_class_weak += num_points_room_weak 140 | 141 | print("finish computing w. success!") 142 | # print(f"num_points for raw:{num_points_each_class_raw}\n") 143 | # print(f"num_points for sub:{num_points_each_class_sub}\n") 144 | # print(f"num_points for weak:{num_points_each_class_weak}") 145 | 146 | rows = [ 147 | ['type'] + list(gt_class2label.values()), 148 | ['The number of raw points'] + num_points_each_class_raw.tolist(), 149 | ['The number of sub points'] + num_points_each_class_sub.tolist(), 150 | ['The number of weak points'] + num_points_each_class_weak.tolist(), 151 | ] 152 | filename_xls = join(dirname(dataset_path), f'S3DIS_point_distribution_sub_{sub_grid_size:.3f}_weak_{weak_label_ratio}.xlsx') 153 | print('Save file to {}'.format(filename_xls)) 154 | with xlsxwriter.Workbook(filename_xls) as workbook: 155 | worksheet = workbook.add_worksheet() 156 | for i, data in enumerate(rows): 157 | worksheet.write_row(i, 0, data) 158 | print('Finished.') 159 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/compile_wrappers.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Compile cpp subsampling 4 | cd cpp_subsampling 5 | python3 setup.py build_ext --inplace 6 | cd .. 7 | 8 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "grid_subsampling.h" 3 | 4 | 5 | void grid_subsampling(vector& original_points, 6 | vector& subsampled_points, 7 | vector& original_features, 8 | vector& subsampled_features, 9 | vector& original_classes, 10 | vector& subsampled_classes, 11 | float sampleDl, 12 | int verbose) { 13 | 14 | // Initiate variables 15 | // ****************** 16 | 17 | // Number of points in the cloud 18 | size_t N = original_points.size(); 19 | 20 | // Dimension of the features 21 | size_t fdim = original_features.size() / N; 22 | size_t ldim = original_classes.size() / N; 23 | 24 | // Limits of the cloud 25 | PointXYZ minCorner = min_point(original_points); 26 | PointXYZ maxCorner = max_point(original_points); 27 | PointXYZ originCorner = floor(minCorner * (1/sampleDl)) * sampleDl; 28 | 29 | // Dimensions of the grid 30 | size_t sampleNX = (size_t)floor((maxCorner.x - originCorner.x) / sampleDl) + 1; 31 | size_t sampleNY = (size_t)floor((maxCorner.y - originCorner.y) / sampleDl) + 1; 32 | //size_t sampleNZ = (size_t)floor((maxCorner.z - originCorner.z) / sampleDl) + 1; 33 | 34 | // Check if features and classes need to be processed 35 | bool use_feature = original_features.size() > 0; 36 | bool use_classes = original_classes.size() > 0; 37 | 38 | 39 | // Create the sampled map 40 | // ********************** 41 | 42 | // Verbose parameters 43 | int i = 0; 44 | int nDisp = N / 100; 45 | 46 | // Initiate variables 47 | size_t iX, iY, iZ, mapIdx; 48 | unordered_map data; 49 | 50 | for (auto& p : original_points) 51 | { 52 | // Position of point in sample map 53 | iX = (size_t)floor((p.x - originCorner.x) / sampleDl); 54 | iY = (size_t)floor((p.y - originCorner.y) / sampleDl); 55 | iZ = (size_t)floor((p.z - originCorner.z) / sampleDl); 56 | mapIdx = iX + sampleNX*iY + sampleNX*sampleNY*iZ; 57 | 58 | // If not already created, create key 59 | if (data.count(mapIdx) < 1) 60 | data.emplace(mapIdx, SampledData(fdim, ldim)); 61 | 62 | // Fill the sample map 63 | if (use_feature && use_classes) 64 | data[mapIdx].update_all(p, original_features.begin() + i * fdim, original_classes.begin() + i * ldim); 65 | else if (use_feature) 66 | data[mapIdx].update_features(p, original_features.begin() + i * fdim); 67 | else if (use_classes) 68 | data[mapIdx].update_classes(p, original_classes.begin() + i * ldim); 69 | else 70 | data[mapIdx].update_points(p); 71 | 72 | // Display 73 | i++; 74 | if (verbose > 1 && i%nDisp == 0) 75 | std::cout << "\rSampled Map : " << std::setw(3) << i / nDisp << "%"; 76 | 77 | } 78 | 79 | // Divide for barycentre and transfer to a vector 80 | subsampled_points.reserve(data.size()); 81 | if (use_feature) 82 | subsampled_features.reserve(data.size() * fdim); 83 | if (use_classes) 84 | subsampled_classes.reserve(data.size() * ldim); 85 | for (auto& v : data) 86 | { 87 | subsampled_points.push_back(v.second.point * (1.0 / v.second.count)); 88 | if (use_feature) 89 | { 90 | float count = (float)v.second.count; 91 | transform(v.second.features.begin(), 92 | v.second.features.end(), 93 | v.second.features.begin(), 94 | [count](float f) { return f / count;}); 95 | subsampled_features.insert(subsampled_features.end(),v.second.features.begin(),v.second.features.end()); 96 | } 97 | if (use_classes) 98 | { 99 | for (int i = 0; i < ldim; i++) 100 | subsampled_classes.push_back(max_element(v.second.labels[i].begin(), v.second.labels[i].end(), 101 | [](const pair&a, const pair&b){return a.second < b.second;})->first); 102 | } 103 | } 104 | 105 | return; 106 | } 107 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h: -------------------------------------------------------------------------------- 1 | 2 | 3 | #include "../../cpp_utils/cloud/cloud.h" 4 | 5 | #include 6 | #include 7 | 8 | using namespace std; 9 | 10 | class SampledData 11 | { 12 | public: 13 | 14 | // Elements 15 | // ******** 16 | 17 | int count; 18 | PointXYZ point; 19 | vector features; 20 | vector> labels; 21 | 22 | 23 | // Methods 24 | // ******* 25 | 26 | // Constructor 27 | SampledData() 28 | { 29 | count = 0; 30 | point = PointXYZ(); 31 | } 32 | 33 | SampledData(const size_t fdim, const size_t ldim) 34 | { 35 | count = 0; 36 | point = PointXYZ(); 37 | features = vector(fdim); 38 | labels = vector>(ldim); 39 | } 40 | 41 | // Method Update 42 | void update_all(const PointXYZ p, vector::iterator f_begin, vector::iterator l_begin) 43 | { 44 | count += 1; 45 | point += p; 46 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 47 | int i = 0; 48 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 49 | { 50 | labels[i][*it] += 1; 51 | i++; 52 | } 53 | return; 54 | } 55 | void update_features(const PointXYZ p, vector::iterator f_begin) 56 | { 57 | count += 1; 58 | point += p; 59 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 60 | return; 61 | } 62 | void update_classes(const PointXYZ p, vector::iterator l_begin) 63 | { 64 | count += 1; 65 | point += p; 66 | int i = 0; 67 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 68 | { 69 | labels[i][*it] += 1; 70 | i++; 71 | } 72 | return; 73 | } 74 | void update_points(const PointXYZ p) 75 | { 76 | count += 1; 77 | point += p; 78 | return; 79 | } 80 | }; 81 | 82 | 83 | 84 | void grid_subsampling(vector& original_points, 85 | vector& subsampled_points, 86 | vector& original_features, 87 | vector& subsampled_features, 88 | vector& original_classes, 89 | vector& subsampled_classes, 90 | float sampleDl, 91 | int verbose); 92 | 93 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup, Extension 2 | import numpy.distutils.misc_util 3 | 4 | # Adding OpenCV to project 5 | # ************************ 6 | 7 | # Adding sources of the project 8 | # ***************************** 9 | 10 | m_name = "grid_subsampling" 11 | 12 | SOURCES = ["../cpp_utils/cloud/cloud.cpp", 13 | "grid_subsampling/grid_subsampling.cpp", 14 | "wrapper.cpp"] 15 | 16 | module = Extension(m_name, 17 | sources=SOURCES, 18 | extra_compile_args=['-std=c++11', 19 | '-D_GLIBCXX_USE_CXX11_ABI=0']) 20 | 21 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/wrapper.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include "grid_subsampling/grid_subsampling.h" 4 | #include 5 | 6 | 7 | 8 | // docstrings for our module 9 | // ************************* 10 | 11 | static char module_docstring[] = "This module provides an interface for the subsampling of a pointcloud"; 12 | 13 | static char compute_docstring[] = "function subsampling a pointcloud"; 14 | 15 | 16 | // Declare the functions 17 | // ********************* 18 | 19 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds); 20 | 21 | 22 | // Specify the members of the module 23 | // ********************************* 24 | 25 | static PyMethodDef module_methods[] = 26 | { 27 | { "compute", (PyCFunction)grid_subsampling_compute, METH_VARARGS | METH_KEYWORDS, compute_docstring }, 28 | {NULL, NULL, 0, NULL} 29 | }; 30 | 31 | 32 | // Initialize the module 33 | // ********************* 34 | 35 | static struct PyModuleDef moduledef = 36 | { 37 | PyModuleDef_HEAD_INIT, 38 | "grid_subsampling", // m_name 39 | module_docstring, // m_doc 40 | -1, // m_size 41 | module_methods, // m_methods 42 | NULL, // m_reload 43 | NULL, // m_traverse 44 | NULL, // m_clear 45 | NULL, // m_free 46 | }; 47 | 48 | PyMODINIT_FUNC PyInit_grid_subsampling(void) 49 | { 50 | import_array(); 51 | return PyModule_Create(&moduledef); 52 | } 53 | 54 | 55 | // Actual wrapper 56 | // ************** 57 | 58 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds) 59 | { 60 | 61 | // Manage inputs 62 | // ************* 63 | 64 | // Args containers 65 | PyObject *points_obj = NULL; 66 | PyObject *features_obj = NULL; 67 | PyObject *classes_obj = NULL; 68 | 69 | // Keywords containers 70 | static char *kwlist[] = {"points", "features", "classes", "sampleDl", "method", "verbose", NULL }; 71 | float sampleDl = 0.1; 72 | const char *method_buffer = "barycenters"; 73 | int verbose = 0; 74 | 75 | // Parse the input 76 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "O|$OOfsi", kwlist, &points_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &verbose)) 77 | { 78 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); 79 | return NULL; 80 | } 81 | 82 | // Get the method argument 83 | string method(method_buffer); 84 | 85 | // Interpret method 86 | if (method.compare("barycenters") && method.compare("voxelcenters")) 87 | { 88 | PyErr_SetString(PyExc_RuntimeError, "Error parsing method. Valid method names are \"barycenters\" and \"voxelcenters\" "); 89 | return NULL; 90 | } 91 | 92 | // Check if using features or classes 93 | bool use_feature = true, use_classes = true; 94 | if (features_obj == NULL) 95 | use_feature = false; 96 | if (classes_obj == NULL) 97 | use_classes = false; 98 | 99 | // Interpret the input objects as numpy arrays. 100 | PyObject *points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY); 101 | PyObject *features_array = NULL; 102 | PyObject *classes_array = NULL; 103 | if (use_feature) 104 | features_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY); 105 | if (use_classes) 106 | classes_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY); 107 | 108 | // Verify data was load correctly. 109 | if (points_array == NULL) 110 | { 111 | Py_XDECREF(points_array); 112 | Py_XDECREF(classes_array); 113 | Py_XDECREF(features_array); 114 | PyErr_SetString(PyExc_RuntimeError, "Error converting input points to numpy arrays of type float32"); 115 | return NULL; 116 | } 117 | if (use_feature && features_array == NULL) 118 | { 119 | Py_XDECREF(points_array); 120 | Py_XDECREF(classes_array); 121 | Py_XDECREF(features_array); 122 | PyErr_SetString(PyExc_RuntimeError, "Error converting input features to numpy arrays of type float32"); 123 | return NULL; 124 | } 125 | if (use_classes && classes_array == NULL) 126 | { 127 | Py_XDECREF(points_array); 128 | Py_XDECREF(classes_array); 129 | Py_XDECREF(features_array); 130 | PyErr_SetString(PyExc_RuntimeError, "Error converting input classes to numpy arrays of type int32"); 131 | return NULL; 132 | } 133 | 134 | // Check that the input array respect the dims 135 | if ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3) 136 | { 137 | Py_XDECREF(points_array); 138 | Py_XDECREF(classes_array); 139 | Py_XDECREF(features_array); 140 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : points.shape is not (N, 3)"); 141 | return NULL; 142 | } 143 | if (use_feature && ((int)PyArray_NDIM(features_array) != 2)) 144 | { 145 | Py_XDECREF(points_array); 146 | Py_XDECREF(classes_array); 147 | Py_XDECREF(features_array); 148 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 149 | return NULL; 150 | } 151 | 152 | if (use_classes && (int)PyArray_NDIM(classes_array) > 2) 153 | { 154 | Py_XDECREF(points_array); 155 | Py_XDECREF(classes_array); 156 | Py_XDECREF(features_array); 157 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 158 | return NULL; 159 | } 160 | 161 | // Number of points 162 | int N = (int)PyArray_DIM(points_array, 0); 163 | 164 | // Dimension of the features 165 | int fdim = 0; 166 | if (use_feature) 167 | fdim = (int)PyArray_DIM(features_array, 1); 168 | 169 | //Dimension of labels 170 | int ldim = 1; 171 | if (use_classes && (int)PyArray_NDIM(classes_array) == 2) 172 | ldim = (int)PyArray_DIM(classes_array, 1); 173 | 174 | // Check that the input array respect the number of points 175 | if (use_feature && (int)PyArray_DIM(features_array, 0) != N) 176 | { 177 | Py_XDECREF(points_array); 178 | Py_XDECREF(classes_array); 179 | Py_XDECREF(features_array); 180 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 181 | return NULL; 182 | } 183 | if (use_classes && (int)PyArray_DIM(classes_array, 0) != N) 184 | { 185 | Py_XDECREF(points_array); 186 | Py_XDECREF(classes_array); 187 | Py_XDECREF(features_array); 188 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 189 | return NULL; 190 | } 191 | 192 | 193 | // Call the C++ function 194 | // ********************* 195 | 196 | // Create pyramid 197 | if (verbose > 0) 198 | cout << "Computing cloud pyramid with support points: " << endl; 199 | 200 | 201 | // Convert PyArray to Cloud C++ class 202 | vector original_points; 203 | vector original_features; 204 | vector original_classes; 205 | original_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N); 206 | if (use_feature) 207 | original_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N*fdim); 208 | if (use_classes) 209 | original_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N*ldim); 210 | 211 | // Subsample 212 | vector subsampled_points; 213 | vector subsampled_features; 214 | vector subsampled_classes; 215 | grid_subsampling(original_points, 216 | subsampled_points, 217 | original_features, 218 | subsampled_features, 219 | original_classes, 220 | subsampled_classes, 221 | sampleDl, 222 | verbose); 223 | 224 | // Check result 225 | if (subsampled_points.size() < 1) 226 | { 227 | PyErr_SetString(PyExc_RuntimeError, "Error"); 228 | return NULL; 229 | } 230 | 231 | // Manage outputs 232 | // ************** 233 | 234 | // Dimension of input containers 235 | npy_intp* point_dims = new npy_intp[2]; 236 | point_dims[0] = subsampled_points.size(); 237 | point_dims[1] = 3; 238 | npy_intp* feature_dims = new npy_intp[2]; 239 | feature_dims[0] = subsampled_points.size(); 240 | feature_dims[1] = fdim; 241 | npy_intp* classes_dims = new npy_intp[2]; 242 | classes_dims[0] = subsampled_points.size(); 243 | classes_dims[1] = ldim; 244 | 245 | // Create output array 246 | PyObject *res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT); 247 | PyObject *res_features_obj = NULL; 248 | PyObject *res_classes_obj = NULL; 249 | PyObject *ret = NULL; 250 | 251 | // Fill output array with values 252 | size_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float); 253 | memcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes); 254 | if (use_feature) 255 | { 256 | size_in_bytes = subsampled_points.size() * fdim * sizeof(float); 257 | res_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT); 258 | memcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes); 259 | } 260 | if (use_classes) 261 | { 262 | size_in_bytes = subsampled_points.size() * ldim * sizeof(int); 263 | res_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT); 264 | memcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes); 265 | } 266 | 267 | 268 | // Merge results 269 | if (use_feature && use_classes) 270 | ret = Py_BuildValue("NNN", res_points_obj, res_features_obj, res_classes_obj); 271 | else if (use_feature) 272 | ret = Py_BuildValue("NN", res_points_obj, res_features_obj); 273 | else if (use_classes) 274 | ret = Py_BuildValue("NN", res_points_obj, res_classes_obj); 275 | else 276 | ret = Py_BuildValue("N", res_points_obj); 277 | 278 | // Clean up 279 | // ******** 280 | 281 | Py_DECREF(points_array); 282 | Py_XDECREF(features_array); 283 | Py_XDECREF(classes_array); 284 | 285 | return ret; 286 | } -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_utils/cloud/cloud.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud source : 13 | // Define usefull Functions/Methods 14 | // 15 | //---------------------------------------------------- 16 | // 17 | // Hugues THOMAS - 10/02/2017 18 | // 19 | 20 | 21 | #include "cloud.h" 22 | 23 | 24 | // Getters 25 | // ******* 26 | 27 | PointXYZ max_point(std::vector points) 28 | { 29 | // Initiate limits 30 | PointXYZ maxP(points[0]); 31 | 32 | // Loop over all points 33 | for (auto p : points) 34 | { 35 | if (p.x > maxP.x) 36 | maxP.x = p.x; 37 | 38 | if (p.y > maxP.y) 39 | maxP.y = p.y; 40 | 41 | if (p.z > maxP.z) 42 | maxP.z = p.z; 43 | } 44 | 45 | return maxP; 46 | } 47 | 48 | PointXYZ min_point(std::vector points) 49 | { 50 | // Initiate limits 51 | PointXYZ minP(points[0]); 52 | 53 | // Loop over all points 54 | for (auto p : points) 55 | { 56 | if (p.x < minP.x) 57 | minP.x = p.x; 58 | 59 | if (p.y < minP.y) 60 | minP.y = p.y; 61 | 62 | if (p.z < minP.z) 63 | minP.z = p.z; 64 | } 65 | 66 | return minP; 67 | } -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_utils/cloud/cloud.h: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud header 13 | // 14 | //---------------------------------------------------- 15 | // 16 | // Hugues THOMAS - 10/02/2017 17 | // 18 | 19 | 20 | # pragma once 21 | 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | 31 | #include 32 | 33 | 34 | 35 | 36 | // Point class 37 | // *********** 38 | 39 | 40 | class PointXYZ 41 | { 42 | public: 43 | 44 | // Elements 45 | // ******** 46 | 47 | float x, y, z; 48 | 49 | 50 | // Methods 51 | // ******* 52 | 53 | // Constructor 54 | PointXYZ() { x = 0; y = 0; z = 0; } 55 | PointXYZ(float x0, float y0, float z0) { x = x0; y = y0; z = z0; } 56 | 57 | // array type accessor 58 | float operator [] (int i) const 59 | { 60 | if (i == 0) return x; 61 | else if (i == 1) return y; 62 | else return z; 63 | } 64 | 65 | // opperations 66 | float dot(const PointXYZ P) const 67 | { 68 | return x * P.x + y * P.y + z * P.z; 69 | } 70 | 71 | float sq_norm() 72 | { 73 | return x*x + y*y + z*z; 74 | } 75 | 76 | PointXYZ cross(const PointXYZ P) const 77 | { 78 | return PointXYZ(y*P.z - z*P.y, z*P.x - x*P.z, x*P.y - y*P.x); 79 | } 80 | 81 | PointXYZ& operator+=(const PointXYZ& P) 82 | { 83 | x += P.x; 84 | y += P.y; 85 | z += P.z; 86 | return *this; 87 | } 88 | 89 | PointXYZ& operator-=(const PointXYZ& P) 90 | { 91 | x -= P.x; 92 | y -= P.y; 93 | z -= P.z; 94 | return *this; 95 | } 96 | 97 | PointXYZ& operator*=(const float& a) 98 | { 99 | x *= a; 100 | y *= a; 101 | z *= a; 102 | return *this; 103 | } 104 | }; 105 | 106 | 107 | // Point Opperations 108 | // ***************** 109 | 110 | inline PointXYZ operator + (const PointXYZ A, const PointXYZ B) 111 | { 112 | return PointXYZ(A.x + B.x, A.y + B.y, A.z + B.z); 113 | } 114 | 115 | inline PointXYZ operator - (const PointXYZ A, const PointXYZ B) 116 | { 117 | return PointXYZ(A.x - B.x, A.y - B.y, A.z - B.z); 118 | } 119 | 120 | inline PointXYZ operator * (const PointXYZ P, const float a) 121 | { 122 | return PointXYZ(P.x * a, P.y * a, P.z * a); 123 | } 124 | 125 | inline PointXYZ operator * (const float a, const PointXYZ P) 126 | { 127 | return PointXYZ(P.x * a, P.y * a, P.z * a); 128 | } 129 | 130 | inline std::ostream& operator << (std::ostream& os, const PointXYZ P) 131 | { 132 | return os << "[" << P.x << ", " << P.y << ", " << P.z << "]"; 133 | } 134 | 135 | inline bool operator == (const PointXYZ A, const PointXYZ B) 136 | { 137 | return A.x == B.x && A.y == B.y && A.z == B.z; 138 | } 139 | 140 | inline PointXYZ floor(const PointXYZ P) 141 | { 142 | return PointXYZ(std::floor(P.x), std::floor(P.y), std::floor(P.z)); 143 | } 144 | 145 | 146 | PointXYZ max_point(std::vector points); 147 | PointXYZ min_point(std::vector points); 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | -------------------------------------------------------------------------------- /utils/data_prepare_s3dis.py: -------------------------------------------------------------------------------- 1 | from sklearn.neighbors import KDTree 2 | from os.path import join, exists, dirname, abspath 3 | import numpy as np 4 | import pandas as pd 5 | import os, sys, glob, pickle 6 | 7 | BASE_DIR = dirname(abspath(__file__)) 8 | ROOT_DIR = dirname(BASE_DIR) 9 | sys.path.append(BASE_DIR) 10 | sys.path.append(ROOT_DIR) 11 | from helper_ply import write_ply 12 | from helper_tool import DataProcessing as DP 13 | 14 | # dataset_path = '/data/S3DIS/Stanford3dDataset_v1.2_Aligned_Version' 15 | dataset_path = './data/S3DIS/Stanford3dDataset_v1.2_Aligned_Version' 16 | 17 | # obtain room paths for training areas -yc 18 | anno_paths = [line.rstrip() for line in open(join(BASE_DIR, 'meta/anno_paths.txt'))] 19 | anno_paths = [join(dataset_path, p) for p in anno_paths] 20 | 21 | # object categories for the S3DIS dataset 22 | gt_class = [x.rstrip() for x in open(join(BASE_DIR, 'meta/class_names.txt'))] 23 | gt_class2label = {cls: i for i, cls in enumerate(gt_class)} 24 | 25 | sub_grid_size = 0.04 26 | original_pc_folder = join(dirname(dataset_path), 'original_ply') 27 | sub_pc_folder = join(dirname(dataset_path), 'input_{:.3f}'.format(sub_grid_size)) 28 | os.mkdir(original_pc_folder) if not exists(original_pc_folder) else None 29 | os.mkdir(sub_pc_folder) if not exists(sub_pc_folder) else None 30 | out_format = '.ply' 31 | 32 | 33 | def convert_pc2ply(anno_path, save_path): 34 | """ 35 | Convert original dataset files to ply file (each line is XYZRGBL). 36 | We aggregated all the points from each instance in the room. 37 | :param anno_path: path to annotations. e.g. Area_1/office_2/Annotations/ 38 | :param save_path: path to save original point clouds (each line is XYZRGBL) 39 | :return: None; 40 | note: Physically, each room will generate four files, including raw_pc.ply, sub_pc.ply, sub_pc.pkl for the kdtree and proj_idx.pkl for each raw point's nearest neighbor in the sub_pc ) 41 | """ 42 | 43 | # store points and labels for the room(correspond to the anno_path), yc 44 | data_list = [] 45 | 46 | for f in glob.glob(join(anno_path, '*.txt')): 47 | class_name = os.path.basename(f).split('_')[0] 48 | if class_name not in gt_class: # note: in some room there is 'staris' class.. 49 | class_name = 'clutter' 50 | pc = pd.read_csv(f, header=None, delim_whitespace=True).values 51 | labels = np.ones((pc.shape[0], 1)) * gt_class2label[class_name] 52 | data_list.append(np.concatenate([pc, labels], 1)) # Nx7 53 | 54 | # translate the data by xyz_min--yc 55 | pc_label = np.concatenate(data_list, 0) # Nx7 as a np object 56 | xyz_min = np.amin(pc_label, axis=0)[0:3] 57 | pc_label[:, 0:3] -= xyz_min 58 | 59 | # manage data types and save in PLY format--yc 60 | xyz = pc_label[:, :3].astype(np.float32) 61 | colors = pc_label[:, 3:6].astype(np.uint8) 62 | labels = pc_label[:, 6].astype(np.uint8) 63 | write_ply(save_path, (xyz, colors, labels), ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 64 | 65 | # save sub_cloud and KDTree file 66 | sub_xyz, sub_colors, sub_labels = DP.grid_sub_sampling(xyz, colors, labels, sub_grid_size) 67 | sub_colors = sub_colors / 255.0 68 | sub_ply_file = join(sub_pc_folder, save_path.split('/')[-1][:-4] + '.ply') 69 | write_ply(sub_ply_file, [sub_xyz, sub_colors, sub_labels], ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 70 | 71 | search_tree = KDTree(sub_xyz) 72 | kd_tree_file = join(sub_pc_folder, str(save_path.split('/')[-1][:-4]) + '_KDTree.pkl') 73 | with open(kd_tree_file, 'wb') as f: 74 | pickle.dump(search_tree, f) 75 | 76 | # nearest nb index list for xyz when searching using the sub-sampled PC generated kdtree--yc 77 | # https://scikit-learn.org/stable/modules/generated/sklearn.neighbors.KDTree.html 78 | proj_idx = np.squeeze(search_tree.query(xyz, return_distance=False)) 79 | proj_idx = proj_idx.astype(np.int32) 80 | proj_save = join(sub_pc_folder, str(save_path.split('/')[-1][:-4]) + '_proj.pkl') 81 | with open(proj_save, 'wb') as f: 82 | pickle.dump([proj_idx, labels], f) 83 | 84 | 85 | """ 86 | Prepare the S3DIS dataset for training by generating new info from each room's point cloud(PC) 87 | - input: each room's PC 88 | - output: 1) raw_pc in ply, 2)sub_pc in ply, 3) KDTree for the sub_pc, and 4)projection indices for each raw point over the sub_pc(used for DL validation/inference as the learning process only occur on the sub_pc, therefore by relating raw point's relations to the sub_pc can help propagate their semantics.) 89 | """ 90 | if __name__ == '__main__': 91 | # Note: there is an extra character in the v1.2 data in Area_5/hallway_6. It's fixed manually. 92 | for annotation_path in anno_paths: 93 | 94 | # e.g.: data/S3DIS/Stanford3dDataset_v1.2_Aligned_Version/Area_1/conferenceRoom_1/Annotations 95 | print(annotation_path) 96 | elements = str(annotation_path).split('/') 97 | # e.g.: Area_1_conferenceRoom_1.ply 98 | out_file_name = elements[-3] + '_' + elements[-2] + out_format 99 | 100 | # convert each room's pc to ply and more(kdtree and projection indices for raw points over its corresponding sub_pc) 101 | convert_pc2ply(annotation_path, join(original_pc_folder, out_file_name)) -------------------------------------------------------------------------------- /utils/data_prepare_s3dis_sqn.py: -------------------------------------------------------------------------------- 1 | """ 2 | prepare S3DIS dataset for SQN model, reproduced based on SQN paper (https://arxiv.org/abs/2104.04891) 3 | Author: Chao YIN 4 | Email: cyinac@connect.ust.hk 5 | 6 | history: 7 | - Oct. 15, 2021, init the file 8 | - Oct. 26, 2021, **fix a fatal bug which is primarily caused by misinterpretation of weak label ration.** 9 | codebase: data_prepare_s3dis.py of the official RandLA-Net 10 | 11 | difference from the codebase (data_prepare_s3dis.py of Official RandLA-Net) 12 | - add CLI arguments (e.g., sub_grid_size, weak_label_ratio) support with argparse 13 | - generate separate weak labels for each room in S3DIS 14 | - refactor the code; if the raw/sub-pc/kdtree/projected indices/weak_labels files exist, then read them into memory. 15 | """ 16 | 17 | import os, sys, glob, pickle, argparse, random 18 | from os.path import join, exists, dirname, abspath 19 | import numpy as np 20 | import pandas as pd 21 | from sklearn.neighbors import KDTree 22 | 23 | BASE_DIR = dirname(abspath(__file__)) 24 | ROOT_DIR = dirname(BASE_DIR) 25 | sys.path.append(BASE_DIR) 26 | sys.path.append(ROOT_DIR) 27 | from helper_ply import write_ply, read_ply 28 | from helper_tool import DataProcessing as DP 29 | 30 | 31 | def convert_pc2plyandweaklabels(anno_path, save_path, sub_pc_folder, 32 | weak_label_folder, weak_label_ratio, sub_grid_size, 33 | gt_class, gt_class2label): 34 | """ 35 | Convert original dataset files (consiting of rooms) to ply file and weak labels. Physically, each room will generate several files, including raw_pc.ply, sub_pc.ply, sub_pc.pkl (for the kdtree), proj_idx.pkl (for each raw point's nearest neighbor in the sub_pc) and weak labels for raw and sub_pc, respectively ) 36 | :param anno_path: path to annotations. e.g. Area_1/office_2/Annotations/ 37 | :param save_path: path to save original point clouds (each line is XYZRGBL), e.g., xx.ply 38 | :return: None 39 | """ 40 | 41 | num_raw_points = 0 # number of raw points for current room 42 | num_sub_points = 0 # number of sub-sampled points for current room 43 | 44 | # save raw_cloud 45 | if not os.path.exists(save_path): 46 | data_list = [] 47 | # aggregate a room's instances into 1 pc 48 | for f in glob.glob(join(anno_path, '*.txt')): 49 | class_name = os.path.basename(f).split('_')[0] 50 | if class_name not in gt_class: # note: in some room there is 'staris' class.. 51 | class_name = 'clutter' 52 | pc = pd.read_csv(f, header=None, delim_whitespace=True).values 53 | labels = np.ones((pc.shape[0], 1)) * gt_class2label[class_name] 54 | data_list.append(np.concatenate([pc, labels], 1)) # Nx7 55 | 56 | # translate the data by xyz_min--yc 57 | pc_label = np.concatenate(data_list, 0) # Nx7 as a np object 58 | xyz_min = np.amin(pc_label, axis=0)[0:3] 59 | pc_label[:, 0:3] -= xyz_min 60 | # manage data types and save in PLY format--yc 61 | xyz = pc_label[:, :3].astype(np.float32) 62 | num_raw_points = xyz.shape[0] 63 | colors = pc_label[:, 3:6].astype(np.uint8) 64 | labels = pc_label[:, 6].astype(np.uint8) 65 | write_ply(save_path, (xyz, colors, labels), ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 66 | else: 67 | # if existed then read this ply file to fill the data/xyz/colors/labels 68 | data = read_ply(save_path) # ply format: x,y,z,red,gree,blue,class 69 | xyz = np.vstack((data['x'], data['y'], data['z'])).T # (N',3), note the transpose symbol 70 | num_raw_points = xyz.shape[0] 71 | colors = np.vstack((data['red'], data['green'], data['blue'])).T # (N',3), note the transpose symbol 72 | labels = data['class'] 73 | pc_label = np.concatenate((xyz, colors, np.expand_dims(labels, axis=1)),axis=1) # (N,7) 74 | 75 | 76 | # save sub_cloud 77 | sub_ply_file = join(sub_pc_folder, save_path.split('/')[-1][:-4] + '.ply') 78 | if not os.path.exists(sub_ply_file): 79 | sub_xyz, sub_colors, sub_labels = DP.grid_sub_sampling(xyz, colors, labels, sub_grid_size) 80 | sub_colors = sub_colors / 255.0 81 | num_sub_points = sub_xyz.shape[0] 82 | write_ply(sub_ply_file, [sub_xyz, sub_colors, sub_labels], ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 83 | else: 84 | data = read_ply(sub_ply_file) # ply format: x,y,z,red,gree,blue,class 85 | sub_xyz = np.vstack((data['x'], data['y'], data['z'])).T # (N',3), note the transpose symbol 86 | num_sub_points = sub_xyz.shape[0] 87 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T # (N',3), note the transpose symbol 88 | sub_labels = data['class'] 89 | 90 | 91 | # save KDTree for sub_pc 92 | kd_tree_file = join(sub_pc_folder, str(save_path.split('/')[-1][:-4]) + '_KDTree.pkl') 93 | if not os.path.exists(kd_tree_file): 94 | search_tree = KDTree(sub_xyz) 95 | with open(kd_tree_file, 'wb') as f: 96 | pickle.dump(search_tree, f) 97 | 98 | 99 | # save projection indcies for all raw points over the corresponding sub_pc 100 | proj_save = join(sub_pc_folder, str(save_path.split('/')[-1][:-4]) + '_proj.pkl') 101 | if not os.path.exists(proj_save): 102 | proj_idx = np.squeeze(search_tree.query(xyz, return_distance=False)) 103 | proj_idx = proj_idx.astype(np.int32) 104 | with open(proj_save, 'wb') as f: 105 | pickle.dump([proj_idx, labels], f) 106 | 107 | 108 | # USED for weakly semantic segmentation, save sub pc's weak labels 109 | # KEY: Randomly select some points to own labels, give them a mask (no need to save weak label mask for raw pc) 110 | weak_label_sub_file = join(weak_label_folder, save_path.split('/')[-1][:-4] + '_sub_weak_label.ply') 111 | if not os.path.exists(weak_label_sub_file): 112 | 113 | # compute weak ratio of weak points w.r.t. #sub_pc 114 | print(f'Current sub-sampled ratio(#sub_points/#raw_points) is {(num_sub_points/num_raw_points)*100:.2f}%') 115 | print(f'Current weak_ratio(#weak_points/#raw_points) is {(weak_label_ratio):.4f}') 116 | 117 | # set weak points by randomly selecting weak_label_ratio*N points(i.e., the number of raw_pc) and denote them w. a mask 118 | weak_label_sub_mask = np.zeros((num_sub_points, 1), dtype=np.uint8) 119 | 120 | # BUG FIXED: fixed already; here, should set replace = True, otherwise a bug will be resulted 121 | # KEY: weak_label_ratio should be multiplied by number of raw points rather sub-sampled points 122 | selected_idx = np.random.choice(num_sub_points, int(num_raw_points*weak_label_ratio),replace=False) 123 | weak_label_sub_mask[selected_idx,:]=1 124 | write_ply(weak_label_sub_file, (weak_label_sub_mask,), ['weak_mask']) 125 | else: 126 | data = read_ply(weak_label_sub_file) 127 | weak_label_mask = data['weak_mask'] 128 | 129 | 130 | """ 131 | Prepare the S3DIS dataset for training SQN by generating new info from each room's point cloud(PC) 132 | - input: each room's PC 133 | - output: 134 | - 1) raw_pc in ply, 135 | - 2) sub_pc in ply, 136 | - 3) KDTree for the sub_pc, 137 | - 4) projection indices for each raw point over the sub_pc(used for DL validation/inference as the learning process only occur on the sub_pc, therefore by relating raw point's relations to the sub_pc can help propagate their semantics.) 138 | - 5) weak labels for raw and sub_pc 139 | """ 140 | if __name__ == '__main__': 141 | 142 | parser = argparse.ArgumentParser() 143 | parser.add_argument("--rng_seed", type=int, default=123, help='manual seed') 144 | parser.add_argument('--dataset_path', type=str, default='./data/S3DIS/Stanford3dDataset_v1.2_Aligned_Version', help='dataset path') 145 | parser.add_argument('--sub_grid_size', type=float, default=0.04, help='grid-sampling size') 146 | parser.add_argument('--weak_label_ratio', type=float, default=0.001, help='the weakly semantic segmentation ratio') 147 | parser.add_argument('--out_format', type=str, default='.ply', help='output format, e.g., ply') 148 | FLAGS = parser.parse_args() 149 | 150 | # set fixed seeds for reproducible results 151 | random.seed(FLAGS.rng_seed) 152 | np.random.seed(FLAGS.rng_seed) 153 | 154 | dataset_path = FLAGS.dataset_path 155 | anno_paths = [line.rstrip() for line in open(join(BASE_DIR, 'meta/anno_paths.txt'))] 156 | anno_paths = [join(dataset_path, p) for p in anno_paths] # each room's path 157 | # object categories for the S3DIS dataset 158 | gt_class = [x.rstrip() for x in open(join(BASE_DIR, 'meta/class_names.txt'))] 159 | gt_class2label = {cls: i for i, cls in enumerate(gt_class)} 160 | sub_grid_size = FLAGS.sub_grid_size # grid_subsampling size 161 | 162 | """ 163 | create 3 folder 164 | - input_0.040, for sub_pc.py, kdtree for sub_pc and the projection indices 165 | - original_ply, raw_pc.ply 166 | - weak_label_0.01, weak labels for raw and sub_pc 167 | """ 168 | original_pc_folder = join(dirname(dataset_path), 'original_ply') 169 | sub_pc_folder = join(dirname(dataset_path), 'input_{:.3f}'.format(sub_grid_size)) 170 | weak_label_ratio = FLAGS.weak_label_ratio 171 | weak_label_folder = join(dirname(dataset_path), 'weak_label_{}'.format(weak_label_ratio)) 172 | os.mkdir(original_pc_folder) if not exists(original_pc_folder) else None 173 | os.mkdir(sub_pc_folder) if not exists(sub_pc_folder) else None 174 | os.mkdir(weak_label_folder) if not exists(weak_label_folder) else None 175 | out_format = FLAGS.out_format 176 | 177 | # Note: there is an extra character in the v1.2 data in Area_5/hallway_6. It's fixed manually. 178 | for annotation_path in anno_paths: 179 | 180 | # e.g.: data/S3DIS/Stanford3dDataset_v1.2_Aligned_Version/Area_1/conferenceRoom_1/Annotations 181 | print(annotation_path) 182 | elements = str(annotation_path).split('/') 183 | # e.g.: Area_1_conferenceRoom_1.ply 184 | out_file_name = elements[-3] + '_' + elements[-2] + out_format 185 | 186 | # convert each room's pc to ply and more(kdtree and projection indices for raw points over its corresponding sub_pc) 187 | convert_pc2plyandweaklabels(annotation_path, join(original_pc_folder, out_file_name), 188 | sub_pc_folder, weak_label_folder, 189 | weak_label_ratio, sub_grid_size, 190 | gt_class, gt_class2label) -------------------------------------------------------------------------------- /utils/knn_interpolation.py: -------------------------------------------------------------------------------- 1 | """ 2 | Author: Chao YIN 3 | Email: cyinac@connect.ust.hk 4 | Date: Oct. 23, 2021 5 | code base: pytorch geometric,https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/unpool/knn_interpolate.html#knn_interpolate 6 | """ 7 | 8 | import numpy as np 9 | import torch 10 | # need install torch geometric 11 | from torch_geometric.nn import knn 12 | from torch_scatter import scatter_add 13 | 14 | 15 | # interpolate over x (N,C with shape (N,3)) located by pos_x to obtain y (located by pos_y with shape (M,3)) leading a tensor (N,C) 16 | def knn_interpolate(support_features, support_points, query_points, support_x=None, support_y=None, k=3, num_workers=1): 17 | """KNN interpolation, copied from , https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/unpool/knn_interpolate.html#knn_interpolate 18 | Args: 19 | support_features ([type]): [description] 20 | support_points ([type]): [description] 21 | query_points ([type]): [description] 22 | support_x ([type], optional): [description]. Defaults to None. 23 | support_y ([type], optional): [description]. Defaults to None. 24 | k (int, optional): [description]. Defaults to 3. 25 | num_workers (int, optional): [description]. Defaults to 1. 26 | Returns: 27 | [type]: [description] 28 | """ 29 | 30 | with torch.no_grad(): 31 | assign_index = knn(support_points, query_points, k, batch_x=support_x, batch_y=support_y, 32 | num_workers=num_workers) 33 | y_idx, x_idx = assign_index 34 | diff = support_points[x_idx] - query_points[y_idx] 35 | squared_distance = (diff * diff).sum(dim=-1, keepdim=True) 36 | weights = 1.0 / torch.clamp(squared_distance, min=1e-16) 37 | 38 | y = scatter_add(support_features[x_idx] * weights, y_idx, dim=0, dim_size=query_points.size(0)) 39 | y = y / scatter_add(weights, y_idx, dim=0, dim_size=query_points.size(0)) 40 | 41 | return y 42 | 43 | def batch_knn_interpolate(batch_support_features, batch_support_points, batch_query_points,k=3,num_workers=1): 44 | """[summary] 45 | Args: 46 | batch_support_features ([type]): (B,N,C) 47 | batch_support_points ([type]): (B,N,3) 48 | batch_query_points ([type]): (B,M,3) 49 | k (int, optional): [description]. Defaults to 3. 50 | num_workers (int, optional): [description]. Defaults to 1. 51 | return query features, (B,M,C) 52 | """ 53 | query_feature_list=[] 54 | for i in range(batch_support_features.shape[0]): 55 | y=knn_interpolate(batch_support_features[i], batch_support_points[i], batch_query_points[i], support_x=None, support_y=None, k=3, num_workers=1) 56 | query_feature_list.append(y) 57 | 58 | query_features=torch.tensor(query_feature_list) 59 | 60 | return query_features # (B,M,C) 61 | 62 | 63 | class BatchKnnInterpolate(object): 64 | 65 | def __init__(self, k=3, num_workers=1): 66 | self.k = k 67 | self.num_workers = num_workers 68 | 69 | def __call__(self, query_points, support_points, support_features): 70 | query_feature_list=[] 71 | batch_size = query_points.size()[0] 72 | for i in range(batch_size): 73 | query_features_cur=knn_interpolate(support_features[i], support_points[i], query_points[i], 74 | k=self.k, num_workers=self.num_workers) 75 | query_feature_list.append(query_features_cur) 76 | 77 | query_features=torch.tensor(query_feature_list) 78 | return query_features # (B,M,C) 79 | 80 | if __name__ == "__main__": 81 | 82 | batch_knn_op = BatchKnnInterpolate(k=3, num_workers=1) 83 | query_points = np.random.rand(4,2,3) 84 | support_points = np.random.rand(4,10,3) 85 | support_features = np.random.rand(4,10,5) 86 | 87 | query_features = batch_knn_op(query_points, support_points, support_features) 88 | 89 | print(query_features.shape, '\n', query_features) -------------------------------------------------------------------------------- /utils/meta/anno_paths.txt: -------------------------------------------------------------------------------- 1 | Area_1/conferenceRoom_1/Annotations 2 | Area_1/conferenceRoom_2/Annotations 3 | Area_1/copyRoom_1/Annotations 4 | Area_1/hallway_1/Annotations 5 | Area_1/hallway_2/Annotations 6 | Area_1/hallway_3/Annotations 7 | Area_1/hallway_4/Annotations 8 | Area_1/hallway_5/Annotations 9 | Area_1/hallway_6/Annotations 10 | Area_1/hallway_7/Annotations 11 | Area_1/hallway_8/Annotations 12 | Area_1/office_10/Annotations 13 | Area_1/office_11/Annotations 14 | Area_1/office_12/Annotations 15 | Area_1/office_13/Annotations 16 | Area_1/office_14/Annotations 17 | Area_1/office_15/Annotations 18 | Area_1/office_16/Annotations 19 | Area_1/office_17/Annotations 20 | Area_1/office_18/Annotations 21 | Area_1/office_19/Annotations 22 | Area_1/office_1/Annotations 23 | Area_1/office_20/Annotations 24 | Area_1/office_21/Annotations 25 | Area_1/office_22/Annotations 26 | Area_1/office_23/Annotations 27 | Area_1/office_24/Annotations 28 | Area_1/office_25/Annotations 29 | Area_1/office_26/Annotations 30 | Area_1/office_27/Annotations 31 | Area_1/office_28/Annotations 32 | Area_1/office_29/Annotations 33 | Area_1/office_2/Annotations 34 | Area_1/office_30/Annotations 35 | Area_1/office_31/Annotations 36 | Area_1/office_3/Annotations 37 | Area_1/office_4/Annotations 38 | Area_1/office_5/Annotations 39 | Area_1/office_6/Annotations 40 | Area_1/office_7/Annotations 41 | Area_1/office_8/Annotations 42 | Area_1/office_9/Annotations 43 | Area_1/pantry_1/Annotations 44 | Area_1/WC_1/Annotations 45 | Area_2/auditorium_1/Annotations 46 | Area_2/auditorium_2/Annotations 47 | Area_2/conferenceRoom_1/Annotations 48 | Area_2/hallway_10/Annotations 49 | Area_2/hallway_11/Annotations 50 | Area_2/hallway_12/Annotations 51 | Area_2/hallway_1/Annotations 52 | Area_2/hallway_2/Annotations 53 | Area_2/hallway_3/Annotations 54 | Area_2/hallway_4/Annotations 55 | Area_2/hallway_5/Annotations 56 | Area_2/hallway_6/Annotations 57 | Area_2/hallway_7/Annotations 58 | Area_2/hallway_8/Annotations 59 | Area_2/hallway_9/Annotations 60 | Area_2/office_10/Annotations 61 | Area_2/office_11/Annotations 62 | Area_2/office_12/Annotations 63 | Area_2/office_13/Annotations 64 | Area_2/office_14/Annotations 65 | Area_2/office_1/Annotations 66 | Area_2/office_2/Annotations 67 | Area_2/office_3/Annotations 68 | Area_2/office_4/Annotations 69 | Area_2/office_5/Annotations 70 | Area_2/office_6/Annotations 71 | Area_2/office_7/Annotations 72 | Area_2/office_8/Annotations 73 | Area_2/office_9/Annotations 74 | Area_2/storage_1/Annotations 75 | Area_2/storage_2/Annotations 76 | Area_2/storage_3/Annotations 77 | Area_2/storage_4/Annotations 78 | Area_2/storage_5/Annotations 79 | Area_2/storage_6/Annotations 80 | Area_2/storage_7/Annotations 81 | Area_2/storage_8/Annotations 82 | Area_2/storage_9/Annotations 83 | Area_2/WC_1/Annotations 84 | Area_2/WC_2/Annotations 85 | Area_3/conferenceRoom_1/Annotations 86 | Area_3/hallway_1/Annotations 87 | Area_3/hallway_2/Annotations 88 | Area_3/hallway_3/Annotations 89 | Area_3/hallway_4/Annotations 90 | Area_3/hallway_5/Annotations 91 | Area_3/hallway_6/Annotations 92 | Area_3/lounge_1/Annotations 93 | Area_3/lounge_2/Annotations 94 | Area_3/office_10/Annotations 95 | Area_3/office_1/Annotations 96 | Area_3/office_2/Annotations 97 | Area_3/office_3/Annotations 98 | Area_3/office_4/Annotations 99 | Area_3/office_5/Annotations 100 | Area_3/office_6/Annotations 101 | Area_3/office_7/Annotations 102 | Area_3/office_8/Annotations 103 | Area_3/office_9/Annotations 104 | Area_3/storage_1/Annotations 105 | Area_3/storage_2/Annotations 106 | Area_3/WC_1/Annotations 107 | Area_3/WC_2/Annotations 108 | Area_4/conferenceRoom_1/Annotations 109 | Area_4/conferenceRoom_2/Annotations 110 | Area_4/conferenceRoom_3/Annotations 111 | Area_4/hallway_10/Annotations 112 | Area_4/hallway_11/Annotations 113 | Area_4/hallway_12/Annotations 114 | Area_4/hallway_13/Annotations 115 | Area_4/hallway_14/Annotations 116 | Area_4/hallway_1/Annotations 117 | Area_4/hallway_2/Annotations 118 | Area_4/hallway_3/Annotations 119 | Area_4/hallway_4/Annotations 120 | Area_4/hallway_5/Annotations 121 | Area_4/hallway_6/Annotations 122 | Area_4/hallway_7/Annotations 123 | Area_4/hallway_8/Annotations 124 | Area_4/hallway_9/Annotations 125 | Area_4/lobby_1/Annotations 126 | Area_4/lobby_2/Annotations 127 | Area_4/office_10/Annotations 128 | Area_4/office_11/Annotations 129 | Area_4/office_12/Annotations 130 | Area_4/office_13/Annotations 131 | Area_4/office_14/Annotations 132 | Area_4/office_15/Annotations 133 | Area_4/office_16/Annotations 134 | Area_4/office_17/Annotations 135 | Area_4/office_18/Annotations 136 | Area_4/office_19/Annotations 137 | Area_4/office_1/Annotations 138 | Area_4/office_20/Annotations 139 | Area_4/office_21/Annotations 140 | Area_4/office_22/Annotations 141 | Area_4/office_2/Annotations 142 | Area_4/office_3/Annotations 143 | Area_4/office_4/Annotations 144 | Area_4/office_5/Annotations 145 | Area_4/office_6/Annotations 146 | Area_4/office_7/Annotations 147 | Area_4/office_8/Annotations 148 | Area_4/office_9/Annotations 149 | Area_4/storage_1/Annotations 150 | Area_4/storage_2/Annotations 151 | Area_4/storage_3/Annotations 152 | Area_4/storage_4/Annotations 153 | Area_4/WC_1/Annotations 154 | Area_4/WC_2/Annotations 155 | Area_4/WC_3/Annotations 156 | Area_4/WC_4/Annotations 157 | Area_5/conferenceRoom_1/Annotations 158 | Area_5/conferenceRoom_2/Annotations 159 | Area_5/conferenceRoom_3/Annotations 160 | Area_5/hallway_10/Annotations 161 | Area_5/hallway_11/Annotations 162 | Area_5/hallway_12/Annotations 163 | Area_5/hallway_13/Annotations 164 | Area_5/hallway_14/Annotations 165 | Area_5/hallway_15/Annotations 166 | Area_5/hallway_1/Annotations 167 | Area_5/hallway_2/Annotations 168 | Area_5/hallway_3/Annotations 169 | Area_5/hallway_4/Annotations 170 | Area_5/hallway_5/Annotations 171 | Area_5/hallway_6/Annotations 172 | Area_5/hallway_7/Annotations 173 | Area_5/hallway_8/Annotations 174 | Area_5/hallway_9/Annotations 175 | Area_5/lobby_1/Annotations 176 | Area_5/office_10/Annotations 177 | Area_5/office_11/Annotations 178 | Area_5/office_12/Annotations 179 | Area_5/office_13/Annotations 180 | Area_5/office_14/Annotations 181 | Area_5/office_15/Annotations 182 | Area_5/office_16/Annotations 183 | Area_5/office_17/Annotations 184 | Area_5/office_18/Annotations 185 | Area_5/office_19/Annotations 186 | Area_5/office_1/Annotations 187 | Area_5/office_20/Annotations 188 | Area_5/office_21/Annotations 189 | Area_5/office_22/Annotations 190 | Area_5/office_23/Annotations 191 | Area_5/office_24/Annotations 192 | Area_5/office_25/Annotations 193 | Area_5/office_26/Annotations 194 | Area_5/office_27/Annotations 195 | Area_5/office_28/Annotations 196 | Area_5/office_29/Annotations 197 | Area_5/office_2/Annotations 198 | Area_5/office_30/Annotations 199 | Area_5/office_31/Annotations 200 | Area_5/office_32/Annotations 201 | Area_5/office_33/Annotations 202 | Area_5/office_34/Annotations 203 | Area_5/office_35/Annotations 204 | Area_5/office_36/Annotations 205 | Area_5/office_37/Annotations 206 | Area_5/office_38/Annotations 207 | Area_5/office_39/Annotations 208 | Area_5/office_3/Annotations 209 | Area_5/office_40/Annotations 210 | Area_5/office_41/Annotations 211 | Area_5/office_42/Annotations 212 | Area_5/office_4/Annotations 213 | Area_5/office_5/Annotations 214 | Area_5/office_6/Annotations 215 | Area_5/office_7/Annotations 216 | Area_5/office_8/Annotations 217 | Area_5/office_9/Annotations 218 | Area_5/pantry_1/Annotations 219 | Area_5/storage_1/Annotations 220 | Area_5/storage_2/Annotations 221 | Area_5/storage_3/Annotations 222 | Area_5/storage_4/Annotations 223 | Area_5/WC_1/Annotations 224 | Area_5/WC_2/Annotations 225 | Area_6/conferenceRoom_1/Annotations 226 | Area_6/copyRoom_1/Annotations 227 | Area_6/hallway_1/Annotations 228 | Area_6/hallway_2/Annotations 229 | Area_6/hallway_3/Annotations 230 | Area_6/hallway_4/Annotations 231 | Area_6/hallway_5/Annotations 232 | Area_6/hallway_6/Annotations 233 | Area_6/lounge_1/Annotations 234 | Area_6/office_10/Annotations 235 | Area_6/office_11/Annotations 236 | Area_6/office_12/Annotations 237 | Area_6/office_13/Annotations 238 | Area_6/office_14/Annotations 239 | Area_6/office_15/Annotations 240 | Area_6/office_16/Annotations 241 | Area_6/office_17/Annotations 242 | Area_6/office_18/Annotations 243 | Area_6/office_19/Annotations 244 | Area_6/office_1/Annotations 245 | Area_6/office_20/Annotations 246 | Area_6/office_21/Annotations 247 | Area_6/office_22/Annotations 248 | Area_6/office_23/Annotations 249 | Area_6/office_24/Annotations 250 | Area_6/office_25/Annotations 251 | Area_6/office_26/Annotations 252 | Area_6/office_27/Annotations 253 | Area_6/office_28/Annotations 254 | Area_6/office_29/Annotations 255 | Area_6/office_2/Annotations 256 | Area_6/office_30/Annotations 257 | Area_6/office_31/Annotations 258 | Area_6/office_32/Annotations 259 | Area_6/office_33/Annotations 260 | Area_6/office_34/Annotations 261 | Area_6/office_35/Annotations 262 | Area_6/office_36/Annotations 263 | Area_6/office_37/Annotations 264 | Area_6/office_3/Annotations 265 | Area_6/office_4/Annotations 266 | Area_6/office_5/Annotations 267 | Area_6/office_6/Annotations 268 | Area_6/office_7/Annotations 269 | Area_6/office_8/Annotations 270 | Area_6/office_9/Annotations 271 | Area_6/openspace_1/Annotations 272 | Area_6/pantry_1/Annotations 273 | -------------------------------------------------------------------------------- /utils/meta/class_names.txt: -------------------------------------------------------------------------------- 1 | ceiling 2 | floor 3 | wall 4 | beam 5 | column 6 | window 7 | door 8 | table 9 | chair 10 | sofa 11 | bookcase 12 | board 13 | clutter 14 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/KDTreeTableAdaptor.h: -------------------------------------------------------------------------------- 1 | /*********************************************************************** 2 | * Software License Agreement (BSD License) 3 | * 4 | * Copyright 2011-16 Jose Luis Blanco (joseluisblancoc@gmail.com). 5 | * All rights reserved. 6 | * 7 | * Redistribution and use in source and binary forms, with or without 8 | * modification, are permitted provided that the following conditions 9 | * are met: 10 | * 11 | * 1. Redistributions of source code must retain the above copyright 12 | * notice, this list of conditions and the following disclaimer. 13 | * 2. Redistributions in binary form must reproduce the above copyright 14 | * notice, this list of conditions and the following disclaimer in the 15 | * documentation and/or other materials provided with the distribution. 16 | * 17 | * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR 18 | * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES 19 | * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. 20 | * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, 21 | * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT 22 | * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 23 | * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 24 | * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 25 | * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF 26 | * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 27 | *************************************************************************/ 28 | 29 | #pragma once 30 | 31 | #include "nanoflann.hpp" 32 | 33 | // #include 34 | 35 | // ===== This example shows how to use nanoflann with these types of containers: ======= 36 | //typedef std::vector > my_vector_of_vectors_t; 37 | //typedef std::vector my_vector_of_vectors_t; // This requires #include 38 | // ===================================================================================== 39 | 40 | 41 | /** A simple vector-of-vectors adaptor for nanoflann, without duplicating the storage. 42 | * The i'th vector represents a point in the state space. 43 | * 44 | * \tparam DIM If set to >0, it specifies a compile-time fixed dimensionality for the points in the data set, allowing more compiler optimizations. 45 | * \tparam num_t The type of the point coordinates (typically, double or float). 46 | * \tparam Distance The distance metric to use: nanoflann::metric_L1, nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. 47 | * \tparam IndexType The type for indices in the KD-tree index (typically, size_t of int) 48 | */ 49 | // template 50 | // struct KDTreeVectorAdaptor 51 | // { 52 | // typedef KDTreeVectorAdaptor self_t; 53 | // typedef typename Distance::template traits::distance_t metric_t; 54 | // typedef nanoflann::KDTreeSingleIndexAdaptor< metric_t,self_t,DIM,IndexType> index_t; 55 | 56 | // index_t* index; //! The kd-tree index for the user to call its methods as usual with any other FLANN index. 57 | // size_t dims; 58 | 59 | // /// Constructor: takes a const ref to the vector of vectors object with the data points 60 | // KDTreeVectorAdaptor(const size_t dims /* dimensionality */, const VectorType &mat, const int leaf_max_size = 10) : m_data(mat) 61 | // { 62 | // assert(mat.size() != 0); 63 | // this->dims= dims; 64 | // index = new index_t( static_cast(dims), *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) ); 65 | // index->buildIndex(); 66 | // } 67 | 68 | // ~KDTreeVectorAdaptor() { 69 | // delete index; 70 | // } 71 | 72 | // const VectorType &m_data; 73 | 74 | // /** Query for the \a num_closest closest points to a given point (entered as query_point[0:dim-1]). 75 | // * Note that this is a short-cut method for index->findNeighbors(). 76 | // * The user can also call index->... methods as desired. 77 | // * \note nChecks_IGNORED is ignored but kept for compatibility with the original FLANN interface. 78 | // */ 79 | // inline void query(const num_t *query_point, const size_t num_closest, IndexType *out_indices, num_t *out_distances_sq, const int nChecks_IGNORED = 10) const 80 | // { 81 | // nanoflann::KNNResultSet resultSet(num_closest); 82 | // resultSet.init(out_indices, out_distances_sq); 83 | // index->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); 84 | // } 85 | 86 | // /** @name Interface expected by KDTreeSingleIndexAdaptor 87 | // * @{ */ 88 | 89 | // const self_t & derived() const { 90 | // return *this; 91 | // } 92 | // self_t & derived() { 93 | // return *this; 94 | // } 95 | 96 | // // Must return the number of data points 97 | // inline size_t kdtree_get_point_count() const { 98 | // return m_data.size()/this->dims; 99 | // } 100 | 101 | // // Returns the dim'th component of the idx'th point in the class: 102 | // inline num_t kdtree_get_pt(const size_t idx, const size_t dim) const { 103 | // return m_data[idx*this->dims + dim]; 104 | // } 105 | 106 | // // Optional bounding-box computation: return false to default to a standard bbox computation loop. 107 | // // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 108 | // // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 109 | // template 110 | // bool kdtree_get_bbox(BBOX & /*bb*/) const { 111 | // return false; 112 | // } 113 | 114 | // /** @} */ 115 | 116 | // }; // end of KDTreeVectorOfVectorsAdaptor 117 | 118 | 119 | 120 | 121 | template 122 | struct KDTreeTableAdaptor 123 | { 124 | typedef KDTreeTableAdaptor self_t; 125 | typedef typename Distance::template traits::distance_t metric_t; 126 | typedef nanoflann::KDTreeSingleIndexAdaptor< metric_t,self_t,DIM,IndexType> index_t; 127 | 128 | index_t* index; //! The kd-tree index for the user to call its methods as usual with any other FLANN index. 129 | size_t dim; 130 | size_t npts; 131 | const TableType* m_data; 132 | 133 | /// Constructor: takes a const ref to the vector of vectors object with the data points 134 | KDTreeTableAdaptor(const size_t npts, const size_t dim, const TableType* mat, const int leaf_max_size = 10) : m_data(mat), dim(dim), npts(npts) 135 | { 136 | assert(npts != 0); 137 | index = new index_t( static_cast(dim), *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) ); 138 | index->buildIndex(); 139 | } 140 | 141 | ~KDTreeTableAdaptor() { 142 | delete index; 143 | } 144 | 145 | 146 | /** Query for the \a num_closest closest points to a given point (entered as query_point[0:dim-1]). 147 | * Note that this is a short-cut method for index->findNeighbors(). 148 | * The user can also call index->... methods as desired. 149 | * \note nChecks_IGNORED is ignored but kept for compatibility with the original FLANN interface. 150 | */ 151 | inline void query(const num_t *query_point, const size_t num_closest, IndexType *out_indices, num_t *out_distances_sq, const int nChecks_IGNORED = 10) const 152 | { 153 | nanoflann::KNNResultSet resultSet(num_closest); 154 | resultSet.init(out_indices, out_distances_sq); 155 | index->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); 156 | } 157 | 158 | /** @name Interface expected by KDTreeSingleIndexAdaptor 159 | * @{ */ 160 | 161 | const self_t & derived() const { 162 | return *this; 163 | } 164 | self_t & derived() { 165 | return *this; 166 | } 167 | 168 | // Must return the number of data points 169 | inline size_t kdtree_get_point_count() const { 170 | return this->npts; 171 | } 172 | 173 | // Returns the dim'th component of the idx'th point in the class: 174 | inline num_t kdtree_get_pt(const size_t pts_id, const size_t coord_id) const { 175 | return m_data[pts_id*this->dim + coord_id]; 176 | } 177 | 178 | // Optional bounding-box computation: return false to default to a standard bbox computation loop. 179 | // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 180 | // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 181 | template 182 | bool kdtree_get_bbox(BBOX & /*bb*/) const { 183 | return false; 184 | } 185 | 186 | /** @} */ 187 | 188 | }; // end of KDTreeVectorOfVectorsAdaptor 189 | 190 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/knn.pyx: -------------------------------------------------------------------------------- 1 | # distutils: language = c++ 2 | # distutils: sources = knn.cxx 3 | 4 | import numpy as np 5 | cimport numpy as np 6 | import cython 7 | 8 | cdef extern from "knn_.h": 9 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 10 | const float* queries, const size_t nqueries, 11 | const size_t K, long* indices) 12 | 13 | void cpp_knn_omp(const float* points, const size_t npts, const size_t dim, 14 | const float* queries, const size_t nqueries, 15 | const size_t K, long* indices) 16 | 17 | void cpp_knn_batch(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 18 | const float* queries, const size_t nqueries, 19 | const size_t K, long* batch_indices) 20 | 21 | void cpp_knn_batch_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 22 | const float* queries, const size_t nqueries, 23 | const size_t K, long* batch_indices) 24 | 25 | void cpp_knn_batch_distance_pick(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 26 | float* queries, const size_t nqueries, 27 | const size_t K, long* batch_indices) 28 | 29 | void cpp_knn_batch_distance_pick_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 30 | float* batch_queries, const size_t nqueries, 31 | const size_t K, long* batch_indices) 32 | 33 | def knn(pts, queries, K, omp=False): 34 | 35 | # define shape parameters 36 | cdef int npts 37 | cdef int dim 38 | cdef int K_cpp 39 | cdef int nqueries 40 | 41 | # define tables 42 | cdef np.ndarray[np.float32_t, ndim=2] pts_cpp 43 | cdef np.ndarray[np.float32_t, ndim=2] queries_cpp 44 | cdef np.ndarray[np.int64_t, ndim=2] indices_cpp 45 | 46 | # set shape values 47 | npts = pts.shape[0] 48 | nqueries = queries.shape[0] 49 | dim = pts.shape[1] 50 | K_cpp = K 51 | 52 | # create indices tensor 53 | indices = np.zeros((queries.shape[0], K), dtype=np.int64) 54 | 55 | pts_cpp = np.ascontiguousarray(pts, dtype=np.float32) 56 | queries_cpp = np.ascontiguousarray(queries, dtype=np.float32) 57 | indices_cpp = indices 58 | 59 | # normal estimation 60 | if omp: 61 | cpp_knn_omp( pts_cpp.data, npts, dim, 62 | queries_cpp.data, nqueries, 63 | K_cpp, indices_cpp.data) 64 | else: 65 | cpp_knn( pts_cpp.data, npts, dim, 66 | queries_cpp.data, nqueries, 67 | K_cpp, indices_cpp.data) 68 | 69 | return indices 70 | 71 | def knn_batch(pts, queries, K, omp=False): 72 | 73 | # define shape parameters 74 | cdef int batch_size 75 | cdef int npts 76 | cdef int nqueries 77 | cdef int K_cpp 78 | cdef int dim 79 | 80 | # define tables 81 | cdef np.ndarray[np.float32_t, ndim=3] pts_cpp 82 | cdef np.ndarray[np.float32_t, ndim=3] queries_cpp 83 | cdef np.ndarray[np.int64_t, ndim=3] indices_cpp 84 | 85 | # set shape values 86 | batch_size = pts.shape[0] 87 | npts = pts.shape[1] 88 | dim = pts.shape[2] 89 | nqueries = queries.shape[1] 90 | K_cpp = K 91 | 92 | # create indices tensor 93 | indices = np.zeros((pts.shape[0], queries.shape[1], K), dtype=np.int64) 94 | 95 | pts_cpp = np.ascontiguousarray(pts, dtype=np.float32) 96 | queries_cpp = np.ascontiguousarray(queries, dtype=np.float32) 97 | indices_cpp = indices 98 | 99 | # normal estimation 100 | if omp: 101 | cpp_knn_batch_omp( pts_cpp.data, batch_size, npts, dim, 102 | queries_cpp.data, nqueries, 103 | K_cpp, indices_cpp.data) 104 | else: 105 | cpp_knn_batch( pts_cpp.data, batch_size, npts, dim, 106 | queries_cpp.data, nqueries, 107 | K_cpp, indices_cpp.data) 108 | 109 | return indices 110 | 111 | def knn_batch_distance_pick(pts, nqueries, K, omp=False): 112 | 113 | # define shape parameters 114 | cdef int batch_size 115 | cdef int npts 116 | cdef int nqueries_cpp 117 | cdef int K_cpp 118 | cdef int dim 119 | 120 | # define tables 121 | cdef np.ndarray[np.float32_t, ndim=3] pts_cpp 122 | cdef np.ndarray[np.float32_t, ndim=3] queries_cpp 123 | cdef np.ndarray[np.int64_t, ndim=3] indices_cpp 124 | 125 | # set shape values 126 | batch_size = pts.shape[0] 127 | npts = pts.shape[1] 128 | dim = pts.shape[2] 129 | nqueries_cpp = nqueries 130 | K_cpp = K 131 | 132 | # create indices tensor 133 | indices = np.zeros((pts.shape[0], nqueries, K), dtype=np.long) 134 | queries = np.zeros((pts.shape[0], nqueries, dim), dtype=np.float32) 135 | 136 | pts_cpp = np.ascontiguousarray(pts, dtype=np.float32) 137 | queries_cpp = np.ascontiguousarray(queries, dtype=np.float32) 138 | indices_cpp = indices 139 | 140 | if omp: 141 | cpp_knn_batch_distance_pick_omp( pts_cpp.data, batch_size, npts, dim, 142 | queries_cpp.data, nqueries, 143 | K_cpp, indices_cpp.data) 144 | else: 145 | cpp_knn_batch_distance_pick( pts_cpp.data, batch_size, npts, dim, 146 | queries_cpp.data, nqueries, 147 | K_cpp, indices_cpp.data) 148 | 149 | return indices, queries -------------------------------------------------------------------------------- /utils/nearest_neighbors/knn_.cxx: -------------------------------------------------------------------------------- 1 | 2 | #include "knn_.h" 3 | #include "nanoflann.hpp" 4 | using namespace nanoflann; 5 | 6 | #include "KDTreeTableAdaptor.h" 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | #include 15 | #include 16 | #include 17 | 18 | using namespace std; 19 | 20 | 21 | 22 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 23 | const float* queries, const size_t nqueries, 24 | const size_t K, long* indices){ 25 | 26 | // create the kdtree 27 | typedef KDTreeTableAdaptor< float, float> KDTree; 28 | KDTree mat_index(npts, dim, points, 10); 29 | mat_index.index->buildIndex(); 30 | 31 | std::vector out_dists_sqr(K); 32 | std::vector out_ids(K); 33 | 34 | // iterate over the points 35 | for(size_t i=0; i resultSet(K); 38 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 39 | mat_index.index->findNeighbors(resultSet, &queries[i*dim], nanoflann::SearchParams(10)); 40 | for(size_t j=0; j KDTree; 52 | KDTree mat_index(npts, dim, points, 10); 53 | mat_index.index->buildIndex(); 54 | 55 | 56 | // iterate over the points 57 | # pragma omp parallel for 58 | for(size_t i=0; i out_ids(K); 60 | std::vector out_dists_sqr(K); 61 | 62 | nanoflann::KNNResultSet resultSet(K); 63 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 64 | mat_index.index->findNeighbors(resultSet, &queries[i*dim], nanoflann::SearchParams(10)); 65 | for(size_t j=0; j KDTree; 83 | KDTree mat_index(npts, dim, points, 10); 84 | 85 | mat_index.index->buildIndex(); 86 | 87 | std::vector out_dists_sqr(K); 88 | std::vector out_ids(K); 89 | 90 | // iterate over the points 91 | for(size_t i=0; i resultSet(K); 93 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 94 | mat_index.index->findNeighbors(resultSet, &queries[bid*nqueries*dim + i*dim], nanoflann::SearchParams(10)); 95 | for(size_t j=0; j KDTree; 116 | KDTree mat_index(npts, dim, points, 10); 117 | 118 | mat_index.index->buildIndex(); 119 | 120 | std::vector out_dists_sqr(K); 121 | std::vector out_ids(K); 122 | 123 | // iterate over the points 124 | for(size_t i=0; i resultSet(K); 126 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 127 | mat_index.index->findNeighbors(resultSet, &queries[bid*nqueries*dim + i*dim], nanoflann::SearchParams(10)); 128 | for(size_t j=0; j KDTree; 153 | KDTree tree(npts, dim, points, 10); 154 | tree.index->buildIndex(); 155 | 156 | vector used(npts, 0); 157 | int current_id = 0; 158 | for(size_t ptid=0; ptid possible_ids; 162 | while(possible_ids.size() == 0){ 163 | for(size_t i=0; i query(3); 178 | for(size_t i=0; i dists(K); 183 | std::vector ids(K); 184 | nanoflann::KNNResultSet resultSet(K); 185 | resultSet.init(&ids[0], &dists[0] ); 186 | tree.index->findNeighbors(resultSet, &query[0], nanoflann::SearchParams(10)); 187 | 188 | for(size_t i=0; i KDTree; 221 | KDTree tree(npts, dim, points, 10); 222 | tree.index->buildIndex(); 223 | 224 | vector used(npts, 0); 225 | int current_id = 0; 226 | for(size_t ptid=0; ptid possible_ids; 230 | while(possible_ids.size() == 0){ 231 | for(size_t i=0; i query(3); 246 | for(size_t i=0; i dists(K); 251 | std::vector ids(K); 252 | nanoflann::KNNResultSet resultSet(K); 253 | resultSet.init(&ids[0], &dists[0] ); 254 | tree.index->findNeighbors(resultSet, &query[0], nanoflann::SearchParams(10)); 255 | 256 | for(size_t i=0; i 4 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 5 | const float* queries, const size_t nqueries, 6 | const size_t K, long* indices); 7 | 8 | void cpp_knn_omp(const float* points, const size_t npts, const size_t dim, 9 | const float* queries, const size_t nqueries, 10 | const size_t K, long* indices); 11 | 12 | 13 | void cpp_knn_batch(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 14 | const float* queries, const size_t nqueries, 15 | const size_t K, long* batch_indices); 16 | 17 | void cpp_knn_batch_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 18 | const float* queries, const size_t nqueries, 19 | const size_t K, long* batch_indices); 20 | 21 | void cpp_knn_batch_distance_pick(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 22 | float* queries, const size_t nqueries, 23 | const size_t K, long* batch_indices); 24 | 25 | void cpp_knn_batch_distance_pick_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 26 | float* batch_queries, const size_t nqueries, 27 | const size_t K, long* batch_indices); -------------------------------------------------------------------------------- /utils/nearest_neighbors/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup 2 | from distutils.extension import Extension 3 | from Cython.Distutils import build_ext 4 | import numpy 5 | 6 | 7 | 8 | ext_modules = [Extension( 9 | "nearest_neighbors", 10 | sources=["knn.pyx", "knn_.cxx",], # source file(s) 11 | include_dirs=["./", numpy.get_include()], 12 | language="c++", 13 | extra_compile_args = [ "-std=c++11", "-fopenmp",], 14 | extra_link_args=["-std=c++11", '-fopenmp'], 15 | )] 16 | 17 | setup( 18 | name = "KNN NanoFLANN", 19 | ext_modules = ext_modules, 20 | cmdclass = {'build_ext': build_ext}, 21 | ) 22 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import lib.python.nearest_neighbors as nearest_neighbors 3 | import time 4 | 5 | batch_size = 16 6 | num_points = 81920 7 | K = 16 8 | pc = np.random.rand(batch_size, num_points, 3).astype(np.float32) 9 | 10 | # nearest neighbours 11 | start = time.time() 12 | neigh_idx = nearest_neighbors.knn_batch(pc, pc, K, omp=True) 13 | print(time.time() - start) 14 | 15 | 16 | --------------------------------------------------------------------------------