├── README.md ├── assets ├── NewFramework-1.png └── README.md ├── environment.yml ├── outputs └── .gitignore └── src ├── __pycache__ ├── bunny_data.cpython-38.pyc ├── data.cpython-36.pyc ├── data.cpython-37.pyc ├── data.cpython-38.pyc ├── data_shapenet.cpython-38.pyc ├── data_util.cpython-38.pyc ├── pretrain.cpython-38.pyc ├── train.cpython-36.pyc ├── train.cpython-37.pyc ├── train.cpython-38.pyc ├── train_keypoint.cpython-38.pyc ├── train_real.cpython-38.pyc ├── train_semkitti.cpython-38.pyc ├── util.cpython-36.pyc ├── util.cpython-37.pyc └── util.cpython-38.pyc ├── chamfer_distance ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── __init__.cpython-38.pyc │ ├── chamfer_distance_cpu.cpython-36.pyc │ ├── chamfer_distance_cpu.cpython-38.pyc │ ├── chamfer_distance_gpu.cpython-36.pyc │ └── chamfer_distance_gpu.cpython-38.pyc ├── chamfer_distance.cu ├── chamfer_distance_cpu.cpp ├── chamfer_distance_cpu.py ├── chamfer_distance_gpu.cpp └── chamfer_distance_gpu.py ├── data ├── RealDataset.py ├── __pycache__ │ ├── RealDataset.cpython-38.pyc │ ├── data.cpython-38.pyc │ └── real_dataset.cpython-38.pyc ├── data.py └── real_dataset.py ├── demo.py ├── extensions ├── __init__.py ├── __pycache__ │ └── __init__.cpython-38.pyc ├── chamfer_dist │ ├── __init__.py │ ├── __pycache__ │ │ └── __init__.cpython-38.pyc │ ├── chamfer.cu │ ├── chamfer_cuda.cpp │ ├── setup.py │ └── test.py ├── cubic_feature_sampling │ ├── __init__.py │ ├── __pycache__ │ │ └── __init__.cpython-38.pyc │ ├── build │ │ ├── lib.linux-x86_64-3.8 │ │ │ └── cubic_feature_sampling.cpython-38-x86_64-linux-gnu.so │ │ └── temp.linux-x86_64-3.8 │ │ │ ├── .ninja_deps │ │ │ ├── .ninja_log │ │ │ ├── build.ninja │ │ │ ├── cubic_feature_sampling.o │ │ │ └── cubic_feature_sampling_cuda.o │ ├── cubic_feature_sampling.cu │ ├── cubic_feature_sampling.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ └── top_level.txt │ ├── cubic_feature_sampling_cuda.cpp │ ├── dist │ │ └── cubic_feature_sampling-1.1.0-py3.8-linux-x86_64.egg │ ├── setup.py │ └── test.py ├── gridding │ ├── __init__.py │ ├── __pycache__ │ │ └── __init__.cpython-38.pyc │ ├── build │ │ ├── lib.linux-x86_64-3.8 │ │ │ └── gridding.cpython-38-x86_64-linux-gnu.so │ │ └── temp.linux-x86_64-3.8 │ │ │ ├── .ninja_deps │ │ │ ├── .ninja_log │ │ │ ├── build.ninja │ │ │ ├── gridding.o │ │ │ ├── gridding_cuda.o │ │ │ └── gridding_reverse.o │ ├── dist │ │ └── gridding-2.1.0-py3.8-linux-x86_64.egg │ ├── gridding.cu │ ├── gridding.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ └── top_level.txt │ ├── gridding_cuda.cpp │ ├── gridding_reverse.cu │ ├── setup.py │ └── test.py └── gridding_loss │ ├── __init__.py │ ├── gridding_distance.cu │ ├── gridding_distance_cuda.cpp │ └── setup.py ├── loss ├── __pycache__ │ ├── loss.cpython-36.pyc │ ├── loss.cpython-37.pyc │ └── loss.cpython-38.pyc └── loss.py ├── main.py ├── models ├── PUGAN.py ├── TopNet.py ├── __pycache__ │ ├── AutoEncoder.cpython-36.pyc │ ├── AutoEncoder.cpython-38.pyc │ ├── PUGAN.cpython-36.pyc │ ├── PUGAN.cpython-38.pyc │ ├── TopNet.cpython-38.pyc │ ├── base_model.cpython-37.pyc │ ├── base_model.cpython-38.pyc │ ├── grnet.cpython-38.pyc │ ├── keypointfinder.cpython-37.pyc │ ├── keypointfinder.cpython-38.pyc │ ├── model.cpython-36.pyc │ ├── model.cpython-38.pyc │ ├── pcn.cpython-38.pyc │ ├── pcnenc_polydec.cpython-38.pyc │ ├── polyconv.cpython-36.pyc │ ├── polyconv.cpython-37.pyc │ ├── polyconv.cpython-38.pyc │ ├── polynet.cpython-36.pyc │ ├── polynet.cpython-37.pyc │ ├── polynet.cpython-38.pyc │ ├── polynet_enc_pcn_dec.cpython-38.pyc │ ├── util.cpython-36.pyc │ ├── util.cpython-38.pyc │ └── utils.cpython-36.pyc ├── base_model.py ├── grnet.py ├── model.py ├── ninja-linux.zip ├── pcn.py ├── pcnenc_polydec.py ├── polyconv.py ├── polynet.py └── polynet_enc_pcn_dec.py ├── pointnet2 ├── __pycache__ │ ├── pointnet2_utils.cpython-36.pyc │ ├── pointnet2_utils.cpython-38.pyc │ ├── pytorch_utils.cpython-36.pyc │ └── pytorch_utils.cpython-38.pyc ├── _ext_src │ ├── include │ │ ├── ball_query.h │ │ ├── cuda_utils.h │ │ ├── group_points.h │ │ ├── interpolate.h │ │ ├── sampling.h │ │ └── utils.h │ └── src │ │ ├── ball_query.cpp │ │ ├── ball_query_gpu.cu │ │ ├── bindings.cpp │ │ ├── group_points.cpp │ │ ├── group_points_gpu.cu │ │ ├── interpolate.cpp │ │ ├── interpolate_gpu.cu │ │ ├── sampling.cpp │ │ └── sampling_gpu.cu ├── build │ ├── lib.linux-x86_64-3.6 │ │ └── pointnet2 │ │ │ └── _ext.cpython-36m-x86_64-linux-gnu.so │ └── temp.linux-x86_64-3.6 │ │ └── _ext_src │ │ └── src │ │ ├── ball_query.o │ │ ├── ball_query_gpu.o │ │ ├── bindings.o │ │ ├── group_points.o │ │ ├── group_points_gpu.o │ │ ├── interpolate.o │ │ ├── interpolate_gpu.o │ │ ├── sampling.o │ │ └── sampling_gpu.o ├── dist │ └── pointnet2-0.0.0-py3.6-linux-x86_64.egg ├── pointnet2.egg-info │ ├── PKG-INFO │ ├── SOURCES.txt │ ├── dependency_links.txt │ └── top_level.txt ├── pointnet2_modules.py ├── pointnet2_utils.py ├── pytorch_utils.py └── setup.py ├── pretrain.py ├── train.py ├── train_real.py ├── train_semkitti.py ├── util.py └── utils ├── Logger.py ├── __pycache__ ├── Logger.cpython-36.pyc ├── Logger.cpython-38.pyc ├── find_Nearest_Neighbor.cpython-36.pyc ├── find_Nearest_Neighbor.cpython-37.pyc ├── find_Nearest_Neighbor.cpython-38.pyc ├── generate_normals.cpython-36.pyc ├── generate_normals.cpython-38.pyc ├── io.cpython-38.pyc ├── output_xyz.cpython-36.pyc ├── output_xyz.cpython-37.pyc ├── output_xyz.cpython-38.pyc ├── pc_transform.cpython-38.pyc ├── pointnet2_utils.cpython-38.pyc ├── random_pc_downsample.cpython-36.pyc ├── random_pc_downsample.cpython-37.pyc ├── random_pc_downsample.cpython-38.pyc └── read_ply.cpython-38.pyc ├── find_Nearest_Neighbor.py ├── io.py ├── output_xyz.py ├── pc_transform.py ├── pointnet2_utils.py └── read_ply.py /README.md: -------------------------------------------------------------------------------- 1 | # ACL-SPC_PyTorch 2 | 3 | This repository contains the official code to reproduce the results from the paper: 4 | 5 | **ACL-SPC: Adaptive Closed-Loop system for Self-Supervised Point Cloud Completion (CVPR 2023)** 6 | 7 | \[[arXiv](https://arxiv.org/abs/2303.01979)\] \[[presentation]()\] 8 | 9 | ![architecture](https://github.com/Sangminhong/ACL-SPC_PyTorch/blob/master/assets/NewFramework-1.png) 10 | 11 | ## Installation 12 | Clone this repository into any place you want. 13 | ``` 14 | git clone https://github.com/Sangminhong/ACL-SPC_PyTorch.git 15 | cd ACL-SPC_PyTorch 16 | ``` 17 | ### Dependencies 18 | You can try downloading the environment.yml and set the conda environment. 19 | ``` 20 | conda env create -f environment.yml --name ACL_SPC 21 | conda activate ACL_SPC 22 | ``` 23 | Or try to set the environment manually. 24 | * Python 3.8.5 25 | * PyTorch 1.7.1 26 | * numpy 27 | * h5py 28 | * numba 29 | * scikit-learn 30 | * open3d 31 | * torchsummary 32 | * pytorch3d 33 | * KNN-CUDA 34 | * pykdtree 35 | * torch_scatter 36 | 37 | ## Quick Start 38 | If you want to test your partial point cloud on pretrained model you can simply run the below command. 39 | ``` 40 | CUDA_VISIBLE_DEVICES=0 python demo.py --input {input_filename} --model_filename {model_filename} --fine_tune False 41 | ``` 42 | or if you need to finetune on your own dataset, you can simply run the below command. 43 | ``` 44 | CUDA_VISIBLE_DEVICES=0 python demo.py --input {input_filename} --model_filename {model_filename} --fine_tune True --dataset_directory {dataset_directory} 45 | ``` 46 | ## Expriments 47 | 48 | #### Pretrained model 49 | Download `model_best.pth` from this [link](https://drive.google.com/drive/folders/1tG3hBXtroHe4iXHb5W8XIfQ8YJEeS3Tp?usp=sharing) and save them. 50 | #### NOTE: The pretrained model is updated at April. 24th 2023. 51 | 52 | You can now go to src folder and test our ACL-SPC: 53 | ``` 54 | CUDA_VISIBLE_DEVICES=0 python main.py --experiment_id {experiment id} --dataset_name {dataset} --class_name {plane/car/chair/table} 55 | ``` 56 | 57 | 58 | ### Citation 59 | If you find our code or paper useful, please consider citing: 60 | ``` 61 | @inproceedings{Hong2023ACLSPC, 62 | title={ACL-SPC: Adaptive Closed-Loop system for Self-Supervised Point Cloud Completion}, 63 | author={Sangmin Hong and Mohsen Yavartanoo and Reyhaneh Neshatavar and Kyoung Mu Lee}, 64 | booktitle={IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}, 65 | year={2023} 66 | } 67 | ``` 68 | -------------------------------------------------------------------------------- /assets/NewFramework-1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/assets/NewFramework-1.png -------------------------------------------------------------------------------- /assets/README.md: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /outputs/.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore everything in this directory 2 | * 3 | # Except this file 4 | !.gitignore 5 | -------------------------------------------------------------------------------- /src/__pycache__/bunny_data.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/bunny_data.cpython-38.pyc -------------------------------------------------------------------------------- /src/__pycache__/data.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/data.cpython-36.pyc -------------------------------------------------------------------------------- /src/__pycache__/data.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/data.cpython-37.pyc -------------------------------------------------------------------------------- /src/__pycache__/data.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/data.cpython-38.pyc -------------------------------------------------------------------------------- /src/__pycache__/data_shapenet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/data_shapenet.cpython-38.pyc -------------------------------------------------------------------------------- /src/__pycache__/data_util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/data_util.cpython-38.pyc -------------------------------------------------------------------------------- /src/__pycache__/pretrain.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/pretrain.cpython-38.pyc -------------------------------------------------------------------------------- /src/__pycache__/train.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/train.cpython-36.pyc -------------------------------------------------------------------------------- /src/__pycache__/train.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/train.cpython-37.pyc -------------------------------------------------------------------------------- /src/__pycache__/train.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/train.cpython-38.pyc -------------------------------------------------------------------------------- /src/__pycache__/train_keypoint.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/train_keypoint.cpython-38.pyc -------------------------------------------------------------------------------- /src/__pycache__/train_real.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/train_real.cpython-38.pyc -------------------------------------------------------------------------------- /src/__pycache__/train_semkitti.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/train_semkitti.cpython-38.pyc -------------------------------------------------------------------------------- /src/__pycache__/util.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/util.cpython-36.pyc -------------------------------------------------------------------------------- /src/__pycache__/util.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/util.cpython-37.pyc -------------------------------------------------------------------------------- /src/__pycache__/util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/__pycache__/util.cpython-38.pyc -------------------------------------------------------------------------------- /src/chamfer_distance/__init__.py: -------------------------------------------------------------------------------- 1 | from .chamfer_distance_cpu import ChamferDistance 2 | -------------------------------------------------------------------------------- /src/chamfer_distance/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/chamfer_distance/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /src/chamfer_distance/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/chamfer_distance/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /src/chamfer_distance/__pycache__/chamfer_distance_cpu.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/chamfer_distance/__pycache__/chamfer_distance_cpu.cpython-36.pyc -------------------------------------------------------------------------------- /src/chamfer_distance/__pycache__/chamfer_distance_cpu.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/chamfer_distance/__pycache__/chamfer_distance_cpu.cpython-38.pyc -------------------------------------------------------------------------------- /src/chamfer_distance/__pycache__/chamfer_distance_gpu.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/chamfer_distance/__pycache__/chamfer_distance_gpu.cpython-36.pyc -------------------------------------------------------------------------------- /src/chamfer_distance/__pycache__/chamfer_distance_gpu.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/chamfer_distance/__pycache__/chamfer_distance_gpu.cpython-38.pyc -------------------------------------------------------------------------------- /src/chamfer_distance/chamfer_distance.cu: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | #include 5 | 6 | __global__ 7 | void ChamferDistanceKernel( 8 | int b, 9 | int n, 10 | const float* xyz, 11 | int m, 12 | const float* xyz2, 13 | float* result, 14 | int* result_i) 15 | { 16 | const int batch=512; 17 | __shared__ float buf[batch*3]; 18 | for (int i=blockIdx.x;ibest){ 130 | result[(i*n+j)]=best; 131 | result_i[(i*n+j)]=best_i; 132 | } 133 | } 134 | __syncthreads(); 135 | } 136 | } 137 | } 138 | 139 | void ChamferDistanceKernelLauncher( 140 | const int b, const int n, 141 | const float* xyz, 142 | const int m, 143 | const float* xyz2, 144 | float* result, 145 | int* result_i, 146 | float* result2, 147 | int* result2_i) 148 | { 149 | ChamferDistanceKernel<<>>(b, n, xyz, m, xyz2, result, result_i); 150 | ChamferDistanceKernel<<>>(b, m, xyz2, n, xyz, result2, result2_i); 151 | 152 | cudaError_t err = cudaGetLastError(); 153 | if (err != cudaSuccess) 154 | printf("error in chamfer distance updateOutput: %s\n", cudaGetErrorString(err)); 155 | } 156 | 157 | 158 | __global__ 159 | void ChamferDistanceGradKernel( 160 | int b, int n, 161 | const float* xyz1, 162 | int m, 163 | const float* xyz2, 164 | const float* grad_dist1, 165 | const int* idx1, 166 | float* grad_xyz1, 167 | float* grad_xyz2) 168 | { 169 | for (int i = blockIdx.x; i>>(b, n, xyz1, m, xyz2, grad_dist1, idx1, grad_xyz1, grad_xyz2); 204 | ChamferDistanceGradKernel<<>>(b, m, xyz2, n, xyz1, grad_dist2, idx2, grad_xyz2, grad_xyz1); 205 | 206 | cudaError_t err = cudaGetLastError(); 207 | if (err != cudaSuccess) 208 | printf("error in chamfer distance get grad: %s\n", cudaGetErrorString(err)); 209 | } 210 | -------------------------------------------------------------------------------- /src/chamfer_distance/chamfer_distance_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | // CUDA forward declarations 4 | int ChamferDistanceKernelLauncher( 5 | const int b, const int n, 6 | const float* xyz, 7 | const int m, 8 | const float* xyz2, 9 | float* result, 10 | int* result_i, 11 | float* result2, 12 | int* result2_i); 13 | 14 | int ChamferDistanceGradKernelLauncher( 15 | const int b, const int n, 16 | const float* xyz1, 17 | const int m, 18 | const float* xyz2, 19 | const float* grad_dist1, 20 | const int* idx1, 21 | const float* grad_dist2, 22 | const int* idx2, 23 | float* grad_xyz1, 24 | float* grad_xyz2); 25 | 26 | 27 | // void chamfer_distance_forward_cuda( 28 | // const at::Tensor xyz1, 29 | // const at::Tensor xyz2, 30 | // const at::Tensor dist1, 31 | // const at::Tensor dist2, 32 | // const at::Tensor idx1, 33 | // const at::Tensor idx2) 34 | // { 35 | // ChamferDistanceKernelLauncher(xyz1.size(0), xyz1.size(1), xyz1.data(), 36 | // xyz2.size(1), xyz2.data(), 37 | // dist1.data(), idx1.data(), 38 | // dist2.data(), idx2.data()); 39 | // } 40 | 41 | // void chamfer_distance_backward_cuda( 42 | // const at::Tensor xyz1, 43 | // const at::Tensor xyz2, 44 | // at::Tensor gradxyz1, 45 | // at::Tensor gradxyz2, 46 | // at::Tensor graddist1, 47 | // at::Tensor graddist2, 48 | // at::Tensor idx1, 49 | // at::Tensor idx2) 50 | // { 51 | // ChamferDistanceGradKernelLauncher(xyz1.size(0), xyz1.size(1), xyz1.data(), 52 | // xyz2.size(1), xyz2.data(), 53 | // graddist1.data(), idx1.data(), 54 | // graddist2.data(), idx2.data(), 55 | // gradxyz1.data(), gradxyz2.data()); 56 | // } 57 | 58 | 59 | void nnsearch( 60 | const int b, const int n, const int m, 61 | const float* xyz1, 62 | const float* xyz2, 63 | float* dist, 64 | int* idx) 65 | { 66 | for (int i = 0; i < b; i++) { 67 | for (int j = 0; j < n; j++) { 68 | const float x1 = xyz1[(i*n+j)*3+0]; 69 | const float y1 = xyz1[(i*n+j)*3+1]; 70 | const float z1 = xyz1[(i*n+j)*3+2]; 71 | double best = 0; 72 | int besti = 0; 73 | for (int k = 0; k < m; k++) { 74 | const float x2 = xyz2[(i*m+k)*3+0] - x1; 75 | const float y2 = xyz2[(i*m+k)*3+1] - y1; 76 | const float z2 = xyz2[(i*m+k)*3+2] - z1; 77 | const double d=x2*x2+y2*y2+z2*z2; 78 | if (k==0 || d < best){ 79 | best = d; 80 | besti = k; 81 | } 82 | } 83 | dist[i*n+j] = best; 84 | idx[i*n+j] = besti; 85 | } 86 | } 87 | } 88 | 89 | 90 | void chamfer_distance_forward( 91 | const at::Tensor xyz1, 92 | const at::Tensor xyz2, 93 | const at::Tensor dist1, 94 | const at::Tensor dist2, 95 | const at::Tensor idx1, 96 | const at::Tensor idx2) 97 | { 98 | const int batchsize = xyz1.size(0); 99 | const int n = xyz1.size(1); 100 | const int m = xyz2.size(1); 101 | 102 | const float* xyz1_data = xyz1.data(); 103 | const float* xyz2_data = xyz2.data(); 104 | float* dist1_data = dist1.data(); 105 | float* dist2_data = dist2.data(); 106 | int* idx1_data = idx1.data(); 107 | int* idx2_data = idx2.data(); 108 | 109 | nnsearch(batchsize, n, m, xyz1_data, xyz2_data, dist1_data, idx1_data); 110 | nnsearch(batchsize, m, n, xyz2_data, xyz1_data, dist2_data, idx2_data); 111 | } 112 | 113 | 114 | void chamfer_distance_backward( 115 | const at::Tensor xyz1, 116 | const at::Tensor xyz2, 117 | at::Tensor gradxyz1, 118 | at::Tensor gradxyz2, 119 | at::Tensor graddist1, 120 | at::Tensor graddist2, 121 | at::Tensor idx1, 122 | at::Tensor idx2) 123 | { 124 | const int b = xyz1.size(0); 125 | const int n = xyz1.size(1); 126 | const int m = xyz2.size(1); 127 | 128 | const float* xyz1_data = xyz1.data(); 129 | const float* xyz2_data = xyz2.data(); 130 | float* gradxyz1_data = gradxyz1.data(); 131 | float* gradxyz2_data = gradxyz2.data(); 132 | float* graddist1_data = graddist1.data(); 133 | float* graddist2_data = graddist2.data(); 134 | const int* idx1_data = idx1.data(); 135 | const int* idx2_data = idx2.data(); 136 | 137 | for (int i = 0; i < b*n*3; i++) 138 | gradxyz1_data[i] = 0; 139 | for (int i = 0; i < b*m*3; i++) 140 | gradxyz2_data[i] = 0; 141 | for (int i = 0;i < b; i++) { 142 | for (int j = 0; j < n; j++) { 143 | const float x1 = xyz1_data[(i*n+j)*3+0]; 144 | const float y1 = xyz1_data[(i*n+j)*3+1]; 145 | const float z1 = xyz1_data[(i*n+j)*3+2]; 146 | const int j2 = idx1_data[i*n+j]; 147 | 148 | const float x2 = xyz2_data[(i*m+j2)*3+0]; 149 | const float y2 = xyz2_data[(i*m+j2)*3+1]; 150 | const float z2 = xyz2_data[(i*m+j2)*3+2]; 151 | const float g = graddist1_data[i*n+j]*2; 152 | 153 | gradxyz1_data[(i*n+j)*3+0] += g*(x1-x2); 154 | gradxyz1_data[(i*n+j)*3+1] += g*(y1-y2); 155 | gradxyz1_data[(i*n+j)*3+2] += g*(z1-z2); 156 | gradxyz2_data[(i*m+j2)*3+0] -= (g*(x1-x2)); 157 | gradxyz2_data[(i*m+j2)*3+1] -= (g*(y1-y2)); 158 | gradxyz2_data[(i*m+j2)*3+2] -= (g*(z1-z2)); 159 | } 160 | for (int j = 0; j < m; j++) { 161 | const float x1 = xyz2_data[(i*m+j)*3+0]; 162 | const float y1 = xyz2_data[(i*m+j)*3+1]; 163 | const float z1 = xyz2_data[(i*m+j)*3+2]; 164 | const int j2 = idx2_data[i*m+j]; 165 | const float x2 = xyz1_data[(i*n+j2)*3+0]; 166 | const float y2 = xyz1_data[(i*n+j2)*3+1]; 167 | const float z2 = xyz1_data[(i*n+j2)*3+2]; 168 | const float g = graddist2_data[i*m+j]*2; 169 | gradxyz2_data[(i*m+j)*3+0] += g*(x1-x2); 170 | gradxyz2_data[(i*m+j)*3+1] += g*(y1-y2); 171 | gradxyz2_data[(i*m+j)*3+2] += g*(z1-z2); 172 | gradxyz1_data[(i*n+j2)*3+0] -= (g*(x1-x2)); 173 | gradxyz1_data[(i*n+j2)*3+1] -= (g*(y1-y2)); 174 | gradxyz1_data[(i*n+j2)*3+2] -= (g*(z1-z2)); 175 | } 176 | } 177 | } 178 | 179 | 180 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 181 | m.def("forward", &chamfer_distance_forward, "ChamferDistance forward"); 182 | // m.def("forward_cuda", &chamfer_distance_forward_cuda, "ChamferDistance forward (CUDA)"); 183 | m.def("backward", &chamfer_distance_backward, "ChamferDistance backward"); 184 | // m.def("backward_cuda", &chamfer_distance_backward_cuda, "ChamferDistance backward (CUDA)"); 185 | } 186 | -------------------------------------------------------------------------------- /src/chamfer_distance/chamfer_distance_cpu.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | 4 | from torch.utils.cpp_extension import load 5 | cd = load(name="cd", 6 | sources=["chamfer_distance/chamfer_distance_cpu.cpp" 7 | #,"chamfer_distance/chamfer_distance.cu" 8 | ], verbose=True) 9 | 10 | class ChamferDistanceFunction(torch.autograd.Function): 11 | @staticmethod 12 | def forward(ctx, xyz1, xyz2): 13 | batchsize, n, _ = xyz1.size() 14 | _, m, _ = xyz2.size() 15 | xyz1 = xyz1.contiguous() 16 | xyz2 = xyz2.contiguous() 17 | dist1 = torch.zeros(batchsize, n) 18 | dist2 = torch.zeros(batchsize, m) 19 | 20 | idx1 = torch.zeros(batchsize, n, dtype=torch.int) 21 | idx2 = torch.zeros(batchsize, m, dtype=torch.int) 22 | 23 | if not xyz1.is_cuda: 24 | cd.forward(xyz1, xyz2, dist1, dist2, idx1, idx2) 25 | # else: 26 | # dist1 = dist1.cuda() 27 | # dist2 = dist2.cuda() 28 | # idx1 = idx1.cuda() 29 | # idx2 = idx2.cuda() 30 | # cd.forward_cuda(xyz1, xyz2, dist1, dist2, idx1, idx2) 31 | 32 | ctx.save_for_backward(xyz1, xyz2, idx1, idx2) 33 | 34 | return dist1, dist2 35 | 36 | @staticmethod 37 | def backward(ctx, graddist1, graddist2): 38 | xyz1, xyz2, idx1, idx2 = ctx.saved_tensors 39 | 40 | graddist1 = graddist1.contiguous() 41 | graddist2 = graddist2.contiguous() 42 | 43 | gradxyz1 = torch.zeros(xyz1.size()) 44 | gradxyz2 = torch.zeros(xyz2.size()) 45 | 46 | if not graddist1.is_cuda: 47 | cd.backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2) 48 | # else: 49 | # gradxyz1 = gradxyz1.cuda() 50 | # gradxyz2 = gradxyz2.cuda() 51 | # cd.backward_cuda(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2) 52 | 53 | return gradxyz1, gradxyz2 54 | 55 | 56 | class ChamferDistance(torch.nn.Module): 57 | def forward(self, xyz1, xyz2): 58 | return ChamferDistanceFunction.apply(xyz1, xyz2) 59 | 60 | 61 | 62 | if __name__ == '__main__': 63 | 64 | 65 | import os 66 | import sys 67 | sys.path.insert(1, os.path.abspath('.')) 68 | 69 | from utils.read_xyz import getpts_fromXYZ 70 | 71 | batch_size = 1 72 | filename = '/mnt/hdd1/mchiash2/PUdataset/all_testset/4/input/camel.xyz' 73 | filename2 = '/mnt/hdd1/mchiash2/PUdataset/all_testset/4/gt/camel.xyz' 74 | preds = getpts_fromXYZ(filename) 75 | gts = getpts_fromXYZ(filename2) 76 | preds = torch.reshape(preds, (1,3, -1)) 77 | gts = torch.reshape(gts, (1,3, -1)) 78 | 79 | loss = ChamferDistance() 80 | 81 | dist1, dist2 = loss(preds,gts) 82 | 83 | train_loss = (torch.mean(dist1)) + (torch.mean(dist2)) 84 | 85 | print(train_loss) -------------------------------------------------------------------------------- /src/chamfer_distance/chamfer_distance_gpu.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | // CUDA forward declarations 4 | int ChamferDistanceKernelLauncher( 5 | const int b, const int n, 6 | const float* xyz, 7 | const int m, 8 | const float* xyz2, 9 | float* result, 10 | int* result_i, 11 | float* result2, 12 | int* result2_i); 13 | 14 | int ChamferDistanceGradKernelLauncher( 15 | const int b, const int n, 16 | const float* xyz1, 17 | const int m, 18 | const float* xyz2, 19 | const float* grad_dist1, 20 | const int* idx1, 21 | const float* grad_dist2, 22 | const int* idx2, 23 | float* grad_xyz1, 24 | float* grad_xyz2); 25 | 26 | 27 | void chamfer_distance_forward_cuda( 28 | const at::Tensor xyz1, 29 | const at::Tensor xyz2, 30 | const at::Tensor dist1, 31 | const at::Tensor dist2, 32 | const at::Tensor idx1, 33 | const at::Tensor idx2) 34 | { 35 | ChamferDistanceKernelLauncher(xyz1.size(0), xyz1.size(1), xyz1.data(), 36 | xyz2.size(1), xyz2.data(), 37 | dist1.data(), idx1.data(), 38 | dist2.data(), idx2.data()); 39 | } 40 | 41 | void chamfer_distance_backward_cuda( 42 | const at::Tensor xyz1, 43 | const at::Tensor xyz2, 44 | at::Tensor gradxyz1, 45 | at::Tensor gradxyz2, 46 | at::Tensor graddist1, 47 | at::Tensor graddist2, 48 | at::Tensor idx1, 49 | at::Tensor idx2) 50 | { 51 | ChamferDistanceGradKernelLauncher(xyz1.size(0), xyz1.size(1), xyz1.data(), 52 | xyz2.size(1), xyz2.data(), 53 | graddist1.data(), idx1.data(), 54 | graddist2.data(), idx2.data(), 55 | gradxyz1.data(), gradxyz2.data()); 56 | } 57 | 58 | 59 | void nnsearch( 60 | const int b, const int n, const int m, 61 | const float* xyz1, 62 | const float* xyz2, 63 | float* dist, 64 | int* idx) 65 | { 66 | for (int i = 0; i < b; i++) { 67 | for (int j = 0; j < n; j++) { 68 | const float x1 = xyz1[(i*n+j)*3+0]; 69 | const float y1 = xyz1[(i*n+j)*3+1]; 70 | const float z1 = xyz1[(i*n+j)*3+2]; 71 | double best = 0; 72 | int besti = 0; 73 | for (int k = 0; k < m; k++) { 74 | const float x2 = xyz2[(i*m+k)*3+0] - x1; 75 | const float y2 = xyz2[(i*m+k)*3+1] - y1; 76 | const float z2 = xyz2[(i*m+k)*3+2] - z1; 77 | const double d=x2*x2+y2*y2+z2*z2; 78 | if (k==0 || d < best){ 79 | best = d; 80 | besti = k; 81 | } 82 | } 83 | dist[i*n+j] = best; 84 | idx[i*n+j] = besti; 85 | } 86 | } 87 | } 88 | 89 | 90 | void chamfer_distance_forward( 91 | const at::Tensor xyz1, 92 | const at::Tensor xyz2, 93 | const at::Tensor dist1, 94 | const at::Tensor dist2, 95 | const at::Tensor idx1, 96 | const at::Tensor idx2) 97 | { 98 | const int batchsize = xyz1.size(0); 99 | const int n = xyz1.size(1); 100 | const int m = xyz2.size(1); 101 | 102 | const float* xyz1_data = xyz1.data(); 103 | const float* xyz2_data = xyz2.data(); 104 | float* dist1_data = dist1.data(); 105 | float* dist2_data = dist2.data(); 106 | int* idx1_data = idx1.data(); 107 | int* idx2_data = idx2.data(); 108 | 109 | nnsearch(batchsize, n, m, xyz1_data, xyz2_data, dist1_data, idx1_data); 110 | nnsearch(batchsize, m, n, xyz2_data, xyz1_data, dist2_data, idx2_data); 111 | } 112 | 113 | 114 | void chamfer_distance_backward( 115 | const at::Tensor xyz1, 116 | const at::Tensor xyz2, 117 | at::Tensor gradxyz1, 118 | at::Tensor gradxyz2, 119 | at::Tensor graddist1, 120 | at::Tensor graddist2, 121 | at::Tensor idx1, 122 | at::Tensor idx2) 123 | { 124 | const int b = xyz1.size(0); 125 | const int n = xyz1.size(1); 126 | const int m = xyz2.size(1); 127 | 128 | const float* xyz1_data = xyz1.data(); 129 | const float* xyz2_data = xyz2.data(); 130 | float* gradxyz1_data = gradxyz1.data(); 131 | float* gradxyz2_data = gradxyz2.data(); 132 | float* graddist1_data = graddist1.data(); 133 | float* graddist2_data = graddist2.data(); 134 | const int* idx1_data = idx1.data(); 135 | const int* idx2_data = idx2.data(); 136 | 137 | for (int i = 0; i < b*n*3; i++) 138 | gradxyz1_data[i] = 0; 139 | for (int i = 0; i < b*m*3; i++) 140 | gradxyz2_data[i] = 0; 141 | for (int i = 0;i < b; i++) { 142 | for (int j = 0; j < n; j++) { 143 | const float x1 = xyz1_data[(i*n+j)*3+0]; 144 | const float y1 = xyz1_data[(i*n+j)*3+1]; 145 | const float z1 = xyz1_data[(i*n+j)*3+2]; 146 | const int j2 = idx1_data[i*n+j]; 147 | 148 | const float x2 = xyz2_data[(i*m+j2)*3+0]; 149 | const float y2 = xyz2_data[(i*m+j2)*3+1]; 150 | const float z2 = xyz2_data[(i*m+j2)*3+2]; 151 | const float g = graddist1_data[i*n+j]*2; 152 | 153 | gradxyz1_data[(i*n+j)*3+0] += g*(x1-x2); 154 | gradxyz1_data[(i*n+j)*3+1] += g*(y1-y2); 155 | gradxyz1_data[(i*n+j)*3+2] += g*(z1-z2); 156 | gradxyz2_data[(i*m+j2)*3+0] -= (g*(x1-x2)); 157 | gradxyz2_data[(i*m+j2)*3+1] -= (g*(y1-y2)); 158 | gradxyz2_data[(i*m+j2)*3+2] -= (g*(z1-z2)); 159 | } 160 | for (int j = 0; j < m; j++) { 161 | const float x1 = xyz2_data[(i*m+j)*3+0]; 162 | const float y1 = xyz2_data[(i*m+j)*3+1]; 163 | const float z1 = xyz2_data[(i*m+j)*3+2]; 164 | const int j2 = idx2_data[i*m+j]; 165 | const float x2 = xyz1_data[(i*n+j2)*3+0]; 166 | const float y2 = xyz1_data[(i*n+j2)*3+1]; 167 | const float z2 = xyz1_data[(i*n+j2)*3+2]; 168 | const float g = graddist2_data[i*m+j]*2; 169 | gradxyz2_data[(i*m+j)*3+0] += g*(x1-x2); 170 | gradxyz2_data[(i*m+j)*3+1] += g*(y1-y2); 171 | gradxyz2_data[(i*m+j)*3+2] += g*(z1-z2); 172 | gradxyz1_data[(i*n+j2)*3+0] -= (g*(x1-x2)); 173 | gradxyz1_data[(i*n+j2)*3+1] -= (g*(y1-y2)); 174 | gradxyz1_data[(i*n+j2)*3+2] -= (g*(z1-z2)); 175 | } 176 | } 177 | } 178 | 179 | 180 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 181 | m.def("forward", &chamfer_distance_forward, "ChamferDistance forward"); 182 | m.def("forward_cuda", &chamfer_distance_forward_cuda, "ChamferDistance forward (CUDA)"); 183 | m.def("backward", &chamfer_distance_backward, "ChamferDistance backward"); 184 | m.def("backward_cuda", &chamfer_distance_backward_cuda, "ChamferDistance backward (CUDA)"); 185 | } 186 | -------------------------------------------------------------------------------- /src/chamfer_distance/chamfer_distance_gpu.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | 4 | from torch.utils.cpp_extension import load 5 | cd = load(name="cd", 6 | sources=["chamfer_distance/chamfer_distance_gpu.cpp" 7 | ,"chamfer_distance/chamfer_distance.cu" 8 | ]) 9 | 10 | class ChamferDistanceFunction(torch.autograd.Function): 11 | @staticmethod 12 | def forward(ctx, xyz1, xyz2): 13 | 14 | device = torch.device(xyz1.get_device()) 15 | torch.cuda.set_device(device) 16 | batchsize, n, _ = xyz1.size() 17 | _, m, _ = xyz2.size() 18 | xyz1 = xyz1.contiguous() 19 | xyz2 = xyz2.contiguous() 20 | dist1 = torch.zeros(batchsize, n) 21 | dist2 = torch.zeros(batchsize, m) 22 | 23 | idx1 = torch.zeros(batchsize, n, dtype=torch.int) 24 | idx2 = torch.zeros(batchsize, m, dtype=torch.int) 25 | 26 | if not xyz1.is_cuda: 27 | cd.forward(xyz1, xyz2, dist1, dist2, idx1, idx2) 28 | else: 29 | dist1 = dist1.cuda() 30 | dist2 = dist2.cuda() 31 | idx1 = idx1.cuda() 32 | idx2 = idx2.cuda() 33 | cd.forward_cuda(xyz1, xyz2, dist1, dist2, idx1, idx2) 34 | 35 | ctx.save_for_backward(xyz1, xyz2, idx1, idx2) 36 | 37 | return dist1, dist2 38 | 39 | @staticmethod 40 | def backward(ctx, graddist1, graddist2): 41 | 42 | xyz1, xyz2, idx1, idx2 = ctx.saved_tensors 43 | device = torch.device(xyz1.get_device()) 44 | torch.cuda.set_device(device) 45 | graddist1 = graddist1.contiguous() 46 | graddist2 = graddist2.contiguous() 47 | 48 | gradxyz1 = torch.zeros(xyz1.size()) 49 | gradxyz2 = torch.zeros(xyz2.size()) 50 | 51 | if not graddist1.is_cuda: 52 | cd.backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2) 53 | else: 54 | gradxyz1 = gradxyz1.cuda() 55 | gradxyz2 = gradxyz2.cuda() 56 | cd.backward_cuda(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2) 57 | 58 | return gradxyz1 ,gradxyz2 59 | 60 | 61 | class ChamferDistance(torch.nn.Module): 62 | def forward(self, xyz1, xyz2): 63 | return ChamferDistanceFunction.apply(xyz1, xyz2) 64 | 65 | 66 | if __name__== '__main__': 67 | 68 | import torch 69 | 70 | device = torch.device(5) 71 | 72 | pc1 = torch.Tensor([[[1,2,3],[4,5,6],[7,8,9],[10,11,12]]]).to(device) 73 | pc2 = torch.Tensor([[[4,5,6],[1,2,3],[7,8,9],[10,11,12],[100,1,2],[2,3,4], [3,4,5]]]).to(device) 74 | 75 | print(pc1.is_cuda) 76 | loss = ChamferDistance() 77 | dist1, dist2 = loss(pc1, pc2, device) 78 | print("dist1: ", dist1) 79 | print("dist2: ", dist2) 80 | -------------------------------------------------------------------------------- /src/data/RealDataset.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import torch.utils.data as data 3 | import os 4 | import os.path 5 | import torch 6 | import numpy as np 7 | from numpy.random import RandomState 8 | import h5py 9 | import random 10 | from tqdm import tqdm 11 | import pickle 12 | import glob 13 | from utils.io import read_ply_xyz, read_ply_from_file_list 14 | from utils.pc_transform import swap_axis 15 | from data.real_dataset import RealWorldPointsDataset 16 | from plyfile import PlyData 17 | 18 | def random_pose(): 19 | angle_y = -(np.random.uniform() * np.pi/3 - np.pi/9) 20 | angle_z = np.random.uniform() * 2 * np.pi 21 | Ry = np.array([[np.cos(angle_y), 0, np.sin(angle_y)], 22 | [0, 1, 0], 23 | [-np.sin(angle_y), 0, np.cos(angle_y)]]) 24 | Rz = np.array([[np.cos(angle_z), -np.sin(angle_z), 0], 25 | [np.sin(angle_z), np.cos(angle_z), 0], 26 | [0, 0, 1]]) 27 | R = np.matmul(Rz,Ry) 28 | R = torch.from_numpy(R).float() 29 | return R 30 | 31 | class RealDataset(data.Dataset): 32 | """ 33 | datasets that with Ply format 34 | without GT: MatterPort, ScanNet, KITTI 35 | Datasets provided by pcl2pcl 36 | with GT: PartNet, each subdir under args.dataset_path contains 37 | the partial shape raw.ply and complete shape ply-2048.txt. 38 | Dataset provided by MPC 39 | """ 40 | def __init__(self, args, class_choice, split): 41 | #self.dataset = args.dataset 42 | #self.dataset_path = args.dataset_path 43 | self.root = args.root 44 | self.dataset = args.dataset_name 45 | self.random_seed = 0 46 | self.rand_gen = RandomState(self.random_seed) 47 | 48 | if self.dataset in ['MatterPort', 'ScanNet', 'KITTI']: 49 | if self.dataset == 'ScanNet': 50 | REALDATASET = RealWorldPointsDataset(self.root+'real_data/combined/data/scannet_v2_'+class_choice+'s_aligned/point_cloud', batch_size=6, npoint=2048, shuffle=False, split=split, random_seed=0) 51 | elif self.dataset == 'MatterPort': 52 | if split in ['train', 'trainval']: 53 | REALDATASET = RealWorldPointsDataset(self.root+'real_data/combined/data/MatterPort_v1_'+class_choice+'_Yup_aligned/point_cloud', batch_size=6, npoint=2048, shuffle=False, split='test', random_seed=0) 54 | else: 55 | REALDATASET = RealWorldPointsDataset(self.root+'real_data/combined/data/MatterPort_v1_'+class_choice+'_Yup_aligned/point_cloud', batch_size=6, npoint=2048, shuffle=False, split=split, random_seed=0) 56 | elif self.dataset == 'KITTI': 57 | if split in ['train']: 58 | REALDATASET = KITTIDataset(self.root+'real_data/combined/data/KITTI_frustum_data_for_pcl2pcl/point_cloud_train/') 59 | elif split in ['test', 'val']: 60 | REALDATASET = KITTIDataset(self.root+'real_data/combined/data/KITTI_frustum_data_for_pcl2pcl/point_cloud_val/') 61 | input_ls = REALDATASET.point_clouds 62 | # swap axis as pcl2pcl and ShapeInversion have different canonical pose 63 | input_ls_swapped = [np.float32(swap_axis(itm, swap_mode='n210')) for itm in input_ls] 64 | self.input_ls = input_ls_swapped 65 | self.stems = range(len(self.input_ls)) 66 | 67 | elif self.dataset in ['PartNet']: 68 | pathnames = sorted(glob.glob(self.dataset_path+'/*')) 69 | basenames = [os.path.basename(itm) for itm in pathnames] 70 | 71 | self.stems = [int(itm) for itm in basenames] 72 | 73 | input_ls = [read_ply_xyz(os.path.join(itm,'raw.ply')) for itm in pathnames] 74 | gt_ls = [np.loadtxt(os.path.join(itm,'ply-2048.txt'),delimiter=';').astype(np.float32) for itm in pathnames] 75 | 76 | # swap axis as multimodal and ShapeInversion have different canonical pose 77 | self.input_ls = [swap_axis(itm, swap_mode='210') for itm in input_ls] 78 | self.gt_ls = [swap_axis(itm, swap_mode='210') for itm in gt_ls] 79 | else: 80 | raise NotImplementedError 81 | 82 | def __getitem__(self, index): 83 | if self.dataset in ['MatterPort','ScanNet','KITTI']: 84 | stem = self.stems[index] 85 | choice = self.rand_gen.choice(self.input_ls[index].shape[0], 2048, replace=True) 86 | input_pcd = self.input_ls[index][choice,:] 87 | R = torch.cat([random_pose().unsqueeze(0) for i in range(8)],0) 88 | return (input_pcd, stem, R) 89 | elif self.dataset in ['PartNet']: 90 | stem = self.stems[index] 91 | input_pcd = self.input_ls[index] 92 | gt_pcd = self.gt_ls[index] 93 | return (gt_pcd, input_pcd, stem) 94 | 95 | def __len__(self): 96 | return len(self.input_ls) 97 | 98 | 99 | class KITTIDataset(): 100 | def __init__(self, load_path): 101 | self.point_clouds = [] 102 | file_list = glob.glob(load_path + '*.ply') 103 | total_num = len(file_list) 104 | for i in range(total_num): 105 | file_name = load_path + str(i) + '.ply' 106 | ply_file = PlyData.read(file_name) 107 | pc = np.array([ply_file['vertex']['x'], ply_file['vertex']['y'], ply_file['vertex']['z']]) 108 | pc = np.transpose(pc,(1,0)) 109 | self.point_clouds.append(pc) 110 | return 111 | 112 | 113 | 114 | -------------------------------------------------------------------------------- /src/data/__pycache__/RealDataset.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/data/__pycache__/RealDataset.cpython-38.pyc -------------------------------------------------------------------------------- /src/data/__pycache__/data.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/data/__pycache__/data.cpython-38.pyc -------------------------------------------------------------------------------- /src/data/__pycache__/real_dataset.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/data/__pycache__/real_dataset.cpython-38.pyc -------------------------------------------------------------------------------- /src/extensions/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/__init__.py -------------------------------------------------------------------------------- /src/extensions/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /src/extensions/chamfer_dist/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Thibault GROUEIX 3 | # @Date: 2019-08-07 20:54:24 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-18 15:06:25 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | import chamfer 11 | 12 | 13 | class ChamferFunction(torch.autograd.Function): 14 | @staticmethod 15 | def forward(ctx, xyz1, xyz2): 16 | dist1, dist2, idx1, idx2 = chamfer.forward(xyz1, xyz2) 17 | ctx.save_for_backward(xyz1, xyz2, idx1, idx2) 18 | 19 | return dist1, dist2 20 | 21 | @staticmethod 22 | def backward(ctx, grad_dist1, grad_dist2): 23 | xyz1, xyz2, idx1, idx2 = ctx.saved_tensors 24 | grad_xyz1, grad_xyz2 = chamfer.backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2) 25 | return grad_xyz1, grad_xyz2 26 | 27 | 28 | class ChamferDistance(torch.nn.Module): 29 | def __init__(self, ignore_zeros=False): 30 | super(ChamferDistance, self).__init__() 31 | self.ignore_zeros = ignore_zeros 32 | 33 | def forward(self, xyz1, xyz2): 34 | batch_size = xyz1.size(0) 35 | if batch_size == 1 and self.ignore_zeros: 36 | non_zeros1 = torch.sum(xyz1, dim=2).ne(0) 37 | non_zeros2 = torch.sum(xyz2, dim=2).ne(0) 38 | xyz1 = xyz1[non_zeros1].unsqueeze(dim=0) 39 | xyz2 = xyz2[non_zeros2].unsqueeze(dim=0) 40 | 41 | dist1, dist2 = ChamferFunction.apply(xyz1, xyz2) 42 | return torch.mean(dist1) + torch.mean(dist2) 43 | -------------------------------------------------------------------------------- /src/extensions/chamfer_dist/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/chamfer_dist/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /src/extensions/chamfer_dist/chamfer_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-08-07 20:54:24 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2019-12-10 10:33:50 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | std::vector chamfer_cuda_forward(torch::Tensor xyz1, 13 | torch::Tensor xyz2); 14 | 15 | std::vector chamfer_cuda_backward(torch::Tensor xyz1, 16 | torch::Tensor xyz2, 17 | torch::Tensor idx1, 18 | torch::Tensor idx2, 19 | torch::Tensor grad_dist1, 20 | torch::Tensor grad_dist2); 21 | 22 | std::vector chamfer_forward(torch::Tensor xyz1, 23 | torch::Tensor xyz2) { 24 | return chamfer_cuda_forward(xyz1, xyz2); 25 | } 26 | 27 | std::vector chamfer_backward(torch::Tensor xyz1, 28 | torch::Tensor xyz2, 29 | torch::Tensor idx1, 30 | torch::Tensor idx2, 31 | torch::Tensor grad_dist1, 32 | torch::Tensor grad_dist2) { 33 | return chamfer_cuda_backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2); 34 | } 35 | 36 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 37 | m.def("forward", &chamfer_forward, "Chamfer forward (CUDA)"); 38 | m.def("backward", &chamfer_backward, "Chamfer backward (CUDA)"); 39 | } 40 | -------------------------------------------------------------------------------- /src/extensions/chamfer_dist/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-08-07 20:54:24 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-10 10:04:25 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup(name='chamfer', 12 | version='2.0.0', 13 | ext_modules=[ 14 | CUDAExtension('chamfer', [ 15 | 'chamfer_cuda.cpp', 16 | 'chamfer.cu', 17 | ]), 18 | ], 19 | cmdclass={'build_ext': BuildExtension}) 20 | -------------------------------------------------------------------------------- /src/extensions/chamfer_dist/test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-10 10:38:01 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 14:21:36 6 | # @Email: cshzxie@gmail.com 7 | # 8 | # Note: 9 | # - Replace float -> double, kFloat -> kDouble in chamfer.cu 10 | 11 | import os 12 | import sys 13 | import torch 14 | import unittest 15 | 16 | from torch.autograd import gradcheck 17 | 18 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) 19 | from extensions.chamfer_dist import ChamferFunction 20 | 21 | 22 | class ChamferDistanceTestCase(unittest.TestCase): 23 | def test_chamfer_dist(self): 24 | x = torch.rand(4, 64, 3).double() 25 | y = torch.rand(4, 128, 3).double() 26 | x.requires_grad = True 27 | y.requires_grad = True 28 | print(gradcheck(ChamferFunction.apply, [x.cuda(), y.cuda()])) 29 | 30 | 31 | if __name__ == '__main__': 32 | unittest.main() 33 | -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-19 16:55:15 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 13:15:14 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | import cubic_feature_sampling 11 | 12 | 13 | class CubicFeatureSamplingFunction(torch.autograd.Function): 14 | @staticmethod 15 | def forward(ctx, ptcloud, cubic_features, neighborhood_size=1): 16 | scale = cubic_features.size(2) 17 | point_features, grid_pt_indexes = cubic_feature_sampling.forward(scale, neighborhood_size, ptcloud, 18 | cubic_features) 19 | ctx.save_for_backward(torch.Tensor([scale]), torch.Tensor([neighborhood_size]), grid_pt_indexes) 20 | return point_features 21 | 22 | @staticmethod 23 | def backward(ctx, grad_point_features): 24 | scale, neighborhood_size, grid_pt_indexes = ctx.saved_tensors 25 | scale = int(scale.item()) 26 | neighborhood_size = int(neighborhood_size.item()) 27 | grad_point_features = grad_point_features.contiguous() 28 | grad_ptcloud, grad_cubic_features = cubic_feature_sampling.backward(scale, neighborhood_size, 29 | grad_point_features, grid_pt_indexes) 30 | return grad_ptcloud, grad_cubic_features, None 31 | 32 | 33 | class CubicFeatureSampling(torch.nn.Module): 34 | def __init__(self): 35 | super(CubicFeatureSampling, self).__init__() 36 | 37 | def forward(self, ptcloud, cubic_features, neighborhood_size=1): 38 | h_scale = cubic_features.size(2) / 2 39 | ptcloud = ptcloud * h_scale + h_scale 40 | return CubicFeatureSamplingFunction.apply(ptcloud, cubic_features, neighborhood_size) 41 | -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/cubic_feature_sampling/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/build/lib.linux-x86_64-3.8/cubic_feature_sampling.cpython-38-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/cubic_feature_sampling/build/lib.linux-x86_64-3.8/cubic_feature_sampling.cpython-38-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/build/temp.linux-x86_64-3.8/.ninja_deps: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/cubic_feature_sampling/build/temp.linux-x86_64-3.8/.ninja_deps -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/build/temp.linux-x86_64-3.8/.ninja_log: -------------------------------------------------------------------------------- 1 | # ninja log v5 2 | 0 18801 1675063635726647629 /home/mchiash2/ACL-SPC/src_shapenet_simple/extensions/cubic_feature_sampling/build/temp.linux-x86_64-3.8/cubic_feature_sampling_cuda.o b962bd7e04734ab0 3 | 0 58130 1675063675054603318 /home/mchiash2/ACL-SPC/src_shapenet_simple/extensions/cubic_feature_sampling/build/temp.linux-x86_64-3.8/cubic_feature_sampling.o 958f5910033e1fa7 4 | -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/build/temp.linux-x86_64-3.8/build.ninja: -------------------------------------------------------------------------------- 1 | ninja_required_version = 1.3 2 | cxx = c++ 3 | nvcc = /usr/local/cuda-10.0/bin/nvcc 4 | 5 | cflags = -pthread -B /home/mchiash2/anaconda3/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include/TH -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include/THC -I/usr/local/cuda-10.0/include -I/home/mchiash2/anaconda3/include/python3.8 -c 6 | post_cflags = -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=cubic_feature_sampling -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++14 7 | cuda_cflags = -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include/TH -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include/THC -I/usr/local/cuda-10.0/include -I/home/mchiash2/anaconda3/include/python3.8 -c 8 | cuda_post_cflags = -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options ''"'"'-fPIC'"'"'' -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=cubic_feature_sampling -D_GLIBCXX_USE_CXX11_ABI=0 -gencode=arch=compute_75,code=sm_75 -std=c++14 9 | ldflags = 10 | 11 | rule compile 12 | command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags 13 | depfile = $out.d 14 | deps = gcc 15 | 16 | rule cuda_compile 17 | command = $nvcc $cuda_cflags -c $in -o $out $cuda_post_cflags 18 | 19 | 20 | 21 | build /home/mchiash2/ACL-SPC/src_shapenet_simple/extensions/cubic_feature_sampling/build/temp.linux-x86_64-3.8/cubic_feature_sampling.o: cuda_compile /home/mchiash2/ACL-SPC/src_shapenet_simple/extensions/cubic_feature_sampling/cubic_feature_sampling.cu 22 | build /home/mchiash2/ACL-SPC/src_shapenet_simple/extensions/cubic_feature_sampling/build/temp.linux-x86_64-3.8/cubic_feature_sampling_cuda.o: compile /home/mchiash2/ACL-SPC/src_shapenet_simple/extensions/cubic_feature_sampling/cubic_feature_sampling_cuda.cpp 23 | 24 | 25 | 26 | 27 | 28 | -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/build/temp.linux-x86_64-3.8/cubic_feature_sampling.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/cubic_feature_sampling/build/temp.linux-x86_64-3.8/cubic_feature_sampling.o -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/build/temp.linux-x86_64-3.8/cubic_feature_sampling_cuda.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/cubic_feature_sampling/build/temp.linux-x86_64-3.8/cubic_feature_sampling_cuda.o -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/cubic_feature_sampling.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-12-19 20:36:36 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2020-06-17 14:55:41 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | #define CUDA_NUM_THREADS 512 15 | 16 | // Computer the number of threads needed in GPU 17 | inline int get_n_threads(int n) { 18 | const int pow_2 = std::log(static_cast(n)) / std::log(2.0); 19 | return max(min(1 << pow_2, CUDA_NUM_THREADS), 1); 20 | } 21 | 22 | __device__ int compute_index(int offset_x, 23 | int offset_y, 24 | int offset_z, 25 | int scale) { 26 | return offset_x * scale * scale + offset_y * scale + offset_z; 27 | } 28 | 29 | __global__ void cubic_feature_sampling_kernel( 30 | int scale, 31 | int neighborhood_size, 32 | int n_vertices, 33 | int n_pts, 34 | int n_cubic_channels, 35 | const float *__restrict__ ptcloud, 36 | const float *__restrict__ cubic_features, 37 | float *__restrict__ point_features, 38 | int *__restrict__ grid_pt_indexes) { 39 | int batch_index = blockIdx.x; 40 | int index = threadIdx.x; 41 | int stride = blockDim.x; 42 | int cub_scale = scale * scale * scale; 43 | 44 | ptcloud += batch_index * n_pts * 3; 45 | cubic_features += batch_index * n_cubic_channels * cub_scale; 46 | point_features += batch_index * n_pts * n_vertices * n_cubic_channels; 47 | grid_pt_indexes += batch_index * n_pts * n_vertices; 48 | 49 | for (int i = index; i < n_pts; i += stride) { 50 | float pt_x = ptcloud[i * 3 + 0]; 51 | float pt_y = ptcloud[i * 3 + 1]; 52 | float pt_z = ptcloud[i * 3 + 2]; 53 | 54 | int lower_x = std::floor(pt_x); 55 | int upper_x = std::ceil(pt_x); 56 | if (lower_x == upper_x) { 57 | upper_x += 1; 58 | } 59 | int lower_y = std::floor(pt_y); 60 | int upper_y = std::ceil(pt_y); 61 | if (lower_y == upper_y) { 62 | upper_y += 1; 63 | } 64 | int lower_z = std::floor(pt_z); 65 | int upper_z = std::ceil(pt_z); 66 | if (lower_z == upper_z) { 67 | upper_z += 1; 68 | } 69 | 70 | int ns = neighborhood_size - 1; 71 | int vertex_idx = 0; 72 | for (int j = lower_x - ns; j <= upper_x + ns; ++j) { 73 | for (int k = lower_y - ns; k <= upper_y + ns; ++k) { 74 | for (int m = lower_z - ns; m <= upper_z + ns; ++m) { 75 | if (j < 0 || j >= scale || k < 0 || k >= scale || m < 0 || 76 | m >= scale) { 77 | // Ignore points lies out of the grid 78 | grid_pt_indexes[i * n_vertices + vertex_idx++] = -1; 79 | } else { 80 | // Calcuating indexes for adjacent vertices 81 | grid_pt_indexes[i * n_vertices + vertex_idx++] = 82 | compute_index(j, k, m, scale); 83 | } 84 | } 85 | } 86 | } 87 | 88 | // Gather Features 89 | for (int j = 0; j < n_vertices; ++j) { 90 | for (int k = 0; k < n_cubic_channels; ++k) { 91 | int vertex_idx = grid_pt_indexes[i * n_vertices + j]; 92 | if (vertex_idx == -1) { 93 | continue; 94 | } 95 | int feature_idx = 96 | i * n_vertices * n_cubic_channels + j * n_cubic_channels + k; 97 | float feature_val = cubic_features[k * cub_scale + vertex_idx]; 98 | point_features[feature_idx] = feature_val; 99 | } 100 | } 101 | } 102 | } 103 | 104 | std::vector cubic_feature_sampling_cuda_forward( 105 | int scale, 106 | int neighborhood_size, 107 | torch::Tensor ptcloud, 108 | torch::Tensor cubic_features, 109 | cudaStream_t stream) { 110 | int batch_size = ptcloud.size(0); 111 | int n_pts = ptcloud.size(1); 112 | int n_cubic_channels = cubic_features.size(1); 113 | 114 | int n_vertices = std::pow(neighborhood_size * 2, 3); 115 | torch::Tensor point_features = 116 | torch::zeros({batch_size, n_pts, n_vertices, n_cubic_channels}, 117 | torch::CUDA(torch::kFloat)); 118 | torch::Tensor grid_pt_indexes = 119 | torch::zeros({batch_size, n_pts, n_vertices}, torch::CUDA(torch::kInt)); 120 | 121 | cubic_feature_sampling_kernel<<>>( 123 | scale, neighborhood_size, n_vertices, n_pts, n_cubic_channels, 124 | ptcloud.data_ptr(), cubic_features.data_ptr(), 125 | point_features.data_ptr(), grid_pt_indexes.data_ptr()); 126 | 127 | cudaError_t err = cudaGetLastError(); 128 | if (err != cudaSuccess) { 129 | printf("Error in cubic_feature_sampling_cuda_forward: %s\n", 130 | cudaGetErrorString(err)); 131 | } 132 | return {point_features, grid_pt_indexes}; 133 | } 134 | 135 | __global__ void cubic_feature_sampling_grad_kernel( 136 | int scale, 137 | int neighborhood_size, 138 | int n_vertices, 139 | int n_pts, 140 | int n_cubic_channels, 141 | const float *__restrict__ grad_point_features, 142 | const int *__restrict__ grid_pt_indexes, 143 | float *__restrict__ grad_ptcloud, 144 | float *__restrict__ grad_cubic_features) { 145 | int batch_index = blockIdx.x; 146 | int index = threadIdx.x; 147 | int stride = blockDim.x; 148 | int cub_scale = scale * scale * scale; 149 | 150 | grad_point_features += batch_index * n_pts * n_vertices * n_cubic_channels; 151 | grid_pt_indexes += batch_index * n_pts * n_vertices; 152 | grad_ptcloud += batch_index * n_pts * 3; 153 | grad_cubic_features += batch_index * n_cubic_channels * cub_scale; 154 | 155 | for (int i = index; i < n_pts; i += stride) { 156 | for (int j = 0; j < n_vertices; ++j) { 157 | int vertex_idx = grid_pt_indexes[i * n_vertices + j]; 158 | if (vertex_idx == -1) { 159 | continue; 160 | } 161 | for (int k = 0; k < n_cubic_channels; ++k) { 162 | int grad_idx = 163 | i * n_vertices * n_cubic_channels + j * n_cubic_channels + k; 164 | float grad_val = grad_point_features[grad_idx]; 165 | // Fix bugs: the gradients of ceil and floor functions are zeros. 166 | // Ref: https://github.com/tensorflow/tensorflow/issues/897 167 | // atomicAdd(&(grad_ptcloud[i * 3 + 0]), grad_val); 168 | // atomicAdd(&(grad_ptcloud[i * 3 + 1]), grad_val); 169 | // atomicAdd(&(grad_ptcloud[i * 3 + 2]), grad_val); 170 | atomicAdd(&(grad_cubic_features[k * cub_scale + vertex_idx]), grad_val); 171 | } 172 | } 173 | } 174 | } 175 | 176 | std::vector cubic_feature_sampling_cuda_backward( 177 | int scale, 178 | int neighborhood_size, 179 | torch::Tensor grad_point_features, 180 | torch::Tensor grid_pt_indexes, 181 | cudaStream_t stream) { 182 | int batch_size = grad_point_features.size(0); 183 | int n_cubic_channels = grad_point_features.size(3); 184 | int n_pts = grid_pt_indexes.size(1); 185 | int n_vertices = std::pow(neighborhood_size * 2, 3); 186 | 187 | torch::Tensor grad_ptcloud = 188 | torch::zeros({batch_size, n_pts, 3}, torch::CUDA(torch::kFloat)); 189 | torch::Tensor grad_cubic_features = 190 | torch::zeros({batch_size, n_cubic_channels, scale, scale, scale}, 191 | torch::CUDA(torch::kFloat)); 192 | 193 | cubic_feature_sampling_grad_kernel<<>>( 195 | scale, neighborhood_size, n_vertices, n_pts, n_cubic_channels, 196 | grad_point_features.data_ptr(), grid_pt_indexes.data_ptr(), 197 | grad_ptcloud.data_ptr(), grad_cubic_features.data_ptr()); 198 | 199 | cudaError_t err = cudaGetLastError(); 200 | if (err != cudaSuccess) { 201 | printf("Error in cubic_feature_sampling_cuda_backward: %s\n", 202 | cudaGetErrorString(err)); 203 | } 204 | return {grad_ptcloud, grad_cubic_features}; 205 | } -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 2.1 2 | Name: cubic-feature-sampling 3 | Version: 1.1.0 4 | Summary: UNKNOWN 5 | License: UNKNOWN 6 | Platform: UNKNOWN 7 | 8 | UNKNOWN 9 | 10 | -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | cubic_feature_sampling.cu 2 | cubic_feature_sampling_cuda.cpp 3 | setup.py 4 | cubic_feature_sampling.egg-info/PKG-INFO 5 | cubic_feature_sampling.egg-info/SOURCES.txt 6 | cubic_feature_sampling.egg-info/dependency_links.txt 7 | cubic_feature_sampling.egg-info/top_level.txt -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/cubic_feature_sampling.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | cubic_feature_sampling 2 | -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/cubic_feature_sampling_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-12-19 17:04:38 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2020-06-17 14:50:22 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | // NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. 13 | #define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") 14 | #define CHECK_CONTIGUOUS(x) \ 15 | AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) \ 17 | CHECK_CUDA(x); \ 18 | CHECK_CONTIGUOUS(x) 19 | 20 | std::vector cubic_feature_sampling_cuda_forward( 21 | int scale, 22 | int neighborhood_size, 23 | torch::Tensor ptcloud, 24 | torch::Tensor cubic_features, 25 | cudaStream_t stream); 26 | 27 | std::vector cubic_feature_sampling_cuda_backward( 28 | int scale, 29 | int neighborhood_size, 30 | torch::Tensor grad_point_features, 31 | torch::Tensor grid_pt_indexes, 32 | cudaStream_t stream); 33 | 34 | std::vector cubic_feature_sampling_forward( 35 | int scale, 36 | int neighborhood_size, 37 | torch::Tensor ptcloud, 38 | torch::Tensor cubic_features) { 39 | CHECK_INPUT(ptcloud); 40 | CHECK_INPUT(cubic_features); 41 | 42 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 43 | return cubic_feature_sampling_cuda_forward(scale, neighborhood_size, ptcloud, 44 | cubic_features, stream); 45 | } 46 | 47 | std::vector cubic_feature_sampling_backward( 48 | int scale, 49 | int neighborhood_size, 50 | torch::Tensor grad_point_features, 51 | torch::Tensor grid_pt_indexes) { 52 | CHECK_INPUT(grad_point_features); 53 | CHECK_INPUT(grid_pt_indexes); 54 | 55 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 56 | return cubic_feature_sampling_cuda_backward( 57 | scale, neighborhood_size, grad_point_features, grid_pt_indexes, stream); 58 | } 59 | 60 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 61 | m.def("forward", &cubic_feature_sampling_forward, 62 | "Cubic Feature Sampling forward (CUDA)"); 63 | m.def("backward", &cubic_feature_sampling_backward, 64 | "Cubic Feature Sampling backward (CUDA)"); 65 | } -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py3.8-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/cubic_feature_sampling/dist/cubic_feature_sampling-1.1.0-py3.8-linux-x86_64.egg -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-19 17:03:06 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 14:02:06 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup(name='cubic_feature_sampling', 12 | version='1.1.0', 13 | ext_modules=[ 14 | CUDAExtension('cubic_feature_sampling', ['cubic_feature_sampling_cuda.cpp', 'cubic_feature_sampling.cu']), 15 | ], 16 | cmdclass={'build_ext': BuildExtension}) 17 | -------------------------------------------------------------------------------- /src/extensions/cubic_feature_sampling/test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-20 11:50:50 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 13:52:33 6 | # @Email: cshzxie@gmail.com 7 | # 8 | # Note: 9 | # - Replace float -> double, kFloat -> kDouble in cubic_feature_sampling.cu 10 | 11 | import os 12 | import sys 13 | import torch 14 | import unittest 15 | 16 | from torch.autograd import gradcheck 17 | 18 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) 19 | from extensions.cubic_feature_sampling import CubicFeatureSamplingFunction 20 | 21 | 22 | class CubicFeatureSamplingTestCase(unittest.TestCase): 23 | def test_neighborhood_size_1(self): 24 | ptcloud = torch.rand(2, 64, 3) * 2 - 1 25 | cubic_features = torch.rand(2, 4, 8, 8, 8) 26 | ptcloud.requires_grad = True 27 | cubic_features.requires_grad = True 28 | self.assertTrue( 29 | gradcheck(CubicFeatureSamplingFunction.apply, 30 | [ptcloud.double().cuda(), cubic_features.double().cuda()])) 31 | 32 | def test_neighborhood_size_2(self): 33 | ptcloud = torch.rand(2, 32, 3) * 2 - 1 34 | cubic_features = torch.rand(2, 2, 8, 8, 8) 35 | ptcloud.requires_grad = True 36 | cubic_features.requires_grad = True 37 | self.assertTrue( 38 | gradcheck(CubicFeatureSamplingFunction.apply, 39 | [ptcloud.double().cuda(), cubic_features.double().cuda(), 2])) 40 | 41 | def test_neighborhood_size_3(self): 42 | ptcloud = torch.rand(1, 32, 3) * 2 - 1 43 | cubic_features = torch.rand(1, 2, 16, 16, 16) 44 | ptcloud.requires_grad = True 45 | cubic_features.requires_grad = True 46 | self.assertTrue( 47 | gradcheck(CubicFeatureSamplingFunction.apply, 48 | [ptcloud.double().cuda(), cubic_features.double().cuda(), 3])) 49 | 50 | 51 | if __name__ == '__main__': 52 | unittest.main() 53 | -------------------------------------------------------------------------------- /src/extensions/gridding/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-11-15 20:33:52 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-30 09:55:53 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | import gridding 11 | 12 | 13 | class GriddingFunction(torch.autograd.Function): 14 | @staticmethod 15 | def forward(ctx, scale, ptcloud): 16 | grid, grid_pt_weights, grid_pt_indexes = gridding.forward(-scale, scale - 1, -scale, scale - 1, -scale, 17 | scale - 1, ptcloud) 18 | # print(grid.size()) # torch.Size(batch_size, n_grid_vertices) 19 | # print(grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) 20 | # print(grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) 21 | ctx.save_for_backward(grid_pt_weights, grid_pt_indexes) 22 | 23 | return grid 24 | 25 | @staticmethod 26 | def backward(ctx, grad_grid): 27 | grid_pt_weights, grid_pt_indexes = ctx.saved_tensors 28 | grad_ptcloud = gridding.backward(grid_pt_weights, grid_pt_indexes, grad_grid) 29 | # print(grad_ptcloud.size()) # torch.Size(batch_size, n_pts, 3) 30 | 31 | return None, grad_ptcloud 32 | 33 | 34 | class Gridding(torch.nn.Module): 35 | def __init__(self, scale=1): 36 | super(Gridding, self).__init__() 37 | self.scale = scale // 2 38 | 39 | def forward(self, ptcloud): 40 | ptcloud = ptcloud * self.scale 41 | _ptcloud = torch.split(ptcloud, 1, dim=0) 42 | grids = [] 43 | for p in _ptcloud: 44 | non_zeros = torch.sum(p, dim=2).ne(0) 45 | p = p[non_zeros].unsqueeze(dim=0) 46 | grids.append(GriddingFunction.apply(self.scale, p)) 47 | 48 | return torch.cat(grids, dim=0).contiguous() 49 | 50 | 51 | class GriddingReverseFunction(torch.autograd.Function): 52 | @staticmethod 53 | def forward(ctx, scale, grid): 54 | ptcloud = gridding.rev_forward(scale, grid) 55 | ctx.save_for_backward(torch.Tensor([scale]), grid, ptcloud) 56 | return ptcloud 57 | 58 | @staticmethod 59 | def backward(ctx, grad_ptcloud): 60 | scale, grid, ptcloud = ctx.saved_tensors 61 | scale = int(scale.item()) 62 | grad_grid = gridding.rev_backward(ptcloud, grid, grad_ptcloud) 63 | grad_grid = grad_grid.view(-1, scale, scale, scale) 64 | return None, grad_grid 65 | 66 | 67 | class GriddingReverse(torch.nn.Module): 68 | def __init__(self, scale=1): 69 | super(GriddingReverse, self).__init__() 70 | self.scale = scale 71 | 72 | def forward(self, grid): 73 | ptcloud = GriddingReverseFunction.apply(self.scale, grid) 74 | return ptcloud / self.scale * 2 75 | -------------------------------------------------------------------------------- /src/extensions/gridding/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/gridding/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /src/extensions/gridding/build/lib.linux-x86_64-3.8/gridding.cpython-38-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/gridding/build/lib.linux-x86_64-3.8/gridding.cpython-38-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /src/extensions/gridding/build/temp.linux-x86_64-3.8/.ninja_deps: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/gridding/build/temp.linux-x86_64-3.8/.ninja_deps -------------------------------------------------------------------------------- /src/extensions/gridding/build/temp.linux-x86_64-3.8/.ninja_log: -------------------------------------------------------------------------------- 1 | # ninja log v5 2 | 1 19245 1675062624582204376 /home/mchiash2/ACL-SPC/GRNet/extensions/gridding/build/temp.linux-x86_64-3.8/gridding_cuda.o 711cc07d7cfedae 3 | 0 55866 1675062661206073355 /home/mchiash2/ACL-SPC/GRNet/extensions/gridding/build/temp.linux-x86_64-3.8/gridding.o f80859354c358fd8 4 | 1 59583 1675062664922060013 /home/mchiash2/ACL-SPC/GRNet/extensions/gridding/build/temp.linux-x86_64-3.8/gridding_reverse.o 87523a610059b491 5 | -------------------------------------------------------------------------------- /src/extensions/gridding/build/temp.linux-x86_64-3.8/build.ninja: -------------------------------------------------------------------------------- 1 | ninja_required_version = 1.3 2 | cxx = c++ 3 | nvcc = /usr/local/cuda-10.0/bin/nvcc 4 | 5 | cflags = -pthread -B /home/mchiash2/anaconda3/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include/TH -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include/THC -I/usr/local/cuda-10.0/include -I/home/mchiash2/anaconda3/include/python3.8 -c 6 | post_cflags = -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=gridding -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++14 7 | cuda_cflags = -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include/TH -I/home/mchiash2/anaconda3/lib/python3.8/site-packages/torch/include/THC -I/usr/local/cuda-10.0/include -I/home/mchiash2/anaconda3/include/python3.8 -c 8 | cuda_post_cflags = -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options ''"'"'-fPIC'"'"'' -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=gridding -D_GLIBCXX_USE_CXX11_ABI=0 -gencode=arch=compute_75,code=sm_75 -std=c++14 9 | ldflags = 10 | 11 | rule compile 12 | command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags 13 | depfile = $out.d 14 | deps = gcc 15 | 16 | rule cuda_compile 17 | command = $nvcc $cuda_cflags -c $in -o $out $cuda_post_cflags 18 | 19 | 20 | 21 | build /home/mchiash2/ACL-SPC/GRNet/extensions/gridding/build/temp.linux-x86_64-3.8/gridding.o: cuda_compile /home/mchiash2/ACL-SPC/GRNet/extensions/gridding/gridding.cu 22 | build /home/mchiash2/ACL-SPC/GRNet/extensions/gridding/build/temp.linux-x86_64-3.8/gridding_cuda.o: compile /home/mchiash2/ACL-SPC/GRNet/extensions/gridding/gridding_cuda.cpp 23 | build /home/mchiash2/ACL-SPC/GRNet/extensions/gridding/build/temp.linux-x86_64-3.8/gridding_reverse.o: cuda_compile /home/mchiash2/ACL-SPC/GRNet/extensions/gridding/gridding_reverse.cu 24 | 25 | 26 | 27 | 28 | 29 | -------------------------------------------------------------------------------- /src/extensions/gridding/build/temp.linux-x86_64-3.8/gridding.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/gridding/build/temp.linux-x86_64-3.8/gridding.o -------------------------------------------------------------------------------- /src/extensions/gridding/build/temp.linux-x86_64-3.8/gridding_cuda.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/gridding/build/temp.linux-x86_64-3.8/gridding_cuda.o -------------------------------------------------------------------------------- /src/extensions/gridding/build/temp.linux-x86_64-3.8/gridding_reverse.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/gridding/build/temp.linux-x86_64-3.8/gridding_reverse.o -------------------------------------------------------------------------------- /src/extensions/gridding/dist/gridding-2.1.0-py3.8-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/extensions/gridding/dist/gridding-2.1.0-py3.8-linux-x86_64.egg -------------------------------------------------------------------------------- /src/extensions/gridding/gridding.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 2.1 2 | Name: gridding 3 | Version: 2.1.0 4 | Summary: UNKNOWN 5 | License: UNKNOWN 6 | Platform: UNKNOWN 7 | 8 | UNKNOWN 9 | 10 | -------------------------------------------------------------------------------- /src/extensions/gridding/gridding.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | gridding.cu 2 | gridding_cuda.cpp 3 | gridding_reverse.cu 4 | setup.py 5 | gridding.egg-info/PKG-INFO 6 | gridding.egg-info/SOURCES.txt 7 | gridding.egg-info/dependency_links.txt 8 | gridding.egg-info/top_level.txt -------------------------------------------------------------------------------- /src/extensions/gridding/gridding.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /src/extensions/gridding/gridding.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | gridding 2 | -------------------------------------------------------------------------------- /src/extensions/gridding/gridding_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-11-13 10:52:53 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2020-06-17 14:52:32 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | // NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. 13 | #define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") 14 | #define CHECK_CONTIGUOUS(x) \ 15 | AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) \ 17 | CHECK_CUDA(x); \ 18 | CHECK_CONTIGUOUS(x) 19 | 20 | std::vector gridding_cuda_forward(float min_x, 21 | float max_x, 22 | float min_y, 23 | float max_y, 24 | float min_z, 25 | float max_z, 26 | torch::Tensor ptcloud, 27 | cudaStream_t stream); 28 | 29 | torch::Tensor gridding_cuda_backward(torch::Tensor grid_pt_weights, 30 | torch::Tensor grid_pt_indexes, 31 | torch::Tensor grad_grid, 32 | cudaStream_t stream); 33 | 34 | torch::Tensor gridding_reverse_cuda_forward(int scale, 35 | torch::Tensor grid, 36 | cudaStream_t stream); 37 | 38 | torch::Tensor gridding_reverse_cuda_backward(torch::Tensor ptcloud, 39 | torch::Tensor grid, 40 | torch::Tensor grad_ptcloud, 41 | cudaStream_t stream); 42 | 43 | std::vector gridding_forward(float min_x, 44 | float max_x, 45 | float min_y, 46 | float max_y, 47 | float min_z, 48 | float max_z, 49 | torch::Tensor ptcloud) { 50 | CHECK_INPUT(ptcloud); 51 | 52 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 53 | return gridding_cuda_forward(min_x, max_x, min_y, max_y, min_z, max_z, 54 | ptcloud, stream); 55 | } 56 | 57 | torch::Tensor gridding_backward(torch::Tensor grid_pt_weights, 58 | torch::Tensor grid_pt_indexes, 59 | torch::Tensor grad_grid) { 60 | CHECK_INPUT(grid_pt_weights); 61 | CHECK_INPUT(grid_pt_indexes); 62 | CHECK_INPUT(grad_grid); 63 | 64 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 65 | return gridding_cuda_backward(grid_pt_weights, grid_pt_indexes, grad_grid, 66 | stream); 67 | } 68 | 69 | torch::Tensor gridding_reverse_forward(int scale, torch::Tensor grid) { 70 | CHECK_INPUT(grid); 71 | 72 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 73 | return gridding_reverse_cuda_forward(scale, grid, stream); 74 | } 75 | 76 | torch::Tensor gridding_reverse_backward(torch::Tensor ptcloud, 77 | torch::Tensor grid, 78 | torch::Tensor grad_ptcloud) { 79 | CHECK_INPUT(ptcloud); 80 | CHECK_INPUT(grid); 81 | CHECK_INPUT(grad_ptcloud); 82 | 83 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 84 | return gridding_reverse_cuda_backward(ptcloud, grid, grad_ptcloud, stream); 85 | } 86 | 87 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 88 | m.def("forward", &gridding_forward, "Gridding forward (CUDA)"); 89 | m.def("backward", &gridding_backward, "Gridding backward (CUDA)"); 90 | m.def("rev_forward", &gridding_reverse_forward, 91 | "Gridding Reverse forward (CUDA)"); 92 | m.def("rev_backward", &gridding_reverse_backward, 93 | "Gridding Reverse backward (CUDA)"); 94 | } 95 | -------------------------------------------------------------------------------- /src/extensions/gridding/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-11-13 10:51:33 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-02 17:02:16 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup(name='gridding', 12 | version='2.1.0', 13 | ext_modules=[ 14 | CUDAExtension('gridding', ['gridding_cuda.cpp', 'gridding.cu', 'gridding_reverse.cu']), 15 | ], 16 | cmdclass={'build_ext': BuildExtension}) 17 | -------------------------------------------------------------------------------- /src/extensions/gridding/test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-10 10:48:55 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 14:20:42 6 | # @Email: cshzxie@gmail.com 7 | # 8 | # Note: 9 | # - Replace float -> double, kFloat -> kDouble in gridding.cu and gridding_reverse.cu 10 | 11 | import os 12 | import sys 13 | import torch 14 | import unittest 15 | 16 | from torch.autograd import gradcheck 17 | 18 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) 19 | from extensions.gridding import GriddingFunction, GriddingReverseFunction 20 | 21 | 22 | class GriddingTestCase(unittest.TestCase): 23 | def test_gridding_reverse_function_4(self): 24 | x = torch.rand(2, 4, 4, 4) 25 | x.requires_grad = True 26 | self.assertTrue(gradcheck(GriddingReverseFunction.apply, [4, x.double().cuda()])) 27 | 28 | def test_gridding_reverse_function_8(self): 29 | x = torch.rand(4, 8, 8, 8) 30 | x.requires_grad = True 31 | self.assertTrue(gradcheck(GriddingReverseFunction.apply, [8, x.double().cuda()])) 32 | 33 | def test_gridding_reverse_function_16(self): 34 | x = torch.rand(1, 16, 16, 16) 35 | x.requires_grad = True 36 | self.assertTrue(gradcheck(GriddingReverseFunction.apply, [16, x.double().cuda()])) 37 | 38 | def test_gridding_function_32pts(self): 39 | x = torch.rand(1, 32, 3) 40 | x.requires_grad = True 41 | self.assertTrue(gradcheck(GriddingFunction.apply, [x.double().cuda()])) 42 | 43 | 44 | if __name__ == '__main__': 45 | unittest.main() 46 | -------------------------------------------------------------------------------- /src/extensions/gridding_loss/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-30 09:56:06 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2020-02-22 19:19:43 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | import gridding_distance 11 | 12 | 13 | class GriddingDistanceFunction(torch.autograd.Function): 14 | @staticmethod 15 | def forward(ctx, min_x, max_x, min_y, max_y, min_z, max_z, pred_cloud, gt_cloud): 16 | pred_grid, pred_grid_pt_weights, pred_grid_pt_indexes = gridding_distance.forward( 17 | min_x, max_x, min_y, max_y, min_z, max_z, pred_cloud) 18 | # print(pred_grid.size()) # torch.Size(batch_size, n_grid_vertices, 8) 19 | # print(pred_grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) 20 | # print(pred_grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) 21 | gt_grid, gt_grid_pt_weights, gt_grid_pt_indexes = gridding_distance.forward( 22 | min_x, max_x, min_y, max_y, min_z, max_z, gt_cloud) 23 | # print(gt_grid.size()) # torch.Size(batch_size, n_grid_vertices, 8) 24 | # print(gt_grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) 25 | # print(gt_grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) 26 | 27 | ctx.save_for_backward(pred_grid_pt_weights, pred_grid_pt_indexes, gt_grid_pt_weights, gt_grid_pt_indexes) 28 | return pred_grid, gt_grid 29 | 30 | @staticmethod 31 | def backward(ctx, grad_pred_grid, grad_gt_grid): 32 | pred_grid_pt_weights, pred_grid_pt_indexes, gt_grid_pt_weights, gt_grid_pt_indexes = ctx.saved_tensors 33 | 34 | grad_pred_cloud = gridding_distance.backward(pred_grid_pt_weights, pred_grid_pt_indexes, grad_pred_grid) 35 | # print(grad_pred_cloud.size()) # torch.Size(batch_size, n_pts, 3) 36 | grad_gt_cloud = gridding_distance.backward(gt_grid_pt_weights, gt_grid_pt_indexes, grad_gt_grid) 37 | # print(grad_gt_cloud.size()) # torch.Size(batch_size, n_pts, 3) 38 | 39 | return None, None, None, None, None, None, grad_pred_cloud, grad_gt_cloud 40 | 41 | 42 | class GriddingDistance(torch.nn.Module): 43 | def __init__(self, scale=1): 44 | super(GriddingDistance, self).__init__() 45 | self.scale = scale 46 | 47 | def forward(self, pred_cloud, gt_cloud): 48 | ''' 49 | pred_cloud(b, n_pts1, 3) 50 | gt_cloud(b, n_pts2, 3) 51 | ''' 52 | pred_cloud = pred_cloud * self.scale / 2 53 | gt_cloud = gt_cloud * self.scale / 2 54 | 55 | min_pred_x = torch.min(pred_cloud[:, :, 0]) 56 | max_pred_x = torch.max(pred_cloud[:, :, 0]) 57 | min_pred_y = torch.min(pred_cloud[:, :, 1]) 58 | max_pred_y = torch.max(pred_cloud[:, :, 1]) 59 | min_pred_z = torch.min(pred_cloud[:, :, 2]) 60 | max_pred_z = torch.max(pred_cloud[:, :, 2]) 61 | 62 | min_gt_x = torch.min(gt_cloud[:, :, 0]) 63 | max_gt_x = torch.max(gt_cloud[:, :, 0]) 64 | min_gt_y = torch.min(gt_cloud[:, :, 1]) 65 | max_gt_y = torch.max(gt_cloud[:, :, 1]) 66 | min_gt_z = torch.min(gt_cloud[:, :, 2]) 67 | max_gt_z = torch.max(gt_cloud[:, :, 2]) 68 | 69 | min_x = torch.floor(torch.min(min_pred_x, min_gt_x)) - 1 70 | max_x = torch.ceil(torch.max(max_pred_x, max_gt_x)) + 1 71 | min_y = torch.floor(torch.min(min_pred_y, min_gt_y)) - 1 72 | max_y = torch.ceil(torch.max(max_pred_y, max_gt_y)) + 1 73 | min_z = torch.floor(torch.min(min_pred_z, min_gt_z)) - 1 74 | max_z = torch.ceil(torch.max(max_pred_z, max_gt_z)) + 1 75 | 76 | _pred_clouds = torch.split(pred_cloud, 1, dim=0) 77 | _gt_clouds = torch.split(gt_cloud, 1, dim=0) 78 | pred_grids = [] 79 | gt_grids = [] 80 | for pc, gc in zip(_pred_clouds, _gt_clouds): 81 | non_zeros = torch.sum(pc, dim=2).ne(0) 82 | pc = pc[non_zeros].unsqueeze(dim=0) 83 | non_zeros = torch.sum(gc, dim=2).ne(0) 84 | gc = gc[non_zeros].unsqueeze(dim=0) 85 | pred_grid, gt_grid = GriddingDistanceFunction.apply(min_x, max_x, min_y, max_y, min_z, max_z, pc, gc) 86 | pred_grids.append(pred_grid) 87 | gt_grids.append(gt_grid) 88 | 89 | return torch.cat(pred_grids, dim=0).contiguous(), torch.cat(gt_grids, dim=0).contiguous() 90 | 91 | 92 | class GriddingLoss(torch.nn.Module): 93 | def __init__(self, scales=[], alphas=[]): 94 | super(GriddingLoss, self).__init__() 95 | self.scales = scales 96 | self.alphas = alphas 97 | self.gridding_dists = [GriddingDistance(scale=s) for s in scales] 98 | self.l1_loss = torch.nn.L1Loss() 99 | 100 | def forward(self, pred_cloud, gt_cloud): 101 | gridding_loss = None 102 | n_dists = len(self.scales) 103 | 104 | for i in range(n_dists): 105 | alpha = self.alphas[i] 106 | gdist = self.gridding_dists[i] 107 | pred_grid, gt_grid = gdist(pred_cloud, gt_cloud) 108 | 109 | if gridding_loss is None: 110 | gridding_loss = alpha * self.l1_loss(pred_grid, gt_grid) 111 | else: 112 | gridding_loss += alpha * self.l1_loss(pred_grid, gt_grid) 113 | 114 | return gridding_loss 115 | -------------------------------------------------------------------------------- /src/extensions/gridding_loss/gridding_distance_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-12-30 10:59:31 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2020-06-17 14:52:52 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | // NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. 13 | #define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") 14 | #define CHECK_CONTIGUOUS(x) \ 15 | AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) \ 17 | CHECK_CUDA(x); \ 18 | CHECK_CONTIGUOUS(x) 19 | 20 | std::vector gridding_distance_cuda_forward(float min_x, 21 | float max_x, 22 | float min_y, 23 | float max_y, 24 | float min_z, 25 | float max_z, 26 | torch::Tensor ptcloud, 27 | cudaStream_t stream); 28 | 29 | torch::Tensor gridding_distance_cuda_backward(torch::Tensor grid_pt_weights, 30 | torch::Tensor grid_pt_indexes, 31 | torch::Tensor grad_grid, 32 | cudaStream_t stream); 33 | 34 | std::vector gridding_distance_forward(float min_x, 35 | float max_x, 36 | float min_y, 37 | float max_y, 38 | float min_z, 39 | float max_z, 40 | torch::Tensor ptcloud) { 41 | CHECK_INPUT(ptcloud); 42 | 43 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 44 | return gridding_distance_cuda_forward(min_x, max_x, min_y, max_y, min_z, 45 | max_z, ptcloud, stream); 46 | } 47 | 48 | torch::Tensor gridding_distance_backward(torch::Tensor grid_pt_weights, 49 | torch::Tensor grid_pt_indexes, 50 | torch::Tensor grad_grid) { 51 | CHECK_INPUT(grid_pt_weights); 52 | CHECK_INPUT(grid_pt_indexes); 53 | CHECK_INPUT(grad_grid); 54 | 55 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 56 | return gridding_distance_cuda_backward(grid_pt_weights, grid_pt_indexes, 57 | grad_grid, stream); 58 | } 59 | 60 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 61 | m.def("forward", &gridding_distance_forward, 62 | "Gridding Distance Forward (CUDA)"); 63 | m.def("backward", &gridding_distance_backward, 64 | "Gridding Distance Backward (CUDA)"); 65 | } 66 | -------------------------------------------------------------------------------- /src/extensions/gridding_loss/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-30 11:03:55 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-30 11:13:39 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup(name='gridding_distance', 12 | version='1.0.0', 13 | ext_modules=[ 14 | CUDAExtension('gridding_distance', ['gridding_distance_cuda.cpp', 'gridding_distance.cu']), 15 | ], 16 | cmdclass={'build_ext': BuildExtension}) 17 | -------------------------------------------------------------------------------- /src/loss/__pycache__/loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/loss/__pycache__/loss.cpython-36.pyc -------------------------------------------------------------------------------- /src/loss/__pycache__/loss.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/loss/__pycache__/loss.cpython-37.pyc -------------------------------------------------------------------------------- /src/loss/__pycache__/loss.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/loss/__pycache__/loss.cpython-38.pyc -------------------------------------------------------------------------------- /src/loss/loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as f 4 | import sys 5 | sys.path.append('../') 6 | import numpy as np 7 | from knn_cuda import KNN 8 | from pykdtree.kdtree import KDTree 9 | import open3d as o3d 10 | from pytorch3d.ops.knn import knn_gather, knn_points 11 | import torch.nn.functional as F 12 | from torch_scatter import scatter_min, scatter_mean 13 | 14 | l1_loss = nn.L1Loss() 15 | l2_loss = nn.MSELoss() 16 | 17 | 18 | 19 | def PolyConv_loss(x,Input,indices,conv_p,conv_w): 20 | Input_adj = knn_gather(Input,indices) 21 | x = x.unsqueeze(2).unsqueeze(-1) #Input_adj[:,0,:].unsqueeze(1).unsqueeze(-1) 22 | y = Input_adj[:,:,:].unsqueeze(-1) 23 | x_repeat = x*torch.ones_like(y) 24 | fxy = (torch.cat([torch.ones_like(y),x_repeat,y,torch.pow(x_repeat,2),x_repeat*y,torch.pow(y,2)],-1))*(conv_p.detach().unsqueeze(0).unsqueeze(0).unsqueeze(0)) 25 | yfxy = fxy * y 26 | yfxy = torch.mean(torch.sum(yfxy,-1),dim=2) 27 | x_prime = torch.cat([2*torch.ones_like(x),2*x,torch.zeros_like(x),2*torch.pow(x,2),torch.zeros_like(x),(2/3)*torch.ones_like(x)],-1) 28 | fx = x_prime*(conv_p.unsqueeze(0).unsqueeze(0).unsqueeze(0)) 29 | fx = torch.sum(fx,-1)[:,:,0] 30 | fxy_cond = torch.true_divide(yfxy,fx+0.0001) 31 | fxy_cond = F.linear(fxy_cond, conv_w.weight.detach(), conv_w.bias.detach()) 32 | return fxy_cond 33 | 34 | 35 | 36 | def estimate_normals(points,k): 37 | batch_size = points.size(0) 38 | points_mean = points.mean(1,keepdim=True) 39 | points_centered = points - points_mean 40 | knn_indices, _ = get_nearest_neighbors_indices_batch(points_centered.detach().cpu().numpy(),points_centered.detach().cpu().numpy(),k) 41 | knn_indices = torch.LongTensor(knn_indices).to(points.device) 42 | knn_points = [] 43 | for b in range(batch_size): 44 | knn_points.append(points[b,knn_indices[b]].unsqueeze(0)) 45 | knn_points = torch.cat(knn_points,0) 46 | pt_mean = knn_points.mean(2, keepdim=True) 47 | central_diff = knn_points - pt_mean 48 | per_pt_cov = central_diff.unsqueeze(4) * central_diff.unsqueeze(3) 49 | covariances = per_pt_cov.mean(2) 50 | print(covariances.device) 51 | curvatures, local_coord_frames = torch.symeig(covariances, eigenvectors=True) 52 | print(curvatures.shape) 53 | normals = local_coord_frames[:, :, :, 0] 54 | print(normals.shape) 55 | return normals, knn_indices 56 | 57 | 58 | 59 | 60 | def get_nearest_neighbors_indices_batch(points_src, points_tgt, k=1): 61 | ''' Returns the nearest neighbors for point sets batchwise. 62 | Args: 63 | points_src (numpy array): source points 64 | points_tgt (numpy array): target points 65 | k (int): number of nearest neighbors to return 66 | ''' 67 | indices = [] 68 | distances = [] 69 | 70 | for (p1, p2) in zip(points_src, points_tgt): 71 | kdtree = KDTree(p2) 72 | dist, idx = kdtree.query(p1, k=k) 73 | indices.append(idx) 74 | distances.append(dist) 75 | 76 | return indices, distances 77 | 78 | 79 | 80 | def chamfer_distance_naive(points1, points2): 81 | ''' Naive implementation of the Chamfer distance. 82 | Args: 83 | points1 (batch, 3, num_on_points) 84 | points2 (batch, 3, num_on_points) 85 | ''' 86 | points1 = torch.transpose(points1,1,2) 87 | points2 = torch.transpose(points2,1,2) 88 | batch_size, T1, _ = points1.size() 89 | _, T2, _ = points2.size() 90 | 91 | points1 = points1.view(batch_size, T1, 1, 3) 92 | points2 = points2.view(batch_size, 1, T2, 3) 93 | 94 | distances = (points1 - points2).pow(2).sum(-1) 95 | 96 | chamfer1 = distances.min(dim=1)[0].mean(dim=1) 97 | chamfer2 = distances.min(dim=2)[0].mean(dim=1) 98 | return chamfer1, chamfer2 99 | 100 | 101 | def directed_hausdorff(point_cloud1:torch.Tensor, point_cloud2:torch.Tensor, reduce_mean=True): 102 | """ 103 | # UHD from MPC: https://github.com/ChrisWu1997/Multimodal-Shape-Completion/blob/master/evaluation/completeness.py 104 | :param point_cloud1: (B, 3, N) 105 | :param point_cloud2: (B, 3, M) 106 | :return: directed hausdorff distance, A -> B 107 | """ 108 | 109 | 110 | n_pts1 = point_cloud1.shape[2] 111 | n_pts2 = point_cloud2.shape[2] 112 | 113 | pc1 = point_cloud1.unsqueeze(3) 114 | pc1 = pc1.repeat((1, 1, 1, n_pts2)) # (B, 3, N, M) 115 | pc2 = point_cloud2.unsqueeze(2) 116 | pc2 = pc2.repeat((1, 1, n_pts1, 1)) # (B, 3, N, M) 117 | 118 | l2_dist = torch.sqrt(torch.sum((pc1 - pc2) ** 2, dim=1)) # (B, N, M) 119 | 120 | shortest_dist, _ = torch.min(l2_dist, dim=2) 121 | 122 | hausdorff_dist, _ = torch.max(shortest_dist, dim=1) # (B, ) 123 | 124 | if reduce_mean: 125 | hausdorff_dist = torch.mean(hausdorff_dist) 126 | 127 | return hausdorff_dist 128 | 129 | def distChamfer(a, b): 130 | """ 131 | :param a: Pointclouds Batch x nul_points x dim 132 | :param b: Pointclouds Batch x nul_points x dim 133 | :return: 134 | -closest point on b of points from a 135 | -closest point on a of points from b 136 | -idx of closest point on b of points from a 137 | -idx of closest point on a of points from b 138 | Works for pointcloud of any dimension 139 | """ 140 | x, y = a.double(), b.double() 141 | bs, num_points_x, points_dim = x.size() 142 | bs, num_points_y, points_dim = y.size() 143 | 144 | xx = torch.pow(x, 2).sum(2) 145 | yy = torch.pow(y, 2).sum(2) 146 | zz = torch.bmm(x, y.transpose(2, 1)) 147 | rx = xx.unsqueeze(1).expand(bs, num_points_y, num_points_x) # Diagonal elements xx 148 | ry = yy.unsqueeze(1).expand(bs, num_points_x, num_points_y) # Diagonal elements yy 149 | P = rx.transpose(2, 1) + ry - 2 * zz 150 | 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() 151 | 152 | 153 | def chamfer_distance_kdtree(x, y): 154 | ''' KD-tree based implementation of the Chamfer distance. 155 | Args: 156 | points1 (batch, 3, num_on_points) 157 | points2 (batch, 3, num_on_points) 158 | ''' 159 | 160 | x = torch.transpose(x,1,2) 161 | y = torch.transpose(y,1,2) 162 | 163 | 164 | x_nn = knn_points(x, y, K=1) 165 | y_nn = knn_points(y, x, K=1) 166 | 167 | cham_x = x_nn.dists[..., 0] # (N, P1) 168 | cham_y = y_nn.dists[..., 0] # (N, P2) 169 | 170 | 171 | # Apply point reduction 172 | cham_x = torch.sqrt(cham_x) # (N,) 173 | cham_y = torch.sqrt(cham_y) # (N,) 174 | 175 | return cham_x.mean(1), cham_y.mean(1) 176 | 177 | 178 | def normal_loss(points,k=6, epsilon=0.001): 179 | points = torch.transpose(points,1,2) 180 | batch_size = points.size(0) 181 | points_mean = points.mean(1,keepdim=True) 182 | points_centered = points - points_mean 183 | knn_indices, _ = get_nearest_neighbors_indices_batch(points_centered.detach().cpu().numpy(),points_centered.detach().cpu().numpy(),k) 184 | knn_indices = torch.LongTensor(knn_indices).to(points.device) 185 | knn_points = [] 186 | for b in range(batch_size): 187 | knn_points.append(points[b,knn_indices[b]].unsqueeze(0)) 188 | 189 | knn_points = torch.cat(knn_points,0) 190 | central_diff = knn_points 191 | per_pt_cov = central_diff.unsqueeze(4) * central_diff.unsqueeze(3) 192 | covariances = per_pt_cov.mean(2) 193 | curvatures, local_coord_frames = torch.linalg.eigh(covariances) 194 | vertex_normal = local_coord_frames[:, :, :, 0] 195 | 196 | knn_normals = [] 197 | for b in range(batch_size): 198 | knn_normals.append(vertex_normal[b,knn_indices[b]].unsqueeze(0)) 199 | 200 | knn_normals = torch.cat(knn_normals,0) 201 | loss = l1_loss(knn_normals[:,:,1:,:]-knn_normals[:,:,0,:].unsqueeze(2),torch.zeros_like(knn_normals[:,:,1:,:])) 202 | return loss 203 | 204 | 205 | class MSELoss(nn.Module): 206 | def __init__(self): 207 | super(MSELoss, self).__init__() 208 | self.loss = nn.MSELoss() 209 | 210 | def forward(self, input, target): 211 | 212 | return self.loss(input, target) 213 | 214 | 215 | 216 | class ChamferLoss(nn.Module): 217 | def __init__(self, device): 218 | super(ChamferLoss, self).__init__() 219 | self.loss = ChamferDistance() 220 | self.device = device 221 | 222 | def forward(self, xyz1, xyz2): 223 | return self.loss(xyz1, xyz2) 224 | 225 | 226 | 227 | -------------------------------------------------------------------------------- /src/main.py: -------------------------------------------------------------------------------- 1 | from symbol import arglist 2 | import torch 3 | import numpy as np 4 | import argparse 5 | import os 6 | from train import trainer 7 | from train_semkitti import trainer as trainer_semkitt 8 | from train_real import trainer as trainer_real 9 | from pretrain import pretrainer 10 | import torch.multiprocessing as mp 11 | import os 12 | os.environ['CUDA_LAUNCH_BLOCKING'] = "1" 13 | 14 | def get_parser(): 15 | parser = argparse.ArgumentParser(description='Point cloud Completion') 16 | parser.add_argument('--exp_name', type=str, default=None, metavar='N', 17 | help='Name of the experiment') 18 | parser.add_argument('--experiment_id', type=str, default= 'experiment1', help='experiment id ') 19 | parser.add_argument('--dataset_name', type=str, default='ShapeNet', help='The name of the dataset', choices = ['ShapeNet', 'SemanticKITTI', 'MatterPort', 'ScanNet', 'KITTI']) 20 | parser.add_argument('--root', type=str, default='/mnt/disk2/mchiash2/', help='The directory that contains dataset and experiment') 21 | parser.add_argument('--class_name', type=str, default='plane', help='class of dataset', choices = ['plane', 'car', 'chair', 'table']) 22 | parser.add_argument('--lr', type=float, default=0.001) 23 | parser.add_argument('--batch_size', type=int, default=4, metavar='batch_size', 24 | help='Size of batch)') 25 | parser.add_argument('--k', type=int, default=None, metavar='N', 26 | help='Num of nearest neighbors to use for KNN') 27 | parser.add_argument('--workers', type=int, help='Number of data loading workers', default=16) 28 | parser.add_argument('--epochs', type=int, default=1000, metavar='N', 29 | help='Number of episode to train ') 30 | parser.add_argument('--snapshot_interval', type=int, default=1, metavar='N', 31 | help='Save snapshot interval ') 32 | parser.add_argument('--pretrain', type = bool, default = False, choices = [True, False]) 33 | parser.add_argument('--num_points', type=int, default=8192, metavar='N', 34 | help='Num_points before removing (Original num_points)') 35 | parser.add_argument('--resume', type=str, default=None, metavar='N', 36 | help='checkpoint address') 37 | parser.add_argument('--remov_ratio', type=int, default=8, metavar='N', 38 | help='How much part of point cloud is goint to be removed') 39 | args = parser.parse_args() 40 | return args 41 | 42 | 43 | if __name__ == '__main__': 44 | args = get_parser() 45 | ExpNum_dir = os.path.join("/mnt/disk2/mchiash2/experiments/", args.experiment_id ) 46 | if os.path.isdir(ExpNum_dir) is False: 47 | os.mkdir(ExpNum_dir) 48 | if os.path.isdir(os.path.join(ExpNum_dir, "models")) is False: 49 | os.mkdir(os.path.join(ExpNum_dir, "models")) 50 | isdir = os.path.isdir("../outputs/"+str(args.experiment_id)) 51 | if isdir is False: 52 | os.mkdir("../outputs/"+str(args.experiment_id)) 53 | os.mkdir("../outputs/"+str(args.experiment_id)+'/train') 54 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test') 55 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/gt') 56 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/ipc') 57 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/ipc/0') 58 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/ipc/1') 59 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/ipc/2') 60 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/ipc/3') 61 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/ipc/4') 62 | 63 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/cpc') 64 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/cpc/0') 65 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/cpc/1') 66 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/cpc/2') 67 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/cpc/3') 68 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/cpc/4') 69 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/coarse') 70 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/coarse/0') 71 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/coarse/1') 72 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/coarse/2') 73 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/coarse/3') 74 | os.mkdir("../outputs/"+str(args.experiment_id)+'/test/coarse/4') 75 | np.save('../outputs/'+str(args.experiment_id)+'/result.npy',100*np.ones(4)) 76 | 77 | if args.pretrain is True: 78 | pretrain = pretrainer(args) 79 | pretrain.run(args) 80 | 81 | elif (args.dataset_name == 'ShapeNet'): 82 | train = trainer(args) 83 | train.run(args) 84 | 85 | elif args.dataset_name == 'SemanticKITTI': 86 | train = trainer_semkitt(args) 87 | train.run(args) 88 | 89 | elif (args.dataset_name == 'MatterPort') or (args.dataset_name == 'ScanNet') or (args.dataset_name == 'KITTI'): 90 | train = trainer_real(args) 91 | train.run(args) 92 | 93 | 94 | -------------------------------------------------------------------------------- /src/models/TopNet.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import torch.nn as nn 4 | import math 5 | import numpy as np 6 | 7 | 8 | 9 | tree_arch = {} 10 | tree_arch[2] = [32, 64] 11 | tree_arch[4] = [4, 8, 8, 8] 12 | tree_arch[6] = [2, 4, 4, 4, 4, 4] 13 | tree_arch[8] = [2, 2, 2, 2, 2, 4, 4, 4] 14 | 15 | def get_arch(nlevels, npts): 16 | logmult = int(math.log2(npts/2048)) 17 | assert 2048*(2**(logmult)) == npts, "Number of points is %d, expected 2048x(2^n)" % (npts) 18 | arch = tree_arch[nlevels] 19 | while logmult > 0: 20 | last_min_pos = np.where(arch==np.min(arch))[0][-1] 21 | arch[last_min_pos]*=2 22 | logmult -= 1 23 | return arch 24 | 25 | def build_model_from_cfg(cfg, **kwargs): 26 | """ 27 | Build a dataset, defined by `dataset_name`. 28 | Args: 29 | cfg (eDICT): 30 | Returns: 31 | Dataset: a constructed dataset specified by dataset_name. 32 | """ 33 | return MODELS.build(cfg, **kwargs) 34 | 35 | class TopNet(nn.Module): 36 | def __init__(self): # node_feature = 8, encoder_feature = 1024, nlevels = 8, num_pred = 2048 37 | super().__init__() 38 | self.node_feature = 8 39 | self.encoder_feature = 512 40 | self.nlevels = 8 41 | self.num_pred = 8192 42 | 43 | self.tarch = get_arch(self.nlevels, self.num_pred) 44 | self.Top_in_channel = self.encoder_feature + self.node_feature 45 | self.Top_out_channel = self.node_feature 46 | self.first_conv = nn.Sequential( 47 | nn.Conv1d(3,128,1), 48 | nn.BatchNorm1d(128), 49 | nn.ReLU(inplace=True), 50 | nn.Conv1d(128,256,1) 51 | ) 52 | self.second_conv = nn.Sequential( 53 | nn.Conv1d(512,512,1), 54 | nn.BatchNorm1d(512), 55 | nn.ReLU(inplace=True), 56 | nn.Conv1d(512,self.encoder_feature,1) 57 | ) 58 | self.root_layer = nn.Sequential( 59 | nn.Linear(self.encoder_feature,256), 60 | nn.BatchNorm1d(256), 61 | nn.ReLU(inplace=True), 62 | nn.Linear(256,64), 63 | nn.BatchNorm1d(64), 64 | nn.ReLU(inplace=True), 65 | nn.Linear(64 , self.node_feature * int(self.tarch[0])), 66 | nn.Tanh() 67 | ) 68 | self.leaf_layer = self.get_tree_layer(self.Top_in_channel, 3, int(self.tarch[-1])) 69 | self.feature_layers = nn.ModuleList([self.get_tree_layer(self.Top_in_channel, self.Top_out_channel, int(self.tarch[d]) ) for d in range(1, self.nlevels - 1)]) 70 | 71 | 72 | @staticmethod 73 | def get_tree_layer(in_channel, out_channel, node): 74 | return nn.Sequential( 75 | nn.Conv1d(in_channel, in_channel//2, 1), 76 | nn.BatchNorm1d(in_channel//2), 77 | nn.ReLU(inplace=True), 78 | nn.Conv1d(in_channel//2, in_channel//4, 1), 79 | nn.BatchNorm1d(in_channel//4), 80 | nn.ReLU(inplace=True), 81 | nn.Conv1d(in_channel//4, in_channel//8, 1), 82 | nn.BatchNorm1d(in_channel//8), 83 | nn.ReLU(inplace=True), 84 | nn.Conv1d(in_channel//8, out_channel * node, 1), 85 | ) 86 | 87 | def forward(self, xyz): 88 | xyz = xyz.transpose(1,2) 89 | bs , n , _ = xyz.shape 90 | # encoder 91 | feature = self.first_conv(xyz.transpose(2,1)) # B 256 n 92 | feature_global = torch.max(feature,dim=2,keepdim=True)[0] # B 256 1 93 | feature = torch.cat([feature_global.expand(-1,-1,n), feature], dim=1)# B 512 n 94 | feature = self.second_conv(feature) # B 1024 n 95 | feature_global = torch.max(feature,dim=2,keepdim=False)[0] # B 1024 96 | # decoder 97 | level10 = self.root_layer(feature_global).reshape(-1, self.node_feature ,int(self.tarch[0])) # B 8 node 98 | outs = [level10,] 99 | for i in range(1, self.nlevels): 100 | last_level = outs[-1] 101 | expand_feature = feature_global.unsqueeze(2).expand(-1,-1,last_level.shape[2]) 102 | if i == self.nlevels - 1: 103 | layer_feature = self.leaf_layer(torch.cat([expand_feature,last_level],dim=1)).reshape(bs, 3 ,-1) 104 | else: 105 | layer_feature = self.feature_layers[i-1](torch.cat([expand_feature,last_level],dim=1)).reshape(bs, self.node_feature, -1) 106 | outs.append(nn.Tanh()(layer_feature)) 107 | return outs[-1].contiguous() -------------------------------------------------------------------------------- /src/models/__pycache__/AutoEncoder.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/AutoEncoder.cpython-36.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/AutoEncoder.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/AutoEncoder.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/PUGAN.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/PUGAN.cpython-36.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/PUGAN.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/PUGAN.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/TopNet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/TopNet.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/base_model.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/base_model.cpython-37.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/base_model.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/base_model.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/grnet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/grnet.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/keypointfinder.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/keypointfinder.cpython-37.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/keypointfinder.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/keypointfinder.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/model.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/model.cpython-36.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/model.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/model.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/pcn.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/pcn.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/pcnenc_polydec.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/pcnenc_polydec.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/polyconv.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/polyconv.cpython-36.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/polyconv.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/polyconv.cpython-37.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/polyconv.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/polyconv.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/polynet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/polynet.cpython-36.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/polynet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/polynet.cpython-37.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/polynet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/polynet.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/polynet_enc_pcn_dec.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/polynet_enc_pcn_dec.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/util.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/util.cpython-36.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/util.cpython-38.pyc -------------------------------------------------------------------------------- /src/models/__pycache__/utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/__pycache__/utils.cpython-36.pyc -------------------------------------------------------------------------------- /src/models/base_model.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import numpy as np 3 | from abc import abstractmethod 4 | 5 | 6 | class BaseModel(nn.Module): 7 | """ 8 | Base class for all models 9 | """ 10 | @abstractmethod 11 | def forward(self, *inputs): 12 | """ 13 | Forward pass logic 14 | 15 | :return: Model output 16 | """ 17 | raise NotImplementedError 18 | 19 | def __str__(self): 20 | """ 21 | Model prints with number of trainable parameters 22 | """ 23 | model_parameters = filter(lambda p: p.requires_grad, self.parameters()) 24 | params = sum([np.prod(p.size()) for p in model_parameters]) 25 | return super().__str__() + '\nTrainable parameters: {}'.format(params) -------------------------------------------------------------------------------- /src/models/grnet.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-09-06 11:35:30 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2020-02-22 19:20:36 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | from extensions.gridding import Gridding, GriddingReverse 11 | from extensions.cubic_feature_sampling import CubicFeatureSampling 12 | 13 | 14 | class RandomPointSampling(torch.nn.Module): 15 | def __init__(self, n_points): 16 | super(RandomPointSampling, self).__init__() 17 | self.n_points = n_points 18 | 19 | def forward(self, pred_cloud, partial_cloud=None): 20 | if partial_cloud is not None: 21 | pred_cloud = torch.cat([partial_cloud, pred_cloud], dim=1) 22 | 23 | _ptcloud = torch.split(pred_cloud, 1, dim=0) 24 | ptclouds = [] 25 | for p in _ptcloud: 26 | non_zeros = torch.sum(p, dim=2).ne(0) 27 | p = p[non_zeros].unsqueeze(dim=0) 28 | n_pts = p.size(1) 29 | if n_pts < self.n_points: 30 | rnd_idx = torch.cat([torch.randint(0, n_pts, (self.n_points, ))]) 31 | else: 32 | rnd_idx = torch.randperm(p.size(1))[:self.n_points] 33 | ptclouds.append(p[:, rnd_idx, :]) 34 | 35 | return torch.cat(ptclouds, dim=0).contiguous() 36 | 37 | 38 | class GRNet(torch.nn.Module): 39 | def __init__(self): 40 | super(GRNet, self).__init__() 41 | self.gridding = Gridding(scale=64) 42 | self.fc0 = torch.nn.Sequential( 43 | torch.nn.Linear(3096, 2048), 44 | torch.nn.ReLU() 45 | ) 46 | self.conv1 = torch.nn.Sequential( 47 | torch.nn.Conv3d(1, 32, kernel_size=4, padding=2), 48 | torch.nn.BatchNorm3d(32), 49 | torch.nn.LeakyReLU(0.2), 50 | torch.nn.MaxPool3d(kernel_size=2) 51 | ) 52 | self.conv2 = torch.nn.Sequential( 53 | torch.nn.Conv3d(32, 64, kernel_size=4, padding=2), 54 | torch.nn.BatchNorm3d(64), 55 | torch.nn.LeakyReLU(0.2), 56 | torch.nn.MaxPool3d(kernel_size=2) 57 | ) 58 | self.conv3 = torch.nn.Sequential( 59 | torch.nn.Conv3d(64, 128, kernel_size=4, padding=2), 60 | torch.nn.BatchNorm3d(128), 61 | torch.nn.LeakyReLU(0.2), 62 | torch.nn.MaxPool3d(kernel_size=2) 63 | ) 64 | self.conv4 = torch.nn.Sequential( 65 | torch.nn.Conv3d(128, 256, kernel_size=4, padding=2), 66 | torch.nn.BatchNorm3d(256), 67 | torch.nn.LeakyReLU(0.2), 68 | torch.nn.MaxPool3d(kernel_size=2) 69 | ) 70 | self.fc5 = torch.nn.Sequential( 71 | torch.nn.Linear(16384, 2048), 72 | torch.nn.ReLU() 73 | ) 74 | self.fc6 = torch.nn.Sequential( 75 | torch.nn.Linear(2048, 16384), 76 | torch.nn.ReLU() 77 | ) 78 | self.dconv7 = torch.nn.Sequential( 79 | torch.nn.ConvTranspose3d(256, 128, kernel_size=4, stride=2, bias=False, padding=1), 80 | torch.nn.BatchNorm3d(128), 81 | torch.nn.ReLU() 82 | ) 83 | self.dconv8 = torch.nn.Sequential( 84 | torch.nn.ConvTranspose3d(128, 64, kernel_size=4, stride=2, bias=False, padding=1), 85 | torch.nn.BatchNorm3d(64), 86 | torch.nn.ReLU() 87 | ) 88 | self.dconv9 = torch.nn.Sequential( 89 | torch.nn.ConvTranspose3d(64, 32, kernel_size=4, stride=2, bias=False, padding=1), 90 | torch.nn.BatchNorm3d(32), 91 | torch.nn.ReLU() 92 | ) 93 | self.dconv10 = torch.nn.Sequential( 94 | torch.nn.ConvTranspose3d(32, 1, kernel_size=4, stride=2, bias=False, padding=1), 95 | torch.nn.BatchNorm3d(1), 96 | torch.nn.ReLU() 97 | ) 98 | self.gridding_rev = GriddingReverse(scale=64) 99 | self.point_sampling = RandomPointSampling(n_points=2048) 100 | self.point_sampling2 = RandomPointSampling(n_points=2048) 101 | self.feature_sampling = CubicFeatureSampling() 102 | self.fc11 = torch.nn.Sequential( 103 | torch.nn.Linear(1792, 1792), 104 | torch.nn.ReLU() 105 | ) 106 | self.fc12 = torch.nn.Sequential( 107 | torch.nn.Linear(1792, 448), 108 | torch.nn.ReLU() 109 | ) 110 | self.fc13 = torch.nn.Sequential( 111 | torch.nn.Linear(448, 112), 112 | torch.nn.ReLU() 113 | ) 114 | self.fc14 = torch.nn.Linear(112, 12) 115 | 116 | def forward(self, data): 117 | partial_cloud = self.fc0(data).transpose(1,2) # torch.Size([batch_size, 3096, 3]) 118 | # print(partial_cloud.size()) # torch.Size([batch_size, 2048, 3]) 119 | pt_features_64_l = self.gridding(partial_cloud).view(-1, 1, 64, 64, 64) 120 | # print(pt_features_64_l.size()) # torch.Size([batch_size, 1, 64, 64, 64]) 121 | pt_features_32_l = self.conv1(pt_features_64_l) 122 | # print(pt_features_32_l.size()) # torch.Size([batch_size, 32, 32, 32, 32]) 123 | pt_features_16_l = self.conv2(pt_features_32_l) 124 | # print(pt_features_16_l.size()) # torch.Size([batch_size, 64, 16, 16, 16]) 125 | pt_features_8_l = self.conv3(pt_features_16_l) 126 | # print(pt_features_8_l.size()) # torch.Size([batch_size, 128, 8, 8, 8]) 127 | pt_features_4_l = self.conv4(pt_features_8_l) 128 | # print(pt_features_4_l.size()) # torch.Size([batch_size, 256, 4, 4, 4]) 129 | features = self.fc5(pt_features_4_l.view(-1, 16384)) 130 | # print(features.size()) # torch.Size([batch_size, 2048]) 131 | pt_features_4_r = self.fc6(features).view(-1, 256, 4, 4, 4) + pt_features_4_l 132 | # print(pt_features_4_r.size()) # torch.Size([batch_size, 256, 4, 4, 4]) 133 | pt_features_8_r = self.dconv7(pt_features_4_r) + pt_features_8_l 134 | # print(pt_features_8_r.size()) # torch.Size([batch_size, 128, 8, 8, 8]) 135 | pt_features_16_r = self.dconv8(pt_features_8_r) + pt_features_16_l 136 | # print(pt_features_16_r.size()) # torch.Size([batch_size, 64, 16, 16, 16]) 137 | pt_features_32_r = self.dconv9(pt_features_16_r) + pt_features_32_l 138 | # print(pt_features_32_r.size()) # torch.Size([batch_size, 32, 32, 32, 32]) 139 | pt_features_64_r = self.dconv10(pt_features_32_r) + pt_features_64_l 140 | # print(pt_features_64_r.size()) # torch.Size([batch_size, 1, 64, 64, 64]) 141 | sparse_cloud = self.gridding_rev(pt_features_64_r.squeeze(dim=1)) 142 | # print(sparse_cloud.size()) # torch.Size([batch_size, 262144, 3]) 143 | sparse_cloud = self.point_sampling(sparse_cloud, partial_cloud) 144 | # print(sparse_cloud.size()) # torch.Size([batch_size, 2048, 3]) 145 | point_features_32 = self.feature_sampling(sparse_cloud, pt_features_32_r).view(-1, 2048, 256) 146 | # print(point_features_32.size()) # torch.Size([batch_size, 2048, 256]) 147 | point_features_16 = self.feature_sampling(sparse_cloud, pt_features_16_r).view(-1, 2048, 512) 148 | # print(point_features_16.size()) # torch.Size([batch_size, 2048, 512]) 149 | point_features_8 = self.feature_sampling(sparse_cloud, pt_features_8_r).view(-1, 2048, 1024) 150 | # print(point_features_8.size()) # torch.Size([batch_size, 2048, 1024]) 151 | point_features = torch.cat([point_features_32, point_features_16, point_features_8], dim=2) 152 | # print(point_features.size()) # torch.Size([batch_size, 2048, 1792]) 153 | point_features = self.fc11(point_features) 154 | # print(point_features.size()) # torch.Size([batch_size, 2048, 1792]) 155 | point_features = self.fc12(point_features) 156 | # print(point_features.size()) # torch.Size([batch_size, 2048, 448]) 157 | point_features = self.fc13(point_features) 158 | # print(point_features.size()) # torch.Size([batch_size, 2048, 112]) 159 | point_offset = self.fc14(point_features).view(-1, 8192, 3) 160 | # print(point_features.size()) # torch.Size([batch_size, 16384, 3]) 161 | dense_cloud = sparse_cloud.unsqueeze(dim=2).repeat(1, 1, 4, 1).view(-1, 8192, 3) + point_offset 162 | # print(dense_cloud.size()) # torch.Size([batch_size, 16384, 3]) 163 | 164 | dense_cloud = dense_cloud.transpose(1,2) 165 | return dense_cloud -------------------------------------------------------------------------------- /src/models/model.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import torch.nn as nn 4 | 5 | class mlp(nn.Module): 6 | 7 | def __init__(self): 8 | super(mlp, self).__init__() 9 | n_layers=6 10 | n_in= 2048 11 | n_out= 4096 12 | layers = [] 13 | for i in range(n_layers-1): 14 | layers.append(nn.Linear(n_in, n_in)) 15 | layers.append(nn.ReLU(inplace=True)) 16 | 17 | layers.append(nn.Linear(n_in, n_out)) 18 | self.layers = nn.Sequential(*layers) 19 | 20 | 21 | def forward(self, pc): 22 | out = self.layers(pc) 23 | x = torch.linspace(-1, 1, steps=16, device=out.device) 24 | y = torch.linspace(-1, 1, steps=16, device=out.device) 25 | z = torch.linspace(-1, 1, steps=16, device=out.device) 26 | x, y, z = torch.meshgrid(x,y,z) 27 | grid = torch.cat([x.unsqueeze(0),y.unsqueeze(0),z.unsqueeze(0)],0).unsqueeze(0).view(1,3,-1) 28 | 29 | return out 30 | -------------------------------------------------------------------------------- /src/models/ninja-linux.zip: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/models/ninja-linux.zip -------------------------------------------------------------------------------- /src/models/pcn.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import torch.nn as nn 4 | 5 | 6 | class PCN(nn.Module): 7 | """ 8 | "PCN: Point Cloud Completion Network" 9 | (https://arxiv.org/pdf/1808.00671.pdf) 10 | Attributes: 11 | num_dense: 16384 12 | latent_dim: 1024 13 | grid_size: 4 14 | num_coarse: 1024 15 | """ 16 | 17 | def __init__(self, num_dense=8192, latent_dim=512, grid_size=4): 18 | super().__init__() 19 | 20 | self.num_dense = num_dense 21 | self.latent_dim = latent_dim 22 | self.grid_size = grid_size 23 | 24 | assert self.num_dense % self.grid_size ** 2 == 0 25 | 26 | self.num_coarse = self.num_dense // (self.grid_size ** 2) 27 | self.first_conv = nn.Sequential( 28 | nn.Conv1d(3, 128, 1), 29 | nn.BatchNorm1d(128), 30 | nn.ReLU(inplace=True), 31 | nn.Conv1d(128, 256, 1) 32 | ) 33 | 34 | self.second_conv = nn.Sequential( 35 | nn.Conv1d(512, 512, 1), 36 | nn.BatchNorm1d(512), 37 | nn.ReLU(inplace=True), 38 | nn.Conv1d(512, self.latent_dim, 1) 39 | ) 40 | 41 | self.mlp = nn.Sequential( 42 | nn.Linear(self.latent_dim, 1024), 43 | nn.ReLU(inplace=True), 44 | nn.Linear(1024, 1024), 45 | nn.ReLU(inplace=True), 46 | nn.Linear(1024, 3 * self.num_coarse) 47 | ) 48 | 49 | self.final_conv = nn.Sequential( 50 | #nn.Conv1d(1024 + 3 + 2, 512, 1) 51 | nn.Conv1d(512 + 3 + 2, 512, 1), 52 | nn.BatchNorm1d(512), 53 | nn.ReLU(inplace=True), 54 | nn.Conv1d(512, 512, 1), 55 | nn.BatchNorm1d(512), 56 | nn.ReLU(inplace=True), 57 | nn.Conv1d(512, 3, 1) 58 | ) 59 | a = torch.linspace(-0.05, 0.05, steps=self.grid_size, dtype=torch.float).view(1, self.grid_size).expand(self.grid_size, self.grid_size).reshape(1, -1) 60 | b = torch.linspace(-0.05, 0.05, steps=self.grid_size, dtype=torch.float).view(self.grid_size, 1).expand(self.grid_size, self.grid_size).reshape(1, -1) 61 | 62 | self.folding_seed = torch.cat([a, b], dim=0).view(1, 2, self.grid_size ** 2).cuda() # (1, 2, S) 63 | 64 | def forward(self, xyz): 65 | xyz = xyz.transpose(1,2) 66 | B, N, _ = xyz.shape 67 | feature = self.first_conv(xyz.transpose(2, 1)) # (B, 256, N) 68 | feature_global = torch.max(feature, dim=2, keepdim=True)[0] # (B, 256, 1) 69 | feature = torch.cat([feature_global.expand(-1, -1, N), feature], dim=1) # (B, 512, N) 70 | feature = self.second_conv(feature) # (B, 1024, N) 71 | feature_global = torch.max(feature,dim=2,keepdim=False)[0] # (B, 1024) 72 | 73 | # decoder 74 | coarse = self.mlp(feature_global).reshape(-1, self.num_coarse, 3) # (B, num_coarse, 3), coarse point cloud 75 | point_feat = coarse.unsqueeze(2).expand(-1, -1, self.grid_size ** 2, -1) # (B, num_coarse, S, 3) 76 | point_feat = point_feat.reshape(-1, self.num_dense, 3).transpose(2, 1) # (B, 3, num_fine) 77 | 78 | seed = self.folding_seed.unsqueeze(2).expand(B, -1, self.num_coarse, -1) # (B, 2, num_coarse, S) 79 | seed = seed.reshape(B, -1, self.num_dense) # (B, 2, num_fine) 80 | 81 | feature_global = feature_global.unsqueeze(2).expand(-1, -1, self.num_dense) # (B, 1024, num_fine) 82 | feat = torch.cat([feature_global, seed, point_feat], dim=1) # (B, 1024+2+3, num_fine) 83 | fine = self.final_conv(feat) + point_feat # (B, 3, num_fine), fine point cloud 84 | #return coarse.contiguous(), fine.contiguous() 85 | return fine.contiguous() -------------------------------------------------------------------------------- /src/models/pcnenc_polydec.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import torch.nn as nn 4 | 5 | 6 | class PCN_Conv(nn.Module): 7 | """ 8 | "PCN: Point Cloud Completion Network" 9 | (https://arxiv.org/pdf/1808.00671.pdf) 10 | Attributes: 11 | num_dense: 16384 12 | latent_dim: 1024 13 | grid_size: 4 14 | num_coarse: 1024 15 | """ 16 | 17 | def __init__(self, num_dense=8192, latent_dim=512, grid_size=4): 18 | super().__init__() 19 | 20 | self.num_dense = num_dense 21 | self.latent_dim = latent_dim 22 | self.grid_size = grid_size 23 | 24 | assert self.num_dense % self.grid_size ** 2 == 0 25 | 26 | self.num_coarse = self.num_dense // (self.grid_size ** 2) 27 | self.first_conv = nn.Sequential( 28 | nn.Conv1d(3, 128, 1), 29 | nn.BatchNorm1d(128), 30 | nn.ReLU(inplace=True), 31 | nn.Conv1d(128, 256, 1) 32 | ) 33 | 34 | self.second_conv = nn.Sequential( 35 | nn.Conv1d(512, 512, 1), 36 | nn.BatchNorm1d(512), 37 | nn.ReLU(inplace=True), 38 | nn.Conv1d(512, self.latent_dim, 1) 39 | ) 40 | 41 | layers = [] 42 | layers.append(nn.Linear(512, 1024)) 43 | layers.append(nn.ReLU(inplace=True)) 44 | for i in range(1): 45 | layers.append(nn.Linear(1024, 1024)) 46 | layers.append(nn.ReLU(inplace=True)) 47 | layers.append(nn.Linear(1024, 8192*3)) 48 | self.decoder = nn.Sequential(*layers) 49 | 50 | 51 | 52 | 53 | def forward(self, xyz): 54 | xyz = xyz.transpose(1,2) 55 | B, N, _ = xyz.shape 56 | feature = self.first_conv(xyz.transpose(2, 1)) # (B, 256, N) 57 | feature_global = torch.max(feature, dim=2, keepdim=True)[0] # (B, 256, 1) 58 | feature = torch.cat([feature_global.expand(-1, -1, N), feature], dim=1) # (B, 512, N) 59 | feature = self.second_conv(feature) # (B, 1024, N) 60 | feature_global = torch.max(feature,dim=2,keepdim=False)[0] # (B, 1024) 61 | 62 | # decoder 63 | x = self.decoder(feature_global).view(-1,3,8192) 64 | 65 | return x -------------------------------------------------------------------------------- /src/models/polyconv.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import time 5 | import numpy as np 6 | from pytorch3d.ops.knn import knn_gather, knn_points 7 | 8 | dtype = torch.double 9 | device = torch.device("cuda") 10 | 11 | 12 | 13 | def PolyConv(Input,indices,conv_p,conv_w): 14 | 15 | Input_adj = knn_gather(Input,indices) 16 | x = Input.unsqueeze(2).unsqueeze(-1) #Input_adj[:,0,:].unsqueeze(1).unsqueeze(-1) 17 | y = Input_adj[:,:,:].unsqueeze(-1) 18 | x_repeat = x*torch.ones_like(y) 19 | fxy = (torch.cat([torch.ones_like(y),x_repeat,y,torch.pow(x_repeat,2),x_repeat*y,torch.pow(y,2)],-1))*(conv_p.unsqueeze(0).unsqueeze(0).unsqueeze(0)) 20 | 21 | yfxy = fxy * y 22 | yfxy = torch.mean(torch.sum(yfxy,-1),dim=2) 23 | 24 | 25 | x_prime = torch.cat([2*torch.ones_like(x),2*x,torch.zeros_like(x),2*torch.pow(x,2),torch.zeros_like(x),(2/3)*torch.ones_like(x)],-1) 26 | 27 | fx = x_prime*(conv_p.unsqueeze(0).unsqueeze(0).unsqueeze(0)) 28 | 29 | 30 | fx = torch.sum(fx,-1)[:,:,0] 31 | 32 | fxy_cond = torch.true_divide(yfxy,fx+0.0001) 33 | fxy_cond = conv_w(fxy_cond) 34 | return fxy_cond 35 | 36 | 37 | def pool_max(Input, Adj, C, pool_num, conv_num,b): 38 | x = Input[:,Adj[b, :conv_num[1]]] 39 | x = scatter_max(x,C[b,:conv_num[1]],dim=1)[0] 40 | x = x[:, :pool_num[0]].float() 41 | return x 42 | 43 | def pool_mean(Input, Adj, C, pool_num, conv_num,b): 44 | x = Input[:,Adj[b, :conv_num[1]]] 45 | x = scatter_mean(x,C[b,:conv_num[1]],dim=1) 46 | x = x[:, :pool_num[0]].float() 47 | return x 48 | 49 | 50 | def network(Input,CONV, IN): 51 | Input = torch.transpose(Input,2,1) 52 | 53 | k = 6 54 | pool = True 55 | 56 | # First layer 57 | _,indices,_= knn_points(Input,Input,K=k) 58 | #x1 = torch.tanh(PolyConv(Input.float(),indices,CONV[0],CONV[1])) 59 | x1 = torch.transpose(torch.tanh(IN[0](torch.transpose(PolyConv(Input.float(),indices,CONV[0],CONV[1]),1,2))),1,2) 60 | pool_ind1 = torch.randint(x1.shape[1], (512,)) 61 | x1 = torch.max(knn_gather(x1,indices[:,pool_ind1,:]),2)[0] 62 | 63 | 64 | 65 | # Second layer 66 | _,indices,_= knn_points(Input[:,pool_ind1],Input[:,pool_ind1],K=k) 67 | #x2 = torch.tanh(PolyConv(x1,indices,CONV[2],CONV[3])) 68 | x2 = torch.transpose(torch.tanh(IN[1](torch.transpose(PolyConv(x1,indices,CONV[2],CONV[3]),1,2))),1,2) 69 | pool_ind2 = torch.randint(x2.shape[1], (128,)) 70 | x2 = torch.max(knn_gather(x2,indices[:,pool_ind2,:]),2)[0] 71 | 72 | 73 | # Third layer 74 | _,indices,_= knn_points(Input[:,pool_ind1[pool_ind2]],Input[:,pool_ind1[pool_ind2]],K=k) 75 | #x3 = torch.tanh(PolyConv(x2,indices,CONV[4],CONV[5])) 76 | x3 = torch.transpose(torch.tanh(IN[2](torch.transpose(PolyConv(x2,indices,CONV[4],CONV[5]),1,2))),1,2) 77 | pool_ind3 = torch.randint(x3.shape[1], (32,)) 78 | x3 = torch.max(knn_gather(x3,indices[:,pool_ind3,:]),2)[0] 79 | 80 | 81 | # Forth layer 82 | _,indices,_= knn_points(Input[:,pool_ind1[pool_ind2[pool_ind3]]],Input[:,pool_ind1[pool_ind2[pool_ind3]]],K=k) 83 | #x4 = torch.tanh(PolyConv(x3,indices,CONV[6],CONV[7])) 84 | x4 = torch.transpose(torch.tanh(IN[3](torch.transpose(PolyConv(x3,indices,CONV[6],CONV[7]),1,2))),1,2) 85 | 86 | Output = torch.mean(x4,1) 87 | return Output 88 | -------------------------------------------------------------------------------- /src/models/polynet.py: -------------------------------------------------------------------------------- 1 | if __name__ == "__main__": 2 | import os, sys 3 | 4 | sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__)))) 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | from torchsummary import summary 10 | from models.base_model import BaseModel 11 | # import torchvision.models as models 12 | from models.polyconv import network, PolyConv 13 | from torch.autograd import Variable 14 | import numpy as np 15 | device = torch.device("cuda") 16 | from pytorch3d.ops.knn import knn_points, knn_gather 17 | 18 | import pdb 19 | 20 | 21 | 22 | ''' 23 | phi = np.linspace(0, np.pi, 20) 24 | theta = np.linspace(0, 2 * np.pi, 40) 25 | x = np.outer(np.sin(theta), np.cos(phi)) 26 | y = np.outer(np.sin(theta), np.sin(phi)) 27 | z = np.outer(np.cos(theta), np.ones_like(phi)) 28 | 29 | xi, yi, zi = sample_spherical(64) 30 | x = torch.Tensor(xi) 31 | y = torch.Tensor(yi) 32 | z = torch.Tensor(zi) 33 | grid = torch.cat([x.unsqueeze(0),y.unsqueeze(0),z.unsqueeze(0)],0).unsqueeze(0).view(1,1,3,-1)/10 34 | grid = torch.transpose(grid,3,2) 35 | grid = torch.zeros_like(grid) 36 | ''' 37 | 38 | 39 | 40 | 41 | class PolyNet(BaseModel): 42 | def __init__(self): 43 | super().__init__() 44 | 45 | if True: 46 | self.Conv1_p = torch.nn.Parameter(torch.randn(3,6)) 47 | self.Conv1_w = nn.Linear(3, 64) 48 | 49 | self.Conv2_p = torch.nn.Parameter(torch.randn(64,6)) 50 | self.Conv2_w = nn.Linear(64, 128) 51 | 52 | self.Conv3_p = torch.nn.Parameter(torch.randn(128,6)) 53 | self.Conv3_w = nn.Linear(128, 256) 54 | 55 | self.Conv4_p = torch.nn.Parameter(torch.randn(256,6)) 56 | self.Conv4_w = nn.Linear(256, 512) 57 | 58 | 59 | 60 | 61 | 62 | self.in1 = nn.InstanceNorm1d(64,affine=True) 63 | self.in2 = nn.InstanceNorm1d(128,affine=True) 64 | self.in3 = nn.InstanceNorm1d(256,affine=True) 65 | self.in4 = nn.InstanceNorm1d(512,affine=True) 66 | 67 | #self.bn5 = nn.BatchNorm1d(1024) 68 | #self.bn2 = nn.BatchNorm1d(64) 69 | #self.bn0 = nn.BatchNorm1d(1024) 70 | 71 | #self.drop1 = nn.Dropout(p=0.2) 72 | #self.drop2 = nn.Dropout(p=0.5) 73 | layers = [] 74 | layers.append(nn.Linear(512, 1024)) 75 | layers.append(nn.ReLU(inplace=True)) 76 | for i in range(1): 77 | layers.append(nn.Linear(1024, 1024)) 78 | layers.append(nn.ReLU(inplace=True)) 79 | layers.append(nn.Linear(1024, 8192*3)) 80 | self.decoder = nn.Sequential(*layers) 81 | 82 | 83 | 84 | else: 85 | print("Check model name in config.json") 86 | raise 87 | 88 | def forward(self, input): 89 | def _closeness_constraint(netcoeff): # (num_funcs, 6) 90 | # 6 parameters 91 | B = torch.zeros((netcoeff.shape[0], 3, 3), device='cuda') 92 | triu_idcs = torch.triu_indices(row=3, col=3, offset=0).to('cuda') 93 | B[:, triu_idcs[0], triu_idcs[1]] = netcoeff # vector to upper triangular matrix 94 | B[:, triu_idcs[1], triu_idcs[0]] = netcoeff # B: symm. matrix 95 | A = torch.bmm(B, B) # A = B**2 // A: symm. positive definite (num_funcs, 3,3) 96 | 97 | # [1, x, y, x^2, xy, y^2] 98 | p4coeff = torch.zeros((netcoeff.shape[0], 6), device='cuda') 99 | p4coeff[:, 0] = A[:, 0,0] # 1 100 | p4coeff[:, 3] = A[:, 1,1] # x^2 101 | p4coeff[:, 5] = A[:, 2,2] # y^2 102 | 103 | p4coeff[:, 1] = A[:, 1,0]+A[:, 0,1] # x 104 | p4coeff[:, 2] = A[:, 2,0]+A[:, 0,2] # y 105 | p4coeff[:, 4] = A[:, 1,2]+A[:, 2,1] # xy 106 | return p4coeff 107 | 108 | 109 | if True: 110 | Conv1_p = _closeness_constraint(self.Conv1_p) 111 | Conv2_p = _closeness_constraint(self.Conv2_p) 112 | Conv3_p = _closeness_constraint(self.Conv3_p) 113 | Conv4_p = _closeness_constraint(self.Conv4_p) 114 | 115 | 116 | 117 | 118 | CONV = [Conv1_p,self.Conv1_w,Conv2_p,self.Conv2_w,Conv3_p,self.Conv3_w,Conv4_p,self.Conv4_w] 119 | IN = [self.in1,self.in2,self.in3,self.in4] 120 | 121 | 122 | feat = network(input, CONV, IN) 123 | #x = torch.mean(x,1) 124 | 125 | 126 | x = self.decoder(feat).view(-1,3,8192) 127 | 128 | return x 129 | 130 | 131 | if __name__ == "__main__": 132 | 133 | a = torch.rand(4,3,3096).to('cuda:0') 134 | net = PolyNet().to('cuda:0') 135 | output = net(a) 136 | -------------------------------------------------------------------------------- /src/models/polynet_enc_pcn_dec.py: -------------------------------------------------------------------------------- 1 | if __name__ == "__main__": 2 | import os, sys 3 | 4 | sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__)))) 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | from torchsummary import summary 10 | from models.base_model import BaseModel 11 | # import torchvision.models as models 12 | from models.polyconv import network, PolyConv 13 | from torch.autograd import Variable 14 | import numpy as np 15 | device = torch.device("cuda") 16 | from pytorch3d.ops.knn import knn_points, knn_gather 17 | 18 | import pdb 19 | 20 | 21 | 22 | ''' 23 | phi = np.linspace(0, np.pi, 20) 24 | theta = np.linspace(0, 2 * np.pi, 40) 25 | x = np.outer(np.sin(theta), np.cos(phi)) 26 | y = np.outer(np.sin(theta), np.sin(phi)) 27 | z = np.outer(np.cos(theta), np.ones_like(phi)) 28 | 29 | xi, yi, zi = sample_spherical(64) 30 | x = torch.Tensor(xi) 31 | y = torch.Tensor(yi) 32 | z = torch.Tensor(zi) 33 | grid = torch.cat([x.unsqueeze(0),y.unsqueeze(0),z.unsqueeze(0)],0).unsqueeze(0).view(1,1,3,-1)/10 34 | grid = torch.transpose(grid,3,2) 35 | grid = torch.zeros_like(grid) 36 | ''' 37 | 38 | 39 | 40 | 41 | class PolyPCN(BaseModel): 42 | def __init__(self): 43 | super().__init__() 44 | 45 | self.grid_size=8 46 | self.num_dense = 8192 47 | self.latent_dim = 512 48 | self.num_coarse = self.num_dense // (self.grid_size ** 2) 49 | if True: 50 | self.Conv1_p = torch.nn.Parameter(torch.randn(3,6)) 51 | self.Conv1_w = nn.Linear(3, 64) 52 | 53 | self.Conv2_p = torch.nn.Parameter(torch.randn(64,6)) 54 | self.Conv2_w = nn.Linear(64, 128) 55 | 56 | self.Conv3_p = torch.nn.Parameter(torch.randn(128,6)) 57 | self.Conv3_w = nn.Linear(128, 256) 58 | 59 | self.Conv4_p = torch.nn.Parameter(torch.randn(256,6)) 60 | self.Conv4_w = nn.Linear(256, 512) 61 | 62 | 63 | 64 | 65 | 66 | self.in1 = nn.InstanceNorm1d(64,affine=True) 67 | self.in2 = nn.InstanceNorm1d(128,affine=True) 68 | self.in3 = nn.InstanceNorm1d(256,affine=True) 69 | self.in4 = nn.InstanceNorm1d(512,affine=True) 70 | 71 | #self.bn5 = nn.BatchNorm1d(1024) 72 | #self.bn2 = nn.BatchNorm1d(64) 73 | #self.bn0 = nn.BatchNorm1d(1024) 74 | 75 | #self.drop1 = nn.Dropout(p=0.2) 76 | #self.drop2 = nn.Dropout(p=0.5) 77 | 78 | self.mlp = nn.Sequential( 79 | nn.Linear(self.latent_dim, 1024), 80 | nn.ReLU(inplace=True), 81 | nn.Linear(1024, 1024), 82 | nn.ReLU(inplace=True), 83 | nn.Linear(1024, 3 * self.num_coarse) 84 | ) 85 | 86 | self.final_conv = nn.Sequential( 87 | #nn.Conv1d(1024 + 3 + 2, 512, 1) 88 | nn.Conv1d(512 + 3 + 2, 512, 1), 89 | nn.BatchNorm1d(512), 90 | nn.ReLU(inplace=True), 91 | nn.Conv1d(512, 512, 1), 92 | nn.BatchNorm1d(512), 93 | nn.ReLU(inplace=True), 94 | nn.Conv1d(512, 3, 1) 95 | ) 96 | a = torch.linspace(-0.05, 0.05, steps=self.grid_size, dtype=torch.float).view(1, self.grid_size).expand(self.grid_size, self.grid_size).reshape(1, -1) 97 | b = torch.linspace(-0.05, 0.05, steps=self.grid_size, dtype=torch.float).view(self.grid_size, 1).expand(self.grid_size, self.grid_size).reshape(1, -1) 98 | 99 | self.folding_seed = torch.cat([a, b], dim=0).view(1, 2, self.grid_size ** 2).cuda() # (1, 2, S) 100 | 101 | 102 | 103 | else: 104 | print("Check model name in config.json") 105 | raise 106 | 107 | def forward(self, input): 108 | def _closeness_constraint(netcoeff): # (num_funcs, 6) 109 | # 6 parameters 110 | B = torch.zeros((netcoeff.shape[0], 3, 3), device='cuda') 111 | triu_idcs = torch.triu_indices(row=3, col=3, offset=0).to('cuda') 112 | B[:, triu_idcs[0], triu_idcs[1]] = netcoeff # vector to upper triangular matrix 113 | B[:, triu_idcs[1], triu_idcs[0]] = netcoeff # B: symm. matrix 114 | A = torch.bmm(B, B) # A = B**2 // A: symm. positive definite (num_funcs, 3,3) 115 | 116 | # [1, x, y, x^2, xy, y^2] 117 | p4coeff = torch.zeros((netcoeff.shape[0], 6), device='cuda') 118 | p4coeff[:, 0] = A[:, 0,0] # 1 119 | p4coeff[:, 3] = A[:, 1,1] # x^2 120 | p4coeff[:, 5] = A[:, 2,2] # y^2 121 | 122 | p4coeff[:, 1] = A[:, 1,0]+A[:, 0,1] # x 123 | p4coeff[:, 2] = A[:, 2,0]+A[:, 0,2] # y 124 | p4coeff[:, 4] = A[:, 1,2]+A[:, 2,1] # xy 125 | return p4coeff 126 | 127 | B, _, _ = input.shape 128 | 129 | if True: 130 | Conv1_p = _closeness_constraint(self.Conv1_p) 131 | Conv2_p = _closeness_constraint(self.Conv2_p) 132 | Conv3_p = _closeness_constraint(self.Conv3_p) 133 | Conv4_p = _closeness_constraint(self.Conv4_p) 134 | 135 | 136 | 137 | 138 | CONV = [Conv1_p,self.Conv1_w,Conv2_p,self.Conv2_w,Conv3_p,self.Conv3_w,Conv4_p,self.Conv4_w] 139 | IN = [self.in1,self.in2,self.in3,self.in4] 140 | 141 | 142 | feat = network(input, CONV, IN) 143 | #x = torch.mean(x,1) 144 | 145 | 146 | # decoder 147 | coarse = self.mlp(feat).reshape(-1, self.num_coarse, 3) # (B, num_coarse, 3), coarse point cloud 148 | point_feat = coarse.unsqueeze(2).expand(-1, -1, self.grid_size ** 2, -1) # (B, num_coarse, S, 3) 149 | point_feat = point_feat.reshape(-1, self.num_dense, 3).transpose(2, 1) # (B, 3, num_fine) 150 | seed = self.folding_seed.unsqueeze(2).expand(B, -1, self.num_coarse, -1) # (B, 2, num_coarse, S) 151 | seed = seed.reshape(B, -1, self.num_dense) # (B, 2, num_fine) 152 | 153 | feature_global = feat.unsqueeze(2).expand(-1, -1, self.num_dense) # (B, 1024, num_fine) 154 | feat = torch.cat([feature_global, seed, point_feat], dim=1) # (B, 1024+2+3, num_fine) 155 | fine = self.final_conv(feat) + point_feat # (B, 3, num_fine), fine point cloud 156 | #return coarse.contiguous(), fine.contiguous() 157 | 158 | return fine.contiguous() 159 | 160 | 161 | 162 | 163 | if __name__ == "__main__": 164 | 165 | a = torch.rand(4,3,3096).to('cuda:0') 166 | net = PolyNet().to('cuda:0') 167 | output = net(a) 168 | -------------------------------------------------------------------------------- /src/pointnet2/__pycache__/pointnet2_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/__pycache__/pointnet2_utils.cpython-36.pyc -------------------------------------------------------------------------------- /src/pointnet2/__pycache__/pointnet2_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/__pycache__/pointnet2_utils.cpython-38.pyc -------------------------------------------------------------------------------- /src/pointnet2/__pycache__/pytorch_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/__pycache__/pytorch_utils.cpython-36.pyc -------------------------------------------------------------------------------- /src/pointnet2/__pycache__/pytorch_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/__pycache__/pytorch_utils.cpython-38.pyc -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/include/ball_query.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | 9 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 10 | const int nsample); 11 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/include/cuda_utils.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #ifndef _CUDA_UTILS_H 7 | #define _CUDA_UTILS_H 8 | 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | #include 15 | 16 | #include 17 | 18 | #define TOTAL_THREADS 512 19 | 20 | inline int opt_n_threads(int work_size) { 21 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 22 | 23 | return max(min(1 << pow_2, TOTAL_THREADS), 1); 24 | } 25 | 26 | inline dim3 opt_block_config(int x, int y) { 27 | const int x_threads = opt_n_threads(x); 28 | const int y_threads = 29 | max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 30 | dim3 block_config(x_threads, y_threads, 1); 31 | 32 | return block_config; 33 | } 34 | 35 | #define CUDA_CHECK_ERRORS() \ 36 | do { \ 37 | cudaError_t err = cudaGetLastError(); \ 38 | if (cudaSuccess != err) { \ 39 | fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ 40 | cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ 41 | __FILE__); \ 42 | exit(-1); \ 43 | } \ 44 | } while (0) 45 | 46 | #endif 47 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/include/group_points.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | 9 | at::Tensor group_points(at::Tensor points, at::Tensor idx); 10 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 11 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/include/interpolate.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | 8 | #include 9 | #include 10 | 11 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows); 12 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 13 | at::Tensor weight); 14 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 15 | at::Tensor weight, const int m); 16 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/include/sampling.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | 9 | at::Tensor gather_points(at::Tensor points, at::Tensor idx); 10 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 11 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples); 12 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/include/utils.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | #include 9 | 10 | #define CHECK_CUDA(x) \ 11 | do { \ 12 | AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor"); \ 13 | } while (0) 14 | 15 | #define CHECK_CONTIGUOUS(x) \ 16 | do { \ 17 | AT_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \ 18 | } while (0) 19 | 20 | #define CHECK_IS_INT(x) \ 21 | do { \ 22 | AT_CHECK(x.scalar_type() == at::ScalarType::Int, \ 23 | #x " must be an int tensor"); \ 24 | } while (0) 25 | 26 | #define CHECK_IS_FLOAT(x) \ 27 | do { \ 28 | AT_CHECK(x.scalar_type() == at::ScalarType::Float, \ 29 | #x " must be a float tensor"); \ 30 | } while (0) 31 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "ball_query.h" 7 | #include "utils.h" 8 | 9 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 10 | int nsample, const float *new_xyz, 11 | const float *xyz, int *idx); 12 | 13 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 14 | const int nsample) { 15 | CHECK_CONTIGUOUS(new_xyz); 16 | CHECK_CONTIGUOUS(xyz); 17 | CHECK_IS_FLOAT(new_xyz); 18 | CHECK_IS_FLOAT(xyz); 19 | 20 | if (new_xyz.type().is_cuda()) { 21 | CHECK_CUDA(xyz); 22 | } 23 | 24 | at::Tensor idx = 25 | torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample}, 26 | at::device(new_xyz.device()).dtype(at::ScalarType::Int)); 27 | 28 | if (new_xyz.type().is_cuda()) { 29 | query_ball_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), 30 | radius, nsample, new_xyz.data(), 31 | xyz.data(), idx.data()); 32 | } else { 33 | AT_CHECK(false, "CPU not supported"); 34 | } 35 | 36 | return idx; 37 | } 38 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/src/ball_query_gpu.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include 7 | #include 8 | #include 9 | 10 | #include "cuda_utils.h" 11 | 12 | // input: new_xyz(b, m, 3) xyz(b, n, 3) 13 | // output: idx(b, m, nsample) 14 | __global__ void query_ball_point_kernel(int b, int n, int m, float radius, 15 | int nsample, 16 | const float *__restrict__ new_xyz, 17 | const float *__restrict__ xyz, 18 | int *__restrict__ idx) { 19 | int batch_index = blockIdx.x; 20 | xyz += batch_index * n * 3; 21 | new_xyz += batch_index * m * 3; 22 | idx += m * nsample * batch_index; 23 | 24 | int index = threadIdx.x; 25 | int stride = blockDim.x; 26 | 27 | float radius2 = radius * radius; 28 | for (int j = index; j < m; j += stride) { 29 | float new_x = new_xyz[j * 3 + 0]; 30 | float new_y = new_xyz[j * 3 + 1]; 31 | float new_z = new_xyz[j * 3 + 2]; 32 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { 33 | float x = xyz[k * 3 + 0]; 34 | float y = xyz[k * 3 + 1]; 35 | float z = xyz[k * 3 + 2]; 36 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 37 | (new_z - z) * (new_z - z); 38 | if (d2 < radius2) { 39 | if (cnt == 0) { 40 | for (int l = 0; l < nsample; ++l) { 41 | idx[j * nsample + l] = k; 42 | } 43 | } 44 | idx[j * nsample + cnt] = k; 45 | ++cnt; 46 | } 47 | } 48 | } 49 | } 50 | 51 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 52 | int nsample, const float *new_xyz, 53 | const float *xyz, int *idx) { 54 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 55 | query_ball_point_kernel<<>>( 56 | b, n, m, radius, nsample, new_xyz, xyz, idx); 57 | 58 | CUDA_CHECK_ERRORS(); 59 | } 60 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/src/bindings.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "ball_query.h" 7 | #include "group_points.h" 8 | #include "interpolate.h" 9 | #include "sampling.h" 10 | 11 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 12 | m.def("gather_points", &gather_points); 13 | m.def("gather_points_grad", &gather_points_grad); 14 | m.def("furthest_point_sampling", &furthest_point_sampling); 15 | 16 | m.def("three_nn", &three_nn); 17 | m.def("three_interpolate", &three_interpolate); 18 | m.def("three_interpolate_grad", &three_interpolate_grad); 19 | 20 | m.def("ball_query", &ball_query); 21 | 22 | m.def("group_points", &group_points); 23 | m.def("group_points_grad", &group_points_grad); 24 | } 25 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/src/group_points.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "group_points.h" 7 | #include "utils.h" 8 | 9 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 10 | const float *points, const int *idx, 11 | float *out); 12 | 13 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 14 | int nsample, const float *grad_out, 15 | const int *idx, float *grad_points); 16 | 17 | at::Tensor group_points(at::Tensor points, at::Tensor idx) { 18 | CHECK_CONTIGUOUS(points); 19 | CHECK_CONTIGUOUS(idx); 20 | CHECK_IS_FLOAT(points); 21 | CHECK_IS_INT(idx); 22 | 23 | if (points.type().is_cuda()) { 24 | CHECK_CUDA(idx); 25 | } 26 | 27 | at::Tensor output = 28 | torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)}, 29 | at::device(points.device()).dtype(at::ScalarType::Float)); 30 | 31 | if (points.type().is_cuda()) { 32 | group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 33 | idx.size(1), idx.size(2), points.data(), 34 | idx.data(), output.data()); 35 | } else { 36 | AT_CHECK(false, "CPU not supported"); 37 | } 38 | 39 | return output; 40 | } 41 | 42 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) { 43 | CHECK_CONTIGUOUS(grad_out); 44 | CHECK_CONTIGUOUS(idx); 45 | CHECK_IS_FLOAT(grad_out); 46 | CHECK_IS_INT(idx); 47 | 48 | if (grad_out.type().is_cuda()) { 49 | CHECK_CUDA(idx); 50 | } 51 | 52 | at::Tensor output = 53 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 54 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 55 | 56 | if (grad_out.type().is_cuda()) { 57 | group_points_grad_kernel_wrapper( 58 | grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2), 59 | grad_out.data(), idx.data(), output.data()); 60 | } else { 61 | AT_CHECK(false, "CPU not supported"); 62 | } 63 | 64 | return output; 65 | } 66 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/src/group_points_gpu.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include 7 | #include 8 | 9 | #include "cuda_utils.h" 10 | 11 | // input: points(b, c, n) idx(b, npoints, nsample) 12 | // output: out(b, c, npoints, nsample) 13 | __global__ void group_points_kernel(int b, int c, int n, int npoints, 14 | int nsample, 15 | const float *__restrict__ points, 16 | const int *__restrict__ idx, 17 | float *__restrict__ out) { 18 | int batch_index = blockIdx.x; 19 | points += batch_index * n * c; 20 | idx += batch_index * npoints * nsample; 21 | out += batch_index * npoints * nsample * c; 22 | 23 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 24 | const int stride = blockDim.y * blockDim.x; 25 | for (int i = index; i < c * npoints; i += stride) { 26 | const int l = i / npoints; 27 | const int j = i % npoints; 28 | for (int k = 0; k < nsample; ++k) { 29 | int ii = idx[j * nsample + k]; 30 | out[(l * npoints + j) * nsample + k] = points[l * n + ii]; 31 | } 32 | } 33 | } 34 | 35 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 36 | const float *points, const int *idx, 37 | float *out) { 38 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 39 | 40 | group_points_kernel<<>>( 41 | b, c, n, npoints, nsample, points, idx, out); 42 | 43 | CUDA_CHECK_ERRORS(); 44 | } 45 | 46 | // input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample) 47 | // output: grad_points(b, c, n) 48 | __global__ void group_points_grad_kernel(int b, int c, int n, int npoints, 49 | int nsample, 50 | const float *__restrict__ grad_out, 51 | const int *__restrict__ idx, 52 | float *__restrict__ grad_points) { 53 | int batch_index = blockIdx.x; 54 | grad_out += batch_index * npoints * nsample * c; 55 | idx += batch_index * npoints * nsample; 56 | grad_points += batch_index * n * c; 57 | 58 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 59 | const int stride = blockDim.y * blockDim.x; 60 | for (int i = index; i < c * npoints; i += stride) { 61 | const int l = i / npoints; 62 | const int j = i % npoints; 63 | for (int k = 0; k < nsample; ++k) { 64 | int ii = idx[j * nsample + k]; 65 | atomicAdd(grad_points + l * n + ii, 66 | grad_out[(l * npoints + j) * nsample + k]); 67 | } 68 | } 69 | } 70 | 71 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 72 | int nsample, const float *grad_out, 73 | const int *idx, float *grad_points) { 74 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 75 | 76 | group_points_grad_kernel<<>>( 77 | b, c, n, npoints, nsample, grad_out, idx, grad_points); 78 | 79 | CUDA_CHECK_ERRORS(); 80 | } 81 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/src/interpolate.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "interpolate.h" 7 | #include "utils.h" 8 | 9 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 10 | const float *known, float *dist2, int *idx); 11 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 12 | const float *points, const int *idx, 13 | const float *weight, float *out); 14 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 15 | const float *grad_out, 16 | const int *idx, const float *weight, 17 | float *grad_points); 18 | 19 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows) { 20 | CHECK_CONTIGUOUS(unknowns); 21 | CHECK_CONTIGUOUS(knows); 22 | CHECK_IS_FLOAT(unknowns); 23 | CHECK_IS_FLOAT(knows); 24 | 25 | if (unknowns.type().is_cuda()) { 26 | CHECK_CUDA(knows); 27 | } 28 | 29 | at::Tensor idx = 30 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 31 | at::device(unknowns.device()).dtype(at::ScalarType::Int)); 32 | at::Tensor dist2 = 33 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 34 | at::device(unknowns.device()).dtype(at::ScalarType::Float)); 35 | 36 | if (unknowns.type().is_cuda()) { 37 | three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1), 38 | unknowns.data(), knows.data(), 39 | dist2.data(), idx.data()); 40 | } else { 41 | AT_CHECK(false, "CPU not supported"); 42 | } 43 | 44 | return {dist2, idx}; 45 | } 46 | 47 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 48 | at::Tensor weight) { 49 | CHECK_CONTIGUOUS(points); 50 | CHECK_CONTIGUOUS(idx); 51 | CHECK_CONTIGUOUS(weight); 52 | CHECK_IS_FLOAT(points); 53 | CHECK_IS_INT(idx); 54 | CHECK_IS_FLOAT(weight); 55 | 56 | if (points.type().is_cuda()) { 57 | CHECK_CUDA(idx); 58 | CHECK_CUDA(weight); 59 | } 60 | 61 | at::Tensor output = 62 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 63 | at::device(points.device()).dtype(at::ScalarType::Float)); 64 | 65 | if (points.type().is_cuda()) { 66 | three_interpolate_kernel_wrapper( 67 | points.size(0), points.size(1), points.size(2), idx.size(1), 68 | points.data(), idx.data(), weight.data(), 69 | output.data()); 70 | } else { 71 | AT_CHECK(false, "CPU not supported"); 72 | } 73 | 74 | return output; 75 | } 76 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 77 | at::Tensor weight, const int m) { 78 | CHECK_CONTIGUOUS(grad_out); 79 | CHECK_CONTIGUOUS(idx); 80 | CHECK_CONTIGUOUS(weight); 81 | CHECK_IS_FLOAT(grad_out); 82 | CHECK_IS_INT(idx); 83 | CHECK_IS_FLOAT(weight); 84 | 85 | if (grad_out.type().is_cuda()) { 86 | CHECK_CUDA(idx); 87 | CHECK_CUDA(weight); 88 | } 89 | 90 | at::Tensor output = 91 | torch::zeros({grad_out.size(0), grad_out.size(1), m}, 92 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 93 | 94 | if (grad_out.type().is_cuda()) { 95 | three_interpolate_kernel_wrapper( 96 | grad_out.size(0), grad_out.size(1), grad_out.size(2), m, 97 | grad_out.data(), idx.data(), weight.data(), 98 | output.data()); 99 | } else { 100 | AT_CHECK(false, "CPU not supported"); 101 | } 102 | 103 | return output; 104 | } 105 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/src/interpolate_gpu.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include 7 | #include 8 | #include 9 | 10 | #include "cuda_utils.h" 11 | 12 | // input: unknown(b, n, 3) known(b, m, 3) 13 | // output: dist2(b, n, 3), idx(b, n, 3) 14 | __global__ void three_nn_kernel(int b, int n, int m, 15 | const float *__restrict__ unknown, 16 | const float *__restrict__ known, 17 | float *__restrict__ dist2, 18 | int *__restrict__ idx) { 19 | int batch_index = blockIdx.x; 20 | unknown += batch_index * n * 3; 21 | known += batch_index * m * 3; 22 | dist2 += batch_index * n * 3; 23 | idx += batch_index * n * 3; 24 | 25 | int index = threadIdx.x; 26 | int stride = blockDim.x; 27 | for (int j = index; j < n; j += stride) { 28 | float ux = unknown[j * 3 + 0]; 29 | float uy = unknown[j * 3 + 1]; 30 | float uz = unknown[j * 3 + 2]; 31 | 32 | double best1 = 1e40, best2 = 1e40, best3 = 1e40; 33 | int besti1 = 0, besti2 = 0, besti3 = 0; 34 | for (int k = 0; k < m; ++k) { 35 | float x = known[k * 3 + 0]; 36 | float y = known[k * 3 + 1]; 37 | float z = known[k * 3 + 2]; 38 | float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z); 39 | if (d < best1) { 40 | best3 = best2; 41 | besti3 = besti2; 42 | best2 = best1; 43 | besti2 = besti1; 44 | best1 = d; 45 | besti1 = k; 46 | } else if (d < best2) { 47 | best3 = best2; 48 | besti3 = besti2; 49 | best2 = d; 50 | besti2 = k; 51 | } else if (d < best3) { 52 | best3 = d; 53 | besti3 = k; 54 | } 55 | } 56 | dist2[j * 3 + 0] = best1; 57 | dist2[j * 3 + 1] = best2; 58 | dist2[j * 3 + 2] = best3; 59 | 60 | idx[j * 3 + 0] = besti1; 61 | idx[j * 3 + 1] = besti2; 62 | idx[j * 3 + 2] = besti3; 63 | } 64 | } 65 | 66 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 67 | const float *known, float *dist2, int *idx) { 68 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 69 | three_nn_kernel<<>>(b, n, m, unknown, known, 70 | dist2, idx); 71 | 72 | CUDA_CHECK_ERRORS(); 73 | } 74 | 75 | // input: points(b, c, m), idx(b, n, 3), weight(b, n, 3) 76 | // output: out(b, c, n) 77 | __global__ void three_interpolate_kernel(int b, int c, int m, int n, 78 | const float *__restrict__ points, 79 | const int *__restrict__ idx, 80 | const float *__restrict__ weight, 81 | float *__restrict__ out) { 82 | int batch_index = blockIdx.x; 83 | points += batch_index * m * c; 84 | 85 | idx += batch_index * n * 3; 86 | weight += batch_index * n * 3; 87 | 88 | out += batch_index * n * c; 89 | 90 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 91 | const int stride = blockDim.y * blockDim.x; 92 | for (int i = index; i < c * n; i += stride) { 93 | const int l = i / n; 94 | const int j = i % n; 95 | float w1 = weight[j * 3 + 0]; 96 | float w2 = weight[j * 3 + 1]; 97 | float w3 = weight[j * 3 + 2]; 98 | 99 | int i1 = idx[j * 3 + 0]; 100 | int i2 = idx[j * 3 + 1]; 101 | int i3 = idx[j * 3 + 2]; 102 | 103 | out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 + 104 | points[l * m + i3] * w3; 105 | } 106 | } 107 | 108 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 109 | const float *points, const int *idx, 110 | const float *weight, float *out) { 111 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 112 | three_interpolate_kernel<<>>( 113 | b, c, m, n, points, idx, weight, out); 114 | 115 | CUDA_CHECK_ERRORS(); 116 | } 117 | 118 | // input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3) 119 | // output: grad_points(b, c, m) 120 | 121 | __global__ void three_interpolate_grad_kernel( 122 | int b, int c, int n, int m, const float *__restrict__ grad_out, 123 | const int *__restrict__ idx, const float *__restrict__ weight, 124 | float *__restrict__ grad_points) { 125 | int batch_index = blockIdx.x; 126 | grad_out += batch_index * n * c; 127 | idx += batch_index * n * 3; 128 | weight += batch_index * n * 3; 129 | grad_points += batch_index * m * c; 130 | 131 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 132 | const int stride = blockDim.y * blockDim.x; 133 | for (int i = index; i < c * n; i += stride) { 134 | const int l = i / n; 135 | const int j = i % n; 136 | float w1 = weight[j * 3 + 0]; 137 | float w2 = weight[j * 3 + 1]; 138 | float w3 = weight[j * 3 + 2]; 139 | 140 | int i1 = idx[j * 3 + 0]; 141 | int i2 = idx[j * 3 + 1]; 142 | int i3 = idx[j * 3 + 2]; 143 | 144 | atomicAdd(grad_points + l * m + i1, grad_out[i] * w1); 145 | atomicAdd(grad_points + l * m + i2, grad_out[i] * w2); 146 | atomicAdd(grad_points + l * m + i3, grad_out[i] * w3); 147 | } 148 | } 149 | 150 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 151 | const float *grad_out, 152 | const int *idx, const float *weight, 153 | float *grad_points) { 154 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 155 | three_interpolate_grad_kernel<<>>( 156 | b, c, n, m, grad_out, idx, weight, grad_points); 157 | 158 | CUDA_CHECK_ERRORS(); 159 | } 160 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/src/sampling.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "sampling.h" 7 | #include "utils.h" 8 | 9 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 10 | const float *points, const int *idx, 11 | float *out); 12 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 13 | const float *grad_out, const int *idx, 14 | float *grad_points); 15 | 16 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 17 | const float *dataset, float *temp, 18 | int *idxs); 19 | 20 | at::Tensor gather_points(at::Tensor points, at::Tensor idx) { 21 | CHECK_CONTIGUOUS(points); 22 | CHECK_CONTIGUOUS(idx); 23 | CHECK_IS_FLOAT(points); 24 | CHECK_IS_INT(idx); 25 | 26 | if (points.type().is_cuda()) { 27 | CHECK_CUDA(idx); 28 | } 29 | 30 | at::Tensor output = 31 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 32 | at::device(points.device()).dtype(at::ScalarType::Float)); 33 | 34 | if (points.type().is_cuda()) { 35 | gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 36 | idx.size(1), points.data(), 37 | idx.data(), output.data()); 38 | } else { 39 | AT_CHECK(false, "CPU not supported"); 40 | } 41 | 42 | return output; 43 | } 44 | 45 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, 46 | const int n) { 47 | CHECK_CONTIGUOUS(grad_out); 48 | CHECK_CONTIGUOUS(idx); 49 | CHECK_IS_FLOAT(grad_out); 50 | CHECK_IS_INT(idx); 51 | 52 | if (grad_out.type().is_cuda()) { 53 | CHECK_CUDA(idx); 54 | } 55 | 56 | at::Tensor output = 57 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 58 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 59 | 60 | if (grad_out.type().is_cuda()) { 61 | gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n, 62 | idx.size(1), grad_out.data(), 63 | idx.data(), output.data()); 64 | } else { 65 | AT_CHECK(false, "CPU not supported"); 66 | } 67 | 68 | return output; 69 | } 70 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) { 71 | CHECK_CONTIGUOUS(points); 72 | CHECK_IS_FLOAT(points); 73 | 74 | at::Tensor output = 75 | torch::zeros({points.size(0), nsamples}, 76 | at::device(points.device()).dtype(at::ScalarType::Int)); 77 | 78 | at::Tensor tmp = 79 | torch::full({points.size(0), points.size(1)}, 1e10, 80 | at::device(points.device()).dtype(at::ScalarType::Float)); 81 | 82 | if (points.type().is_cuda()) { 83 | furthest_point_sampling_kernel_wrapper( 84 | points.size(0), points.size(1), nsamples, points.data(), 85 | tmp.data(), output.data()); 86 | } else { 87 | AT_CHECK(false, "CPU not supported"); 88 | } 89 | 90 | return output; 91 | } 92 | -------------------------------------------------------------------------------- /src/pointnet2/_ext_src/src/sampling_gpu.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include 7 | #include 8 | 9 | #include "cuda_utils.h" 10 | 11 | // input: points(b, c, n) idx(b, m) 12 | // output: out(b, c, m) 13 | __global__ void gather_points_kernel(int b, int c, int n, int m, 14 | const float *__restrict__ points, 15 | const int *__restrict__ idx, 16 | float *__restrict__ out) { 17 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 18 | for (int l = blockIdx.y; l < c; l += gridDim.y) { 19 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 20 | int a = idx[i * m + j]; 21 | out[(i * c + l) * m + j] = points[(i * c + l) * n + a]; 22 | } 23 | } 24 | } 25 | } 26 | 27 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 28 | const float *points, const int *idx, 29 | float *out) { 30 | gather_points_kernel<<>>(b, c, n, npoints, 32 | points, idx, out); 33 | 34 | CUDA_CHECK_ERRORS(); 35 | } 36 | 37 | // input: grad_out(b, c, m) idx(b, m) 38 | // output: grad_points(b, c, n) 39 | __global__ void gather_points_grad_kernel(int b, int c, int n, int m, 40 | const float *__restrict__ grad_out, 41 | const int *__restrict__ idx, 42 | float *__restrict__ grad_points) { 43 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 44 | for (int l = blockIdx.y; l < c; l += gridDim.y) { 45 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 46 | int a = idx[i * m + j]; 47 | atomicAdd(grad_points + (i * c + l) * n + a, 48 | grad_out[(i * c + l) * m + j]); 49 | } 50 | } 51 | } 52 | } 53 | 54 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 55 | const float *grad_out, const int *idx, 56 | float *grad_points) { 57 | gather_points_grad_kernel<<>>( 59 | b, c, n, npoints, grad_out, idx, grad_points); 60 | 61 | CUDA_CHECK_ERRORS(); 62 | } 63 | 64 | __device__ void __update(float *__restrict__ dists, int *__restrict__ dists_i, 65 | int idx1, int idx2) { 66 | const float v1 = dists[idx1], v2 = dists[idx2]; 67 | const int i1 = dists_i[idx1], i2 = dists_i[idx2]; 68 | dists[idx1] = max(v1, v2); 69 | dists_i[idx1] = v2 > v1 ? i2 : i1; 70 | } 71 | 72 | // Input dataset: (b, n, 3), tmp: (b, n) 73 | // Ouput idxs (b, m) 74 | template 75 | __global__ void furthest_point_sampling_kernel( 76 | int b, int n, int m, const float *__restrict__ dataset, 77 | float *__restrict__ temp, int *__restrict__ idxs) { 78 | if (m <= 0) return; 79 | __shared__ float dists[block_size]; 80 | __shared__ int dists_i[block_size]; 81 | 82 | int batch_index = blockIdx.x; 83 | dataset += batch_index * n * 3; 84 | temp += batch_index * n; 85 | idxs += batch_index * m; 86 | 87 | int tid = threadIdx.x; 88 | const int stride = block_size; 89 | 90 | int old = 0; 91 | if (threadIdx.x == 0) idxs[0] = old; 92 | 93 | __syncthreads(); 94 | for (int j = 1; j < m; j++) { 95 | int besti = 0; 96 | float best = -1; 97 | float x1 = dataset[old * 3 + 0]; 98 | float y1 = dataset[old * 3 + 1]; 99 | float z1 = dataset[old * 3 + 2]; 100 | for (int k = tid; k < n; k += stride) { 101 | float x2, y2, z2; 102 | x2 = dataset[k * 3 + 0]; 103 | y2 = dataset[k * 3 + 1]; 104 | z2 = dataset[k * 3 + 2]; 105 | float mag = (x2 * x2) + (y2 * y2) + (z2 * z2); 106 | if (mag <= 1e-3) continue; 107 | 108 | float d = 109 | (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1); 110 | 111 | float d2 = min(d, temp[k]); 112 | temp[k] = d2; 113 | besti = d2 > best ? k : besti; 114 | best = d2 > best ? d2 : best; 115 | } 116 | dists[tid] = best; 117 | dists_i[tid] = besti; 118 | __syncthreads(); 119 | 120 | if (block_size >= 512) { 121 | if (tid < 256) { 122 | __update(dists, dists_i, tid, tid + 256); 123 | } 124 | __syncthreads(); 125 | } 126 | if (block_size >= 256) { 127 | if (tid < 128) { 128 | __update(dists, dists_i, tid, tid + 128); 129 | } 130 | __syncthreads(); 131 | } 132 | if (block_size >= 128) { 133 | if (tid < 64) { 134 | __update(dists, dists_i, tid, tid + 64); 135 | } 136 | __syncthreads(); 137 | } 138 | if (block_size >= 64) { 139 | if (tid < 32) { 140 | __update(dists, dists_i, tid, tid + 32); 141 | } 142 | __syncthreads(); 143 | } 144 | if (block_size >= 32) { 145 | if (tid < 16) { 146 | __update(dists, dists_i, tid, tid + 16); 147 | } 148 | __syncthreads(); 149 | } 150 | if (block_size >= 16) { 151 | if (tid < 8) { 152 | __update(dists, dists_i, tid, tid + 8); 153 | } 154 | __syncthreads(); 155 | } 156 | if (block_size >= 8) { 157 | if (tid < 4) { 158 | __update(dists, dists_i, tid, tid + 4); 159 | } 160 | __syncthreads(); 161 | } 162 | if (block_size >= 4) { 163 | if (tid < 2) { 164 | __update(dists, dists_i, tid, tid + 2); 165 | } 166 | __syncthreads(); 167 | } 168 | if (block_size >= 2) { 169 | if (tid < 1) { 170 | __update(dists, dists_i, tid, tid + 1); 171 | } 172 | __syncthreads(); 173 | } 174 | 175 | old = dists_i[0]; 176 | if (tid == 0) idxs[j] = old; 177 | } 178 | } 179 | 180 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 181 | const float *dataset, float *temp, 182 | int *idxs) { 183 | unsigned int n_threads = opt_n_threads(n); 184 | 185 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 186 | 187 | switch (n_threads) { 188 | case 512: 189 | furthest_point_sampling_kernel<512> 190 | <<>>(b, n, m, dataset, temp, idxs); 191 | break; 192 | case 256: 193 | furthest_point_sampling_kernel<256> 194 | <<>>(b, n, m, dataset, temp, idxs); 195 | break; 196 | case 128: 197 | furthest_point_sampling_kernel<128> 198 | <<>>(b, n, m, dataset, temp, idxs); 199 | break; 200 | case 64: 201 | furthest_point_sampling_kernel<64> 202 | <<>>(b, n, m, dataset, temp, idxs); 203 | break; 204 | case 32: 205 | furthest_point_sampling_kernel<32> 206 | <<>>(b, n, m, dataset, temp, idxs); 207 | break; 208 | case 16: 209 | furthest_point_sampling_kernel<16> 210 | <<>>(b, n, m, dataset, temp, idxs); 211 | break; 212 | case 8: 213 | furthest_point_sampling_kernel<8> 214 | <<>>(b, n, m, dataset, temp, idxs); 215 | break; 216 | case 4: 217 | furthest_point_sampling_kernel<4> 218 | <<>>(b, n, m, dataset, temp, idxs); 219 | break; 220 | case 2: 221 | furthest_point_sampling_kernel<2> 222 | <<>>(b, n, m, dataset, temp, idxs); 223 | break; 224 | case 1: 225 | furthest_point_sampling_kernel<1> 226 | <<>>(b, n, m, dataset, temp, idxs); 227 | break; 228 | default: 229 | furthest_point_sampling_kernel<512> 230 | <<>>(b, n, m, dataset, temp, idxs); 231 | } 232 | 233 | CUDA_CHECK_ERRORS(); 234 | } 235 | -------------------------------------------------------------------------------- /src/pointnet2/build/lib.linux-x86_64-3.6/pointnet2/_ext.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/build/lib.linux-x86_64-3.6/pointnet2/_ext.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/ball_query.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/ball_query.o -------------------------------------------------------------------------------- /src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/ball_query_gpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/ball_query_gpu.o -------------------------------------------------------------------------------- /src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/bindings.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/bindings.o -------------------------------------------------------------------------------- /src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/group_points.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/group_points.o -------------------------------------------------------------------------------- /src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/group_points_gpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/group_points_gpu.o -------------------------------------------------------------------------------- /src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/interpolate.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/interpolate.o -------------------------------------------------------------------------------- /src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/interpolate_gpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/interpolate_gpu.o -------------------------------------------------------------------------------- /src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/sampling.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/sampling.o -------------------------------------------------------------------------------- /src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/sampling_gpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/build/temp.linux-x86_64-3.6/_ext_src/src/sampling_gpu.o -------------------------------------------------------------------------------- /src/pointnet2/dist/pointnet2-0.0.0-py3.6-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/pointnet2/dist/pointnet2-0.0.0-py3.6-linux-x86_64.egg -------------------------------------------------------------------------------- /src/pointnet2/pointnet2.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 2.1 2 | Name: pointnet2 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | License: UNKNOWN 7 | Platform: UNKNOWN 8 | 9 | UNKNOWN 10 | 11 | -------------------------------------------------------------------------------- /src/pointnet2/pointnet2.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | setup.py 2 | _ext_src/src/ball_query.cpp 3 | _ext_src/src/ball_query_gpu.cu 4 | _ext_src/src/bindings.cpp 5 | _ext_src/src/group_points.cpp 6 | _ext_src/src/group_points_gpu.cu 7 | _ext_src/src/interpolate.cpp 8 | _ext_src/src/interpolate_gpu.cu 9 | _ext_src/src/sampling.cpp 10 | _ext_src/src/sampling_gpu.cu 11 | pointnet2.egg-info/PKG-INFO 12 | pointnet2.egg-info/SOURCES.txt 13 | pointnet2.egg-info/dependency_links.txt 14 | pointnet2.egg-info/top_level.txt -------------------------------------------------------------------------------- /src/pointnet2/pointnet2.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /src/pointnet2/pointnet2.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | pointnet2 2 | -------------------------------------------------------------------------------- /src/pointnet2/setup.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | from setuptools import setup 7 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 8 | import glob 9 | 10 | _ext_src_root = "_ext_src" 11 | _ext_sources = glob.glob("{}/src/*.cpp".format(_ext_src_root)) + glob.glob( 12 | "{}/src/*.cu".format(_ext_src_root) 13 | ) 14 | _ext_headers = glob.glob("{}/include/*".format(_ext_src_root)) 15 | 16 | setup( 17 | name='pointnet2', 18 | ext_modules=[ 19 | CUDAExtension( 20 | name='pointnet2._ext', 21 | sources=_ext_sources, 22 | extra_compile_args={ 23 | "cxx": ["-O2", "-I{}".format("{}/include".format(_ext_src_root))], 24 | "nvcc": ["-O2", "-I{}".format("{}/include".format(_ext_src_root))], 25 | }, 26 | ) 27 | ], 28 | cmdclass={ 29 | 'build_ext': BuildExtension 30 | } 31 | ) 32 | -------------------------------------------------------------------------------- /src/util.py: -------------------------------------------------------------------------------- 1 | import readline 2 | import rlcompleter 3 | readline.parse_and_bind("tab: complete") 4 | import code 5 | import pdb 6 | 7 | import torch.multiprocessing as mp 8 | 9 | # debugging tools 10 | def interact(local=None): 11 | """interactive console with autocomplete function. Useful for debugging. 12 | interact(locals()) 13 | """ 14 | if local is None: 15 | local=dict(globals(), **locals()) 16 | 17 | readline.set_completer(rlcompleter.Completer(local).complete) 18 | code.interact(local=local) 19 | 20 | def set_trace(local=None): 21 | """debugging with pdb 22 | """ 23 | if local is None: 24 | local=dict(globals(), **locals()) 25 | 26 | pdb.Pdb.complete = rlcompleter.Completer(local).complete 27 | pdb.set_trace() -------------------------------------------------------------------------------- /src/utils/Logger.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | 5 | class Logger(object): 6 | def __init__(self, fpath=None): 7 | self.console = sys.stdout 8 | self.file = None 9 | if fpath is not None: 10 | self.file = open(fpath, 'w') 11 | 12 | def __del__(self): 13 | self.close() 14 | 15 | def __enter__(self): 16 | pass 17 | 18 | def __exit__(self, *args): 19 | self.close() 20 | 21 | def write(self, msg): 22 | self.console.write(msg) 23 | if self.file is not None: 24 | self.file.write(msg) 25 | 26 | def flush(self): 27 | self.console.flush() 28 | if self.file is not None: 29 | self.file.flush() 30 | os.fsync(self.file.fileno()) 31 | 32 | def close(self): 33 | self.console.close() 34 | if self.file is not None: 35 | self.file.close() 36 | 37 | 38 | 39 | -------------------------------------------------------------------------------- /src/utils/__pycache__/Logger.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/Logger.cpython-36.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/Logger.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/Logger.cpython-38.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/find_Nearest_Neighbor.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/find_Nearest_Neighbor.cpython-36.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/find_Nearest_Neighbor.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/find_Nearest_Neighbor.cpython-37.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/find_Nearest_Neighbor.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/find_Nearest_Neighbor.cpython-38.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/generate_normals.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/generate_normals.cpython-36.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/generate_normals.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/generate_normals.cpython-38.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/io.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/io.cpython-38.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/output_xyz.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/output_xyz.cpython-36.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/output_xyz.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/output_xyz.cpython-37.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/output_xyz.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/output_xyz.cpython-38.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/pc_transform.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/pc_transform.cpython-38.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/pointnet2_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/pointnet2_utils.cpython-38.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/random_pc_downsample.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/random_pc_downsample.cpython-36.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/random_pc_downsample.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/random_pc_downsample.cpython-37.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/random_pc_downsample.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/random_pc_downsample.cpython-38.pyc -------------------------------------------------------------------------------- /src/utils/__pycache__/read_ply.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Sangminhong/ACL-SPC_PyTorch/85167b079a8e5b43066f7cccee4da376296b4eca/src/utils/__pycache__/read_ply.cpython-38.pyc -------------------------------------------------------------------------------- /src/utils/find_Nearest_Neighbor.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from numba import jit, int32 4 | import numpy as np 5 | from sklearn.neighbors import KDTree 6 | 7 | import sys 8 | import os 9 | sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__)))) 10 | 11 | import pdb 12 | import util 13 | import time 14 | 15 | 16 | class find_NN(nn.Module): 17 | 18 | def forward(self,pcloud, P1, K): 19 | 20 | """args 21 | pcloud: point cloud with shape (1, num_points, 3) 22 | P1 ; query point with shape (3) 23 | 24 | 25 | 26 | """ 27 | #device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 28 | pcloud = torch.reshape(pcloud, (-1,3)) 29 | P1 = torch.reshape(P1, (1,3)) 30 | dist = torch.norm(pcloud-P1, dim=1, p=None) 31 | knn = dist.topk(K, largest=False) 32 | distances, indices = knn.values, knn.indices 33 | #max_distance = torch.max(distances) 34 | 35 | 36 | #return pcloud[indices].detach().cpu(), torch.reshape(indices,(1,-1)).detach().cpu() 37 | return pcloud[indices], torch.reshape(indices,(1,-1)) 38 | class find_NN_cuda(nn.Module): 39 | 40 | def forward(self,pcloud, P1, K): 41 | 42 | """args 43 | pcloud: point cloud with shape (1, num_points, 3) 44 | P1 ; query point with shape (3) 45 | 46 | 47 | 48 | """ 49 | pcloud = torch.reshape(pcloud, (-1,3)) 50 | P1 = torch.reshape(P1, (1,3)) 51 | dist = torch.norm(pcloud-P1, dim=1, p=None) 52 | knn = dist.topk(K, largest=False) 53 | distances, indices = knn.values, knn.indices 54 | #max_distance = torch.max(distances) 55 | 56 | 57 | return pcloud[indices], torch.reshape(indices,(1,-1)) 58 | 59 | 60 | class find_NN_batch(nn.Module): 61 | 62 | def forward(self,pcloud, P1, K): 63 | 64 | """args 65 | pcloud: point cloud with shape (batch_size, 3, num_points) 66 | P1 ; query point with shape (batch_size,3) 67 | 68 | output a : nearest point pcloud (batch_size, 3, new_npoints) 69 | output b : nearest point indices (batch_size, new_npoints) 70 | 71 | """ 72 | pcloud = torch.transpose(pcloud,1,2) 73 | batch_size = pcloud.shape[0] 74 | 75 | 76 | a = [] 77 | b = [] 78 | 79 | 80 | for i in range(batch_size): 81 | #device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 82 | tmp_pcloud = torch.reshape(pcloud[i], (-1,3)) 83 | tmp_P1 = torch.reshape(P1[i], (1,3)) 84 | 85 | #util.interact(locals()) 86 | #pdb.set_trace() 87 | dist = torch.norm(tmp_pcloud-tmp_P1, dim=1, p=None) 88 | knn = dist.topk(K, largest=False) 89 | distances, indices = knn.values, knn.indices 90 | 91 | #tree = KDTree(tmp_pcloud.detach().cpu()) 92 | #distances, indices = tree.query(tmp_P1.detach().cpu(), K) 93 | #max_distance = torch.max(distances) 94 | a.append(tmp_pcloud[indices].unsqueeze(0)) 95 | b.append(indices.unsqueeze(0)) 96 | b = torch.cat(b,0) 97 | b = torch.reshape(b, (batch_size, -1)) 98 | 99 | return torch.transpose(torch.cat(a,0),1,2), b 100 | 101 | 102 | if __name__ == '__main__': 103 | 104 | s = time.time() 105 | #pcloud = torch.Tensor([[0,1,2], [4,5,6], [3,5,2], [100,1,2]]) 106 | pcloud = torch.rand(120,3000,3) 107 | P1 = torch.rand(120,3) 108 | K = 2000 109 | find_NN_batch1 = find_NN_batch() 110 | pc, indices = find_NN_batch1(pcloud, P1, K) 111 | t = time.time() 112 | print(t-s) 113 | 114 | 115 | 116 | 117 | 118 | 119 | 120 | 121 | 122 | 123 | 124 | -------------------------------------------------------------------------------- /src/utils/io.py: -------------------------------------------------------------------------------- 1 | from plyfile import PlyData, PlyElement 2 | import os 3 | import numpy as np 4 | from tqdm import tqdm 5 | import shutil 6 | 7 | 8 | def read_ply_xyz(filename): 9 | """ read XYZ point cloud from filename PLY file """ 10 | if not os.path.isfile(filename): 11 | print(filename) 12 | assert(os.path.isfile(filename)) 13 | with open(filename, 'rb') as f: 14 | plydata = PlyData.read(f) 15 | num_verts = plydata['vertex'].count 16 | vertices = np.zeros(shape=[num_verts, 3], dtype=np.float32) 17 | vertices[:,0] = plydata['vertex'].data['x'] 18 | vertices[:,1] = plydata['vertex'].data['y'] 19 | vertices[:,2] = plydata['vertex'].data['z'] 20 | return vertices 21 | 22 | def read_all_ply_under_dir(dir): 23 | ''' 24 | read all .ply under a directory 25 | return a list of arrays 26 | ''' 27 | all_filenames = os.listdir(dir) 28 | ply_filenames = [] 29 | for f in all_filenames: 30 | if f.endswith('.ply'): 31 | ply_filenames.append(os.path.join(dir, f)) 32 | ply_filenames.sort() 33 | 34 | point_clouds = [] 35 | stems = [] 36 | for ply_f in tqdm(ply_filenames): 37 | pc = read_ply_xyz(ply_f) 38 | point_clouds.append(pc) 39 | 40 | basename = os.path.basename(ply_f) 41 | stem, ext = os.path.splitext(basename) 42 | stems.append(stem) 43 | 44 | return point_clouds, stems 45 | 46 | def read_ply_from_file_list(file_list): 47 | ''' 48 | read all .ply from a list 49 | return a list of numpy array 50 | ''' 51 | point_clouds = [] 52 | for ply_f in tqdm(file_list): 53 | if not os.path.isfile(ply_f): 54 | print('Warning: skipping. ', ply_f) 55 | continue 56 | pc = read_ply_xyz(ply_f) 57 | point_clouds.append(pc) 58 | 59 | return point_clouds 60 | 61 | def export_ply(pc, filename): 62 | """ 63 | export .ply from a point cloud 64 | """ 65 | vertex = np.zeros(pc.shape[0], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')]) 66 | for i in range(pc.shape[0]): 67 | vertex[i] = (pc[i][0], pc[i][1], pc[i][2]) 68 | ply_out = PlyData([PlyElement.describe(vertex, 'vertex', comments=['vertices'])]) 69 | ply_out.write(filename) 70 | 71 | def copytree(src, dst, symlinks=False, ignore=None): 72 | # ref: https://stackoverflow.com/questions/1868714/how-do-i-copy-an-entire-directory-of-files-into-an-existing-directory-using-pyth 73 | for item in os.listdir(src): 74 | s = os.path.join(src, item) 75 | d = os.path.join(dst, item) 76 | if os.path.isdir(s): 77 | shutil.copytree(s, d, symlinks, ignore) 78 | else: 79 | shutil.copy2(s, d) 80 | 81 | def read_txt_xyz(pathname): 82 | """ 83 | read a point cloud from txt file 84 | """ 85 | try: 86 | pcd = np.loadtxt(pathname,delimiter=';').astype(np.float32) 87 | except: 88 | pcd = np.loadtxt(pathname,delimiter=',').astype(np.float32) 89 | return pcd 90 | 91 | def export_pcd_to_txt(pcd,output_dir,stem): 92 | """ 93 | export .txt from a point cloud 94 | """ 95 | if not isinstance(pcd,np.ndarray): 96 | pcd = pcd.detach().cpu().numpy() 97 | np.savetxt(os.path.join(output_dir,str(stem)+'.txt'), pcd, fmt = "%f;%f;%f") -------------------------------------------------------------------------------- /src/utils/output_xyz.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import open3d as o3d 4 | 5 | def output_xyz(pts, output_path): 6 | """ 7 | args 8 | pts: point cloud tensor to output [num_points, 3] 9 | output_path : 'directory/filename' 10 | 11 | 12 | 13 | """ 14 | 15 | pcd = o3d.geometry.PointCloud() 16 | pcd.points = o3d.utility.Vector3dVector(pts.detach().cpu().numpy()) 17 | o3d.io.write_point_cloud(output_path, pcd) 18 | 19 | 20 | 21 | 22 | if __name__ == '__main__': 23 | a = torch.rand(100,3) 24 | output_xyz(a, '/home/mchiash2/PC_Upsampling/a.xyz') 25 | -------------------------------------------------------------------------------- /src/utils/pc_transform.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | """ 5 | this file contains various functions for point cloud transformation, 6 | some of which are not used in the clean version of code, 7 | but feel free to use them if you have different forms of point clouds. 8 | """ 9 | 10 | 11 | def swap_axis(input_np, swap_mode='n210'): 12 | """ 13 | swap axis for point clouds with different canonical frame 14 | e.g., pcl2pcl and MPC 15 | """ 16 | if swap_mode == '021': 17 | output_np = np.stack([input_np[:,0], input_np[:,2], input_np[:,1]],axis=1) 18 | elif swap_mode == 'n210': 19 | output_np = np.stack([input_np[:,2]*(-1), input_np[:,1], input_np[:,0]],axis=1) 20 | elif swap_mode == '210': 21 | output_np = np.stack([input_np[:,2], input_np[:,1], input_np[:,0]],axis=1) 22 | else: 23 | raise NotImplementedError 24 | 25 | return output_np 26 | 27 | def scale_numpy(input_array, range=0.25,ax_wise=True): 28 | """ 29 | scale point cloud in the form of numpy array 30 | """ 31 | if ax_wise: 32 | max_abs = np.max(np.abs(input_array),axis=0) 33 | d0 = input_array[:,0] * (range/max_abs[0]) 34 | d1 = input_array[:,1] * (range/max_abs[1]) 35 | d2 = input_array[:,2] * (range/max_abs[2]) 36 | scaled_array = np.stack([d0,d1,d2], axis=1) 37 | else: 38 | """ 39 | scale all dimension by the same value, ie the max(abs) 40 | """ 41 | max_abs = np.max(np.abs(input_array)) 42 | scaled_array = input_array * (range/max_abs) 43 | return scaled_array 44 | 45 | def scale_numpy_ls(input_ls, range=0.25): 46 | """ 47 | calling a list of point clouds 48 | """ 49 | output_ls = [] 50 | for itm in input_ls: 51 | output = scale_numpy(itm, range=range) 52 | output_ls.append(output) 53 | return output_ls 54 | 55 | def shift_numpy(input_array, mode='center',additional_limit=None): 56 | """ 57 | shift 58 | """ 59 | if mode == 'center': 60 | ax_max = np.max(input_array,axis=0) 61 | ax_min = np.min(input_array,axis=0) 62 | ax_center = (ax_max+ax_min)/2 63 | shifted_np = input_array - ax_center 64 | elif mode == 'given_some_limit': 65 | print(additional_limit) 66 | if additional_limit[0] != 'yl': 67 | raise NotImplementedError 68 | ax_max = np.max(input_array,axis=0) 69 | ax_min = np.min(input_array,axis=0) 70 | ax_min[1] = additional_limit[1] # addtional step 71 | ax_center = (ax_max+ax_min)/2 72 | shifted_np = input_array - ax_center 73 | else: 74 | raise NotImplementedError # weighted center, pc_mean 75 | 76 | return shifted_np 77 | 78 | def shift_np_one_dim(input_array, dim=2): 79 | max_dim = input_array.max(axis=0) 80 | min_dim = input_array.min(axis=0) 81 | mean_dim = (max_dim[dim]+min_dim[dim])/2 82 | input_array[:,dim] -= mean_dim 83 | return input_array 84 | 85 | def downsample_numpy(input_np, points=1024,seed=0): 86 | if input_np.shape[0] <= points: 87 | return input_np 88 | else: 89 | np.random.seed(seed) 90 | indices = np.array(range(input_np.shape[0])) 91 | np.random.shuffle(indices) 92 | input_downsampled = input_np[indices[:points]] 93 | return input_downsampled 94 | 95 | def voxelize(image, n_bins=32, pcd_limit=0.5, threshold=0): 96 | """ 97 | voxelize a point cloud 98 | """ 99 | if isinstance(image, np.ndarray): 100 | image = torch.from_numpy(image).unsqueeze(0) 101 | pcd_new = image * n_bins + n_bins * 0.5 102 | pcd_new = pcd_new.type(torch.int64) 103 | ls_voxels = pcd_new.squeeze(0).tolist() # 2028 of sublists 104 | try: 105 | tuple_voxels = [tuple(itm) for itm in ls_voxels] 106 | except: 107 | import pdb; pdb.set_trace() 108 | mask_dict = {} 109 | for tuple_voxel in tuple_voxels: 110 | if tuple_voxel not in mask_dict: 111 | mask_dict[tuple_voxel] = 1 112 | else: 113 | mask_dict[tuple_voxel] += 1 114 | for voxel, cnt in mask_dict.items(): 115 | if cnt <= threshold: 116 | del mask_dict[voxel] 117 | return mask_dict 118 | 119 | def return_plot_range(pcd, plot_range): 120 | """ 121 | return a range of point cloud, 122 | to plot Fig.3 in the main paper 123 | """ 124 | pcd_new = [] 125 | x1, x2 = plot_range[0] 126 | y1, y2 = plot_range[1] 127 | z1, z2 = plot_range[2] 128 | for i in range(2048): 129 | xs = pcd[i,0] 130 | ys = pcd[i,2] 131 | zs = pcd[i,1] 132 | if x1 < xs < x2 and y1 < ys < y2 and z1 < zs < z2: 133 | pcd_new.append(pcd[i]) 134 | pcd_new = np.stack(pcd_new) 135 | return pcd_new 136 | 137 | def reverse_normalize(pc, pc_CRN): 138 | """ 139 | scale up by m and relocate 140 | """ 141 | m = np.max(np.sqrt(np.sum(pc_CRN**2, axis=1))) 142 | pc = pc * m 143 | centroid = np.mean(pc_CRN, axis=0) 144 | pc = pc + centroid 145 | 146 | return pc 147 | 148 | def remove_zeros(partial): 149 | """ 150 | remove zeros (masked) from a point cloud 151 | """ 152 | if isinstance(partial, np.ndarray): 153 | partial = torch.from_numpy(partial) 154 | norm = torch.norm(partial,dim=1) 155 | idx = torch.where(norm > 0) 156 | partial = partial[idx[0]] 157 | return partial.numpy() 158 | 159 | def retrieve_region(pcd, retrieve_range): 160 | """ 161 | retrieve a range 162 | input: np.array (N,3) 163 | """ 164 | x1, x2 = retrieve_range[0] 165 | y1, y2 = retrieve_range[1] 166 | z1, z2 = retrieve_range[2] 167 | points = [] 168 | for i in range(pcd.shape[0]): 169 | xs = pcd[i,0] 170 | ys = pcd[i,2] 171 | zs = pcd[i,1] 172 | if x1 < xs < x2 and y1 < ys < y2 and z1 < zs < z2: 173 | points.append(pcd[i]) 174 | new_pcd = np.stack(points) 175 | print('new_pcd shape',new_pcd.shape) 176 | return new_pcd -------------------------------------------------------------------------------- /src/utils/read_ply.py: -------------------------------------------------------------------------------- 1 | import os 2 | def read_ply_xyz(filename): 3 | """ read XYZ RGB normals point cloud from filename PLY file """ 4 | assert(os.path.isfile(filename)) 5 | with open(filename, 'rb') as f: 6 | plydata = PlyData.read(f) 7 | num_verts = plydata['vertex'].count 8 | vertices = np.zeros(shape=[num_verts, 9], dtype=np.float32) 9 | vertices[:,0] = plydata['vertex'].data['x'] 10 | vertices[:,1] = plydata['vertex'].data['y'] 11 | vertices[:,2] = plydata['vertex'].data['z'] 12 | vertices[:,3] = plydata['vertex'].data['confidence'] 13 | vertices[:,4] = plydata['vertex'].data['intensity'] 14 | 15 | 16 | return vertices 17 | 18 | --------------------------------------------------------------------------------