├── .gitignore ├── LICENSE ├── README.md ├── completion ├── .gitignore ├── LICENSE ├── README.md ├── __init__.py ├── cfgs │ ├── __init__.py │ ├── ecg.yaml │ ├── pcn.yaml │ └── vrcnet.yaml ├── data │ └── .gitignore ├── dataset.py ├── images │ ├── complete_pcds.gif │ ├── intro.png │ ├── modules.png │ ├── mvp.png │ ├── overview.png │ └── partial_pcds.gif ├── model_utils.py ├── models │ ├── __init__.py │ ├── ecg.py │ ├── pcn.py │ └── vrcnet.py ├── requirements.txt ├── run_test.sh ├── run_train.sh ├── test.py ├── train.py ├── train_utils.py └── vis_utils.py ├── images ├── complete_pcds.gif ├── logo.png └── partial_pcds.gif ├── registration ├── .gitignore ├── README.md ├── cfgs │ ├── dcp.yaml │ ├── deepgmr.yaml │ └── idam.yaml ├── data │ └── .gitignore ├── dataset.py ├── images │ ├── registration.png │ ├── registration_black.png │ └── registration_raw.png ├── model_utils.py ├── models │ ├── dcp.py │ ├── deepgmr.py │ └── idam.py ├── run_test.sh ├── run_train.sh ├── test.py ├── train.py ├── train_utils.py └── visu_utils.py ├── setup.sh └── utils ├── __init__.py ├── metrics ├── CD │ ├── .gitignore │ ├── LICENSE │ ├── README.md │ ├── __init__.py │ ├── chamfer3D │ │ ├── .gitignore │ │ ├── chamfer3D.cu │ │ ├── chamfer_cuda.cpp │ │ ├── dist_chamfer_3D.py │ │ ├── run_build.sh │ │ └── setup.py │ ├── chamfer_python.py │ ├── fscore.py │ └── unit_test.py ├── EMD │ ├── .gitignore │ ├── CDEMD.png │ ├── README.md │ ├── __init__.py │ ├── clean.sh │ ├── emd.cpp │ ├── emd_cuda.cu │ ├── emd_module.py │ ├── run_build.sh │ ├── run_compile.sh │ └── setup.py └── __init__.py └── mm3d_pn2 ├── .gitignore ├── __init__.py ├── ops ├── __init__.py ├── ball_query │ ├── __init__.py │ ├── ball_query.py │ └── src │ │ ├── ball_query.cpp │ │ └── ball_query_cuda.cu ├── furthest_point_sample │ ├── __init__.py │ ├── furthest_point_sample.py │ ├── points_sampler.py │ ├── src │ │ ├── furthest_point_sample.cpp │ │ └── furthest_point_sample_cuda.cu │ └── utils.py ├── gather_points │ ├── __init__.py │ ├── gather_points.py │ └── src │ │ ├── gather_points.cpp │ │ └── gather_points_cuda.cu ├── group_points │ ├── __init__.py │ ├── group_points.py │ └── src │ │ ├── group_points.cpp │ │ └── group_points_cuda.cu ├── interpolate │ ├── __init__.py │ ├── src │ │ ├── interpolate.cpp │ │ ├── three_interpolate_cuda.cu │ │ └── three_nn_cuda.cu │ ├── three_interpolate.py │ └── three_nn.py ├── iou3d │ ├── __init__.py │ ├── iou3d_utils.py │ └── src │ │ ├── iou3d.cpp │ │ └── iou3d_kernel.cu ├── knn │ ├── __init__.py │ ├── knn.py │ └── src │ │ ├── knn.cpp │ │ └── knn_cuda.cu ├── norm.py ├── paconv │ ├── __init__.py │ ├── assign_score.py │ ├── paconv.py │ ├── src │ │ ├── assign_score_withk.cpp │ │ └── assign_score_withk_cuda.cu │ └── utils.py ├── pointnet_modules │ ├── __init__.py │ ├── builder.py │ ├── point_fp_module.py │ └── point_sa_module.py ├── roiaware_pool3d │ ├── __init__.py │ ├── points_in_boxes.py │ ├── roiaware_pool3d.py │ └── src │ │ ├── points_in_boxes_cpu.cpp │ │ ├── points_in_boxes_cuda.cu │ │ ├── roiaware_pool3d.cpp │ │ └── roiaware_pool3d_kernel.cu ├── sparse_block.py ├── spconv │ ├── __init__.py │ ├── conv.py │ ├── functional.py │ ├── include │ │ ├── paramsgrid.h │ │ ├── prettyprint.h │ │ ├── pybind11_utils.h │ │ ├── spconv │ │ │ ├── fused_spconv_ops.h │ │ │ ├── geometry.h │ │ │ ├── indice.cu.h │ │ │ ├── indice.h │ │ │ ├── maxpool.h │ │ │ ├── mp_helper.h │ │ │ ├── point2voxel.h │ │ │ ├── pool_ops.h │ │ │ ├── reordering.cu.h │ │ │ ├── reordering.h │ │ │ └── spconv_ops.h │ │ ├── tensorview │ │ │ ├── helper_kernel.cu.h │ │ │ ├── helper_launch.h │ │ │ └── tensorview.h │ │ ├── torch_utils.h │ │ └── utility │ │ │ └── timer.h │ ├── modules.py │ ├── ops.py │ ├── pool.py │ ├── src │ │ ├── all.cc │ │ ├── indice.cc │ │ ├── indice_cuda.cu │ │ ├── maxpool.cc │ │ ├── maxpool_cuda.cu │ │ ├── reordering.cc │ │ └── reordering_cuda.cu │ ├── structure.py │ └── test_utils.py └── voxel │ ├── __init__.py │ ├── scatter_points.py │ ├── src │ ├── scatter_points_cpu.cpp │ ├── scatter_points_cuda.cu │ ├── voxelization.cpp │ ├── voxelization.h │ ├── voxelization_cpu.cpp │ └── voxelization_cuda.cu │ └── voxelize.py ├── requirements.txt ├── requirements ├── build.txt ├── docs.txt ├── mminstall.txt ├── optional.txt ├── readthedocs.txt ├── runtime.txt └── tests.txt ├── run_build.sh ├── setup.cfg ├── setup.py ├── setup.sh └── version.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | *.so 3 | *.pyc 4 | -------------------------------------------------------------------------------- /completion/.gitignore: -------------------------------------------------------------------------------- 1 | /.idea/* 2 | /.DS_Store 3 | /log/* 4 | */**/__pycache__ 5 | /__pycache__/* 6 | *.h5 7 | *.backup -------------------------------------------------------------------------------- /completion/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 PL 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /completion/README.md: -------------------------------------------------------------------------------- 1 | # *MVP Benchmark:* Point Cloud Completioin 2 |

3 | 4 |

5 | 6 | 7 | We include the following methods for point cloud completion: 8 | 9 | [1] [PCN](https://github.com/wentaoyuan/pcn);   [2] [ECG](https://github.com/paul007pl/ECG);   [3] [VRCNet](https://github.com/paul007pl/VRCNet) 10 | 11 | 12 | ### MVP Completion Dataset 13 | 17 | Download the MVP completion dataset [Google Drive](https://drive.google.com/drive/folders/1XxZ4M_dOB3_OG1J6PnpNvrGTie5X9Vk_) or [百度网盘](https://pan.baidu.com/s/18pli79KSGGsWQ8FPiSW9qg)  (code: p364) to the folder "data". 18 | 19 | The data structure will be: 20 | ``` 21 | data 22 | ├── MVP_Train_CP.h5 23 | | ├── incomplete_pcds (62400, 2048, 3) 24 | | ├── complete_pcds (2400, 2048, 3) 25 | | └── labels (62400,) 26 | ├── MVP_Test_CP.h5 27 | | ├── incomplete_pcds (41600, 2048, 3) 28 | | ├── complete_pcds (1600, 2048, 3) 29 | | └── labels (41600,) 30 | └── MVP_ExtraTest_Shuffled_CP.h5 31 | ├── incomplete_pcds (59800, 2048, 3) 32 | └── labels (59800,) 33 | ``` 34 | **Attention**: `MVP_ExtraTest_Shuffled_CP.h5` is only available during the competition period. `MVP_Test_CP.h5` is the standard test set if you use MVP dataset in your papers. 35 | 36 | | id | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 37 | |:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:| 38 | | category | airplane | cabinet | car | chair | lamp | sofa | table | watercraft | bed | bench | bookshelf | bus | guitar | motorbike | pistol | skateboard | 39 | | \#train | 5200 | 5200 | 5200 | 5200 | 5200 | 5200 | 5200 | 5200 | 2600 | 2600 | 2600 | 2600 | 2600 | 2600 | 2600 | 2600 | 40 | | \#test | 3900 | 3900 | 3900 | 3900 | 3900 | 3900 | 3900 | 3900 | 1300 | 1300 | 1300 | 1300 | 1300 | 1300 | 1300 | 1300 | 41 | 42 | 43 | 44 | 51 | 52 | 55 | 56 | 57 | ### Usage 58 | + To train a model: run `python train.py -c ./cfgs/*.yaml`, e.g. `python train.py -c ./cfgs/pcn.yaml` 59 | + To test a model: run `python test.py -c ./cfgs/*.yaml`, e.g. `python test.py -c ./cfgs/pcn.yaml` 60 | + Config for each algorithm can be found in `cfgs/`. 61 | + `run_train.sh` and `run_test.sh` are provided for SLURM users. 62 | 63 | 64 | ## Citation 65 | If you find our code useful, please cite our paper: 66 | ```bibtex 67 | @inproceedings{pan2021variational, 68 | title={Variational Relational Point Completion Network}, 69 | author={Pan, Liang and Chen, Xinyi and Cai, Zhongang and Zhang, Junzhe and Zhao, Haiyu and Yi, Shuai and Liu, Ziwei}, 70 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, 71 | pages={8524--8533}, 72 | year={2021} 73 | } 74 | @article{pan2021robust, 75 | title={Robust Partial-to-Partial Point Cloud Registration in a Full Range}, 76 | author={Pan, Liang and Cai, Zhongang and Liu, Ziwei}, 77 | journal={arXiv preprint arXiv:2111.15606}, 78 | year={2021} 79 | } 80 | @article{pan2021mvp, 81 | title={Multi-View Partial (MVP) Point Cloud Challenge 2021 on Completion and Registration: Methods and Results}, 82 | author={Pan, Liang and Wu, Tong and Cai, Zhongang and Liu, Ziwei and Yu, Xumin and Rao, Yongming and Lu, Jiwen and Zhou, Jie and Xu, Mingye and Luo, Xiaoyuan and Fu, Kexue, and Gao, Peng, and Wang, Manning, and Wang, Yali, and Qiao, Yu, and Zhou, Junsheng, and Wen, Xin, and Xiang, Peng, and Liu, Yu-Shen, and Han, Zhizhong, and Yan, Yuanjie, and An, Junyi, and Zhu, Lifa, and Lin, Changwei, and Liu, Dongrui, and Li, Xin, and G ́omez-Fern ́andez, Francisco, and Wang, Qinlong, and Yang, Yang}, 83 | journal={arXiv preprint arXiv:2112.12053}, 84 | year={2021} 85 | } 86 | ``` 87 | 88 | 89 | ## License 90 | Our code is released under Apache-2.0 License. 91 | 92 | 93 | ## Acknowledgement 94 | We include the following algorithms: 95 | [1] [PCN](https://github.com/wentaoyuan/pcn) 96 | [2] [ECG](https://github.com/paul007pl/ECG) 97 | [3] [VRCNet](https://github.com/paul007pl/VRCNet) 98 | 99 | -------------------------------------------------------------------------------- /completion/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/completion/__init__.py -------------------------------------------------------------------------------- /completion/cfgs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/completion/cfgs/__init__.py -------------------------------------------------------------------------------- /completion/cfgs/ecg.yaml: -------------------------------------------------------------------------------- 1 | batch_size: 32 2 | workers: 0 3 | nepoch: 100 4 | model_name: ecg 5 | load_model: null 6 | start_epoch: 0 7 | num_points: 2048 8 | work_dir: log/ 9 | flag: debug 10 | loss: cd 11 | manual_seed: null 12 | use_mean_feature: False 13 | step_interval_to_print: 500 14 | epoch_interval_to_save: 1 15 | epoch_interval_to_val: 1 16 | varying_constant: 0.01, 0.1, 0.5, 1 17 | varying_constant_epochs: 5, 15, 30 18 | 19 | lr: 0.0001 20 | lr_decay: True 21 | lr_decay_interval: 40 22 | lr_decay_rate: 0.7 23 | lr_step_decay_epochs: null 24 | lr_step_decay_rates: null 25 | lr_clip: 1.e-6 26 | optimizer: Adam 27 | weight_decay: 0 28 | betas: 0.9, 0.999 29 | 30 | # test 31 | save_vis: True 32 | eval_emd: False -------------------------------------------------------------------------------- /completion/cfgs/pcn.yaml: -------------------------------------------------------------------------------- 1 | batch_size: 32 2 | workers: 0 3 | nepoch: 100 4 | model_name: pcn 5 | load_model: null 6 | start_epoch: 0 7 | num_points: 2048 8 | work_dir: log/ 9 | flag: debug 10 | loss: cd 11 | manual_seed: null 12 | use_mean_feature: False 13 | step_interval_to_print: 500 14 | epoch_interval_to_save: 1 15 | epoch_interval_to_val: 1 16 | varying_constant: 0.01, 0.1, 0.5, 1 17 | varying_constant_epochs: 5, 15, 30 18 | 19 | lr: 0.0001 20 | lr_decay: True 21 | lr_decay_interval: 40 22 | lr_decay_rate: 0.7 23 | lr_step_decay_epochs: null 24 | lr_step_decay_rates: null 25 | lr_clip: 1.e-6 26 | optimizer: Adam 27 | weight_decay: 0 28 | betas: 0.9, 0.999 29 | 30 | # test 31 | save_vis: True 32 | eval_emd: False -------------------------------------------------------------------------------- /completion/cfgs/vrcnet.yaml: -------------------------------------------------------------------------------- 1 | batch_size: 32 2 | workers: 0 3 | nepoch: 100 4 | model_name: vrcnet 5 | load_model: ./log/vrcnet_cd_debug_2021-06-18T16:36:20/best_cd_t_network.pth 6 | start_epoch: 0 7 | num_points: 2048 8 | work_dir: log/ 9 | flag: debug 10 | loss: cd 11 | manual_seed: null 12 | use_mean_feature: False 13 | step_interval_to_print: 500 14 | epoch_interval_to_save: 1 15 | epoch_interval_to_val: 1 16 | varying_constant: 0.01, 0.1, 0.5, 1 17 | varying_constant_epochs: 5, 15, 30 18 | 19 | lr: 0.0001 20 | lr_decay: True 21 | lr_decay_interval: 40 22 | lr_decay_rate: 0.7 23 | lr_step_decay_epochs: null 24 | lr_step_decay_rates: null 25 | lr_clip: 1.e-6 26 | optimizer: Adam 27 | weight_decay: 0 28 | betas: 0.9, 0.999 29 | 30 | layers: 1, 1, 1, 1 31 | distribution_loss: KLD 32 | knn_list: "16" 33 | pk: 10 34 | local_folding: True 35 | points_label: True 36 | num_coarse_raw: 1024 37 | num_fps: 2048 38 | num_coarse: 2048 39 | 40 | # test 41 | save_vis: False 42 | eval_emd: False -------------------------------------------------------------------------------- /completion/data/.gitignore: -------------------------------------------------------------------------------- 1 | *.h5 -------------------------------------------------------------------------------- /completion/dataset.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.utils.data as data 4 | import h5py 5 | import os 6 | 7 | 8 | class MVP_CP(data.Dataset): 9 | def __init__(self, prefix="train"): 10 | if prefix=="train": 11 | self.file_path = './data/MVP_Train_CP.h5' 12 | elif prefix=="val": 13 | self.file_path = './data/MVP_Test_CP.h5' 14 | elif prefix=="test": 15 | self.file_path = './data/MVP_ExtraTest_Shuffled_CP.h5' 16 | else: 17 | raise ValueError("ValueError prefix should be [train/val/test] ") 18 | 19 | self.prefix = prefix 20 | 21 | input_file = h5py.File(self.file_path, 'r') 22 | self.input_data = np.array(input_file['incomplete_pcds'][()]) 23 | 24 | print(self.input_data.shape) 25 | 26 | if prefix is not "test": 27 | self.gt_data = np.array(input_file['complete_pcds'][()]) 28 | self.labels = np.array(input_file['labels'][()]) 29 | print(self.gt_data.shape, self.labels.shape) 30 | 31 | 32 | input_file.close() 33 | self.len = self.input_data.shape[0] 34 | 35 | def __len__(self): 36 | return self.len 37 | 38 | def __getitem__(self, index): 39 | partial = torch.from_numpy((self.input_data[index])) 40 | 41 | if self.prefix is not "test": 42 | complete = torch.from_numpy((self.gt_data[index // 26])) 43 | label = (self.labels[index]) 44 | return label, partial, complete 45 | else: 46 | return partial 47 | -------------------------------------------------------------------------------- /completion/images/complete_pcds.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/completion/images/complete_pcds.gif -------------------------------------------------------------------------------- /completion/images/intro.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/completion/images/intro.png -------------------------------------------------------------------------------- /completion/images/modules.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/completion/images/modules.png -------------------------------------------------------------------------------- /completion/images/mvp.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/completion/images/mvp.png -------------------------------------------------------------------------------- /completion/images/overview.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/completion/images/overview.png -------------------------------------------------------------------------------- /completion/images/partial_pcds.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/completion/images/partial_pcds.gif -------------------------------------------------------------------------------- /completion/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/completion/models/__init__.py -------------------------------------------------------------------------------- /completion/models/pcn.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.parallel 5 | import torch.utils.data 6 | import torch.nn.functional as F 7 | import math 8 | 9 | # from utils.model_utils import gen_grid_up, calc_emd, calc_cd 10 | from model_utils import gen_grid_up, calc_emd, calc_cd 11 | 12 | 13 | class PCN_encoder(nn.Module): 14 | def __init__(self, output_size=1024): 15 | super(PCN_encoder, self).__init__() 16 | self.conv1 = nn.Conv1d(3, 128, 1) 17 | self.conv2 = nn.Conv1d(128, 256, 1) 18 | self.conv3 = nn.Conv1d(512, 512, 1) 19 | self.conv4 = nn.Conv1d(512, output_size, 1) 20 | 21 | def forward(self, x): 22 | batch_size, _, num_points = x.size() 23 | x = F.relu(self.conv1(x)) 24 | x = self.conv2(x) 25 | global_feature, _ = torch.max(x, 2) 26 | x = torch.cat((x, global_feature.view(batch_size, -1, 1).repeat(1, 1, num_points).contiguous()), 1) 27 | x = F.relu(self.conv3(x)) 28 | x = self.conv4(x) 29 | global_feature, _ = torch.max(x, 2) 30 | return global_feature.view(batch_size, -1) 31 | 32 | 33 | class PCN_decoder(nn.Module): 34 | def __init__(self, num_coarse, num_fine, scale, cat_feature_num): 35 | super(PCN_decoder, self).__init__() 36 | self.num_coarse = num_coarse 37 | self.num_fine = num_fine 38 | self.fc1 = nn.Linear(1024, 1024) 39 | self.fc2 = nn.Linear(1024, 1024) 40 | self.fc3 = nn.Linear(1024, num_coarse * 3) 41 | 42 | self.scale = scale 43 | self.grid = gen_grid_up(2 ** (int(math.log2(scale))), 0.05).cuda().contiguous() 44 | self.conv1 = nn.Conv1d(cat_feature_num, 512, 1) 45 | self.conv2 = nn.Conv1d(512, 512, 1) 46 | self.conv3 = nn.Conv1d(512, 3, 1) 47 | 48 | def forward(self, x): 49 | batch_size = x.size()[0] 50 | coarse = F.relu(self.fc1(x)) 51 | coarse = F.relu(self.fc2(coarse)) 52 | coarse = self.fc3(coarse).view(-1, 3, self.num_coarse) 53 | 54 | grid = self.grid.clone().detach() 55 | grid_feat = grid.unsqueeze(0).repeat(batch_size, 1, self.num_coarse).contiguous().cuda() 56 | 57 | point_feat = ( 58 | (coarse.transpose(1, 2).contiguous()).unsqueeze(2).repeat(1, 1, self.scale, 1).view(-1, self.num_fine, 59 | 3)).transpose(1, 60 | 2).contiguous() 61 | 62 | global_feat = x.unsqueeze(2).repeat(1, 1, self.num_fine) 63 | 64 | feat = torch.cat((grid_feat, point_feat, global_feat), 1) 65 | 66 | center = ((coarse.transpose(1, 2).contiguous()).unsqueeze(2).repeat(1, 1, self.scale, 1).view(-1, self.num_fine, 67 | 3)).transpose(1, 68 | 2).contiguous() 69 | 70 | fine = self.conv3(F.relu(self.conv2(F.relu(self.conv1(feat))))) + center 71 | return coarse, fine 72 | 73 | 74 | class Model(nn.Module): 75 | def __init__(self, args, num_coarse=1024): 76 | super(Model, self).__init__() 77 | 78 | self.num_coarse = num_coarse 79 | self.num_points = args.num_points 80 | self.train_loss = args.loss 81 | self.eval_emd = args.eval_emd 82 | self.scale = self.num_points // num_coarse 83 | self.cat_feature_num = 2 + 3 + 1024 84 | 85 | self.encoder = PCN_encoder() 86 | self.decoder = PCN_decoder(num_coarse, self.num_points, self.scale, self.cat_feature_num) 87 | 88 | def forward(self, x, gt=None, prefix="train", mean_feature=None, alpha=None): 89 | feat = self.encoder(x) 90 | out1, out2 = self.decoder(feat) 91 | out1 = out1.transpose(1, 2).contiguous() 92 | out2 = out2.transpose(1, 2).contiguous() 93 | 94 | if prefix=="train": 95 | if self.train_loss == 'emd': 96 | loss1 = calc_emd(out1, gt) 97 | loss2 = calc_emd(out2, gt) 98 | elif self.train_loss == 'cd': 99 | loss1, _ = calc_cd(out1, gt) 100 | loss2, _ = calc_cd(out2, gt) 101 | else: 102 | raise NotImplementedError('Train loss is either CD or EMD!') 103 | 104 | total_train_loss = loss1.mean() + loss2.mean() * alpha 105 | return out2, loss2, total_train_loss 106 | elif prefix=="val": 107 | if self.eval_emd: 108 | emd = calc_emd(out2, gt, eps=0.004, iterations=3000) 109 | else: 110 | emd = 0 111 | cd_p, cd_t, f1 = calc_cd(out2, gt, calc_f1=True) 112 | return {'out1': out1, 'out2': out2, 'emd': emd, 'cd_p': cd_p, 'cd_t': cd_t, 'f1': f1} 113 | else: 114 | return {'result': out2} -------------------------------------------------------------------------------- /completion/requirements.txt: -------------------------------------------------------------------------------- 1 | h5py==2.10.0 2 | matplotlib==3.0.3 3 | munch==2.5.0 4 | open3d==0.9.0 5 | PyYAML>=5.4 6 | ninja 7 | future -------------------------------------------------------------------------------- /completion/run_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | job_name=mvp_completion_test 4 | gpus=$2 5 | 6 | srun -p dsta --mpi=pmi2 --gres=gpu:${gpus} -n1 \ 7 | --ntasks-per-node=1 --job-name=${job_name}$\ 8 | --kill-on-bad-exit=1 -w SG-IDC1-10-51-2-38 \ 9 | python test.py --config cfgs/$1.yaml 10 | -------------------------------------------------------------------------------- /completion/run_train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | job_name=mvp_completion_train 4 | gpus=$2 5 | 6 | srun -p dsta --mpi=pmi2 --gres=gpu:${gpus} -n1 \ 7 | --ntasks-per-node=1 --job-name=${job_name}$\ 8 | --kill-on-bad-exit=1 -w SG-IDC1-10-51-2-38 \ 9 | python train.py --config cfgs/$1.yaml 10 | -------------------------------------------------------------------------------- /completion/test.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import sys 4 | import importlib 5 | import argparse 6 | import numpy as np 7 | import h5py 8 | import subprocess 9 | 10 | from numpy.lib.index_tricks import AxisConcatenator 11 | import munch 12 | import yaml 13 | # from utils.vis_utils import plot_single_pcd 14 | # from utils.train_utils import * 15 | from vis_utils import plot_single_pcd 16 | from train_utils import * 17 | from dataset import MVP_CP 18 | 19 | import warnings 20 | warnings.filterwarnings("ignore") 21 | 22 | 23 | def test(): 24 | dataset_test = MVP_CP(prefix="test") 25 | dataloader_test = torch.utils.data.DataLoader(dataset_test, batch_size=args.batch_size, 26 | shuffle=False, num_workers=int(args.workers)) 27 | dataset_length = len(dataset_test) 28 | logging.info('Length of test dataset:%d', len(dataset_test)) 29 | 30 | # load model 31 | model_module = importlib.import_module('.%s' % args.model_name, 'models') 32 | net = torch.nn.DataParallel(model_module.Model(args)) 33 | net.cuda() 34 | net.module.load_state_dict(torch.load(args.load_model)['net_state_dict']) 35 | logging.info("%s's previous weights loaded." % args.model_name) 36 | net.eval() 37 | 38 | logging.info('Testing...') 39 | with torch.no_grad(): 40 | results_list = [] 41 | for i, data in enumerate(dataloader_test): 42 | 43 | inputs_cpu = data 44 | 45 | inputs = inputs_cpu.float().cuda() 46 | inputs = inputs.transpose(2, 1).contiguous() 47 | 48 | result_dict = net(inputs, prefix="test") 49 | results_list.append(result_dict['result'].cpu().numpy()) 50 | 51 | if i % args.step_interval_to_print == 0: 52 | logging.info('test [%d/%d]' % (i, dataset_length / args.batch_size)) 53 | 54 | all_results = np.concatenate(results_list, axis=0) 55 | print(all_results.shape) 56 | 57 | with h5py.File(log_dir + '/results.h5', 'w') as f: 58 | f.create_dataset('results', data=all_results) 59 | 60 | cur_dir = os.getcwd() 61 | cmd = "cd %s; zip -r submission.zip results.h5 ; cd %s" % (log_dir, cur_dir) 62 | process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE,stderr=subprocess.PIPE) 63 | _, _ = process.communicate() 64 | print("Submission file has been saved to %s/submission.zip" % (log_dir)) 65 | 66 | 67 | if __name__ == "__main__": 68 | parser = argparse.ArgumentParser(description='Test config file') 69 | parser.add_argument('-c', '--config', help='path to config file', required=True) 70 | arg = parser.parse_args() 71 | config_path = arg.config 72 | args = munch.munchify(yaml.safe_load(open(config_path))) 73 | 74 | if not args.load_model: 75 | raise ValueError('Model path must be provided to load model!') 76 | 77 | exp_name = os.path.basename(args.load_model) 78 | log_dir = os.path.dirname(args.load_model) 79 | logging.basicConfig(level=logging.INFO, handlers=[logging.FileHandler(os.path.join(log_dir, 'test.log')), 80 | logging.StreamHandler(sys.stdout)]) 81 | 82 | test() 83 | -------------------------------------------------------------------------------- /completion/train_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | class AverageValueMeter(object): 4 | def __init__(self): 5 | self.reset() 6 | 7 | def reset(self): 8 | self.val = 0 9 | self.avg = 0 10 | self.sum = 0 11 | self.count = 0.0 12 | 13 | def update(self, val, n=1): 14 | self.val = val 15 | self.sum += val * n 16 | self.count += n 17 | self.avg = self.sum / self.count 18 | 19 | 20 | def set_requires_grad(nets, requires_grad=False): 21 | if not isinstance(nets, list): 22 | nets = [nets] 23 | for net in nets: 24 | if net is not None: 25 | for param in net.parameters(): 26 | param.requires_grad = requires_grad 27 | 28 | 29 | def save_model(path, net, net_d=None): 30 | if net_d is not None: 31 | torch.save({'net_state_dict': net.module.state_dict(), 32 | 'D_state_dict': net_d.module.state_dict()}, path) 33 | else: 34 | torch.save({'net_state_dict': net.module.state_dict()}, path) 35 | 36 | 37 | def generator_step(net_d, out2, net_loss, optimizer): 38 | set_requires_grad(net_d, False) 39 | d_fake = net_d(out2[:, 0:2048, :]) 40 | errG_loss_batch = torch.mean((d_fake - 1) ** 2) 41 | total_gen_loss_batch = errG_loss_batch + net_loss * 200 42 | total_gen_loss_batch.backward(torch.ones(torch.cuda.device_count()).cuda(), retain_graph=True, ) 43 | optimizer.step() 44 | return d_fake 45 | 46 | 47 | def discriminator_step(net_d, gt, d_fake, optimizer_d): 48 | set_requires_grad(net_d, True) 49 | d_real = net_d(gt[:, 0:2048, :]) 50 | d_loss_fake = torch.mean(d_fake ** 2) 51 | d_loss_real = torch.mean((d_real - 1) ** 2) 52 | errD_loss_batch = 0.5 * (d_loss_real + d_loss_fake) 53 | total_dis_loss_batch = errD_loss_batch 54 | total_dis_loss_batch.backward(torch.ones(torch.cuda.device_count()).cuda()) 55 | optimizer_d.step() 56 | 57 | 58 | 59 | 60 | 61 | -------------------------------------------------------------------------------- /completion/vis_utils.py: -------------------------------------------------------------------------------- 1 | import open3d as o3d 2 | import numpy as np 3 | from mpl_toolkits.mplot3d import Axes3D 4 | import matplotlib.pyplot as plt 5 | 6 | 7 | def get_pts(pcd): 8 | points = np.asarray(pcd.points) 9 | X = [] 10 | Y = [] 11 | Z = [] 12 | for pt in range(points.shape[0]): 13 | X.append(points[pt][0]) 14 | Y.append(points[pt][1]) 15 | Z.append(points[pt][2]) 16 | return np.asarray(X), np.asarray(Y), np.asarray(Z) 17 | 18 | 19 | def set_axes_equal(ax): 20 | x_limits = ax.get_xlim3d() 21 | y_limits = ax.get_ylim3d() 22 | z_limits = ax.get_zlim3d() 23 | x_range = abs(x_limits[1] - x_limits[0]) 24 | x_middle = np.mean(x_limits) 25 | y_range = abs(y_limits[1] - y_limits[0]) 26 | y_middle = np.mean(y_limits) 27 | z_range = abs(z_limits[1] - z_limits[0]) 28 | z_middle = np.mean(z_limits) 29 | plot_radius = 0.5*max([x_range, y_range, z_range]) 30 | ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius]) 31 | ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius]) 32 | ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius]) 33 | 34 | 35 | def plot_single_pcd(points, save_path): 36 | fig = plt.figure() 37 | ax = fig.add_subplot(111, projection='3d') 38 | ax.set_aspect('equal') 39 | pcd = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(points)) 40 | rotation_matrix = np.asarray([[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]) 41 | pcd = pcd.transform(rotation_matrix) 42 | X, Y, Z = get_pts(pcd) 43 | t = Z 44 | ax.scatter(X, Y, Z, c=t, cmap='jet', marker='o', s=0.5, linewidths=0) 45 | ax.grid(False) 46 | ax.w_xaxis.set_pane_color((1.0, 1.0, 1.0, 1.0)) 47 | ax.w_yaxis.set_pane_color((1.0, 1.0, 1.0, 1.0)) 48 | ax.w_zaxis.set_pane_color((1.0, 1.0, 1.0, 1.0)) 49 | set_axes_equal(ax) 50 | plt.axis('off') 51 | plt.savefig(save_path, format='png', dpi=600) 52 | plt.close() 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | -------------------------------------------------------------------------------- /images/complete_pcds.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/images/complete_pcds.gif -------------------------------------------------------------------------------- /images/logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/images/logo.png -------------------------------------------------------------------------------- /images/partial_pcds.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/images/partial_pcds.gif -------------------------------------------------------------------------------- /registration/.gitignore: -------------------------------------------------------------------------------- 1 | /.idea/* 2 | /.DS_Store 3 | /log/* 4 | */**/__pycache__ 5 | /__pycache__/* 6 | *.h5 7 | *.backup 8 | gmcnet.py -------------------------------------------------------------------------------- /registration/README.md: -------------------------------------------------------------------------------- 1 | # *MVP Benchmark:* Partial-to-Partial Point Cloud Registration 2 |

3 | 4 |

5 | 6 | We include the following methods for point cloud registration: 7 | 8 | [1] [DCP](https://github.com/WangYueFt/dcp);   [2] [DeepGMR](https://github.com/wentaoyuan/deepgmr);   [3] [IDAM](https://github.com/jiahaowork/idam) 9 | 10 | 11 | ### MVP Registration Dataset 12 | 16 | Download the MVP registration dataset [Google Drive](https://drive.google.com/drive/folders/1RlUW0vmmyqxkBTM_ITVguAjxzIS1MFz4) or [百度网盘](https://pan.baidu.com/s/18pli79KSGGsWQ8FPiSW9qg)  (code: p364) to the folder "data". 17 | 18 | The data structure will be: 19 | ``` 20 | data 21 | ├── MVP_Train_RG.h5 22 | | ├── src (6400, 2048, 3) 23 | | ├── tgt (6400, 2048, 3) 24 | | ├── complete (6400, 2048, 3) 25 | | ├── cat_label (6400,) 26 | | ├── match_id (6400,) 27 | | └── match_level (6400,) 28 | ├── MVP_Test_RG.h5 29 | | ├── rotated_src (1200, 2048, 3) 30 | | ├── rotated_tgt (1200, 2048, 3) 31 | | ├── pose_src (1200, 4, 4) 32 | | ├── pose_tgt (1200, 4, 4) 33 | | ├── rot_level (1200,) 34 | | ├── transforms (1200, 4, 4) 35 | | ├── src (1200, 2048, 3) 36 | | ├── tgt (1200, 2048, 3) 37 | | ├── complete (1200, 2048, 3) 38 | | ├── cat_label (1200,) 39 | | ├── match_id (1200,) 40 | | └── match_level (1200,) 41 | └── MVP_ExtraTest_RG.h5 42 | ├── rotated_src (2000, 2048, 3) 43 | ├── rotated_tgt (2000, 2048, 3) 44 | └── cat_label (2000,) 45 | ``` 46 | **Attention**: `MVP_ExtraTest_Shuffled_RG.h5` is only available during the competition period. `MVP_Test_RG.h5` is the standard test set if you use MVP dataset in your papers. 47 | 48 | We create the registration dataset by ensuring surfficent overlaps between the source point cloud and the target. 49 | Partial point cloud pairs with "match_level = 1" mostly have more correspondences than those with "match_level = 0". 50 | 51 | Most relative rotations are within [0, 45\textdegree], and the rest have unrestricted rotations \in [0, 360\textdegree]. 52 | The ratio is roughly 4 : 1. 53 | 54 | Note that the source and the target are two different incomplete point clouds scanned from the same CAD model. 55 | 56 | 57 | ### Usage 58 | + To train a model: run `python train.py -c ./cfgs/*.yaml`, e.g. `python train.py -c ./cfgs/pcn.yaml` 59 | + To test a model: run `python test.py -c ./cfgs/*.yaml`, e.g. `python test.py -c ./cfgs/pcn.yaml` 60 | + Config for each algorithm can be found in `cfgs/`. 61 | + `run_train.sh` and `run_test.sh` are provided for SLURM users. 62 | 63 | 64 | ## Citation 65 | If you find our code useful, please cite our paper: 66 | ```bibtex 67 | @inproceedings{pan2021variational, 68 | title={Variational Relational Point Completion Network}, 69 | author={Pan, Liang and Chen, Xinyi and Cai, Zhongang and Zhang, Junzhe and Zhao, Haiyu and Yi, Shuai and Liu, Ziwei}, 70 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, 71 | pages={8524--8533}, 72 | year={2021} 73 | } 74 | @article{pan2021robust, 75 | title={Robust Partial-to-Partial Point Cloud Registration in a Full Range}, 76 | author={Pan, Liang and Cai, Zhongang and Liu, Ziwei}, 77 | journal={arXiv preprint arXiv:2111.15606}, 78 | year={2021} 79 | } 80 | @article{pan2021mvp, 81 | title={Multi-View Partial (MVP) Point Cloud Challenge 2021 on Completion and Registration: Methods and Results}, 82 | author={Pan, Liang and Wu, Tong and Cai, Zhongang and Liu, Ziwei and Yu, Xumin and Rao, Yongming and Lu, Jiwen and Zhou, Jie and Xu, Mingye and Luo, Xiaoyuan and Fu, Kexue, and Gao, Peng, and Wang, Manning, and Wang, Yali, and Qiao, Yu, and Zhou, Junsheng, and Wen, Xin, and Xiang, Peng, and Liu, Yu-Shen, and Han, Zhizhong, and Yan, Yuanjie, and An, Junyi, and Zhu, Lifa, and Lin, Changwei, and Liu, Dongrui, and Li, Xin, and G ́omez-Fern ́andez, Francisco, and Wang, Qinlong, and Yang, Yang}, 83 | journal={arXiv preprint arXiv:2112.12053}, 84 | year={2021} 85 | } 86 | ``` 87 | 88 | ## License 89 | Our code is released under Apache-2.0 License. 90 | 91 | 92 | ## Acknowledgement 93 | We include the following algorithms: 94 | [1] [DCP](https://github.com/WangYueFt/dcp) 95 | [2] [DeepGMR](https://github.com/wentaoyuan/deepgmr) 96 | [3] [IDAM](https://github.com/jiahaowork/idam) 97 | -------------------------------------------------------------------------------- /registration/cfgs/dcp.yaml: -------------------------------------------------------------------------------- 1 | batch_size: 32 2 | workers: 0 3 | nepoch: 100 4 | model_name: dcp 5 | load_model: null 6 | start_epoch: 0 7 | work_dir: log/ 8 | flag: debug 9 | 10 | manual_seed: null 11 | 12 | step_interval_to_print: 30 13 | step_interval_to_plot: 250 14 | epoch_interval_to_save: 10 15 | epoch_interval_to_val: 1 16 | 17 | lr: 0.001 18 | lr_decay: True 19 | lr_decay_rate: 0.5 20 | lr_clip: 1.e-6 21 | optimizer: Adam 22 | weight_decay: 0 23 | betas: 0.9, 0.999 24 | 25 | use_rri: False 26 | rri_size: 5 27 | num_clusters: 16 28 | num_points: 2048 29 | use_tnet: False 30 | use_fpfh: False 31 | use_ppf: False 32 | descriptor_size: 512 33 | max_angle: 180 34 | max_trans: 0.5 35 | category: null 36 | benchmark: mvp 37 | num_rot_levels: 2 38 | num_corr_levels: 2 39 | -------------------------------------------------------------------------------- /registration/cfgs/deepgmr.yaml: -------------------------------------------------------------------------------- 1 | batch_size: 32 2 | workers: 0 3 | nepoch: 100 4 | model_name: deepgmr 5 | load_model: null 6 | start_epoch: 0 7 | work_dir: log/ 8 | flag: debug 9 | 10 | manual_seed: null 11 | 12 | step_interval_to_print: 30 13 | step_interval_to_plot: 250 14 | epoch_interval_to_save: 10 15 | epoch_interval_to_val: 1 16 | 17 | lr: 0.001 18 | lr_decay: True 19 | lr_decay_rate: 0.5 20 | lr_clip: 1.e-6 21 | optimizer: Adam 22 | weight_decay: 0 23 | betas: 0.9, 0.999 24 | 25 | use_rri: True 26 | rri_size: 20 27 | num_groups: 16 28 | num_points: 2048 29 | use_tnet: False 30 | use_fpfh: False 31 | use_ppf: False 32 | descriptor_size: 1024 33 | max_angle: 180 34 | max_trans: 0.5 35 | category: null 36 | benchmark: mvp 37 | num_rot_levels: 2 38 | num_corr_levels: 2 39 | -------------------------------------------------------------------------------- /registration/cfgs/idam.yaml: -------------------------------------------------------------------------------- 1 | batch_size: 32 2 | workers: 0 3 | nepoch: 100 4 | model_name: idam 5 | # load_model: ./log/idam_mvp_debug_2021-06-22T02:59:30/best_RotE_network.pth 6 | load_model: null 7 | start_epoch: 0 8 | work_dir: log/ 9 | flag: debug 10 | 11 | manual_seed: null 12 | 13 | step_interval_to_print: 30 14 | step_interval_to_plot: 250 15 | epoch_interval_to_save: 10 16 | epoch_interval_to_val: 1 17 | 18 | lr: 0.001 19 | lr_decay: True 20 | lr_decay_rate: 0.5 21 | lr_clip: 1.e-6 22 | optimizer: Adam 23 | weight_decay: 0 24 | betas: 0.9, 0.999 25 | 26 | use_rri: False 27 | rri_size: 20 28 | num_groups: 16 29 | num_points: 2048 30 | # use_tnet: False 31 | use_fpfh: False 32 | use_ppf: False 33 | descriptor_size: 64 34 | num_iters: 3 35 | 36 | max_angle: 180 37 | max_trans: 0.5 38 | category: null 39 | benchmark: mvp 40 | num_rot_levels: 2 41 | num_corr_levels: 2 42 | -------------------------------------------------------------------------------- /registration/data/.gitignore: -------------------------------------------------------------------------------- 1 | *.h5 -------------------------------------------------------------------------------- /registration/dataset.py: -------------------------------------------------------------------------------- 1 | import h5py 2 | import numpy as np 3 | import os 4 | import open3d as o3d 5 | import torch 6 | from torch.utils.data import Dataset 7 | 8 | 9 | def jitter_pcd(pcd, sigma=0.01, clip=0.05): 10 | pcd += np.clip(sigma * np.random.randn(*pcd.shape), -1 * clip, clip) 11 | return pcd 12 | 13 | 14 | def random_pose(max_angle, max_trans): 15 | R = random_rotation(max_angle) 16 | t = random_translation(max_trans) 17 | return np.concatenate([np.concatenate([R, t], 1), [[0, 0, 0, 1]]], 0) 18 | 19 | 20 | def random_rotation(max_angle): 21 | axis = np.random.randn(3) 22 | axis /= np.linalg.norm(axis) 23 | angle = np.random.rand() * max_angle 24 | A = np.array([[0, -axis[2], axis[1]], 25 | [axis[2], 0, -axis[0]], 26 | [-axis[1], axis[0], 0]]) 27 | R = np.eye(3) + np.sin(angle) * A + (1 - np.cos(angle)) * np.dot(A, A) 28 | return R 29 | 30 | 31 | def random_translation(max_dist): 32 | t = np.random.randn(3) 33 | t /= np.linalg.norm(t) 34 | t *= np.random.rand() * max_dist 35 | return np.expand_dims(t, 1) 36 | 37 | 38 | class MVP_RG(Dataset): 39 | """docstring for MVP_RG""" 40 | def __init__(self, prefix, args): 41 | self.prefix = prefix 42 | 43 | if self.prefix == "train": 44 | f = h5py.File('./data/MVP_Train_RG.h5', 'r') 45 | elif self.prefix == "val": 46 | f = h5py.File('./data/MVP_Test_RG.h5', 'r') 47 | elif self.prefix == "test": 48 | f = h5py.File('./data/MVP_ExtraTest_RG.h5', 'r') 49 | 50 | self.max_angle = args.max_angle / 180 * np.pi 51 | self.max_trans = args.max_trans 52 | 53 | self.label = f['cat_labels'][:].astype('int32') 54 | if self.prefix == "test": 55 | self.src = np.array(f['rotated_src'][:].astype('float32')) 56 | self.tgt = np.array(f['rotated_tgt'][:].astype('float32')) 57 | else: 58 | self.match_level = np.array(f['match_level'][:].astype('int32')) 59 | 60 | match_id = [] 61 | for i in range(len(f['match_id'].keys())): 62 | ds_data = f['match_id'][str(i)][:] 63 | match_id.append(ds_data) 64 | self.match_id = np.array(match_id, dtype=object) 65 | 66 | if self.prefix == "train": 67 | self.src = np.array(f['src'][:].astype('float32')) 68 | self.tgt = np.array(f['tgt'][:].astype('float32')) 69 | if args.max_angle > 45: 70 | self.rot_level = int(1) 71 | else: 72 | self.rot_level = int(0) 73 | elif self.prefix == "val": 74 | self.src = np.array(f['rotated_src'][:].astype('float32')) 75 | self.tgt = np.array(f['rotated_tgt'][:].astype('float32')) 76 | self.transforms = np.array(f['transforms'][:].astype('float32')) 77 | self.rot_level = np.array(f['rot_level'][:].astype('int32')) 78 | 79 | f.close() 80 | 81 | if args.category: 82 | self.src = self.src[self.label==args.category] 83 | self.tgt = self.tgt[self.label==args.category] 84 | 85 | if self.prefix is not "test": 86 | self.match_id = self.match_id[self.label==args.category] 87 | self.match_level = self.match_level[self.label==args.category] 88 | if self.prefix == False: 89 | self.transforms = self.transforms[self.label==args.category] 90 | self.rot_level = self.rot_level[self.label==args.category] 91 | self.label = self.label[self.label==args.category] 92 | 93 | # print(self.src.shape, self.tgt.shape, self.match_id.shape, self.match_level.shape, self.label.shape) 94 | 95 | def __len__(self): 96 | return self.src.shape[0] 97 | 98 | def __getitem__(self, index): 99 | src = self.src[index] 100 | tgt = self.tgt[index] 101 | 102 | if self.prefix == "train": 103 | transform = random_pose(self.max_angle, self.max_trans / 2) 104 | pose1 = random_pose(np.pi, self.max_trans) 105 | pose2 = transform @ pose1 106 | src = src @ pose1[:3, :3].T + pose1[:3, 3] 107 | tgt = tgt @ pose2[:3, :3].T + pose2[:3, 3] 108 | rot_level = self.rot_level 109 | match_level = self.match_level[index] 110 | 111 | elif self.prefix == "val": 112 | transform = self.transforms[index] 113 | rot_level = self.rot_level[index] 114 | match_level = self.match_level[index] 115 | 116 | # src = np.random.permutation(src) 117 | # tgt = np.random.permutation(tgt) 118 | 119 | src = torch.from_numpy(src) 120 | tgt = torch.from_numpy(tgt) 121 | 122 | if self.prefix is not "test": 123 | transform = torch.from_numpy(transform) 124 | match_level = match_level 125 | rot_level = rot_level 126 | return src, tgt, transform, match_level, rot_level 127 | else: 128 | return src, tgt 129 | -------------------------------------------------------------------------------- /registration/images/registration.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/registration/images/registration.png -------------------------------------------------------------------------------- /registration/images/registration_black.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/registration/images/registration_black.png -------------------------------------------------------------------------------- /registration/images/registration_raw.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/registration/images/registration_raw.png -------------------------------------------------------------------------------- /registration/run_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | job_name=mvp_registration_test 4 | gpus=$2 5 | 6 | srun -p dsta --mpi=pmi2 --gres=gpu:${gpus} -n1 \ 7 | --ntasks-per-node=1 --job-name=${job_name}$\ 8 | --kill-on-bad-exit=1 -w SG-IDC1-10-51-2-38 \ 9 | python test.py --config cfgs/$1.yaml 10 | -------------------------------------------------------------------------------- /registration/run_train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | job_name=mvp_registration_train 4 | gpus=$2 5 | 6 | srun -p dsta --mpi=pmi2 --gres=gpu:${gpus} -n1 \ 7 | --ntasks-per-node=1 --job-name=${job_name}$\ 8 | --kill-on-bad-exit=1 -w SG-IDC1-10-51-2-38 \ 9 | python train.py --config cfgs/$1.yaml 10 | -------------------------------------------------------------------------------- /registration/test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | import torch 4 | import torch.nn as nn 5 | import h5py 6 | import subprocess 7 | 8 | import torch.optim as optim 9 | from tqdm import tqdm 10 | import os 11 | import random 12 | import sys 13 | 14 | import logging 15 | import math 16 | import importlib 17 | import datetime 18 | import munch 19 | import yaml 20 | import argparse 21 | import copy 22 | 23 | from train_utils import AverageValueMeter, save_model 24 | from dataset import MVP_RG 25 | 26 | 27 | def test(): 28 | logging.info(str(args)) 29 | 30 | dataset_test = MVP_RG(prefix="test", args=args) 31 | dataloader_test = torch.utils.data.DataLoader(dataset_test, batch_size=args.batch_size, 32 | shuffle=False, num_workers=int(args.workers)) 33 | logging.info('Length of test dataset:%d', len(dataset_test)) 34 | 35 | model_module = importlib.import_module('.%s' % args.model_name, 'models') 36 | net = torch.nn.DataParallel(model_module.Model(args)) 37 | net.cuda() 38 | 39 | if args.load_model: 40 | ckpt = torch.load(args.load_model) 41 | net.module.load_state_dict(ckpt['net_state_dict']) 42 | logging.info("%s's previous weights loaded." % args.model_name) 43 | 44 | logging.info('Testing...') 45 | net.module.eval() 46 | with torch.no_grad(): 47 | result_list = [] 48 | for _, data in enumerate(dataloader_test): 49 | src, tgt = data 50 | src = src.float().cuda() 51 | tgt = tgt.float().cuda() 52 | 53 | result = net(src, tgt, prefix="test") 54 | result_list.append(result.cpu().numpy()) 55 | 56 | all_results = np.concatenate(result_list, axis=0) 57 | print(all_results.shape) 58 | 59 | with h5py.File(log_dir + '/results.h5', 'w') as f: 60 | f.create_dataset('results', data=all_results) 61 | 62 | cur_dir = os.getcwd() 63 | cmd = "cd %s; zip -r submission.zip results.h5 ; cd %s" % (log_dir, cur_dir) 64 | process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE,stderr=subprocess.PIPE) 65 | _, _ = process.communicate() 66 | print("Submission file has been saved to %s/submission.zip" % (log_dir)) 67 | 68 | 69 | if __name__ == "__main__": 70 | parser = argparse.ArgumentParser(description='Train config file') 71 | parser.add_argument('-c', '--config', help='path to config file', required=True) 72 | arg = parser.parse_args() 73 | config_path = arg.config 74 | args = munch.munchify(yaml.safe_load(open(config_path))) 75 | 76 | time = datetime.datetime.now().isoformat()[:19] 77 | if args.load_model: 78 | exp_name = os.path.basename(os.path.dirname(args.load_model)) 79 | log_dir = os.path.dirname(args.load_model) 80 | else: 81 | exp_name = args.model_name + '_' + args.benchmark + '_' + args.flag + '_' + time 82 | log_dir = os.path.join(args.work_dir, exp_name) 83 | if not os.path.exists(log_dir): 84 | os.makedirs(log_dir) 85 | logging.basicConfig(level=logging.INFO, handlers=[logging.FileHandler(os.path.join(log_dir, 'train.log')), 86 | logging.StreamHandler(sys.stdout)]) 87 | test() 88 | 89 | 90 | -------------------------------------------------------------------------------- /registration/train_utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import os 3 | import argparse 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | import torch.optim as optim 8 | import numpy as np 9 | from scipy.spatial.transform import Rotation 10 | import math 11 | 12 | 13 | class AverageValueMeter(object): 14 | def __init__(self): 15 | self.reset() 16 | 17 | def reset(self): 18 | self.val = 0 19 | self.avg = 0 20 | self.sum = 0 21 | self.count = 0.0 22 | 23 | def update(self, val, n=1): 24 | self.val = val 25 | self.sum += val * n 26 | self.count += n 27 | self.avg = self.sum / self.count 28 | 29 | 30 | def save_model(path, net): 31 | torch.save({'net_state_dict': net.module.state_dict()}, path) 32 | 33 | 34 | # Part of the code is referred from: https://github.com/ClementPinard/SfmLearner-Pytorch/blob/master/inverse_warp.py 35 | 36 | def quat2mat(quat): 37 | x, y, z, w = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3] 38 | 39 | B = quat.size(0) 40 | 41 | w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2) 42 | wx, wy, wz = w*x, w*y, w*z 43 | xy, xz, yz = x*y, x*z, y*z 44 | 45 | rotMat = torch.stack([w2 + x2 - y2 - z2, 2*xy - 2*wz, 2*wy + 2*xz, 46 | 2*wz + 2*xy, w2 - x2 + y2 - z2, 2*yz - 2*wx, 47 | 2*xz - 2*wy, 2*wx + 2*yz, w2 - x2 - y2 + z2], dim=1).reshape(B, 3, 3) 48 | return rotMat 49 | 50 | 51 | def transform_point_cloud(point_cloud, rotation, translation): 52 | if len(rotation.size()) == 2: 53 | rot_mat = quat2mat(rotation) 54 | else: 55 | rot_mat = rotation 56 | return torch.matmul(rot_mat, point_cloud) + translation.unsqueeze(2) 57 | 58 | 59 | def npmat2euler(mats, seq='zyx'): 60 | eulers = [] 61 | for i in range(mats.shape[0]): 62 | r = Rotation.from_dcm(mats[i]) 63 | eulers.append(r.as_euler(seq, degrees=True)) 64 | return np.asarray(eulers, dtype='float32') 65 | 66 | 67 | def rt_to_transformation(R, t): 68 | bot_row = torch.Tensor([[[0, 0, 0, 1]]]).repeat(R.shape[0], 1, 1).to(R.device) 69 | T = torch.cat([torch.cat([R, t], dim=2), bot_row], dim=1) 70 | return T 71 | 72 | 73 | def rotation_error(R, R_gt): 74 | cos_theta = (torch.einsum('bij,bij->b', R, R_gt) - 1) / 2 75 | cos_theta = torch.clamp(cos_theta, -1, 1) 76 | return torch.acos(cos_theta) * 180 / math.pi 77 | 78 | 79 | def translation_error(t, t_gt): 80 | return torch.norm(t - t_gt, dim=1) 81 | 82 | 83 | def rmse_loss(pts, T, T_gt): 84 | pts_pred = pts @ T[:, :3, :3].transpose(1, 2) + T[:, :3, 3].unsqueeze(1) 85 | pts_gt = pts @ T_gt[:, :3, :3].transpose(1, 2) + T_gt[:, :3, 3].unsqueeze(1) 86 | return torch.norm(pts_pred - pts_gt, dim=2).mean(dim=1) 87 | 88 | 89 | def rotation_geodesic_error(m1, m2): 90 | batch=m1.shape[0] 91 | m = torch.bmm(m1, m2.transpose(1,2)) #batch*3*3 92 | 93 | cos = ( m[:,0,0] + m[:,1,1] + m[:,2,2] - 1 )/2 94 | cos = torch.min(cos, torch.autograd.Variable(torch.ones(batch).cuda()) ) 95 | cos = torch.max(cos, torch.autograd.Variable(torch.ones(batch).cuda())*-1 ) 96 | 97 | theta = torch.acos(cos) 98 | 99 | #theta = torch.min(theta, 2*np.pi - theta) 100 | 101 | return theta -------------------------------------------------------------------------------- /registration/visu_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from matplotlib import cm 3 | from matplotlib import pyplot as plt 4 | from mpl_toolkits.mplot3d import Axes3D 5 | 6 | 7 | def plot_pcd(ax, pcd, color=None, cmap='viridis', size=4, alpha=0.9, azim=60, elev=0): 8 | if color is None: 9 | color = pcd[:, 0] 10 | vmin = -2 11 | vmax = 1.5 12 | else: 13 | vmin = 0 14 | vmax = 1 15 | ax.view_init(azim=azim, elev=elev) 16 | ax.scatter(pcd[:, 0], pcd[:, 1], pcd[:, 2], c=color, s=size, cmap=cmap, vmin=vmin, vmax=vmax, alpha=alpha) 17 | lims = np.array([ax.get_xlim3d(), ax.get_ylim3d(), ax.get_zlim3d()]) 18 | min_lim = min(pcd.min() * 0.9, lims.min()) 19 | max_lim = max(pcd.max() * 0.9, lims.max()) 20 | for axis in 'xyz': 21 | getattr(ax, 'set_{}lim'.format(axis))((min_lim, max_lim)) 22 | ax.set_axis_off() 23 | 24 | 25 | def plot_matches(ax, mpts1, mpts2, color=None, cmap='viridis', size=4, alpha=0.9, azim=60, elev=0): 26 | if color is None: 27 | color = np.arange(mpts1.shape[0]) / (mpts1.shape[0] - 1) 28 | if cmap is not None: 29 | cmap = cm.get_cmap(cmap) 30 | color = cmap(color) 31 | 32 | ax.view_init(azim=azim, elev=elev) 33 | 34 | for k in range(mpts1.shape[0]): 35 | ptp = np.array([mpts1[k], mpts2[k]]) 36 | ax.plot(ptp[:, 0], ptp[:, 1], ptp[:, 2], color=color[k], marker='o', markersize=12) 37 | 38 | 39 | def plot_gmm(ax, mix, mu, cov, color=None, cmap='viridis', azim=60, elev=0, numWires=15, wireframe=True): 40 | if color is None: 41 | color = np.arange(mix.shape[0]) / (mix.shape[0] - 1) 42 | if cmap is not None: 43 | cmap = cm.get_cmap(cmap) 44 | color = cmap(color) 45 | 46 | u = np.linspace(0.0, 2.0 * np.pi, numWires) 47 | v = np.linspace(0.0, np.pi, numWires) 48 | X = np.outer(np.cos(u), np.sin(v)) 49 | Y = np.outer(np.sin(u), np.sin(v)) 50 | Z = np.outer(np.ones_like(u), np.cos(v)) 51 | XYZ = np.stack([X.flatten(), Y.flatten(), Z.flatten()]) 52 | 53 | alpha = mix / mix.max() 54 | ax.view_init(azim=azim, elev=elev) 55 | 56 | for k in range(mix.shape[0]): 57 | # find the rotation matrix and radii of the axes 58 | U, s, V = np.linalg.svd(cov[k]) 59 | x, y, z = V.T @ (np.sqrt(s)[:, None] * XYZ) + mu[k][:, None] 60 | x = x.reshape(numWires, numWires) 61 | y = y.reshape(numWires, numWires) 62 | z = z.reshape(numWires, numWires) 63 | if wireframe: 64 | ax.plot_wireframe(x, y, z, rstride=1, cstride=1, color=color[k], alpha=alpha[k]) 65 | else: 66 | ax.plot_surface(x, y, z, rstride=1, cstride=1, color=color[k], alpha=alpha[k]) 67 | 68 | 69 | def visualize(inputs): 70 | for i in range(len(inputs)): 71 | inputs[i] = inputs[i].detach().cpu().numpy() 72 | p1, gamma1, pi1, mu1, sigma1, p2, gamma2, pi2, mu2, sigma2, \ 73 | p1_trans, init_r_err, init_t_err, init_rmse, r_err, t_err, rmse = inputs 74 | 75 | fig = plt.figure(figsize=(8, 8)) 76 | title = 'Rotation error {:.2f}\nTranslation error {:.4f}\nRMSE {:.4f}' 77 | 78 | ax = fig.add_subplot(221, projection='3d') 79 | plot_pcd(ax, p1, cmap='Reds') 80 | plot_pcd(ax, p2, cmap='Blues') 81 | ax.set_title(title.format(init_r_err, init_t_err, init_rmse)) 82 | 83 | ax = fig.add_subplot(222, projection='3d') 84 | plot_pcd(ax, p1_trans, cmap='Reds') 85 | plot_pcd(ax, p2, cmap='Blues') 86 | ax.set_title(title.format(r_err, t_err, rmse)) 87 | 88 | ax = fig.add_subplot(223, projection='3d') 89 | color1 = np.argmax(gamma1, axis=1) / (gamma1.shape[1] - 1) 90 | plot_pcd(ax, p1, color1) 91 | plot_gmm(ax, pi1, mu1, sigma1) 92 | ax.set_title('Source GMM') 93 | 94 | ax = fig.add_subplot(224, projection='3d') 95 | color2 = np.argmax(gamma2, axis=1) / (gamma2.shape[1] - 1) 96 | plot_pcd(ax, p2, color2) 97 | plot_gmm(ax, pi2, mu2, sigma2) 98 | ax.set_title('Target GMM') 99 | 100 | plt.tight_layout() 101 | return fig 102 | -------------------------------------------------------------------------------- /setup.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ############## make sure you can write in \tmp; Or you should set TORCH_EXTENSIONS_DIR 4 | # e.g. export TORCH_EXTENSIONS_DIR=/mnt/lustre/$YourUserName$/tmp 5 | 6 | conda create -n mvp python=3.7 -y 7 | conda activate mvp 8 | conda install pytorch==1.5.0 torchvision==0.6.0 cudatoolkit=10.1 -c pytorch -y 9 | 10 | # setup completion 11 | cd completion 12 | pip install -r requirements.txt 13 | 14 | 15 | cd ../utils/mm3d_pn2/ 16 | sh setup.sh 17 | 18 | ############## make sure NVCC is in your environment 19 | # SLURM users 20 | # sh run_build.sh 21 | # or 22 | pip install -v -e . 23 | # python setup.py develop 24 | 25 | cd ../../ 26 | 27 | # setup registration 28 | 29 | 30 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .metrics import (cd, fscore, emd) 2 | from .mm3d_pn2 import (nms, RoIAlign, roi_align, get_compiler_version, get_compiling_cuda_version, 3 | NaiveSyncBatchNorm1d, NaiveSyncBatchNorm2d, sigmoid_focal_loss, SigmoidFocalLoss, ball_query, knn, 4 | furthest_point_sample, furthest_point_sample_with_dist, three_interpolate, three_nn, gather_points, 5 | grouping_operation, group_points, GroupAll, QueryAndGroup, get_compiler_version, get_compiling_cuda_version, 6 | Points_Sampler) 7 | 8 | __all__ = [ 9 | 'cd', 'fscore', 'emd', 10 | 'nms', 11 | 'RoIAlign', 'roi_align', 'get_compiler_version', 12 | 'get_compiling_cuda_version', 'NaiveSyncBatchNorm1d', 13 | 'NaiveSyncBatchNorm2d', 14 | 'sigmoid_focal_loss', 15 | 'SigmoidFocalLoss', 16 | 'ball_query', 'knn', 'furthest_point_sample', 17 | 'furthest_point_sample_with_dist', 'three_interpolate', 'three_nn', 18 | 'gather_points', 'grouping_operation', 'group_points', 'GroupAll', 19 | 'QueryAndGroup', 20 | 'get_compiler_version', 21 | 'get_compiling_cuda_version', 'Points_Sampler', 22 | ] -------------------------------------------------------------------------------- /utils/metrics/CD/.gitignore: -------------------------------------------------------------------------------- 1 | *__pycache__* -------------------------------------------------------------------------------- /utils/metrics/CD/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 ThibaultGROUEIX 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /utils/metrics/CD/README.md: -------------------------------------------------------------------------------- 1 | # Pytorch Chamfer Distance (CD). 2 | 3 | ### What is the Chamfer Distance ? 4 | 5 | [Stanford course](http://graphics.stanford.edu/courses/cs468-17-spring/LectureSlides/L14%20-%203d%20deep%20learning%20on%20point%20cloud%20representation%20(analysis).pdf) on 3D deep Learning 6 | 7 | Include a **CUDA** version, and a **PYTHON** version with pytorch standard operations. 8 | NB : In this depo, dist1 and dist2 are squared pointcloud euclidean distances, so you should adapt thresholds accordingly. 9 | 10 | - [x] F - Score 11 | 12 | 13 | ### Usage 14 | 15 | ```python 16 | import torch, chamfer3D.dist_chamfer_3D, fscore 17 | chamLoss = chamfer3D.dist_chamfer_3D.chamfer_3DDist() 18 | points1 = torch.rand(32, 1000, 3).cuda() 19 | points2 = torch.rand(32, 2000, 3, requires_grad=True).cuda() 20 | dist1, dist2, idx1, idx2 = chamLoss(points1, points2) 21 | f_score, precision, recall = fscore.fscore(dist1, dist2) 22 | ``` 23 | 24 | ### Aknowledgment 25 | 26 | Original backbone from [Fei Xia](https://github.com/fxia22/pointGAN/blob/master/nndistance/src/nnd_cuda.cu). 27 | 28 | JIT cool trick from [Christian Diller](https://github.com/chrdiller) 29 | 30 | Modify from [Thibault GROUEIX](https://github.com/ThibaultGROUEIX/ChamferDistancePytorch) 31 | 32 | 33 | -------------------------------------------------------------------------------- /utils/metrics/CD/__init__.py: -------------------------------------------------------------------------------- 1 | from .chamfer3D.dist_chamfer_3D import chamfer_3DDist as cd 2 | from .fscore import fscore 3 | 4 | __all__ = ['cd', 'fscore'] -------------------------------------------------------------------------------- /utils/metrics/CD/chamfer3D/.gitignore: -------------------------------------------------------------------------------- 1 | /build/* 2 | /chamfer_3D.egg-info/* 3 | /dist/* 4 | /__pycache__/* -------------------------------------------------------------------------------- /utils/metrics/CD/chamfer3D/chamfer_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | ///TMP 5 | //#include "common.h" 6 | /// NOT TMP 7 | 8 | 9 | int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2); 10 | 11 | 12 | int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2); 13 | 14 | 15 | 16 | 17 | int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2) { 18 | return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2); 19 | } 20 | 21 | 22 | int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, 23 | at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) { 24 | 25 | return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2); 26 | } 27 | 28 | 29 | 30 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 31 | m.def("forward", &chamfer_forward, "chamfer forward (CUDA)"); 32 | m.def("backward", &chamfer_backward, "chamfer backward (CUDA)"); 33 | } -------------------------------------------------------------------------------- /utils/metrics/CD/chamfer3D/dist_chamfer_3D.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from torch.autograd import Function 3 | import torch 4 | import importlib 5 | import os 6 | # chamfer_found = importlib.find_loader("chamfer_3D") is not None 7 | # if not chamfer_found: 8 | ## Cool trick from https://github.com/chrdiller 9 | print("Jitting Chamfer 3D") 10 | 11 | from torch.utils.cpp_extension import load 12 | chamfer_3D = load(name="chamfer_3D", 13 | sources=[ 14 | "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer_cuda.cpp"]), 15 | "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer3D.cu"]), 16 | ]) 17 | print("Loaded JIT 3D CUDA chamfer distance") 18 | 19 | # else: 20 | # import chamfer_3D 21 | # print("Loaded compiled 3D CUDA chamfer distance") 22 | 23 | 24 | # Chamfer's distance module @thibaultgroueix 25 | # GPU tensors only 26 | class chamfer_3DFunction(Function): 27 | @staticmethod 28 | def forward(ctx, xyz1, xyz2): 29 | batchsize, n, _ = xyz1.size() 30 | _, m, _ = xyz2.size() 31 | device = xyz1.device 32 | 33 | dist1 = torch.zeros(batchsize, n) 34 | dist2 = torch.zeros(batchsize, m) 35 | 36 | idx1 = torch.zeros(batchsize, n).type(torch.IntTensor) 37 | idx2 = torch.zeros(batchsize, m).type(torch.IntTensor) 38 | 39 | dist1 = dist1.to(device) 40 | dist2 = dist2.to(device) 41 | idx1 = idx1.to(device) 42 | idx2 = idx2.to(device) 43 | torch.cuda.set_device(device) 44 | 45 | chamfer_3D.forward(xyz1, xyz2, dist1, dist2, idx1, idx2) 46 | ctx.save_for_backward(xyz1, xyz2, idx1, idx2) 47 | return dist1, dist2, idx1, idx2 48 | 49 | @staticmethod 50 | def backward(ctx, graddist1, graddist2, gradidx1, gradidx2): 51 | xyz1, xyz2, idx1, idx2 = ctx.saved_tensors 52 | graddist1 = graddist1.contiguous() 53 | graddist2 = graddist2.contiguous() 54 | device = graddist1.device 55 | 56 | gradxyz1 = torch.zeros(xyz1.size()) 57 | gradxyz2 = torch.zeros(xyz2.size()) 58 | 59 | gradxyz1 = gradxyz1.to(device) 60 | gradxyz2 = gradxyz2.to(device) 61 | chamfer_3D.backward( 62 | xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2 63 | ) 64 | return gradxyz1, gradxyz2 65 | 66 | 67 | class chamfer_3DDist(nn.Module): 68 | def __init__(self): 69 | super(chamfer_3DDist, self).__init__() 70 | 71 | def forward(self, input1, input2): 72 | input1 = input1.contiguous() 73 | input2 = input2.contiguous() 74 | return chamfer_3DFunction.apply(input1, input2) -------------------------------------------------------------------------------- /utils/metrics/CD/chamfer3D/run_build.sh: -------------------------------------------------------------------------------- 1 | job_name=build 2 | gpus=1 3 | 4 | srun -p dsta --mpi=pmi2 --gres=gpu:${gpus} -n1 \ 5 | --ntasks-per-node=1 --job-name=${job_name}$\ 6 | --kill-on-bad-exit=1 -w SG-IDC1-10-51-2-38 \ 7 | python setup.py install 8 | -------------------------------------------------------------------------------- /utils/metrics/CD/chamfer3D/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 3 | 4 | setup( 5 | name='chamfer_3D', 6 | ext_modules=[ 7 | CUDAExtension('chamfer_3D', [ 8 | "/".join(__file__.split('/')[:-1] + ['chamfer_cuda.cpp']), 9 | "/".join(__file__.split('/')[:-1] + ['chamfer3D.cu']), 10 | ]), 11 | ], 12 | cmdclass={ 13 | 'build_ext': BuildExtension 14 | }) -------------------------------------------------------------------------------- /utils/metrics/CD/chamfer_python.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def pairwise_dist(x, y): 5 | xx, yy, zz = torch.mm(x, x.t()), torch.mm(y, y.t()), torch.mm(x, y.t()) 6 | rx = xx.diag().unsqueeze(0).expand_as(xx) 7 | ry = yy.diag().unsqueeze(0).expand_as(yy) 8 | P = rx.t() + ry - 2 * zz 9 | return P 10 | 11 | 12 | def NN_loss(x, y, dim=0): 13 | dist = pairwise_dist(x, y) 14 | values, indices = dist.min(dim=dim) 15 | return values.mean() 16 | 17 | 18 | def distChamfer(a, b): 19 | """ 20 | :param a: Pointclouds Batch x nul_points x dim 21 | :param b: Pointclouds Batch x nul_points x dim 22 | :return: 23 | -closest point on b of points from a 24 | -closest point on a of points from b 25 | -idx of closest point on b of points from a 26 | -idx of closest point on a of points from b 27 | Works for pointcloud of any dimension 28 | """ 29 | x, y = a.double(), b.double() 30 | bs, num_points_x, points_dim = x.size() 31 | bs, num_points_y, points_dim = y.size() 32 | 33 | xx = torch.pow(x, 2).sum(2) 34 | yy = torch.pow(y, 2).sum(2) 35 | zz = torch.bmm(x, y.transpose(2, 1)) 36 | rx = xx.unsqueeze(1).expand(bs, num_points_y, num_points_x) # Diagonal elements xx 37 | ry = yy.unsqueeze(1).expand(bs, num_points_x, num_points_y) # Diagonal elements yy 38 | P = rx.transpose(2, 1) + ry - 2 * zz 39 | return torch.min(P, 2)[0].float(), torch.min(P, 1)[0].float(), torch.min(P, 2)[1].int(), torch.min(P, 1)[1].int() 40 | 41 | -------------------------------------------------------------------------------- /utils/metrics/CD/fscore.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | def fscore(dist1, dist2, threshold=0.0001): 4 | """ 5 | Calculates the F-score between two point clouds with the corresponding threshold value. 6 | :param dist1: Batch, N-Points 7 | :param dist2: Batch, N-Points 8 | :param th: float 9 | :return: fscore, precision, recall 10 | """ 11 | # NB : In this depo, dist1 and dist2 are squared pointcloud euclidean distances, so you should adapt the threshold accordingly. 12 | precision_1 = torch.mean((dist1 < threshold).float(), dim=1) 13 | precision_2 = torch.mean((dist2 < threshold).float(), dim=1) 14 | fscore = 2 * precision_1 * precision_2 / (precision_1 + precision_2) 15 | fscore[torch.isnan(fscore)] = 0 16 | return fscore, precision_1, precision_2 17 | 18 | -------------------------------------------------------------------------------- /utils/metrics/CD/unit_test.py: -------------------------------------------------------------------------------- 1 | import torch, time 2 | import chamfer2D.dist_chamfer_2D 3 | import chamfer3D.dist_chamfer_3D 4 | import chamfer5D.dist_chamfer_5D 5 | import chamfer_python 6 | 7 | cham2D = chamfer2D.dist_chamfer_2D.chamfer_2DDist() 8 | cham3D = chamfer3D.dist_chamfer_3D.chamfer_3DDist() 9 | cham5D = chamfer5D.dist_chamfer_5D.chamfer_5DDist() 10 | 11 | from torch.autograd import Variable 12 | from fscore import fscore 13 | 14 | def test_chamfer(distChamfer, dim): 15 | points1 = torch.rand(4, 100, dim).cuda() 16 | points2 = torch.rand(4, 200, dim, requires_grad=True).cuda() 17 | dist1, dist2, idx1, idx2= distChamfer(points1, points2) 18 | 19 | loss = torch.sum(dist1) 20 | loss.backward() 21 | 22 | mydist1, mydist2, myidx1, myidx2 = chamfer_python.distChamfer(points1, points2) 23 | d1 = (dist1 - mydist1) ** 2 24 | d2 = (dist2 - mydist2) ** 2 25 | assert ( 26 | torch.mean(d1) + torch.mean(d2) < 0.00000001 27 | ), "chamfer cuda and chamfer normal are not giving the same results" 28 | 29 | xd1 = idx1 - myidx1 30 | xd2 = idx2 - myidx2 31 | assert ( 32 | torch.norm(xd1.float()) + torch.norm(xd2.float()) == 0 33 | ), "chamfer cuda and chamfer normal are not giving the same results" 34 | print(f"fscore :", fscore(dist1, dist2)) 35 | print("Unit test passed") 36 | 37 | 38 | def timings(distChamfer, dim): 39 | p1 = torch.rand(32, 2000, dim).cuda() 40 | p2 = torch.rand(32, 1000, dim).cuda() 41 | print("Timings : Start CUDA version") 42 | start = time.time() 43 | num_it = 100 44 | for i in range(num_it): 45 | points1 = Variable(p1, requires_grad=True) 46 | points2 = Variable(p2) 47 | mydist1, mydist2, idx1, idx2 = distChamfer(points1, points2) 48 | loss = torch.sum(mydist1) 49 | loss.backward() 50 | print(f"Ellapsed time forward backward is {(time.time() - start)/num_it} seconds.") 51 | 52 | 53 | print("Timings : Start Pythonic version") 54 | start = time.time() 55 | for i in range(num_it): 56 | points1 = Variable(p1, requires_grad=True) 57 | points2 = Variable(p2) 58 | mydist1, mydist2, idx1, idx2 = chamfer_python.distChamfer(points1, points2) 59 | loss = torch.sum(mydist1) 60 | loss.backward() 61 | print(f"Ellapsed time forward backward is {(time.time() - start)/num_it} seconds.") 62 | 63 | 64 | 65 | dims = [2,3,5] 66 | for i,cham in enumerate([cham2D, cham3D, cham5D]): 67 | print(f"testing Chamfer {dims[i]}D") 68 | test_chamfer(cham, dims[i]) 69 | timings(cham, dims[i]) 70 | -------------------------------------------------------------------------------- /utils/metrics/EMD/.gitignore: -------------------------------------------------------------------------------- 1 | /build/* 2 | /emd.egg-info/* 3 | /dist/* 4 | /__pycache__/* -------------------------------------------------------------------------------- /utils/metrics/EMD/CDEMD.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/paul007pl/MVP_Benchmark/36987fe2ceac95e8ae89f2ee6c5e8003fb55f87d/utils/metrics/EMD/CDEMD.png -------------------------------------------------------------------------------- /utils/metrics/EMD/README.md: -------------------------------------------------------------------------------- 1 | # Earth Mover's Distance (EMD). 2 | 3 | ### What is the Earth Mover's Distance ? 4 | 5 | Compared to the Chamfer Distance (CD), the Earth Mover's Distance (EMD) is more reliable to distinguish the visual quality of the point clouds. See our [paper](http://cseweb.ucsd.edu/~mil070/projects/AAAI2020/paper.pdf) for more details. 6 | 7 | The provided EMD implementation for point cloud comparison only needs $O(n)$ memory and thus enables dense point clouds (with 10,000 points or over) and large batch size. It is based on an approximated algorithm (auction algorithm) and cannot guarantee a (but near) bijection assignment. It employs a parameter $\epsilon$ to balance the error rate and the speed of convergence. Smaller $\epsilon$ achieves more accurate results, but needs a longer time for convergence. The time complexity is $O(n^2k)$, where $k$ is the number of iterations. We set a $\epsilon = 0.005, k = 50$ during training and a $\epsilon = 0.002, k = 10000$ during testing. 8 | 9 | ### Compile 10 | 11 | Use JIT cool trick, and it is not necessary to compile. 12 | 13 | ### Usage 14 | See `emd_module.py/test_emd()` for examples. 15 | 16 | **Input** 17 | 18 | - **xyz1, xyz2**: float tensors with shape `[#batch, #points, 3]`. xyz1 is the predicted point cloud and xyz2 is the ground truth point cloud. Two point clouds should have same size and be normalized to [0, 1]. The number of points should be a multiple of 1024. The batch size should be no greater than 512. Since we only calculate gradients for xyz1, please do not swap xyz1 and xyz2. 19 | - **eps**: a float tensor, the parameter balances the error rate and the speed of convergence. 20 | - **iters**: a int tensor, the number of iterations. 21 | 22 | **Output** 23 | 24 | - **dist**: a float tensor with shape `[#batch, #points]`. sqrt(dist) are the L2 distances between the pairs of points. 25 | - **assignment**: a int tensor with shape `[#batch, #points]`. The index of the matched point in the ground truth point cloud. 26 | 27 | 28 | ### Aknowledgment 29 | 30 | Modify from [Minghua LIU](https://github.com/Colin97/MSN-Point-Cloud-Completion/tree/master/emd) -------------------------------------------------------------------------------- /utils/metrics/EMD/__init__.py: -------------------------------------------------------------------------------- 1 | from .emd_module import emdModule as emd 2 | 3 | __all__ = ['emd'] -------------------------------------------------------------------------------- /utils/metrics/EMD/clean.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | rm -rf __pycache__/ 4 | rm -rf dist/ 5 | rm -rf build/ 6 | rm -rf emd.egg-info/ 7 | rm -rf /mnt/lustre/chenxinyi1/.conda/envs/pt/lib/python3.7/site-packages/emd-0.0.0-py3.7-linux-x86_64.egg/ 8 | -------------------------------------------------------------------------------- /utils/metrics/EMD/emd.cpp: -------------------------------------------------------------------------------- 1 | // EMD approximation module (based on auction algorithm) 2 | // author: Minghua Liu 3 | #include 4 | #include 5 | 6 | int emd_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price, 7 | at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments, 8 | at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters); 9 | 10 | int emd_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx); 11 | 12 | 13 | 14 | int emd_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price, 15 | at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments, 16 | at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters) { 17 | return emd_cuda_forward(xyz1, xyz2, dist, assignment, price, assignment_inv, bid, bid_increments, max_increments, unass_idx, unass_cnt, unass_cnt_sum, cnt_tmp, max_idx, eps, iters); 18 | } 19 | 20 | int emd_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx) { 21 | 22 | return emd_cuda_backward(xyz1, xyz2, gradxyz, graddist, idx); 23 | } 24 | 25 | 26 | 27 | 28 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 29 | m.def("forward", &emd_forward, "emd forward (CUDA)"); 30 | m.def("backward", &emd_backward, "emd backward (CUDA)"); 31 | } -------------------------------------------------------------------------------- /utils/metrics/EMD/emd_module.py: -------------------------------------------------------------------------------- 1 | # EMD approximation module (based on auction algorithm) 2 | # memory complexity: O(n) 3 | # time complexity: O(n^2 * iter) 4 | # author: Minghua Liu 5 | 6 | # Input: 7 | # xyz1, xyz2: [#batch, #points, 3] 8 | # where xyz1 is the predicted point cloud and xyz2 is the ground truth point cloud 9 | # two point clouds should have same size and be normalized to [0, 1] 10 | # #points should be a multiple of 1024 11 | # #batch should be no greater than 512 12 | # eps is a parameter which balances the error rate and the speed of convergence 13 | # iters is the number of iteration 14 | # we only calculate gradient for xyz1 15 | 16 | # Output: 17 | # dist: [#batch, #points], sqrt(dist) -> L2 distance 18 | # assignment: [#batch, #points], index of the matched point in the ground truth point cloud 19 | # the result is an approximation and the assignment is not guranteed to be a bijection 20 | 21 | import time 22 | import numpy as np 23 | import torch 24 | from torch import nn 25 | from torch.autograd import Function 26 | 27 | import os 28 | # import importlib 29 | # import emd 30 | 31 | from torch.utils.cpp_extension import load 32 | emd = load(name="emd", 33 | sources=[ 34 | "/".join(os.path.abspath(__file__).split('/')[:-1] + ["emd.cpp"]), 35 | "/".join(os.path.abspath(__file__).split('/')[:-1] + ["emd_cuda.cu"]), 36 | ]) 37 | print("Loaded JIT 3D CUDA emd") 38 | 39 | 40 | class emdFunction(Function): 41 | @staticmethod 42 | def forward(ctx, xyz1, xyz2, eps, iters): 43 | 44 | batchsize, n, _ = xyz1.size() 45 | _, m, _ = xyz2.size() 46 | 47 | assert(n == m) 48 | assert(xyz1.size()[0] == xyz2.size()[0]) 49 | #assert(n % 1024 == 0) 50 | assert(batchsize <= 512) 51 | 52 | xyz1 = xyz1.contiguous().float().cuda() 53 | xyz2 = xyz2.contiguous().float().cuda() 54 | dist = torch.zeros(batchsize, n, device='cuda').contiguous() 55 | assignment = torch.zeros(batchsize, n, device='cuda', dtype=torch.int32).contiguous() - 1 56 | assignment_inv = torch.zeros(batchsize, m, device='cuda', dtype=torch.int32).contiguous() - 1 57 | price = torch.zeros(batchsize, m, device='cuda').contiguous() 58 | bid = torch.zeros(batchsize, n, device='cuda', dtype=torch.int32).contiguous() 59 | bid_increments = torch.zeros(batchsize, n, device='cuda').contiguous() 60 | max_increments = torch.zeros(batchsize, m, device='cuda').contiguous() 61 | unass_idx = torch.zeros(batchsize * n, device='cuda', dtype=torch.int32).contiguous() 62 | max_idx = torch.zeros(batchsize * m, device='cuda', dtype=torch.int32).contiguous() 63 | unass_cnt = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous() 64 | unass_cnt_sum = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous() 65 | cnt_tmp = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous() 66 | 67 | emd.forward(xyz1, xyz2, dist, assignment, price, assignment_inv, bid, bid_increments, max_increments, unass_idx, unass_cnt, unass_cnt_sum, cnt_tmp, max_idx, eps, iters) 68 | 69 | ctx.save_for_backward(xyz1, xyz2, assignment) 70 | return dist, assignment 71 | 72 | @staticmethod 73 | def backward(ctx, graddist, gradidx): 74 | xyz1, xyz2, assignment = ctx.saved_tensors 75 | graddist = graddist.contiguous() 76 | 77 | gradxyz1 = torch.zeros(xyz1.size(), device='cuda').contiguous() 78 | gradxyz2 = torch.zeros(xyz2.size(), device='cuda').contiguous() 79 | 80 | emd.backward(xyz1, xyz2, gradxyz1, graddist, assignment) 81 | return gradxyz1, gradxyz2, None, None 82 | 83 | class emdModule(nn.Module): 84 | def __init__(self): 85 | super(emdModule, self).__init__() 86 | 87 | def forward(self, input1, input2, eps, iters): 88 | return emdFunction.apply(input1, input2, eps, iters) 89 | 90 | def test_emd(): 91 | x1 = torch.rand(20, 8192, 3).cuda() 92 | x2 = torch.rand(20, 8192, 3).cuda() 93 | emd = emdModule() 94 | start_time = time.perf_counter() 95 | dis, assigment = emd(x1, x2, 0.05, 3000) 96 | print("Input_size: ", x1.shape) 97 | print("Runtime: %lfs" % (time.perf_counter() - start_time)) 98 | print("EMD: %lf" % np.sqrt(dis.cpu()).mean()) 99 | print("|set(assignment)|: %d" % assigment.unique().numel()) 100 | assigment = assigment.cpu().numpy() 101 | assigment = np.expand_dims(assigment, -1) 102 | x2 = np.take_along_axis(x2, assigment, axis = 1) 103 | d = (x1 - x2) * (x1 - x2) 104 | print("Verified EMD: %lf" % np.sqrt(d.cpu().sum(-1)).mean()) 105 | 106 | # test_emd() 107 | 108 | 109 | -------------------------------------------------------------------------------- /utils/metrics/EMD/run_build.sh: -------------------------------------------------------------------------------- 1 | job_name=build 2 | gpus=1 3 | 4 | srun -p dsta --mpi=pmi2 --gres=gpu:${gpus} -n1 \ 5 | --ntasks-per-node=1 --job-name=${job_name}$\ 6 | --kill-on-bad-exit=1 -w SG-IDC1-10-51-2-38 \ 7 | python setup.py install 8 | -------------------------------------------------------------------------------- /utils/metrics/EMD/run_compile.sh: -------------------------------------------------------------------------------- 1 | partition=ips_share 2 | job_name=compile 3 | gpus=1 4 | g=$((${gpus}<8?${gpus}:8)) 5 | 6 | 7 | srun -u --partition=${partition} --job-name=${job_name} \ 8 | -n1 --gres=gpu:${gpus} --ntasks-per-node=1 -w 'SH-IDC1-10-198-6-85' \ 9 | python3 emd_module.py 10 | -------------------------------------------------------------------------------- /utils/metrics/EMD/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 3 | 4 | setup( 5 | name='emd', 6 | ext_modules=[ 7 | CUDAExtension('emd', [ 8 | 'emd.cpp', 9 | 'emd_cuda.cu', 10 | ]), 11 | ], 12 | cmdclass={ 13 | 'build_ext': BuildExtension 14 | }) -------------------------------------------------------------------------------- /utils/metrics/__init__.py: -------------------------------------------------------------------------------- 1 | from .CD import (cd, fscore) 2 | from .EMD import emd 3 | 4 | __all__ = [ 5 | 'cd', 'fscore', 'emd', 6 | ] -------------------------------------------------------------------------------- /utils/mm3d_pn2/.gitignore: -------------------------------------------------------------------------------- 1 | /__pycache__/* 2 | /build/* 3 | /mmdet3d.egg-info/* -------------------------------------------------------------------------------- /utils/mm3d_pn2/__init__.py: -------------------------------------------------------------------------------- 1 | from .ops import (nms, RoIAlign, roi_align, get_compiler_version, get_compiling_cuda_version, 2 | NaiveSyncBatchNorm1d, NaiveSyncBatchNorm2d, sigmoid_focal_loss, SigmoidFocalLoss, ball_query, knn, 3 | furthest_point_sample, furthest_point_sample_with_dist, three_interpolate, three_nn, gather_points, 4 | grouping_operation, group_points, GroupAll, QueryAndGroup, get_compiler_version, get_compiling_cuda_version, 5 | Points_Sampler) 6 | 7 | __all__=[ 8 | 'nms', 9 | 'RoIAlign', 'roi_align', 'get_compiler_version', 10 | 'get_compiling_cuda_version', 'NaiveSyncBatchNorm1d', 11 | 'NaiveSyncBatchNorm2d', 12 | 'sigmoid_focal_loss', 13 | 'SigmoidFocalLoss', 14 | 'ball_query', 'knn', 'furthest_point_sample', 15 | 'furthest_point_sample_with_dist', 'three_interpolate', 'three_nn', 16 | 'gather_points', 'grouping_operation', 'group_points', 'GroupAll', 17 | 'QueryAndGroup', 18 | 'get_compiler_version', 19 | 'get_compiling_cuda_version', 'Points_Sampler', 20 | ] -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/__init__.py: -------------------------------------------------------------------------------- 1 | from mmcv.ops import (RoIAlign, SigmoidFocalLoss, get_compiler_version, 2 | get_compiling_cuda_version, nms, roi_align, 3 | sigmoid_focal_loss) 4 | 5 | from .ball_query import ball_query 6 | from .furthest_point_sample import (Points_Sampler, furthest_point_sample, 7 | furthest_point_sample_with_dist) 8 | from .gather_points import gather_points 9 | from .group_points import (GroupAll, QueryAndGroup, group_points, 10 | grouping_operation) 11 | from .interpolate import three_interpolate, three_nn 12 | from .knn import knn 13 | from .norm import NaiveSyncBatchNorm1d, NaiveSyncBatchNorm2d 14 | # from .paconv import PAConv, PAConvCUDA, assign_score_withk 15 | # from .pointnet_modules import (PointFPModule, PointSAModule, PointSAModuleMSG, 16 | # build_sa_module) 17 | # from .roiaware_pool3d import (RoIAwarePool3d, points_in_boxes_batch, 18 | # points_in_boxes_cpu, points_in_boxes_gpu) 19 | # from .sparse_block import (SparseBasicBlock, SparseBottleneck, 20 | # make_sparse_convmodule) 21 | # from .voxel import DynamicScatter, Voxelization, dynamic_scatter, voxelization 22 | 23 | __all__ = [ 24 | 'nms', 25 | # 'soft_nms', 26 | 'RoIAlign', 'roi_align', 'get_compiler_version', 27 | 'get_compiling_cuda_version', 'NaiveSyncBatchNorm1d', 28 | 'NaiveSyncBatchNorm2d', 29 | # 'batched_nms', 'Voxelization', 'voxelization', 30 | # 'dynamic_scatter', 'DynamicScatter', 31 | 'sigmoid_focal_loss', 32 | 'SigmoidFocalLoss', 33 | # 'SparseBasicBlock', 'SparseBottleneck', 34 | # 'RoIAwarePool3d', 'points_in_boxes_gpu', 'points_in_boxes_cpu', 35 | # 'make_sparse_convmodule', 36 | 'ball_query', 'knn', 'furthest_point_sample', 37 | 'furthest_point_sample_with_dist', 'three_interpolate', 'three_nn', 38 | 'gather_points', 'grouping_operation', 'group_points', 'GroupAll', 39 | 'QueryAndGroup', 40 | # 'PointSAModule', 'PointSAModuleMSG', 'PointFPModule', 41 | # 'points_in_boxes_batch', 42 | 'get_compiler_version', 43 | # 'assign_score_withk', 44 | 'get_compiling_cuda_version', 'Points_Sampler', 45 | # 'build_sa_module', 46 | # 'PAConv', 'PAConvCUDA' 47 | ] 48 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/ball_query/__init__.py: -------------------------------------------------------------------------------- 1 | from .ball_query import ball_query 2 | 3 | __all__ = ['ball_query'] 4 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/ball_query/ball_query.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Function 3 | 4 | from . import ball_query_ext 5 | 6 | 7 | class BallQuery(Function): 8 | """Ball Query. 9 | 10 | Find nearby points in spherical space. 11 | """ 12 | 13 | @staticmethod 14 | def forward(ctx, min_radius: float, max_radius: float, sample_num: int, 15 | xyz: torch.Tensor, center_xyz: torch.Tensor) -> torch.Tensor: 16 | """forward. 17 | 18 | Args: 19 | min_radius (float): minimum radius of the balls. 20 | max_radius (float): maximum radius of the balls. 21 | sample_num (int): maximum number of features in the balls. 22 | xyz (Tensor): (B, N, 3) xyz coordinates of the features. 23 | center_xyz (Tensor): (B, npoint, 3) centers of the ball query. 24 | 25 | Returns: 26 | Tensor: (B, npoint, nsample) tensor with the indicies of 27 | the features that form the query balls. 28 | """ 29 | assert center_xyz.is_contiguous() 30 | assert xyz.is_contiguous() 31 | assert min_radius < max_radius 32 | 33 | B, N, _ = xyz.size() 34 | npoint = center_xyz.size(1) 35 | idx = torch.cuda.IntTensor(B, npoint, sample_num).zero_() 36 | 37 | ball_query_ext.ball_query_wrapper(B, N, npoint, min_radius, max_radius, 38 | sample_num, center_xyz, xyz, idx) 39 | ctx.mark_non_differentiable(idx) 40 | return idx 41 | 42 | @staticmethod 43 | def backward(ctx, a=None): 44 | return None, None, None, None 45 | 46 | 47 | ball_query = BallQuery.apply 48 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/ball_query/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/ball_query.cpp 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | #include 11 | 12 | extern THCState *state; 13 | 14 | #define CHECK_CUDA(x) \ 15 | TORCH_CHECK(x.type().is_cuda(), #x, " must be a CUDAtensor ") 16 | #define CHECK_CONTIGUOUS(x) \ 17 | TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ") 18 | #define CHECK_INPUT(x) \ 19 | CHECK_CUDA(x); \ 20 | CHECK_CONTIGUOUS(x) 21 | 22 | int ball_query_wrapper(int b, int n, int m, float min_radius, float max_radius, int nsample, 23 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, 24 | at::Tensor idx_tensor); 25 | 26 | void ball_query_kernel_launcher(int b, int n, int m, float min_radius, float max_radius, 27 | int nsample, const float *xyz, const float *new_xyz, 28 | int *idx, cudaStream_t stream); 29 | 30 | int ball_query_wrapper(int b, int n, int m, float min_radius, float max_radius, int nsample, 31 | at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, 32 | at::Tensor idx_tensor) { 33 | CHECK_INPUT(new_xyz_tensor); 34 | CHECK_INPUT(xyz_tensor); 35 | const float *new_xyz = new_xyz_tensor.data_ptr(); 36 | const float *xyz = xyz_tensor.data_ptr(); 37 | int *idx = idx_tensor.data_ptr(); 38 | 39 | cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 40 | ball_query_kernel_launcher(b, n, m, min_radius, max_radius, 41 | nsample, new_xyz, xyz, idx, stream); 42 | return 1; 43 | } 44 | 45 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 46 | m.def("ball_query_wrapper", &ball_query_wrapper, "ball_query_wrapper"); 47 | } 48 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/ball_query/src/ball_query_cuda.cu: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/ball_query_gpu.cu 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #define THREADS_PER_BLOCK 256 9 | #define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0)) 10 | 11 | __global__ void ball_query_kernel(int b, int n, int m, 12 | float min_radius, 13 | float max_radius, 14 | int nsample, 15 | const float *__restrict__ new_xyz, 16 | const float *__restrict__ xyz, 17 | int *__restrict__ idx) { 18 | // new_xyz: (B, M, 3) 19 | // xyz: (B, N, 3) 20 | // output: 21 | // idx: (B, M, nsample) 22 | int bs_idx = blockIdx.y; 23 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 24 | if (bs_idx >= b || pt_idx >= m) return; 25 | 26 | new_xyz += bs_idx * m * 3 + pt_idx * 3; 27 | xyz += bs_idx * n * 3; 28 | idx += bs_idx * m * nsample + pt_idx * nsample; 29 | 30 | float max_radius2 = max_radius * max_radius; 31 | float min_radius2 = min_radius * min_radius; 32 | float new_x = new_xyz[0]; 33 | float new_y = new_xyz[1]; 34 | float new_z = new_xyz[2]; 35 | 36 | int cnt = 0; 37 | for (int k = 0; k < n; ++k) { 38 | float x = xyz[k * 3 + 0]; 39 | float y = xyz[k * 3 + 1]; 40 | float z = xyz[k * 3 + 2]; 41 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 42 | (new_z - z) * (new_z - z); 43 | if (d2 == 0 || (d2 >= min_radius2 && d2 < max_radius2)) { 44 | if (cnt == 0) { 45 | for (int l = 0; l < nsample; ++l) { 46 | idx[l] = k; 47 | } 48 | } 49 | idx[cnt] = k; 50 | ++cnt; 51 | if (cnt >= nsample) break; 52 | } 53 | } 54 | } 55 | 56 | void ball_query_kernel_launcher(int b, int n, int m, float min_radius, float max_radius, 57 | int nsample, const float *new_xyz, const float *xyz, 58 | int *idx, cudaStream_t stream) { 59 | // new_xyz: (B, M, 3) 60 | // xyz: (B, N, 3) 61 | // output: 62 | // idx: (B, M, nsample) 63 | 64 | cudaError_t err; 65 | 66 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), 67 | b); // blockIdx.x(col), blockIdx.y(row) 68 | dim3 threads(THREADS_PER_BLOCK); 69 | 70 | ball_query_kernel<<>>(b, n, m, min_radius, max_radius, 71 | nsample, new_xyz, xyz, idx); 72 | // cudaDeviceSynchronize(); // for using printf in kernel function 73 | err = cudaGetLastError(); 74 | if (cudaSuccess != err) { 75 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 76 | exit(-1); 77 | } 78 | } 79 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/furthest_point_sample/__init__.py: -------------------------------------------------------------------------------- 1 | from .furthest_point_sample import (furthest_point_sample, 2 | furthest_point_sample_with_dist) 3 | from .points_sampler import Points_Sampler 4 | 5 | __all__ = [ 6 | 'furthest_point_sample', 'furthest_point_sample_with_dist', 7 | 'Points_Sampler' 8 | ] 9 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/furthest_point_sample/furthest_point_sample.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Function 3 | 4 | from . import furthest_point_sample_ext 5 | 6 | 7 | class FurthestPointSampling(Function): 8 | """Furthest Point Sampling. 9 | 10 | Uses iterative furthest point sampling to select a set of features whose 11 | corresponding points have the furthest distance. 12 | """ 13 | 14 | @staticmethod 15 | def forward(ctx, points_xyz: torch.Tensor, 16 | num_points: int) -> torch.Tensor: 17 | """forward. 18 | 19 | Args: 20 | points_xyz (Tensor): (B, N, 3) where N > num_points. 21 | num_points (int): Number of points in the sampled set. 22 | 23 | Returns: 24 | Tensor: (B, num_points) indices of the sampled points. 25 | """ 26 | assert points_xyz.is_contiguous() 27 | 28 | B, N = points_xyz.size()[:2] 29 | output = torch.cuda.IntTensor(B, num_points) 30 | temp = torch.cuda.FloatTensor(B, N).fill_(1e10) 31 | 32 | furthest_point_sample_ext.furthest_point_sampling_wrapper( 33 | B, N, num_points, points_xyz, temp, output) 34 | ctx.mark_non_differentiable(output) 35 | return output 36 | 37 | @staticmethod 38 | def backward(xyz, a=None): 39 | return None, None 40 | 41 | 42 | class FurthestPointSamplingWithDist(Function): 43 | """Furthest Point Sampling With Distance. 44 | 45 | Uses iterative furthest point sampling to select a set of features whose 46 | corresponding points have the furthest distance. 47 | """ 48 | 49 | @staticmethod 50 | def forward(ctx, points_dist: torch.Tensor, 51 | num_points: int) -> torch.Tensor: 52 | """forward. 53 | 54 | Args: 55 | points_dist (Tensor): (B, N, N) Distance between each point pair. 56 | num_points (int): Number of points in the sampled set. 57 | 58 | Returns: 59 | Tensor: (B, num_points) indices of the sampled points. 60 | """ 61 | assert points_dist.is_contiguous() 62 | 63 | B, N, _ = points_dist.size() 64 | output = points_dist.new_zeros([B, num_points], dtype=torch.int32) 65 | temp = points_dist.new_zeros([B, N]).fill_(1e10) 66 | 67 | furthest_point_sample_ext.furthest_point_sampling_with_dist_wrapper( 68 | B, N, num_points, points_dist, temp, output) 69 | ctx.mark_non_differentiable(output) 70 | return output 71 | 72 | @staticmethod 73 | def backward(xyz, a=None): 74 | return None, None 75 | 76 | 77 | furthest_point_sample = FurthestPointSampling.apply 78 | furthest_point_sample_with_dist = FurthestPointSamplingWithDist.apply 79 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/furthest_point_sample/src/furthest_point_sample.cpp: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/sampling.cpp 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | #include 10 | 11 | extern THCState *state; 12 | 13 | int furthest_point_sampling_wrapper(int b, int n, int m, 14 | at::Tensor points_tensor, 15 | at::Tensor temp_tensor, 16 | at::Tensor idx_tensor); 17 | 18 | void furthest_point_sampling_kernel_launcher(int b, int n, int m, 19 | const float *dataset, float *temp, 20 | int *idxs, cudaStream_t stream); 21 | 22 | int furthest_point_sampling_with_dist_wrapper(int b, int n, int m, 23 | at::Tensor points_tensor, 24 | at::Tensor temp_tensor, 25 | at::Tensor idx_tensor); 26 | 27 | void furthest_point_sampling_with_dist_kernel_launcher(int b, int n, int m, 28 | const float *dataset, 29 | float *temp, int *idxs, 30 | cudaStream_t stream); 31 | 32 | int furthest_point_sampling_wrapper(int b, int n, int m, 33 | at::Tensor points_tensor, 34 | at::Tensor temp_tensor, 35 | at::Tensor idx_tensor) { 36 | const float *points = points_tensor.data_ptr(); 37 | float *temp = temp_tensor.data_ptr(); 38 | int *idx = idx_tensor.data_ptr(); 39 | 40 | cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 41 | furthest_point_sampling_kernel_launcher(b, n, m, points, temp, idx, stream); 42 | return 1; 43 | } 44 | 45 | int furthest_point_sampling_with_dist_wrapper(int b, int n, int m, 46 | at::Tensor points_tensor, 47 | at::Tensor temp_tensor, 48 | at::Tensor idx_tensor) { 49 | 50 | const float *points = points_tensor.data(); 51 | float *temp = temp_tensor.data(); 52 | int *idx = idx_tensor.data(); 53 | 54 | cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 55 | furthest_point_sampling_with_dist_kernel_launcher(b, n, m, points, temp, idx, stream); 56 | return 1; 57 | } 58 | 59 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 60 | m.def("furthest_point_sampling_wrapper", &furthest_point_sampling_wrapper, 61 | "furthest_point_sampling_wrapper"); 62 | m.def("furthest_point_sampling_with_dist_wrapper", 63 | &furthest_point_sampling_with_dist_wrapper, 64 | "furthest_point_sampling_with_dist_wrapper"); 65 | } 66 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/furthest_point_sample/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def calc_square_dist(point_feat_a, point_feat_b, norm=True): 5 | """Calculating square distance between a and b. 6 | 7 | Args: 8 | point_feat_a (Tensor): (B, N, C) Feature vector of each point. 9 | point_feat_b (Tensor): (B, M, C) Feature vector of each point. 10 | norm (Bool): Whether to normalize the distance. 11 | Default: True. 12 | 13 | Returns: 14 | Tensor: (B, N, M) Distance between each pair points. 15 | """ 16 | length_a = point_feat_a.shape[1] 17 | length_b = point_feat_b.shape[1] 18 | num_channel = point_feat_a.shape[-1] 19 | # [bs, n, 1] 20 | a_square = torch.sum(point_feat_a.unsqueeze(dim=2).pow(2), dim=-1) 21 | # [bs, 1, m] 22 | b_square = torch.sum(point_feat_b.unsqueeze(dim=1).pow(2), dim=-1) 23 | a_square = a_square.repeat((1, 1, length_b)) # [bs, n, m] 24 | b_square = b_square.repeat((1, length_a, 1)) # [bs, n, m] 25 | 26 | coor = torch.matmul(point_feat_a, point_feat_b.transpose(1, 2)) 27 | 28 | dist = a_square + b_square - 2 * coor 29 | if norm: 30 | dist = torch.sqrt(dist) / num_channel 31 | return dist 32 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/gather_points/__init__.py: -------------------------------------------------------------------------------- 1 | from .gather_points import gather_points 2 | 3 | __all__ = ['gather_points'] 4 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/gather_points/gather_points.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Function 3 | 4 | from . import gather_points_ext 5 | 6 | 7 | class GatherPoints(Function): 8 | """Gather Points. 9 | 10 | Gather points with given index. 11 | """ 12 | 13 | @staticmethod 14 | def forward(ctx, features: torch.Tensor, 15 | indices: torch.Tensor) -> torch.Tensor: 16 | """forward. 17 | 18 | Args: 19 | features (Tensor): (B, C, N) features to gather. 20 | indices (Tensor): (B, M) where M is the number of points. 21 | 22 | Returns: 23 | Tensor: (B, C, M) where M is the number of points. 24 | """ 25 | assert features.is_contiguous() 26 | assert indices.is_contiguous() 27 | 28 | B, npoint = indices.size() 29 | _, C, N = features.size() 30 | output = torch.cuda.FloatTensor(B, C, npoint) 31 | 32 | gather_points_ext.gather_points_wrapper(B, C, N, npoint, features, 33 | indices, output) 34 | 35 | ctx.for_backwards = (indices, C, N) 36 | ctx.mark_non_differentiable(indices) 37 | return output 38 | 39 | @staticmethod 40 | def backward(ctx, grad_out): 41 | idx, C, N = ctx.for_backwards 42 | B, npoint = idx.size() 43 | 44 | grad_features = torch.cuda.FloatTensor(B, C, N).zero_() 45 | grad_out_data = grad_out.data.contiguous() 46 | gather_points_ext.gather_points_grad_wrapper(B, C, N, npoint, 47 | grad_out_data, idx, 48 | grad_features.data) 49 | return grad_features, None 50 | 51 | 52 | gather_points = GatherPoints.apply 53 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/gather_points/src/gather_points.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #include 7 | 8 | extern THCState *state; 9 | 10 | int gather_points_wrapper(int b, int c, int n, int npoints, 11 | at::Tensor points_tensor, at::Tensor idx_tensor, 12 | at::Tensor out_tensor); 13 | 14 | void gather_points_kernel_launcher(int b, int c, int n, int npoints, 15 | const float *points, const int *idx, 16 | float *out, cudaStream_t stream); 17 | 18 | int gather_points_grad_wrapper(int b, int c, int n, int npoints, 19 | at::Tensor grad_out_tensor, 20 | at::Tensor idx_tensor, 21 | at::Tensor grad_points_tensor); 22 | 23 | void gather_points_grad_kernel_launcher(int b, int c, int n, int npoints, 24 | const float *grad_out, const int *idx, 25 | float *grad_points, 26 | cudaStream_t stream); 27 | 28 | int gather_points_wrapper(int b, int c, int n, int npoints, 29 | at::Tensor points_tensor, at::Tensor idx_tensor, 30 | at::Tensor out_tensor) { 31 | const float *points = points_tensor.data_ptr(); 32 | const int *idx = idx_tensor.data_ptr(); 33 | float *out = out_tensor.data_ptr(); 34 | 35 | cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 36 | gather_points_kernel_launcher(b, c, n, npoints, points, idx, out, stream); 37 | return 1; 38 | } 39 | 40 | int gather_points_grad_wrapper(int b, int c, int n, int npoints, 41 | at::Tensor grad_out_tensor, 42 | at::Tensor idx_tensor, 43 | at::Tensor grad_points_tensor) { 44 | const float *grad_out = grad_out_tensor.data_ptr(); 45 | const int *idx = idx_tensor.data_ptr(); 46 | float *grad_points = grad_points_tensor.data_ptr(); 47 | 48 | cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 49 | gather_points_grad_kernel_launcher(b, c, n, npoints, grad_out, idx, 50 | grad_points, stream); 51 | return 1; 52 | } 53 | 54 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 55 | m.def("gather_points_wrapper", &gather_points_wrapper, 56 | "gather_points_wrapper"); 57 | m.def("gather_points_grad_wrapper", &gather_points_grad_wrapper, 58 | "gather_points_grad_wrapper"); 59 | } 60 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/gather_points/src/gather_points_cuda.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #define TOTAL_THREADS 1024 5 | #define THREADS_PER_BLOCK 256 6 | #define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0)) 7 | 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 | // points: (B, C, N) 13 | // idx: (B, M) 14 | // output: 15 | // out: (B, C, M) 16 | 17 | int bs_idx = blockIdx.z; 18 | int c_idx = blockIdx.y; 19 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 20 | if (bs_idx >= b || c_idx >= c || pt_idx >= m) return; 21 | 22 | out += bs_idx * c * m + c_idx * m + pt_idx; 23 | idx += bs_idx * m + pt_idx; 24 | points += bs_idx * c * n + c_idx * n; 25 | out[0] = points[idx[0]]; 26 | } 27 | 28 | void gather_points_kernel_launcher(int b, int c, int n, int npoints, 29 | const float *points, const int *idx, 30 | float *out, cudaStream_t stream) { 31 | // points: (B, C, N) 32 | // idx: (B, npoints) 33 | // output: 34 | // out: (B, C, npoints) 35 | 36 | cudaError_t err; 37 | dim3 blocks(DIVUP(npoints, THREADS_PER_BLOCK), c, 38 | b); // blockIdx.x(col), blockIdx.y(row) 39 | dim3 threads(THREADS_PER_BLOCK); 40 | 41 | gather_points_kernel<<>>(b, c, n, npoints, points, 42 | idx, out); 43 | 44 | err = cudaGetLastError(); 45 | if (cudaSuccess != err) { 46 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 47 | exit(-1); 48 | } 49 | } 50 | 51 | __global__ void gather_points_grad_kernel(int b, int c, int n, int m, 52 | const float *__restrict__ grad_out, 53 | const int *__restrict__ idx, 54 | float *__restrict__ grad_points) { 55 | // grad_out: (B, C, M) 56 | // idx: (B, M) 57 | // output: 58 | // grad_points: (B, C, N) 59 | 60 | int bs_idx = blockIdx.z; 61 | int c_idx = blockIdx.y; 62 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 63 | if (bs_idx >= b || c_idx >= c || pt_idx >= m) return; 64 | 65 | grad_out += bs_idx * c * m + c_idx * m + pt_idx; 66 | idx += bs_idx * m + pt_idx; 67 | grad_points += bs_idx * c * n + c_idx * n; 68 | 69 | atomicAdd(grad_points + idx[0], grad_out[0]); 70 | } 71 | 72 | void gather_points_grad_kernel_launcher(int b, int c, int n, int npoints, 73 | const float *grad_out, const int *idx, 74 | float *grad_points, 75 | cudaStream_t stream) { 76 | // grad_out: (B, C, npoints) 77 | // idx: (B, npoints) 78 | // output: 79 | // grad_points: (B, C, N) 80 | 81 | cudaError_t err; 82 | dim3 blocks(DIVUP(npoints, THREADS_PER_BLOCK), c, 83 | b); // blockIdx.x(col), blockIdx.y(row) 84 | dim3 threads(THREADS_PER_BLOCK); 85 | 86 | gather_points_grad_kernel<<>>( 87 | b, c, n, npoints, grad_out, idx, grad_points); 88 | 89 | err = cudaGetLastError(); 90 | if (cudaSuccess != err) { 91 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 92 | exit(-1); 93 | } 94 | } 95 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/group_points/__init__.py: -------------------------------------------------------------------------------- 1 | from .group_points import GroupAll, QueryAndGroup, grouping_operation 2 | 3 | __all__ = ['QueryAndGroup', 'GroupAll', 'grouping_operation'] 4 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/group_points/src/group_points.cpp: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/group_points.cpp 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | #include 11 | 12 | extern THCState *state; 13 | 14 | int group_points_wrapper(int b, int c, int n, int npoints, int nsample, 15 | at::Tensor points_tensor, at::Tensor idx_tensor, 16 | at::Tensor out_tensor); 17 | 18 | void group_points_kernel_launcher(int b, int c, int n, int npoints, int nsample, 19 | const float *points, const int *idx, 20 | float *out, cudaStream_t stream); 21 | 22 | int group_points_grad_wrapper(int b, int c, int n, int npoints, int nsample, 23 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, 24 | at::Tensor grad_points_tensor); 25 | 26 | void group_points_grad_kernel_launcher(int b, int c, int n, int npoints, 27 | int nsample, const float *grad_out, 28 | const int *idx, float *grad_points, 29 | cudaStream_t stream); 30 | 31 | int group_points_grad_wrapper(int b, int c, int n, int npoints, int nsample, 32 | at::Tensor grad_out_tensor, at::Tensor idx_tensor, 33 | at::Tensor grad_points_tensor) { 34 | float *grad_points = grad_points_tensor.data_ptr(); 35 | const int *idx = idx_tensor.data_ptr(); 36 | const float *grad_out = grad_out_tensor.data_ptr(); 37 | 38 | cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 39 | 40 | group_points_grad_kernel_launcher(b, c, n, npoints, nsample, grad_out, idx, 41 | grad_points, stream); 42 | return 1; 43 | } 44 | 45 | int group_points_wrapper(int b, int c, int n, int npoints, int nsample, 46 | at::Tensor points_tensor, at::Tensor idx_tensor, 47 | at::Tensor out_tensor) { 48 | const float *points = points_tensor.data_ptr(); 49 | const int *idx = idx_tensor.data_ptr(); 50 | float *out = out_tensor.data_ptr(); 51 | 52 | cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 53 | 54 | group_points_kernel_launcher(b, c, n, npoints, nsample, points, idx, out, 55 | stream); 56 | return 1; 57 | } 58 | 59 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 60 | m.def("forward", &group_points_wrapper, "group_points_wrapper"); 61 | m.def("backward", &group_points_grad_wrapper, "group_points_grad_wrapper"); 62 | } 63 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/group_points/src/group_points_cuda.cu: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/group_points_gpu.cu 3 | 4 | #include 5 | #include 6 | 7 | #define THREADS_PER_BLOCK 256 8 | #define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0)) 9 | 10 | __global__ void group_points_grad_kernel(int b, int c, int n, int npoints, 11 | int nsample, 12 | const float *__restrict__ grad_out, 13 | const int *__restrict__ idx, 14 | float *__restrict__ grad_points) { 15 | // grad_out: (B, C, npoints, nsample) 16 | // idx: (B, npoints, nsample) 17 | // output: 18 | // grad_points: (B, C, N) 19 | int bs_idx = blockIdx.z; 20 | int c_idx = blockIdx.y; 21 | int index = blockIdx.x * blockDim.x + threadIdx.x; 22 | int pt_idx = index / nsample; 23 | if (bs_idx >= b || c_idx >= c || pt_idx >= npoints) return; 24 | 25 | int sample_idx = index % nsample; 26 | grad_out += bs_idx * c * npoints * nsample + c_idx * npoints * nsample + 27 | pt_idx * nsample + sample_idx; 28 | idx += bs_idx * npoints * nsample + pt_idx * nsample + sample_idx; 29 | 30 | atomicAdd(grad_points + bs_idx * c * n + c_idx * n + idx[0], grad_out[0]); 31 | } 32 | 33 | void group_points_grad_kernel_launcher(int b, int c, int n, int npoints, 34 | int nsample, const float *grad_out, 35 | const int *idx, float *grad_points, 36 | cudaStream_t stream) { 37 | // grad_out: (B, C, npoints, nsample) 38 | // idx: (B, npoints, nsample) 39 | // output: 40 | // grad_points: (B, C, N) 41 | cudaError_t err; 42 | dim3 blocks(DIVUP(npoints * nsample, THREADS_PER_BLOCK), c, 43 | b); // blockIdx.x(col), blockIdx.y(row) 44 | dim3 threads(THREADS_PER_BLOCK); 45 | 46 | group_points_grad_kernel<<>>( 47 | b, c, n, npoints, nsample, grad_out, idx, grad_points); 48 | 49 | err = cudaGetLastError(); 50 | if (cudaSuccess != err) { 51 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 52 | exit(-1); 53 | } 54 | } 55 | 56 | __global__ void group_points_kernel(int b, int c, int n, int npoints, 57 | int nsample, 58 | const float *__restrict__ points, 59 | const int *__restrict__ idx, 60 | float *__restrict__ out) { 61 | // points: (B, C, N) 62 | // idx: (B, npoints, nsample) 63 | // output: 64 | // out: (B, C, npoints, nsample) 65 | int bs_idx = blockIdx.z; 66 | int c_idx = blockIdx.y; 67 | int index = blockIdx.x * blockDim.x + threadIdx.x; 68 | int pt_idx = index / nsample; 69 | if (bs_idx >= b || c_idx >= c || pt_idx >= npoints) return; 70 | 71 | int sample_idx = index % nsample; 72 | 73 | idx += bs_idx * npoints * nsample + pt_idx * nsample + sample_idx; 74 | int in_idx = bs_idx * c * n + c_idx * n + idx[0]; 75 | int out_idx = bs_idx * c * npoints * nsample + c_idx * npoints * nsample + 76 | pt_idx * nsample + sample_idx; 77 | 78 | out[out_idx] = points[in_idx]; 79 | } 80 | 81 | void group_points_kernel_launcher(int b, int c, int n, int npoints, int nsample, 82 | const float *points, const int *idx, 83 | float *out, cudaStream_t stream) { 84 | // points: (B, C, N) 85 | // idx: (B, npoints, nsample) 86 | // output: 87 | // out: (B, C, npoints, nsample) 88 | cudaError_t err; 89 | dim3 blocks(DIVUP(npoints * nsample, THREADS_PER_BLOCK), c, 90 | b); // blockIdx.x(col), blockIdx.y(row) 91 | dim3 threads(THREADS_PER_BLOCK); 92 | 93 | group_points_kernel<<>>(b, c, n, npoints, nsample, 94 | points, idx, out); 95 | // cudaDeviceSynchronize(); // for using printf in kernel function 96 | err = cudaGetLastError(); 97 | if (cudaSuccess != err) { 98 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 99 | exit(-1); 100 | } 101 | } 102 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/interpolate/__init__.py: -------------------------------------------------------------------------------- 1 | from .three_interpolate import three_interpolate 2 | from .three_nn import three_nn 3 | 4 | __all__ = ['three_nn', 'three_interpolate'] 5 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/interpolate/src/interpolate.cpp: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/interpolate.cpp 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | 15 | extern THCState *state; 16 | 17 | void three_nn_wrapper(int b, int n, int m, at::Tensor unknown_tensor, 18 | at::Tensor known_tensor, at::Tensor dist2_tensor, 19 | at::Tensor idx_tensor); 20 | 21 | void three_nn_kernel_launcher(int b, int n, int m, const float *unknown, 22 | const float *known, float *dist2, int *idx, 23 | cudaStream_t stream); 24 | 25 | void three_interpolate_wrapper(int b, int c, int m, int n, 26 | at::Tensor points_tensor, at::Tensor idx_tensor, 27 | at::Tensor weight_tensor, at::Tensor out_tensor); 28 | 29 | void three_interpolate_kernel_launcher(int b, int c, int m, int n, 30 | const float *points, const int *idx, 31 | const float *weight, float *out, 32 | cudaStream_t stream); 33 | 34 | void three_interpolate_grad_wrapper(int b, int c, int n, int m, 35 | at::Tensor grad_out_tensor, 36 | at::Tensor idx_tensor, 37 | at::Tensor weight_tensor, 38 | at::Tensor grad_points_tensor); 39 | 40 | void three_interpolate_grad_kernel_launcher(int b, int c, int n, int m, 41 | const float *grad_out, 42 | const int *idx, const float *weight, 43 | float *grad_points, 44 | cudaStream_t stream); 45 | 46 | void three_nn_wrapper(int b, int n, int m, at::Tensor unknown_tensor, 47 | at::Tensor known_tensor, at::Tensor dist2_tensor, 48 | at::Tensor idx_tensor) { 49 | const float *unknown = unknown_tensor.data_ptr(); 50 | const float *known = known_tensor.data_ptr(); 51 | float *dist2 = dist2_tensor.data_ptr(); 52 | int *idx = idx_tensor.data_ptr(); 53 | 54 | cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 55 | three_nn_kernel_launcher(b, n, m, unknown, known, dist2, idx, stream); 56 | } 57 | 58 | void three_interpolate_wrapper(int b, int c, int m, int n, 59 | at::Tensor points_tensor, at::Tensor idx_tensor, 60 | at::Tensor weight_tensor, 61 | at::Tensor out_tensor) { 62 | const float *points = points_tensor.data_ptr(); 63 | const float *weight = weight_tensor.data_ptr(); 64 | float *out = out_tensor.data_ptr(); 65 | const int *idx = idx_tensor.data_ptr(); 66 | 67 | cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 68 | three_interpolate_kernel_launcher(b, c, m, n, points, idx, weight, out, 69 | stream); 70 | } 71 | 72 | void three_interpolate_grad_wrapper(int b, int c, int n, int m, 73 | at::Tensor grad_out_tensor, 74 | at::Tensor idx_tensor, 75 | at::Tensor weight_tensor, 76 | at::Tensor grad_points_tensor) { 77 | const float *grad_out = grad_out_tensor.data_ptr(); 78 | const float *weight = weight_tensor.data_ptr(); 79 | float *grad_points = grad_points_tensor.data_ptr(); 80 | const int *idx = idx_tensor.data_ptr(); 81 | 82 | cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 83 | three_interpolate_grad_kernel_launcher(b, c, n, m, grad_out, idx, weight, 84 | grad_points, stream); 85 | } 86 | 87 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 88 | m.def("three_nn_wrapper", &three_nn_wrapper, "three_nn_wrapper"); 89 | m.def("three_interpolate_wrapper", &three_interpolate_wrapper, 90 | "three_interpolate_wrapper"); 91 | m.def("three_interpolate_grad_wrapper", &three_interpolate_grad_wrapper, 92 | "three_interpolate_grad_wrapper"); 93 | } 94 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/interpolate/src/three_interpolate_cuda.cu: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/interpolate_gpu.cu 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #define THREADS_PER_BLOCK 256 9 | #define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0)) 10 | 11 | __global__ void three_interpolate_kernel(int b, int c, int m, int n, 12 | const float *__restrict__ points, 13 | const int *__restrict__ idx, 14 | const float *__restrict__ weight, 15 | float *__restrict__ out) { 16 | // points: (B, C, M) 17 | // idx: (B, N, 3) 18 | // weight: (B, N, 3) 19 | // output: 20 | // out: (B, C, N) 21 | 22 | int bs_idx = blockIdx.z; 23 | int c_idx = blockIdx.y; 24 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 25 | 26 | if (bs_idx >= b || c_idx >= c || pt_idx >= n) return; 27 | 28 | weight += bs_idx * n * 3 + pt_idx * 3; 29 | points += bs_idx * c * m + c_idx * m; 30 | idx += bs_idx * n * 3 + pt_idx * 3; 31 | out += bs_idx * c * n + c_idx * n; 32 | 33 | out[pt_idx] = weight[0] * points[idx[0]] + weight[1] * points[idx[1]] + 34 | weight[2] * points[idx[2]]; 35 | } 36 | 37 | void three_interpolate_kernel_launcher(int b, int c, int m, int n, 38 | const float *points, const int *idx, 39 | const float *weight, float *out, 40 | cudaStream_t stream) { 41 | // points: (B, C, M) 42 | // idx: (B, N, 3) 43 | // weight: (B, N, 3) 44 | // output: 45 | // out: (B, C, N) 46 | 47 | cudaError_t err; 48 | dim3 blocks(DIVUP(n, THREADS_PER_BLOCK), c, 49 | b); // blockIdx.x(col), blockIdx.y(row) 50 | dim3 threads(THREADS_PER_BLOCK); 51 | three_interpolate_kernel<<>>(b, c, m, n, points, 52 | idx, weight, out); 53 | 54 | err = cudaGetLastError(); 55 | if (cudaSuccess != err) { 56 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 57 | exit(-1); 58 | } 59 | } 60 | 61 | __global__ void three_interpolate_grad_kernel( 62 | int b, int c, int n, int m, const float *__restrict__ grad_out, 63 | const int *__restrict__ idx, const float *__restrict__ weight, 64 | float *__restrict__ grad_points) { 65 | // grad_out: (B, C, N) 66 | // weight: (B, N, 3) 67 | // output: 68 | // grad_points: (B, C, M) 69 | 70 | int bs_idx = blockIdx.z; 71 | int c_idx = blockIdx.y; 72 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 73 | 74 | if (bs_idx >= b || c_idx >= c || pt_idx >= n) return; 75 | 76 | grad_out += bs_idx * c * n + c_idx * n + pt_idx; 77 | weight += bs_idx * n * 3 + pt_idx * 3; 78 | grad_points += bs_idx * c * m + c_idx * m; 79 | idx += bs_idx * n * 3 + pt_idx * 3; 80 | 81 | atomicAdd(grad_points + idx[0], grad_out[0] * weight[0]); 82 | atomicAdd(grad_points + idx[1], grad_out[0] * weight[1]); 83 | atomicAdd(grad_points + idx[2], grad_out[0] * weight[2]); 84 | } 85 | 86 | void three_interpolate_grad_kernel_launcher(int b, int c, int n, int m, 87 | const float *grad_out, 88 | const int *idx, const float *weight, 89 | float *grad_points, 90 | cudaStream_t stream) { 91 | // grad_out: (B, C, N) 92 | // weight: (B, N, 3) 93 | // output: 94 | // grad_points: (B, C, M) 95 | 96 | cudaError_t err; 97 | dim3 blocks(DIVUP(n, THREADS_PER_BLOCK), c, 98 | b); // blockIdx.x(col), blockIdx.y(row) 99 | dim3 threads(THREADS_PER_BLOCK); 100 | three_interpolate_grad_kernel<<>>( 101 | b, c, n, m, grad_out, idx, weight, grad_points); 102 | 103 | err = cudaGetLastError(); 104 | if (cudaSuccess != err) { 105 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 106 | exit(-1); 107 | } 108 | } 109 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/interpolate/src/three_nn_cuda.cu: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/sshaoshuai/Pointnet2.PyTorch/tree/master/pointnet2/src/interpolate_gpu.cu 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #define THREADS_PER_BLOCK 256 9 | #define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0)) 10 | 11 | __global__ void three_nn_kernel(int b, int n, int m, 12 | const float *__restrict__ unknown, 13 | const float *__restrict__ known, 14 | float *__restrict__ dist2, 15 | int *__restrict__ idx) { 16 | // unknown: (B, N, 3) 17 | // known: (B, M, 3) 18 | // output: 19 | // dist2: (B, N, 3) 20 | // idx: (B, N, 3) 21 | 22 | int bs_idx = blockIdx.y; 23 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 24 | if (bs_idx >= b || pt_idx >= n) return; 25 | 26 | unknown += bs_idx * n * 3 + pt_idx * 3; 27 | known += bs_idx * m * 3; 28 | dist2 += bs_idx * n * 3 + pt_idx * 3; 29 | idx += bs_idx * n * 3 + pt_idx * 3; 30 | 31 | float ux = unknown[0]; 32 | float uy = unknown[1]; 33 | float uz = unknown[2]; 34 | 35 | double best1 = 1e40, best2 = 1e40, best3 = 1e40; 36 | int besti1 = 0, besti2 = 0, besti3 = 0; 37 | for (int k = 0; k < m; ++k) { 38 | float x = known[k * 3 + 0]; 39 | float y = known[k * 3 + 1]; 40 | float z = known[k * 3 + 2]; 41 | float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z); 42 | if (d < best1) { 43 | best3 = best2; 44 | besti3 = besti2; 45 | best2 = best1; 46 | besti2 = besti1; 47 | best1 = d; 48 | besti1 = k; 49 | } else if (d < best2) { 50 | best3 = best2; 51 | besti3 = besti2; 52 | best2 = d; 53 | besti2 = k; 54 | } else if (d < best3) { 55 | best3 = d; 56 | besti3 = k; 57 | } 58 | } 59 | dist2[0] = best1; 60 | dist2[1] = best2; 61 | dist2[2] = best3; 62 | idx[0] = besti1; 63 | idx[1] = besti2; 64 | idx[2] = besti3; 65 | } 66 | 67 | void three_nn_kernel_launcher(int b, int n, int m, const float *unknown, 68 | const float *known, float *dist2, int *idx, 69 | cudaStream_t stream) { 70 | // unknown: (B, N, 3) 71 | // known: (B, M, 3) 72 | // output: 73 | // dist2: (B, N, 3) 74 | // idx: (B, N, 3) 75 | 76 | cudaError_t err; 77 | dim3 blocks(DIVUP(n, THREADS_PER_BLOCK), 78 | b); // blockIdx.x(col), blockIdx.y(row) 79 | dim3 threads(THREADS_PER_BLOCK); 80 | 81 | three_nn_kernel<<>>(b, n, m, unknown, known, 82 | dist2, idx); 83 | 84 | err = cudaGetLastError(); 85 | if (cudaSuccess != err) { 86 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 87 | exit(-1); 88 | } 89 | } 90 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/interpolate/three_interpolate.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Function 3 | from typing import Tuple 4 | 5 | from . import interpolate_ext 6 | 7 | 8 | class ThreeInterpolate(Function): 9 | 10 | @staticmethod 11 | def forward(ctx, features: torch.Tensor, indices: torch.Tensor, 12 | weight: torch.Tensor) -> torch.Tensor: 13 | """Performs weighted linear interpolation on 3 features. 14 | 15 | Args: 16 | features (Tensor): (B, C, M) Features descriptors to be 17 | interpolated from 18 | indices (Tensor): (B, n, 3) index three nearest neighbors 19 | of the target features in features 20 | weight (Tensor): (B, n, 3) weights of interpolation 21 | 22 | Returns: 23 | Tensor: (B, C, N) tensor of the interpolated features 24 | """ 25 | assert features.is_contiguous() 26 | assert indices.is_contiguous() 27 | assert weight.is_contiguous() 28 | 29 | B, c, m = features.size() 30 | n = indices.size(1) 31 | ctx.three_interpolate_for_backward = (indices, weight, m) 32 | output = torch.cuda.FloatTensor(B, c, n) 33 | 34 | interpolate_ext.three_interpolate_wrapper(B, c, m, n, features, 35 | indices, weight, output) 36 | return output 37 | 38 | @staticmethod 39 | def backward( 40 | ctx, grad_out: torch.Tensor 41 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 42 | """Backward of three interpolate. 43 | 44 | Args: 45 | grad_out (Tensor): (B, C, N) tensor with gradients of outputs 46 | 47 | Returns: 48 | Tensor: (B, C, M) tensor with gradients of features 49 | """ 50 | idx, weight, m = ctx.three_interpolate_for_backward 51 | B, c, n = grad_out.size() 52 | 53 | grad_features = torch.cuda.FloatTensor(B, c, m).zero_() 54 | grad_out_data = grad_out.data.contiguous() 55 | 56 | interpolate_ext.three_interpolate_grad_wrapper(B, c, n, m, 57 | grad_out_data, idx, 58 | weight, 59 | grad_features.data) 60 | return grad_features, None, None 61 | 62 | 63 | three_interpolate = ThreeInterpolate.apply 64 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/interpolate/three_nn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Function 3 | from typing import Tuple 4 | 5 | from . import interpolate_ext 6 | 7 | 8 | class ThreeNN(Function): 9 | 10 | @staticmethod 11 | def forward(ctx, target: torch.Tensor, 12 | source: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: 13 | """Find the top-3 nearest neighbors of the target set from the source 14 | set. 15 | 16 | Args: 17 | target (Tensor): shape (B, N, 3), points set that needs to 18 | find the nearest neighbors. 19 | source (Tensor): shape (B, M, 3), points set that is used 20 | to find the nearest neighbors of points in target set. 21 | 22 | Returns: 23 | Tensor: shape (B, N, 3), L2 distance of each point in target 24 | set to their corresponding nearest neighbors. 25 | """ 26 | assert target.is_contiguous() 27 | assert source.is_contiguous() 28 | 29 | B, N, _ = target.size() 30 | m = source.size(1) 31 | dist2 = torch.cuda.FloatTensor(B, N, 3) 32 | idx = torch.cuda.IntTensor(B, N, 3) 33 | 34 | interpolate_ext.three_nn_wrapper(B, N, m, target, source, dist2, idx) 35 | 36 | ctx.mark_non_differentiable(idx) 37 | 38 | return torch.sqrt(dist2), idx 39 | 40 | @staticmethod 41 | def backward(ctx, a=None, b=None): 42 | return None, None 43 | 44 | 45 | three_nn = ThreeNN.apply 46 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/iou3d/__init__.py: -------------------------------------------------------------------------------- 1 | from .iou3d_utils import boxes_iou_bev, nms_gpu, nms_normal_gpu 2 | 3 | __all__ = ['boxes_iou_bev', 'nms_gpu', 'nms_normal_gpu'] 4 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/iou3d/iou3d_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from . import iou3d_cuda 4 | 5 | 6 | def boxes_iou_bev(boxes_a, boxes_b): 7 | """Calculate boxes IoU in the bird view. 8 | 9 | Args: 10 | boxes_a (torch.Tensor): Input boxes a with shape (M, 5). 11 | boxes_b (torch.Tensor): Input boxes b with shape (N, 5). 12 | 13 | Returns: 14 | ans_iou (torch.Tensor): IoU result with shape (M, N). 15 | """ 16 | ans_iou = boxes_a.new_zeros( 17 | torch.Size((boxes_a.shape[0], boxes_b.shape[0]))) 18 | 19 | iou3d_cuda.boxes_iou_bev_gpu(boxes_a.contiguous(), boxes_b.contiguous(), 20 | ans_iou) 21 | 22 | return ans_iou 23 | 24 | 25 | def nms_gpu(boxes, scores, thresh, pre_maxsize=None, post_max_size=None): 26 | """Nms function with gpu implementation. 27 | 28 | Args: 29 | boxes (torch.Tensor): Input boxes with the shape of [N, 5] 30 | ([x1, y1, x2, y2, ry]). 31 | scores (torch.Tensor): Scores of boxes with the shape of [N]. 32 | thresh (int): Threshold. 33 | pre_maxsize (int): Max size of boxes before nms. Default: None. 34 | post_maxsize (int): Max size of boxes after nms. Default: None. 35 | 36 | Returns: 37 | torch.Tensor: Indexes after nms. 38 | """ 39 | order = scores.sort(0, descending=True)[1] 40 | 41 | if pre_maxsize is not None: 42 | order = order[:pre_maxsize] 43 | boxes = boxes[order].contiguous() 44 | 45 | keep = torch.zeros(boxes.size(0), dtype=torch.long) 46 | num_out = iou3d_cuda.nms_gpu(boxes, keep, thresh, boxes.device.index) 47 | keep = order[keep[:num_out].cuda(boxes.device)].contiguous() 48 | if post_max_size is not None: 49 | keep = keep[:post_max_size] 50 | return keep 51 | 52 | 53 | def nms_normal_gpu(boxes, scores, thresh): 54 | """Normal non maximum suppression on GPU. 55 | 56 | Args: 57 | boxes (torch.Tensor): Input boxes with shape (N, 5). 58 | scores (torch.Tensor): Scores of predicted boxes with shape (N). 59 | thresh (torch.Tensor): Threshold of non maximum suppression. 60 | 61 | Returns: 62 | torch.Tensor: Remaining indices with scores in descending order. 63 | """ 64 | order = scores.sort(0, descending=True)[1] 65 | 66 | boxes = boxes[order].contiguous() 67 | 68 | keep = torch.zeros(boxes.size(0), dtype=torch.long) 69 | num_out = iou3d_cuda.nms_normal_gpu(boxes, keep, thresh, 70 | boxes.device.index) 71 | return order[keep[:num_out].cuda(boxes.device)].contiguous() 72 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/knn/__init__.py: -------------------------------------------------------------------------------- 1 | from .knn import knn 2 | 3 | __all__ = ['knn'] 4 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/knn/knn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Function 3 | 4 | from . import knn_ext 5 | 6 | 7 | class KNN(Function): 8 | r"""KNN (CUDA) based on heap data structure. 9 | Modified from `PAConv `_. 11 | 12 | Find k-nearest points. 13 | """ 14 | 15 | @staticmethod 16 | def forward(ctx, 17 | k: int, 18 | xyz: torch.Tensor, 19 | center_xyz: torch.Tensor = None, 20 | transposed: bool = False) -> torch.Tensor: 21 | """Forward. 22 | 23 | Args: 24 | k (int): number of nearest neighbors. 25 | xyz (Tensor): (B, N, 3) if transposed == False, else (B, 3, N). 26 | xyz coordinates of the features. 27 | center_xyz (Tensor): (B, npoint, 3) if transposed == False, 28 | else (B, 3, npoint). centers of the knn query. 29 | transposed (bool): whether the input tensors are transposed. 30 | defaults to False. Should not expicitly use this keyword 31 | when calling knn (=KNN.apply), just add the fourth param. 32 | 33 | Returns: 34 | Tensor: (B, k, npoint) tensor with the indicies of 35 | the features that form k-nearest neighbours. 36 | """ 37 | assert k > 0 38 | 39 | if center_xyz is None: 40 | center_xyz = xyz 41 | 42 | if transposed: 43 | xyz = xyz.transpose(2, 1).contiguous() 44 | center_xyz = center_xyz.transpose(2, 1).contiguous() 45 | 46 | assert xyz.is_contiguous() # [B, N, 3] 47 | assert center_xyz.is_contiguous() # [B, npoint, 3] 48 | 49 | center_xyz_device = center_xyz.get_device() 50 | assert center_xyz_device == xyz.get_device(), \ 51 | 'center_xyz and xyz should be put on the same device' 52 | if torch.cuda.current_device() != center_xyz_device: 53 | torch.cuda.set_device(center_xyz_device) 54 | 55 | B, npoint, _ = center_xyz.shape 56 | N = xyz.shape[1] 57 | 58 | idx = center_xyz.new_zeros((B, npoint, k)).int() 59 | dist2 = center_xyz.new_zeros((B, npoint, k)).float() 60 | 61 | knn_ext.knn_wrapper(B, N, npoint, k, xyz, center_xyz, idx, dist2) 62 | # idx shape to [B, k, npoint] 63 | idx = idx.transpose(2, 1).contiguous() 64 | ctx.mark_non_differentiable(idx) 65 | return idx 66 | 67 | @staticmethod 68 | def backward(ctx, a=None): 69 | return None, None, None 70 | 71 | 72 | knn = KNN.apply 73 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/knn/src/knn.cpp: -------------------------------------------------------------------------------- 1 | // Modified from https://github.com/CVMI-Lab/PAConv/tree/main/scene_seg/lib/pointops/src/knnquery_heap 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | extern THCState *state; 10 | 11 | #define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x, " must be a CUDAtensor ") 12 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ") 13 | #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x) 14 | 15 | 16 | void knn_kernel_launcher( 17 | int b, 18 | int n, 19 | int m, 20 | int nsample, 21 | const float *xyz, 22 | const float *new_xyz, 23 | int *idx, 24 | float *dist2, 25 | cudaStream_t stream 26 | ); 27 | 28 | void knn_wrapper(int b, int n, int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor) 29 | { 30 | CHECK_INPUT(new_xyz_tensor); 31 | CHECK_INPUT(xyz_tensor); 32 | 33 | const float *new_xyz = new_xyz_tensor.data_ptr(); 34 | const float *xyz = xyz_tensor.data_ptr(); 35 | int *idx = idx_tensor.data_ptr(); 36 | float *dist2 = dist2_tensor.data_ptr(); 37 | 38 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 39 | 40 | knn_kernel_launcher(b, n, m, nsample, xyz, new_xyz, idx, dist2, stream); 41 | } 42 | 43 | 44 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 45 | m.def("knn_wrapper", &knn_wrapper, "knn_wrapper"); 46 | } 47 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/knn/src/knn_cuda.cu: -------------------------------------------------------------------------------- 1 | // Modified from https://github.com/CVMI-Lab/PAConv/tree/main/scene_seg/lib/pointops/src/knnquery_heap 2 | 3 | #include 4 | #include 5 | 6 | #define THREADS_PER_BLOCK 256 7 | #define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0)) 8 | 9 | 10 | __device__ void swap_float(float *x, float *y) 11 | { 12 | float tmp = *x; 13 | *x = *y; 14 | *y = tmp; 15 | } 16 | 17 | 18 | __device__ void swap_int(int *x, int *y) 19 | { 20 | int tmp = *x; 21 | *x = *y; 22 | *y = tmp; 23 | } 24 | 25 | 26 | __device__ void reheap(float *dist, int *idx, int k) 27 | { 28 | int root = 0; 29 | int child = root * 2 + 1; 30 | while (child < k) 31 | { 32 | if(child + 1 < k && dist[child+1] > dist[child]) 33 | child++; 34 | if(dist[root] > dist[child]) 35 | return; 36 | swap_float(&dist[root], &dist[child]); 37 | swap_int(&idx[root], &idx[child]); 38 | root = child; 39 | child = root * 2 + 1; 40 | } 41 | } 42 | 43 | 44 | __device__ void heap_sort(float *dist, int *idx, int k) 45 | { 46 | int i; 47 | for (i = k - 1; i > 0; i--) 48 | { 49 | swap_float(&dist[0], &dist[i]); 50 | swap_int(&idx[0], &idx[i]); 51 | reheap(dist, idx, i); 52 | } 53 | } 54 | 55 | 56 | // input: xyz (b, n, 3) new_xyz (b, m, 3) 57 | // output: idx (b, m, nsample) dist2 (b, m, nsample) 58 | __global__ void knn_kernel(int b, int n, int m, int nsample, const float *__restrict__ xyz, const float *__restrict__ new_xyz, int *__restrict__ idx, float *__restrict__ dist2) { 59 | int bs_idx = blockIdx.y; 60 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 61 | if (bs_idx >= b || pt_idx >= m) return; 62 | 63 | new_xyz += bs_idx * m * 3 + pt_idx * 3; 64 | xyz += bs_idx * n * 3; 65 | idx += bs_idx * m * nsample + pt_idx * nsample; 66 | dist2 += bs_idx * m * nsample + pt_idx * nsample; 67 | 68 | float new_x = new_xyz[0]; 69 | float new_y = new_xyz[1]; 70 | float new_z = new_xyz[2]; 71 | 72 | float best_dist[100]; 73 | int best_idx[100]; 74 | for(int i = 0; i < nsample; i++){ 75 | best_dist[i] = 1e10; 76 | best_idx[i] = 0; 77 | } 78 | for(int i = 0; i < n; i++){ 79 | float x = xyz[i * 3 + 0]; 80 | float y = xyz[i * 3 + 1]; 81 | float z = xyz[i * 3 + 2]; 82 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); 83 | if (d2 < best_dist[0]){ 84 | best_dist[0] = d2; 85 | best_idx[0] = i; 86 | reheap(best_dist, best_idx, nsample); 87 | } 88 | } 89 | heap_sort(best_dist, best_idx, nsample); 90 | for(int i = 0; i < nsample; i++){ 91 | idx[i] = best_idx[i]; 92 | dist2[i] = best_dist[i]; 93 | } 94 | } 95 | 96 | 97 | void knn_kernel_launcher(int b, int n, int m, int nsample, const float *xyz, const float *new_xyz, int *idx, float *dist2, cudaStream_t stream) { 98 | // param new_xyz: (B, m, 3) 99 | // param xyz: (B, n, 3) 100 | // param idx: (B, m, nsample) 101 | 102 | cudaError_t err; 103 | 104 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), b); // blockIdx.x(col), blockIdx.y(row) 105 | dim3 threads(THREADS_PER_BLOCK); 106 | 107 | knn_kernel<<>>(b, n, m, nsample, xyz, new_xyz, idx, dist2); 108 | // cudaDeviceSynchronize(); // for using printf in kernel function 109 | 110 | err = cudaGetLastError(); 111 | if (cudaSuccess != err) { 112 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 113 | exit(-1); 114 | } 115 | } 116 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/paconv/__init__.py: -------------------------------------------------------------------------------- 1 | from .assign_score import assign_score_withk 2 | from .paconv import PAConv, PAConvCUDA 3 | 4 | __all__ = ['assign_score_withk', 'PAConv', 'PAConvCUDA'] 5 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/paconv/assign_score.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Function 2 | 3 | from . import assign_score_withk_ext 4 | 5 | 6 | class AssignScoreWithK(Function): 7 | r"""Perform weighted sum to generate output features according to scores. 8 | Modified from `PAConv `_. 10 | 11 | This is a memory-efficient CUDA implementation of assign_scores operation, 12 | which first transform all point feature with weight bank, then assemble 13 | neighbor features with `knn_idx` and perform weighted sum of `scores`. 14 | See the `paper `_ appendix Sec. D for 15 | more detailed descriptions. 16 | 17 | Note: 18 | This implementation assumes using ``neighbor`` kernel input, which is 19 | (point_features - center_features, point_features). 20 | See https://github.com/CVMI-Lab/PAConv/blob/main/scene_seg/model/ 21 | pointnet2/paconv.py#L128 for more details. 22 | """ 23 | 24 | @staticmethod 25 | def forward(ctx, 26 | scores, 27 | point_features, 28 | center_features, 29 | knn_idx, 30 | aggregate='sum'): 31 | """Forward. 32 | 33 | Args: 34 | scores (torch.Tensor): (B, npoint, K, M), predicted scores to 35 | aggregate weight matrices in the weight bank. 36 | ``npoint`` is the number of sampled centers. 37 | ``K`` is the number of queried neighbors. 38 | ``M`` is the number of weight matrices in the weight bank. 39 | point_features (torch.Tensor): (B, N, M, out_dim) 40 | Pre-computed point features to be aggregated. 41 | center_features (torch.Tensor): (B, N, M, out_dim) 42 | Pre-computed center features to be aggregated. 43 | knn_idx (torch.Tensor): (B, npoint, K), index of sampled kNN. 44 | We assume the first idx in each row is the idx of the center. 45 | aggregate (str, optional): Aggregation method. 46 | Can be 'sum', 'avg' or 'max'. Defaults to 'sum'. 47 | 48 | Returns: 49 | torch.Tensor: (B, out_dim, npoint, K), the aggregated features. 50 | """ 51 | agg = {'sum': 0, 'avg': 1, 'max': 2} 52 | 53 | B, N, M, out_dim = point_features.size() 54 | _, npoint, K, _ = scores.size() 55 | 56 | output = point_features.new_zeros((B, out_dim, npoint, K)) 57 | assign_score_withk_ext.assign_score_withk_forward_wrapper( 58 | B, N, npoint, M, K, out_dim, agg[aggregate], 59 | point_features.contiguous(), center_features.contiguous(), 60 | scores.contiguous(), knn_idx.contiguous(), output) 61 | 62 | ctx.save_for_backward(output, point_features, center_features, scores, 63 | knn_idx) 64 | ctx.agg = agg[aggregate] 65 | 66 | return output 67 | 68 | @staticmethod 69 | def backward(ctx, grad_out): 70 | """Backward. 71 | 72 | Args: 73 | grad_out (torch.Tensor): (B, out_dim, npoint, K) 74 | 75 | Returns: 76 | grad_scores (torch.Tensor): (B, npoint, K, M) 77 | grad_point_features (torch.Tensor): (B, N, M, out_dim) 78 | grad_center_features (torch.Tensor): (B, N, M, out_dim) 79 | """ 80 | _, point_features, center_features, scores, knn_idx = ctx.saved_tensors 81 | 82 | agg = ctx.agg 83 | 84 | B, N, M, out_dim = point_features.size() 85 | _, npoint, K, _ = scores.size() 86 | 87 | grad_point_features = point_features.new_zeros(point_features.shape) 88 | grad_center_features = center_features.new_zeros(center_features.shape) 89 | grad_scores = scores.new_zeros(scores.shape) 90 | 91 | assign_score_withk_ext.assign_score_withk_backward_wrapper( 92 | B, N, npoint, M, K, out_dim, agg, grad_out.contiguous(), 93 | point_features.contiguous(), center_features.contiguous(), 94 | scores.contiguous(), knn_idx.contiguous(), grad_point_features, 95 | grad_center_features, grad_scores) 96 | 97 | return grad_scores, grad_point_features, \ 98 | grad_center_features, None, None 99 | 100 | 101 | assign_score_withk = AssignScoreWithK.apply 102 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/paconv/src/assign_score_withk.cpp: -------------------------------------------------------------------------------- 1 | // Modified from https://github.com/CVMI-Lab/PAConv/tree/main/scene_seg/lib/paconv_lib/src/gpu 2 | 3 | #include 4 | #include 5 | 6 | void assign_score_withk_forward_wrapper( 7 | int B, int N0, int N1, int M, 8 | int K, int O, int aggregate, 9 | const at::Tensor& points, 10 | const at::Tensor& centers, 11 | const at::Tensor& scores, 12 | const at::Tensor& knn_idx, 13 | at::Tensor& output 14 | ); 15 | 16 | void assign_score_withk_backward_wrapper( 17 | int B, int N0, int N1, int M, 18 | int K, int O, int aggregate, 19 | const at::Tensor& grad_out, 20 | const at::Tensor& points, 21 | const at::Tensor& centers, 22 | const at::Tensor& scores, 23 | const at::Tensor& knn_idx, 24 | at::Tensor& grad_points, 25 | at::Tensor& grad_centers, 26 | at::Tensor& grad_scores 27 | ); 28 | 29 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 30 | m.def("assign_score_withk_forward_wrapper", 31 | &assign_score_withk_forward_wrapper, 32 | "Assign score kernel forward (GPU), save memory version"); 33 | m.def("assign_score_withk_backward_wrapper", 34 | &assign_score_withk_backward_wrapper, 35 | "Assign score kernel backward (GPU), save memory version"); 36 | } 37 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/paconv/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def calc_euclidian_dist(xyz1, xyz2): 5 | """Calculate the Euclidian distance between two sets of points. 6 | 7 | Args: 8 | xyz1 (torch.Tensor): (N, 3), the first set of points. 9 | xyz2 (torch.Tensor): (N, 3), the second set of points. 10 | 11 | Returns: 12 | torch.Tensor: (N, ), the Euclidian distance between each point pair. 13 | """ 14 | assert xyz1.shape[0] == xyz2.shape[0], 'number of points are not the same' 15 | assert xyz1.shape[1] == xyz2.shape[1] == 3, \ 16 | 'points coordinates dimension is not 3' 17 | return torch.norm(xyz1 - xyz2, dim=-1) 18 | 19 | 20 | def assign_score(scores, point_features): 21 | """Perform weighted sum to aggregate output features according to scores. 22 | This function is used in non-CUDA version of PAConv. 23 | 24 | Compared to the cuda op assigh_score_withk, this pytorch implementation 25 | pre-computes output features for the neighbors of all centers, and then 26 | performs aggregation. It consumes more GPU memories. 27 | 28 | Args: 29 | scores (torch.Tensor): (B, npoint, K, M), predicted scores to 30 | aggregate weight matrices in the weight bank. 31 | `npoint` is the number of sampled centers. 32 | `K` is the number of queried neighbors. 33 | `M` is the number of weight matrices in the weight bank. 34 | point_features (torch.Tensor): (B, npoint, K, M, out_dim) 35 | Pre-computed point features to be aggregated. 36 | 37 | Returns: 38 | torch.Tensor: (B, npoint, K, out_dim), the aggregated features. 39 | """ 40 | B, npoint, K, M = scores.size() 41 | scores = scores.view(B, npoint, K, 1, M) 42 | output = torch.matmul(scores, point_features).view(B, npoint, K, -1) 43 | return output 44 | 45 | 46 | def assign_kernel_withoutk(features, kernels, M): 47 | """Pre-compute features with weight matrices in weight bank. This function 48 | is used before cuda op assign_score_withk in CUDA version PAConv. 49 | 50 | Args: 51 | features (torch.Tensor): (B, in_dim, N), input features of all points. 52 | `N` is the number of points in current point cloud. 53 | kernels (torch.Tensor): (2 * in_dim, M * out_dim), weight matrices in 54 | the weight bank, transformed from (M, 2 * in_dim, out_dim). 55 | `2 * in_dim` is because the input features are concatenation of 56 | (point_features - center_features, point_features). 57 | M (int): Number of weight matrices in the weight bank. 58 | 59 | Returns: 60 | Tuple[torch.Tensor]: both of shape (B, N, M, out_dim) 61 | point_features: Pre-computed features for points. 62 | center_features: Pre-computed features for centers. 63 | """ 64 | B, in_dim, N = features.size() 65 | feat_trans = features.permute(0, 2, 1) # [B, N, in_dim] 66 | out_feat_half1 = torch.matmul(feat_trans, kernels[:in_dim]).view( 67 | B, N, M, -1) # [B, N, M, out_dim] 68 | out_feat_half2 = torch.matmul(feat_trans, kernels[in_dim:]).view( 69 | B, N, M, -1) # [B, N, M, out_dim] 70 | 71 | # TODO: why this hard-coded if condition? 72 | # when the network input is only xyz without additional features 73 | # xyz will be used as features, so that features.size(1) == 3 % 2 != 0 74 | # we need to compensate center_features because otherwise 75 | # `point_features - center_features` will result in all zeros? 76 | if features.size(1) % 2 != 0: 77 | out_feat_half_coord = torch.matmul( 78 | feat_trans[:, :, :3], # [B, N, 3] 79 | kernels[in_dim:in_dim + 3]).view(B, N, M, -1) # [B, N, M, out_dim] 80 | else: 81 | out_feat_half_coord = torch.zeros_like(out_feat_half2) 82 | 83 | point_features = out_feat_half1 + out_feat_half2 84 | center_features = out_feat_half1 + out_feat_half_coord 85 | return point_features, center_features 86 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/pointnet_modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import build_sa_module 2 | from .point_fp_module import PointFPModule 3 | from .point_sa_module import PointSAModule, PointSAModuleMSG 4 | 5 | __all__ = [ 6 | 'build_sa_module', 'PointSAModuleMSG', 'PointSAModule', 'PointFPModule' 7 | ] 8 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/pointnet_modules/builder.py: -------------------------------------------------------------------------------- 1 | from mmcv.utils import Registry 2 | 3 | SA_MODULES = Registry('point_sa_module') 4 | 5 | 6 | def build_sa_module(cfg, *args, **kwargs): 7 | """Build PointNet2 set abstraction (SA) module. 8 | 9 | Args: 10 | cfg (None or dict): The SA module config, which should contain: 11 | - type (str): Module type. 12 | - module args: Args needed to instantiate an SA module. 13 | args (argument list): Arguments passed to the `__init__` 14 | method of the corresponding module. 15 | kwargs (keyword arguments): Keyword arguments passed to the `__init__` 16 | method of the corresponding SA module . 17 | 18 | Returns: 19 | nn.Module: Created SA module. 20 | """ 21 | if cfg is None: 22 | cfg_ = dict(type='PointSAModule') 23 | else: 24 | if not isinstance(cfg, dict): 25 | raise TypeError('cfg must be a dict') 26 | if 'type' not in cfg: 27 | raise KeyError('the cfg dict must contain the key "type"') 28 | cfg_ = cfg.copy() 29 | 30 | module_type = cfg_.pop('type') 31 | if module_type not in SA_MODULES: 32 | raise KeyError(f'Unrecognized module type {module_type}') 33 | else: 34 | sa_module = SA_MODULES.get(module_type) 35 | 36 | module = sa_module(*args, **kwargs, **cfg_) 37 | 38 | return module 39 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/pointnet_modules/point_fp_module.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from mmcv.cnn import ConvModule 3 | from mmcv.runner import force_fp32 4 | from torch import nn as nn 5 | from typing import List 6 | 7 | from mmdet3d.ops import three_interpolate, three_nn 8 | 9 | 10 | class PointFPModule(nn.Module): 11 | """Point feature propagation module used in PointNets. 12 | 13 | Propagate the features from one set to another. 14 | 15 | Args: 16 | mlp_channels (list[int]): List of mlp channels. 17 | norm_cfg (dict): Type of normalization method. 18 | Default: dict(type='BN2d'). 19 | """ 20 | 21 | def __init__(self, 22 | mlp_channels: List[int], 23 | norm_cfg: dict = dict(type='BN2d')): 24 | super().__init__() 25 | self.fp16_enabled = False 26 | self.mlps = nn.Sequential() 27 | for i in range(len(mlp_channels) - 1): 28 | self.mlps.add_module( 29 | f'layer{i}', 30 | ConvModule( 31 | mlp_channels[i], 32 | mlp_channels[i + 1], 33 | kernel_size=(1, 1), 34 | stride=(1, 1), 35 | conv_cfg=dict(type='Conv2d'), 36 | norm_cfg=norm_cfg)) 37 | 38 | @force_fp32() 39 | def forward(self, target: torch.Tensor, source: torch.Tensor, 40 | target_feats: torch.Tensor, 41 | source_feats: torch.Tensor) -> torch.Tensor: 42 | """forward. 43 | 44 | Args: 45 | target (Tensor): (B, n, 3) tensor of the xyz positions of 46 | the target features. 47 | source (Tensor): (B, m, 3) tensor of the xyz positions of 48 | the source features. 49 | target_feats (Tensor): (B, C1, n) tensor of the features to be 50 | propagated to. 51 | source_feats (Tensor): (B, C2, m) tensor of features 52 | to be propagated. 53 | 54 | Return: 55 | Tensor: (B, M, N) M = mlp[-1], tensor of the target features. 56 | """ 57 | if source is not None: 58 | dist, idx = three_nn(target, source) 59 | dist_reciprocal = 1.0 / (dist + 1e-8) 60 | norm = torch.sum(dist_reciprocal, dim=2, keepdim=True) 61 | weight = dist_reciprocal / norm 62 | 63 | interpolated_feats = three_interpolate(source_feats, idx, weight) 64 | else: 65 | interpolated_feats = source_feats.expand(*source_feats.size()[0:2], 66 | target.size(1)) 67 | 68 | if target_feats is not None: 69 | new_features = torch.cat([interpolated_feats, target_feats], 70 | dim=1) # (B, C2 + C1, n) 71 | else: 72 | new_features = interpolated_feats 73 | 74 | new_features = new_features.unsqueeze(-1) 75 | new_features = self.mlps(new_features) 76 | 77 | return new_features.squeeze(-1) 78 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/roiaware_pool3d/__init__.py: -------------------------------------------------------------------------------- 1 | from .points_in_boxes import (points_in_boxes_batch, points_in_boxes_cpu, 2 | points_in_boxes_gpu) 3 | from .roiaware_pool3d import RoIAwarePool3d 4 | 5 | __all__ = [ 6 | 'RoIAwarePool3d', 'points_in_boxes_gpu', 'points_in_boxes_cpu', 7 | 'points_in_boxes_batch' 8 | ] 9 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/roiaware_pool3d/roiaware_pool3d.py: -------------------------------------------------------------------------------- 1 | import mmcv 2 | import torch 3 | from torch import nn as nn 4 | from torch.autograd import Function 5 | 6 | from . import roiaware_pool3d_ext 7 | 8 | 9 | class RoIAwarePool3d(nn.Module): 10 | 11 | def __init__(self, out_size, max_pts_per_voxel=128, mode='max'): 12 | super().__init__() 13 | """RoIAwarePool3d module 14 | 15 | Args: 16 | out_size (int or tuple): n or [n1, n2, n3] 17 | max_pts_per_voxel (int): m 18 | mode (str): 'max' or 'avg' 19 | """ 20 | self.out_size = out_size 21 | self.max_pts_per_voxel = max_pts_per_voxel 22 | assert mode in ['max', 'avg'] 23 | pool_method_map = {'max': 0, 'avg': 1} 24 | self.mode = pool_method_map[mode] 25 | 26 | def forward(self, rois, pts, pts_feature): 27 | """RoIAwarePool3d module forward. 28 | 29 | Args: 30 | rois (torch.Tensor): [N, 7],in LiDAR coordinate, 31 | (x, y, z) is the bottom center of rois 32 | pts (torch.Tensor): [npoints, 3] 33 | pts_feature (torch.Tensor): [npoints, C] 34 | 35 | Returns: 36 | pooled_features (torch.Tensor): [N, out_x, out_y, out_z, C] 37 | """ 38 | 39 | return RoIAwarePool3dFunction.apply(rois, pts, pts_feature, 40 | self.out_size, 41 | self.max_pts_per_voxel, self.mode) 42 | 43 | 44 | class RoIAwarePool3dFunction(Function): 45 | 46 | @staticmethod 47 | def forward(ctx, rois, pts, pts_feature, out_size, max_pts_per_voxel, 48 | mode): 49 | """RoIAwarePool3d function forward. 50 | 51 | Args: 52 | rois (torch.Tensor): [N, 7], in LiDAR coordinate, 53 | (x, y, z) is the bottom center of rois 54 | pts (torch.Tensor): [npoints, 3] 55 | pts_feature (torch.Tensor): [npoints, C] 56 | out_size (int or tuple): n or [n1, n2, n3] 57 | max_pts_per_voxel (int): m 58 | mode (int): 0 (max pool) or 1 (average pool) 59 | 60 | Returns: 61 | pooled_features (torch.Tensor): [N, out_x, out_y, out_z, C] 62 | """ 63 | 64 | if isinstance(out_size, int): 65 | out_x = out_y = out_z = out_size 66 | else: 67 | assert len(out_size) == 3 68 | assert mmcv.is_tuple_of(out_size, int) 69 | out_x, out_y, out_z = out_size 70 | 71 | num_rois = rois.shape[0] 72 | num_channels = pts_feature.shape[-1] 73 | num_pts = pts.shape[0] 74 | 75 | pooled_features = pts_feature.new_zeros( 76 | (num_rois, out_x, out_y, out_z, num_channels)) 77 | argmax = pts_feature.new_zeros( 78 | (num_rois, out_x, out_y, out_z, num_channels), dtype=torch.int) 79 | pts_idx_of_voxels = pts_feature.new_zeros( 80 | (num_rois, out_x, out_y, out_z, max_pts_per_voxel), 81 | dtype=torch.int) 82 | 83 | roiaware_pool3d_ext.forward(rois, pts, pts_feature, argmax, 84 | pts_idx_of_voxels, pooled_features, mode) 85 | 86 | ctx.roiaware_pool3d_for_backward = (pts_idx_of_voxels, argmax, mode, 87 | num_pts, num_channels) 88 | return pooled_features 89 | 90 | @staticmethod 91 | def backward(ctx, grad_out): 92 | """RoIAwarePool3d function forward. 93 | 94 | Args: 95 | grad_out (torch.Tensor): [N, out_x, out_y, out_z, C] 96 | Returns: 97 | grad_in (torch.Tensor): [npoints, C] 98 | """ 99 | ret = ctx.roiaware_pool3d_for_backward 100 | pts_idx_of_voxels, argmax, mode, num_pts, num_channels = ret 101 | 102 | grad_in = grad_out.new_zeros((num_pts, num_channels)) 103 | roiaware_pool3d_ext.backward(pts_idx_of_voxels, argmax, 104 | grad_out.contiguous(), grad_in, mode) 105 | 106 | return None, None, grad_in, None, None, None 107 | 108 | 109 | if __name__ == '__main__': 110 | pass 111 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/roiaware_pool3d/src/points_in_boxes_cpu.cpp: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/sshaoshuai/PCDet/blob/master/pcdet/ops/roiaware_pool3d/src/roiaware_pool3d_kernel.cu 3 | // Written by Shaoshuai Shi 4 | // All Rights Reserved 2019. 5 | 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | 12 | #define CHECK_CONTIGUOUS(x) \ 13 | TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ") 14 | // #define DEBUG 15 | 16 | inline void lidar_to_local_coords_cpu(float shift_x, float shift_y, float rz, 17 | float &local_x, float &local_y) { 18 | // should rotate pi/2 + alpha to translate LiDAR to local 19 | float rot_angle = rz + M_PI / 2; 20 | float cosa = cos(rot_angle), sina = sin(rot_angle); 21 | local_x = shift_x * cosa + shift_y * (-sina); 22 | local_y = shift_x * sina + shift_y * cosa; 23 | } 24 | 25 | inline int check_pt_in_box3d_cpu(const float *pt, const float *box3d, 26 | float &local_x, float &local_y) { 27 | // param pt: (x, y, z) 28 | // param box3d: (cx, cy, cz, w, l, h, rz) in LiDAR coordinate, cz in the 29 | // bottom center 30 | float x = pt[0], y = pt[1], z = pt[2]; 31 | float cx = box3d[0], cy = box3d[1], cz = box3d[2]; 32 | float w = box3d[3], l = box3d[4], h = box3d[5], rz = box3d[6]; 33 | cz += h / 2.0; // shift to the center since cz in box3d is the bottom center 34 | 35 | if (fabsf(z - cz) > h / 2.0) return 0; 36 | lidar_to_local_coords_cpu(x - cx, y - cy, rz, local_x, local_y); 37 | float in_flag = (local_x > -l / 2.0) & (local_x < l / 2.0) & 38 | (local_y > -w / 2.0) & (local_y < w / 2.0); 39 | return in_flag; 40 | } 41 | 42 | int points_in_boxes_cpu(at::Tensor boxes_tensor, at::Tensor pts_tensor, 43 | at::Tensor pts_indices_tensor) { 44 | // params boxes: (N, 7) [x, y, z, w, l, h, rz] in LiDAR coordinate, z is the 45 | // bottom center, each box DO NOT overlaps params pts: (npoints, 3) [x, y, z] 46 | // in LiDAR coordinate params pts_indices: (N, npoints) 47 | 48 | CHECK_CONTIGUOUS(boxes_tensor); 49 | CHECK_CONTIGUOUS(pts_tensor); 50 | CHECK_CONTIGUOUS(pts_indices_tensor); 51 | 52 | int boxes_num = boxes_tensor.size(0); 53 | int pts_num = pts_tensor.size(0); 54 | 55 | const float *boxes = boxes_tensor.data_ptr(); 56 | const float *pts = pts_tensor.data_ptr(); 57 | int *pts_indices = pts_indices_tensor.data_ptr(); 58 | 59 | float local_x = 0, local_y = 0; 60 | for (int i = 0; i < boxes_num; i++) { 61 | for (int j = 0; j < pts_num; j++) { 62 | int cur_in_flag = 63 | check_pt_in_box3d_cpu(pts + j * 3, boxes + i * 7, local_x, local_y); 64 | pts_indices[i * pts_num + j] = cur_in_flag; 65 | } 66 | } 67 | 68 | return 1; 69 | } 70 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/spconv/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 Yan Yan 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from .conv import (SparseConv2d, SparseConv3d, SparseConvTranspose2d, 16 | SparseConvTranspose3d, SparseInverseConv2d, 17 | SparseInverseConv3d, SubMConv2d, SubMConv3d) 18 | from .modules import SparseModule, SparseSequential 19 | from .pool import SparseMaxPool2d, SparseMaxPool3d 20 | from .structure import SparseConvTensor, scatter_nd 21 | 22 | __all__ = [ 23 | 'SparseConv2d', 24 | 'SparseConv3d', 25 | 'SubMConv2d', 26 | 'SubMConv3d', 27 | 'SparseConvTranspose2d', 28 | 'SparseConvTranspose3d', 29 | 'SparseInverseConv2d', 30 | 'SparseInverseConv3d', 31 | 'SparseModule', 32 | 'SparseSequential', 33 | 'SparseMaxPool2d', 34 | 'SparseMaxPool3d', 35 | 'SparseConvTensor', 36 | 'scatter_nd', 37 | ] 38 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/spconv/functional.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 Yan Yan 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from torch.autograd import Function 16 | 17 | from . import ops as ops 18 | 19 | 20 | class SparseConvFunction(Function): 21 | 22 | @staticmethod 23 | def forward(ctx, features, filters, indice_pairs, indice_pair_num, 24 | num_activate_out): 25 | ctx.save_for_backward(indice_pairs, indice_pair_num, features, filters) 26 | return ops.indice_conv(features, filters, indice_pairs, 27 | indice_pair_num, num_activate_out, False) 28 | 29 | @staticmethod 30 | def backward(ctx, grad_output): 31 | indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors 32 | input_bp, filters_bp = ops.indice_conv_backward( 33 | features, filters, grad_output, indice_pairs, indice_pair_num, 34 | False) 35 | 36 | return input_bp, filters_bp, None, None, None 37 | 38 | 39 | class SparseInverseConvFunction(Function): 40 | 41 | @staticmethod 42 | def forward(ctx, features, filters, indice_pairs, indice_pair_num, 43 | num_activate_out): 44 | ctx.save_for_backward(indice_pairs, indice_pair_num, features, filters) 45 | return ops.indice_conv(features, filters, indice_pairs, 46 | indice_pair_num, num_activate_out, True, False) 47 | 48 | @staticmethod 49 | def backward(ctx, grad_output): 50 | indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors 51 | input_bp, filters_bp = ops.indice_conv_backward( 52 | features, filters, grad_output, indice_pairs, indice_pair_num, 53 | True, False) 54 | 55 | return input_bp, filters_bp, None, None, None 56 | 57 | 58 | class SubMConvFunction(Function): 59 | 60 | @staticmethod 61 | def forward(ctx, features, filters, indice_pairs, indice_pair_num, 62 | num_activate_out): 63 | ctx.save_for_backward(indice_pairs, indice_pair_num, features, filters) 64 | return ops.indice_conv(features, filters, indice_pairs, 65 | indice_pair_num, num_activate_out, False, True) 66 | 67 | @staticmethod 68 | def backward(ctx, grad_output): 69 | indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors 70 | input_bp, filters_bp = ops.indice_conv_backward( 71 | features, filters, grad_output, indice_pairs, indice_pair_num, 72 | False, True) 73 | 74 | return input_bp, filters_bp, None, None, None 75 | 76 | 77 | class SparseMaxPoolFunction(Function): 78 | 79 | @staticmethod 80 | def forward(ctx, features, indice_pairs, indice_pair_num, 81 | num_activate_out): 82 | out = ops.indice_maxpool(features, indice_pairs, indice_pair_num, 83 | num_activate_out) 84 | ctx.save_for_backward(indice_pairs, indice_pair_num, features, out) 85 | return out 86 | 87 | @staticmethod 88 | def backward(ctx, grad_output): 89 | indice_pairs, indice_pair_num, features, out = ctx.saved_tensors 90 | input_bp = ops.indice_maxpool_backward(features, out, grad_output, 91 | indice_pairs, indice_pair_num) 92 | return input_bp, None, None, None 93 | 94 | 95 | indice_conv = SparseConvFunction.apply 96 | indice_inverse_conv = SparseInverseConvFunction.apply 97 | indice_subm_conv = SubMConvFunction.apply 98 | indice_maxpool = SparseMaxPoolFunction.apply 99 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/spconv/include/paramsgrid.h: -------------------------------------------------------------------------------- 1 | // Copyright 2019 Yan Yan 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #ifndef PARAMS_GRID_H_ 16 | #define PARAMS_GRID_H_ 17 | #include 18 | #include 19 | 20 | namespace detail { 21 | template 22 | int getTotalSize(std::vector arg) { 23 | return arg.size(); 24 | } 25 | 26 | template 27 | int getTotalSize(std::vector arg, std::vector... args) { 28 | return arg.size() * getTotalSize(args...); 29 | } 30 | template 31 | int getSize(std::vector arg) { 32 | return arg.size(); 33 | } 34 | 35 | template 36 | void assigner(TT &src, std::vector counter, std::vector &arg) { 37 | std::get(src) = arg[counter[Idx]]; 38 | } 39 | 40 | template 41 | void assigner(TT &src, std::vector counter, std::vector &arg, 42 | std::vector &... args) { 43 | std::get(src) = arg[counter[Idx]]; 44 | assigner(src, counter, args...); 45 | } 46 | } // namespace detail 47 | template 48 | std::vector> paramsGrid(std::vector... args) { 49 | int length = detail::getTotalSize(args...); 50 | std::vector sizes = {detail::getSize(args)...}; 51 | int size = sizes.size(); 52 | 53 | std::vector> params(length); 54 | std::vector counter(size); 55 | for (int i = 0; i < length; ++i) { 56 | detail::assigner<0>(params[i], counter, args...); 57 | counter[size - 1] += 1; 58 | for (int c = size - 1; c >= 0; --c) { 59 | if (counter[c] == sizes[c] && c > 0) { 60 | counter[c - 1] += 1; 61 | counter[c] = 0; 62 | } 63 | } 64 | } 65 | return params; 66 | } 67 | 68 | #endif 69 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/spconv/include/pybind11_utils.h: -------------------------------------------------------------------------------- 1 | // Copyright 2019 Yan Yan 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #pragma once 16 | #include 17 | #include 18 | #include // everything needed for embedding 19 | #include 20 | #include 21 | #include 22 | #include 23 | 24 | #include 25 | 26 | namespace py = pybind11; 27 | 28 | template 29 | std::vector array2Vector(TPyObject arr){ 30 | py::array arr_np = arr; 31 | size_t size = arr.attr("size").template cast(); 32 | py::array_t arr_cc = arr_np; 33 | std::vector data(arr_cc.data(), arr_cc.data() + size); 34 | return data; 35 | } 36 | 37 | template 38 | std::vector arrayT2Vector(py::array_t arr) 39 | { 40 | std::vector data(arr.data(), arr.data() + arr.size()); 41 | return data; 42 | } 43 | 44 | template 45 | tv::TensorView array2TensorView(TPyObject arr){ 46 | py::array arr_np = arr; 47 | py::array_t arr_cc = arr_np; 48 | tv::Shape shape; 49 | for (int i = 0; i < arr_cc.ndim(); ++i){ 50 | shape.push_back(arr_cc.shape(i)); 51 | } 52 | return tv::TensorView(arr_cc.mutable_data(), shape); 53 | } 54 | template 55 | tv::TensorView arrayT2TensorView(py::array_t arr){ 56 | tv::Shape shape; 57 | for (int i = 0; i < arr.ndim(); ++i){ 58 | shape.push_back(arr.shape(i)); 59 | } 60 | return tv::TensorView(arr.mutable_data(), shape); 61 | } 62 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/spconv/include/spconv/indice.h: -------------------------------------------------------------------------------- 1 | // Copyright 2019 Yan Yan 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #ifndef SPARSE_CONV_INDICE_FUNCTOR_H_ 16 | #define SPARSE_CONV_INDICE_FUNCTOR_H_ 17 | #include 18 | 19 | namespace spconv { 20 | namespace functor { 21 | template 22 | struct CreateConvIndicePairFunctorP1 { 23 | Index operator()(const Device& d, tv::TensorView indicesIn, 24 | tv::TensorView indicesOut, 25 | tv::TensorView gridsOut, 26 | tv::TensorView indicePairs, 27 | tv::TensorView indiceNum, 28 | tv::TensorView indicePairUnique, 29 | const tv::SimpleVector kernelSize, 30 | const tv::SimpleVector stride, 31 | const tv::SimpleVector padding, 32 | const tv::SimpleVector dilation, 33 | const tv::SimpleVector outSpatialShape, 34 | bool transpose); 35 | }; 36 | 37 | template 38 | struct CreateConvIndicePairFunctorP2 { 39 | Index operator()(const Device& d, tv::TensorView indicesIn, 40 | tv::TensorView indicesOut, 41 | tv::TensorView gridsOut, 42 | tv::TensorView indicePairs, 43 | tv::TensorView indiceNum, 44 | tv::TensorView indicePairUnique, 45 | const tv::SimpleVector outSpatialShape, 46 | bool transpose, bool resetGrid = false); 47 | }; 48 | 49 | template 50 | struct CreateConvIndicePairFunctor { 51 | Index operator()(const Device& d, tv::TensorView indicesIn, 52 | tv::TensorView indicesOut, 53 | tv::TensorView gridsOut, 54 | tv::TensorView indicePairs, 55 | tv::TensorView indiceNum, 56 | const tv::SimpleVector kernelSize, 57 | const tv::SimpleVector stride, 58 | const tv::SimpleVector padding, 59 | const tv::SimpleVector dilation, 60 | const tv::SimpleVector outSpatialShape, 61 | bool transpose, bool resetGrid = false); 62 | }; 63 | 64 | template 65 | struct CreateSubMIndicePairFunctor { 66 | Index operator()(const Device& d, tv::TensorView indicesIn, 67 | tv::TensorView gridsOut, 68 | tv::TensorView indicePairs, 69 | tv::TensorView indiceNum, 70 | const tv::SimpleVector kernelSize, 71 | const tv::SimpleVector stride, 72 | const tv::SimpleVector padding, 73 | const tv::SimpleVector dilation, 74 | const tv::SimpleVector outSpatialShape, 75 | bool transpose, bool resetGrid = false); 76 | }; 77 | } // namespace functor 78 | } // namespace spconv 79 | 80 | #endif 81 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/spconv/include/spconv/maxpool.h: -------------------------------------------------------------------------------- 1 | // Copyright 2019 Yan Yan 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #ifndef SPARSE_MAXPOOL_FUNCTOR_H_ 16 | #define SPARSE_MAXPOOL_FUNCTOR_H_ 17 | #include 18 | 19 | namespace spconv { 20 | namespace functor { 21 | template 22 | struct SparseMaxPoolForwardFunctor { 23 | void operator()(const Device& d, tv::TensorView outFeatures, 24 | tv::TensorView inFeatures, 25 | tv::TensorView indices, int size); 26 | }; 27 | 28 | template 29 | struct SparseMaxPoolBackwardFunctor { 30 | void operator()(const Device& d, tv::TensorView outFeatures, 31 | tv::TensorView inFeatures, 32 | tv::TensorView dout, tv::TensorView din, 33 | tv::TensorView indices, int size); 34 | }; 35 | 36 | } // namespace functor 37 | } // namespace spconv 38 | 39 | #endif 40 | -------------------------------------------------------------------------------- /utils/mm3d_pn2/ops/spconv/include/spconv/mp_helper.h: -------------------------------------------------------------------------------- 1 | #ifndef MP_HELPER_H_ 2 | #define MP_HELPER_H_ 3 | #include 4 | #include 5 | 6 | namespace spconv { 7 | template 8 | struct mp_list {}; 9 | 10 | template 11 | using mp_list_c = mp_list...>; 12 | 13 | namespace detail { 14 | 15 | template 16 | constexpr F mp_for_each_impl(mp_list, F &&f) { 17 | return std::initializer_list{(f(T()), 0)...}, std::forward(f); 18 | } 19 | 20 | template 21 | constexpr F mp_for_each_impl(mp_list<>, F &&f) { 22 | return std::forward(f); 23 | } 24 | 25 | } // namespace detail 26 | 27 | namespace detail { 28 | 29 | template class B> 30 | struct mp_rename_impl { 31 | // An error "no type named 'type'" here means that the first argument to 32 | // mp_rename is not a list 33 | }; 34 | 35 | template