├── .gitattributes ├── .gitignore ├── LICENSE ├── README.md ├── RandLANet.py ├── compile_op.sh ├── helper_ply.py ├── helper_requirements.txt ├── helper_tf_util.py ├── helper_tool.py ├── imgs ├── S3DIS_area2.gif ├── S3DIS_area3.gif ├── Semantic3D-1.gif ├── Semantic3D-3.gif ├── Semantic3D-4.gif └── SemanticKITTI-2.gif ├── jobs_6_fold_cv_s3dis.sh ├── jobs_test_semantickitti.sh ├── main_S3DIS.py ├── main_Semantic3D.py ├── main_SemanticKITTI.py ├── tester_S3DIS.py ├── tester_Semantic3D.py ├── tester_SemanticKITTI.py └── utils ├── 6_fold_cv.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_semantic3d.py ├── data_prepare_semantickitti.py ├── download_semantic3d.sh ├── meta ├── anno_paths.txt └── class_names.txt ├── nearest_neighbors ├── KDTreeTableAdaptor.h ├── knn.pyx ├── knn_.cxx ├── knn_.h ├── nanoflann.hpp ├── setup.py └── test.py └── semantic-kitti.yaml /.gitattributes: -------------------------------------------------------------------------------- 1 | *.sh linguist-language=python 2 | *.h linguist-language=python 3 | *.hpp linguist-language=python 4 | 5 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | # Attribution-NonCommercial-ShareAlike 4.0 International 3 | 4 | Creative Commons Corporation (“Creative Commons”) is not a law firm and does not provide legal services or legal advice. Distribution of Creative Commons public licenses does not create a lawyer-client or other relationship. Creative Commons makes its licenses and related information available on an “as-is” basis. Creative Commons gives no warranties regarding its licenses, any material licensed under their terms and conditions, or any related information. Creative Commons disclaims all liability for damages resulting from their use to the fullest extent possible. 5 | 6 | ### Using Creative Commons Public Licenses 7 | 8 | Creative Commons public licenses provide a standard set of terms and conditions that creators and other rights holders may use to share original works of authorship and other material subject to copyright and certain other rights specified in the public license below. The following considerations are for informational purposes only, are not exhaustive, and do not form part of our licenses. 9 | 10 | * __Considerations for licensors:__ Our public licenses are intended for use by those authorized to give the public permission to use material in ways otherwise restricted by copyright and certain other rights. Our licenses are irrevocable. Licensors should read and understand the terms and conditions of the license they choose before applying it. Licensors should also secure all rights necessary before applying our licenses so that the public can reuse the material as expected. Licensors should clearly mark any material not subject to the license. This includes other CC-licensed material, or material used under an exception or limitation to copyright. [More considerations for licensors](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensors). 11 | 12 | * __Considerations for the public:__ By using one of our public licenses, a licensor grants the public permission to use the licensed material under specified terms and conditions. If the licensor’s permission is not necessary for any reason–for example, because of any applicable exception or limitation to copyright–then that use is not regulated by the license. Our licenses grant only permissions under copyright and certain other rights that a licensor has authority to grant. Use of the licensed material may still be restricted for other reasons, including because others have copyright or other rights in the material. A licensor may make special requests, such as asking that all changes be marked or described. Although not required by our licenses, you are encouraged to respect those requests where reasonable. [More considerations for the public](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensees). 13 | 14 | ## Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License 15 | 16 | By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License ("Public License"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions. 17 | 18 | ### Section 1 – Definitions. 19 | 20 | a. __Adapted Material__ means material subject to Copyright and Similar Rights that is derived from or based upon the Licensed Material and in which the Licensed Material is translated, altered, arranged, transformed, or otherwise modified in a manner requiring permission under the Copyright and Similar Rights held by the Licensor. For purposes of this Public License, where the Licensed Material is a musical work, performance, or sound recording, Adapted Material is always produced where the Licensed Material is synched in timed relation with a moving image. 21 | 22 | b. __Adapter's License__ means the license You apply to Your Copyright and Similar Rights in Your contributions to Adapted Material in accordance with the terms and conditions of this Public License. 23 | 24 | c. __BY-NC-SA Compatible License__ means a license listed at [creativecommons.org/compatiblelicenses](http://creativecommons.org/compatiblelicenses), approved by Creative Commons as essentially the equivalent of this Public License. 25 | 26 | d. __Copyright and Similar Rights__ means copyright and/or similar rights closely related to copyright including, without limitation, performance, broadcast, sound recording, and Sui Generis Database Rights, without regard to how the rights are labeled or categorized. For purposes of this Public License, the rights specified in Section 2(b)(1)-(2) are not Copyright and Similar Rights. 27 | 28 | e. __Effective Technological Measures__ means those measures that, in the absence of proper authority, may not be circumvented under laws fulfilling obligations under Article 11 of the WIPO Copyright Treaty adopted on December 20, 1996, and/or similar international agreements. 29 | 30 | f. __Exceptions and Limitations__ means fair use, fair dealing, and/or any other exception or limitation to Copyright and Similar Rights that applies to Your use of the Licensed Material. 31 | 32 | g. __License Elements__ means the license attributes listed in the name of a Creative Commons Public License. The License Elements of this Public License are Attribution, NonCommercial, and ShareAlike. 33 | 34 | h. __Licensed Material__ means the artistic or literary work, database, or other material to which the Licensor applied this Public License. 35 | 36 | i. __Licensed Rights__ means the rights granted to You subject to the terms and conditions of this Public License, which are limited to all Copyright and Similar Rights that apply to Your use of the Licensed Material and that the Licensor has authority to license. 37 | 38 | j. __Licensor__ means the individual(s) or entity(ies) granting rights under this Public License. 39 | 40 | k. __NonCommercial__ means not primarily intended for or directed towards commercial advantage or monetary compensation. For purposes of this Public License, the exchange of the Licensed Material for other material subject to Copyright and Similar Rights by digital file-sharing or similar means is NonCommercial provided there is no payment of monetary compensation in connection with the exchange. 41 | 42 | l. __Share__ means to provide material to the public by any means or process that requires permission under the Licensed Rights, such as reproduction, public display, public performance, distribution, dissemination, communication, or importation, and to make material available to the public including in ways that members of the public may access the material from a place and at a time individually chosen by them. 43 | 44 | m. __Sui Generis Database Rights__ means rights other than copyright resulting from Directive 96/9/EC of the European Parliament and of the Council of 11 March 1996 on the legal protection of databases, as amended and/or succeeded, as well as other essentially equivalent rights anywhere in the world. 45 | 46 | n. __You__ means the individual or entity exercising the Licensed Rights under this Public License. Your has a corresponding meaning. 47 | 48 | ### Section 2 – Scope. 49 | 50 | a. ___License grant.___ 51 | 52 | 1. Subject to the terms and conditions of this Public License, the Licensor hereby grants You a worldwide, royalty-free, non-sublicensable, non-exclusive, irrevocable license to exercise the Licensed Rights in the Licensed Material to: 53 | 54 | A. reproduce and Share the Licensed Material, in whole or in part, for NonCommercial purposes only; and 55 | 56 | B. produce, reproduce, and Share Adapted Material for NonCommercial purposes only. 57 | 58 | 2. __Exceptions and Limitations.__ For the avoidance of doubt, where Exceptions and Limitations apply to Your use, this Public License does not apply, and You do not need to comply with its terms and conditions. 59 | 60 | 3. __Term.__ The term of this Public License is specified in Section 6(a). 61 | 62 | 4. __Media and formats; technical modifications allowed.__ The Licensor authorizes You to exercise the Licensed Rights in all media and formats whether now known or hereafter created, and to make technical modifications necessary to do so. The Licensor waives and/or agrees not to assert any right or authority to forbid You from making technical modifications necessary to exercise the Licensed Rights, including technical modifications necessary to circumvent Effective Technological Measures. For purposes of this Public License, simply making modifications authorized by this Section 2(a)(4) never produces Adapted Material. 63 | 64 | 5. __Downstream recipients.__ 65 | 66 | A. __Offer from the Licensor – Licensed Material.__ Every recipient of the Licensed Material automatically receives an offer from the Licensor to exercise the Licensed Rights under the terms and conditions of this Public License. 67 | 68 | B. __Additional offer from the Licensor – Adapted Material.__ Every recipient of Adapted Material from You automatically receives an offer from the Licensor to exercise the Licensed Rights in the Adapted Material under the conditions of the Adapter’s License You apply. 69 | 70 | C. __No downstream restrictions.__ You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, the Licensed Material if doing so restricts exercise of the Licensed Rights by any recipient of the Licensed Material. 71 | 72 | 6. __No endorsement.__ Nothing in this Public License constitutes or may be construed as permission to assert or imply that You are, or that Your use of the Licensed Material is, connected with, or sponsored, endorsed, or granted official status by, the Licensor or others designated to receive attribution as provided in Section 3(a)(1)(A)(i). 73 | 74 | b. ___Other rights.___ 75 | 76 | 1. Moral rights, such as the right of integrity, are not licensed under this Public License, nor are publicity, privacy, and/or other similar personality rights; however, to the extent possible, the Licensor waives and/or agrees not to assert any such rights held by the Licensor to the limited extent necessary to allow You to exercise the Licensed Rights, but not otherwise. 77 | 78 | 2. Patent and trademark rights are not licensed under this Public License. 79 | 80 | 3. To the extent possible, the Licensor waives any right to collect royalties from You for the exercise of the Licensed Rights, whether directly or through a collecting society under any voluntary or waivable statutory or compulsory licensing scheme. In all other cases the Licensor expressly reserves any right to collect such royalties, including when the Licensed Material is used other than for NonCommercial purposes. 81 | 82 | ### Section 3 – License Conditions. 83 | 84 | Your exercise of the Licensed Rights is expressly made subject to the following conditions. 85 | 86 | a. ___Attribution.___ 87 | 88 | 1. If You Share the Licensed Material (including in modified form), You must: 89 | 90 | A. retain the following if it is supplied by the Licensor with the Licensed Material: 91 | 92 | i. identification of the creator(s) of the Licensed Material and any others designated to receive attribution, in any reasonable manner requested by the Licensor (including by pseudonym if designated); 93 | 94 | ii. a copyright notice; 95 | 96 | iii. a notice that refers to this Public License; 97 | 98 | iv. a notice that refers to the disclaimer of warranties; 99 | 100 | v. a URI or hyperlink to the Licensed Material to the extent reasonably practicable; 101 | 102 | B. indicate if You modified the Licensed Material and retain an indication of any previous modifications; and 103 | 104 | C. indicate the Licensed Material is licensed under this Public License, and include the text of, or the URI or hyperlink to, this Public License. 105 | 106 | 2. You may satisfy the conditions in Section 3(a)(1) in any reasonable manner based on the medium, means, and context in which You Share the Licensed Material. For example, it may be reasonable to satisfy the conditions by providing a URI or hyperlink to a resource that includes the required information. 107 | 108 | 3. If requested by the Licensor, You must remove any of the information required by Section 3(a)(1)(A) to the extent reasonably practicable. 109 | 110 | b. ___ShareAlike.___ 111 | 112 | In addition to the conditions in Section 3(a), if You Share Adapted Material You produce, the following conditions also apply. 113 | 114 | 1. The Adapter’s License You apply must be a Creative Commons license with the same License Elements, this version or later, or a BY-NC-SA Compatible License. 115 | 116 | 2. You must include the text of, or the URI or hyperlink to, the Adapter's License You apply. You may satisfy this condition in any reasonable manner based on the medium, means, and context in which You Share Adapted Material. 117 | 118 | 3. You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, Adapted Material that restrict exercise of the rights granted under the Adapter's License You apply. 119 | 120 | ### Section 4 – Sui Generis Database Rights. 121 | 122 | Where the Licensed Rights include Sui Generis Database Rights that apply to Your use of the Licensed Material: 123 | 124 | a. for the avoidance of doubt, Section 2(a)(1) grants You the right to extract, reuse, reproduce, and Share all or a substantial portion of the contents of the database for NonCommercial purposes only; 125 | 126 | b. if You include all or a substantial portion of the database contents in a database in which You have Sui Generis Database Rights, then the database in which You have Sui Generis Database Rights (but not its individual contents) is Adapted Material, including for purposes of Section 3(b); and 127 | 128 | c. You must comply with the conditions in Section 3(a) if You Share all or a substantial portion of the contents of the database. 129 | 130 | For the avoidance of doubt, this Section 4 supplements and does not replace Your obligations under this Public License where the Licensed Rights include other Copyright and Similar Rights. 131 | 132 | ### Section 5 – Disclaimer of Warranties and Limitation of Liability. 133 | 134 | a. __Unless otherwise separately undertaken by the Licensor, to the extent possible, the Licensor offers the Licensed Material as-is and as-available, and makes no representations or warranties of any kind concerning the Licensed Material, whether express, implied, statutory, or other. This includes, without limitation, warranties of title, merchantability, fitness for a particular purpose, non-infringement, absence of latent or other defects, accuracy, or the presence or absence of errors, whether or not known or discoverable. Where disclaimers of warranties are not allowed in full or in part, this disclaimer may not apply to You.__ 135 | 136 | b. __To the extent possible, in no event will the Licensor be liable to You on any legal theory (including, without limitation, negligence) or otherwise for any direct, special, indirect, incidental, consequential, punitive, exemplary, or other losses, costs, expenses, or damages arising out of this Public License or use of the Licensed Material, even if the Licensor has been advised of the possibility of such losses, costs, expenses, or damages. Where a limitation of liability is not allowed in full or in part, this limitation may not apply to You.__ 137 | 138 | c. The disclaimer of warranties and limitation of liability provided above shall be interpreted in a manner that, to the extent possible, most closely approximates an absolute disclaimer and waiver of all liability. 139 | 140 | ### Section 6 – Term and Termination. 141 | 142 | a. This Public License applies for the term of the Copyright and Similar Rights licensed here. However, if You fail to comply with this Public License, then Your rights under this Public License terminate automatically. 143 | 144 | b. Where Your right to use the Licensed Material has terminated under Section 6(a), it reinstates: 145 | 146 | 1. automatically as of the date the violation is cured, provided it is cured within 30 days of Your discovery of the violation; or 147 | 148 | 2. automatically as of the date the violation is cured, provided it is cured within 30 days of Your discovery of the violation; or 149 | 150 | For the avoidance of doubt, this Section 6(b) does not affect any right the Licensor may have to seek remedies for Your violations of this Public License. 151 | 152 | c. For the avoidance of doubt, the Licensor may also offer the Licensed Material under separate terms or conditions or stop distributing the Licensed Material at any time; however, doing so will not terminate this Public License. 153 | 154 | d. Sections 1, 5, 6, 7, and 8 survive termination of this Public License. 155 | 156 | ### Section 7 – Other Terms and Conditions. 157 | 158 | a. The Licensor shall not be bound by any additional or different terms or conditions communicated by You unless expressly agreed. 159 | 160 | b. Any arrangements, understandings, or agreements regarding the Licensed Material not stated herein are separate from and independent of the terms and conditions of this Public License. 161 | 162 | ### Section 8 – Interpretation. 163 | 164 | a. For the avoidance of doubt, this Public License does not, and shall not be interpreted to, reduce, limit, restrict, or impose conditions on any use of the Licensed Material that could lawfully be made without permission under this Public License. 165 | 166 | b. To the extent possible, if any provision of this Public License is deemed unenforceable, it shall be automatically reformed to the minimum extent necessary to make it enforceable. If the provision cannot be reformed, it shall be severed from this Public License without affecting the enforceability of the remaining terms and conditions. 167 | 168 | c. No term or condition of this Public License will be waived and no failure to comply consented to unless expressly agreed to by the Licensor. 169 | 170 | d. Nothing in this Public License constitutes or may be interpreted as a limitation upon, or waiver of, any privileges and immunities that apply to the Licensor or You, including from the legal processes of any jurisdiction or authority. 171 | 172 | ``` 173 | Creative Commons is not a party to its public licenses. Notwithstanding, Creative Commons may elect to apply one of its public licenses to material it publishes and in those instances will be considered the “Licensor.” Except for the limited purpose of indicating that material is shared under a Creative Commons public license or as otherwise permitted by the Creative Commons policies published at [creativecommons.org/policies](http://creativecommons.org/policies), Creative Commons does not authorize the use of the trademark “Creative Commons” or any other trademark or logo of Creative Commons without its prior written consent including, without limitation, in connection with any unauthorized modifications to any of its public licenses or any other arrangements, understandings, or agreements concerning use of licensed material. For the avoidance of doubt, this paragraph does not form part of the public licenses. 174 | 175 | Creative Commons may be contacted at [creativecommons.org](http://creativecommons.org/). 176 | ``` 177 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/191111236/semantic-segmentation-on-semantic3d)](https://paperswithcode.com/sota/semantic-segmentation-on-semantic3d?p=191111236) 2 | [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/191111236/3d-semantic-segmentation-on-semantickitti)](https://paperswithcode.com/sota/3d-semantic-segmentation-on-semantickitti?p=191111236) 3 | [![License CC BY-NC-SA 4.0](https://img.shields.io/badge/license-CC4.0-blue.svg)](https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode) 4 | 5 | # RandLA-Net: Efficient Semantic Segmentation of Large-Scale Point Clouds (CVPR 2020) 6 | 7 | This is the official implementation of **RandLA-Net** (CVPR2020, Oral presentation), a simple and efficient neural architecture for semantic segmentation of large-scale 3D point clouds. For technical details, please refer to: 8 | 9 | **RandLA-Net: Efficient Semantic Segmentation of Large-Scale Point Clouds**
10 | [Qingyong Hu](https://www.cs.ox.ac.uk/people/qingyong.hu/), [Bo Yang*](https://yang7879.github.io/), [Linhai Xie](https://www.cs.ox.ac.uk/people/linhai.xie/), [Stefano Rosa](https://www.cs.ox.ac.uk/people/stefano.rosa/), [Yulan Guo](http://yulanguo.me/), [Zhihua Wang](https://www.cs.ox.ac.uk/people/zhihua.wang/), [Niki Trigoni](https://www.cs.ox.ac.uk/people/niki.trigoni/), [Andrew Markham](https://www.cs.ox.ac.uk/people/andrew.markham/).
11 | **[[Paper](https://arxiv.org/abs/1911.11236)] [[Video](https://youtu.be/Ar3eY_lwzMk)] [[Blog](https://zhuanlan.zhihu.com/p/105433460)] [[Project page](http://randla-net.cs.ox.ac.uk/)]**
12 | 13 | 14 |

15 | 16 | 17 | 18 | ### (1) Setup 19 | This code has been tested with Python 3.5, Tensorflow 1.11, CUDA 9.0 and cuDNN 7.4.1 on Ubuntu 16.04. 20 | 21 | - Clone the repository 22 | ``` 23 | git clone --depth=1 https://github.com/QingyongHu/RandLA-Net && cd RandLA-Net 24 | ``` 25 | - Setup python environment 26 | ``` 27 | conda create -n randlanet python=3.5 28 | source activate randlanet 29 | pip install -r helper_requirements.txt 30 | sh compile_op.sh 31 | ``` 32 | 33 | **Update 03/21/2020, pre-trained models and results are available now.** 34 | You can download the pre-trained models and results [here](https://drive.google.com/open?id=1iU8yviO3TP87-IexBXsu13g6NklwEkXB). 35 | Note that, please specify the model path in the main function (e.g., `main_S3DIS.py`) if you want to use the pre-trained model and have a quick try of our RandLA-Net. 36 | 37 | ### (2) S3DIS 38 | S3DIS dataset can be found 39 | here. 40 | Download the files named "Stanford3dDataset_v1.2_Aligned_Version.zip". Uncompress the folder and move it to 41 | `/data/S3DIS`. 42 | 43 | - Preparing the dataset: 44 | ``` 45 | python utils/data_prepare_s3dis.py 46 | ``` 47 | - Start 6-fold cross validation: 48 | ``` 49 | sh jobs_6_fold_cv_s3dis.sh 50 | ``` 51 | - Move all the generated results (*.ply) in `/test` folder to `/data/S3DIS/results`, calculate the final mean IoU results: 52 | ``` 53 | python utils/6_fold_cv.py 54 | ``` 55 | 56 | Quantitative results of different approaches on S3DIS dataset (6-fold cross-validation): 57 | 58 | ![a](http://randla-net.cs.ox.ac.uk/imgs/S3DIS_table.png) 59 | 60 | Qualitative results of our RandLA-Net: 61 | 62 | | ![2](imgs/S3DIS_area2.gif) | ![z](imgs/S3DIS_area3.gif) | 63 | | ------------------------------ | ---------------------------- | 64 | 65 | 66 | 67 | ### (3) Semantic3D 68 | 7zip is required to uncompress the raw data in this dataset, to install p7zip: 69 | ``` 70 | sudo apt-get install p7zip-full 71 | ``` 72 | - Download and extract the dataset. First, please specify the path of the dataset by changing the `BASE_DIR` in "download_semantic3d.sh" 73 | ``` 74 | sh utils/download_semantic3d.sh 75 | ``` 76 | - Preparing the dataset: 77 | ``` 78 | python utils/data_prepare_semantic3d.py 79 | ``` 80 | - Start training: 81 | ``` 82 | python main_Semantic3D.py --mode train --gpu 0 83 | ``` 84 | - Evaluation: 85 | ``` 86 | python main_Semantic3D.py --mode test --gpu 0 87 | ``` 88 | Quantitative results of different approaches on Semantic3D (reduced-8): 89 | 90 | ![a](http://randla-net.cs.ox.ac.uk/imgs/Semantic3D_table.png) 91 | 92 | Qualitative results of our RandLA-Net: 93 | 94 | | ![z](imgs/Semantic3D-1.gif) | ![z](http://randla-net.cs.ox.ac.uk/imgs/Semantic3D-2.gif) | 95 | | -------------------------------- | ------------------------------- | 96 | | ![z](imgs/Semantic3D-3.gif) | ![z](imgs/Semantic3D-4.gif) | 97 | 98 | 99 | 100 | **Note:** 101 | - Preferably with more than 64G RAM to process this dataset due to the large volume of point cloud 102 | 103 | 104 | ### (4) SemanticKITTI 105 | 106 | SemanticKITTI dataset can be found here. Download the files 107 | related to semantic segmentation and extract everything into the same folder. Uncompress the folder and move it to 108 | `/data/semantic_kitti/dataset`. 109 | 110 | - Preparing the dataset: 111 | ``` 112 | python utils/data_prepare_semantickitti.py 113 | ``` 114 | 115 | - Start training: 116 | ``` 117 | python main_SemanticKITTI.py --mode train --gpu 0 118 | ``` 119 | 120 | - Evaluation: 121 | ``` 122 | sh jobs_test_semantickitti.sh 123 | ``` 124 | 125 | Quantitative results of different approaches on SemanticKITTI dataset: 126 | 127 | ![s](http://randla-net.cs.ox.ac.uk/imgs/SemanticKITTI_table.png) 128 | 129 | Qualitative results of our RandLA-Net: 130 | 131 | ![zzz](imgs/SemanticKITTI-2.gif) 132 | 133 | 134 | ### (5) Demo 135 | 136 |

137 | 138 | 139 | ### Citation 140 | If you find our work useful in your research, please consider citing: 141 | 142 | @article{hu2019randla, 143 | title={RandLA-Net: Efficient Semantic Segmentation of Large-Scale Point Clouds}, 144 | author={Hu, Qingyong and Yang, Bo and Xie, Linhai and Rosa, Stefano and Guo, Yulan and Wang, Zhihua and Trigoni, Niki and Markham, Andrew}, 145 | journal={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition}, 146 | year={2020} 147 | } 148 | 149 | @article{hu2021learning, 150 | title={Learning Semantic Segmentation of Large-Scale Point Clouds with Random Sampling}, 151 | author={Hu, Qingyong and Yang, Bo and Xie, Linhai and Rosa, Stefano and Guo, Yulan and Wang, Zhihua and Trigoni, Niki and Markham, Andrew}, 152 | journal={IEEE Transactions on Pattern Analysis and Machine Intelligence}, 153 | year={2021}, 154 | publisher={IEEE} 155 | } 156 | 157 | 158 | ### Acknowledgment 159 | - Part of our code refers to nanoflann library and the the recent work KPConv. 160 | - We use blender to make the video demo. 161 | 162 | 163 | ### License 164 | Licensed under the CC BY-NC-SA 4.0 license, see [LICENSE](./LICENSE). 165 | 166 | 167 | ### Updates 168 | * 21/03/2020: Updating all experimental results 169 | * 21/03/2020: Adding pretrained models and results 170 | * 02/03/2020: Code available! 171 | * 15/11/2019: Initial release！ 172 | 173 | ## Related Repos 174 | 1. [SoTA-Point-Cloud: Deep Learning for 3D Point Clouds: A Survey](https://github.com/QingyongHu/SoTA-Point-Cloud) ![GitHub stars](https://img.shields.io/github/stars/QingyongHu/SoTA-Point-Cloud.svg?style=flat&label=Star) 175 | 2. [SensatUrban: Learning Semantics from Urban-Scale Photogrammetric Point Clouds](https://github.com/QingyongHu/SpinNet) ![GitHub stars](https://img.shields.io/github/stars/QingyongHu/SensatUrban.svg?style=flat&label=Star) 176 | 3. [3D-BoNet: Learning Object Bounding Boxes for 3D Instance Segmentation on Point Clouds](https://github.com/Yang7879/3D-BoNet) ![GitHub stars](https://img.shields.io/github/stars/Yang7879/3D-BoNet.svg?style=flat&label=Star) 177 | 4. [SpinNet: Learning a General Surface Descriptor for 3D Point Cloud Registration](https://github.com/QingyongHu/SpinNet) ![GitHub stars](https://img.shields.io/github/stars/QingyongHu/SpinNet.svg?style=flat&label=Star) 178 | 5. [SQN: Weakly-Supervised Semantic Segmentation of Large-Scale 3D Point Clouds with 1000x Fewer Labels](https://github.com/QingyongHu/SQN) ![GitHub stars](https://img.shields.io/github/stars/QingyongHu/SQN.svg?style=flat&label=Star) 179 | 180 | 181 | -------------------------------------------------------------------------------- /RandLANet.py: -------------------------------------------------------------------------------- 1 | from os.path import exists, join 2 | from os import makedirs 3 | from sklearn.metrics import confusion_matrix 4 | from helper_tool import DataProcessing as DP 5 | import tensorflow as tf 6 | import numpy as np 7 | import helper_tf_util 8 | import time 9 | 10 | 11 | def log_out(out_str, f_out): 12 | f_out.write(out_str + '\n') 13 | f_out.flush() 14 | print(out_str) 15 | 16 | 17 | class Network: 18 | def __init__(self, dataset, config): 19 | flat_inputs = dataset.flat_inputs 20 | self.config = config 21 | # Path of the result folder 22 | if self.config.saving: 23 | if self.config.saving_path is None: 24 | self.saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 25 | else: 26 | self.saving_path = self.config.saving_path 27 | makedirs(self.saving_path) if not exists(self.saving_path) else None 28 | 29 | with tf.variable_scope('inputs'): 30 | self.inputs = dict() 31 | num_layers = self.config.num_layers 32 | self.inputs['xyz'] = flat_inputs[:num_layers] 33 | self.inputs['neigh_idx'] = flat_inputs[num_layers: 2 * num_layers] 34 | self.inputs['sub_idx'] = flat_inputs[2 * num_layers:3 * num_layers] 35 | self.inputs['interp_idx'] = flat_inputs[3 * num_layers:4 * num_layers] 36 | self.inputs['features'] = flat_inputs[4 * num_layers] 37 | self.inputs['labels'] = flat_inputs[4 * num_layers + 1] 38 | self.inputs['input_inds'] = flat_inputs[4 * num_layers + 2] 39 | self.inputs['cloud_inds'] = flat_inputs[4 * num_layers + 3] 40 | 41 | self.labels = self.inputs['labels'] 42 | self.is_training = tf.placeholder(tf.bool, shape=()) 43 | self.training_step = 1 44 | self.training_epoch = 0 45 | self.correct_prediction = 0 46 | self.accuracy = 0 47 | self.mIou_list = [0] 48 | self.class_weights = DP.get_class_weights(dataset.name) 49 | self.Log_file = open('log_train_' + dataset.name + str(dataset.val_split) + '.txt', 'a') 50 | 51 | with tf.variable_scope('layers'): 52 | self.logits = self.inference(self.inputs, self.is_training) 53 | 54 | ##################################################################### 55 | # Ignore the invalid point (unlabeled) when calculating the loss # 56 | ##################################################################### 57 | with tf.variable_scope('loss'): 58 | self.logits = tf.reshape(self.logits, [-1, config.num_classes]) 59 | self.labels = tf.reshape(self.labels, [-1]) 60 | 61 | # Boolean mask of points that should be ignored 62 | ignored_bool = tf.zeros_like(self.labels, dtype=tf.bool) 63 | for ign_label in self.config.ignored_label_inds: 64 | ignored_bool = tf.logical_or(ignored_bool, tf.equal(self.labels, ign_label)) 65 | 66 | # Collect logits and labels that are not ignored 67 | valid_idx = tf.squeeze(tf.where(tf.logical_not(ignored_bool))) 68 | valid_logits = tf.gather(self.logits, valid_idx, axis=0) 69 | valid_labels_init = tf.gather(self.labels, valid_idx, axis=0) 70 | 71 | # Reduce label values in the range of logit shape 72 | reducing_list = tf.range(self.config.num_classes, dtype=tf.int32) 73 | inserted_value = tf.zeros((1,), dtype=tf.int32) 74 | for ign_label in self.config.ignored_label_inds: 75 | reducing_list = tf.concat([reducing_list[:ign_label], inserted_value, reducing_list[ign_label:]], 0) 76 | valid_labels = tf.gather(reducing_list, valid_labels_init) 77 | 78 | self.loss = self.get_loss(valid_logits, valid_labels, self.class_weights) 79 | 80 | with tf.variable_scope('optimizer'): 81 | self.learning_rate = tf.Variable(config.learning_rate, trainable=False, name='learning_rate') 82 | self.train_op = tf.train.AdamOptimizer(self.learning_rate).minimize(self.loss) 83 | self.extra_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 84 | 85 | with tf.variable_scope('results'): 86 | self.correct_prediction = tf.nn.in_top_k(valid_logits, valid_labels, 1) 87 | self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32)) 88 | self.prob_logits = tf.nn.softmax(self.logits) 89 | 90 | tf.summary.scalar('learning_rate', self.learning_rate) 91 | tf.summary.scalar('loss', self.loss) 92 | tf.summary.scalar('accuracy', self.accuracy) 93 | 94 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 95 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 96 | c_proto = tf.ConfigProto() 97 | c_proto.gpu_options.allow_growth = True 98 | self.sess = tf.Session(config=c_proto) 99 | self.merged = tf.summary.merge_all() 100 | self.train_writer = tf.summary.FileWriter(config.train_sum_dir, self.sess.graph) 101 | self.sess.run(tf.global_variables_initializer()) 102 | 103 | def inference(self, inputs, is_training): 104 | 105 | d_out = self.config.d_out 106 | feature = inputs['features'] 107 | feature = tf.layers.dense(feature, 8, activation=None, name='fc0') 108 | feature = tf.nn.leaky_relu(tf.layers.batch_normalization(feature, -1, 0.99, 1e-6, training=is_training)) 109 | feature = tf.expand_dims(feature, axis=2) 110 | 111 | # ###########################Encoder############################ 112 | f_encoder_list = [] 113 | for i in range(self.config.num_layers): 114 | f_encoder_i = self.dilated_res_block(feature, inputs['xyz'][i], inputs['neigh_idx'][i], d_out[i], 115 | 'Encoder_layer_' + str(i), is_training) 116 | f_sampled_i = self.random_sample(f_encoder_i, inputs['sub_idx'][i]) 117 | feature = f_sampled_i 118 | if i == 0: 119 | f_encoder_list.append(f_encoder_i) 120 | f_encoder_list.append(f_sampled_i) 121 | # ###########################Encoder############################ 122 | 123 | feature = helper_tf_util.conv2d(f_encoder_list[-1], f_encoder_list[-1].get_shape()[3].value, [1, 1], 124 | 'decoder_0', 125 | [1, 1], 'VALID', True, is_training) 126 | 127 | # ###########################Decoder############################ 128 | f_decoder_list = [] 129 | for j in range(self.config.num_layers): 130 | f_interp_i = self.nearest_interpolation(feature, inputs['interp_idx'][-j - 1]) 131 | f_decoder_i = helper_tf_util.conv2d_transpose(tf.concat([f_encoder_list[-j - 2], f_interp_i], axis=3), 132 | f_encoder_list[-j - 2].get_shape()[-1].value, [1, 1], 133 | 'Decoder_layer_' + str(j), [1, 1], 'VALID', bn=True, 134 | is_training=is_training) 135 | feature = f_decoder_i 136 | f_decoder_list.append(f_decoder_i) 137 | # ###########################Decoder############################ 138 | 139 | f_layer_fc1 = helper_tf_util.conv2d(f_decoder_list[-1], 64, [1, 1], 'fc1', [1, 1], 'VALID', True, is_training) 140 | f_layer_fc2 = helper_tf_util.conv2d(f_layer_fc1, 32, [1, 1], 'fc2', [1, 1], 'VALID', True, is_training) 141 | f_layer_drop = helper_tf_util.dropout(f_layer_fc2, keep_prob=0.5, is_training=is_training, scope='dp1') 142 | f_layer_fc3 = helper_tf_util.conv2d(f_layer_drop, self.config.num_classes, [1, 1], 'fc', [1, 1], 'VALID', False, 143 | is_training, activation_fn=None) 144 | f_out = tf.squeeze(f_layer_fc3, [2]) 145 | return f_out 146 | 147 | def train(self, dataset): 148 | log_out('****EPOCH {}****'.format(self.training_epoch), self.Log_file) 149 | self.sess.run(dataset.train_init_op) 150 | while self.training_epoch < self.config.max_epoch: 151 | t_start = time.time() 152 | try: 153 | ops = [self.train_op, 154 | self.extra_update_ops, 155 | self.merged, 156 | self.loss, 157 | self.logits, 158 | self.labels, 159 | self.accuracy] 160 | _, _, summary, l_out, probs, labels, acc = self.sess.run(ops, {self.is_training: True}) 161 | self.train_writer.add_summary(summary, self.training_step) 162 | t_end = time.time() 163 | if self.training_step % 50 == 0: 164 | message = 'Step {:08d} L_out={:5.3f} Acc={:4.2f} ''---{:8.2f} ms/batch' 165 | log_out(message.format(self.training_step, l_out, acc, 1000 * (t_end - t_start)), self.Log_file) 166 | self.training_step += 1 167 | 168 | except tf.errors.OutOfRangeError: 169 | 170 | m_iou = self.evaluate(dataset) 171 | if m_iou > np.max(self.mIou_list): 172 | # Save the best model 173 | snapshot_directory = join(self.saving_path, 'snapshots') 174 | makedirs(snapshot_directory) if not exists(snapshot_directory) else None 175 | self.saver.save(self.sess, snapshot_directory + '/snap', global_step=self.training_step) 176 | self.mIou_list.append(m_iou) 177 | log_out('Best m_IoU is: {:5.3f}'.format(max(self.mIou_list)), self.Log_file) 178 | 179 | self.training_epoch += 1 180 | self.sess.run(dataset.train_init_op) 181 | # Update learning rate 182 | op = self.learning_rate.assign(tf.multiply(self.learning_rate, 183 | self.config.lr_decays[self.training_epoch])) 184 | self.sess.run(op) 185 | log_out('****EPOCH {}****'.format(self.training_epoch), self.Log_file) 186 | 187 | except tf.errors.InvalidArgumentError as e: 188 | 189 | print('Caught a NaN error :') 190 | print(e.error_code) 191 | print(e.message) 192 | print(e.op) 193 | print(e.op.name) 194 | print([t.name for t in e.op.inputs]) 195 | print([t.name for t in e.op.outputs]) 196 | 197 | a = 1 / 0 198 | 199 | print('finished') 200 | self.sess.close() 201 | 202 | def evaluate(self, dataset): 203 | 204 | # Initialise iterator with validation data 205 | self.sess.run(dataset.val_init_op) 206 | 207 | gt_classes = [0 for _ in range(self.config.num_classes)] 208 | positive_classes = [0 for _ in range(self.config.num_classes)] 209 | true_positive_classes = [0 for _ in range(self.config.num_classes)] 210 | val_total_correct = 0 211 | val_total_seen = 0 212 | 213 | for step_id in range(self.config.val_steps): 214 | if step_id % 50 == 0: 215 | print(str(step_id) + ' / ' + str(self.config.val_steps)) 216 | try: 217 | ops = (self.prob_logits, self.labels, self.accuracy) 218 | stacked_prob, labels, acc = self.sess.run(ops, {self.is_training: False}) 219 | pred = np.argmax(stacked_prob, 1) 220 | if not self.config.ignored_label_inds: 221 | pred_valid = pred 222 | labels_valid = labels 223 | else: 224 | invalid_idx = np.where(labels == self.config.ignored_label_inds)[0] 225 | labels_valid = np.delete(labels, invalid_idx) 226 | labels_valid = labels_valid - 1 227 | pred_valid = np.delete(pred, invalid_idx) 228 | 229 | correct = np.sum(pred_valid == labels_valid) 230 | val_total_correct += correct 231 | val_total_seen += len(labels_valid) 232 | 233 | conf_matrix = confusion_matrix(labels_valid, pred_valid, np.arange(0, self.config.num_classes, 1)) 234 | gt_classes += np.sum(conf_matrix, axis=1) 235 | positive_classes += np.sum(conf_matrix, axis=0) 236 | true_positive_classes += np.diagonal(conf_matrix) 237 | 238 | except tf.errors.OutOfRangeError: 239 | break 240 | 241 | iou_list = [] 242 | for n in range(0, self.config.num_classes, 1): 243 | iou = true_positive_classes[n] / float(gt_classes[n] + positive_classes[n] - true_positive_classes[n]) 244 | iou_list.append(iou) 245 | mean_iou = sum(iou_list) / float(self.config.num_classes) 246 | 247 | log_out('eval accuracy: {}'.format(val_total_correct / float(val_total_seen)), self.Log_file) 248 | log_out('mean IOU:{}'.format(mean_iou), self.Log_file) 249 | 250 | mean_iou = 100 * mean_iou 251 | log_out('Mean IoU = {:.1f}%'.format(mean_iou), self.Log_file) 252 | s = '{:5.2f} | '.format(mean_iou) 253 | for IoU in iou_list: 254 | s += '{:5.2f} '.format(100 * IoU) 255 | log_out('-' * len(s), self.Log_file) 256 | log_out(s, self.Log_file) 257 | log_out('-' * len(s) + '\n', self.Log_file) 258 | return mean_iou 259 | 260 | def get_loss(self, logits, labels, pre_cal_weights): 261 | # calculate the weighted cross entropy according to the inverse frequency 262 | class_weights = tf.convert_to_tensor(pre_cal_weights, dtype=tf.float32) 263 | one_hot_labels = tf.one_hot(labels, depth=self.config.num_classes) 264 | weights = tf.reduce_sum(class_weights * one_hot_labels, axis=1) 265 | unweighted_losses = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=one_hot_labels) 266 | weighted_losses = unweighted_losses * weights 267 | output_loss = tf.reduce_mean(weighted_losses) 268 | return output_loss 269 | 270 | def dilated_res_block(self, feature, xyz, neigh_idx, d_out, name, is_training): 271 | f_pc = helper_tf_util.conv2d(feature, d_out // 2, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) 272 | f_pc = self.building_block(xyz, f_pc, neigh_idx, d_out, name + 'LFA', is_training) 273 | f_pc = helper_tf_util.conv2d(f_pc, d_out * 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training, 274 | activation_fn=None) 275 | shortcut = helper_tf_util.conv2d(feature, d_out * 2, [1, 1], name + 'shortcut', [1, 1], 'VALID', 276 | activation_fn=None, bn=True, is_training=is_training) 277 | return tf.nn.leaky_relu(f_pc + shortcut) 278 | 279 | def building_block(self, xyz, feature, neigh_idx, d_out, name, is_training): 280 | d_in = feature.get_shape()[-1].value 281 | f_xyz = self.relative_pos_encoding(xyz, neigh_idx) 282 | f_xyz = helper_tf_util.conv2d(f_xyz, d_in, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) 283 | f_neighbours = self.gather_neighbour(tf.squeeze(feature, axis=2), neigh_idx) 284 | f_concat = tf.concat([f_neighbours, f_xyz], axis=-1) 285 | f_pc_agg = self.att_pooling(f_concat, d_out // 2, name + 'att_pooling_1', is_training) 286 | 287 | f_xyz = helper_tf_util.conv2d(f_xyz, d_out // 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training) 288 | f_neighbours = self.gather_neighbour(tf.squeeze(f_pc_agg, axis=2), neigh_idx) 289 | f_concat = tf.concat([f_neighbours, f_xyz], axis=-1) 290 | f_pc_agg = self.att_pooling(f_concat, d_out, name + 'att_pooling_2', is_training) 291 | return f_pc_agg 292 | 293 | def relative_pos_encoding(self, xyz, neigh_idx): 294 | neighbor_xyz = self.gather_neighbour(xyz, neigh_idx) 295 | xyz_tile = tf.tile(tf.expand_dims(xyz, axis=2), [1, 1, tf.shape(neigh_idx)[-1], 1]) 296 | relative_xyz = xyz_tile - neighbor_xyz 297 | relative_dis = tf.sqrt(tf.reduce_sum(tf.square(relative_xyz), axis=-1, keepdims=True)) 298 | relative_feature = tf.concat([relative_dis, relative_xyz, xyz_tile, neighbor_xyz], axis=-1) 299 | return relative_feature 300 | 301 | @staticmethod 302 | def random_sample(feature, pool_idx): 303 | """ 304 | :param feature: [B, N, d] input features matrix 305 | :param pool_idx: [B, N', max_num] N' < N, N' is the selected position after pooling 306 | :return: pool_features = [B, N', d] pooled features matrix 307 | """ 308 | feature = tf.squeeze(feature, axis=2) 309 | num_neigh = tf.shape(pool_idx)[-1] 310 | d = feature.get_shape()[-1] 311 | batch_size = tf.shape(pool_idx)[0] 312 | pool_idx = tf.reshape(pool_idx, [batch_size, -1]) 313 | pool_features = tf.batch_gather(feature, pool_idx) 314 | pool_features = tf.reshape(pool_features, [batch_size, -1, num_neigh, d]) 315 | pool_features = tf.reduce_max(pool_features, axis=2, keepdims=True) 316 | return pool_features 317 | 318 | @staticmethod 319 | def nearest_interpolation(feature, interp_idx): 320 | """ 321 | :param feature: [B, N, d] input features matrix 322 | :param interp_idx: [B, up_num_points, 1] nearest neighbour index 323 | :return: [B, up_num_points, d] interpolated features matrix 324 | """ 325 | feature = tf.squeeze(feature, axis=2) 326 | batch_size = tf.shape(interp_idx)[0] 327 | up_num_points = tf.shape(interp_idx)[1] 328 | interp_idx = tf.reshape(interp_idx, [batch_size, up_num_points]) 329 | interpolated_features = tf.batch_gather(feature, interp_idx) 330 | interpolated_features = tf.expand_dims(interpolated_features, axis=2) 331 | return interpolated_features 332 | 333 | @staticmethod 334 | def gather_neighbour(pc, neighbor_idx): 335 | # gather the coordinates or features of neighboring points 336 | batch_size = tf.shape(pc)[0] 337 | num_points = tf.shape(pc)[1] 338 | d = pc.get_shape()[2].value 339 | index_input = tf.reshape(neighbor_idx, shape=[batch_size, -1]) 340 | features = tf.batch_gather(pc, index_input) 341 | features = tf.reshape(features, [batch_size, num_points, tf.shape(neighbor_idx)[-1], d]) 342 | return features 343 | 344 | @staticmethod 345 | def att_pooling(feature_set, d_out, name, is_training): 346 | batch_size = tf.shape(feature_set)[0] 347 | num_points = tf.shape(feature_set)[1] 348 | num_neigh = tf.shape(feature_set)[2] 349 | d = feature_set.get_shape()[3].value 350 | f_reshaped = tf.reshape(feature_set, shape=[-1, num_neigh, d]) 351 | att_activation = tf.layers.dense(f_reshaped, d, activation=None, use_bias=False, name=name + 'fc') 352 | att_scores = tf.nn.softmax(att_activation, axis=1) 353 | f_agg = f_reshaped * att_scores 354 | f_agg = tf.reduce_sum(f_agg, axis=1) 355 | f_agg = tf.reshape(f_agg, [batch_size, num_points, 1, d]) 356 | f_agg = helper_tf_util.conv2d(f_agg, d_out, [1, 1], name + 'mlp', [1, 1], 'VALID', True, is_training) 357 | return f_agg 358 | -------------------------------------------------------------------------------- /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 ../../../ -------------------------------------------------------------------------------- /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 cpp_wrappers.cpp_subsampling.grid_subsampling as cpp_subsampling 15 | import 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 | 44 | class ConfigS3DIS: 45 | k_n = 16 # KNN 46 | num_layers = 5 # Number of layers 47 | num_points = 40960 # Number of input points 48 | num_classes = 13 # Number of valid classes 49 | sub_grid_size = 0.04 # preprocess_parameter 50 | 51 | batch_size = 6 # batch_size during training 52 | val_batch_size = 20 # batch_size during validation and test 53 | train_steps = 500 # Number of steps per epochs 54 | val_steps = 100 # Number of validation steps per epoch 55 | 56 | sub_sampling_ratio = [4, 4, 4, 4, 2] # sampling ratio of random sampling at each layer 57 | d_out = [16, 64, 128, 256, 512] # feature dimension 58 | 59 | noise_init = 3.5 # noise initial parameter 60 | max_epoch = 100 # maximum epoch during training 61 | learning_rate = 1e-2 # initial learning rate 62 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 63 | 64 | train_sum_dir = 'train_log' 65 | saving = True 66 | saving_path = None 67 | 68 | 69 | class ConfigSemantic3D: 70 | k_n = 16 # KNN 71 | num_layers = 5 # Number of layers 72 | num_points = 65536 # Number of input points 73 | num_classes = 8 # Number of valid classes 74 | sub_grid_size = 0.06 # preprocess_parameter 75 | 76 | batch_size = 4 # batch_size during training 77 | val_batch_size = 16 # batch_size during validation and test 78 | train_steps = 500 # Number of steps per epochs 79 | val_steps = 100 # Number of validation steps per epoch 80 | 81 | sub_sampling_ratio = [4, 4, 4, 4, 2] # sampling ratio of random sampling at each layer 82 | d_out = [16, 64, 128, 256, 512] # feature dimension 83 | 84 | noise_init = 3.5 # noise initial parameter 85 | max_epoch = 100 # maximum epoch during training 86 | learning_rate = 1e-2 # initial learning rate 87 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 88 | 89 | train_sum_dir = 'train_log' 90 | saving = True 91 | saving_path = None 92 | 93 | augment_scale_anisotropic = True 94 | augment_symmetries = [True, False, False] 95 | augment_rotation = 'vertical' 96 | augment_scale_min = 0.8 97 | augment_scale_max = 1.2 98 | augment_noise = 0.001 99 | augment_occlusion = 'none' 100 | augment_color = 0.8 101 | 102 | 103 | class DataProcessing: 104 | @staticmethod 105 | def load_pc_semantic3d(filename): 106 | pc_pd = pd.read_csv(filename, header=None, delim_whitespace=True, dtype=np.float16) 107 | pc = pc_pd.values 108 | return pc 109 | 110 | @staticmethod 111 | def load_label_semantic3d(filename): 112 | label_pd = pd.read_csv(filename, header=None, delim_whitespace=True, dtype=np.uint8) 113 | cloud_labels = label_pd.values 114 | return cloud_labels 115 | 116 | @staticmethod 117 | def load_pc_kitti(pc_path): 118 | scan = np.fromfile(pc_path, dtype=np.float32) 119 | scan = scan.reshape((-1, 4)) 120 | points = scan[:, 0:3] # get xyz 121 | return points 122 | 123 | @staticmethod 124 | def load_label_kitti(label_path, remap_lut): 125 | label = np.fromfile(label_path, dtype=np.uint32) 126 | label = label.reshape((-1)) 127 | sem_label = label & 0xFFFF # semantic label in lower half 128 | inst_label = label >> 16 # instance id in upper half 129 | assert ((sem_label + (inst_label << 16) == label).all()) 130 | sem_label = remap_lut[sem_label] 131 | return sem_label.astype(np.int32) 132 | 133 | @staticmethod 134 | def get_file_list(dataset_path, test_scan_num): 135 | seq_list = np.sort(os.listdir(dataset_path)) 136 | 137 | train_file_list = [] 138 | test_file_list = [] 139 | val_file_list = [] 140 | for seq_id in seq_list: 141 | seq_path = join(dataset_path, seq_id) 142 | pc_path = join(seq_path, 'velodyne') 143 | if seq_id == '08': 144 | val_file_list.append([join(pc_path, f) for f in np.sort(os.listdir(pc_path))]) 145 | if seq_id == test_scan_num: 146 | test_file_list.append([join(pc_path, f) for f in np.sort(os.listdir(pc_path))]) 147 | elif int(seq_id) >= 11 and seq_id == test_scan_num: 148 | test_file_list.append([join(pc_path, f) for f in np.sort(os.listdir(pc_path))]) 149 | elif seq_id in ['00', '01', '02', '03', '04', '05', '06', '07', '09', '10']: 150 | train_file_list.append([join(pc_path, f) for f in np.sort(os.listdir(pc_path))]) 151 | 152 | train_file_list = np.concatenate(train_file_list, axis=0) 153 | val_file_list = np.concatenate(val_file_list, axis=0) 154 | test_file_list = np.concatenate(test_file_list, axis=0) 155 | return train_file_list, val_file_list, test_file_list 156 | 157 | @staticmethod 158 | def knn_search(support_pts, query_pts, k): 159 | """ 160 | :param support_pts: points you have, B*N1*3 161 | :param query_pts: points you want to know the neighbour index, B*N2*3 162 | :param k: Number of neighbours in knn search 163 | :return: neighbor_idx: neighboring points indexes, B*N2*k 164 | """ 165 | 166 | neighbor_idx = nearest_neighbors.knn_batch(support_pts, query_pts, k, omp=True) 167 | return neighbor_idx.astype(np.int32) 168 | 169 | @staticmethod 170 | def data_aug(xyz, color, labels, idx, num_out): 171 | num_in = len(xyz) 172 | dup = np.random.choice(num_in, num_out - num_in) 173 | xyz_dup = xyz[dup, ...] 174 | xyz_aug = np.concatenate([xyz, xyz_dup], 0) 175 | color_dup = color[dup, ...] 176 | color_aug = np.concatenate([color, color_dup], 0) 177 | idx_dup = list(range(num_in)) + list(dup) 178 | idx_aug = idx[idx_dup] 179 | label_aug = labels[idx_dup] 180 | return xyz_aug, color_aug, idx_aug, label_aug 181 | 182 | @staticmethod 183 | def shuffle_idx(x): 184 | # random shuffle the index 185 | idx = np.arange(len(x)) 186 | np.random.shuffle(idx) 187 | return x[idx] 188 | 189 | @staticmethod 190 | def shuffle_list(data_list): 191 | indices = np.arange(np.shape(data_list)[0]) 192 | np.random.shuffle(indices) 193 | data_list = data_list[indices] 194 | return data_list 195 | 196 | @staticmethod 197 | def grid_sub_sampling(points, features=None, labels=None, grid_size=0.1, verbose=0): 198 | """ 199 | CPP wrapper for a grid sub_sampling (method = barycenter for points and features 200 | :param points: (N, 3) matrix of input points 201 | :param features: optional (N, d) matrix of features (floating number) 202 | :param labels: optional (N,) matrix of integer labels 203 | :param grid_size: parameter defining the size of grid voxels 204 | :param verbose: 1 to display 205 | :return: sub_sampled points, with features and/or labels depending of the input 206 | """ 207 | 208 | if (features is None) and (labels is None): 209 | return cpp_subsampling.compute(points, sampleDl=grid_size, verbose=verbose) 210 | elif labels is None: 211 | return cpp_subsampling.compute(points, features=features, sampleDl=grid_size, verbose=verbose) 212 | elif features is None: 213 | return cpp_subsampling.compute(points, classes=labels, sampleDl=grid_size, verbose=verbose) 214 | else: 215 | return cpp_subsampling.compute(points, features=features, classes=labels, sampleDl=grid_size, 216 | verbose=verbose) 217 | 218 | @staticmethod 219 | def IoU_from_confusions(confusions): 220 | """ 221 | Computes IoU from confusion matrices. 222 | :param confusions: ([..., n_c, n_c] np.int32). Can be any dimension, the confusion matrices should be described by 223 | the last axes. n_c = number of classes 224 | :return: ([..., n_c] np.float32) IoU score 225 | """ 226 | 227 | # Compute TP, FP, FN. This assume that the second to last axis counts the truths (like the first axis of a 228 | # confusion matrix), and that the last axis counts the predictions (like the second axis of a confusion matrix) 229 | TP = np.diagonal(confusions, axis1=-2, axis2=-1) 230 | TP_plus_FN = np.sum(confusions, axis=-1) 231 | TP_plus_FP = np.sum(confusions, axis=-2) 232 | 233 | # Compute IoU 234 | IoU = TP / (TP_plus_FP + TP_plus_FN - TP + 1e-6) 235 | 236 | # Compute mIoU with only the actual classes 237 | mask = TP_plus_FN < 1e-3 238 | counts = np.sum(1 - mask, axis=-1, keepdims=True) 239 | mIoU = np.sum(IoU, axis=-1, keepdims=True) / (counts + 1e-6) 240 | 241 | # If class is absent, place mIoU in place of 0 IoU to get the actual mean later 242 | IoU += mask * mIoU 243 | return IoU 244 | 245 | @staticmethod 246 | def get_class_weights(dataset_name): 247 | # pre-calculate the number of points in each category 248 | num_per_class = [] 249 | if dataset_name is 'S3DIS': 250 | num_per_class = np.array([3370714, 2856755, 4919229, 318158, 375640, 478001, 974733, 251 | 650464, 791496, 88727, 1284130, 229758, 2272837], dtype=np.int32) 252 | elif dataset_name is 'Semantic3D': 253 | num_per_class = np.array([5181602, 5012952, 6830086, 1311528, 10476365, 946982, 334860, 269353], 254 | dtype=np.int32) 255 | elif dataset_name is 'SemanticKITTI': 256 | num_per_class = np.array([55437630, 320797, 541736, 2578735, 3274484, 552662, 184064, 78858, 257 | 240942562, 17294618, 170599734, 6369672, 230413074, 101130274, 476491114, 258 | 9833174, 129609852, 4506626, 1168181]) 259 | weight = num_per_class / float(sum(num_per_class)) 260 | ce_label_weight = 1 / (weight + 0.02) 261 | return np.expand_dims(ce_label_weight, axis=0) 262 | 263 | 264 | class Plot: 265 | @staticmethod 266 | def random_colors(N, bright=True, seed=0): 267 | brightness = 1.0 if bright else 0.7 268 | hsv = [(0.15 + i / float(N), 1, brightness) for i in range(N)] 269 | colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv)) 270 | random.seed(seed) 271 | random.shuffle(colors) 272 | return colors 273 | 274 | @staticmethod 275 | def draw_pc(pc_xyzrgb): 276 | pc = open3d.PointCloud() 277 | pc.points = open3d.Vector3dVector(pc_xyzrgb[:, 0:3]) 278 | if pc_xyzrgb.shape[1] == 3: 279 | open3d.draw_geometries([pc]) 280 | return 0 281 | if np.max(pc_xyzrgb[:, 3:6]) > 20: ## 0-255 282 | pc.colors = open3d.Vector3dVector(pc_xyzrgb[:, 3:6] / 255.) 283 | else: 284 | pc.colors = open3d.Vector3dVector(pc_xyzrgb[:, 3:6]) 285 | open3d.draw_geometries([pc]) 286 | return 0 287 | 288 | @staticmethod 289 | def draw_pc_sem_ins(pc_xyz, pc_sem_ins, plot_colors=None): 290 | """ 291 | pc_xyz: 3D coordinates of point clouds 292 | pc_sem_ins: semantic or instance labels 293 | plot_colors: custom color list 294 | """ 295 | if plot_colors is not None: 296 | ins_colors = plot_colors 297 | else: 298 | ins_colors = Plot.random_colors(len(np.unique(pc_sem_ins)) + 1, seed=2) 299 | 300 | ############################## 301 | sem_ins_labels = np.unique(pc_sem_ins) 302 | sem_ins_bbox = [] 303 | Y_colors = np.zeros((pc_sem_ins.shape[0], 3)) 304 | for id, semins in enumerate(sem_ins_labels): 305 | valid_ind = np.argwhere(pc_sem_ins == semins)[:, 0] 306 | if semins <= -1: 307 | tp = [0, 0, 0] 308 | else: 309 | if plot_colors is not None: 310 | tp = ins_colors[semins] 311 | else: 312 | tp = ins_colors[id] 313 | 314 | Y_colors[valid_ind] = tp 315 | 316 | ### bbox 317 | valid_xyz = pc_xyz[valid_ind] 318 | 319 | xmin = np.min(valid_xyz[:, 0]); 320 | xmax = np.max(valid_xyz[:, 0]) 321 | ymin = np.min(valid_xyz[:, 1]); 322 | ymax = np.max(valid_xyz[:, 1]) 323 | zmin = np.min(valid_xyz[:, 2]); 324 | zmax = np.max(valid_xyz[:, 2]) 325 | sem_ins_bbox.append( 326 | [[xmin, ymin, zmin], [xmax, ymax, zmax], [min(tp[0], 1.), min(tp[1], 1.), min(tp[2], 1.)]]) 327 | 328 | Y_semins = np.concatenate([pc_xyz[:, 0:3], Y_colors], axis=-1) 329 | Plot.draw_pc(Y_semins) 330 | return Y_semins 331 | -------------------------------------------------------------------------------- /imgs/S3DIS_area2.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/RandLA-Net/6b5445f5f279d33d2335e85ed39ca8b68cb1c57e/imgs/S3DIS_area2.gif -------------------------------------------------------------------------------- /imgs/S3DIS_area3.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/RandLA-Net/6b5445f5f279d33d2335e85ed39ca8b68cb1c57e/imgs/S3DIS_area3.gif -------------------------------------------------------------------------------- /imgs/Semantic3D-1.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/RandLA-Net/6b5445f5f279d33d2335e85ed39ca8b68cb1c57e/imgs/Semantic3D-1.gif -------------------------------------------------------------------------------- /imgs/Semantic3D-3.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/RandLA-Net/6b5445f5f279d33d2335e85ed39ca8b68cb1c57e/imgs/Semantic3D-3.gif -------------------------------------------------------------------------------- /imgs/Semantic3D-4.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/RandLA-Net/6b5445f5f279d33d2335e85ed39ca8b68cb1c57e/imgs/Semantic3D-4.gif -------------------------------------------------------------------------------- /imgs/SemanticKITTI-2.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QingyongHu/RandLA-Net/6b5445f5f279d33d2335e85ed39ca8b68cb1c57e/imgs/SemanticKITTI-2.gif -------------------------------------------------------------------------------- /jobs_6_fold_cv_s3dis.sh: -------------------------------------------------------------------------------- 1 | python -B main_S3DIS.py --gpu 0 --mode train --test_area 1 2 | python -B main_S3DIS.py --gpu 0 --mode test --test_area 1 3 | python -B main_S3DIS.py --gpu 0 --mode train --test_area 2 4 | python -B main_S3DIS.py --gpu 0 --mode test --test_area 2 5 | python -B main_S3DIS.py --gpu 0 --mode train --test_area 3 6 | python -B main_S3DIS.py --gpu 0 --mode test --test_area 3 7 | python -B main_S3DIS.py --gpu 0 --mode train --test_area 4 8 | python -B main_S3DIS.py --gpu 0 --mode test --test_area 4 9 | python -B main_S3DIS.py --gpu 0 --mode train --test_area 5 10 | python -B main_S3DIS.py --gpu 0 --mode test --test_area 5 11 | python -B main_S3DIS.py --gpu 0 --mode train --test_area 6 12 | python -B main_S3DIS.py --gpu 0 --mode test --test_area 6 13 | 14 | 15 | -------------------------------------------------------------------------------- /jobs_test_semantickitti.sh: -------------------------------------------------------------------------------- 1 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 11 2 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 12 3 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 13 4 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 14 5 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 15 6 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 16 7 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 17 8 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 18 9 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 19 10 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 20 11 | python -B main_SemanticKITTI.py --gpu 0 --mode test --test_area 21 12 | 13 | -------------------------------------------------------------------------------- /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 | from helper_tool import ConfigS3DIS as cfg 6 | from helper_tool import DataProcessing as DP 7 | from helper_tool import Plot 8 | import tensorflow as tf 9 | import numpy as np 10 | import time, pickle, argparse, glob, os 11 | 12 | 13 | class S3DIS: 14 | def __init__(self, test_area_idx): 15 | self.name = 'S3DIS' 16 | self.path = '/data/S3DIS' 17 | self.label_to_names = {0: 'ceiling', 18 | 1: 'floor', 19 | 2: 'wall', 20 | 3: 'beam', 21 | 4: 'column', 22 | 5: 'window', 23 | 6: 'door', 24 | 7: 'table', 25 | 8: 'chair', 26 | 9: 'sofa', 27 | 10: 'bookcase', 28 | 11: 'board', 29 | 12: 'clutter'} 30 | self.num_classes = len(self.label_to_names) 31 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 32 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 33 | self.ignored_labels = np.array([]) 34 | 35 | self.val_split = 'Area_' + str(test_area_idx) 36 | self.all_files = glob.glob(join(self.path, 'original_ply', '*.ply')) 37 | 38 | # Initiate containers 39 | self.val_proj = [] 40 | self.val_labels = [] 41 | self.possibility = {} 42 | self.min_possibility = {} 43 | self.input_trees = {'training': [], 'validation': []} 44 | self.input_colors = {'training': [], 'validation': []} 45 | self.input_labels = {'training': [], 'validation': []} 46 | self.input_names = {'training': [], 'validation': []} 47 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 48 | 49 | def load_sub_sampled_clouds(self, sub_grid_size): 50 | tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size)) 51 | for i, file_path in enumerate(self.all_files): 52 | t0 = time.time() 53 | cloud_name = file_path.split('/')[-1][:-4] 54 | if self.val_split in cloud_name: 55 | cloud_split = 'validation' 56 | else: 57 | cloud_split = 'training' 58 | 59 | # Name of the input files 60 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 61 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 62 | 63 | data = read_ply(sub_ply_file) 64 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 65 | sub_labels = data['class'] 66 | 67 | # Read pkl with search tree 68 | with open(kd_tree_file, 'rb') as f: 69 | search_tree = pickle.load(f) 70 | 71 | self.input_trees[cloud_split] += [search_tree] 72 | self.input_colors[cloud_split] += [sub_colors] 73 | self.input_labels[cloud_split] += [sub_labels] 74 | self.input_names[cloud_split] += [cloud_name] 75 | 76 | size = sub_colors.shape[0] * 4 * 7 77 | print('{:s} {:.1f} MB loaded in {:.1f}s'.format(kd_tree_file.split('/')[-1], size * 1e-6, time.time() - t0)) 78 | 79 | print('\nPreparing reprojected indices for testing') 80 | 81 | # Get validation and test reprojected indices 82 | for i, file_path in enumerate(self.all_files): 83 | t0 = time.time() 84 | cloud_name = file_path.split('/')[-1][:-4] 85 | 86 | # Validation projection and labels 87 | if self.val_split in cloud_name: 88 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 89 | with open(proj_file, 'rb') as f: 90 | proj_idx, labels = pickle.load(f) 91 | self.val_proj += [proj_idx] 92 | self.val_labels += [labels] 93 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 94 | 95 | # Generate the input data flow 96 | def get_batch_gen(self, split): 97 | if split == 'training': 98 | num_per_epoch = cfg.train_steps * cfg.batch_size 99 | elif split == 'validation': 100 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 101 | 102 | self.possibility[split] = [] 103 | self.min_possibility[split] = [] 104 | # Random initialize 105 | for i, tree in enumerate(self.input_colors[split]): 106 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 107 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 108 | 109 | def spatially_regular_gen(): 110 | # Generator loop 111 | for i in range(num_per_epoch): 112 | 113 | # Choose the cloud with the lowest probability 114 | cloud_idx = int(np.argmin(self.min_possibility[split])) 115 | 116 | # choose the point with the minimum of possibility in the cloud as query point 117 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 118 | 119 | # Get all points within the cloud from tree structure 120 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 121 | 122 | # Center point of input region 123 | center_point = points[point_ind, :].reshape(1, -1) 124 | 125 | # Add noise to the center point 126 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 127 | pick_point = center_point + noise.astype(center_point.dtype) 128 | 129 | # Check if the number of points in the selected cloud is less than the predefined num_points 130 | if len(points) < cfg.num_points: 131 | # Query all points within the cloud 132 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=len(points))[1][0] 133 | else: 134 | # Query the predefined number of points 135 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 136 | 137 | # Shuffle index 138 | queried_idx = DP.shuffle_idx(queried_idx) 139 | # Get corresponding points and colors based on the index 140 | queried_pc_xyz = points[queried_idx] 141 | queried_pc_xyz = queried_pc_xyz - pick_point 142 | queried_pc_colors = self.input_colors[split][cloud_idx][queried_idx] 143 | queried_pc_labels = self.input_labels[split][cloud_idx][queried_idx] 144 | 145 | # Update the possibility of the selected points 146 | dists = np.sum(np.square((points[queried_idx] - pick_point).astype(np.float32)), axis=1) 147 | delta = np.square(1 - dists / np.max(dists)) 148 | self.possibility[split][cloud_idx][queried_idx] += delta 149 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 150 | 151 | # up_sampled with replacement 152 | if len(points) < cfg.num_points: 153 | queried_pc_xyz, queried_pc_colors, queried_idx, queried_pc_labels = \ 154 | DP.data_aug(queried_pc_xyz, queried_pc_colors, queried_pc_labels, queried_idx, cfg.num_points) 155 | 156 | if True: 157 | yield (queried_pc_xyz.astype(np.float32), 158 | queried_pc_colors.astype(np.float32), 159 | queried_pc_labels, 160 | queried_idx.astype(np.int32), 161 | np.array([cloud_idx], dtype=np.int32)) 162 | 163 | gen_func = spatially_regular_gen 164 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32) 165 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None]) 166 | return gen_func, gen_types, gen_shapes 167 | 168 | @staticmethod 169 | def get_tf_mapping2(): 170 | # Collect flat inputs 171 | def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 172 | batch_features = tf.concat([batch_xyz, batch_features], axis=-1) 173 | input_points = [] 174 | input_neighbors = [] 175 | input_pools = [] 176 | input_up_samples = [] 177 | 178 | for i in range(cfg.num_layers): 179 | neighbour_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) 180 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 181 | pool_i = neighbour_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 182 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) 183 | input_points.append(batch_xyz) 184 | input_neighbors.append(neighbour_idx) 185 | input_pools.append(pool_i) 186 | input_up_samples.append(up_i) 187 | batch_xyz = sub_points 188 | 189 | input_list = input_points + input_neighbors + input_pools + input_up_samples 190 | input_list += [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx] 191 | 192 | return input_list 193 | 194 | return tf_map 195 | 196 | def init_input_pipeline(self): 197 | print('Initiating input pipelines') 198 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 199 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 200 | gen_function_val, _, _ = self.get_batch_gen('validation') 201 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 202 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 203 | 204 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 205 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 206 | map_func = self.get_tf_mapping2() 207 | 208 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 209 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 210 | 211 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 212 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 213 | 214 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 215 | self.flat_inputs = iter.get_next() 216 | self.train_init_op = iter.make_initializer(self.batch_train_data) 217 | self.val_init_op = iter.make_initializer(self.batch_val_data) 218 | 219 | 220 | if __name__ == '__main__': 221 | parser = argparse.ArgumentParser() 222 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 223 | parser.add_argument('--test_area', type=int, default=5, help='Which area to use for test, option: 1-6 [default: 5]') 224 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 225 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 226 | FLAGS = parser.parse_args() 227 | 228 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 229 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 230 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 231 | Mode = FLAGS.mode 232 | 233 | test_area = FLAGS.test_area 234 | dataset = S3DIS(test_area) 235 | dataset.init_input_pipeline() 236 | 237 | if Mode == 'train': 238 | model = Network(dataset, cfg) 239 | model.train(dataset) 240 | elif Mode == 'test': 241 | cfg.saving = False 242 | model = Network(dataset, cfg) 243 | if FLAGS.model_path is not 'None': 244 | chosen_snap = FLAGS.model_path 245 | else: 246 | chosen_snapshot = -1 247 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 248 | chosen_folder = logs[-1] 249 | snap_path = join(chosen_folder, 'snapshots') 250 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 251 | chosen_step = np.sort(snap_steps)[-1] 252 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 253 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 254 | tester.test(model, dataset) 255 | else: 256 | ################## 257 | # Visualize data # 258 | ################## 259 | 260 | with tf.Session() as sess: 261 | sess.run(tf.global_variables_initializer()) 262 | sess.run(dataset.train_init_op) 263 | while True: 264 | flat_inputs = sess.run(dataset.flat_inputs) 265 | pc_xyz = flat_inputs[0] 266 | sub_pc_xyz = flat_inputs[1] 267 | labels = flat_inputs[21] 268 | Plot.draw_pc_sem_ins(pc_xyz[0, :, :], labels[0, :]) 269 | Plot.draw_pc_sem_ins(sub_pc_xyz[0, :, :], labels[0, 0:np.shape(sub_pc_xyz)[1]]) 270 | -------------------------------------------------------------------------------- /main_Semantic3D.py: -------------------------------------------------------------------------------- 1 | from os.path import join, exists 2 | from RandLANet import Network 3 | from tester_Semantic3D import ModelTester 4 | from helper_ply import read_ply 5 | from helper_tool import Plot 6 | from helper_tool import DataProcessing as DP 7 | from helper_tool import ConfigSemantic3D as cfg 8 | import tensorflow as tf 9 | import numpy as np 10 | import pickle, argparse, os 11 | 12 | 13 | class Semantic3D: 14 | def __init__(self): 15 | self.name = 'Semantic3D' 16 | self.path = '/data/semantic3d' 17 | self.label_to_names = {0: 'unlabeled', 18 | 1: 'man-made terrain', 19 | 2: 'natural terrain', 20 | 3: 'high vegetation', 21 | 4: 'low vegetation', 22 | 5: 'buildings', 23 | 6: 'hard scape', 24 | 7: 'scanning artefacts', 25 | 8: 'cars'} 26 | self.num_classes = len(self.label_to_names) 27 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 28 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 29 | self.ignored_labels = np.sort([0]) 30 | 31 | self.original_folder = join(self.path, 'original_data') 32 | self.full_pc_folder = join(self.path, 'original_ply') 33 | self.sub_pc_folder = join(self.path, 'input_{:.3f}'.format(cfg.sub_grid_size)) 34 | 35 | # Following KPConv to do the train-validation split 36 | self.all_splits = [0, 1, 4, 5, 3, 4, 3, 0, 1, 2, 3, 4, 2, 0, 5] 37 | self.val_split = 1 38 | 39 | # Initial training-validation-testing files 40 | self.train_files = [] 41 | self.val_files = [] 42 | self.test_files = [] 43 | cloud_names = [file_name[:-4] for file_name in os.listdir(self.original_folder) if file_name[-4:] == '.txt'] 44 | for pc_name in cloud_names: 45 | if exists(join(self.original_folder, pc_name + '.labels')): 46 | self.train_files.append(join(self.sub_pc_folder, pc_name + '.ply')) 47 | else: 48 | self.test_files.append(join(self.full_pc_folder, pc_name + '.ply')) 49 | 50 | self.train_files = np.sort(self.train_files) 51 | self.test_files = np.sort(self.test_files) 52 | 53 | for i, file_path in enumerate(self.train_files): 54 | if self.all_splits[i] == self.val_split: 55 | self.val_files.append(file_path) 56 | 57 | self.train_files = np.sort([x for x in self.train_files if x not in self.val_files]) 58 | 59 | # Initiate containers 60 | self.val_proj = [] 61 | self.val_labels = [] 62 | self.test_proj = [] 63 | self.test_labels = [] 64 | 65 | self.possibility = {} 66 | self.min_possibility = {} 67 | self.class_weight = {} 68 | self.input_trees = {'training': [], 'validation': [], 'test': []} 69 | self.input_colors = {'training': [], 'validation': [], 'test': []} 70 | self.input_labels = {'training': [], 'validation': []} 71 | 72 | # Ascii files dict for testing 73 | self.ascii_files = { 74 | 'MarketplaceFeldkirch_Station4_rgb_intensity-reduced.ply': 'marketsquarefeldkirch4-reduced.labels', 75 | 'sg27_station10_rgb_intensity-reduced.ply': 'sg27_10-reduced.labels', 76 | 'sg28_Station2_rgb_intensity-reduced.ply': 'sg28_2-reduced.labels', 77 | 'StGallenCathedral_station6_rgb_intensity-reduced.ply': 'stgallencathedral6-reduced.labels', 78 | 'birdfountain_station1_xyz_intensity_rgb.ply': 'birdfountain1.labels', 79 | 'castleblatten_station1_intensity_rgb.ply': 'castleblatten1.labels', 80 | 'castleblatten_station5_xyz_intensity_rgb.ply': 'castleblatten5.labels', 81 | 'marketplacefeldkirch_station1_intensity_rgb.ply': 'marketsquarefeldkirch1.labels', 82 | 'marketplacefeldkirch_station4_intensity_rgb.ply': 'marketsquarefeldkirch4.labels', 83 | 'marketplacefeldkirch_station7_intensity_rgb.ply': 'marketsquarefeldkirch7.labels', 84 | 'sg27_station10_intensity_rgb.ply': 'sg27_10.labels', 85 | 'sg27_station3_intensity_rgb.ply': 'sg27_3.labels', 86 | 'sg27_station6_intensity_rgb.ply': 'sg27_6.labels', 87 | 'sg27_station8_intensity_rgb.ply': 'sg27_8.labels', 88 | 'sg28_station2_intensity_rgb.ply': 'sg28_2.labels', 89 | 'sg28_station5_xyz_intensity_rgb.ply': 'sg28_5.labels', 90 | 'stgallencathedral_station1_intensity_rgb.ply': 'stgallencathedral1.labels', 91 | 'stgallencathedral_station3_intensity_rgb.ply': 'stgallencathedral3.labels', 92 | 'stgallencathedral_station6_intensity_rgb.ply': 'stgallencathedral6.labels'} 93 | 94 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 95 | 96 | def load_sub_sampled_clouds(self, sub_grid_size): 97 | 98 | tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size)) 99 | files = np.hstack((self.train_files, self.val_files, self.test_files)) 100 | 101 | for i, file_path in enumerate(files): 102 | cloud_name = file_path.split('/')[-1][:-4] 103 | print('Load_pc_' + str(i) + ': ' + cloud_name) 104 | if file_path in self.val_files: 105 | cloud_split = 'validation' 106 | elif file_path in self.train_files: 107 | cloud_split = 'training' 108 | else: 109 | cloud_split = 'test' 110 | 111 | # Name of the input files 112 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 113 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 114 | 115 | # read ply with data 116 | data = read_ply(sub_ply_file) 117 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 118 | if cloud_split == 'test': 119 | sub_labels = None 120 | else: 121 | sub_labels = data['class'] 122 | 123 | # Read pkl with search tree 124 | with open(kd_tree_file, 'rb') as f: 125 | search_tree = pickle.load(f) 126 | 127 | self.input_trees[cloud_split] += [search_tree] 128 | self.input_colors[cloud_split] += [sub_colors] 129 | if cloud_split in ['training', 'validation']: 130 | self.input_labels[cloud_split] += [sub_labels] 131 | 132 | # Get validation and test re_projection indices 133 | print('\nPreparing reprojection indices for validation and test') 134 | 135 | for i, file_path in enumerate(files): 136 | 137 | # get cloud name and split 138 | cloud_name = file_path.split('/')[-1][:-4] 139 | 140 | # Validation projection and labels 141 | if file_path in self.val_files: 142 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 143 | with open(proj_file, 'rb') as f: 144 | proj_idx, labels = pickle.load(f) 145 | self.val_proj += [proj_idx] 146 | self.val_labels += [labels] 147 | 148 | # Test projection 149 | if file_path in self.test_files: 150 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 151 | with open(proj_file, 'rb') as f: 152 | proj_idx, labels = pickle.load(f) 153 | self.test_proj += [proj_idx] 154 | self.test_labels += [labels] 155 | print('finished') 156 | return 157 | 158 | # Generate the input data flow 159 | def get_batch_gen(self, split): 160 | if split == 'training': 161 | num_per_epoch = cfg.train_steps * cfg.batch_size 162 | elif split == 'validation': 163 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 164 | elif split == 'test': 165 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 166 | 167 | # Reset possibility 168 | self.possibility[split] = [] 169 | self.min_possibility[split] = [] 170 | self.class_weight[split] = [] 171 | 172 | # Random initialize 173 | for i, tree in enumerate(self.input_trees[split]): 174 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 175 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 176 | 177 | if split != 'test': 178 | _, num_class_total = np.unique(np.hstack(self.input_labels[split]), return_counts=True) 179 | self.class_weight[split] += [np.squeeze([num_class_total / np.sum(num_class_total)], axis=0)] 180 | 181 | def spatially_regular_gen(): 182 | 183 | # Generator loop 184 | for i in range(num_per_epoch): # num_per_epoch 185 | 186 | # Choose the cloud with the lowest probability 187 | cloud_idx = int(np.argmin(self.min_possibility[split])) 188 | 189 | # choose the point with the minimum of possibility in the cloud as query point 190 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 191 | 192 | # Get all points within the cloud from tree structure 193 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 194 | 195 | # Center point of input region 196 | center_point = points[point_ind, :].reshape(1, -1) 197 | 198 | # Add noise to the center point 199 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 200 | pick_point = center_point + noise.astype(center_point.dtype) 201 | query_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 202 | 203 | # Shuffle index 204 | query_idx = DP.shuffle_idx(query_idx) 205 | 206 | # Get corresponding points and colors based on the index 207 | queried_pc_xyz = points[query_idx] 208 | queried_pc_xyz[:, 0:2] = queried_pc_xyz[:, 0:2] - pick_point[:, 0:2] 209 | queried_pc_colors = self.input_colors[split][cloud_idx][query_idx] 210 | if split == 'test': 211 | queried_pc_labels = np.zeros(queried_pc_xyz.shape[0]) 212 | queried_pt_weight = 1 213 | else: 214 | queried_pc_labels = self.input_labels[split][cloud_idx][query_idx] 215 | queried_pc_labels = np.array([self.label_to_idx[l] for l in queried_pc_labels]) 216 | queried_pt_weight = np.array([self.class_weight[split][0][n] for n in queried_pc_labels]) 217 | 218 | # Update the possibility of the selected points 219 | dists = np.sum(np.square((points[query_idx] - pick_point).astype(np.float32)), axis=1) 220 | delta = np.square(1 - dists / np.max(dists)) * queried_pt_weight 221 | self.possibility[split][cloud_idx][query_idx] += delta 222 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 223 | 224 | if True: 225 | yield (queried_pc_xyz, 226 | queried_pc_colors.astype(np.float32), 227 | queried_pc_labels, 228 | query_idx.astype(np.int32), 229 | np.array([cloud_idx], dtype=np.int32)) 230 | 231 | gen_func = spatially_regular_gen 232 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32) 233 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None]) 234 | return gen_func, gen_types, gen_shapes 235 | 236 | def get_tf_mapping(self): 237 | # Collect flat inputs 238 | def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 239 | batch_features = tf.map_fn(self.tf_augment_input, [batch_xyz, batch_features], dtype=tf.float32) 240 | input_points = [] 241 | input_neighbors = [] 242 | input_pools = [] 243 | input_up_samples = [] 244 | 245 | for i in range(cfg.num_layers): 246 | neigh_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) 247 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 248 | pool_i = neigh_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 249 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) 250 | input_points.append(batch_xyz) 251 | input_neighbors.append(neigh_idx) 252 | input_pools.append(pool_i) 253 | input_up_samples.append(up_i) 254 | batch_xyz = sub_points 255 | 256 | input_list = input_points + input_neighbors + input_pools + input_up_samples 257 | input_list += [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx] 258 | 259 | return input_list 260 | 261 | return tf_map 262 | 263 | # data augmentation 264 | @staticmethod 265 | def tf_augment_input(inputs): 266 | xyz = inputs[0] 267 | features = inputs[1] 268 | theta = tf.random_uniform((1,), minval=0, maxval=2 * np.pi) 269 | # Rotation matrices 270 | c, s = tf.cos(theta), tf.sin(theta) 271 | cs0 = tf.zeros_like(c) 272 | cs1 = tf.ones_like(c) 273 | R = tf.stack([c, -s, cs0, s, c, cs0, cs0, cs0, cs1], axis=1) 274 | stacked_rots = tf.reshape(R, (3, 3)) 275 | 276 | # Apply rotations 277 | transformed_xyz = tf.reshape(tf.matmul(xyz, stacked_rots), [-1, 3]) 278 | # Choose random scales for each example 279 | min_s = cfg.augment_scale_min 280 | max_s = cfg.augment_scale_max 281 | if cfg.augment_scale_anisotropic: 282 | s = tf.random_uniform((1, 3), minval=min_s, maxval=max_s) 283 | else: 284 | s = tf.random_uniform((1, 1), minval=min_s, maxval=max_s) 285 | 286 | symmetries = [] 287 | for i in range(3): 288 | if cfg.augment_symmetries[i]: 289 | symmetries.append(tf.round(tf.random_uniform((1, 1))) * 2 - 1) 290 | else: 291 | symmetries.append(tf.ones([1, 1], dtype=tf.float32)) 292 | s *= tf.concat(symmetries, 1) 293 | 294 | # Create N x 3 vector of scales to multiply with stacked_points 295 | stacked_scales = tf.tile(s, [tf.shape(transformed_xyz)[0], 1]) 296 | 297 | # Apply scales 298 | transformed_xyz = transformed_xyz * stacked_scales 299 | 300 | noise = tf.random_normal(tf.shape(transformed_xyz), stddev=cfg.augment_noise) 301 | transformed_xyz = transformed_xyz + noise 302 | rgb = features[:, :3] 303 | stacked_features = tf.concat([transformed_xyz, rgb], axis=-1) 304 | return stacked_features 305 | 306 | def init_input_pipeline(self): 307 | print('Initiating input pipelines') 308 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 309 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 310 | gen_function_val, _, _ = self.get_batch_gen('validation') 311 | gen_function_test, _, _ = self.get_batch_gen('test') 312 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 313 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 314 | self.test_data = tf.data.Dataset.from_generator(gen_function_test, gen_types, gen_shapes) 315 | 316 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 317 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 318 | self.batch_test_data = self.test_data.batch(cfg.val_batch_size) 319 | map_func = self.get_tf_mapping() 320 | 321 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 322 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 323 | self.batch_test_data = self.batch_test_data.map(map_func=map_func) 324 | 325 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 326 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 327 | self.batch_test_data = self.batch_test_data.prefetch(cfg.val_batch_size) 328 | 329 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 330 | self.flat_inputs = iter.get_next() 331 | self.train_init_op = iter.make_initializer(self.batch_train_data) 332 | self.val_init_op = iter.make_initializer(self.batch_val_data) 333 | self.test_init_op = iter.make_initializer(self.batch_test_data) 334 | 335 | 336 | if __name__ == '__main__': 337 | parser = argparse.ArgumentParser() 338 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 339 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 340 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 341 | FLAGS = parser.parse_args() 342 | 343 | GPU_ID = FLAGS.gpu 344 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 345 | os.environ['CUDA_VISIBLE_DEVICES'] = str(GPU_ID) 346 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 347 | 348 | Mode = FLAGS.mode 349 | dataset = Semantic3D() 350 | dataset.init_input_pipeline() 351 | 352 | if Mode == 'train': 353 | model = Network(dataset, cfg) 354 | model.train(dataset) 355 | elif Mode == 'test': 356 | cfg.saving = False 357 | model = Network(dataset, cfg) 358 | if FLAGS.model_path is not 'None': 359 | chosen_snap = FLAGS.model_path 360 | else: 361 | chosen_snapshot = -1 362 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 363 | chosen_folder = logs[-1] 364 | snap_path = join(chosen_folder, 'snapshots') 365 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 366 | chosen_step = np.sort(snap_steps)[-1] 367 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 368 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 369 | tester.test(model, dataset) 370 | 371 | else: 372 | ################## 373 | # Visualize data # 374 | ################## 375 | 376 | with tf.Session() as sess: 377 | sess.run(tf.global_variables_initializer()) 378 | sess.run(dataset.train_init_op) 379 | while True: 380 | flat_inputs = sess.run(dataset.flat_inputs) 381 | pc_xyz = flat_inputs[0] 382 | sub_pc_xyz = flat_inputs[1] 383 | labels = flat_inputs[21] 384 | Plot.draw_pc_sem_ins(pc_xyz[0, :, :], labels[0, :]) 385 | Plot.draw_pc_sem_ins(sub_pc_xyz[0, :, :], labels[0, 0:np.shape(sub_pc_xyz)[1]]) 386 | -------------------------------------------------------------------------------- /main_SemanticKITTI.py: -------------------------------------------------------------------------------- 1 | from helper_tool import DataProcessing as DP 2 | from helper_tool import ConfigSemanticKITTI as cfg 3 | from helper_tool import Plot 4 | from os.path import join 5 | from RandLANet import Network 6 | from tester_SemanticKITTI import ModelTester 7 | import tensorflow as tf 8 | import numpy as np 9 | import os, argparse, pickle 10 | 11 | 12 | class SemanticKITTI: 13 | def __init__(self, test_id): 14 | self.name = 'SemanticKITTI' 15 | self.dataset_path = '/data/semantic_kitti/dataset/sequences_0.06' 16 | self.label_to_names = {0: 'unlabeled', 17 | 1: 'car', 18 | 2: 'bicycle', 19 | 3: 'motorcycle', 20 | 4: 'truck', 21 | 5: 'other-vehicle', 22 | 6: 'person', 23 | 7: 'bicyclist', 24 | 8: 'motorcyclist', 25 | 9: 'road', 26 | 10: 'parking', 27 | 11: 'sidewalk', 28 | 12: 'other-ground', 29 | 13: 'building', 30 | 14: 'fence', 31 | 15: 'vegetation', 32 | 16: 'trunk', 33 | 17: 'terrain', 34 | 18: 'pole', 35 | 19: 'traffic-sign'} 36 | self.num_classes = len(self.label_to_names) 37 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 38 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 39 | self.ignored_labels = np.sort([0]) 40 | 41 | self.val_split = '08' 42 | 43 | self.seq_list = np.sort(os.listdir(self.dataset_path)) 44 | self.test_scan_number = str(test_id) 45 | self.train_list, self.val_list, self.test_list = DP.get_file_list(self.dataset_path, 46 | self.test_scan_number) 47 | self.train_list = DP.shuffle_list(self.train_list) 48 | self.val_list = DP.shuffle_list(self.val_list) 49 | 50 | self.possibility = [] 51 | self.min_possibility = [] 52 | 53 | # Generate the input data flow 54 | def get_batch_gen(self, split): 55 | if split == 'training': 56 | num_per_epoch = int(len(self.train_list) / cfg.batch_size) * cfg.batch_size 57 | path_list = self.train_list 58 | elif split == 'validation': 59 | num_per_epoch = int(len(self.val_list) / cfg.val_batch_size) * cfg.val_batch_size 60 | cfg.val_steps = int(len(self.val_list) / cfg.batch_size) 61 | path_list = self.val_list 62 | elif split == 'test': 63 | num_per_epoch = int(len(self.test_list) / cfg.val_batch_size) * cfg.val_batch_size * 4 64 | path_list = self.test_list 65 | for test_file_name in path_list: 66 | points = np.load(test_file_name) 67 | self.possibility += [np.random.rand(points.shape[0]) * 1e-3] 68 | self.min_possibility += [float(np.min(self.possibility[-1]))] 69 | 70 | def spatially_regular_gen(): 71 | # Generator loop 72 | for i in range(num_per_epoch): 73 | if split != 'test': 74 | cloud_ind = i 75 | pc_path = path_list[cloud_ind] 76 | pc, tree, labels = self.get_data(pc_path) 77 | # crop a small point cloud 78 | pick_idx = np.random.choice(len(pc), 1) 79 | selected_pc, selected_labels, selected_idx = self.crop_pc(pc, labels, tree, pick_idx) 80 | else: 81 | cloud_ind = int(np.argmin(self.min_possibility)) 82 | pick_idx = np.argmin(self.possibility[cloud_ind]) 83 | pc_path = path_list[cloud_ind] 84 | pc, tree, labels = self.get_data(pc_path) 85 | selected_pc, selected_labels, selected_idx = self.crop_pc(pc, labels, tree, pick_idx) 86 | 87 | # update the possibility of the selected pc 88 | dists = np.sum(np.square((selected_pc - pc[pick_idx]).astype(np.float32)), axis=1) 89 | delta = np.square(1 - dists / np.max(dists)) 90 | self.possibility[cloud_ind][selected_idx] += delta 91 | self.min_possibility[cloud_ind] = np.min(self.possibility[cloud_ind]) 92 | 93 | if True: 94 | yield (selected_pc.astype(np.float32), 95 | selected_labels.astype(np.int32), 96 | selected_idx.astype(np.int32), 97 | np.array([cloud_ind], dtype=np.int32)) 98 | 99 | gen_func = spatially_regular_gen 100 | gen_types = (tf.float32, tf.int32, tf.int32, tf.int32) 101 | gen_shapes = ([None, 3], [None], [None], [None]) 102 | 103 | return gen_func, gen_types, gen_shapes 104 | 105 | def get_data(self, file_path): 106 | seq_id = file_path.split('/')[-3] 107 | frame_id = file_path.split('/')[-1][:-4] 108 | kd_tree_path = join(self.dataset_path, seq_id, 'KDTree', frame_id + '.pkl') 109 | # Read pkl with search tree 110 | with open(kd_tree_path, 'rb') as f: 111 | search_tree = pickle.load(f) 112 | points = np.array(search_tree.data, copy=False) 113 | # Load labels 114 | if int(seq_id) >= 11: 115 | labels = np.zeros(np.shape(points)[0], dtype=np.uint8) 116 | else: 117 | label_path = join(self.dataset_path, seq_id, 'labels', frame_id + '.npy') 118 | labels = np.squeeze(np.load(label_path)) 119 | return points, search_tree, labels 120 | 121 | @staticmethod 122 | def crop_pc(points, labels, search_tree, pick_idx): 123 | # crop a fixed size point cloud for training 124 | center_point = points[pick_idx, :].reshape(1, -1) 125 | select_idx = search_tree.query(center_point, k=cfg.num_points)[1][0] 126 | select_idx = DP.shuffle_idx(select_idx) 127 | select_points = points[select_idx] 128 | select_labels = labels[select_idx] 129 | return select_points, select_labels, select_idx 130 | 131 | @staticmethod 132 | def get_tf_mapping2(): 133 | 134 | def tf_map(batch_pc, batch_label, batch_pc_idx, batch_cloud_idx): 135 | features = batch_pc 136 | input_points = [] 137 | input_neighbors = [] 138 | input_pools = [] 139 | input_up_samples = [] 140 | 141 | for i in range(cfg.num_layers): 142 | neighbour_idx = tf.py_func(DP.knn_search, [batch_pc, batch_pc, cfg.k_n], tf.int32) 143 | sub_points = batch_pc[:, :tf.shape(batch_pc)[1] // cfg.sub_sampling_ratio[i], :] 144 | pool_i = neighbour_idx[:, :tf.shape(batch_pc)[1] // cfg.sub_sampling_ratio[i], :] 145 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_pc, 1], tf.int32) 146 | input_points.append(batch_pc) 147 | input_neighbors.append(neighbour_idx) 148 | input_pools.append(pool_i) 149 | input_up_samples.append(up_i) 150 | batch_pc = sub_points 151 | 152 | input_list = input_points + input_neighbors + input_pools + input_up_samples 153 | input_list += [features, batch_label, batch_pc_idx, batch_cloud_idx] 154 | 155 | return input_list 156 | 157 | return tf_map 158 | 159 | def init_input_pipeline(self): 160 | print('Initiating input pipelines') 161 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 162 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 163 | gen_function_val, _, _ = self.get_batch_gen('validation') 164 | gen_function_test, _, _ = self.get_batch_gen('test') 165 | 166 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 167 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 168 | self.test_data = tf.data.Dataset.from_generator(gen_function_test, gen_types, gen_shapes) 169 | 170 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 171 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 172 | self.batch_test_data = self.test_data.batch(cfg.val_batch_size) 173 | 174 | map_func = self.get_tf_mapping2() 175 | 176 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 177 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 178 | self.batch_test_data = self.batch_test_data.map(map_func=map_func) 179 | 180 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 181 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 182 | self.batch_test_data = self.batch_test_data.prefetch(cfg.val_batch_size) 183 | 184 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 185 | self.flat_inputs = iter.get_next() 186 | self.train_init_op = iter.make_initializer(self.batch_train_data) 187 | self.val_init_op = iter.make_initializer(self.batch_val_data) 188 | self.test_init_op = iter.make_initializer(self.batch_test_data) 189 | 190 | 191 | if __name__ == '__main__': 192 | parser = argparse.ArgumentParser() 193 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 194 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 195 | parser.add_argument('--test_area', type=str, default='14', help='options: 08, 11,12,13,14,15,16,17,18,19,20,21') 196 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 197 | FLAGS = parser.parse_args() 198 | 199 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 200 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 201 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 202 | Mode = FLAGS.mode 203 | 204 | test_area = FLAGS.test_area 205 | dataset = SemanticKITTI(test_area) 206 | dataset.init_input_pipeline() 207 | 208 | if Mode == 'train': 209 | model = Network(dataset, cfg) 210 | model.train(dataset) 211 | elif Mode == 'test': 212 | cfg.saving = False 213 | model = Network(dataset, cfg) 214 | if FLAGS.model_path is not 'None': 215 | chosen_snap = FLAGS.model_path 216 | else: 217 | chosen_snapshot = -1 218 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 219 | chosen_folder = logs[-1] 220 | snap_path = join(chosen_folder, 'snapshots') 221 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 222 | chosen_step = np.sort(snap_steps)[-1] 223 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 224 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 225 | tester.test(model, dataset) 226 | else: 227 | ################## 228 | # Visualize data # 229 | ################## 230 | 231 | with tf.Session() as sess: 232 | sess.run(tf.global_variables_initializer()) 233 | sess.run(dataset.train_init_op) 234 | while True: 235 | flat_inputs = sess.run(dataset.flat_inputs) 236 | pc_xyz = flat_inputs[0] 237 | sub_pc_xyz = flat_inputs[1] 238 | labels = flat_inputs[17] 239 | Plot.draw_pc_sem_ins(pc_xyz[0, :, :], labels[0, :]) 240 | Plot.draw_pc_sem_ins(sub_pc_xyz[0, :, :], labels[0, 0:np.shape(sub_pc_xyz)[1]]) 241 | -------------------------------------------------------------------------------- /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_Semantic3D.py: -------------------------------------------------------------------------------- 1 | from os import makedirs 2 | from os.path import exists, join 3 | from helper_ply import read_ply, write_ply 4 | import tensorflow as tf 5 | import numpy as np 6 | import time 7 | 8 | 9 | def log_string(out_str, log_out): 10 | log_out.write(out_str + '\n') 11 | log_out.flush() 12 | print(out_str) 13 | 14 | 15 | class ModelTester: 16 | def __init__(self, model, dataset, restore_snap=None): 17 | # Tensorflow Saver definition 18 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 19 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 20 | 21 | # Create a session for running Ops on the Graph. 22 | on_cpu = False 23 | if on_cpu: 24 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 25 | else: 26 | c_proto = tf.ConfigProto() 27 | c_proto.gpu_options.allow_growth = True 28 | self.sess = tf.Session(config=c_proto) 29 | self.sess.run(tf.global_variables_initializer()) 30 | 31 | if restore_snap is not None: 32 | self.saver.restore(self.sess, restore_snap) 33 | print("Model restored from " + restore_snap) 34 | 35 | # Add a softmax operation for predictions 36 | self.prob_logits = tf.nn.softmax(model.logits) 37 | self.test_probs = [np.zeros((l.data.shape[0], model.config.num_classes), dtype=np.float16) 38 | for l in dataset.input_trees['test']] 39 | 40 | self.log_out = open('log_test_' + dataset.name + '.txt', 'a') 41 | 42 | def test(self, model, dataset, num_votes=100): 43 | 44 | # Smoothing parameter for votes 45 | test_smooth = 0.98 46 | 47 | # Initialise iterator with train data 48 | self.sess.run(dataset.test_init_op) 49 | 50 | # Test saving path 51 | saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 52 | test_path = join('test', saving_path.split('/')[-1]) 53 | makedirs(test_path) if not exists(test_path) else None 54 | makedirs(join(test_path, 'predictions')) if not exists(join(test_path, 'predictions')) else None 55 | makedirs(join(test_path, 'probs')) if not exists(join(test_path, 'probs')) else None 56 | 57 | ##################### 58 | # Network predictions 59 | ##################### 60 | 61 | step_id = 0 62 | epoch_id = 0 63 | last_min = -0.5 64 | 65 | while last_min < num_votes: 66 | 67 | try: 68 | ops = (self.prob_logits, 69 | model.labels, 70 | model.inputs['input_inds'], 71 | model.inputs['cloud_inds'],) 72 | 73 | stacked_probs, stacked_labels, point_idx, cloud_idx = self.sess.run(ops, {model.is_training: False}) 74 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, model.config.num_points, 75 | model.config.num_classes]) 76 | 77 | for j in range(np.shape(stacked_probs)[0]): 78 | probs = stacked_probs[j, :, :] 79 | inds = point_idx[j, :] 80 | c_i = cloud_idx[j][0] 81 | self.test_probs[c_i][inds] = test_smooth * self.test_probs[c_i][inds] + (1 - test_smooth) * probs 82 | step_id += 1 83 | log_string('Epoch {:3d}, step {:3d}. min possibility = {:.1f}'.format(epoch_id, step_id, np.min( 84 | dataset.min_possibility['test'])), self.log_out) 85 | 86 | except tf.errors.OutOfRangeError: 87 | 88 | # Save predicted cloud 89 | new_min = np.min(dataset.min_possibility['test']) 90 | log_string('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_id, new_min), self.log_out) 91 | 92 | if last_min + 4 < new_min: 93 | 94 | print('Saving clouds') 95 | 96 | # Update last_min 97 | last_min = new_min 98 | 99 | # Project predictions 100 | print('\nReproject Vote #{:d}'.format(int(np.floor(new_min)))) 101 | t1 = time.time() 102 | files = dataset.test_files 103 | i_test = 0 104 | for i, file_path in enumerate(files): 105 | # Get file 106 | points = self.load_evaluation_points(file_path) 107 | points = points.astype(np.float16) 108 | 109 | # Reproject probs 110 | probs = np.zeros(shape=[np.shape(points)[0], 8], dtype=np.float16) 111 | proj_index = dataset.test_proj[i_test] 112 | 113 | probs = self.test_probs[i_test][proj_index, :] 114 | 115 | # Insert false columns for ignored labels 116 | probs2 = probs 117 | for l_ind, label_value in enumerate(dataset.label_values): 118 | if label_value in dataset.ignored_labels: 119 | probs2 = np.insert(probs2, l_ind, 0, axis=1) 120 | 121 | # Get the predicted labels 122 | preds = dataset.label_values[np.argmax(probs2, axis=1)].astype(np.uint8) 123 | 124 | # Save plys 125 | cloud_name = file_path.split('/')[-1] 126 | 127 | # Save ascii preds 128 | ascii_name = join(test_path, 'predictions', dataset.ascii_files[cloud_name]) 129 | np.savetxt(ascii_name, preds, fmt='%d') 130 | log_string(ascii_name + 'has saved', self.log_out) 131 | i_test += 1 132 | 133 | t2 = time.time() 134 | print('Done in {:.1f} s\n'.format(t2 - t1)) 135 | self.sess.close() 136 | return 137 | 138 | self.sess.run(dataset.test_init_op) 139 | epoch_id += 1 140 | step_id = 0 141 | continue 142 | return 143 | 144 | @staticmethod 145 | def load_evaluation_points(file_path): 146 | data = read_ply(file_path) 147 | return np.vstack((data['x'], data['y'], data['z'])).T 148 | -------------------------------------------------------------------------------- /tester_SemanticKITTI.py: -------------------------------------------------------------------------------- 1 | from os import makedirs 2 | from os.path import exists, join, isfile, dirname, abspath 3 | from helper_tool import DataProcessing as DP 4 | from sklearn.metrics import confusion_matrix 5 | import tensorflow as tf 6 | import numpy as np 7 | import yaml 8 | import pickle 9 | 10 | BASE_DIR = dirname(abspath(__file__)) 11 | 12 | data_config = join(BASE_DIR, 'utils', 'semantic-kitti.yaml') 13 | DATA = yaml.safe_load(open(data_config, 'r')) 14 | remap_dict = DATA["learning_map_inv"] 15 | 16 | # make lookup table for mapping 17 | max_key = max(remap_dict.keys()) 18 | remap_lut = np.zeros((max_key + 100), dtype=np.int32) 19 | remap_lut[list(remap_dict.keys())] = list(remap_dict.values()) 20 | 21 | remap_dict_val = DATA["learning_map"] 22 | max_key = max(remap_dict_val.keys()) 23 | remap_lut_val = np.zeros((max_key + 100), dtype=np.int32) 24 | remap_lut_val[list(remap_dict_val.keys())] = list(remap_dict_val.values()) 25 | 26 | 27 | def log_out(out_str, f_out): 28 | f_out.write(out_str + '\n') 29 | f_out.flush() 30 | print(out_str) 31 | 32 | 33 | class ModelTester: 34 | def __init__(self, model, dataset, restore_snap=None): 35 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 36 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 37 | self.Log_file = open('log_test_' + dataset.name + '.txt', 'a') 38 | 39 | # Create a session for running Ops on the Graph. 40 | on_cpu = False 41 | if on_cpu: 42 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 43 | else: 44 | c_proto = tf.ConfigProto() 45 | c_proto.gpu_options.allow_growth = True 46 | self.sess = tf.Session(config=c_proto) 47 | self.sess.run(tf.global_variables_initializer()) 48 | 49 | # Name of the snapshot to restore to (None if you want to start from beginning) 50 | if restore_snap is not None: 51 | self.saver.restore(self.sess, restore_snap) 52 | print("Model restored from " + restore_snap) 53 | 54 | self.prob_logits = tf.nn.softmax(model.logits) 55 | self.test_probs = 0 56 | self.idx = 0 57 | 58 | def test(self, model, dataset): 59 | 60 | # Initialise iterator with train data 61 | self.sess.run(dataset.test_init_op) 62 | self.test_probs = [np.zeros(shape=[len(l), model.config.num_classes], dtype=np.float16) 63 | for l in dataset.possibility] 64 | 65 | test_path = join('test', 'sequences') 66 | makedirs(test_path) if not exists(test_path) else None 67 | save_path = join(test_path, dataset.test_scan_number, 'predictions') 68 | makedirs(save_path) if not exists(save_path) else None 69 | test_smooth = 0.98 70 | epoch_ind = 0 71 | 72 | while True: 73 | try: 74 | ops = (self.prob_logits, 75 | model.labels, 76 | model.inputs['input_inds'], 77 | model.inputs['cloud_inds']) 78 | stacked_probs, labels, point_inds, cloud_inds = self.sess.run(ops, {model.is_training: False}) 79 | if self.idx % 10 == 0: 80 | print('step ' + str(self.idx)) 81 | self.idx += 1 82 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, 83 | model.config.num_points, 84 | model.config.num_classes]) 85 | for j in range(np.shape(stacked_probs)[0]): 86 | probs = stacked_probs[j, :, :] 87 | inds = point_inds[j, :] 88 | c_i = cloud_inds[j][0] 89 | self.test_probs[c_i][inds] = test_smooth * self.test_probs[c_i][inds] + (1 - test_smooth) * probs 90 | 91 | except tf.errors.OutOfRangeError: 92 | new_min = np.min(dataset.min_possibility) 93 | log_out('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_ind, new_min), self.Log_file) 94 | if np.min(dataset.min_possibility) > 0.5: # 0.5 95 | log_out(' Min possibility = {:.1f}'.format(np.min(dataset.min_possibility)), self.Log_file) 96 | print('\nReproject Vote #{:d}'.format(int(np.floor(new_min)))) 97 | 98 | # For validation set 99 | num_classes = 19 100 | gt_classes = [0 for _ in range(num_classes)] 101 | positive_classes = [0 for _ in range(num_classes)] 102 | true_positive_classes = [0 for _ in range(num_classes)] 103 | val_total_correct = 0 104 | val_total_seen = 0 105 | 106 | for j in range(len(self.test_probs)): 107 | test_file_name = dataset.test_list[j] 108 | frame = test_file_name.split('/')[-1][:-4] 109 | proj_path = join(dataset.dataset_path, dataset.test_scan_number, 'proj') 110 | proj_file = join(proj_path, str(frame) + '_proj.pkl') 111 | if isfile(proj_file): 112 | with open(proj_file, 'rb') as f: 113 | proj_inds = pickle.load(f) 114 | probs = self.test_probs[j][proj_inds[0], :] 115 | pred = np.argmax(probs, 1) 116 | if dataset.test_scan_number == '08': 117 | label_path = join(dirname(dataset.dataset_path), 'sequences', dataset.test_scan_number, 118 | 'labels') 119 | label_file = join(label_path, str(frame) + '.label') 120 | labels = DP.load_label_kitti(label_file, remap_lut_val) 121 | invalid_idx = np.where(labels == 0)[0] 122 | labels_valid = np.delete(labels, invalid_idx) 123 | pred_valid = np.delete(pred, invalid_idx) 124 | labels_valid = labels_valid - 1 125 | correct = np.sum(pred_valid == labels_valid) 126 | val_total_correct += correct 127 | val_total_seen += len(labels_valid) 128 | conf_matrix = confusion_matrix(labels_valid, pred_valid, np.arange(0, num_classes, 1)) 129 | gt_classes += np.sum(conf_matrix, axis=1) 130 | positive_classes += np.sum(conf_matrix, axis=0) 131 | true_positive_classes += np.diagonal(conf_matrix) 132 | else: 133 | store_path = join(test_path, dataset.test_scan_number, 'predictions', 134 | str(frame) + '.label') 135 | pred = pred + 1 136 | pred = pred.astype(np.uint32) 137 | upper_half = pred >> 16 # get upper half for instances 138 | lower_half = pred & 0xFFFF # get lower half for semantics 139 | lower_half = remap_lut[lower_half] # do the remapping of semantics 140 | pred = (upper_half << 16) + lower_half # reconstruct full label 141 | pred = pred.astype(np.uint32) 142 | pred.tofile(store_path) 143 | log_out(str(dataset.test_scan_number) + ' finished', self.Log_file) 144 | if dataset.test_scan_number=='08': 145 | iou_list = [] 146 | for n in range(0, num_classes, 1): 147 | iou = true_positive_classes[n] / float( 148 | gt_classes[n] + positive_classes[n] - true_positive_classes[n]) 149 | iou_list.append(iou) 150 | mean_iou = sum(iou_list) / float(num_classes) 151 | 152 | log_out('eval accuracy: {}'.format(val_total_correct / float(val_total_seen)), self.Log_file) 153 | log_out('mean IOU:{}'.format(mean_iou), self.Log_file) 154 | 155 | mean_iou = 100 * mean_iou 156 | print('Mean IoU = {:.1f}%'.format(mean_iou)) 157 | s = '{:5.2f} | '.format(mean_iou) 158 | for IoU in iou_list: 159 | s += '{:5.2f} '.format(100 * IoU) 160 | print('-' * len(s)) 161 | print(s) 162 | print('-' * len(s) + '\n') 163 | self.sess.close() 164 | return 165 | self.sess.run(dataset.test_init_op) 166 | epoch_ind += 1 167 | continue 168 | -------------------------------------------------------------------------------- /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/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 | anno_paths = [line.rstrip() for line in open(join(BASE_DIR, 'meta/anno_paths.txt'))] 16 | anno_paths = [join(dataset_path, p) for p in anno_paths] 17 | 18 | gt_class = [x.rstrip() for x in open(join(BASE_DIR, 'meta/class_names.txt'))] 19 | gt_class2label = {cls: i for i, cls in enumerate(gt_class)} 20 | 21 | sub_grid_size = 0.04 22 | original_pc_folder = join(dirname(dataset_path), 'original_ply') 23 | sub_pc_folder = join(dirname(dataset_path), 'input_{:.3f}'.format(sub_grid_size)) 24 | os.mkdir(original_pc_folder) if not exists(original_pc_folder) else None 25 | os.mkdir(sub_pc_folder) if not exists(sub_pc_folder) else None 26 | out_format = '.ply' 27 | 28 | 29 | def convert_pc2ply(anno_path, save_path): 30 | """ 31 | Convert original dataset files to ply file (each line is XYZRGBL). 32 | We aggregated all the points from each instance in the room. 33 | :param anno_path: path to annotations. e.g. Area_1/office_2/Annotations/ 34 | :param save_path: path to save original point clouds (each line is XYZRGBL) 35 | :return: None 36 | """ 37 | data_list = [] 38 | 39 | for f in glob.glob(join(anno_path, '*.txt')): 40 | class_name = os.path.basename(f).split('_')[0] 41 | if class_name not in gt_class: # note: in some room there is 'staris' class.. 42 | class_name = 'clutter' 43 | pc = pd.read_csv(f, header=None, delim_whitespace=True).values 44 | labels = np.ones((pc.shape[0], 1)) * gt_class2label[class_name] 45 | data_list.append(np.concatenate([pc, labels], 1)) # Nx7 46 | 47 | pc_label = np.concatenate(data_list, 0) 48 | xyz_min = np.amin(pc_label, axis=0)[0:3] 49 | pc_label[:, 0:3] -= xyz_min 50 | 51 | xyz = pc_label[:, :3].astype(np.float32) 52 | colors = pc_label[:, 3:6].astype(np.uint8) 53 | labels = pc_label[:, 6].astype(np.uint8) 54 | write_ply(save_path, (xyz, colors, labels), ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 55 | 56 | # save sub_cloud and KDTree file 57 | sub_xyz, sub_colors, sub_labels = DP.grid_sub_sampling(xyz, colors, labels, sub_grid_size) 58 | sub_colors = sub_colors / 255.0 59 | sub_ply_file = join(sub_pc_folder, save_path.split('/')[-1][:-4] + '.ply') 60 | write_ply(sub_ply_file, [sub_xyz, sub_colors, sub_labels], ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 61 | 62 | search_tree = KDTree(sub_xyz) 63 | kd_tree_file = join(sub_pc_folder, str(save_path.split('/')[-1][:-4]) + '_KDTree.pkl') 64 | with open(kd_tree_file, 'wb') as f: 65 | pickle.dump(search_tree, f) 66 | 67 | proj_idx = np.squeeze(search_tree.query(xyz, return_distance=False)) 68 | proj_idx = proj_idx.astype(np.int32) 69 | proj_save = join(sub_pc_folder, str(save_path.split('/')[-1][:-4]) + '_proj.pkl') 70 | with open(proj_save, 'wb') as f: 71 | pickle.dump([proj_idx, labels], f) 72 | 73 | 74 | if __name__ == '__main__': 75 | # Note: there is an extra character in the v1.2 data in Area_5/hallway_6. It's fixed manually. 76 | for annotation_path in anno_paths: 77 | print(annotation_path) 78 | elements = str(annotation_path).split('/') 79 | out_file_name = elements[-3] + '_' + elements[-2] + out_format 80 | convert_pc2ply(annotation_path, join(original_pc_folder, out_file_name)) 81 | -------------------------------------------------------------------------------- /utils/data_prepare_semantic3d.py: -------------------------------------------------------------------------------- 1 | from sklearn.neighbors import KDTree 2 | from os.path import join, exists, dirname, abspath 3 | import numpy as np 4 | import os, glob, pickle 5 | import sys 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 | grid_size = 0.06 15 | dataset_path = '/data/semantic3d/original_data' 16 | original_pc_folder = join(dirname(dataset_path), 'original_ply') 17 | sub_pc_folder = join(dirname(dataset_path), 'input_{:.3f}'.format(grid_size)) 18 | os.mkdir(original_pc_folder) if not exists(original_pc_folder) else None 19 | os.mkdir(sub_pc_folder) if not exists(sub_pc_folder) else None 20 | 21 | for pc_path in glob.glob(join(dataset_path, '*.txt')): 22 | print(pc_path) 23 | file_name = pc_path.split('/')[-1][:-4] 24 | 25 | # check if it has already calculated 26 | if exists(join(sub_pc_folder, file_name + '_KDTree.pkl')): 27 | continue 28 | 29 | pc = DP.load_pc_semantic3d(pc_path) 30 | # check if label exists 31 | label_path = pc_path[:-4] + '.labels' 32 | if exists(label_path): 33 | labels = DP.load_label_semantic3d(label_path) 34 | full_ply_path = join(original_pc_folder, file_name + '.ply') 35 | 36 | #  Subsample to save space 37 | sub_points, sub_colors, sub_labels = DP.grid_sub_sampling(pc[:, :3].astype(np.float32), 38 | pc[:, 4:7].astype(np.uint8), labels, 0.01) 39 | sub_labels = np.squeeze(sub_labels) 40 | 41 | write_ply(full_ply_path, (sub_points, sub_colors, sub_labels), ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 42 | 43 | # save sub_cloud and KDTree file 44 | sub_xyz, sub_colors, sub_labels = DP.grid_sub_sampling(sub_points, sub_colors, sub_labels, grid_size) 45 | sub_colors = sub_colors / 255.0 46 | sub_labels = np.squeeze(sub_labels) 47 | sub_ply_file = join(sub_pc_folder, file_name + '.ply') 48 | write_ply(sub_ply_file, [sub_xyz, sub_colors, sub_labels], ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 49 | 50 | search_tree = KDTree(sub_xyz, leaf_size=50) 51 | kd_tree_file = join(sub_pc_folder, file_name + '_KDTree.pkl') 52 | with open(kd_tree_file, 'wb') as f: 53 | pickle.dump(search_tree, f) 54 | 55 | proj_idx = np.squeeze(search_tree.query(sub_points, return_distance=False)) 56 | proj_idx = proj_idx.astype(np.int32) 57 | proj_save = join(sub_pc_folder, file_name + '_proj.pkl') 58 | with open(proj_save, 'wb') as f: 59 | pickle.dump([proj_idx, labels], f) 60 | 61 | else: 62 | full_ply_path = join(original_pc_folder, file_name + '.ply') 63 | write_ply(full_ply_path, (pc[:, :3].astype(np.float32), pc[:, 4:7].astype(np.uint8)), 64 | ['x', 'y', 'z', 'red', 'green', 'blue']) 65 | 66 | # save sub_cloud and KDTree file 67 | sub_xyz, sub_colors = DP.grid_sub_sampling(pc[:, :3].astype(np.float32), pc[:, 4:7].astype(np.uint8), 68 | grid_size=grid_size) 69 | sub_colors = sub_colors / 255.0 70 | sub_ply_file = join(sub_pc_folder, file_name + '.ply') 71 | write_ply(sub_ply_file, [sub_xyz, sub_colors], ['x', 'y', 'z', 'red', 'green', 'blue']) 72 | labels = np.zeros(pc.shape[0], dtype=np.uint8) 73 | 74 | search_tree = KDTree(sub_xyz, leaf_size=50) 75 | kd_tree_file = join(sub_pc_folder, file_name + '_KDTree.pkl') 76 | with open(kd_tree_file, 'wb') as f: 77 | pickle.dump(search_tree, f) 78 | 79 | proj_idx = np.squeeze(search_tree.query(pc[:, :3].astype(np.float32), return_distance=False)) 80 | proj_idx = proj_idx.astype(np.int32) 81 | proj_save = join(sub_pc_folder, file_name + '_proj.pkl') 82 | with open(proj_save, 'wb') as f: 83 | pickle.dump([proj_idx, labels], f) 84 | -------------------------------------------------------------------------------- /utils/data_prepare_semantickitti.py: -------------------------------------------------------------------------------- 1 | import pickle, yaml, os, sys 2 | import numpy as np 3 | from os.path import join, exists, dirname, abspath 4 | from sklearn.neighbors import KDTree 5 | 6 | BASE_DIR = dirname(abspath(__file__)) 7 | ROOT_DIR = dirname(BASE_DIR) 8 | sys.path.append(BASE_DIR) 9 | sys.path.append(ROOT_DIR) 10 | from helper_tool import DataProcessing as DP 11 | 12 | data_config = os.path.join(BASE_DIR, 'semantic-kitti.yaml') 13 | DATA = yaml.safe_load(open(data_config, 'r')) 14 | remap_dict = DATA["learning_map"] 15 | max_key = max(remap_dict.keys()) 16 | remap_lut = np.zeros((max_key + 100), dtype=np.int32) 17 | remap_lut[list(remap_dict.keys())] = list(remap_dict.values()) 18 | 19 | grid_size = 0.06 20 | dataset_path = '/data/semantic_kitti/dataset/sequences' 21 | output_path = '/data/semantic_kitti/dataset/sequences' + '_' + str(grid_size) 22 | seq_list = np.sort(os.listdir(dataset_path)) 23 | 24 | for seq_id in seq_list: 25 | print('sequence' + seq_id + ' start') 26 | seq_path = join(dataset_path, seq_id) 27 | seq_path_out = join(output_path, seq_id) 28 | pc_path = join(seq_path, 'velodyne') 29 | pc_path_out = join(seq_path_out, 'velodyne') 30 | KDTree_path_out = join(seq_path_out, 'KDTree') 31 | os.makedirs(seq_path_out) if not exists(seq_path_out) else None 32 | os.makedirs(pc_path_out) if not exists(pc_path_out) else None 33 | os.makedirs(KDTree_path_out) if not exists(KDTree_path_out) else None 34 | 35 | if int(seq_id) < 11: 36 | label_path = join(seq_path, 'labels') 37 | label_path_out = join(seq_path_out, 'labels') 38 | os.makedirs(label_path_out) if not exists(label_path_out) else None 39 | scan_list = np.sort(os.listdir(pc_path)) 40 | for scan_id in scan_list: 41 | print(scan_id) 42 | points = DP.load_pc_kitti(join(pc_path, scan_id)) 43 | labels = DP.load_label_kitti(join(label_path, str(scan_id[:-4]) + '.label'), remap_lut) 44 | sub_points, sub_labels = DP.grid_sub_sampling(points, labels=labels, grid_size=grid_size) 45 | search_tree = KDTree(sub_points) 46 | KDTree_save = join(KDTree_path_out, str(scan_id[:-4]) + '.pkl') 47 | np.save(join(pc_path_out, scan_id)[:-4], sub_points) 48 | np.save(join(label_path_out, scan_id)[:-4], sub_labels) 49 | with open(KDTree_save, 'wb') as f: 50 | pickle.dump(search_tree, f) 51 | if seq_id == '08': 52 | proj_path = join(seq_path_out, 'proj') 53 | os.makedirs(proj_path) if not exists(proj_path) else None 54 | proj_inds = np.squeeze(search_tree.query(points, return_distance=False)) 55 | proj_inds = proj_inds.astype(np.int32) 56 | proj_save = join(proj_path, str(scan_id[:-4]) + '_proj.pkl') 57 | with open(proj_save, 'wb') as f: 58 | pickle.dump([proj_inds], f) 59 | else: 60 | proj_path = join(seq_path_out, 'proj') 61 | os.makedirs(proj_path) if not exists(proj_path) else None 62 | scan_list = np.sort(os.listdir(pc_path)) 63 | for scan_id in scan_list: 64 | print(scan_id) 65 | points = DP.load_pc_kitti(join(pc_path, scan_id)) 66 | sub_points = DP.grid_sub_sampling(points, grid_size=0.06) 67 | search_tree = KDTree(sub_points) 68 | proj_inds = np.squeeze(search_tree.query(points, return_distance=False)) 69 | proj_inds = proj_inds.astype(np.int32) 70 | KDTree_save = join(KDTree_path_out, str(scan_id[:-4]) + '.pkl') 71 | proj_save = join(proj_path, str(scan_id[:-4]) + '_proj.pkl') 72 | np.save(join(pc_path_out, scan_id)[:-4], sub_points) 73 | with open(KDTree_save, 'wb') as f: 74 | pickle.dump(search_tree, f) 75 | with open(proj_save, 'wb') as f: 76 | pickle.dump([proj_inds], f) 77 | -------------------------------------------------------------------------------- /utils/download_semantic3d.sh: -------------------------------------------------------------------------------- 1 | BASE_DIR=${1-/data/semantic3d/original_data} 2 | 3 | # Training data 4 | wget -c -N http://semantic3d.net/data/point-clouds/training1/bildstein_station1_xyz_intensity_rgb.7z -P $BASE_DIR 5 | wget -c -N http://semantic3d.net/data/point-clouds/training1/bildstein_station3_xyz_intensity_rgb.7z -P $BASE_DIR 6 | wget -c -N http://semantic3d.net/data/point-clouds/training1/bildstein_station5_xyz_intensity_rgb.7z -P $BASE_DIR 7 | wget -c -N http://semantic3d.net/data/point-clouds/training1/domfountain_station1_xyz_intensity_rgb.7z -P $BASE_DIR 8 | wget -c -N http://semantic3d.net/data/point-clouds/training1/domfountain_station2_xyz_intensity_rgb.7z -P $BASE_DIR/ 9 | wget -c -N http://semantic3d.net/data/point-clouds/training1/domfountain_station3_xyz_intensity_rgb.7z -P $BASE_DIR 10 | wget -c -N http://semantic3d.net/data/point-clouds/training1/neugasse_station1_xyz_intensity_rgb.7z -P $BASE_DIR 11 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station1_intensity_rgb.7z -P $BASE_DIR 12 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station2_intensity_rgb.7z -P $BASE_DIR 13 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station4_intensity_rgb.7z -P $BASE_DIR/ 14 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station5_intensity_rgb.7z -P $BASE_DIR 15 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station9_intensity_rgb.7z -P $BASE_DIR 16 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg28_station4_intensity_rgb.7z -P $BASE_DIR 17 | wget -c -N http://semantic3d.net/data/point-clouds/training1/untermaederbrunnen_station1_xyz_intensity_rgb.7z -P $BASE_DIR 18 | wget -c -N http://semantic3d.net/data/point-clouds/training1/untermaederbrunnen_station3_xyz_intensity_rgb.7z -P $BASE_DIR/ 19 | wget -c -N http://semantic3d.net/data/sem8_labels_training.7z -P $BASE_DIR 20 | 21 | 22 | # Test data 23 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/birdfountain_station1_xyz_intensity_rgb.7z -P $BASE_DIR 24 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/castleblatten_station1_intensity_rgb.7z -P $BASE_DIR 25 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/castleblatten_station5_xyz_intensity_rgb.7z -P $BASE_DIR 26 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/marketplacefeldkirch_station1_intensity_rgb.7z -P $BASE_DIR 27 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/marketplacefeldkirch_station4_intensity_rgb.7z -P $BASE_DIR 28 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/marketplacefeldkirch_station7_intensity_rgb.7z -P $BASE_DIR 29 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg27_station10_intensity_rgb.7z -P $BASE_DIR 30 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg27_station3_intensity_rgb.7z -P $BASE_DIR 31 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg27_station6_intensity_rgb.7z -P $BASE_DIR 32 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg27_station8_intensity_rgb.7z -P $BASE_DIR 33 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg28_station2_intensity_rgb.7z -P $BASE_DIR 34 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg28_station5_xyz_intensity_rgb.7z -P $BASE_DIR 35 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/stgallencathedral_station1_intensity_rgb.7z -P $BASE_DIR 36 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/stgallencathedral_station3_intensity_rgb.7z -P $BASE_DIR 37 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/stgallencathedral_station6_intensity_rgb.7z -P $BASE_DIR 38 | 39 | # reduced-8 40 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/MarketplaceFeldkirch_Station4_rgb_intensity-reduced.txt.7z -P $BASE_DIR 41 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/StGallenCathedral_station6_rgb_intensity-reduced.txt.7z -P $BASE_DIR 42 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/sg27_station10_rgb_intensity-reduced.txt.7z -P $BASE_DIR 43 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/sg28_Station2_rgb_intensity-reduced.txt.7z -P $BASE_DIR 44 | 45 | 46 | 47 | for entry in "$BASE_DIR"/* 48 | do 49 | 7z x "$entry" -o$(dirname "$entry") -y 50 | done 51 | 52 | mv $BASE_DIR/station1_xyz_intensity_rgb.txt $BASE_DIR/neugasse_station1_xyz_intensity_rgb.txt 53 | 54 | for entry in "$BASE_DIR"/*.7z 55 | do 56 | rm "$entry" 57 | done -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /utils/semantic-kitti.yaml: -------------------------------------------------------------------------------- 1 | # This file is covered by the LICENSE file in the root of this project. 2 | labels: 3 | 0 : "unlabeled" 4 | 1 : "outlier" 5 | 10: "car" 6 | 11: "bicycle" 7 | 13: "bus" 8 | 15: "motorcycle" 9 | 16: "on-rails" 10 | 18: "truck" 11 | 20: "other-vehicle" 12 | 30: "person" 13 | 31: "bicyclist" 14 | 32: "motorcyclist" 15 | 40: "road" 16 | 44: "parking" 17 | 48: "sidewalk" 18 | 49: "other-ground" 19 | 50: "building" 20 | 51: "fence" 21 | 52: "other-structure" 22 | 60: "lane-marking" 23 | 70: "vegetation" 24 | 71: "trunk" 25 | 72: "terrain" 26 | 80: "pole" 27 | 81: "traffic-sign" 28 | 99: "other-object" 29 | 252: "moving-car" 30 | 253: "moving-bicyclist" 31 | 254: "moving-person" 32 | 255: "moving-motorcyclist" 33 | 256: "moving-on-rails" 34 | 257: "moving-bus" 35 | 258: "moving-truck" 36 | 259: "moving-other-vehicle" 37 | color_map: # bgr 38 | 0 : [0, 0, 0] 39 | 1 : [0, 0, 255] 40 | 10: [245, 150, 100] 41 | 11: [245, 230, 100] 42 | 13: [250, 80, 100] 43 | 15: [150, 60, 30] 44 | 16: [255, 0, 0] 45 | 18: [180, 30, 80] 46 | 20: [255, 0, 0] 47 | 30: [30, 30, 255] 48 | 31: [200, 40, 255] 49 | 32: [90, 30, 150] 50 | 40: [255, 0, 255] 51 | 44: [255, 150, 255] 52 | 48: [75, 0, 75] 53 | 49: [75, 0, 175] 54 | 50: [0, 200, 255] 55 | 51: [50, 120, 255] 56 | 52: [0, 150, 255] 57 | 60: [170, 255, 150] 58 | 70: [0, 175, 0] 59 | 71: [0, 60, 135] 60 | 72: [80, 240, 150] 61 | 80: [150, 240, 255] 62 | 81: [0, 0, 255] 63 | 99: [255, 255, 50] 64 | 252: [245, 150, 100] 65 | 256: [255, 0, 0] 66 | 253: [200, 40, 255] 67 | 254: [30, 30, 255] 68 | 255: [90, 30, 150] 69 | 257: [250, 80, 100] 70 | 258: [180, 30, 80] 71 | 259: [255, 0, 0] 72 | content: # as a ratio with the total number of points 73 | 0: 0.018889854628292943 74 | 1: 0.0002937197336781505 75 | 10: 0.040818519255974316 76 | 11: 0.00016609538710764618 77 | 13: 2.7879693665067774e-05 78 | 15: 0.00039838616015114444 79 | 16: 0.0 80 | 18: 0.0020633612104619787 81 | 20: 0.0016218197275284021 82 | 30: 0.00017698551338515307 83 | 31: 1.1065903904919655e-08 84 | 32: 5.532951952459828e-09 85 | 40: 0.1987493871255525 86 | 44: 0.014717169549888214 87 | 48: 0.14392298360372 88 | 49: 0.0039048553037472045 89 | 50: 0.1326861944777486 90 | 51: 0.0723592229456223 91 | 52: 0.002395131480328884 92 | 60: 4.7084144280367186e-05 93 | 70: 0.26681502148037506 94 | 71: 0.006035012012626033 95 | 72: 0.07814222006271769 96 | 80: 0.002855498193863172 97 | 81: 0.0006155958086189918 98 | 99: 0.009923127583046915 99 | 252: 0.001789309418528068 100 | 253: 0.00012709999297008662 101 | 254: 0.00016059776092534436 102 | 255: 3.745553104802113e-05 103 | 256: 0.0 104 | 257: 0.00011351574470342043 105 | 258: 0.00010157861367183268 106 | 259: 4.3840131989471124e-05 107 | # classes that are indistinguishable from single scan or inconsistent in 108 | # ground truth are mapped to their closest equivalent 109 | learning_map: 110 | 0 : 0 # "unlabeled" 111 | 1 : 0 # "outlier" mapped to "unlabeled" --------------------------mapped 112 | 10: 1 # "car" 113 | 11: 2 # "bicycle" 114 | 13: 5 # "bus" mapped to "other-vehicle" --------------------------mapped 115 | 15: 3 # "motorcycle" 116 | 16: 5 # "on-rails" mapped to "other-vehicle" ---------------------mapped 117 | 18: 4 # "truck" 118 | 20: 5 # "other-vehicle" 119 | 30: 6 # "person" 120 | 31: 7 # "bicyclist" 121 | 32: 8 # "motorcyclist" 122 | 40: 9 # "road" 123 | 44: 10 # "parking" 124 | 48: 11 # "sidewalk" 125 | 49: 12 # "other-ground" 126 | 50: 13 # "building" 127 | 51: 14 # "fence" 128 | 52: 0 # "other-structure" mapped to "unlabeled" ------------------mapped 129 | 60: 9 # "lane-marking" to "road" ---------------------------------mapped 130 | 70: 15 # "vegetation" 131 | 71: 16 # "trunk" 132 | 72: 17 # "terrain" 133 | 80: 18 # "pole" 134 | 81: 19 # "traffic-sign" 135 | 99: 0 # "other-object" to "unlabeled" ----------------------------mapped 136 | 252: 1 # "moving-car" to "car" ------------------------------------mapped 137 | 253: 7 # "moving-bicyclist" to "bicyclist" ------------------------mapped 138 | 254: 6 # "moving-person" to "person" ------------------------------mapped 139 | 255: 8 # "moving-motorcyclist" to "motorcyclist" ------------------mapped 140 | 256: 5 # "moving-on-rails" mapped to "other-vehicle" --------------mapped 141 | 257: 5 # "moving-bus" mapped to "other-vehicle" -------------------mapped 142 | 258: 4 # "moving-truck" to "truck" --------------------------------mapped 143 | 259: 5 # "moving-other"-vehicle to "other-vehicle" ----------------mapped 144 | learning_map_inv: # inverse of previous map 145 | 0: 0 # "unlabeled", and others ignored 146 | 1: 10 # "car" 147 | 2: 11 # "bicycle" 148 | 3: 15 # "motorcycle" 149 | 4: 18 # "truck" 150 | 5: 20 # "other-vehicle" 151 | 6: 30 # "person" 152 | 7: 31 # "bicyclist" 153 | 8: 32 # "motorcyclist" 154 | 9: 40 # "road" 155 | 10: 44 # "parking" 156 | 11: 48 # "sidewalk" 157 | 12: 49 # "other-ground" 158 | 13: 50 # "building" 159 | 14: 51 # "fence" 160 | 15: 70 # "vegetation" 161 | 16: 71 # "trunk" 162 | 17: 72 # "terrain" 163 | 18: 80 # "pole" 164 | 19: 81 # "traffic-sign" 165 | learning_ignore: # Ignore classes 166 | 0: True # "unlabeled", and others ignored 167 | 1: False # "car" 168 | 2: False # "bicycle" 169 | 3: False # "motorcycle" 170 | 4: False # "truck" 171 | 5: False # "other-vehicle" 172 | 6: False # "person" 173 | 7: False # "bicyclist" 174 | 8: False # "motorcyclist" 175 | 9: False # "road" 176 | 10: False # "parking" 177 | 11: False # "sidewalk" 178 | 12: False # "other-ground" 179 | 13: False # "building" 180 | 14: False # "fence" 181 | 15: False # "vegetation" 182 | 16: False # "trunk" 183 | 17: False # "terrain" 184 | 18: False # "pole" 185 | 19: False # "traffic-sign" 186 | split: # sequence numbers 187 | train: 188 | - 0 189 | - 1 190 | - 2 191 | - 3 192 | - 4 193 | - 5 194 | - 6 195 | - 7 196 | - 9 197 | - 10 198 | valid: 199 | - 8 200 | test: 201 | - 11 202 | - 12 203 | - 13 204 | - 14 205 | - 15 206 | - 16 207 | - 17 208 | - 18 209 | - 19 210 | - 20 211 | - 21 212 | --------------------------------------------------------------------------------