├── LICENSE ├── README.md ├── cfgs ├── modelnet40 │ └── pointstack.yaml ├── partnormal │ └── pointstack.yaml └── scanobjectnn │ └── pointstack.yaml ├── core ├── builders.py ├── datasets │ ├── __init__.py │ ├── dataset_template.py │ ├── modelnet40.py │ ├── partnormal.py │ └── scanobjectnn.py └── networks │ ├── encoders │ ├── __init__.py │ ├── pointMLP.py │ ├── pointstack_cls_encoder.py │ └── pointstack_seg_encoder.py │ ├── heads │ ├── __init__.py │ ├── linear_classifier.py │ └── linear_segmentator.py │ └── networks │ ├── __init__.py │ ├── network_template.py │ ├── pointmlp.py │ └── pointstack.py ├── docs └── overview.png ├── pointnet2_ops_lib ├── MANIFEST.in ├── build │ ├── lib.linux-x86_64-3.7 │ │ └── pointnet2_ops │ │ │ ├── __init__.py │ │ │ ├── _ext-src │ │ │ ├── include │ │ │ │ ├── ball_query.h │ │ │ │ ├── cuda_utils.h │ │ │ │ ├── group_points.h │ │ │ │ ├── interpolate.h │ │ │ │ ├── sampling.h │ │ │ │ └── utils.h │ │ │ └── src │ │ │ │ ├── ball_query.cpp │ │ │ │ ├── ball_query_gpu.cu │ │ │ │ ├── bindings.cpp │ │ │ │ ├── group_points.cpp │ │ │ │ ├── group_points_gpu.cu │ │ │ │ ├── interpolate.cpp │ │ │ │ ├── interpolate_gpu.cu │ │ │ │ ├── sampling.cpp │ │ │ │ └── sampling_gpu.cu │ │ │ ├── _ext.cpython-37m-x86_64-linux-gnu.so │ │ │ ├── _version.py │ │ │ ├── pointnet2_modules.py │ │ │ └── pointnet2_utils.py │ └── temp.linux-x86_64-3.7 │ │ └── pointnet2_ops │ │ └── _ext-src │ │ └── src │ │ ├── ball_query.o │ │ ├── ball_query_gpu.o │ │ ├── bindings.o │ │ ├── group_points.o │ │ ├── group_points_gpu.o │ │ ├── interpolate.o │ │ ├── interpolate_gpu.o │ │ ├── sampling.o │ │ └── sampling_gpu.o ├── pointnet2_ops.egg-info │ ├── PKG-INFO │ ├── SOURCES.txt │ ├── dependency_links.txt │ ├── requires.txt │ └── top_level.txt ├── pointnet2_ops │ ├── __init__.py │ ├── _ext-src │ │ ├── include │ │ │ ├── ball_query.h │ │ │ ├── cuda_utils.h │ │ │ ├── group_points.h │ │ │ ├── interpolate.h │ │ │ ├── sampling.h │ │ │ └── utils.h │ │ └── src │ │ │ ├── ball_query.cpp │ │ │ ├── ball_query_gpu.cu │ │ │ ├── bindings.cpp │ │ │ ├── group_points.cpp │ │ │ ├── group_points_gpu.cu │ │ │ ├── interpolate.cpp │ │ │ ├── interpolate_gpu.cu │ │ │ ├── sampling.cpp │ │ │ └── sampling_gpu.cu │ ├── _version.py │ ├── pointnet2_modules.py │ └── pointnet2_utils.py └── setup.py ├── requirements.txt ├── test.py ├── train.py ├── utils ├── runtime_utils.py └── vis_utils.py └── visualize_part.py /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # PointStack 2 | 3 |

4 | overview 5 |

6 | 7 | This repository provides the official PyTorch implementation for the following paper: 8 | 9 | **Advanced Feature Learning on Point Clouds using Multi-resolution Features and Learnable Pooling**
10 | [Kevin Tirta Wijaya](https://www.ktirta.xyz), [Dong-Hee Paek](http://ave.kaist.ac.kr/bbs/board.php?bo_table=sub1_2&wr_id=5), and [Seung-Hyun Kong](http://ave.kaist.ac.kr/)
11 | [\[**arXiv**\]](https://arxiv.org/abs/2205.09962) 12 | > **Abstract:** *Existing point cloud feature learning networks often incorporate sequences of sampling, neighborhood grouping, neighborhood-wise feature learning, and feature aggregation to learn high-semantic point features that represent the global context of a point cloud. 13 | Unfortunately, such a process may result in a substantial loss of granular information due to the sampling operation. 14 | Moreover, the widely-used max-pooling feature aggregation may exacerbate the loss since it completely neglects information from non-maximum point features. 15 | Due to the compounded loss of information concerning granularity and non-maximum point features, the resulting high-semantic point features from existing networks could be insufficient to represent the local context of a point cloud, which in turn may hinder the network in distinguishing fine shapes. 16 | To cope with this problem, we propose a novel point cloud feature learning network, PointStack, using multi-resolution feature learning and learnable pooling (LP). 17 | The multi-resolution feature learning is realized by aggregating point features of various resolutions in the multiple layers, so that the final point features contain both high-semantic and high-resolution information. 18 | On the other hand, the LP is used as a generalized pooling function that calculates the weighted sum of multi-resolution point features through the attention mechanism with learnable queries, in order to extract all possible information from all available point features. 19 | Consequently, PointStack is capable of extracting high-semantic point features with minimal loss of information concerning granularity and non-maximum point features. 20 | Therefore, the final aggregated point features can effectively represent both global and local contexts of a point cloud. 21 | In addition, both the global structure and the local shape details of a point cloud can be well comprehended by the network head, which enables PointStack to advance the state-of-the-art of feature learning on point clouds. 22 | Specifically, PointStack outperforms various existing feature learning networks for shape classification and part segmentation on the ScanObjectNN and ShapeNetPart datasets.* 23 | 24 | ## Preparations 25 | 26 | **To install the requirements**: 27 | 28 | ```setup 29 | # 1. Create new environment 30 | conda create -n python=3.7 31 | 32 | # 2. Install PyTorch with corresponding cudatoolkit, for example, 33 | pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113 34 | 35 | # 3. Install other requirements 36 | pip install -r requirements.txt 37 | ``` 38 | 39 | **To prepare the dataset**: 40 | 41 | 1. Create empty directories for 'modelnet40', 'partnormal', and 'scanobjectnn' inside './data' directory. 42 | 2. Download the corresponding dataset for [Modelnet40](https://github.com/AnTao97/PointCloudDatasets), [ShapeNetPart](https://shapenet.cs.stanford.edu/media/shapenetcore_partanno_segmentation_benchmark_v0_normal.zip), and [ScanObjectNN](https://hkust-vgd.github.io/scanobjectnn/). 43 | 3. Create 'data' and 'experiments' directories on the workspace, and organize as follows, 44 | 45 | ``` 46 | ./ 47 | |-- cfgs 48 | |-- core 49 | |-- data 50 | |-- modelnet40 51 | |-- .json files 52 | |-- .txt files 53 | |-- partnormal 54 | |-- dirs 55 | |-- synsetoffset2category.txt 56 | |-- scanobjectnn 57 | |-- dirs 58 | |-- experiments 59 | ``` 60 | 61 | ## Training 62 | 63 | To train the network, run this command: 64 | 65 | ```train 66 | python train.py --cfg_file --exp_name --val_steps 67 | ``` 68 | 69 | ## Evaluation 70 | 71 | To evaluate the network with pre-trained weights, run: 72 | 73 | ```eval 74 | python test.py --cfg_file --ckpt 75 | ``` 76 | 77 | ## Results 78 | 79 | Our pretrained model achieves the following performances on : 80 | 81 | ### [3D Point Cloud Shape Classification on ModelNet40](https://paperswithcode.com/sota/3d-point-cloud-classification-on-modelnet40) 82 | 83 | | Model name | Overall Accuracy | Class Mean Accuracy | 84 | | ------------------ |---------------- | -------------- | 85 | | [PointStack](https://drive.google.com/file/d/1Emw8kR48htvPSNZ7e2UjO1CKy8W857s6/view?usp=sharing) | 93.3 | 89.6% | 86 | 87 | ### [3D Point Cloud Shape Classification on ScanObjectNN](https://paperswithcode.com/sota/3d-point-cloud-classification-on-scanobjectnn) 88 | 89 | | Model name | Overall Accuracy | Class Mean Accuracy | 90 | | ------------------ |---------------- | -------------- | 91 | | [PointStack](https://drive.google.com/file/d/1XTfYSkc0m4GKEhcV0wLAsb8GJ3-hBaiz/view?usp=sharing) | 87.2% | 86.2% | 92 | 93 | ### [3D Part Segmentation on ShapeNetPart](https://paperswithcode.com/sota/3d-part-segmentation-on-shapenet-part) 94 | 95 | | Model name | Instance mIoU | 96 | | ------------------ |---------------- | 97 | | [PointStack](https://drive.google.com/file/d/1Gab0_Cmdc-QFDgdMnWNMCP1NnE4_Uex1/view?usp=sharing) | 87.2% | 98 | -------------------------------------------------------------------------------- /cfgs/modelnet40/pointstack.yaml: -------------------------------------------------------------------------------- 1 | RANDOM_SEED: 0 2 | DATASET: 3 | NAME: ModelNet40 4 | NUM_CLASS: 40 5 | NUM_POINTS: 1024 6 | IS_SEGMENTATION: False 7 | USE_AUG_JIT: False 8 | USE_AUG_ROT: False 9 | USE_AUG_TRANS: True 10 | USE_RANDOM_SHUFFLE: True 11 | 12 | NETWORK: 13 | NAME: PointStack 14 | 15 | ENCODER: 16 | NAME: PointStackCls 17 | 18 | NUM_POINTS: 1024 19 | EMBED_DIM: 64 20 | GROUPS: 1 21 | RES_EXP: 1.0 22 | 23 | DIM_EXP: [2, 2, 2, 2] 24 | PRE_BLOCKS: [2, 2, 2, 2] 25 | POS_BLOCKS: [2, 2, 2, 2] 26 | K_NEIGHBORS: [24, 24, 24, 24] 27 | REDUCERS: [2, 2, 2, 2] 28 | 29 | HEAD: 30 | CLASSIFIER: 31 | NAME: LinearClassifier 32 | IN_CHANNELS: [4096] 33 | ACT: relu 34 | DIMS: [512, 256] 35 | 36 | OPTIMIZER: 37 | NAME: SGD 38 | MAX_EPOCH: 300 39 | BATCH_SIZE: 48 40 | GRAD_ACCUMULATION: 1 41 | GRAD_CLIP: 100 42 | LR: 0.01 43 | MIN_LR: 0.005 44 | MOMENTUM: 0.9 45 | NESTEROV: True 46 | WEIGHT_DECAY: 0.0002 47 | BETAS: [0.9, 0.999] # For Adam 48 | SCHEDULER: cosine_annealing 49 | WARM_RESTART_EVERY: 300 -------------------------------------------------------------------------------- /cfgs/partnormal/pointstack.yaml: -------------------------------------------------------------------------------- 1 | RANDOM_SEED: 2 2 | DATASET: 3 | NAME: PartNormal 4 | NUM_CLASS: 50 5 | NUM_POINTS: 2048 6 | IS_SEGMENTATION: True 7 | USE_AUG_JIT: False 8 | USE_AUG_ROT: False 9 | USE_AUG_TRANS: False 10 | USE_RANDOM_SHUFFLE: True 11 | 12 | NETWORK: 13 | NAME: PointStack 14 | 15 | ENCODER: 16 | NAME: PointStackSeg 17 | NUM_POINTS: 2048 18 | EMBED_DIM: 64 19 | GROUPS: 1 20 | RES_EXP: 1.0 21 | 22 | DIM_EXP: [2, 2, 2, 2] 23 | PRE_BLOCKS: [2, 2, 2, 2] 24 | POS_BLOCKS: [2, 2, 2, 2] 25 | K_NEIGHBORS: [24, 24, 24, 24] 26 | REDUCERS: [2, 2, 2, 2] 27 | 28 | HEAD: 29 | SEGMENTATOR: 30 | NAME: LinearSegmentator 31 | IN_CHANNELS: [4288] 32 | ACT: relu 33 | DIMS: [512, 256] 34 | 35 | OPTIMIZER: 36 | NAME: SGD 37 | MAX_EPOCH: 400 38 | BATCH_SIZE: 24 39 | GRAD_ACCUMULATION: 1 40 | GRAD_CLIP: 100 41 | LR: 0.01 42 | MIN_LR: 0.0001 43 | MOMENTUM: 0.9 44 | NESTEROV: True 45 | WEIGHT_DECAY: 0.0002 46 | BETAS: [0.9, 0.999] # For Adam 47 | SCHEDULER: cosine_annealing 48 | WARM_RESTART_EVERY: 400 -------------------------------------------------------------------------------- /cfgs/scanobjectnn/pointstack.yaml: -------------------------------------------------------------------------------- 1 | RANDOM_SEED: 2 2 | DATASET: 3 | NAME: ScanObjectNN 4 | DIRSPLIT: main # For ScanObjectNN, can be {main, supplementary, all} 5 | AUGSPLIT: PB_T50_RS # For ScanObjectNN, can be {OBJ_ONLY, PB_T25, PB_T25_R, PB_T50_R, PB_T50_RS} 6 | NUM_CLASS: 15 7 | NUM_POINTS: 1024 8 | IS_SEGMENTATION: False 9 | USE_AUG_JIT: False 10 | USE_AUG_ROT: True 11 | USE_AUG_TRANS: True 12 | USE_RANDOM_SHUFFLE: True 13 | LIMIT_TO_MODELNET40: False 14 | 15 | NETWORK: 16 | NAME: PointStack 17 | 18 | ENCODER: 19 | NAME: PointStackCls 20 | 21 | NUM_POINTS: 1024 22 | EMBED_DIM: 64 23 | GROUPS: 1 24 | RES_EXP: 1.0 25 | 26 | DIM_EXP: [2, 2, 2, 2] 27 | PRE_BLOCKS: [2, 2, 2, 2] 28 | POS_BLOCKS: [2, 2, 2, 2] 29 | K_NEIGHBORS: [24, 24, 24, 24] 30 | REDUCERS: [2, 2, 2, 2] 31 | 32 | HEAD: 33 | CLASSIFIER: 34 | NAME: LinearClassifier 35 | IN_CHANNELS: [4096] 36 | ACT: relu 37 | DIMS: [1024, 512, 256] 38 | 39 | OPTIMIZER: 40 | NAME: SGD 41 | MAX_EPOCH: 200 42 | BATCH_SIZE: 48 43 | GRAD_ACCUMULATION: 1 44 | GRAD_CLIP: 100 45 | LR: 0.01 46 | MIN_LR: 0.0001 47 | MOMENTUM: 0.9 48 | NESTEROV: True 49 | WEIGHT_DECAY: 0.0005 50 | BETAS: [0.9, 0.999] # For Adam 51 | SCHEDULER: cosine_annealing 52 | WARM_RESTART_EVERY: 200 -------------------------------------------------------------------------------- /core/builders.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | 4 | from . import datasets 5 | from .networks import networks 6 | 7 | def build_network(cfg, topology = None): 8 | net = networks.__all__[cfg.NETWORK.NAME](cfg, topology = topology) 9 | 10 | return net 11 | 12 | def build_dataset(cfg, split = 'train'): 13 | Dataset = datasets.__all__[cfg.DATASET.NAME](cfg, split = split) 14 | return Dataset 15 | 16 | def build_optimizer(cfg, params, data_loader_length, mode = None): 17 | opt_cfg = cfg.OPTIMIZER 18 | 19 | if mode == None: 20 | if (opt_cfg.NAME.lower() == 'adam'): 21 | opt = torch.optim.Adam(params, lr=opt_cfg.LR, betas=opt_cfg.BETAS, weight_decay=opt_cfg.WEIGHT_DECAY) 22 | elif (opt_cfg.NAME.lower() == 'sgd'): 23 | opt = torch.optim.SGD(params, lr=opt_cfg.LR, momentum=opt_cfg.MOMENTUM, weight_decay=opt_cfg.WEIGHT_DECAY, nesterov=opt_cfg.NESTEROV) 24 | 25 | if (opt_cfg.SCHEDULER is not None): 26 | if opt_cfg.SCHEDULER == 'cosine_annealing': 27 | scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0 = opt_cfg.WARM_RESTART_EVERY * data_loader_length, eta_min = opt_cfg.MIN_LR) 28 | else: 29 | 30 | if (getattr(opt_cfg, mode).NAME.lower() == 'adam'): 31 | opt = torch.optim.Adam(params, lr=getattr(opt_cfg, mode).LR, betas=getattr(opt_cfg, mode).BETAS, weight_decay=getattr(opt_cfg, mode).WEIGHT_DECAY) 32 | elif (getattr(opt_cfg, mode).NAME.lower() == 'sgd'): 33 | opt = torch.optim.SGD(params, lr=getattr(opt_cfg, mode).LR, momentum=getattr(opt_cfg, mode).MOMENTUM, weight_decay=getattr(opt_cfg, mode).WEIGHT_DECAY, nesterov=getattr(opt_cfg, mode).NESTEROV) 34 | 35 | if (getattr(opt_cfg, mode).SCHEDULER is not None): 36 | if getattr(opt_cfg, mode).SCHEDULER.lower()== 'cosine_annealing': 37 | scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max = opt_cfg.MAX_EPOCH//2) 38 | 39 | return opt, scheduler 40 | -------------------------------------------------------------------------------- /core/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .modelnet40 import ModelNet40 2 | from .scanobjectnn import ScanObjectNN 3 | from .partnormal import PartNormal 4 | 5 | __all__ = { 6 | 'ModelNet40': ModelNet40, 7 | 'ScanObjectNN': ScanObjectNN, 8 | 'PartNormal': PartNormal 9 | } -------------------------------------------------------------------------------- /core/datasets/dataset_template.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Based on codes originally written by An Tao (ta19@mails.tsinghua.edu.cn) 5 | """ 6 | 7 | import os 8 | import torch 9 | import json 10 | import h5py 11 | from glob import glob 12 | import numpy as np 13 | import torch.utils.data as data 14 | 15 | from utils.vis_utils import visualize_numpy 16 | 17 | shapenetpart_cat2id = {'airplane': 0, 'bag': 1, 'cap': 2, 'car': 3, 'chair': 4, 18 | 'earphone': 5, 'guitar': 6, 'knife': 7, 'lamp': 8, 'laptop': 9, 19 | 'motor': 10, 'mug': 11, 'pistol': 12, 'rocket': 13, 'skateboard': 14, 'table': 15} 20 | shapenetpart_seg_num = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3] 21 | shapenetpart_seg_start_index = [0, 4, 6, 8, 12, 16, 19, 22, 24, 28, 30, 36, 38, 41, 44, 47] 22 | 23 | 24 | class DatasetTemplate(data.Dataset): 25 | def __init__(self, cfg, class_choice=None, split='train', load_name=False, load_file=False, random_rotate=False, random_jitter=False, 26 | random_translate=False): 27 | assert cfg.DATASET.NUM_POINTS <= 2048 28 | if cfg.DATASET.USE_RANDOM_SHUFFLE: 29 | print('using random shuffle per point set') 30 | 31 | self.cfg = cfg 32 | self.data_dir = os.path.join(cfg.ROOT_DIR, 'data', cfg.DATASET.NAME.lower()) 33 | print('Loading dataset: ', self.data_dir) 34 | self.class_choice = class_choice 35 | self.split = split 36 | self.load_name = load_name 37 | self.load_file = load_file 38 | 39 | self.dataset_name = cfg.DATASET.NAME 40 | self.num_points = cfg.DATASET.NUM_POINTS 41 | self.segmentation = cfg.DATASET.IS_SEGMENTATION 42 | self.random_rotate = cfg.DATASET.USE_AUG_ROT if self.split in ['train'] else False 43 | self.random_jitter = cfg.DATASET.USE_AUG_JIT if self.split in ['train'] else False 44 | self.random_translate = cfg.DATASET.USE_AUG_TRANS if self.split in ['train'] else False 45 | 46 | def get_path(self, type): 47 | path_h5py = os.path.join(self.data_dir, '*%s*.h5'%type) 48 | self.path_h5py_all += glob(path_h5py) 49 | if self.load_name: 50 | path_json = os.path.join(self.data_dir, '%s*_id2name.json'%type) 51 | self.path_name_all += glob(path_json) 52 | if self.load_file: 53 | path_json = os.path.join(self.data_dir, '%s*_id2file.json'%type) 54 | self.path_file_all += glob(path_json) 55 | return 56 | 57 | def load_h5py(self, path): 58 | all_data = [] 59 | all_label = [] 60 | all_seg = [] 61 | for h5_name in path: 62 | f = h5py.File(h5_name, 'r+') 63 | 64 | data = f['data'][:].astype('float32') 65 | label = f['label'][:].astype('int64') 66 | if self.segmentation: 67 | if 'mask' in f.keys(): 68 | seg = f['mask'][:].astype('int64') 69 | seg[np.where(seg > -1)] = 1 # Foreground 70 | seg[np.where(seg == -1)] = 0 # Background 71 | else: 72 | seg = f['seg'][:].astype('int64') 73 | f.close() 74 | all_data.append(data) 75 | all_label.append(label) 76 | if self.segmentation: 77 | all_seg.append(seg) 78 | return all_data, all_label, all_seg 79 | 80 | def load_json(self, path): 81 | all_data = [] 82 | for json_name in path: 83 | j = open(json_name, 'r+') 84 | data = json.load(j) 85 | all_data += data 86 | return all_data 87 | 88 | def __getitem__(self, item): 89 | point_set = self.data[item][:self.num_points] 90 | # visualize_numpy(point_set) 91 | label = self.label[item] 92 | if self.segmentation: 93 | seg = self.seg[item][:self.num_points] 94 | 95 | if self.random_rotate: 96 | point_set, theta = self.rotate_pointcloud(point_set) 97 | if self.random_jitter: 98 | point_set, sigma, clip = self.jitter_pointcloud(point_set) 99 | if self.random_translate: 100 | point_set, xyz1, xyz2 = self.translate_pointcloud(point_set) 101 | 102 | if self.split == 'train': 103 | if self.cfg.DATASET.USE_RANDOM_SHUFFLE: 104 | p = np.random.permutation(len(point_set)) 105 | # np.random.shuffle(point_set) 106 | point_set = point_set[p] 107 | 108 | if self.segmentation: 109 | assert len(point_set) == len(seg) 110 | seg = seg[p] 111 | 112 | # convert numpy array to pytorch Tensor 113 | point_set = torch.from_numpy(point_set) 114 | label = torch.from_numpy(np.array([label]).astype(np.int64)) 115 | label = label.squeeze(0) 116 | 117 | data_dic = { 118 | 'points': point_set.cuda(), 119 | 'cls_id': label.cuda(), 120 | } 121 | 122 | if self.segmentation: 123 | seg = torch.from_numpy(seg).cuda() 124 | data_dic['seg_id'] = seg 125 | 126 | return data_dic 127 | 128 | def __len__(self): 129 | return self.data.shape[0] 130 | 131 | def __repr__(self): 132 | return ('Name=%s, Class Choice=%s, Num Points=%s, Split=%s, Random Rotate=%s, Random Jitter=%s, Random Translate=%s, )' 133 | % (repr(self.dataset_name), repr(self.class_choice), repr(self.num_points), 134 | repr(self.split), repr(self.random_rotate), repr(self.random_jitter), repr(self.random_translate))) 135 | 136 | def translate_pointcloud(self, pointcloud, xyz1 = None, xyz2 = None): 137 | if xyz1 == None: 138 | xyz1 = np.random.uniform(low=2./3., high=3./2., size=[3]) 139 | if xyz2 == None: 140 | xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3]) 141 | 142 | translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype('float32') 143 | return translated_pointcloud, xyz1, xyz2 144 | 145 | 146 | def jitter_pointcloud(self, pointcloud, sigma=0.01, clip=0.02): 147 | N, C = pointcloud.shape 148 | pointcloud += np.clip(sigma * np.random.randn(N, C), -1*clip, clip) 149 | return pointcloud, sigma, clip 150 | 151 | 152 | def rotate_pointcloud(self, pointcloud, theta = None): 153 | if theta == None: 154 | theta = np.pi*2 * np.random.rand() 155 | rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]]) 156 | pointcloud[:,[0,2]] = pointcloud[:,[0,2]].dot(rotation_matrix) # random rotation (x,z) 157 | return pointcloud, theta -------------------------------------------------------------------------------- /core/datasets/modelnet40.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Based on codes originally written by An Tao (ta19@mails.tsinghua.edu.cn) 5 | """ 6 | 7 | import os 8 | import torch 9 | import json 10 | import h5py 11 | from glob import glob 12 | import numpy as np 13 | 14 | from .dataset_template import DatasetTemplate 15 | 16 | class ModelNet40(DatasetTemplate): 17 | def __init__(self, cfg, class_choice=None, split='train', load_name=True, load_file=True, random_rotate=False, random_jitter=False, random_translate=False): 18 | super().__init__(cfg = cfg, class_choice=None, split=split, load_name=True, load_file=True, random_rotate=False, random_jitter=False, random_translate=False) 19 | 20 | self.path_h5py_all = [] 21 | self.path_name_all = [] 22 | self.path_file_all = [] 23 | 24 | if self.split in ['train','trainval','all']: 25 | self.get_path('train') 26 | if self.split in ['val', 'test', 'all']: 27 | self.get_path('test') 28 | 29 | self.path_h5py_all.sort() 30 | data, label, seg = self.load_h5py(self.path_h5py_all) 31 | 32 | if self.load_name or self.class_choice != None: 33 | self.path_name_all.sort() 34 | self.name = np.array(self.load_json(self.path_name_all)) # load label name 35 | 36 | if self.load_file: 37 | self.path_file_all.sort() 38 | self.file = np.array(self.load_json(self.path_file_all)) # load file name 39 | 40 | self.data = np.concatenate(data, axis=0) 41 | self.label = np.concatenate(label, axis=0) 42 | 43 | 44 | if self.segmentation: 45 | self.seg = np.concatenate(seg, axis=0) 46 | 47 | if self.class_choice != None: 48 | indices = (self.name == class_choice) 49 | self.data = self.data[indices] 50 | self.label = self.label[indices] 51 | self.name = self.name[indices] 52 | if self.load_file: 53 | self.file = self.file[indices] -------------------------------------------------------------------------------- /core/datasets/partnormal.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import os 4 | import torch 5 | import json 6 | import h5py 7 | from glob import glob 8 | import numpy as np 9 | 10 | from .dataset_template import DatasetTemplate 11 | 12 | class PartNormal(DatasetTemplate): 13 | def __init__(self, cfg, class_choice=None, split='train', load_name=True, load_file=True, random_rotate=False, random_jitter=False, random_translate=False): 14 | super().__init__(cfg = cfg, class_choice=None, split=split, load_name=True, load_file=True, random_rotate=False, random_jitter=False, random_translate=False) 15 | self.npoints = cfg.DATASET.NUM_POINTS 16 | self.root = self.data_dir 17 | self.catfile = os.path.join(self.root, 'synsetoffset2category.txt') 18 | self.cat = {} 19 | self.normalize = False 20 | 21 | self.split = split 22 | 23 | with open(self.catfile, 'r') as f: 24 | for line in f: 25 | ls = line.strip().split() 26 | self.cat[ls[0]] = ls[1] 27 | self.cat = {k: v for k, v in self.cat.items()} 28 | 29 | self.meta = {} 30 | with open(os.path.join(self.root, 'train_test_split', 'shuffled_train_file_list.json'), 'r') as f: 31 | train_ids = set([str(d.split('/')[2]) for d in json.load(f)]) 32 | with open(os.path.join(self.root, 'train_test_split', 'shuffled_val_file_list.json'), 'r') as f: 33 | val_ids = set([str(d.split('/')[2]) for d in json.load(f)]) 34 | with open(os.path.join(self.root, 'train_test_split', 'shuffled_test_file_list.json'), 'r') as f: 35 | test_ids = set([str(d.split('/')[2]) for d in json.load(f)]) 36 | for item in self.cat: 37 | self.meta[item] = [] 38 | dir_point = os.path.join(self.root, self.cat[item]) 39 | fns = sorted(os.listdir(dir_point)) 40 | 41 | if split == 'trainval': 42 | fns = [fn for fn in fns if ((fn[0:-4] in train_ids) or (fn[0:-4] in val_ids))] 43 | elif split == 'train': 44 | fns = [fn for fn in fns if fn[0:-4] in train_ids] 45 | elif split == 'val': 46 | fns = [fn for fn in fns if fn[0:-4] in val_ids] 47 | elif split == 'test': 48 | fns = [fn for fn in fns if fn[0:-4] in test_ids] 49 | else: 50 | print('Unknown split: %s. Exiting..' % (split)) 51 | exit(-1) 52 | 53 | for fn in fns: 54 | token = (os.path.splitext(os.path.basename(fn))[0]) 55 | self.meta[item].append(os.path.join(dir_point, token + '.txt')) 56 | 57 | self.datapath = [] 58 | for item in self.cat: 59 | for fn in self.meta[item]: 60 | self.datapath.append((item, fn)) 61 | 62 | self.classes = dict(zip(self.cat, range(len(self.cat)))) 63 | # Mapping from category ('Chair') to a list of int [10,11,12,13] as segmentation labels 64 | self.seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43], 65 | 'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 66 | 'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 67 | 'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 68 | 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]} 69 | 70 | self.cache = {} # from index to (point_set, cls, seg) tuple 71 | self.cache_size = 20000 72 | 73 | def __len__(self): 74 | return len(self.datapath) 75 | 76 | def __getitem__(self, index): 77 | if index in self.cache: 78 | point_set, normal, seg, cls = self.cache[index] 79 | else: 80 | fn = self.datapath[index] 81 | cat = self.datapath[index][0] 82 | cls = self.classes[cat] 83 | cls = np.array([cls]).astype(np.int32) 84 | data = np.loadtxt(fn[1]).astype(np.float32) 85 | point_set = data[:, 0:3] 86 | normal = data[:, 3:6] 87 | seg = data[:, -1].astype(np.int32) 88 | if len(self.cache) < self.cache_size: 89 | self.cache[index] = (point_set, normal, seg, cls) 90 | 91 | # if self.normalize: 92 | # point_set = pc_normalize(point_set) 93 | 94 | choice = np.random.choice(min(len(seg), self.npoints), self.npoints, replace=True) 95 | #choice = np.linspace(0, self.num_points, num=self.num_points).astype(int) 96 | # resample 97 | # note that the number of points in some points clouds is less than 2048, thus use random.choice 98 | # remember to use the same seed during train and test for a getting stable result 99 | 100 | point_set = point_set[choice, :] 101 | seg = seg[choice] 102 | normal = normal[choice, :] 103 | 104 | data_dic = { 105 | 'points' : torch.from_numpy(point_set).cuda(), 106 | 'seg_id' : torch.from_numpy(seg).cuda(), 107 | 'cls_tokens': torch.from_numpy(cls).cuda(), 108 | 'norms' : torch.from_numpy(normal).cuda() 109 | } 110 | 111 | return data_dic -------------------------------------------------------------------------------- /core/datasets/scanobjectnn.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Based on codes originally written by An Tao (ta19@mails.tsinghua.edu.cn) 5 | """ 6 | 7 | import os 8 | import torch 9 | import json 10 | import h5py 11 | from glob import glob 12 | import numpy as np 13 | 14 | from .dataset_template import DatasetTemplate 15 | 16 | CLS_NAME_TO_ID_DIC = {'bag': 0, 'bin': 1, 'box': 2, 'cabinet': 3, 'chair': 4, 'desk': 5, 'display': 6, 'door': 7, 'shelf': 8, 'table': 9, 'bed': 10, 'pillow': 11, 'sink': 12, 'sofa': 13, 'toilet': 14} 17 | CLS_ID_TO_NAME_DIC = {0: 'bag', 1: 'bin', 2: 'box', 3: 'cabinet', 4: 'chair', 5: 'desk', 6: 'display', 7: 'door', 8: 'shelf', 9: 'table', 10: 'bed', 11: 'pillow', 12: 'sink', 13: 'sofa', 14: 'toilet'} 18 | MODELNET40_LIMITER = [246, 246, 246, 246, 246, 246, 246, 246, 246, 246, 246, 246, 246, 246, 246] 19 | 20 | class ScanObjectNN(DatasetTemplate): 21 | def __init__(self, cfg, class_choice=None, split='train', load_name=True, load_file=True, random_rotate=False, random_jitter=False, random_translate=False): 22 | super().__init__(cfg = cfg, class_choice=None, split=split, load_name=True, load_file=True, random_rotate=False, random_jitter=False, random_translate=False) 23 | 24 | self.path_h5py_all = [] 25 | self.path_h5py_nobg = [] 26 | 27 | self.path_name_all = [] 28 | 29 | self.path_file_all = [] 30 | self.path_file_nobg = [] 31 | 32 | if self.cfg.DATASET.DIRSPLIT == 'main': 33 | self.dirsplit_all = ['main_split'] 34 | self.dirsplit_nobg = ['main_split_nobg'] 35 | elif self.cfg.DATASET.DIRSPLIT == 'supplementary': 36 | self.dirsplit_all = ['split1', 'split2', 'split3', 'split4'] 37 | self.dirsplit_nobg = ['split1_nobg', 'split2_nobg', 'split3_nobg', 'split4_nobg'] 38 | else: 39 | self.dirsplit_all = ['main_split', 'split1', 'split2', 'split3', 'split4'] 40 | self.dirsplit_nobg = ['main_split_nobg', 'split1_nobg', 'split2_nobg', 'split3_nobg', 'split4_nobg'] 41 | 42 | if self.cfg.DATASET.AUGSPLIT == 'OBJ_ONLY': 43 | self.augsplit = 'objectdataset' 44 | elif self.cfg.DATASET.AUGSPLIT == 'PB_T25': 45 | self.augsplit = 'objectdataset_augmented25_norot' 46 | elif self.cfg.DATASET.AUGSPLIT == 'PB_T25_R': 47 | self.augsplit = 'objectdataset_augmented25rot' 48 | elif self.cfg.DATASET.AUGSPLIT == 'PB_T50_R': 49 | self.augsplit = 'objectdataset_augmentedrot' 50 | elif self.cfg.DATASET.AUGSPLIT == 'PB_T50_RS': 51 | self.augsplit = 'objectdataset_augmentedrot_scale75' 52 | 53 | if self.split in ['train','trainval','all']: 54 | self.get_path('training', self.dirsplit_all, self.augsplit) 55 | if self.split in ['val', 'test', 'all']: 56 | self.get_path('test', self.dirsplit_all, self.augsplit) 57 | 58 | self.path_h5py_all.sort() 59 | self.path_h5py_nobg.sort() 60 | 61 | 62 | data, label, seg = self.load_h5py(self.path_h5py_all) 63 | 64 | if self.load_name or self.class_choice != None: 65 | self.path_name_all.sort() 66 | self.name = np.array(self.load_json(self.path_name_all)) # load label name 67 | 68 | if self.load_file: 69 | self.path_file_all.sort() 70 | self.file = np.array(self.load_json(self.path_file_all)) # load file name 71 | 72 | self.data = np.concatenate(data, axis=0) 73 | 74 | self.label = np.concatenate(label, axis=0) 75 | 76 | if self.segmentation: 77 | self.seg = np.concatenate(seg, axis=0) 78 | 79 | if self.class_choice != None: 80 | indices = (self.name == class_choice) 81 | self.data = self.data[indices] 82 | self.label = self.label[indices] 83 | 84 | self.name = self.name[indices] 85 | if self.load_file: 86 | self.file = self.file[indices] 87 | 88 | if (cfg.DATASET.LIMIT_TO_MODELNET40) and (split == 'train'): 89 | chosen_data = [] 90 | chosen_label = [] 91 | 92 | for cls_id in range(15): 93 | indices = np.where(self.label == cls_id) 94 | label = self.label[indices] 95 | data = self.data[indices] 96 | 97 | chosen_data.extend(data[:MODELNET40_LIMITER[cls_id]]) 98 | chosen_label.extend(label[:MODELNET40_LIMITER[cls_id]]) 99 | 100 | self.data = np.array(chosen_data) 101 | self.label = np.array(chosen_label) 102 | 103 | 104 | def get_path(self, type, dirsplits, augsplit): 105 | for dirsplit in dirsplits: 106 | path_h5py = os.path.join(self.data_dir, dirsplit, type+'*'+augsplit+'.h5') 107 | path_h5py_nobg = os.path.join(self.data_dir, dirsplit + '_nobg', type+'*'+augsplit+'.h5') 108 | 109 | self.path_h5py_all += glob(path_h5py) 110 | self.path_h5py_nobg += glob(path_h5py_nobg) 111 | 112 | return 113 | 114 | def __getitem__(self, item): 115 | point_set = self.data[item][:self.num_points] 116 | label = self.label[item] 117 | 118 | if self.segmentation: 119 | seg = self.seg[item][:self.num_points] 120 | 121 | pose_target = { 122 | 'theta': 0, 123 | 'dx': 0, 124 | 'dy': 0, 125 | 'dz': 0 126 | } 127 | 128 | if self.random_rotate: 129 | point_set, theta = self.rotate_pointcloud(point_set) 130 | pose_target['theta'] = theta 131 | 132 | if self.random_jitter: 133 | point_set, sigma, clip = self.jitter_pointcloud(point_set) 134 | 135 | if self.random_translate: 136 | point_set, xyz1, xyz2 = self.translate_pointcloud(point_set) 137 | pose_target['dx'], pose_target['dy'], pose_target['dz'] = xyz2 138 | 139 | if self.split == 'train': 140 | if self.cfg.DATASET.USE_RANDOM_SHUFFLE: 141 | p = np.random.permutation(len(point_set)) 142 | point_set = point_set[p] 143 | 144 | if self.segmentation: 145 | assert len(point_set) == len(seg) 146 | seg = seg[p] 147 | point_set = torch.from_numpy(point_set) 148 | 149 | label = torch.from_numpy(np.array([label]).astype(np.int64)) 150 | label = label.squeeze(0) 151 | 152 | data_dic = { 153 | 'points': point_set.cuda(), 154 | 'cls_id': label.cuda(), 155 | } 156 | 157 | if self.segmentation: 158 | seg = torch.from_numpy(seg).cuda() 159 | data_dic['seg_id'] = seg 160 | 161 | return data_dic -------------------------------------------------------------------------------- /core/networks/encoders/__init__.py: -------------------------------------------------------------------------------- 1 | from .pointMLP import PointMLP 2 | from .pointstack_seg_encoder import PointStackSeg 3 | from .pointstack_cls_encoder import PointStackCls 4 | 5 | 6 | __all__ = { 7 | 'PointMLP': PointMLP, 8 | 'PointStackSeg': PointStackSeg, 9 | 'PointStackCls': PointStackCls, 10 | } -------------------------------------------------------------------------------- /core/networks/encoders/pointMLP.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | from pointnet2_ops import pointnet2_utils 6 | 7 | 8 | def get_activation(activation): 9 | if activation.lower() == 'gelu': 10 | return nn.GELU() 11 | elif activation.lower() == 'rrelu': 12 | return nn.RReLU(inplace=True) 13 | elif activation.lower() == 'selu': 14 | return nn.SELU(inplace=True) 15 | elif activation.lower() == 'silu': 16 | return nn.SiLU(inplace=True) 17 | elif activation.lower() == 'hardswish': 18 | return nn.Hardswish(inplace=True) 19 | elif activation.lower() == 'leakyrelu': 20 | return nn.LeakyReLU(inplace=True) 21 | else: 22 | return nn.ReLU(inplace=True) 23 | 24 | 25 | def square_distance(src, dst): 26 | """ 27 | Calculate Euclid distance between each two points. 28 | src^T * dst = xn * xm + yn * ym + zn * zm； 29 | sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn; 30 | sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm; 31 | dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2 32 | = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst 33 | Input: 34 | src: source points, [B, N, C] 35 | dst: target points, [B, M, C] 36 | Output: 37 | dist: per-point square distance, [B, N, M] 38 | """ 39 | B, N, _ = src.shape 40 | _, M, _ = dst.shape 41 | dist = -2 * torch.matmul(src, dst.permute(0, 2, 1)) 42 | dist += torch.sum(src ** 2, -1).view(B, N, 1) 43 | dist += torch.sum(dst ** 2, -1).view(B, 1, M) 44 | return dist 45 | 46 | 47 | def index_points(points, idx): 48 | """ 49 | Input: 50 | points: input points data, [B, N, C] 51 | idx: sample index data, [B, S] 52 | Return: 53 | new_points:, indexed points data, [B, S, C] 54 | """ 55 | device = points.device 56 | B = points.shape[0] 57 | view_shape = list(idx.shape) 58 | view_shape[1:] = [1] * (len(view_shape) - 1) 59 | repeat_shape = list(idx.shape) 60 | repeat_shape[0] = 1 61 | batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape) 62 | new_points = points[batch_indices, idx, :] 63 | return new_points 64 | 65 | 66 | def farthest_point_sample(xyz, npoint): 67 | """ 68 | Input: 69 | xyz: pointcloud data, [B, N, 3] 70 | npoint: number of samples 71 | Return: 72 | centroids: sampled pointcloud index, [B, npoint] 73 | """ 74 | device = xyz.device 75 | B, N, C = xyz.shape 76 | centroids = torch.zeros(B, npoint, dtype=torch.long).to(device) 77 | distance = torch.ones(B, N).to(device) * 1e10 78 | farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device) 79 | batch_indices = torch.arange(B, dtype=torch.long).to(device) 80 | for i in range(npoint): 81 | centroids[:, i] = farthest 82 | centroid = xyz[batch_indices, farthest, :].view(B, 1, 3) 83 | dist = torch.sum((xyz - centroid) ** 2, -1) 84 | distance = torch.min(distance, dist) 85 | farthest = torch.max(distance, -1)[1] 86 | return centroids 87 | 88 | 89 | def query_ball_point(radius, nsample, xyz, new_xyz): 90 | """ 91 | Input: 92 | radius: local region radius 93 | nsample: max sample number in local region 94 | xyz: all points, [B, N, 3] 95 | new_xyz: query points, [B, S, 3] 96 | Return: 97 | group_idx: grouped points index, [B, S, nsample] 98 | """ 99 | device = xyz.device 100 | B, N, C = xyz.shape 101 | _, S, _ = new_xyz.shape 102 | group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1]) 103 | sqrdists = square_distance(new_xyz, xyz) 104 | group_idx[sqrdists > radius ** 2] = N 105 | group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample] 106 | group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample]) 107 | mask = group_idx == N 108 | group_idx[mask] = group_first[mask] 109 | return group_idx 110 | 111 | 112 | def knn_point(nsample, xyz, new_xyz): 113 | """ 114 | Input: 115 | nsample: max sample number in local region 116 | xyz: all points, [B, N, C] 117 | new_xyz: query points, [B, S, C] 118 | Return: 119 | group_idx: grouped points index, [B, S, nsample] 120 | """ 121 | sqrdists = square_distance(new_xyz, xyz) 122 | _, group_idx = torch.topk(sqrdists, nsample, dim=-1, largest=False, sorted=False) 123 | return group_idx 124 | 125 | 126 | class LocalGrouper(nn.Module): 127 | def __init__(self, channel, groups, kneighbors, use_xyz=True, normalize="center", **kwargs): 128 | """ 129 | Give xyz[b,p,3] and fea[b,p,d], return new_xyz[b,g,3] and new_fea[b,g,k,d] 130 | :param groups: groups number 131 | :param kneighbors: k-nerighbors 132 | :param kwargs: others 133 | """ 134 | super(LocalGrouper, self).__init__() 135 | self.groups = groups 136 | self.kneighbors = kneighbors 137 | self.use_xyz = use_xyz 138 | if normalize is not None: 139 | self.normalize = normalize.lower() 140 | else: 141 | self.normalize = None 142 | if self.normalize not in ["center", "anchor"]: 143 | print(f"Unrecognized normalize parameter (self.normalize), set to None. Should be one of [center, anchor].") 144 | self.normalize = None 145 | if self.normalize is not None: 146 | add_channel=3 if self.use_xyz else 0 147 | self.affine_alpha = nn.Parameter(torch.ones([1,1,1,channel + add_channel])) 148 | self.affine_beta = nn.Parameter(torch.zeros([1, 1, 1, channel + add_channel])) 149 | 150 | def forward(self, xyz, points): 151 | B, N, C = xyz.shape 152 | S = self.groups 153 | xyz = xyz.contiguous() # xyz [btach, points, xyz] 154 | 155 | # fps_idx = torch.multinomial(torch.linspace(0, N - 1, steps=N).repeat(B, 1).to(xyz.device), num_samples=self.groups, replacement=False).long() 156 | # fps_idx = farthest_point_sample(xyz, self.groups).long() 157 | fps_idx = pointnet2_utils.furthest_point_sample(xyz, self.groups).long() # [B, npoint] 158 | new_xyz = index_points(xyz, fps_idx) # [B, npoint, 3] 159 | new_points = index_points(points, fps_idx) # [B, npoint, d] 160 | 161 | idx = knn_point(self.kneighbors, xyz, new_xyz) 162 | # idx = query_ball_point(radius, nsample, xyz, new_xyz) 163 | grouped_xyz = index_points(xyz, idx) # [B, npoint, k, 3] 164 | grouped_points = index_points(points, idx) # [B, npoint, k, d] 165 | if self.use_xyz: 166 | grouped_points = torch.cat([grouped_points, grouped_xyz],dim=-1) # [B, npoint, k, d+3] 167 | if self.normalize is not None: 168 | if self.normalize =="center": 169 | mean = torch.mean(grouped_points, dim=2, keepdim=True) 170 | if self.normalize =="anchor": 171 | mean = torch.cat([new_points, new_xyz],dim=-1) if self.use_xyz else new_points 172 | mean = mean.unsqueeze(dim=-2) # [B, npoint, 1, d+3] 173 | std = torch.std((grouped_points-mean).reshape(B,-1),dim=-1,keepdim=True).unsqueeze(dim=-1).unsqueeze(dim=-1) 174 | grouped_points = (grouped_points-mean)/(std + 1e-5) 175 | grouped_points = self.affine_alpha*grouped_points + self.affine_beta 176 | 177 | new_points = torch.cat([grouped_points, new_points.view(B, S, 1, -1).repeat(1, 1, self.kneighbors, 1)], dim=-1) 178 | return new_xyz, new_points 179 | 180 | 181 | class ConvBNReLU1D(nn.Module): 182 | def __init__(self, in_channels, out_channels, kernel_size=1, bias=True, activation='relu'): 183 | super(ConvBNReLU1D, self).__init__() 184 | self.act = get_activation(activation) 185 | self.net = nn.Sequential( 186 | nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, bias=bias), 187 | nn.BatchNorm1d(out_channels), 188 | self.act 189 | ) 190 | 191 | def forward(self, x): 192 | return self.net(x) 193 | 194 | 195 | class ConvBNReLURes1D(nn.Module): 196 | def __init__(self, channel, kernel_size=1, groups=1, res_expansion=1.0, bias=True, activation='relu'): 197 | super(ConvBNReLURes1D, self).__init__() 198 | self.act = get_activation(activation) 199 | self.net1 = nn.Sequential( 200 | nn.Conv1d(in_channels=channel, out_channels=int(channel * res_expansion), 201 | kernel_size=kernel_size, groups=groups, bias=bias), 202 | nn.BatchNorm1d(int(channel * res_expansion)), 203 | self.act 204 | ) 205 | if groups > 1: 206 | self.net2 = nn.Sequential( 207 | nn.Conv1d(in_channels=int(channel * res_expansion), out_channels=channel, 208 | kernel_size=kernel_size, groups=groups, bias=bias), 209 | nn.BatchNorm1d(channel), 210 | self.act, 211 | nn.Conv1d(in_channels=channel, out_channels=channel, 212 | kernel_size=kernel_size, bias=bias), 213 | nn.BatchNorm1d(channel), 214 | ) 215 | else: 216 | self.net2 = nn.Sequential( 217 | nn.Conv1d(in_channels=int(channel * res_expansion), out_channels=channel, 218 | kernel_size=kernel_size, bias=bias), 219 | nn.BatchNorm1d(channel) 220 | ) 221 | 222 | def forward(self, x): 223 | return self.act(self.net2(self.net1(x)) + x) 224 | 225 | 226 | class PreExtraction(nn.Module): 227 | def __init__(self, channels, out_channels, blocks=1, groups=1, res_expansion=1, bias=True, 228 | activation='relu', use_xyz=True): 229 | """ 230 | input: [b,g,k,d]: output:[b,d,g] 231 | :param channels: 232 | :param blocks: 233 | """ 234 | super(PreExtraction, self).__init__() 235 | in_channels = 3+2*channels if use_xyz else 2*channels 236 | self.transfer = ConvBNReLU1D(in_channels, out_channels, bias=bias, activation=activation) 237 | operation = [] 238 | for _ in range(blocks): 239 | operation.append( 240 | ConvBNReLURes1D(out_channels, groups=groups, res_expansion=res_expansion, 241 | bias=bias, activation=activation) 242 | ) 243 | self.operation = nn.Sequential(*operation) 244 | 245 | def forward(self, x): 246 | b, n, s, d = x.size() # torch.Size([32, 512, 32, 6]) 247 | x = x.permute(0, 1, 3, 2) 248 | x = x.reshape(-1, d, s) 249 | x = self.transfer(x) 250 | batch_size, _, _ = x.size() 251 | x = self.operation(x) # [b, d, k] 252 | x = F.adaptive_max_pool1d(x, 1).view(batch_size, -1) 253 | x = x.reshape(b, n, -1).permute(0, 2, 1) 254 | return x 255 | 256 | 257 | class PosExtraction(nn.Module): 258 | def __init__(self, channels, blocks=1, groups=1, res_expansion=1, bias=True, activation='relu'): 259 | """ 260 | input[b,d,g]; output[b,d,g] 261 | :param channels: 262 | :param blocks: 263 | """ 264 | super(PosExtraction, self).__init__() 265 | operation = [] 266 | for _ in range(blocks): 267 | operation.append( 268 | ConvBNReLURes1D(channels, groups=groups, res_expansion=res_expansion, bias=bias, activation=activation) 269 | ) 270 | self.operation = nn.Sequential(*operation) 271 | 272 | def forward(self, x): # [b, d, g] 273 | return self.operation(x) 274 | 275 | 276 | class PointMLP(nn.Module): 277 | def __init__(self, cfg, activation="relu", bias=False, use_xyz=False, normalize="anchor", **kwargs): 278 | super(PointMLP, self).__init__() 279 | self.stages = len(cfg.NETWORK.ENCODER.PRE_BLOCKS) 280 | self.class_num = cfg.DATASET.NUM_CLASS 281 | self.points = cfg.DATASET.NUM_POINTS 282 | self.embedding = ConvBNReLU1D(3, cfg.NETWORK.ENCODER.EMBED_DIM, bias=bias, activation=activation) 283 | assert len(cfg.NETWORK.ENCODER.PRE_BLOCKS) == len(cfg.NETWORK.ENCODER.K_NEIGHBORS) == len(cfg.NETWORK.ENCODER.REDUCERS) == len(cfg.NETWORK.ENCODER.POS_BLOCKS) == len(cfg.NETWORK.ENCODER.DIM_EXP), \ 284 | "Please check stage number consistent for pre_blocks, pos_blocks k_neighbors, reducers." 285 | self.local_grouper_list = nn.ModuleList() 286 | self.pre_blocks_list = nn.ModuleList() 287 | self.pos_blocks_list = nn.ModuleList() 288 | last_channel = cfg.NETWORK.ENCODER.EMBED_DIM 289 | anchor_points = self.points 290 | for i in range(len(cfg.NETWORK.ENCODER.PRE_BLOCKS)): 291 | out_channel = last_channel * cfg.NETWORK.ENCODER.DIM_EXP[i] 292 | pre_block_num = cfg.NETWORK.ENCODER.PRE_BLOCKS[i] 293 | pos_block_num = cfg.NETWORK.ENCODER.POS_BLOCKS[i] 294 | kneighbor = cfg.NETWORK.ENCODER.K_NEIGHBORS[i] 295 | reduce = cfg.NETWORK.ENCODER.REDUCERS[i] 296 | anchor_points = anchor_points // reduce 297 | # append local_grouper_list 298 | local_grouper = LocalGrouper(last_channel, anchor_points, kneighbor, use_xyz, normalize) # [b,g,k,d] 299 | self.local_grouper_list.append(local_grouper) 300 | # append pre_block_list 301 | pre_block_module = PreExtraction(last_channel, out_channel, pre_block_num, groups=cfg.NETWORK.ENCODER.GROUPS, 302 | res_expansion=cfg.NETWORK.ENCODER.RES_EXP, 303 | bias=bias, activation=activation, use_xyz=use_xyz) 304 | self.pre_blocks_list.append(pre_block_module) 305 | # append pos_block_list 306 | pos_block_module = PosExtraction(out_channel, pos_block_num, groups=cfg.NETWORK.ENCODER.GROUPS, 307 | res_expansion=cfg.NETWORK.ENCODER.RES_EXP, bias=bias, activation=activation) 308 | self.pos_blocks_list.append(pos_block_module) 309 | 310 | last_channel = out_channel 311 | 312 | def forward(self, data_dic): 313 | x = data_dic['points'] 314 | x = x.permute(0, 2, 1) 315 | xyz = x.permute(0, 2, 1) 316 | batch_size, _, _ = x.size() 317 | x = self.embedding(x) # B,D,N 318 | for i in range(self.stages): 319 | # Give xyz[b, p, 3] and fea[b, p, d], return new_xyz[b, g, 3] and new_fea[b, g, k, d] 320 | xyz, x = self.local_grouper_list[i](xyz, x.permute(0, 2, 1)) # [b,g,3] [b,g,k,d] 321 | x = self.pre_blocks_list[i](x) # [b,d,g] 322 | x = self.pos_blocks_list[i](x) # [b,d,g] 323 | 324 | x = F.adaptive_max_pool1d(x, 1).squeeze(dim=-1) 325 | data_dic['point_features'] = x 326 | return data_dic 327 | 328 | 329 | # def pointMLP(num_classes=40, **kwargs) -> PointMLP: 330 | # return PointMLP(points=1024, class_num=num_classes, embed_dim=64, groups=1, res_expansion=1.0, 331 | # activation="relu", bias=False, use_xyz=False, normalize="anchor", 332 | # dim_expansion=[2, 2, 2, 2], pre_blocks=[2, 2, 2, 2], pos_blocks=[2, 2, 2, 2], 333 | # k_neighbors=[24, 24, 24, 24], reducers=[2, 2, 2, 2], **kwargs) 334 | 335 | 336 | # def pointMLPElite(num_classes=40, **kwargs) -> PointMLP: 337 | # return PointMLP(points=1024, class_num=num_classes, embed_dim=32, groups=1, res_expansion=0.25, 338 | # activation="relu", bias=False, use_xyz=False, normalize="anchor", 339 | # dim_expansion=[2, 2, 2, 1], pre_blocks=[1, 1, 2, 1], pos_blocks=[1, 1, 2, 1], 340 | # k_neighbors=[24,24,24,24], reducers=[2, 2, 2, 2], **kwargs) 341 | 342 | if __name__ == '__main__': 343 | data = torch.rand(2, 1024, 3).cuda() 344 | model = PointMLP(None).cuda() 345 | data_dic = { 346 | 'points': data 347 | } 348 | out = model(data_dic) 349 | print(out) 350 | -------------------------------------------------------------------------------- /core/networks/heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .linear_classifier import LinearClassifier 2 | from .linear_segmentator import LinearSegmentator 3 | 4 | __all__ = { 5 | 'LinearClassifier': LinearClassifier, 6 | 'LinearSegmentator': LinearSegmentator, 7 | } -------------------------------------------------------------------------------- /core/networks/heads/linear_classifier.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | class LinearClassifier(nn.Module): 5 | def __init__(self, cfg): 6 | super(LinearClassifier, self).__init__() 7 | self.cfg = cfg 8 | self.in_channels = self.cfg.NETWORK.HEAD.CLASSIFIER.IN_CHANNELS + self.cfg.NETWORK.HEAD.CLASSIFIER.DIMS[0:-1] 9 | self.out_channels = self.cfg.NETWORK.HEAD.CLASSIFIER.DIMS 10 | 11 | self.classifier = [] 12 | for c_in, c_out in zip(self.in_channels, self.out_channels): 13 | self.classifier.extend(nn.Sequential( 14 | nn.Linear(c_in, c_out), 15 | nn.BatchNorm1d(c_out), 16 | nn.ReLU(), 17 | nn.Dropout(0.5) 18 | )) 19 | 20 | self.classifier.extend(nn.Sequential( 21 | nn.Linear(self.cfg.NETWORK.HEAD.CLASSIFIER.DIMS[-1], self.cfg.DATASET.NUM_CLASS) 22 | )) 23 | 24 | self.classifier = nn.ModuleList(self.classifier) 25 | 26 | def forward(self, data_dic): 27 | logits = data_dic['point_features'] 28 | 29 | for cur_module in self.classifier: 30 | logits = cur_module(logits) 31 | 32 | data_dic['pred_score_logits'] = logits 33 | return data_dic -------------------------------------------------------------------------------- /core/networks/heads/linear_segmentator.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | def to_categorical(y, num_classes): 5 | """ 1-hot encodes a tensor """ 6 | new_y = torch.eye(num_classes)[y.cpu().data.numpy(),] 7 | if (y.is_cuda): 8 | return new_y.cuda(non_blocking=True) 9 | return new_y 10 | 11 | class LinearSegmentator(nn.Module): 12 | def __init__(self, cfg): 13 | super(LinearSegmentator, self).__init__() 14 | self.cfg = cfg 15 | self.in_channels = self.cfg.NETWORK.HEAD.SEGMENTATOR.IN_CHANNELS + self.cfg.NETWORK.HEAD.SEGMENTATOR.DIMS[0:-1] 16 | self.out_channels = self.cfg.NETWORK.HEAD.SEGMENTATOR.DIMS 17 | 18 | self.segmentator = [] 19 | for c_in, c_out in zip(self.in_channels, self.out_channels): 20 | self.segmentator.extend(nn.Sequential( 21 | nn.Conv1d(c_in, c_out, kernel_size=1), 22 | nn.BatchNorm1d(c_out), 23 | nn.ReLU(), 24 | nn.Dropout(0.5) 25 | )) 26 | 27 | self.segmentator.extend(nn.Sequential( 28 | nn.Conv1d(self.cfg.NETWORK.HEAD.SEGMENTATOR.DIMS[-1], self.cfg.DATASET.NUM_CLASS, kernel_size=1) 29 | )) 30 | 31 | self.segmentator = nn.ModuleList(self.segmentator) 32 | 33 | self.cls_map = nn.Sequential( 34 | nn.Conv1d(in_channels=16, out_channels=64, kernel_size=1, bias=True), 35 | nn.BatchNorm1d(64), 36 | nn.ReLU(), 37 | nn.Conv1d(in_channels=64, out_channels=64, kernel_size=1, bias=True), 38 | nn.BatchNorm1d(64), 39 | nn.ReLU(), 40 | ) 41 | 42 | def forward(self, data_dic): 43 | cls_tokens = data_dic['cls_tokens'] # (B, 1, 1) 44 | cls_tokens = to_categorical(cls_tokens, 16) # (B, 1, 16) 45 | cls_tokens = self.cls_map(cls_tokens.permute(0, 2, 1)) # B, 1, 64 46 | 47 | logits = data_dic['point_features'] 48 | logits = torch.cat([logits, cls_tokens.repeat(1, 1, logits.shape[-1])], dim = 1) 49 | 50 | for cur_module in self.segmentator: 51 | logits = cur_module(logits) 52 | data_dic['pred_score_logits'] = logits.permute(0,2,1) 53 | return data_dic 54 | -------------------------------------------------------------------------------- /core/networks/networks/__init__.py: -------------------------------------------------------------------------------- 1 | from .pointmlp import PointMLP 2 | from .pointstack import PointStack 3 | 4 | __all__ = { 5 | 'PointMLP': PointMLP, 6 | 'PointStack': PointStack 7 | } -------------------------------------------------------------------------------- /core/networks/networks/network_template.py: -------------------------------------------------------------------------------- 1 | ''' 2 | * author: Kevin Tirta Wijaya, AVELab, KAIST 3 | * e-mail: kevin.tirta@kaist.ac.kr 4 | ''' 5 | 6 | from .. import heads, encoders 7 | 8 | import torch 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | 12 | class NetworkTemplate(nn.Module): 13 | def __init__(self, cfg, topology = None): 14 | super().__init__() 15 | self.cfg = cfg 16 | self.num_class = cfg.DATASET.NUM_CLASS 17 | 18 | if topology == None: 19 | self.module_topology = ['encoder', 'head'] 20 | else: 21 | self.module_topology = topology 22 | 23 | self.module_list = self.build_networks() 24 | 25 | def build_networks(self): 26 | model_info_dict = { 27 | 'module_list': [], 28 | } 29 | for module_name in self.module_topology: 30 | modules, model_info_dict = getattr(self, 'build_%s' % module_name)( 31 | model_info_dict=model_info_dict 32 | ) 33 | 34 | 35 | return nn.ModuleList(model_info_dict['module_list']) 36 | 37 | def build_encoder(self, model_info_dict): 38 | if self.cfg.NETWORK.get('ENCODER', None) is None: 39 | return None, model_info_dict 40 | encoder_modules = [] 41 | encoder_module = encoders.__all__[self.cfg.NETWORK.ENCODER.NAME](cfg=self.cfg) 42 | model_info_dict['module_list'].append(encoder_module) 43 | encoder_modules.append(encoder_module) 44 | return nn.ModuleList(encoder_modules), model_info_dict 45 | 46 | 47 | def build_head(self, model_info_dict): 48 | if self.cfg.NETWORK.get('HEAD', None) is None: 49 | return None, model_info_dict 50 | 51 | head_modules = [] 52 | if self.cfg.NETWORK.HEAD.get('CLASSIFIER', None) is None: 53 | classifier_module = None 54 | else: 55 | classifier_module = heads.__all__[self.cfg.NETWORK.HEAD.CLASSIFIER.NAME](cfg=self.cfg) 56 | model_info_dict['module_list'].append(classifier_module) 57 | head_modules.append(classifier_module) 58 | 59 | if self.cfg.NETWORK.HEAD.get('SEGMENTATOR', None) is None: 60 | segmentator_module = None 61 | else: 62 | segmentator_module = heads.__all__[self.cfg.NETWORK.HEAD.SEGMENTATOR.NAME](cfg=self.cfg) 63 | model_info_dict['module_list'].append(segmentator_module) 64 | head_modules.append(segmentator_module) 65 | 66 | return nn.ModuleList(head_modules), model_info_dict 67 | 68 | def forward(self, data_dic): 69 | for curr_module in self.module_list: 70 | data_dic = curr_module(data_dic) 71 | 72 | return data_dic -------------------------------------------------------------------------------- /core/networks/networks/pointmlp.py: -------------------------------------------------------------------------------- 1 | from .. import heads, encoders 2 | from .network_template import NetworkTemplate 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | 8 | import numpy as np 9 | 10 | class PointMLP(NetworkTemplate): 11 | def __init__(self, cfg, topology = None): 12 | super().__init__(cfg, topology=topology) 13 | self.build_networks() 14 | self.cfg = cfg 15 | 16 | 17 | def get_loss(self, data_dic, smoothing=True, is_segmentation = False): 18 | ''' Calculate cross entropy loss, apply label smoothing if needed. ''' 19 | 20 | if is_segmentation: 21 | pred_logits = data_dic['pred_score_logits'].contiguous().view(-1, self.cfg.DATASET.NUM_CLASS) # (Batch size * Num Points, Num Classes) 22 | gt_cls_id = data_dic['seg_id'].contiguous().view(-1, 1).long() # (Batch size * Num Points, 1) 23 | 24 | else: 25 | pred_logits = data_dic['pred_score_logits'] # (Batch size, Num Classes) 26 | gt_cls_id = data_dic['cls_id'] # (Batch size, 1) 27 | 28 | if smoothing: 29 | eps = 0.2 30 | n_class = pred_logits.size(1) 31 | 32 | one_hot = torch.zeros_like(pred_logits).scatter(1, gt_cls_id.view(-1, 1), 1) 33 | one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1) 34 | log_prb = F.log_softmax(pred_logits, dim=1) 35 | loss = -(one_hot * log_prb).sum(dim=1) 36 | loss = loss[torch.isfinite(loss)].mean() 37 | else: 38 | loss = F.cross_entropy(pred_logits, gt_cls_id, reduction='mean') 39 | 40 | loss_dict = { 41 | 'Cls': loss.item(), 42 | } 43 | 44 | return loss, loss_dict 45 | 46 | def compute_overall_iou(self, pred, target, num_classes): 47 | shape_ious = [] 48 | pred = pred.max(dim=2)[1] # (batch_size, num_points) 49 | pred_np = pred.cpu().data.numpy() 50 | 51 | target_np = target.cpu().data.numpy() 52 | for shape_idx in range(pred.size(0)): # sample_idx 53 | part_ious = [] 54 | for part in range(num_classes): 55 | I = np.sum(np.logical_and(pred_np[shape_idx] == part, target_np[shape_idx] == part)) 56 | U = np.sum(np.logical_or(pred_np[shape_idx] == part, target_np[shape_idx] == part)) 57 | 58 | F = np.sum(target_np[shape_idx] == part) 59 | 60 | if F != 0: 61 | iou = I / float(U) 62 | part_ious.append(iou) 63 | shape_ious.append(np.mean(part_ious)) 64 | return shape_ious # [batch_size] -------------------------------------------------------------------------------- /core/networks/networks/pointstack.py: -------------------------------------------------------------------------------- 1 | ''' 2 | * author: Kevin Tirta Wijaya, AVELab, KAIST 3 | * e-mail: kevin.tirta@kaist.ac.kr 4 | ''' 5 | 6 | from .. import heads, encoders 7 | from .network_template import NetworkTemplate 8 | 9 | import torch 10 | import torch.nn as nn 11 | import torch.nn.functional as F 12 | 13 | import numpy as np 14 | 15 | class PointStack(NetworkTemplate): 16 | def __init__(self, cfg, topology = None): 17 | super().__init__(cfg, topology=topology) 18 | self.build_networks() 19 | self.cfg = cfg 20 | 21 | 22 | def get_loss(self, data_dic, smoothing=True, is_segmentation = False): 23 | ''' Calculate cross entropy loss, apply label smoothing if needed. ''' 24 | 25 | if is_segmentation: 26 | pred_logits = data_dic['pred_score_logits'].contiguous().view(-1, self.cfg.DATASET.NUM_CLASS) # (Batch size * Num Points, Num Classes) 27 | gt_cls_id = data_dic['seg_id'].contiguous().view(-1, 1).long() # (Batch size * Num Points, 1) 28 | 29 | else: 30 | pred_logits = data_dic['pred_score_logits'] # (Batch size, Num Classes) 31 | gt_cls_id = data_dic['cls_id'] # (Batch size, 1) 32 | 33 | if smoothing: 34 | eps = 0.2 35 | n_class = pred_logits.size(1) 36 | 37 | one_hot = torch.zeros_like(pred_logits).scatter(1, gt_cls_id.view(-1, 1), 1) 38 | one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1) 39 | log_prb = F.log_softmax(pred_logits, dim=1) 40 | loss = -(one_hot * log_prb).sum(dim=1) 41 | loss = loss[torch.isfinite(loss)].mean() 42 | else: 43 | loss = F.cross_entropy(pred_logits, gt_cls_id, reduction='mean') 44 | 45 | loss_dict = { 46 | 'Cls': loss.item(), 47 | } 48 | 49 | return loss, loss_dict 50 | 51 | def compute_overall_iou(self, pred, target, num_classes): 52 | shape_ious = [] 53 | pred = pred.max(dim=2)[1] # (batch_size, num_points) 54 | pred_np = pred.cpu().data.numpy() 55 | 56 | target_np = target.cpu().data.numpy() 57 | for shape_idx in range(pred.size(0)): # sample_idx 58 | part_ious = [] 59 | for part in range(num_classes): 60 | I = np.sum(np.logical_and(pred_np[shape_idx] == part, target_np[shape_idx] == part)) 61 | U = np.sum(np.logical_or(pred_np[shape_idx] == part, target_np[shape_idx] == part)) 62 | 63 | F = np.sum(target_np[shape_idx] == part) 64 | 65 | if F != 0: 66 | iou = I / float(U) 67 | part_ious.append(iou) 68 | shape_ious.append(np.mean(part_ious)) 69 | return shape_ious # [batch_size] -------------------------------------------------------------------------------- /docs/overview.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LongerVision/PointStack/9fe227bc654008452283a075d50f7406a66016be/docs/overview.png -------------------------------------------------------------------------------- /pointnet2_ops_lib/MANIFEST.in: -------------------------------------------------------------------------------- 1 | graft pointnet2_ops/_ext-src 2 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/__init__.py: -------------------------------------------------------------------------------- 1 | import pointnet2_ops.pointnet2_modules 2 | import pointnet2_ops.pointnet2_utils 3 | from pointnet2_ops._version import __version__ 4 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/include/ball_query.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 5 | const int nsample); 6 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/include/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef _CUDA_UTILS_H 2 | #define _CUDA_UTILS_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include 9 | #include 10 | 11 | #include 12 | 13 | #define TOTAL_THREADS 512 14 | 15 | inline int opt_n_threads(int work_size) { 16 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 17 | 18 | return max(min(1 << pow_2, TOTAL_THREADS), 1); 19 | } 20 | 21 | inline dim3 opt_block_config(int x, int y) { 22 | const int x_threads = opt_n_threads(x); 23 | const int y_threads = 24 | max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 25 | dim3 block_config(x_threads, y_threads, 1); 26 | 27 | return block_config; 28 | } 29 | 30 | #define CUDA_CHECK_ERRORS() \ 31 | do { \ 32 | cudaError_t err = cudaGetLastError(); \ 33 | if (cudaSuccess != err) { \ 34 | fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ 35 | cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ 36 | __FILE__); \ 37 | exit(-1); \ 38 | } \ 39 | } while (0) 40 | 41 | #endif 42 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/include/group_points.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor group_points(at::Tensor points, at::Tensor idx); 5 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 6 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/include/interpolate.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows); 7 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 8 | at::Tensor weight); 9 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 10 | at::Tensor weight, const int m); 11 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/include/sampling.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor gather_points(at::Tensor points, at::Tensor idx); 5 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 6 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples); 7 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/include/utils.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | 5 | #define CHECK_CUDA(x) \ 6 | do { \ 7 | AT_ASSERT(x.is_cuda(), #x " must be a CUDA tensor"); \ 8 | } while (0) 9 | 10 | #define CHECK_CONTIGUOUS(x) \ 11 | do { \ 12 | AT_ASSERT(x.is_contiguous(), #x " must be a contiguous tensor"); \ 13 | } while (0) 14 | 15 | #define CHECK_IS_INT(x) \ 16 | do { \ 17 | AT_ASSERT(x.scalar_type() == at::ScalarType::Int, \ 18 | #x " must be an int tensor"); \ 19 | } while (0) 20 | 21 | #define CHECK_IS_FLOAT(x) \ 22 | do { \ 23 | AT_ASSERT(x.scalar_type() == at::ScalarType::Float, \ 24 | #x " must be a float tensor"); \ 25 | } while (0) 26 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query.h" 2 | #include "utils.h" 3 | 4 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 5 | int nsample, const float *new_xyz, 6 | const float *xyz, int *idx); 7 | 8 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 9 | const int nsample) { 10 | CHECK_CONTIGUOUS(new_xyz); 11 | CHECK_CONTIGUOUS(xyz); 12 | CHECK_IS_FLOAT(new_xyz); 13 | CHECK_IS_FLOAT(xyz); 14 | 15 | if (new_xyz.is_cuda()) { 16 | CHECK_CUDA(xyz); 17 | } 18 | 19 | at::Tensor idx = 20 | torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample}, 21 | at::device(new_xyz.device()).dtype(at::ScalarType::Int)); 22 | 23 | if (new_xyz.is_cuda()) { 24 | query_ball_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), 25 | radius, nsample, new_xyz.data_ptr(), 26 | xyz.data_ptr(), idx.data_ptr()); 27 | } else { 28 | AT_ASSERT(false, "CPU not supported"); 29 | } 30 | 31 | return idx; 32 | } 33 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "cuda_utils.h" 6 | 7 | // input: new_xyz(b, m, 3) xyz(b, n, 3) 8 | // output: idx(b, m, nsample) 9 | __global__ void query_ball_point_kernel(int b, int n, int m, float radius, 10 | int nsample, 11 | const float *__restrict__ new_xyz, 12 | const float *__restrict__ xyz, 13 | int *__restrict__ idx) { 14 | int batch_index = blockIdx.x; 15 | xyz += batch_index * n * 3; 16 | new_xyz += batch_index * m * 3; 17 | idx += m * nsample * batch_index; 18 | 19 | int index = threadIdx.x; 20 | int stride = blockDim.x; 21 | 22 | float radius2 = radius * radius; 23 | for (int j = index; j < m; j += stride) { 24 | float new_x = new_xyz[j * 3 + 0]; 25 | float new_y = new_xyz[j * 3 + 1]; 26 | float new_z = new_xyz[j * 3 + 2]; 27 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { 28 | float x = xyz[k * 3 + 0]; 29 | float y = xyz[k * 3 + 1]; 30 | float z = xyz[k * 3 + 2]; 31 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 32 | (new_z - z) * (new_z - z); 33 | if (d2 < radius2) { 34 | if (cnt == 0) { 35 | for (int l = 0; l < nsample; ++l) { 36 | idx[j * nsample + l] = k; 37 | } 38 | } 39 | idx[j * nsample + cnt] = k; 40 | ++cnt; 41 | } 42 | } 43 | } 44 | } 45 | 46 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 47 | int nsample, const float *new_xyz, 48 | const float *xyz, int *idx) { 49 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 50 | query_ball_point_kernel<<>>( 51 | b, n, m, radius, nsample, new_xyz, xyz, idx); 52 | 53 | CUDA_CHECK_ERRORS(); 54 | } 55 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/bindings.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query.h" 2 | #include "group_points.h" 3 | #include "interpolate.h" 4 | #include "sampling.h" 5 | 6 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 7 | m.def("gather_points", &gather_points); 8 | m.def("gather_points_grad", &gather_points_grad); 9 | m.def("furthest_point_sampling", &furthest_point_sampling); 10 | 11 | m.def("three_nn", &three_nn); 12 | m.def("three_interpolate", &three_interpolate); 13 | m.def("three_interpolate_grad", &three_interpolate_grad); 14 | 15 | m.def("ball_query", &ball_query); 16 | 17 | m.def("group_points", &group_points); 18 | m.def("group_points_grad", &group_points_grad); 19 | } 20 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points.cpp: -------------------------------------------------------------------------------- 1 | #include "group_points.h" 2 | #include "utils.h" 3 | 4 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 5 | const float *points, const int *idx, 6 | float *out); 7 | 8 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 9 | int nsample, const float *grad_out, 10 | const int *idx, float *grad_points); 11 | 12 | at::Tensor group_points(at::Tensor points, at::Tensor idx) { 13 | CHECK_CONTIGUOUS(points); 14 | CHECK_CONTIGUOUS(idx); 15 | CHECK_IS_FLOAT(points); 16 | CHECK_IS_INT(idx); 17 | 18 | if (points.is_cuda()) { 19 | CHECK_CUDA(idx); 20 | } 21 | 22 | at::Tensor output = 23 | torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)}, 24 | at::device(points.device()).dtype(at::ScalarType::Float)); 25 | 26 | if (points.is_cuda()) { 27 | group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 28 | idx.size(1), idx.size(2), 29 | points.data_ptr(), idx.data_ptr(), 30 | output.data_ptr()); 31 | } else { 32 | AT_ASSERT(false, "CPU not supported"); 33 | } 34 | 35 | return output; 36 | } 37 | 38 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) { 39 | CHECK_CONTIGUOUS(grad_out); 40 | CHECK_CONTIGUOUS(idx); 41 | CHECK_IS_FLOAT(grad_out); 42 | CHECK_IS_INT(idx); 43 | 44 | if (grad_out.is_cuda()) { 45 | CHECK_CUDA(idx); 46 | } 47 | 48 | at::Tensor output = 49 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 50 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 51 | 52 | if (grad_out.is_cuda()) { 53 | group_points_grad_kernel_wrapper( 54 | grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2), 55 | grad_out.data_ptr(), idx.data_ptr(), 56 | output.data_ptr()); 57 | } else { 58 | AT_ASSERT(false, "CPU not supported"); 59 | } 60 | 61 | return output; 62 | } 63 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "cuda_utils.h" 5 | 6 | // input: points(b, c, n) idx(b, npoints, nsample) 7 | // output: out(b, c, npoints, nsample) 8 | __global__ void group_points_kernel(int b, int c, int n, int npoints, 9 | int nsample, 10 | const float *__restrict__ points, 11 | const int *__restrict__ idx, 12 | float *__restrict__ out) { 13 | int batch_index = blockIdx.x; 14 | points += batch_index * n * c; 15 | idx += batch_index * npoints * nsample; 16 | out += batch_index * npoints * nsample * c; 17 | 18 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 19 | const int stride = blockDim.y * blockDim.x; 20 | for (int i = index; i < c * npoints; i += stride) { 21 | const int l = i / npoints; 22 | const int j = i % npoints; 23 | for (int k = 0; k < nsample; ++k) { 24 | int ii = idx[j * nsample + k]; 25 | out[(l * npoints + j) * nsample + k] = points[l * n + ii]; 26 | } 27 | } 28 | } 29 | 30 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 31 | const float *points, const int *idx, 32 | float *out) { 33 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 34 | 35 | group_points_kernel<<>>( 36 | b, c, n, npoints, nsample, points, idx, out); 37 | 38 | CUDA_CHECK_ERRORS(); 39 | } 40 | 41 | // input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample) 42 | // output: grad_points(b, c, n) 43 | __global__ void group_points_grad_kernel(int b, int c, int n, int npoints, 44 | int nsample, 45 | const float *__restrict__ grad_out, 46 | const int *__restrict__ idx, 47 | float *__restrict__ grad_points) { 48 | int batch_index = blockIdx.x; 49 | grad_out += batch_index * npoints * nsample * c; 50 | idx += batch_index * npoints * nsample; 51 | grad_points += batch_index * n * c; 52 | 53 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 54 | const int stride = blockDim.y * blockDim.x; 55 | for (int i = index; i < c * npoints; i += stride) { 56 | const int l = i / npoints; 57 | const int j = i % npoints; 58 | for (int k = 0; k < nsample; ++k) { 59 | int ii = idx[j * nsample + k]; 60 | atomicAdd(grad_points + l * n + ii, 61 | grad_out[(l * npoints + j) * nsample + k]); 62 | } 63 | } 64 | } 65 | 66 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 67 | int nsample, const float *grad_out, 68 | const int *idx, float *grad_points) { 69 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 70 | 71 | group_points_grad_kernel<<>>( 72 | b, c, n, npoints, nsample, grad_out, idx, grad_points); 73 | 74 | CUDA_CHECK_ERRORS(); 75 | } 76 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate.cpp: -------------------------------------------------------------------------------- 1 | #include "interpolate.h" 2 | #include "utils.h" 3 | 4 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 5 | const float *known, float *dist2, int *idx); 6 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 7 | const float *points, const int *idx, 8 | const float *weight, float *out); 9 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 10 | const float *grad_out, 11 | const int *idx, const float *weight, 12 | float *grad_points); 13 | 14 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows) { 15 | CHECK_CONTIGUOUS(unknowns); 16 | CHECK_CONTIGUOUS(knows); 17 | CHECK_IS_FLOAT(unknowns); 18 | CHECK_IS_FLOAT(knows); 19 | 20 | if (unknowns.is_cuda()) { 21 | CHECK_CUDA(knows); 22 | } 23 | 24 | at::Tensor idx = 25 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 26 | at::device(unknowns.device()).dtype(at::ScalarType::Int)); 27 | at::Tensor dist2 = 28 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 29 | at::device(unknowns.device()).dtype(at::ScalarType::Float)); 30 | 31 | if (unknowns.is_cuda()) { 32 | three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1), 33 | unknowns.data_ptr(), knows.data_ptr(), 34 | dist2.data_ptr(), idx.data_ptr()); 35 | } else { 36 | AT_ASSERT(false, "CPU not supported"); 37 | } 38 | 39 | return {dist2, idx}; 40 | } 41 | 42 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 43 | at::Tensor weight) { 44 | CHECK_CONTIGUOUS(points); 45 | CHECK_CONTIGUOUS(idx); 46 | CHECK_CONTIGUOUS(weight); 47 | CHECK_IS_FLOAT(points); 48 | CHECK_IS_INT(idx); 49 | CHECK_IS_FLOAT(weight); 50 | 51 | if (points.is_cuda()) { 52 | CHECK_CUDA(idx); 53 | CHECK_CUDA(weight); 54 | } 55 | 56 | at::Tensor output = 57 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 58 | at::device(points.device()).dtype(at::ScalarType::Float)); 59 | 60 | if (points.is_cuda()) { 61 | three_interpolate_kernel_wrapper( 62 | points.size(0), points.size(1), points.size(2), idx.size(1), 63 | points.data_ptr(), idx.data_ptr(), weight.data_ptr(), 64 | output.data_ptr()); 65 | } else { 66 | AT_ASSERT(false, "CPU not supported"); 67 | } 68 | 69 | return output; 70 | } 71 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 72 | at::Tensor weight, const int m) { 73 | CHECK_CONTIGUOUS(grad_out); 74 | CHECK_CONTIGUOUS(idx); 75 | CHECK_CONTIGUOUS(weight); 76 | CHECK_IS_FLOAT(grad_out); 77 | CHECK_IS_INT(idx); 78 | CHECK_IS_FLOAT(weight); 79 | 80 | if (grad_out.is_cuda()) { 81 | CHECK_CUDA(idx); 82 | CHECK_CUDA(weight); 83 | } 84 | 85 | at::Tensor output = 86 | torch::zeros({grad_out.size(0), grad_out.size(1), m}, 87 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 88 | 89 | if (grad_out.is_cuda()) { 90 | three_interpolate_grad_kernel_wrapper( 91 | grad_out.size(0), grad_out.size(1), grad_out.size(2), m, 92 | grad_out.data_ptr(), idx.data_ptr(), 93 | weight.data_ptr(), output.data_ptr()); 94 | } else { 95 | AT_ASSERT(false, "CPU not supported"); 96 | } 97 | 98 | return output; 99 | } 100 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "cuda_utils.h" 6 | 7 | // input: unknown(b, n, 3) known(b, m, 3) 8 | // output: dist2(b, n, 3), idx(b, n, 3) 9 | __global__ void three_nn_kernel(int b, int n, int m, 10 | const float *__restrict__ unknown, 11 | const float *__restrict__ known, 12 | float *__restrict__ dist2, 13 | int *__restrict__ idx) { 14 | int batch_index = blockIdx.x; 15 | unknown += batch_index * n * 3; 16 | known += batch_index * m * 3; 17 | dist2 += batch_index * n * 3; 18 | idx += batch_index * n * 3; 19 | 20 | int index = threadIdx.x; 21 | int stride = blockDim.x; 22 | for (int j = index; j < n; j += stride) { 23 | float ux = unknown[j * 3 + 0]; 24 | float uy = unknown[j * 3 + 1]; 25 | float uz = unknown[j * 3 + 2]; 26 | 27 | double best1 = 1e40, best2 = 1e40, best3 = 1e40; 28 | int besti1 = 0, besti2 = 0, besti3 = 0; 29 | for (int k = 0; k < m; ++k) { 30 | float x = known[k * 3 + 0]; 31 | float y = known[k * 3 + 1]; 32 | float z = known[k * 3 + 2]; 33 | float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z); 34 | if (d < best1) { 35 | best3 = best2; 36 | besti3 = besti2; 37 | best2 = best1; 38 | besti2 = besti1; 39 | best1 = d; 40 | besti1 = k; 41 | } else if (d < best2) { 42 | best3 = best2; 43 | besti3 = besti2; 44 | best2 = d; 45 | besti2 = k; 46 | } else if (d < best3) { 47 | best3 = d; 48 | besti3 = k; 49 | } 50 | } 51 | dist2[j * 3 + 0] = best1; 52 | dist2[j * 3 + 1] = best2; 53 | dist2[j * 3 + 2] = best3; 54 | 55 | idx[j * 3 + 0] = besti1; 56 | idx[j * 3 + 1] = besti2; 57 | idx[j * 3 + 2] = besti3; 58 | } 59 | } 60 | 61 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 62 | const float *known, float *dist2, int *idx) { 63 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 64 | three_nn_kernel<<>>(b, n, m, unknown, known, 65 | dist2, idx); 66 | 67 | CUDA_CHECK_ERRORS(); 68 | } 69 | 70 | // input: points(b, c, m), idx(b, n, 3), weight(b, n, 3) 71 | // output: out(b, c, n) 72 | __global__ void three_interpolate_kernel(int b, int c, int m, int n, 73 | const float *__restrict__ points, 74 | const int *__restrict__ idx, 75 | const float *__restrict__ weight, 76 | float *__restrict__ out) { 77 | int batch_index = blockIdx.x; 78 | points += batch_index * m * c; 79 | 80 | idx += batch_index * n * 3; 81 | weight += batch_index * n * 3; 82 | 83 | out += batch_index * n * c; 84 | 85 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 86 | const int stride = blockDim.y * blockDim.x; 87 | for (int i = index; i < c * n; i += stride) { 88 | const int l = i / n; 89 | const int j = i % n; 90 | float w1 = weight[j * 3 + 0]; 91 | float w2 = weight[j * 3 + 1]; 92 | float w3 = weight[j * 3 + 2]; 93 | 94 | int i1 = idx[j * 3 + 0]; 95 | int i2 = idx[j * 3 + 1]; 96 | int i3 = idx[j * 3 + 2]; 97 | 98 | out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 + 99 | points[l * m + i3] * w3; 100 | } 101 | } 102 | 103 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 104 | const float *points, const int *idx, 105 | const float *weight, float *out) { 106 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 107 | three_interpolate_kernel<<>>( 108 | b, c, m, n, points, idx, weight, out); 109 | 110 | CUDA_CHECK_ERRORS(); 111 | } 112 | 113 | // input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3) 114 | // output: grad_points(b, c, m) 115 | 116 | __global__ void three_interpolate_grad_kernel( 117 | int b, int c, int n, int m, const float *__restrict__ grad_out, 118 | const int *__restrict__ idx, const float *__restrict__ weight, 119 | float *__restrict__ grad_points) { 120 | int batch_index = blockIdx.x; 121 | grad_out += batch_index * n * c; 122 | idx += batch_index * n * 3; 123 | weight += batch_index * n * 3; 124 | grad_points += batch_index * m * c; 125 | 126 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 127 | const int stride = blockDim.y * blockDim.x; 128 | for (int i = index; i < c * n; i += stride) { 129 | const int l = i / n; 130 | const int j = i % n; 131 | float w1 = weight[j * 3 + 0]; 132 | float w2 = weight[j * 3 + 1]; 133 | float w3 = weight[j * 3 + 2]; 134 | 135 | int i1 = idx[j * 3 + 0]; 136 | int i2 = idx[j * 3 + 1]; 137 | int i3 = idx[j * 3 + 2]; 138 | 139 | atomicAdd(grad_points + l * m + i1, grad_out[i] * w1); 140 | atomicAdd(grad_points + l * m + i2, grad_out[i] * w2); 141 | atomicAdd(grad_points + l * m + i3, grad_out[i] * w3); 142 | } 143 | } 144 | 145 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 146 | const float *grad_out, 147 | const int *idx, const float *weight, 148 | float *grad_points) { 149 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 150 | three_interpolate_grad_kernel<<>>( 151 | b, c, n, m, grad_out, idx, weight, grad_points); 152 | 153 | CUDA_CHECK_ERRORS(); 154 | } 155 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling.cpp: -------------------------------------------------------------------------------- 1 | #include "sampling.h" 2 | #include "utils.h" 3 | 4 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 5 | const float *points, const int *idx, 6 | float *out); 7 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 8 | const float *grad_out, const int *idx, 9 | float *grad_points); 10 | 11 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 12 | const float *dataset, float *temp, 13 | int *idxs); 14 | 15 | at::Tensor gather_points(at::Tensor points, at::Tensor idx) { 16 | CHECK_CONTIGUOUS(points); 17 | CHECK_CONTIGUOUS(idx); 18 | CHECK_IS_FLOAT(points); 19 | CHECK_IS_INT(idx); 20 | 21 | if (points.is_cuda()) { 22 | CHECK_CUDA(idx); 23 | } 24 | 25 | at::Tensor output = 26 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 27 | at::device(points.device()).dtype(at::ScalarType::Float)); 28 | 29 | if (points.is_cuda()) { 30 | gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 31 | idx.size(1), points.data_ptr(), 32 | idx.data_ptr(), output.data_ptr()); 33 | } else { 34 | AT_ASSERT(false, "CPU not supported"); 35 | } 36 | 37 | return output; 38 | } 39 | 40 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, 41 | const int n) { 42 | CHECK_CONTIGUOUS(grad_out); 43 | CHECK_CONTIGUOUS(idx); 44 | CHECK_IS_FLOAT(grad_out); 45 | CHECK_IS_INT(idx); 46 | 47 | if (grad_out.is_cuda()) { 48 | CHECK_CUDA(idx); 49 | } 50 | 51 | at::Tensor output = 52 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 53 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 54 | 55 | if (grad_out.is_cuda()) { 56 | gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n, 57 | idx.size(1), grad_out.data_ptr(), 58 | idx.data_ptr(), 59 | output.data_ptr()); 60 | } else { 61 | AT_ASSERT(false, "CPU not supported"); 62 | } 63 | 64 | return output; 65 | } 66 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) { 67 | CHECK_CONTIGUOUS(points); 68 | CHECK_IS_FLOAT(points); 69 | 70 | at::Tensor output = 71 | torch::zeros({points.size(0), nsamples}, 72 | at::device(points.device()).dtype(at::ScalarType::Int)); 73 | 74 | at::Tensor tmp = 75 | torch::full({points.size(0), points.size(1)}, 1e10, 76 | at::device(points.device()).dtype(at::ScalarType::Float)); 77 | 78 | if (points.is_cuda()) { 79 | furthest_point_sampling_kernel_wrapper( 80 | points.size(0), points.size(1), nsamples, points.data_ptr(), 81 | tmp.data_ptr(), output.data_ptr()); 82 | } else { 83 | AT_ASSERT(false, "CPU not supported"); 84 | } 85 | 86 | return output; 87 | } 88 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "cuda_utils.h" 5 | 6 | // input: points(b, c, n) idx(b, m) 7 | // output: out(b, c, m) 8 | __global__ void gather_points_kernel(int b, int c, int n, int m, 9 | const float *__restrict__ points, 10 | const int *__restrict__ idx, 11 | float *__restrict__ out) { 12 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 13 | for (int l = blockIdx.y; l < c; l += gridDim.y) { 14 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 15 | int a = idx[i * m + j]; 16 | out[(i * c + l) * m + j] = points[(i * c + l) * n + a]; 17 | } 18 | } 19 | } 20 | } 21 | 22 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 23 | const float *points, const int *idx, 24 | float *out) { 25 | gather_points_kernel<<>>(b, c, n, npoints, 27 | points, idx, out); 28 | 29 | CUDA_CHECK_ERRORS(); 30 | } 31 | 32 | // input: grad_out(b, c, m) idx(b, m) 33 | // output: grad_points(b, c, n) 34 | __global__ void gather_points_grad_kernel(int b, int c, int n, int m, 35 | const float *__restrict__ grad_out, 36 | const int *__restrict__ idx, 37 | float *__restrict__ grad_points) { 38 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 39 | for (int l = blockIdx.y; l < c; l += gridDim.y) { 40 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 41 | int a = idx[i * m + j]; 42 | atomicAdd(grad_points + (i * c + l) * n + a, 43 | grad_out[(i * c + l) * m + j]); 44 | } 45 | } 46 | } 47 | } 48 | 49 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 50 | const float *grad_out, const int *idx, 51 | float *grad_points) { 52 | gather_points_grad_kernel<<>>( 54 | b, c, n, npoints, grad_out, idx, grad_points); 55 | 56 | CUDA_CHECK_ERRORS(); 57 | } 58 | 59 | __device__ void __update(float *__restrict__ dists, int *__restrict__ dists_i, 60 | int idx1, int idx2) { 61 | const float v1 = dists[idx1], v2 = dists[idx2]; 62 | const int i1 = dists_i[idx1], i2 = dists_i[idx2]; 63 | dists[idx1] = max(v1, v2); 64 | dists_i[idx1] = v2 > v1 ? i2 : i1; 65 | } 66 | 67 | // Input dataset: (b, n, 3), tmp: (b, n) 68 | // Ouput idxs (b, m) 69 | template 70 | __global__ void furthest_point_sampling_kernel( 71 | int b, int n, int m, const float *__restrict__ dataset, 72 | float *__restrict__ temp, int *__restrict__ idxs) { 73 | if (m <= 0) return; 74 | __shared__ float dists[block_size]; 75 | __shared__ int dists_i[block_size]; 76 | 77 | int batch_index = blockIdx.x; 78 | dataset += batch_index * n * 3; 79 | temp += batch_index * n; 80 | idxs += batch_index * m; 81 | 82 | int tid = threadIdx.x; 83 | const int stride = block_size; 84 | 85 | int old = 0; 86 | if (threadIdx.x == 0) idxs[0] = old; 87 | 88 | __syncthreads(); 89 | for (int j = 1; j < m; j++) { 90 | int besti = 0; 91 | float best = -1; 92 | float x1 = dataset[old * 3 + 0]; 93 | float y1 = dataset[old * 3 + 1]; 94 | float z1 = dataset[old * 3 + 2]; 95 | for (int k = tid; k < n; k += stride) { 96 | float x2, y2, z2; 97 | x2 = dataset[k * 3 + 0]; 98 | y2 = dataset[k * 3 + 1]; 99 | z2 = dataset[k * 3 + 2]; 100 | float mag = (x2 * x2) + (y2 * y2) + (z2 * z2); 101 | if (mag <= 1e-3) continue; 102 | 103 | float d = 104 | (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1); 105 | 106 | float d2 = min(d, temp[k]); 107 | temp[k] = d2; 108 | besti = d2 > best ? k : besti; 109 | best = d2 > best ? d2 : best; 110 | } 111 | dists[tid] = best; 112 | dists_i[tid] = besti; 113 | __syncthreads(); 114 | 115 | if (block_size >= 512) { 116 | if (tid < 256) { 117 | __update(dists, dists_i, tid, tid + 256); 118 | } 119 | __syncthreads(); 120 | } 121 | if (block_size >= 256) { 122 | if (tid < 128) { 123 | __update(dists, dists_i, tid, tid + 128); 124 | } 125 | __syncthreads(); 126 | } 127 | if (block_size >= 128) { 128 | if (tid < 64) { 129 | __update(dists, dists_i, tid, tid + 64); 130 | } 131 | __syncthreads(); 132 | } 133 | if (block_size >= 64) { 134 | if (tid < 32) { 135 | __update(dists, dists_i, tid, tid + 32); 136 | } 137 | __syncthreads(); 138 | } 139 | if (block_size >= 32) { 140 | if (tid < 16) { 141 | __update(dists, dists_i, tid, tid + 16); 142 | } 143 | __syncthreads(); 144 | } 145 | if (block_size >= 16) { 146 | if (tid < 8) { 147 | __update(dists, dists_i, tid, tid + 8); 148 | } 149 | __syncthreads(); 150 | } 151 | if (block_size >= 8) { 152 | if (tid < 4) { 153 | __update(dists, dists_i, tid, tid + 4); 154 | } 155 | __syncthreads(); 156 | } 157 | if (block_size >= 4) { 158 | if (tid < 2) { 159 | __update(dists, dists_i, tid, tid + 2); 160 | } 161 | __syncthreads(); 162 | } 163 | if (block_size >= 2) { 164 | if (tid < 1) { 165 | __update(dists, dists_i, tid, tid + 1); 166 | } 167 | __syncthreads(); 168 | } 169 | 170 | old = dists_i[0]; 171 | if (tid == 0) idxs[j] = old; 172 | } 173 | } 174 | 175 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 176 | const float *dataset, float *temp, 177 | int *idxs) { 178 | unsigned int n_threads = opt_n_threads(n); 179 | 180 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 181 | 182 | switch (n_threads) { 183 | case 512: 184 | furthest_point_sampling_kernel<512> 185 | <<>>(b, n, m, dataset, temp, idxs); 186 | break; 187 | case 256: 188 | furthest_point_sampling_kernel<256> 189 | <<>>(b, n, m, dataset, temp, idxs); 190 | break; 191 | case 128: 192 | furthest_point_sampling_kernel<128> 193 | <<>>(b, n, m, dataset, temp, idxs); 194 | break; 195 | case 64: 196 | furthest_point_sampling_kernel<64> 197 | <<>>(b, n, m, dataset, temp, idxs); 198 | break; 199 | case 32: 200 | furthest_point_sampling_kernel<32> 201 | <<>>(b, n, m, dataset, temp, idxs); 202 | break; 203 | case 16: 204 | furthest_point_sampling_kernel<16> 205 | <<>>(b, n, m, dataset, temp, idxs); 206 | break; 207 | case 8: 208 | furthest_point_sampling_kernel<8> 209 | <<>>(b, n, m, dataset, temp, idxs); 210 | break; 211 | case 4: 212 | furthest_point_sampling_kernel<4> 213 | <<>>(b, n, m, dataset, temp, idxs); 214 | break; 215 | case 2: 216 | furthest_point_sampling_kernel<2> 217 | <<>>(b, n, m, dataset, temp, idxs); 218 | break; 219 | case 1: 220 | furthest_point_sampling_kernel<1> 221 | <<>>(b, n, m, dataset, temp, idxs); 222 | break; 223 | default: 224 | furthest_point_sampling_kernel<512> 225 | <<>>(b, n, m, dataset, temp, idxs); 226 | } 227 | 228 | CUDA_CHECK_ERRORS(); 229 | } 230 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LongerVision/PointStack/9fe227bc654008452283a075d50f7406a66016be/pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_version.py: -------------------------------------------------------------------------------- 1 | __version__ = "3.0.0" 2 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/pointnet2_modules.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Tuple 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | from pointnet2_ops import pointnet2_utils 7 | 8 | 9 | def build_shared_mlp(mlp_spec: List[int], bn: bool = True): 10 | layers = [] 11 | for i in range(1, len(mlp_spec)): 12 | layers.append( 13 | nn.Conv2d(mlp_spec[i - 1], mlp_spec[i], kernel_size=1, bias=not bn) 14 | ) 15 | if bn: 16 | layers.append(nn.BatchNorm2d(mlp_spec[i])) 17 | layers.append(nn.ReLU(True)) 18 | 19 | return nn.Sequential(*layers) 20 | 21 | 22 | class _PointnetSAModuleBase(nn.Module): 23 | def __init__(self): 24 | super(_PointnetSAModuleBase, self).__init__() 25 | self.npoint = None 26 | self.groupers = None 27 | self.mlps = None 28 | 29 | def forward( 30 | self, xyz: torch.Tensor, features: Optional[torch.Tensor] 31 | ) -> Tuple[torch.Tensor, torch.Tensor]: 32 | r""" 33 | Parameters 34 | ---------- 35 | xyz : torch.Tensor 36 | (B, N, 3) tensor of the xyz coordinates of the features 37 | features : torch.Tensor 38 | (B, C, N) tensor of the descriptors of the the features 39 | 40 | Returns 41 | ------- 42 | new_xyz : torch.Tensor 43 | (B, npoint, 3) tensor of the new features' xyz 44 | new_features : torch.Tensor 45 | (B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors 46 | """ 47 | 48 | new_features_list = [] 49 | 50 | xyz_flipped = xyz.transpose(1, 2).contiguous() 51 | new_xyz = ( 52 | pointnet2_utils.gather_operation( 53 | xyz_flipped, pointnet2_utils.furthest_point_sample(xyz, self.npoint) 54 | ) 55 | .transpose(1, 2) 56 | .contiguous() 57 | if self.npoint is not None 58 | else None 59 | ) 60 | 61 | for i in range(len(self.groupers)): 62 | new_features = self.groupers[i]( 63 | xyz, new_xyz, features 64 | ) # (B, C, npoint, nsample) 65 | 66 | new_features = self.mlps[i](new_features) # (B, mlp[-1], npoint, nsample) 67 | new_features = F.max_pool2d( 68 | new_features, kernel_size=[1, new_features.size(3)] 69 | ) # (B, mlp[-1], npoint, 1) 70 | new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) 71 | 72 | new_features_list.append(new_features) 73 | 74 | return new_xyz, torch.cat(new_features_list, dim=1) 75 | 76 | 77 | class PointnetSAModuleMSG(_PointnetSAModuleBase): 78 | r"""Pointnet set abstrction layer with multiscale grouping 79 | 80 | Parameters 81 | ---------- 82 | npoint : int 83 | Number of features 84 | radii : list of float32 85 | list of radii to group with 86 | nsamples : list of int32 87 | Number of samples in each ball query 88 | mlps : list of list of int32 89 | Spec of the pointnet before the global max_pool for each scale 90 | bn : bool 91 | Use batchnorm 92 | """ 93 | 94 | def __init__(self, npoint, radii, nsamples, mlps, bn=True, use_xyz=True): 95 | # type: (PointnetSAModuleMSG, int, List[float], List[int], List[List[int]], bool, bool) -> None 96 | super(PointnetSAModuleMSG, self).__init__() 97 | 98 | assert len(radii) == len(nsamples) == len(mlps) 99 | 100 | self.npoint = npoint 101 | self.groupers = nn.ModuleList() 102 | self.mlps = nn.ModuleList() 103 | for i in range(len(radii)): 104 | radius = radii[i] 105 | nsample = nsamples[i] 106 | self.groupers.append( 107 | pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz) 108 | if npoint is not None 109 | else pointnet2_utils.GroupAll(use_xyz) 110 | ) 111 | mlp_spec = mlps[i] 112 | if use_xyz: 113 | mlp_spec[0] += 3 114 | 115 | self.mlps.append(build_shared_mlp(mlp_spec, bn)) 116 | 117 | 118 | class PointnetSAModule(PointnetSAModuleMSG): 119 | r"""Pointnet set abstrction layer 120 | 121 | Parameters 122 | ---------- 123 | npoint : int 124 | Number of features 125 | radius : float 126 | Radius of ball 127 | nsample : int 128 | Number of samples in the ball query 129 | mlp : list 130 | Spec of the pointnet before the global max_pool 131 | bn : bool 132 | Use batchnorm 133 | """ 134 | 135 | def __init__( 136 | self, mlp, npoint=None, radius=None, nsample=None, bn=True, use_xyz=True 137 | ): 138 | # type: (PointnetSAModule, List[int], int, float, int, bool, bool) -> None 139 | super(PointnetSAModule, self).__init__( 140 | mlps=[mlp], 141 | npoint=npoint, 142 | radii=[radius], 143 | nsamples=[nsample], 144 | bn=bn, 145 | use_xyz=use_xyz, 146 | ) 147 | 148 | 149 | class PointnetFPModule(nn.Module): 150 | r"""Propigates the features of one set to another 151 | 152 | Parameters 153 | ---------- 154 | mlp : list 155 | Pointnet module parameters 156 | bn : bool 157 | Use batchnorm 158 | """ 159 | 160 | def __init__(self, mlp, bn=True): 161 | # type: (PointnetFPModule, List[int], bool) -> None 162 | super(PointnetFPModule, self).__init__() 163 | self.mlp = build_shared_mlp(mlp, bn=bn) 164 | 165 | def forward(self, unknown, known, unknow_feats, known_feats): 166 | # type: (PointnetFPModule, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor 167 | r""" 168 | Parameters 169 | ---------- 170 | unknown : torch.Tensor 171 | (B, n, 3) tensor of the xyz positions of the unknown features 172 | known : torch.Tensor 173 | (B, m, 3) tensor of the xyz positions of the known features 174 | unknow_feats : torch.Tensor 175 | (B, C1, n) tensor of the features to be propigated to 176 | known_feats : torch.Tensor 177 | (B, C2, m) tensor of features to be propigated 178 | 179 | Returns 180 | ------- 181 | new_features : torch.Tensor 182 | (B, mlp[-1], n) tensor of the features of the unknown features 183 | """ 184 | 185 | if known is not None: 186 | dist, idx = pointnet2_utils.three_nn(unknown, known) 187 | dist_recip = 1.0 / (dist + 1e-8) 188 | norm = torch.sum(dist_recip, dim=2, keepdim=True) 189 | weight = dist_recip / norm 190 | 191 | interpolated_feats = pointnet2_utils.three_interpolate( 192 | known_feats, idx, weight 193 | ) 194 | else: 195 | interpolated_feats = known_feats.expand( 196 | *(known_feats.size()[0:2] + [unknown.size(1)]) 197 | ) 198 | 199 | if unknow_feats is not None: 200 | new_features = torch.cat( 201 | [interpolated_feats, unknow_feats], dim=1 202 | ) # (B, C2 + C1, n) 203 | else: 204 | new_features = interpolated_feats 205 | 206 | new_features = new_features.unsqueeze(-1) 207 | new_features = self.mlp(new_features) 208 | 209 | return new_features.squeeze(-1) 210 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/pointnet2_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import warnings 4 | from torch.autograd import Function 5 | from typing import * 6 | 7 | try: 8 | import pointnet2_ops._ext as _ext 9 | except ImportError: 10 | from torch.utils.cpp_extension import load 11 | import glob 12 | import os.path as osp 13 | import os 14 | 15 | warnings.warn("Unable to load pointnet2_ops cpp extension. JIT Compiling.") 16 | 17 | _ext_src_root = osp.join(osp.dirname(__file__), "_ext-src") 18 | _ext_sources = glob.glob(osp.join(_ext_src_root, "src", "*.cpp")) + glob.glob( 19 | osp.join(_ext_src_root, "src", "*.cu") 20 | ) 21 | _ext_headers = glob.glob(osp.join(_ext_src_root, "include", "*")) 22 | 23 | os.environ["TORCH_CUDA_ARCH_LIST"] = "3.7+PTX;5.0;6.0;6.1;6.2;7.0;7.5" 24 | _ext = load( 25 | "_ext", 26 | sources=_ext_sources, 27 | extra_include_paths=[osp.join(_ext_src_root, "include")], 28 | extra_cflags=["-O3"], 29 | extra_cuda_cflags=["-O3", "-Xfatbin", "-compress-all"], 30 | with_cuda=True, 31 | ) 32 | 33 | 34 | class FurthestPointSampling(Function): 35 | @staticmethod 36 | def forward(ctx, xyz, npoint): 37 | # type: (Any, torch.Tensor, int) -> torch.Tensor 38 | r""" 39 | Uses iterative furthest point sampling to select a set of npoint features that have the largest 40 | minimum distance 41 | 42 | Parameters 43 | ---------- 44 | xyz : torch.Tensor 45 | (B, N, 3) tensor where N > npoint 46 | npoint : int32 47 | number of features in the sampled set 48 | 49 | Returns 50 | ------- 51 | torch.Tensor 52 | (B, npoint) tensor containing the set 53 | """ 54 | out = _ext.furthest_point_sampling(xyz, npoint) 55 | 56 | ctx.mark_non_differentiable(out) 57 | 58 | return out 59 | 60 | @staticmethod 61 | def backward(ctx, grad_out): 62 | return () 63 | 64 | 65 | furthest_point_sample = FurthestPointSampling.apply 66 | 67 | 68 | class GatherOperation(Function): 69 | @staticmethod 70 | def forward(ctx, features, idx): 71 | # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor 72 | r""" 73 | 74 | Parameters 75 | ---------- 76 | features : torch.Tensor 77 | (B, C, N) tensor 78 | 79 | idx : torch.Tensor 80 | (B, npoint) tensor of the features to gather 81 | 82 | Returns 83 | ------- 84 | torch.Tensor 85 | (B, C, npoint) tensor 86 | """ 87 | 88 | ctx.save_for_backward(idx, features) 89 | 90 | return _ext.gather_points(features, idx) 91 | 92 | @staticmethod 93 | def backward(ctx, grad_out): 94 | idx, features = ctx.saved_tensors 95 | N = features.size(2) 96 | 97 | grad_features = _ext.gather_points_grad(grad_out.contiguous(), idx, N) 98 | return grad_features, None 99 | 100 | 101 | gather_operation = GatherOperation.apply 102 | 103 | 104 | class ThreeNN(Function): 105 | @staticmethod 106 | def forward(ctx, unknown, known): 107 | # type: (Any, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] 108 | r""" 109 | Find the three nearest neighbors of unknown in known 110 | Parameters 111 | ---------- 112 | unknown : torch.Tensor 113 | (B, n, 3) tensor of known features 114 | known : torch.Tensor 115 | (B, m, 3) tensor of unknown features 116 | 117 | Returns 118 | ------- 119 | dist : torch.Tensor 120 | (B, n, 3) l2 distance to the three nearest neighbors 121 | idx : torch.Tensor 122 | (B, n, 3) index of 3 nearest neighbors 123 | """ 124 | dist2, idx = _ext.three_nn(unknown, known) 125 | dist = torch.sqrt(dist2) 126 | 127 | ctx.mark_non_differentiable(dist, idx) 128 | 129 | return dist, idx 130 | 131 | @staticmethod 132 | def backward(ctx, grad_dist, grad_idx): 133 | return () 134 | 135 | 136 | three_nn = ThreeNN.apply 137 | 138 | 139 | class ThreeInterpolate(Function): 140 | @staticmethod 141 | def forward(ctx, features, idx, weight): 142 | # type(Any, torch.Tensor, torch.Tensor, torch.Tensor) -> Torch.Tensor 143 | r""" 144 | Performs weight linear interpolation on 3 features 145 | Parameters 146 | ---------- 147 | features : torch.Tensor 148 | (B, c, m) Features descriptors to be interpolated from 149 | idx : torch.Tensor 150 | (B, n, 3) three nearest neighbors of the target features in features 151 | weight : torch.Tensor 152 | (B, n, 3) weights 153 | 154 | Returns 155 | ------- 156 | torch.Tensor 157 | (B, c, n) tensor of the interpolated features 158 | """ 159 | ctx.save_for_backward(idx, weight, features) 160 | 161 | return _ext.three_interpolate(features, idx, weight) 162 | 163 | @staticmethod 164 | def backward(ctx, grad_out): 165 | # type: (Any, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor] 166 | r""" 167 | Parameters 168 | ---------- 169 | grad_out : torch.Tensor 170 | (B, c, n) tensor with gradients of ouputs 171 | 172 | Returns 173 | ------- 174 | grad_features : torch.Tensor 175 | (B, c, m) tensor with gradients of features 176 | 177 | None 178 | 179 | None 180 | """ 181 | idx, weight, features = ctx.saved_tensors 182 | m = features.size(2) 183 | 184 | grad_features = _ext.three_interpolate_grad( 185 | grad_out.contiguous(), idx, weight, m 186 | ) 187 | 188 | return grad_features, torch.zeros_like(idx), torch.zeros_like(weight) 189 | 190 | 191 | three_interpolate = ThreeInterpolate.apply 192 | 193 | 194 | class GroupingOperation(Function): 195 | @staticmethod 196 | def forward(ctx, features, idx): 197 | # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor 198 | r""" 199 | 200 | Parameters 201 | ---------- 202 | features : torch.Tensor 203 | (B, C, N) tensor of features to group 204 | idx : torch.Tensor 205 | (B, npoint, nsample) tensor containing the indicies of features to group with 206 | 207 | Returns 208 | ------- 209 | torch.Tensor 210 | (B, C, npoint, nsample) tensor 211 | """ 212 | ctx.save_for_backward(idx, features) 213 | 214 | return _ext.group_points(features, idx) 215 | 216 | @staticmethod 217 | def backward(ctx, grad_out): 218 | # type: (Any, torch.tensor) -> Tuple[torch.Tensor, torch.Tensor] 219 | r""" 220 | 221 | Parameters 222 | ---------- 223 | grad_out : torch.Tensor 224 | (B, C, npoint, nsample) tensor of the gradients of the output from forward 225 | 226 | Returns 227 | ------- 228 | torch.Tensor 229 | (B, C, N) gradient of the features 230 | None 231 | """ 232 | idx, features = ctx.saved_tensors 233 | N = features.size(2) 234 | 235 | grad_features = _ext.group_points_grad(grad_out.contiguous(), idx, N) 236 | 237 | return grad_features, torch.zeros_like(idx) 238 | 239 | 240 | grouping_operation = GroupingOperation.apply 241 | 242 | 243 | class BallQuery(Function): 244 | @staticmethod 245 | def forward(ctx, radius, nsample, xyz, new_xyz): 246 | # type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor 247 | r""" 248 | 249 | Parameters 250 | ---------- 251 | radius : float 252 | radius of the balls 253 | nsample : int 254 | maximum number of features in the balls 255 | xyz : torch.Tensor 256 | (B, N, 3) xyz coordinates of the features 257 | new_xyz : torch.Tensor 258 | (B, npoint, 3) centers of the ball query 259 | 260 | Returns 261 | ------- 262 | torch.Tensor 263 | (B, npoint, nsample) tensor with the indicies of the features that form the query balls 264 | """ 265 | output = _ext.ball_query(new_xyz, xyz, radius, nsample) 266 | 267 | ctx.mark_non_differentiable(output) 268 | 269 | return output 270 | 271 | @staticmethod 272 | def backward(ctx, grad_out): 273 | return () 274 | 275 | 276 | ball_query = BallQuery.apply 277 | 278 | 279 | class QueryAndGroup(nn.Module): 280 | r""" 281 | Groups with a ball query of radius 282 | 283 | Parameters 284 | --------- 285 | radius : float32 286 | Radius of ball 287 | nsample : int32 288 | Maximum number of features to gather in the ball 289 | """ 290 | 291 | def __init__(self, radius, nsample, use_xyz=True): 292 | # type: (QueryAndGroup, float, int, bool) -> None 293 | super(QueryAndGroup, self).__init__() 294 | self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz 295 | 296 | def forward(self, xyz, new_xyz, features=None): 297 | # type: (QueryAndGroup, torch.Tensor. torch.Tensor, torch.Tensor) -> Tuple[Torch.Tensor] 298 | r""" 299 | Parameters 300 | ---------- 301 | xyz : torch.Tensor 302 | xyz coordinates of the features (B, N, 3) 303 | new_xyz : torch.Tensor 304 | centriods (B, npoint, 3) 305 | features : torch.Tensor 306 | Descriptors of the features (B, C, N) 307 | 308 | Returns 309 | ------- 310 | new_features : torch.Tensor 311 | (B, 3 + C, npoint, nsample) tensor 312 | """ 313 | 314 | idx = ball_query(self.radius, self.nsample, xyz, new_xyz) 315 | xyz_trans = xyz.transpose(1, 2).contiguous() 316 | grouped_xyz = grouping_operation(xyz_trans, idx) # (B, 3, npoint, nsample) 317 | grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1) 318 | 319 | if features is not None: 320 | grouped_features = grouping_operation(features, idx) 321 | if self.use_xyz: 322 | new_features = torch.cat( 323 | [grouped_xyz, grouped_features], dim=1 324 | ) # (B, C + 3, npoint, nsample) 325 | else: 326 | new_features = grouped_features 327 | else: 328 | assert ( 329 | self.use_xyz 330 | ), "Cannot have not features and not use xyz as a feature!" 331 | new_features = grouped_xyz 332 | 333 | return new_features 334 | 335 | 336 | class GroupAll(nn.Module): 337 | r""" 338 | Groups all features 339 | 340 | Parameters 341 | --------- 342 | """ 343 | 344 | def __init__(self, use_xyz=True): 345 | # type: (GroupAll, bool) -> None 346 | super(GroupAll, self).__init__() 347 | self.use_xyz = use_xyz 348 | 349 | def forward(self, xyz, new_xyz, features=None): 350 | # type: (GroupAll, torch.Tensor, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor] 351 | r""" 352 | Parameters 353 | ---------- 354 | xyz : torch.Tensor 355 | xyz coordinates of the features (B, N, 3) 356 | new_xyz : torch.Tensor 357 | Ignored 358 | features : torch.Tensor 359 | Descriptors of the features (B, C, N) 360 | 361 | Returns 362 | ------- 363 | new_features : torch.Tensor 364 | (B, C + 3, 1, N) tensor 365 | """ 366 | 367 | grouped_xyz = xyz.transpose(1, 2).unsqueeze(2) 368 | if features is not None: 369 | grouped_features = features.unsqueeze(2) 370 | if self.use_xyz: 371 | new_features = torch.cat( 372 | [grouped_xyz, grouped_features], dim=1 373 | ) # (B, 3 + C, 1, N) 374 | else: 375 | new_features = grouped_features 376 | else: 377 | new_features = grouped_xyz 378 | 379 | return new_features 380 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LongerVision/PointStack/9fe227bc654008452283a075d50f7406a66016be/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query.o -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query_gpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LongerVision/PointStack/9fe227bc654008452283a075d50f7406a66016be/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query_gpu.o -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/bindings.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LongerVision/PointStack/9fe227bc654008452283a075d50f7406a66016be/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/bindings.o -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LongerVision/PointStack/9fe227bc654008452283a075d50f7406a66016be/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points.o -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points_gpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LongerVision/PointStack/9fe227bc654008452283a075d50f7406a66016be/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points_gpu.o -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LongerVision/PointStack/9fe227bc654008452283a075d50f7406a66016be/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate.o -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate_gpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LongerVision/PointStack/9fe227bc654008452283a075d50f7406a66016be/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate_gpu.o -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LongerVision/PointStack/9fe227bc654008452283a075d50f7406a66016be/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling.o -------------------------------------------------------------------------------- /pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling_gpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LongerVision/PointStack/9fe227bc654008452283a075d50f7406a66016be/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling_gpu.o -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 2.1 2 | Name: pointnet2-ops 3 | Version: 3.0.0 4 | Summary: UNKNOWN 5 | Author: Erik Wijmans 6 | License: UNKNOWN 7 | Platform: UNKNOWN 8 | 9 | UNKNOWN 10 | 11 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | MANIFEST.in 2 | setup.py 3 | pointnet2_ops/__init__.py 4 | pointnet2_ops/_version.py 5 | pointnet2_ops/pointnet2_modules.py 6 | pointnet2_ops/pointnet2_utils.py 7 | pointnet2_ops.egg-info/PKG-INFO 8 | pointnet2_ops.egg-info/SOURCES.txt 9 | pointnet2_ops.egg-info/dependency_links.txt 10 | pointnet2_ops.egg-info/requires.txt 11 | pointnet2_ops.egg-info/top_level.txt 12 | pointnet2_ops/_ext-src/include/ball_query.h 13 | pointnet2_ops/_ext-src/include/cuda_utils.h 14 | pointnet2_ops/_ext-src/include/group_points.h 15 | pointnet2_ops/_ext-src/include/interpolate.h 16 | pointnet2_ops/_ext-src/include/sampling.h 17 | pointnet2_ops/_ext-src/include/utils.h 18 | pointnet2_ops/_ext-src/src/ball_query.cpp 19 | pointnet2_ops/_ext-src/src/ball_query_gpu.cu 20 | pointnet2_ops/_ext-src/src/bindings.cpp 21 | pointnet2_ops/_ext-src/src/group_points.cpp 22 | pointnet2_ops/_ext-src/src/group_points_gpu.cu 23 | pointnet2_ops/_ext-src/src/interpolate.cpp 24 | pointnet2_ops/_ext-src/src/interpolate_gpu.cu 25 | pointnet2_ops/_ext-src/src/sampling.cpp 26 | pointnet2_ops/_ext-src/src/sampling_gpu.cu -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops.egg-info/requires.txt: -------------------------------------------------------------------------------- 1 | torch>=1.4 2 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | pointnet2_ops 2 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/__init__.py: -------------------------------------------------------------------------------- 1 | import pointnet2_ops.pointnet2_modules 2 | import pointnet2_ops.pointnet2_utils 3 | from pointnet2_ops._version import __version__ 4 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/include/ball_query.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 5 | const int nsample); 6 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/include/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef _CUDA_UTILS_H 2 | #define _CUDA_UTILS_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include 9 | #include 10 | 11 | #include 12 | 13 | #define TOTAL_THREADS 512 14 | 15 | inline int opt_n_threads(int work_size) { 16 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 17 | 18 | return max(min(1 << pow_2, TOTAL_THREADS), 1); 19 | } 20 | 21 | inline dim3 opt_block_config(int x, int y) { 22 | const int x_threads = opt_n_threads(x); 23 | const int y_threads = 24 | max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 25 | dim3 block_config(x_threads, y_threads, 1); 26 | 27 | return block_config; 28 | } 29 | 30 | #define CUDA_CHECK_ERRORS() \ 31 | do { \ 32 | cudaError_t err = cudaGetLastError(); \ 33 | if (cudaSuccess != err) { \ 34 | fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ 35 | cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ 36 | __FILE__); \ 37 | exit(-1); \ 38 | } \ 39 | } while (0) 40 | 41 | #endif 42 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/include/group_points.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor group_points(at::Tensor points, at::Tensor idx); 5 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 6 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/include/interpolate.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows); 7 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 8 | at::Tensor weight); 9 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 10 | at::Tensor weight, const int m); 11 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/include/sampling.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor gather_points(at::Tensor points, at::Tensor idx); 5 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 6 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples); 7 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/include/utils.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | 5 | #define CHECK_CUDA(x) \ 6 | do { \ 7 | AT_ASSERT(x.is_cuda(), #x " must be a CUDA tensor"); \ 8 | } while (0) 9 | 10 | #define CHECK_CONTIGUOUS(x) \ 11 | do { \ 12 | AT_ASSERT(x.is_contiguous(), #x " must be a contiguous tensor"); \ 13 | } while (0) 14 | 15 | #define CHECK_IS_INT(x) \ 16 | do { \ 17 | AT_ASSERT(x.scalar_type() == at::ScalarType::Int, \ 18 | #x " must be an int tensor"); \ 19 | } while (0) 20 | 21 | #define CHECK_IS_FLOAT(x) \ 22 | do { \ 23 | AT_ASSERT(x.scalar_type() == at::ScalarType::Float, \ 24 | #x " must be a float tensor"); \ 25 | } while (0) 26 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query.h" 2 | #include "utils.h" 3 | 4 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 5 | int nsample, const float *new_xyz, 6 | const float *xyz, int *idx); 7 | 8 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 9 | const int nsample) { 10 | CHECK_CONTIGUOUS(new_xyz); 11 | CHECK_CONTIGUOUS(xyz); 12 | CHECK_IS_FLOAT(new_xyz); 13 | CHECK_IS_FLOAT(xyz); 14 | 15 | if (new_xyz.is_cuda()) { 16 | CHECK_CUDA(xyz); 17 | } 18 | 19 | at::Tensor idx = 20 | torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample}, 21 | at::device(new_xyz.device()).dtype(at::ScalarType::Int)); 22 | 23 | if (new_xyz.is_cuda()) { 24 | query_ball_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), 25 | radius, nsample, new_xyz.data_ptr(), 26 | xyz.data_ptr(), idx.data_ptr()); 27 | } else { 28 | AT_ASSERT(false, "CPU not supported"); 29 | } 30 | 31 | return idx; 32 | } 33 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "cuda_utils.h" 6 | 7 | // input: new_xyz(b, m, 3) xyz(b, n, 3) 8 | // output: idx(b, m, nsample) 9 | __global__ void query_ball_point_kernel(int b, int n, int m, float radius, 10 | int nsample, 11 | const float *__restrict__ new_xyz, 12 | const float *__restrict__ xyz, 13 | int *__restrict__ idx) { 14 | int batch_index = blockIdx.x; 15 | xyz += batch_index * n * 3; 16 | new_xyz += batch_index * m * 3; 17 | idx += m * nsample * batch_index; 18 | 19 | int index = threadIdx.x; 20 | int stride = blockDim.x; 21 | 22 | float radius2 = radius * radius; 23 | for (int j = index; j < m; j += stride) { 24 | float new_x = new_xyz[j * 3 + 0]; 25 | float new_y = new_xyz[j * 3 + 1]; 26 | float new_z = new_xyz[j * 3 + 2]; 27 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { 28 | float x = xyz[k * 3 + 0]; 29 | float y = xyz[k * 3 + 1]; 30 | float z = xyz[k * 3 + 2]; 31 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 32 | (new_z - z) * (new_z - z); 33 | if (d2 < radius2) { 34 | if (cnt == 0) { 35 | for (int l = 0; l < nsample; ++l) { 36 | idx[j * nsample + l] = k; 37 | } 38 | } 39 | idx[j * nsample + cnt] = k; 40 | ++cnt; 41 | } 42 | } 43 | } 44 | } 45 | 46 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 47 | int nsample, const float *new_xyz, 48 | const float *xyz, int *idx) { 49 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 50 | query_ball_point_kernel<<>>( 51 | b, n, m, radius, nsample, new_xyz, xyz, idx); 52 | 53 | CUDA_CHECK_ERRORS(); 54 | } 55 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/src/bindings.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query.h" 2 | #include "group_points.h" 3 | #include "interpolate.h" 4 | #include "sampling.h" 5 | 6 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 7 | m.def("gather_points", &gather_points); 8 | m.def("gather_points_grad", &gather_points_grad); 9 | m.def("furthest_point_sampling", &furthest_point_sampling); 10 | 11 | m.def("three_nn", &three_nn); 12 | m.def("three_interpolate", &three_interpolate); 13 | m.def("three_interpolate_grad", &three_interpolate_grad); 14 | 15 | m.def("ball_query", &ball_query); 16 | 17 | m.def("group_points", &group_points); 18 | m.def("group_points_grad", &group_points_grad); 19 | } 20 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points.cpp: -------------------------------------------------------------------------------- 1 | #include "group_points.h" 2 | #include "utils.h" 3 | 4 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 5 | const float *points, const int *idx, 6 | float *out); 7 | 8 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 9 | int nsample, const float *grad_out, 10 | const int *idx, float *grad_points); 11 | 12 | at::Tensor group_points(at::Tensor points, at::Tensor idx) { 13 | CHECK_CONTIGUOUS(points); 14 | CHECK_CONTIGUOUS(idx); 15 | CHECK_IS_FLOAT(points); 16 | CHECK_IS_INT(idx); 17 | 18 | if (points.is_cuda()) { 19 | CHECK_CUDA(idx); 20 | } 21 | 22 | at::Tensor output = 23 | torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)}, 24 | at::device(points.device()).dtype(at::ScalarType::Float)); 25 | 26 | if (points.is_cuda()) { 27 | group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 28 | idx.size(1), idx.size(2), 29 | points.data_ptr(), idx.data_ptr(), 30 | output.data_ptr()); 31 | } else { 32 | AT_ASSERT(false, "CPU not supported"); 33 | } 34 | 35 | return output; 36 | } 37 | 38 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) { 39 | CHECK_CONTIGUOUS(grad_out); 40 | CHECK_CONTIGUOUS(idx); 41 | CHECK_IS_FLOAT(grad_out); 42 | CHECK_IS_INT(idx); 43 | 44 | if (grad_out.is_cuda()) { 45 | CHECK_CUDA(idx); 46 | } 47 | 48 | at::Tensor output = 49 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 50 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 51 | 52 | if (grad_out.is_cuda()) { 53 | group_points_grad_kernel_wrapper( 54 | grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2), 55 | grad_out.data_ptr(), idx.data_ptr(), 56 | output.data_ptr()); 57 | } else { 58 | AT_ASSERT(false, "CPU not supported"); 59 | } 60 | 61 | return output; 62 | } 63 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "cuda_utils.h" 5 | 6 | // input: points(b, c, n) idx(b, npoints, nsample) 7 | // output: out(b, c, npoints, nsample) 8 | __global__ void group_points_kernel(int b, int c, int n, int npoints, 9 | int nsample, 10 | const float *__restrict__ points, 11 | const int *__restrict__ idx, 12 | float *__restrict__ out) { 13 | int batch_index = blockIdx.x; 14 | points += batch_index * n * c; 15 | idx += batch_index * npoints * nsample; 16 | out += batch_index * npoints * nsample * c; 17 | 18 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 19 | const int stride = blockDim.y * blockDim.x; 20 | for (int i = index; i < c * npoints; i += stride) { 21 | const int l = i / npoints; 22 | const int j = i % npoints; 23 | for (int k = 0; k < nsample; ++k) { 24 | int ii = idx[j * nsample + k]; 25 | out[(l * npoints + j) * nsample + k] = points[l * n + ii]; 26 | } 27 | } 28 | } 29 | 30 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 31 | const float *points, const int *idx, 32 | float *out) { 33 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 34 | 35 | group_points_kernel<<>>( 36 | b, c, n, npoints, nsample, points, idx, out); 37 | 38 | CUDA_CHECK_ERRORS(); 39 | } 40 | 41 | // input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample) 42 | // output: grad_points(b, c, n) 43 | __global__ void group_points_grad_kernel(int b, int c, int n, int npoints, 44 | int nsample, 45 | const float *__restrict__ grad_out, 46 | const int *__restrict__ idx, 47 | float *__restrict__ grad_points) { 48 | int batch_index = blockIdx.x; 49 | grad_out += batch_index * npoints * nsample * c; 50 | idx += batch_index * npoints * nsample; 51 | grad_points += batch_index * n * c; 52 | 53 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 54 | const int stride = blockDim.y * blockDim.x; 55 | for (int i = index; i < c * npoints; i += stride) { 56 | const int l = i / npoints; 57 | const int j = i % npoints; 58 | for (int k = 0; k < nsample; ++k) { 59 | int ii = idx[j * nsample + k]; 60 | atomicAdd(grad_points + l * n + ii, 61 | grad_out[(l * npoints + j) * nsample + k]); 62 | } 63 | } 64 | } 65 | 66 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 67 | int nsample, const float *grad_out, 68 | const int *idx, float *grad_points) { 69 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 70 | 71 | group_points_grad_kernel<<>>( 72 | b, c, n, npoints, nsample, grad_out, idx, grad_points); 73 | 74 | CUDA_CHECK_ERRORS(); 75 | } 76 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate.cpp: -------------------------------------------------------------------------------- 1 | #include "interpolate.h" 2 | #include "utils.h" 3 | 4 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 5 | const float *known, float *dist2, int *idx); 6 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 7 | const float *points, const int *idx, 8 | const float *weight, float *out); 9 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 10 | const float *grad_out, 11 | const int *idx, const float *weight, 12 | float *grad_points); 13 | 14 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows) { 15 | CHECK_CONTIGUOUS(unknowns); 16 | CHECK_CONTIGUOUS(knows); 17 | CHECK_IS_FLOAT(unknowns); 18 | CHECK_IS_FLOAT(knows); 19 | 20 | if (unknowns.is_cuda()) { 21 | CHECK_CUDA(knows); 22 | } 23 | 24 | at::Tensor idx = 25 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 26 | at::device(unknowns.device()).dtype(at::ScalarType::Int)); 27 | at::Tensor dist2 = 28 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 29 | at::device(unknowns.device()).dtype(at::ScalarType::Float)); 30 | 31 | if (unknowns.is_cuda()) { 32 | three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1), 33 | unknowns.data_ptr(), knows.data_ptr(), 34 | dist2.data_ptr(), idx.data_ptr()); 35 | } else { 36 | AT_ASSERT(false, "CPU not supported"); 37 | } 38 | 39 | return {dist2, idx}; 40 | } 41 | 42 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 43 | at::Tensor weight) { 44 | CHECK_CONTIGUOUS(points); 45 | CHECK_CONTIGUOUS(idx); 46 | CHECK_CONTIGUOUS(weight); 47 | CHECK_IS_FLOAT(points); 48 | CHECK_IS_INT(idx); 49 | CHECK_IS_FLOAT(weight); 50 | 51 | if (points.is_cuda()) { 52 | CHECK_CUDA(idx); 53 | CHECK_CUDA(weight); 54 | } 55 | 56 | at::Tensor output = 57 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 58 | at::device(points.device()).dtype(at::ScalarType::Float)); 59 | 60 | if (points.is_cuda()) { 61 | three_interpolate_kernel_wrapper( 62 | points.size(0), points.size(1), points.size(2), idx.size(1), 63 | points.data_ptr(), idx.data_ptr(), weight.data_ptr(), 64 | output.data_ptr()); 65 | } else { 66 | AT_ASSERT(false, "CPU not supported"); 67 | } 68 | 69 | return output; 70 | } 71 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 72 | at::Tensor weight, const int m) { 73 | CHECK_CONTIGUOUS(grad_out); 74 | CHECK_CONTIGUOUS(idx); 75 | CHECK_CONTIGUOUS(weight); 76 | CHECK_IS_FLOAT(grad_out); 77 | CHECK_IS_INT(idx); 78 | CHECK_IS_FLOAT(weight); 79 | 80 | if (grad_out.is_cuda()) { 81 | CHECK_CUDA(idx); 82 | CHECK_CUDA(weight); 83 | } 84 | 85 | at::Tensor output = 86 | torch::zeros({grad_out.size(0), grad_out.size(1), m}, 87 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 88 | 89 | if (grad_out.is_cuda()) { 90 | three_interpolate_grad_kernel_wrapper( 91 | grad_out.size(0), grad_out.size(1), grad_out.size(2), m, 92 | grad_out.data_ptr(), idx.data_ptr(), 93 | weight.data_ptr(), output.data_ptr()); 94 | } else { 95 | AT_ASSERT(false, "CPU not supported"); 96 | } 97 | 98 | return output; 99 | } 100 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "cuda_utils.h" 6 | 7 | // input: unknown(b, n, 3) known(b, m, 3) 8 | // output: dist2(b, n, 3), idx(b, n, 3) 9 | __global__ void three_nn_kernel(int b, int n, int m, 10 | const float *__restrict__ unknown, 11 | const float *__restrict__ known, 12 | float *__restrict__ dist2, 13 | int *__restrict__ idx) { 14 | int batch_index = blockIdx.x; 15 | unknown += batch_index * n * 3; 16 | known += batch_index * m * 3; 17 | dist2 += batch_index * n * 3; 18 | idx += batch_index * n * 3; 19 | 20 | int index = threadIdx.x; 21 | int stride = blockDim.x; 22 | for (int j = index; j < n; j += stride) { 23 | float ux = unknown[j * 3 + 0]; 24 | float uy = unknown[j * 3 + 1]; 25 | float uz = unknown[j * 3 + 2]; 26 | 27 | double best1 = 1e40, best2 = 1e40, best3 = 1e40; 28 | int besti1 = 0, besti2 = 0, besti3 = 0; 29 | for (int k = 0; k < m; ++k) { 30 | float x = known[k * 3 + 0]; 31 | float y = known[k * 3 + 1]; 32 | float z = known[k * 3 + 2]; 33 | float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z); 34 | if (d < best1) { 35 | best3 = best2; 36 | besti3 = besti2; 37 | best2 = best1; 38 | besti2 = besti1; 39 | best1 = d; 40 | besti1 = k; 41 | } else if (d < best2) { 42 | best3 = best2; 43 | besti3 = besti2; 44 | best2 = d; 45 | besti2 = k; 46 | } else if (d < best3) { 47 | best3 = d; 48 | besti3 = k; 49 | } 50 | } 51 | dist2[j * 3 + 0] = best1; 52 | dist2[j * 3 + 1] = best2; 53 | dist2[j * 3 + 2] = best3; 54 | 55 | idx[j * 3 + 0] = besti1; 56 | idx[j * 3 + 1] = besti2; 57 | idx[j * 3 + 2] = besti3; 58 | } 59 | } 60 | 61 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 62 | const float *known, float *dist2, int *idx) { 63 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 64 | three_nn_kernel<<>>(b, n, m, unknown, known, 65 | dist2, idx); 66 | 67 | CUDA_CHECK_ERRORS(); 68 | } 69 | 70 | // input: points(b, c, m), idx(b, n, 3), weight(b, n, 3) 71 | // output: out(b, c, n) 72 | __global__ void three_interpolate_kernel(int b, int c, int m, int n, 73 | const float *__restrict__ points, 74 | const int *__restrict__ idx, 75 | const float *__restrict__ weight, 76 | float *__restrict__ out) { 77 | int batch_index = blockIdx.x; 78 | points += batch_index * m * c; 79 | 80 | idx += batch_index * n * 3; 81 | weight += batch_index * n * 3; 82 | 83 | out += batch_index * n * c; 84 | 85 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 86 | const int stride = blockDim.y * blockDim.x; 87 | for (int i = index; i < c * n; i += stride) { 88 | const int l = i / n; 89 | const int j = i % n; 90 | float w1 = weight[j * 3 + 0]; 91 | float w2 = weight[j * 3 + 1]; 92 | float w3 = weight[j * 3 + 2]; 93 | 94 | int i1 = idx[j * 3 + 0]; 95 | int i2 = idx[j * 3 + 1]; 96 | int i3 = idx[j * 3 + 2]; 97 | 98 | out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 + 99 | points[l * m + i3] * w3; 100 | } 101 | } 102 | 103 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 104 | const float *points, const int *idx, 105 | const float *weight, float *out) { 106 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 107 | three_interpolate_kernel<<>>( 108 | b, c, m, n, points, idx, weight, out); 109 | 110 | CUDA_CHECK_ERRORS(); 111 | } 112 | 113 | // input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3) 114 | // output: grad_points(b, c, m) 115 | 116 | __global__ void three_interpolate_grad_kernel( 117 | int b, int c, int n, int m, const float *__restrict__ grad_out, 118 | const int *__restrict__ idx, const float *__restrict__ weight, 119 | float *__restrict__ grad_points) { 120 | int batch_index = blockIdx.x; 121 | grad_out += batch_index * n * c; 122 | idx += batch_index * n * 3; 123 | weight += batch_index * n * 3; 124 | grad_points += batch_index * m * c; 125 | 126 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 127 | const int stride = blockDim.y * blockDim.x; 128 | for (int i = index; i < c * n; i += stride) { 129 | const int l = i / n; 130 | const int j = i % n; 131 | float w1 = weight[j * 3 + 0]; 132 | float w2 = weight[j * 3 + 1]; 133 | float w3 = weight[j * 3 + 2]; 134 | 135 | int i1 = idx[j * 3 + 0]; 136 | int i2 = idx[j * 3 + 1]; 137 | int i3 = idx[j * 3 + 2]; 138 | 139 | atomicAdd(grad_points + l * m + i1, grad_out[i] * w1); 140 | atomicAdd(grad_points + l * m + i2, grad_out[i] * w2); 141 | atomicAdd(grad_points + l * m + i3, grad_out[i] * w3); 142 | } 143 | } 144 | 145 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 146 | const float *grad_out, 147 | const int *idx, const float *weight, 148 | float *grad_points) { 149 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 150 | three_interpolate_grad_kernel<<>>( 151 | b, c, n, m, grad_out, idx, weight, grad_points); 152 | 153 | CUDA_CHECK_ERRORS(); 154 | } 155 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling.cpp: -------------------------------------------------------------------------------- 1 | #include "sampling.h" 2 | #include "utils.h" 3 | 4 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 5 | const float *points, const int *idx, 6 | float *out); 7 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 8 | const float *grad_out, const int *idx, 9 | float *grad_points); 10 | 11 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 12 | const float *dataset, float *temp, 13 | int *idxs); 14 | 15 | at::Tensor gather_points(at::Tensor points, at::Tensor idx) { 16 | CHECK_CONTIGUOUS(points); 17 | CHECK_CONTIGUOUS(idx); 18 | CHECK_IS_FLOAT(points); 19 | CHECK_IS_INT(idx); 20 | 21 | if (points.is_cuda()) { 22 | CHECK_CUDA(idx); 23 | } 24 | 25 | at::Tensor output = 26 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 27 | at::device(points.device()).dtype(at::ScalarType::Float)); 28 | 29 | if (points.is_cuda()) { 30 | gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 31 | idx.size(1), points.data_ptr(), 32 | idx.data_ptr(), output.data_ptr()); 33 | } else { 34 | AT_ASSERT(false, "CPU not supported"); 35 | } 36 | 37 | return output; 38 | } 39 | 40 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, 41 | const int n) { 42 | CHECK_CONTIGUOUS(grad_out); 43 | CHECK_CONTIGUOUS(idx); 44 | CHECK_IS_FLOAT(grad_out); 45 | CHECK_IS_INT(idx); 46 | 47 | if (grad_out.is_cuda()) { 48 | CHECK_CUDA(idx); 49 | } 50 | 51 | at::Tensor output = 52 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 53 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 54 | 55 | if (grad_out.is_cuda()) { 56 | gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n, 57 | idx.size(1), grad_out.data_ptr(), 58 | idx.data_ptr(), 59 | output.data_ptr()); 60 | } else { 61 | AT_ASSERT(false, "CPU not supported"); 62 | } 63 | 64 | return output; 65 | } 66 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) { 67 | CHECK_CONTIGUOUS(points); 68 | CHECK_IS_FLOAT(points); 69 | 70 | at::Tensor output = 71 | torch::zeros({points.size(0), nsamples}, 72 | at::device(points.device()).dtype(at::ScalarType::Int)); 73 | 74 | at::Tensor tmp = 75 | torch::full({points.size(0), points.size(1)}, 1e10, 76 | at::device(points.device()).dtype(at::ScalarType::Float)); 77 | 78 | if (points.is_cuda()) { 79 | furthest_point_sampling_kernel_wrapper( 80 | points.size(0), points.size(1), nsamples, points.data_ptr(), 81 | tmp.data_ptr(), output.data_ptr()); 82 | } else { 83 | AT_ASSERT(false, "CPU not supported"); 84 | } 85 | 86 | return output; 87 | } 88 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "cuda_utils.h" 5 | 6 | // input: points(b, c, n) idx(b, m) 7 | // output: out(b, c, m) 8 | __global__ void gather_points_kernel(int b, int c, int n, int m, 9 | const float *__restrict__ points, 10 | const int *__restrict__ idx, 11 | float *__restrict__ out) { 12 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 13 | for (int l = blockIdx.y; l < c; l += gridDim.y) { 14 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 15 | int a = idx[i * m + j]; 16 | out[(i * c + l) * m + j] = points[(i * c + l) * n + a]; 17 | } 18 | } 19 | } 20 | } 21 | 22 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 23 | const float *points, const int *idx, 24 | float *out) { 25 | gather_points_kernel<<>>(b, c, n, npoints, 27 | points, idx, out); 28 | 29 | CUDA_CHECK_ERRORS(); 30 | } 31 | 32 | // input: grad_out(b, c, m) idx(b, m) 33 | // output: grad_points(b, c, n) 34 | __global__ void gather_points_grad_kernel(int b, int c, int n, int m, 35 | const float *__restrict__ grad_out, 36 | const int *__restrict__ idx, 37 | float *__restrict__ grad_points) { 38 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 39 | for (int l = blockIdx.y; l < c; l += gridDim.y) { 40 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 41 | int a = idx[i * m + j]; 42 | atomicAdd(grad_points + (i * c + l) * n + a, 43 | grad_out[(i * c + l) * m + j]); 44 | } 45 | } 46 | } 47 | } 48 | 49 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 50 | const float *grad_out, const int *idx, 51 | float *grad_points) { 52 | gather_points_grad_kernel<<>>( 54 | b, c, n, npoints, grad_out, idx, grad_points); 55 | 56 | CUDA_CHECK_ERRORS(); 57 | } 58 | 59 | __device__ void __update(float *__restrict__ dists, int *__restrict__ dists_i, 60 | int idx1, int idx2) { 61 | const float v1 = dists[idx1], v2 = dists[idx2]; 62 | const int i1 = dists_i[idx1], i2 = dists_i[idx2]; 63 | dists[idx1] = max(v1, v2); 64 | dists_i[idx1] = v2 > v1 ? i2 : i1; 65 | } 66 | 67 | // Input dataset: (b, n, 3), tmp: (b, n) 68 | // Ouput idxs (b, m) 69 | template 70 | __global__ void furthest_point_sampling_kernel( 71 | int b, int n, int m, const float *__restrict__ dataset, 72 | float *__restrict__ temp, int *__restrict__ idxs) { 73 | if (m <= 0) return; 74 | __shared__ float dists[block_size]; 75 | __shared__ int dists_i[block_size]; 76 | 77 | int batch_index = blockIdx.x; 78 | dataset += batch_index * n * 3; 79 | temp += batch_index * n; 80 | idxs += batch_index * m; 81 | 82 | int tid = threadIdx.x; 83 | const int stride = block_size; 84 | 85 | int old = 0; 86 | if (threadIdx.x == 0) idxs[0] = old; 87 | 88 | __syncthreads(); 89 | for (int j = 1; j < m; j++) { 90 | int besti = 0; 91 | float best = -1; 92 | float x1 = dataset[old * 3 + 0]; 93 | float y1 = dataset[old * 3 + 1]; 94 | float z1 = dataset[old * 3 + 2]; 95 | for (int k = tid; k < n; k += stride) { 96 | float x2, y2, z2; 97 | x2 = dataset[k * 3 + 0]; 98 | y2 = dataset[k * 3 + 1]; 99 | z2 = dataset[k * 3 + 2]; 100 | float mag = (x2 * x2) + (y2 * y2) + (z2 * z2); 101 | if (mag <= 1e-3) continue; 102 | 103 | float d = 104 | (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1); 105 | 106 | float d2 = min(d, temp[k]); 107 | temp[k] = d2; 108 | besti = d2 > best ? k : besti; 109 | best = d2 > best ? d2 : best; 110 | } 111 | dists[tid] = best; 112 | dists_i[tid] = besti; 113 | __syncthreads(); 114 | 115 | if (block_size >= 512) { 116 | if (tid < 256) { 117 | __update(dists, dists_i, tid, tid + 256); 118 | } 119 | __syncthreads(); 120 | } 121 | if (block_size >= 256) { 122 | if (tid < 128) { 123 | __update(dists, dists_i, tid, tid + 128); 124 | } 125 | __syncthreads(); 126 | } 127 | if (block_size >= 128) { 128 | if (tid < 64) { 129 | __update(dists, dists_i, tid, tid + 64); 130 | } 131 | __syncthreads(); 132 | } 133 | if (block_size >= 64) { 134 | if (tid < 32) { 135 | __update(dists, dists_i, tid, tid + 32); 136 | } 137 | __syncthreads(); 138 | } 139 | if (block_size >= 32) { 140 | if (tid < 16) { 141 | __update(dists, dists_i, tid, tid + 16); 142 | } 143 | __syncthreads(); 144 | } 145 | if (block_size >= 16) { 146 | if (tid < 8) { 147 | __update(dists, dists_i, tid, tid + 8); 148 | } 149 | __syncthreads(); 150 | } 151 | if (block_size >= 8) { 152 | if (tid < 4) { 153 | __update(dists, dists_i, tid, tid + 4); 154 | } 155 | __syncthreads(); 156 | } 157 | if (block_size >= 4) { 158 | if (tid < 2) { 159 | __update(dists, dists_i, tid, tid + 2); 160 | } 161 | __syncthreads(); 162 | } 163 | if (block_size >= 2) { 164 | if (tid < 1) { 165 | __update(dists, dists_i, tid, tid + 1); 166 | } 167 | __syncthreads(); 168 | } 169 | 170 | old = dists_i[0]; 171 | if (tid == 0) idxs[j] = old; 172 | } 173 | } 174 | 175 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 176 | const float *dataset, float *temp, 177 | int *idxs) { 178 | unsigned int n_threads = opt_n_threads(n); 179 | 180 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 181 | 182 | switch (n_threads) { 183 | case 512: 184 | furthest_point_sampling_kernel<512> 185 | <<>>(b, n, m, dataset, temp, idxs); 186 | break; 187 | case 256: 188 | furthest_point_sampling_kernel<256> 189 | <<>>(b, n, m, dataset, temp, idxs); 190 | break; 191 | case 128: 192 | furthest_point_sampling_kernel<128> 193 | <<>>(b, n, m, dataset, temp, idxs); 194 | break; 195 | case 64: 196 | furthest_point_sampling_kernel<64> 197 | <<>>(b, n, m, dataset, temp, idxs); 198 | break; 199 | case 32: 200 | furthest_point_sampling_kernel<32> 201 | <<>>(b, n, m, dataset, temp, idxs); 202 | break; 203 | case 16: 204 | furthest_point_sampling_kernel<16> 205 | <<>>(b, n, m, dataset, temp, idxs); 206 | break; 207 | case 8: 208 | furthest_point_sampling_kernel<8> 209 | <<>>(b, n, m, dataset, temp, idxs); 210 | break; 211 | case 4: 212 | furthest_point_sampling_kernel<4> 213 | <<>>(b, n, m, dataset, temp, idxs); 214 | break; 215 | case 2: 216 | furthest_point_sampling_kernel<2> 217 | <<>>(b, n, m, dataset, temp, idxs); 218 | break; 219 | case 1: 220 | furthest_point_sampling_kernel<1> 221 | <<>>(b, n, m, dataset, temp, idxs); 222 | break; 223 | default: 224 | furthest_point_sampling_kernel<512> 225 | <<>>(b, n, m, dataset, temp, idxs); 226 | } 227 | 228 | CUDA_CHECK_ERRORS(); 229 | } 230 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/_version.py: -------------------------------------------------------------------------------- 1 | __version__ = "3.0.0" 2 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/pointnet2_modules.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Tuple 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | from pointnet2_ops import pointnet2_utils 7 | 8 | 9 | def build_shared_mlp(mlp_spec: List[int], bn: bool = True): 10 | layers = [] 11 | for i in range(1, len(mlp_spec)): 12 | layers.append( 13 | nn.Conv2d(mlp_spec[i - 1], mlp_spec[i], kernel_size=1, bias=not bn) 14 | ) 15 | if bn: 16 | layers.append(nn.BatchNorm2d(mlp_spec[i])) 17 | layers.append(nn.ReLU(True)) 18 | 19 | return nn.Sequential(*layers) 20 | 21 | 22 | class _PointnetSAModuleBase(nn.Module): 23 | def __init__(self): 24 | super(_PointnetSAModuleBase, self).__init__() 25 | self.npoint = None 26 | self.groupers = None 27 | self.mlps = None 28 | 29 | def forward( 30 | self, xyz: torch.Tensor, features: Optional[torch.Tensor] 31 | ) -> Tuple[torch.Tensor, torch.Tensor]: 32 | r""" 33 | Parameters 34 | ---------- 35 | xyz : torch.Tensor 36 | (B, N, 3) tensor of the xyz coordinates of the features 37 | features : torch.Tensor 38 | (B, C, N) tensor of the descriptors of the the features 39 | 40 | Returns 41 | ------- 42 | new_xyz : torch.Tensor 43 | (B, npoint, 3) tensor of the new features' xyz 44 | new_features : torch.Tensor 45 | (B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors 46 | """ 47 | 48 | new_features_list = [] 49 | 50 | xyz_flipped = xyz.transpose(1, 2).contiguous() 51 | new_xyz = ( 52 | pointnet2_utils.gather_operation( 53 | xyz_flipped, pointnet2_utils.furthest_point_sample(xyz, self.npoint) 54 | ) 55 | .transpose(1, 2) 56 | .contiguous() 57 | if self.npoint is not None 58 | else None 59 | ) 60 | 61 | for i in range(len(self.groupers)): 62 | new_features = self.groupers[i]( 63 | xyz, new_xyz, features 64 | ) # (B, C, npoint, nsample) 65 | 66 | new_features = self.mlps[i](new_features) # (B, mlp[-1], npoint, nsample) 67 | new_features = F.max_pool2d( 68 | new_features, kernel_size=[1, new_features.size(3)] 69 | ) # (B, mlp[-1], npoint, 1) 70 | new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) 71 | 72 | new_features_list.append(new_features) 73 | 74 | return new_xyz, torch.cat(new_features_list, dim=1) 75 | 76 | 77 | class PointnetSAModuleMSG(_PointnetSAModuleBase): 78 | r"""Pointnet set abstrction layer with multiscale grouping 79 | 80 | Parameters 81 | ---------- 82 | npoint : int 83 | Number of features 84 | radii : list of float32 85 | list of radii to group with 86 | nsamples : list of int32 87 | Number of samples in each ball query 88 | mlps : list of list of int32 89 | Spec of the pointnet before the global max_pool for each scale 90 | bn : bool 91 | Use batchnorm 92 | """ 93 | 94 | def __init__(self, npoint, radii, nsamples, mlps, bn=True, use_xyz=True): 95 | # type: (PointnetSAModuleMSG, int, List[float], List[int], List[List[int]], bool, bool) -> None 96 | super(PointnetSAModuleMSG, self).__init__() 97 | 98 | assert len(radii) == len(nsamples) == len(mlps) 99 | 100 | self.npoint = npoint 101 | self.groupers = nn.ModuleList() 102 | self.mlps = nn.ModuleList() 103 | for i in range(len(radii)): 104 | radius = radii[i] 105 | nsample = nsamples[i] 106 | self.groupers.append( 107 | pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz) 108 | if npoint is not None 109 | else pointnet2_utils.GroupAll(use_xyz) 110 | ) 111 | mlp_spec = mlps[i] 112 | if use_xyz: 113 | mlp_spec[0] += 3 114 | 115 | self.mlps.append(build_shared_mlp(mlp_spec, bn)) 116 | 117 | 118 | class PointnetSAModule(PointnetSAModuleMSG): 119 | r"""Pointnet set abstrction layer 120 | 121 | Parameters 122 | ---------- 123 | npoint : int 124 | Number of features 125 | radius : float 126 | Radius of ball 127 | nsample : int 128 | Number of samples in the ball query 129 | mlp : list 130 | Spec of the pointnet before the global max_pool 131 | bn : bool 132 | Use batchnorm 133 | """ 134 | 135 | def __init__( 136 | self, mlp, npoint=None, radius=None, nsample=None, bn=True, use_xyz=True 137 | ): 138 | # type: (PointnetSAModule, List[int], int, float, int, bool, bool) -> None 139 | super(PointnetSAModule, self).__init__( 140 | mlps=[mlp], 141 | npoint=npoint, 142 | radii=[radius], 143 | nsamples=[nsample], 144 | bn=bn, 145 | use_xyz=use_xyz, 146 | ) 147 | 148 | 149 | class PointnetFPModule(nn.Module): 150 | r"""Propigates the features of one set to another 151 | 152 | Parameters 153 | ---------- 154 | mlp : list 155 | Pointnet module parameters 156 | bn : bool 157 | Use batchnorm 158 | """ 159 | 160 | def __init__(self, mlp, bn=True): 161 | # type: (PointnetFPModule, List[int], bool) -> None 162 | super(PointnetFPModule, self).__init__() 163 | self.mlp = build_shared_mlp(mlp, bn=bn) 164 | 165 | def forward(self, unknown, known, unknow_feats, known_feats): 166 | # type: (PointnetFPModule, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor 167 | r""" 168 | Parameters 169 | ---------- 170 | unknown : torch.Tensor 171 | (B, n, 3) tensor of the xyz positions of the unknown features 172 | known : torch.Tensor 173 | (B, m, 3) tensor of the xyz positions of the known features 174 | unknow_feats : torch.Tensor 175 | (B, C1, n) tensor of the features to be propigated to 176 | known_feats : torch.Tensor 177 | (B, C2, m) tensor of features to be propigated 178 | 179 | Returns 180 | ------- 181 | new_features : torch.Tensor 182 | (B, mlp[-1], n) tensor of the features of the unknown features 183 | """ 184 | 185 | if known is not None: 186 | dist, idx = pointnet2_utils.three_nn(unknown, known) 187 | dist_recip = 1.0 / (dist + 1e-8) 188 | norm = torch.sum(dist_recip, dim=2, keepdim=True) 189 | weight = dist_recip / norm 190 | 191 | interpolated_feats = pointnet2_utils.three_interpolate( 192 | known_feats, idx, weight 193 | ) 194 | else: 195 | interpolated_feats = known_feats.expand( 196 | *(known_feats.size()[0:2] + [unknown.size(1)]) 197 | ) 198 | 199 | if unknow_feats is not None: 200 | new_features = torch.cat( 201 | [interpolated_feats, unknow_feats], dim=1 202 | ) # (B, C2 + C1, n) 203 | else: 204 | new_features = interpolated_feats 205 | 206 | new_features = new_features.unsqueeze(-1) 207 | new_features = self.mlp(new_features) 208 | 209 | return new_features.squeeze(-1) 210 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/pointnet2_ops/pointnet2_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import warnings 4 | from torch.autograd import Function 5 | from typing import * 6 | 7 | try: 8 | import pointnet2_ops._ext as _ext 9 | except ImportError: 10 | from torch.utils.cpp_extension import load 11 | import glob 12 | import os.path as osp 13 | import os 14 | 15 | warnings.warn("Unable to load pointnet2_ops cpp extension. JIT Compiling.") 16 | 17 | _ext_src_root = osp.join(osp.dirname(__file__), "_ext-src") 18 | _ext_sources = glob.glob(osp.join(_ext_src_root, "src", "*.cpp")) + glob.glob( 19 | osp.join(_ext_src_root, "src", "*.cu") 20 | ) 21 | _ext_headers = glob.glob(osp.join(_ext_src_root, "include", "*")) 22 | 23 | os.environ["TORCH_CUDA_ARCH_LIST"] = "3.7+PTX;5.0;6.0;6.1;6.2;7.0;7.5" 24 | _ext = load( 25 | "_ext", 26 | sources=_ext_sources, 27 | extra_include_paths=[osp.join(_ext_src_root, "include")], 28 | extra_cflags=["-O3"], 29 | extra_cuda_cflags=["-O3", "-Xfatbin", "-compress-all"], 30 | with_cuda=True, 31 | ) 32 | 33 | 34 | class FurthestPointSampling(Function): 35 | @staticmethod 36 | def forward(ctx, xyz, npoint): 37 | # type: (Any, torch.Tensor, int) -> torch.Tensor 38 | r""" 39 | Uses iterative furthest point sampling to select a set of npoint features that have the largest 40 | minimum distance 41 | 42 | Parameters 43 | ---------- 44 | xyz : torch.Tensor 45 | (B, N, 3) tensor where N > npoint 46 | npoint : int32 47 | number of features in the sampled set 48 | 49 | Returns 50 | ------- 51 | torch.Tensor 52 | (B, npoint) tensor containing the set 53 | """ 54 | out = _ext.furthest_point_sampling(xyz, npoint) 55 | 56 | ctx.mark_non_differentiable(out) 57 | 58 | return out 59 | 60 | @staticmethod 61 | def backward(ctx, grad_out): 62 | return () 63 | 64 | 65 | furthest_point_sample = FurthestPointSampling.apply 66 | 67 | 68 | class GatherOperation(Function): 69 | @staticmethod 70 | def forward(ctx, features, idx): 71 | # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor 72 | r""" 73 | 74 | Parameters 75 | ---------- 76 | features : torch.Tensor 77 | (B, C, N) tensor 78 | 79 | idx : torch.Tensor 80 | (B, npoint) tensor of the features to gather 81 | 82 | Returns 83 | ------- 84 | torch.Tensor 85 | (B, C, npoint) tensor 86 | """ 87 | 88 | ctx.save_for_backward(idx, features) 89 | 90 | return _ext.gather_points(features, idx) 91 | 92 | @staticmethod 93 | def backward(ctx, grad_out): 94 | idx, features = ctx.saved_tensors 95 | N = features.size(2) 96 | 97 | grad_features = _ext.gather_points_grad(grad_out.contiguous(), idx, N) 98 | return grad_features, None 99 | 100 | 101 | gather_operation = GatherOperation.apply 102 | 103 | 104 | class ThreeNN(Function): 105 | @staticmethod 106 | def forward(ctx, unknown, known): 107 | # type: (Any, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] 108 | r""" 109 | Find the three nearest neighbors of unknown in known 110 | Parameters 111 | ---------- 112 | unknown : torch.Tensor 113 | (B, n, 3) tensor of known features 114 | known : torch.Tensor 115 | (B, m, 3) tensor of unknown features 116 | 117 | Returns 118 | ------- 119 | dist : torch.Tensor 120 | (B, n, 3) l2 distance to the three nearest neighbors 121 | idx : torch.Tensor 122 | (B, n, 3) index of 3 nearest neighbors 123 | """ 124 | dist2, idx = _ext.three_nn(unknown, known) 125 | dist = torch.sqrt(dist2) 126 | 127 | ctx.mark_non_differentiable(dist, idx) 128 | 129 | return dist, idx 130 | 131 | @staticmethod 132 | def backward(ctx, grad_dist, grad_idx): 133 | return () 134 | 135 | 136 | three_nn = ThreeNN.apply 137 | 138 | 139 | class ThreeInterpolate(Function): 140 | @staticmethod 141 | def forward(ctx, features, idx, weight): 142 | # type(Any, torch.Tensor, torch.Tensor, torch.Tensor) -> Torch.Tensor 143 | r""" 144 | Performs weight linear interpolation on 3 features 145 | Parameters 146 | ---------- 147 | features : torch.Tensor 148 | (B, c, m) Features descriptors to be interpolated from 149 | idx : torch.Tensor 150 | (B, n, 3) three nearest neighbors of the target features in features 151 | weight : torch.Tensor 152 | (B, n, 3) weights 153 | 154 | Returns 155 | ------- 156 | torch.Tensor 157 | (B, c, n) tensor of the interpolated features 158 | """ 159 | ctx.save_for_backward(idx, weight, features) 160 | 161 | return _ext.three_interpolate(features, idx, weight) 162 | 163 | @staticmethod 164 | def backward(ctx, grad_out): 165 | # type: (Any, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor] 166 | r""" 167 | Parameters 168 | ---------- 169 | grad_out : torch.Tensor 170 | (B, c, n) tensor with gradients of ouputs 171 | 172 | Returns 173 | ------- 174 | grad_features : torch.Tensor 175 | (B, c, m) tensor with gradients of features 176 | 177 | None 178 | 179 | None 180 | """ 181 | idx, weight, features = ctx.saved_tensors 182 | m = features.size(2) 183 | 184 | grad_features = _ext.three_interpolate_grad( 185 | grad_out.contiguous(), idx, weight, m 186 | ) 187 | 188 | return grad_features, torch.zeros_like(idx), torch.zeros_like(weight) 189 | 190 | 191 | three_interpolate = ThreeInterpolate.apply 192 | 193 | 194 | class GroupingOperation(Function): 195 | @staticmethod 196 | def forward(ctx, features, idx): 197 | # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor 198 | r""" 199 | 200 | Parameters 201 | ---------- 202 | features : torch.Tensor 203 | (B, C, N) tensor of features to group 204 | idx : torch.Tensor 205 | (B, npoint, nsample) tensor containing the indicies of features to group with 206 | 207 | Returns 208 | ------- 209 | torch.Tensor 210 | (B, C, npoint, nsample) tensor 211 | """ 212 | ctx.save_for_backward(idx, features) 213 | 214 | return _ext.group_points(features, idx) 215 | 216 | @staticmethod 217 | def backward(ctx, grad_out): 218 | # type: (Any, torch.tensor) -> Tuple[torch.Tensor, torch.Tensor] 219 | r""" 220 | 221 | Parameters 222 | ---------- 223 | grad_out : torch.Tensor 224 | (B, C, npoint, nsample) tensor of the gradients of the output from forward 225 | 226 | Returns 227 | ------- 228 | torch.Tensor 229 | (B, C, N) gradient of the features 230 | None 231 | """ 232 | idx, features = ctx.saved_tensors 233 | N = features.size(2) 234 | 235 | grad_features = _ext.group_points_grad(grad_out.contiguous(), idx, N) 236 | 237 | return grad_features, torch.zeros_like(idx) 238 | 239 | 240 | grouping_operation = GroupingOperation.apply 241 | 242 | 243 | class BallQuery(Function): 244 | @staticmethod 245 | def forward(ctx, radius, nsample, xyz, new_xyz): 246 | # type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor 247 | r""" 248 | 249 | Parameters 250 | ---------- 251 | radius : float 252 | radius of the balls 253 | nsample : int 254 | maximum number of features in the balls 255 | xyz : torch.Tensor 256 | (B, N, 3) xyz coordinates of the features 257 | new_xyz : torch.Tensor 258 | (B, npoint, 3) centers of the ball query 259 | 260 | Returns 261 | ------- 262 | torch.Tensor 263 | (B, npoint, nsample) tensor with the indicies of the features that form the query balls 264 | """ 265 | output = _ext.ball_query(new_xyz, xyz, radius, nsample) 266 | 267 | ctx.mark_non_differentiable(output) 268 | 269 | return output 270 | 271 | @staticmethod 272 | def backward(ctx, grad_out): 273 | return () 274 | 275 | 276 | ball_query = BallQuery.apply 277 | 278 | 279 | class QueryAndGroup(nn.Module): 280 | r""" 281 | Groups with a ball query of radius 282 | 283 | Parameters 284 | --------- 285 | radius : float32 286 | Radius of ball 287 | nsample : int32 288 | Maximum number of features to gather in the ball 289 | """ 290 | 291 | def __init__(self, radius, nsample, use_xyz=True): 292 | # type: (QueryAndGroup, float, int, bool) -> None 293 | super(QueryAndGroup, self).__init__() 294 | self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz 295 | 296 | def forward(self, xyz, new_xyz, features=None): 297 | # type: (QueryAndGroup, torch.Tensor. torch.Tensor, torch.Tensor) -> Tuple[Torch.Tensor] 298 | r""" 299 | Parameters 300 | ---------- 301 | xyz : torch.Tensor 302 | xyz coordinates of the features (B, N, 3) 303 | new_xyz : torch.Tensor 304 | centriods (B, npoint, 3) 305 | features : torch.Tensor 306 | Descriptors of the features (B, C, N) 307 | 308 | Returns 309 | ------- 310 | new_features : torch.Tensor 311 | (B, 3 + C, npoint, nsample) tensor 312 | """ 313 | 314 | idx = ball_query(self.radius, self.nsample, xyz, new_xyz) 315 | xyz_trans = xyz.transpose(1, 2).contiguous() 316 | grouped_xyz = grouping_operation(xyz_trans, idx) # (B, 3, npoint, nsample) 317 | grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1) 318 | 319 | if features is not None: 320 | grouped_features = grouping_operation(features, idx) 321 | if self.use_xyz: 322 | new_features = torch.cat( 323 | [grouped_xyz, grouped_features], dim=1 324 | ) # (B, C + 3, npoint, nsample) 325 | else: 326 | new_features = grouped_features 327 | else: 328 | assert ( 329 | self.use_xyz 330 | ), "Cannot have not features and not use xyz as a feature!" 331 | new_features = grouped_xyz 332 | 333 | return new_features 334 | 335 | 336 | class GroupAll(nn.Module): 337 | r""" 338 | Groups all features 339 | 340 | Parameters 341 | --------- 342 | """ 343 | 344 | def __init__(self, use_xyz=True): 345 | # type: (GroupAll, bool) -> None 346 | super(GroupAll, self).__init__() 347 | self.use_xyz = use_xyz 348 | 349 | def forward(self, xyz, new_xyz, features=None): 350 | # type: (GroupAll, torch.Tensor, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor] 351 | r""" 352 | Parameters 353 | ---------- 354 | xyz : torch.Tensor 355 | xyz coordinates of the features (B, N, 3) 356 | new_xyz : torch.Tensor 357 | Ignored 358 | features : torch.Tensor 359 | Descriptors of the features (B, C, N) 360 | 361 | Returns 362 | ------- 363 | new_features : torch.Tensor 364 | (B, C + 3, 1, N) tensor 365 | """ 366 | 367 | grouped_xyz = xyz.transpose(1, 2).unsqueeze(2) 368 | if features is not None: 369 | grouped_features = features.unsqueeze(2) 370 | if self.use_xyz: 371 | new_features = torch.cat( 372 | [grouped_xyz, grouped_features], dim=1 373 | ) # (B, 3 + C, 1, N) 374 | else: 375 | new_features = grouped_features 376 | else: 377 | new_features = grouped_xyz 378 | 379 | return new_features 380 | -------------------------------------------------------------------------------- /pointnet2_ops_lib/setup.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import os 3 | import os.path as osp 4 | 5 | from setuptools import find_packages, setup 6 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 7 | 8 | this_dir = osp.dirname(osp.abspath(__file__)) 9 | _ext_src_root = osp.join("pointnet2_ops", "_ext-src") 10 | _ext_sources = glob.glob(osp.join(_ext_src_root, "src", "*.cpp")) + glob.glob( 11 | osp.join(_ext_src_root, "src", "*.cu") 12 | ) 13 | _ext_headers = glob.glob(osp.join(_ext_src_root, "include", "*")) 14 | 15 | requirements = ["torch>=1.4"] 16 | 17 | exec(open(osp.join("pointnet2_ops", "_version.py")).read()) 18 | 19 | os.environ["TORCH_CUDA_ARCH_LIST"] = "3.7+PTX;5.0;6.0;6.1;6.2;7.0;7.5" 20 | setup( 21 | name="pointnet2_ops", 22 | version=__version__, 23 | author="Erik Wijmans", 24 | packages=find_packages(), 25 | install_requires=requirements, 26 | ext_modules=[ 27 | CUDAExtension( 28 | name="pointnet2_ops._ext", 29 | sources=_ext_sources, 30 | extra_compile_args={ 31 | "cxx": ["-O3"], 32 | "nvcc": ["-O3", "-Xfatbin", "-compress-all"], 33 | }, 34 | include_dirs=[osp.join(this_dir, _ext_src_root, "include")], 35 | ) 36 | ], 37 | cmdclass={"build_ext": BuildExtension}, 38 | include_package_data=True, 39 | ) 40 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | cycler==0.10.0 2 | einops==0.3.0 3 | h5py==3.2.1 4 | matplotlib==3.4.2 5 | numpy>=1.21 6 | pyyaml==5.4.1 7 | scikit-learn==0.24.2 8 | scipy==1.6.3 9 | tqdm==4.61.1 10 | open3d 11 | easydict 12 | tensorboard 13 | pointnet2_ops_lib/. 14 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import argparse 4 | import datetime 5 | import numpy as np 6 | import random 7 | import shutil 8 | 9 | from glob import glob 10 | from tqdm import tqdm 11 | from torch.utils.data import DataLoader 12 | 13 | from core.builders import build_dataset, build_network, build_optimizer 14 | from utils.runtime_utils import cfg, cfg_from_yaml_file, validate 15 | from utils.vis_utils import visualize_numpy 16 | 17 | def parse_config(): 18 | parser = argparse.ArgumentParser(description='arg parser') 19 | parser.add_argument('--cfg_file', type=str, default=None, help='specify the config for training') 20 | parser.add_argument('--ckpt', type=str, default=None, help='checkpoint to start from') 21 | 22 | args = parser.parse_args() 23 | 24 | cfg_from_yaml_file(args.cfg_file, cfg) 25 | 26 | return args, cfg 27 | 28 | args, cfg = parse_config() 29 | exp_dir = ('/').join(args.ckpt.split('/')[:-2]) 30 | 31 | random_seed = cfg.RANDOM_SEED # Setup seed for reproducibility 32 | torch.manual_seed(random_seed) 33 | torch.cuda.manual_seed(random_seed) 34 | np.random.seed(random_seed) 35 | random.seed(random_seed) 36 | 37 | # Build Dataloader 38 | val_dataset = build_dataset(cfg, split='val') 39 | val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, drop_last=False) 40 | 41 | # Build Network and Optimizer 42 | net = build_network(cfg) 43 | state_dict = torch.load(args.ckpt) 44 | epoch = state_dict['epoch'] 45 | net.load_state_dict(state_dict['model_state_dict']) 46 | net = net.cuda() 47 | net.eval() 48 | 49 | print('Evaluating Epoch: ', epoch) 50 | val_dict = validate(net, val_dataloader, net.get_loss, 'cuda', is_segmentation = cfg.DATASET.IS_SEGMENTATION) 51 | 52 | if cfg.DATASET.IS_SEGMENTATION: 53 | miou = np.round(val_dict['miou'], 4) 54 | print('miou', miou) 55 | 56 | else: 57 | val_loss = np.round(val_dict['loss'], 4) 58 | val_acc = np.round(val_dict['acc'], 2) 59 | val_acc_avg = np.round(val_dict['acc_avg'], 2) 60 | 61 | print('val_loss', val_loss) 62 | print('val_acc', val_acc) 63 | print('val_acc_avg', val_acc_avg) 64 | 65 | 66 | if cfg.DATASET.IS_SEGMENTATION: 67 | with open(exp_dir + '/eval_best.txt', 'w') as f: 68 | f.write('Best Epoch: ' + str(epoch)) 69 | f.write('\nBest miou: ' + str(miou)) 70 | 71 | else: 72 | with open(exp_dir + '/eval_best.txt', 'w') as f: 73 | f.write('Best Epoch: ' + str(epoch)) 74 | f.write('\nBest Acc: ' + str(val_acc)) 75 | f.write('\nBest Mean Acc: ' + str(val_acc_avg)) 76 | f.write('\nBest Loss: ' + str(val_loss)) 77 | 78 | 79 | torch.save(state_dict['model_state_dict'], exp_dir + '/ckpt_model_only.pth') 80 | 81 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import argparse 4 | import datetime 5 | import numpy as np 6 | import random 7 | import shutil 8 | 9 | from glob import glob 10 | from tqdm import tqdm 11 | from torch.utils.data import DataLoader 12 | 13 | from core.builders import build_dataset, build_network, build_optimizer 14 | from utils.runtime_utils import cfg, cfg_from_yaml_file, validate 15 | from utils.vis_utils import visualize_numpy 16 | 17 | def parse_config(): 18 | parser = argparse.ArgumentParser(description='arg parser') 19 | parser.add_argument('--cfg_file', type=str, default=None, help='specify the config for training') 20 | parser.add_argument('--exp_name', type=str, default=None, help='specify experiment name for saving outputs') 21 | parser.add_argument('--ckpt', type=str, default=None, help='checkpoint to start from') 22 | parser.add_argument('--random_seed', type=int, default=0, help='random seed number') 23 | parser.add_argument('--val_steps', type=int, default=1, help='perform validation every n steps') 24 | parser.add_argument('--pretrained_ckpt', type = str, default = None, help='path to pretrained ckpt') 25 | args = parser.parse_args() 26 | 27 | cfg_from_yaml_file(args.cfg_file, cfg) 28 | exp_dir = cfg.ROOT_DIR / 'experiments' / cfg.DATASET.NAME / args.exp_name 29 | os.makedirs(exp_dir, exist_ok=True) 30 | shutil.copy2(args.cfg_file, exp_dir) 31 | 32 | return args, cfg 33 | 34 | args, cfg = parse_config() 35 | 36 | random_seed = cfg.RANDOM_SEED # Setup seed for reproducibility 37 | torch.manual_seed(random_seed) 38 | torch.cuda.manual_seed(random_seed) 39 | np.random.seed(random_seed) 40 | random.seed(random_seed) 41 | 42 | # Build Dataloader 43 | train_dataset = build_dataset(cfg, split = 'train') 44 | train_dataloader = DataLoader(train_dataset, batch_size=cfg.OPTIMIZER.BATCH_SIZE, shuffle=True, drop_last=True) 45 | 46 | val_dataset = build_dataset(cfg, split='val') 47 | val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, drop_last=False) 48 | 49 | # Build Network and Optimizer 50 | net = build_network(cfg) 51 | if args.pretrained_ckpt is not None: 52 | pretrained_state_dict = torch.load(args.pretrained_ckpt)['model_state_dict'] 53 | 54 | for k, v in net.state_dict().items(): 55 | if (v.shape != pretrained_state_dict[k].shape): 56 | del pretrained_state_dict[k] 57 | 58 | net.load_state_dict(pretrained_state_dict, strict = False) 59 | 60 | net = net.cuda() 61 | opt, scheduler = build_optimizer(cfg, net.parameters(), len(train_dataloader)) 62 | 63 | 64 | from torch.utils.tensorboard import SummaryWriter 65 | ckpt_dir = cfg.ROOT_DIR / 'experiments' / cfg.DATASET.NAME / args.exp_name / 'ckpt' 66 | tensorboard_dir = cfg.ROOT_DIR / 'experiments' / cfg.DATASET.NAME / args.exp_name / 'tensorboard' 67 | 68 | os.makedirs(ckpt_dir, exist_ok=True) 69 | os.makedirs(tensorboard_dir, exist_ok=True) 70 | 71 | writer = SummaryWriter(tensorboard_dir) 72 | 73 | min_loss = 1e20 74 | max_acc = 0 75 | 76 | steps_cnt = 0 77 | epoch_cnt = 0 78 | 79 | 80 | for epoch in tqdm(range(1, cfg.OPTIMIZER.MAX_EPOCH + 1)): 81 | opt.zero_grad() 82 | net.zero_grad() 83 | net.train() 84 | loss = 0 85 | for data_dic in tqdm(train_dataloader): 86 | 87 | data_dic = net(data_dic) 88 | loss, loss_dict = net.get_loss(data_dic, smoothing = True, is_segmentation = cfg.DATASET.IS_SEGMENTATION) 89 | loss = loss 90 | loss.backward() 91 | steps_cnt += 1 92 | 93 | # if (steps_cnt)%(cfg.OPTIMIZER.GRAD_ACCUMULATION) == 0: 94 | torch.nn.utils.clip_grad_norm_(net.parameters(), cfg.OPTIMIZER.GRAD_CLIP) 95 | opt.step() 96 | opt.zero_grad() 97 | lr = scheduler.get_last_lr()[0] 98 | scheduler.step() 99 | writer.add_scalar('steps/loss', loss, steps_cnt) 100 | writer.add_scalar('steps/lr', lr, steps_cnt) 101 | 102 | for k,v in loss_dict.items(): 103 | writer.add_scalar('steps/loss_' + k, v, steps_cnt) 104 | 105 | if (epoch % args.val_steps) == 0: 106 | val_dict = validate(net, val_dataloader, net.get_loss, 'cuda', is_segmentation = cfg.DATASET.IS_SEGMENTATION) 107 | 108 | print('='*20, 'Epoch ' + str(epoch+1), '='*20) 109 | 110 | if cfg.DATASET.IS_SEGMENTATION: 111 | writer.add_scalar('epochs/val_miou', val_dict['miou'], epoch_cnt) 112 | print('Val mIoU: ', val_dict['miou']) 113 | 114 | else: 115 | writer.add_scalar('epochs/val_loss', val_dict['loss'], epoch_cnt) 116 | writer.add_scalar('epochs/val_acc', val_dict['acc'], epoch_cnt) 117 | writer.add_scalar('epochs/val_acc_avg', val_dict['acc_avg'], epoch_cnt) 118 | print('Val Loss: ', val_dict['loss'], 'Val Accuracy: ', val_dict['acc'], 'Val Avg Accuracy: ', val_dict['acc_avg']) 119 | 120 | for k,v in val_dict['loss_dic'].items(): 121 | writer.add_scalar('epochs/val_loss_'+ k, v, epoch_cnt) 122 | 123 | epoch_cnt += 1 124 | 125 | 126 | if cfg.DATASET.IS_SEGMENTATION: 127 | if val_dict['miou'] > max_acc: 128 | torch.save({ 129 | 'epoch': epoch, 130 | 'model_state_dict': net.state_dict(), 131 | 'optimizer_state_dict': opt.state_dict(), 132 | }, ckpt_dir / 'ckpt-best.pth') 133 | 134 | max_acc = val_dict['miou'] 135 | else: 136 | 137 | if val_dict['acc'] > max_acc: 138 | torch.save({ 139 | 'epoch': epoch, 140 | 'model_state_dict': net.state_dict(), 141 | 'optimizer_state_dict': opt.state_dict(), 142 | 'loss': val_dict['loss'], 143 | }, ckpt_dir / 'ckpt-best.pth') 144 | 145 | max_acc = val_dict['acc'] 146 | 147 | torch.save({ 148 | 'epoch': epoch, 149 | 'model_state_dict': net.state_dict(), 150 | 'optimizer_state_dict': opt.state_dict(), 151 | }, ckpt_dir / 'ckpt-last.pth') 152 | -------------------------------------------------------------------------------- /utils/runtime_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import yaml 3 | import torch 4 | import datetime 5 | 6 | import sklearn.metrics as metrics 7 | import numpy as np 8 | from tqdm import tqdm 9 | from easydict import EasyDict 10 | from pathlib import Path 11 | 12 | import torch.nn.functional as F 13 | 14 | def merge_new_config(config, new_config): 15 | if '_BASE_CONFIG_' in new_config: 16 | with open(new_config['_BASE_CONFIG_'], 'r') as f: 17 | try: 18 | yaml_config = yaml.safe_load(f, Loader=yaml.FullLoader) 19 | except: 20 | yaml_config = yaml.safe_load(f) 21 | config.update(EasyDict(yaml_config)) 22 | 23 | for key, val in new_config.items(): 24 | if not isinstance(val, dict): 25 | config[key] = val 26 | continue 27 | if key not in config: 28 | config[key] = EasyDict() 29 | merge_new_config(config[key], val) 30 | 31 | return config 32 | 33 | def cfg_from_yaml_file(cfg_file, config): 34 | with open(cfg_file, 'r') as f: 35 | try: 36 | new_config = yaml.safe_load(f, Loader=yaml.FullLoader) 37 | except: 38 | new_config = yaml.safe_load(f) 39 | 40 | merge_new_config(config=config, new_config=new_config) 41 | 42 | return config 43 | 44 | 45 | cfg = EasyDict() 46 | cfg.ROOT_DIR = (Path(__file__).resolve().parent / '../').resolve() 47 | cfg.LOCAL_RANK = 0 48 | 49 | class AverageMeter(object): 50 | """Computes and stores the average and current value""" 51 | def __init__(self): 52 | self.reset() 53 | 54 | def reset(self): 55 | self.val = 0 56 | self.avg = 0 57 | self.sum = 0 58 | self.count = 0 59 | 60 | def update(self, val, n=1): 61 | self.val = val 62 | self.sum += val * n 63 | self.count += n 64 | self.avg = self.sum / self.count 65 | 66 | def get_n_params(model): 67 | pp=0 68 | for p in list(model.parameters()): 69 | nn=1 70 | for s in list(p.size()): 71 | nn = nn*s 72 | pp += nn 73 | return pp 74 | 75 | def validate(net, testloader, criterion, device, is_segmentation = False): 76 | net.eval() 77 | num_params = get_n_params(net) 78 | 79 | test_loss = 0 80 | correct = 0 81 | total = 0 82 | test_true = [] 83 | test_pred = [] 84 | time_cost = [] 85 | mious = [] 86 | 87 | accuracy = [] 88 | shape_ious = 0.0 89 | count = 0.0 90 | with torch.no_grad(): 91 | for batch_idx, data_dic in enumerate(tqdm(testloader)): 92 | start_time = datetime.datetime.now() 93 | data_dic = net(data_dic) 94 | time_cost.append(float((datetime.datetime.now() - start_time).total_seconds())) 95 | 96 | 97 | if is_segmentation: 98 | miou = net.compute_overall_iou(data_dic['pred_score_logits'], data_dic['seg_id'], num_classes = 50) 99 | # total iou of current batch in each process: 100 | batch_ious = data_dic['pred_score_logits'].new_tensor([np.sum(miou)], dtype=torch.float64) # same device with seg_pred 101 | 102 | # prepare seg_pred and target for later calculating loss and acc: 103 | seg_pred = data_dic['pred_score_logits'].contiguous().view(-1, 50) # ShapeNetPart has 50 classes 104 | 105 | target = data_dic['seg_id'].view(-1, 1)[:, 0] 106 | # Loss 107 | 108 | pred_choice = seg_pred.data.max(1)[1] # b*n 109 | correct = pred_choice.eq(target.data).sum() # torch.int64: total number of correct-predict pts 110 | 111 | shape_ious += batch_ious.item() # count the sum of ious in each iteration 112 | count += data_dic['pred_score_logits'].shape[0] # count the total number of samples in each iteration 113 | accuracy.append(correct.item() / (data_dic['pred_score_logits'].shape[0] * data_dic['points'].shape[1])) # append the accuracy of each iteration 114 | mious.append(miou) 115 | 116 | else: 117 | preds = data_dic['pred_score_logits'].max(dim=1)[1] 118 | test_true.append(data_dic['cls_id'].cpu().numpy()) 119 | test_pred.append(preds.detach().cpu().numpy()) 120 | total += data_dic['cls_id'].size(0) 121 | correct += preds.eq(data_dic['cls_id']).sum().item() 122 | loss, loss_dict = criterion(data_dic) 123 | test_loss += loss.item() 124 | 125 | 126 | 127 | total_samples = len(time_cost) 128 | time_cost = np.mean(time_cost[total_samples//5:total_samples*4//5]) 129 | 130 | if is_segmentation: 131 | overall_miou = np.mean(mious) 132 | overall_acc = np.mean(accuracy) 133 | overall_ins_acc = shape_ious * 1.0 / count 134 | return { 135 | "miou": float("%.3f" % (100. * overall_miou)), 136 | "time": time_cost, 137 | "num_params": num_params, 138 | 'peak_memory': torch.cuda.memory_stats('cuda')['allocated_bytes.all.peak'] 139 | } 140 | else: 141 | test_true = np.concatenate(test_true) 142 | test_pred = np.concatenate(test_pred) 143 | 144 | 145 | return { 146 | "loss": float("%.3f" % (test_loss / (batch_idx + 1))), 147 | "loss_dic": loss_dict, 148 | "acc": float("%.3f" % (100. * metrics.accuracy_score(test_true, test_pred))), 149 | "acc_avg": float("%.3f" % (100. * metrics.balanced_accuracy_score(test_true, test_pred))), 150 | "time": time_cost, 151 | "num_params": num_params, 152 | 'peak_memory': torch.cuda.memory_stats('cuda')['allocated_bytes.all.peak'] 153 | } 154 | 155 | def validate_voting(net, testloader, device = 'cuda', num_repeat = 300, num_vote = 10): 156 | net.eval() 157 | best_acc = 0 158 | best_mean_acc = 0 159 | pointscale = PointcloudScale(scale_low=0.85, scale_high=1.15) 160 | 161 | for i in tqdm(range(num_repeat)): 162 | test_true = [] 163 | test_pred = [] 164 | 165 | for batch_idx, data_dic in enumerate(tqdm(testloader)): 166 | pred = 0 167 | for v in range(num_vote): 168 | new_data = data_dic 169 | if v > 0: 170 | new_data['points'] = pointscale(new_data['points']) 171 | with torch.no_grad(): 172 | pred += F.softmax(net(data_dic)['pred_score_logits'], dim=1) # sum 10 preds 173 | pred /= num_vote # avg the preds! 174 | pred_choice = pred.max(dim=1)[1] 175 | test_true.append(data_dic['cls_id'].cpu().numpy()) 176 | test_pred.append(pred_choice.detach().cpu().numpy()) 177 | test_true = np.concatenate(test_true) 178 | test_pred = np.concatenate(test_pred) 179 | test_acc = 100. * metrics.accuracy_score(test_true, test_pred) 180 | test_mean_acc = 100. * metrics.balanced_accuracy_score(test_true, test_pred) 181 | if test_acc > best_acc: 182 | best_acc = test_acc 183 | if test_mean_acc > best_mean_acc: 184 | best_mean_acc = test_mean_acc 185 | outstr = 'Voting %d, test acc: %.3f, test mean acc: %.3f, [current best(all_acc: %.3f mean_acc: %.3f)]' % \ 186 | (i, test_acc, test_mean_acc, best_acc, best_mean_acc) 187 | print(outstr) 188 | best_acc = best_acc*100 189 | best_mean_acc = best_mean_acc*100 190 | print('Final Test Acc: ', best_acc) 191 | print('Final Test Avg Acc: ', best_mean_acc) 192 | 193 | val_dict = { 194 | 'acc': best_acc, 195 | 'acc_avg': best_mean_acc 196 | } 197 | 198 | return val_dict 199 | 200 | class PointcloudScale(object): # input random scaling 201 | def __init__(self, scale_low=2. / 3., scale_high=3. / 2.): 202 | self.scale_low = scale_low 203 | self.scale_high = scale_high 204 | 205 | def __call__(self, pc): 206 | bsize = pc.size()[0] 207 | for i in range(bsize): 208 | xyz1 = np.random.uniform(low=self.scale_low, high=self.scale_high, size=[3]) 209 | pc[i, :, 0:3] = torch.mul(pc[i, :, 0:3], torch.from_numpy(xyz1).float().cuda()) 210 | 211 | return pc 212 | -------------------------------------------------------------------------------- /utils/vis_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import open3d as o3d 3 | 4 | import torch 5 | from tqdm import tqdm 6 | import datetime 7 | 8 | def visualize_numpy(pc_numpy, colors = None): 9 | point_cloud = o3d.geometry.PointCloud() 10 | point_cloud.points = o3d.utility.Vector3dVector(pc_numpy) 11 | try: 12 | point_cloud.colors = o3d.utility.Vector3dVector(colors) 13 | except: 14 | pass 15 | 16 | vis = o3d.visualization.Visualizer() 17 | vis.create_window() 18 | vis.add_geometry(point_cloud) 19 | ctr = vis.get_view_control() 20 | ctr.set_up((1, 0, 0)) 21 | ctr.set_front((0, 1, 0)) 22 | 23 | vis.run() 24 | 25 | # o3d.visualization.draw_geometries([point_cloud]) 26 | 27 | def visualize_part(net, testloader): 28 | color_choices = [(1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), #3 29 | (1, 0, 0), (0, 1, 0), #5 30 | (1, 0, 0), (0, 1, 0), #7 31 | (1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), #11 32 | (1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), #15 33 | (1, 0, 0), (0, 1, 0), (0, 0, 1), # 18 34 | (1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), #21 35 | (1, 0, 0), (0, 1, 0), #23 36 | (1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), #27 37 | (1, 0, 0), (0, 1, 0), #29 38 | (1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), (0, 1, 1), (1, 0, 1), # 35 39 | (1, 0, 0), (0, 1, 0), #37 40 | (1, 0, 0), (0, 1, 0), (0, 0, 1), #40 41 | (1, 0, 0), (0, 1, 0), (0, 0, 1), #43 42 | (1, 0, 0), (0, 1, 0), (0, 0, 1), #46 43 | (1, 0, 0), (0, 1, 0), (0, 0, 1), #49 44 | ] 45 | 46 | net.eval() 47 | 48 | with torch.no_grad(): 49 | for batch_idx, data_dic in enumerate(tqdm(testloader)): 50 | if ((batch_idx % 10 == 0) and (batch_idx > 1260)): 51 | data_dic = net(data_dic) 52 | points = data_dic['points'].squeeze(0).cpu().numpy() 53 | label = data_dic['seg_id'].squeeze(0).cpu().numpy() 54 | pred = torch.argmax(data_dic['pred_score_logits'], dim = -1).squeeze(0).cpu().numpy() 55 | 56 | color_list = np.zeros_like(points) 57 | for i, pred_id in enumerate(pred): 58 | color_list[i] = color_choices[int(pred_id)] 59 | visualize_numpy(points, colors=color_list) 60 | 61 | color_list = np.zeros_like(points) 62 | for i, label_id in enumerate(label): 63 | color_list[i] = color_choices[int(label_id)] 64 | visualize_numpy(points, colors=color_list) 65 | -------------------------------------------------------------------------------- /visualize_part.py: -------------------------------------------------------------------------------- 1 | import os 2 | from idna import valid_contextj 3 | import torch 4 | import argparse 5 | import datetime 6 | import numpy as np 7 | import random 8 | import shutil 9 | 10 | from glob import glob 11 | from tqdm import tqdm 12 | from torch.utils.data import DataLoader 13 | 14 | from core.builders import build_dataset, build_network, build_optimizer 15 | from utils.runtime_utils import cfg, cfg_from_yaml_file, validate 16 | from utils.vis_utils import visualize_numpy, visualize_part 17 | 18 | def parse_config(): 19 | parser = argparse.ArgumentParser(description='arg parser') 20 | parser.add_argument('--cfg_file', type=str, default=None, help='specify the config for training') 21 | parser.add_argument('--ckpt', type=str, default=None, help='checkpoint to start from') 22 | 23 | args = parser.parse_args() 24 | 25 | cfg_from_yaml_file(args.cfg_file, cfg) 26 | 27 | return args, cfg 28 | 29 | args, cfg = parse_config() 30 | exp_dir = ('/').join(args.ckpt.split('/')[:-2]) 31 | 32 | 33 | # Build Dataloader 34 | val_dataset = build_dataset(cfg, split='val') 35 | val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, drop_last=False) 36 | 37 | # Build Network 38 | net = build_network(cfg) 39 | state_dict = torch.load(args.ckpt) 40 | epoch = state_dict['epoch'] 41 | net.load_state_dict(state_dict['model_state_dict']) 42 | net = net.cuda() 43 | net.eval() 44 | 45 | print('Evaluating Epoch: ', epoch) 46 | visualize_part(net, val_dataloader) --------------------------------------------------------------------------------