├── LICENSE ├── README.md ├── convert_to_pointcloud ├── README └── bunny │ ├── H5DataGenerator.py │ ├── generate_test_dataset.py │ ├── generate_train_dataset.py │ └── tf_ops │ ├── 3d_interpolation │ ├── __pycache__ │ │ └── tf_interpolate.cpython-36.pyc │ ├── interpolate.cpp │ ├── tf_interpolate.cpp │ ├── tf_interpolate.py │ ├── tf_interpolate.pyc │ ├── tf_interpolate_compile.sh │ ├── tf_interpolate_op_test.py │ ├── tf_interpolate_so.so │ └── visu_interpolation.py │ ├── grouping │ ├── __pycache__ │ │ └── tf_grouping.cpython-36.pyc │ ├── test │ │ ├── compile.sh │ │ ├── query_ball_point │ │ ├── query_ball_point.cpp │ │ ├── query_ball_point.cu │ │ ├── query_ball_point_block │ │ ├── query_ball_point_block.cu │ │ ├── query_ball_point_cuda │ │ ├── query_ball_point_grid │ │ ├── query_ball_point_grid.cu │ │ ├── selection_sort │ │ ├── selection_sort.cpp │ │ ├── selection_sort.cu │ │ ├── selection_sort_const.cu │ │ └── selection_sort_cuda │ ├── tf_grouping.cpp │ ├── tf_grouping.py │ ├── tf_grouping.pyc │ ├── tf_grouping_compile.sh │ ├── tf_grouping_g.cu │ ├── tf_grouping_g.cu.o │ ├── tf_grouping_op_test.py │ └── tf_grouping_so.so │ ├── pointSIFT_op │ ├── __pycache__ │ │ └── pointSIFT_op.cpython-36.pyc │ ├── main.cpp │ ├── pointSIFT.cu │ ├── pointSIFT_g.cu.o │ ├── pointSIFT_op.py │ ├── pointSIFT_op.pyc │ ├── tf_pointSIFT_compile.sh │ └── tf_pointSIFT_so.so │ └── sampling │ ├── tf_sampling.cpp │ ├── tf_sampling.py │ ├── tf_sampling.pyc │ ├── tf_sampling_compile.sh │ ├── tf_sampling_g.cu │ ├── tf_sampling_g.cu.o │ └── tf_sampling_so.so ├── h5_dataset └── bunny │ ├── README │ ├── bunny_pprplus.tar │ └── cycle_0000.zip ├── logs └── README ├── models ├── SileaneBunny.obj └── calc_lambda.py ├── pprnet ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-35.pyc │ ├── __init__.cpython-36.pyc │ ├── __init__.cpython-37.pyc │ ├── __init__.cpython-38.pyc │ ├── object_type.cpython-36.pyc │ ├── pose_loss.cpython-36.pyc │ ├── pprnet.cpython-35.pyc │ ├── pprnet.cpython-36.pyc │ ├── pprnet.cpython-37.pyc │ ├── pprnet_test_1.cpython-36.pyc │ └── pprnet_test_1.cpython-38.pyc ├── backbone │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-35.pyc │ │ ├── __init__.cpython-36.pyc │ │ └── __init__.cpython-37.pyc │ ├── pointnet2 │ │ ├── __pycache__ │ │ │ ├── pointnet2_backbone.cpython-35.pyc │ │ │ ├── pointnet2_backbone.cpython-36.pyc │ │ │ ├── pointnet2_backbone.cpython-37.pyc │ │ │ ├── pointnet2_modules.cpython-35.pyc │ │ │ ├── pointnet2_modules.cpython-36.pyc │ │ │ ├── pointnet2_modules.cpython-37.pyc │ │ │ ├── pointnet2_utils.cpython-35.pyc │ │ │ ├── pointnet2_utils.cpython-36.pyc │ │ │ ├── pointnet2_utils.cpython-37.pyc │ │ │ ├── pytorch_utils.cpython-35.pyc │ │ │ ├── pytorch_utils.cpython-36.pyc │ │ │ ├── pytorch_utils.cpython-37.pyc │ │ │ └── test_backbone.cpython-36.pyc │ │ ├── _ext.cpython-36m-x86_64-linux-gnu.so │ │ ├── _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 │ │ ├── pointnet2_backbone.py │ │ ├── pointnet2_modules.py │ │ ├── pointnet2_utils.py │ │ ├── pytorch_utils.py │ │ └── setup.py │ └── rscnn │ │ ├── __pycache__ │ │ └── rscnn_backbone.cpython-36.pyc │ │ ├── rscnn_backbone.py │ │ └── utils │ │ ├── __init__.py │ │ ├── __pycache__ │ │ ├── linalg_utils.cpython-36.pyc │ │ ├── rscnn_pointnet2_modules.cpython-36.pyc │ │ └── rscnn_pointnet2_utils.cpython-36.pyc │ │ ├── _ext │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ └── __init__.cpython-36.pyc │ │ └── pointnet2 │ │ │ ├── __init__.py │ │ │ ├── __pycache__ │ │ │ └── __init__.cpython-36.pyc │ │ │ └── _pointnet2.so │ │ ├── build_ffi.py │ │ ├── cinclude │ │ ├── ball_query_gpu.h │ │ ├── ball_query_wrapper.h │ │ ├── cuda_utils.h │ │ ├── group_points_gpu.h │ │ ├── group_points_wrapper.h │ │ ├── interpolate_gpu.h │ │ ├── interpolate_wrapper.h │ │ ├── sampling_gpu.h │ │ └── sampling_wrapper.h │ │ ├── csrc │ │ ├── ball_query.c │ │ ├── ball_query_gpu.cu │ │ ├── group_points.c │ │ ├── group_points_gpu.cu │ │ ├── interpolate.c │ │ ├── interpolate_gpu.cu │ │ ├── sampling.c │ │ └── sampling_gpu.cu │ │ ├── linalg_utils.py │ │ ├── rscnn_pointnet2_modules.py │ │ ├── rscnn_pointnet2_utils.py │ │ └── rscnn_pytorch_utils │ │ ├── __init__.py │ │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ └── pytorch_utils.cpython-36.pyc │ │ └── pytorch_utils.py ├── data │ ├── IDL_pose_dataset.py │ ├── IPA_pose_dataset.py │ ├── __pycache__ │ │ ├── IDL_pose_dataset.cpython-36.pyc │ │ ├── IPA_pose_dataset.cpython-36.pyc │ │ ├── IPE_pose_dataset.cpython-36.pyc │ │ └── pointcloud_transforms.cpython-36.pyc │ └── pointcloud_transforms.py ├── object_type.py ├── pose_loss.py ├── pprnet.py ├── pprnet_plus.py └── utils │ ├── __pycache__ │ ├── dataset_util.cpython-36.pyc │ ├── dataset_util_ipe.cpython-36.pyc │ ├── eulerangles.cpython-36.pyc │ ├── eval_util.cpython-36.pyc │ ├── show3d_balls.cpython-36.pyc │ ├── train_helper.cpython-36.pyc │ └── visualize_util.cpython-36.pyc │ ├── dataset_util.py │ ├── dataset_util_ipe.py │ ├── eulerangles.py │ ├── eval_util.py │ ├── point_visualize_utils.py │ ├── render_balls_so.so │ ├── show3d_balls.py │ ├── train_helper.py │ └── visualize_util.py └── tools ├── IPABunny_conf_msg ├── evaluate.py ├── train.py └── train_only_conf.py └── IPABunny_msg ├── evaluate.py └── train.py /README.md: -------------------------------------------------------------------------------- 1 | ## PPR-Net++: Accurate 6D Pose Estimation in Stacked Scenarios 2 | This is the code of pytorch version for our IROS2019 paper and TASE2021 journal paper: [**PPR-Net: point-wise pose regression network for instance segmentation and 6d pose estimation in bin-picking scenarios**](https://ieeexplore.ieee.org/abstract/document/8967895); [**PPR-Net++: Accurate 6D Pose Estimation in Stacked Scenarios**](https://ieeexplore.ieee.org/abstract/document/9537584). 3 | 4 | 5 | ## Environment 6 | Ubuntu 16.04/18.04 7 | 8 | python3.6, torch 1.1.0, torchvision 0.3.0, opencv-python, sklearn, h5py, nibabel, et al. 9 | 10 | Our backbone PointNet++ is borrowed from [pointnet2](https://github.com/erikwijmans/Pointnet2_PyTorch). 11 | 12 | ## Dataset 13 | Siléane dataset is available at [here](http://rbregier.github.io/dataset2017). 14 | 15 | Fraunhofer IPA Bin-Picking dataset is available at [here](https://owncloud.fraunhofer.de/index.php/s/AacICuOWQVWDDfP?path=%2F). 16 | 17 | ## Evaluation metric 18 | The python code of evaluation metric is available at [here](https://github.com/rbregier/pose_recovery_evaluation). 19 | 20 | ## Citation 21 | If you use this codebase in your research, please cite: 22 | ``` 23 | @inproceedings{pprnet19IROS, 24 | title={PPR-Net: point-wise pose regression network for instance segmentation and 6d pose estimation in bin-picking scenarios}, 25 | author={Dong, Zhikai and Liu, Sicheng and Zhou, Tao and Cheng, Hui and Zeng, Long and Yu, Xingyao and Liu, Houde}, 26 | booktitle={2019 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)}, 27 | pages={1773--1780}, 28 | year={2019}, 29 | organization={IEEE} 30 | } 31 | 32 | @article{zeng2021ppr, 33 | title={PPR-Net++: accurate 6-D pose estimation in stacked scenarios}, 34 | author={Zeng, Long and Lv, Wei Jie and Dong, Zhi Kai and Liu, Yong Jin}, 35 | journal={IEEE Transactions on Automation Science and Engineering}, 36 | volume={19}, 37 | number={4}, 38 | pages={3139--3151}, 39 | year={2021}, 40 | publisher={IEEE} 41 | } 42 | ``` 43 | -------------------------------------------------------------------------------- /convert_to_pointcloud/README: -------------------------------------------------------------------------------- 1 | This is the code to convert the point cloud to h5 format 2 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/generate_test_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | from H5DataGenerator import * 3 | 4 | # output dirs 5 | OUT_ROOT_DIR = '../../h5_dataset/bunny' 6 | if not os.path.exists( OUT_ROOT_DIR ): 7 | os.mkdir(OUT_ROOT_DIR) 8 | TRAIN_SET_DIR = os.path.join(OUT_ROOT_DIR, 'test') 9 | if not os.path.exists( TRAIN_SET_DIR ): 10 | os.mkdir(TRAIN_SET_DIR) 11 | 12 | # input dirs 13 | IN_ROOT_DIR = '../../data_simulation/training/SileaneBunny_part_2' 14 | DEPTH_DIR = os.path.join(IN_ROOT_DIR, 'p_depth') 15 | 16 | 17 | if __name__ == "__main__": 18 | cycle_idx_list = range(480, 500) 19 | scene_idx_list = range(1, 81) 20 | g = H5DataGenerator(os.path.join(IN_ROOT_DIR, 'parameter.json')) 21 | for cycle_id in cycle_idx_list: 22 | # load background image 23 | bg_depth_image_path = os.path.join(DEPTH_DIR, 'cycle_{:0>4}'.format(cycle_id), '000_depth_uint16.png') 24 | bg_depth_image = cv2.imread(bg_depth_image_path,-1) 25 | 26 | out_cycle_dir = os.path.join(TRAIN_SET_DIR, 'cycle_{:0>4}'.format(cycle_id)) 27 | if not os.path.exists(out_cycle_dir): 28 | os.mkdir(out_cycle_dir) 29 | for scene_id in scene_idx_list: 30 | # load inputs 31 | depth_image_path = os.path.join(DEPTH_DIR, 'cycle_{:0>4}'.format(cycle_id), '{:0>3}_depth_uint16.png'.format(scene_id)) 32 | output_h5_path = os.path.join(out_cycle_dir, '{:0>3}.h5'.format(scene_id)) 33 | depth_image = cv2.imread(depth_image_path,-1) 34 | 35 | # g.process_train_set(depth_image, bg_depth_image, segment_image, gt_file_path, output_h5_path) 36 | g.process_test_set(depth_image, bg_depth_image, output_h5_path) 37 | 38 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/generate_train_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | from H5DataGenerator import * 3 | 4 | # output dirs 5 | OUT_ROOT_DIR = '../../h5_dataset/bunny' 6 | if not os.path.exists( OUT_ROOT_DIR ): 7 | os.mkdir(OUT_ROOT_DIR) 8 | TRAIN_SET_DIR = os.path.join(OUT_ROOT_DIR, 'train') 9 | if not os.path.exists( TRAIN_SET_DIR ): 10 | os.mkdir(TRAIN_SET_DIR) 11 | 12 | # input dirs 13 | IN_ROOT_DIR = '../../data_simulation/training/SileaneBunny_part_2' 14 | GT_DIR = os.path.join(IN_ROOT_DIR, 'gt') 15 | SEGMENT_DIR = os.path.join(IN_ROOT_DIR, 'p_segmentation') 16 | DEPTH_DIR = os.path.join(IN_ROOT_DIR, 'p_depth') 17 | 18 | 19 | if __name__ == "__main__": 20 | cycle_idx_list = range(250, 500) 21 | scene_idx_list = range(1, 81) 22 | g = H5DataGenerator(os.path.join(IN_ROOT_DIR, 'parameter.json')) 23 | for cycle_id in cycle_idx_list: 24 | # load background image 25 | bg_depth_image_path = os.path.join(DEPTH_DIR, 'cycle_{:0>4}'.format(cycle_id), '000_depth_uint16.png') 26 | bg_depth_image = cv2.imread(bg_depth_image_path,-1) 27 | 28 | out_cycle_dir = os.path.join(TRAIN_SET_DIR, 'cycle_{:0>4}'.format(cycle_id)) 29 | if not os.path.exists(out_cycle_dir): 30 | os.mkdir(out_cycle_dir) 31 | for scene_id in scene_idx_list: 32 | # load inputs 33 | depth_image_path = os.path.join(DEPTH_DIR, 'cycle_{:0>4}'.format(cycle_id), '{:0>3}_depth_uint16.png'.format(scene_id)) 34 | depth_image = cv2.imread(depth_image_path,-1) 35 | seg_img_path = os.path.join(SEGMENT_DIR, 'cycle_{:0>4}'.format(cycle_id), '{:0>3}_segmentation.png'.format(scene_id)) 36 | segment_image = cv2.imread(seg_img_path,-1) 37 | gt_file_path = os.path.join(GT_DIR, 'cycle_{:0>4}'.format(cycle_id), '{:0>3}.csv'.format(scene_id)) 38 | output_h5_path = os.path.join(out_cycle_dir, '{:0>3}.h5'.format(scene_id)) 39 | 40 | g.process_train_set(depth_image, bg_depth_image, segment_image, gt_file_path, output_h5_path) 41 | 42 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/3d_interpolation/__pycache__/tf_interpolate.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/3d_interpolation/__pycache__/tf_interpolate.cpython-36.pyc -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/3d_interpolation/interpolate.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // Find three nearest neigbors with square distance 19 | // input: xyz1 (b,n,3), xyz2(b,m,3) 20 | // output: dist (b,n,3), idx (b,n,3) 21 | void threenn_cpu(int b, int n, int m, const float *xyz1, const float *xyz2, float *dist, int *idx) { 22 | for (int i=0;i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | void query_ball_point_cpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | for (int i=0;i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | for (int i=0;i>>(b,n,m,radius,nsample,xyz1,xyz2,idx); 113 | cudaDeviceSynchronize(); 114 | printf("query_ball_point gpu time %f\n",get_time()-t0); 115 | 116 | t0=get_time(); 117 | group_point_gpu<<<1,1>>>(b,n,c,m,nsample,points,idx,out); 118 | cudaDeviceSynchronize(); 119 | printf("grou_point gpu time %f\n",get_time()-t0); 120 | 121 | t0=get_time(); 122 | group_point_grad_gpu<<<1,1>>>(b,n,c,m,nsample,grad_out,idx,grad_points); 123 | cudaDeviceSynchronize(); 124 | printf("grou_point_grad gpu time %f\n",get_time()-t0); 125 | 126 | cudaFree(xyz1); 127 | cudaFree(xyz2); 128 | cudaFree(points); 129 | cudaFree(idx); 130 | cudaFree(out); 131 | cudaFree(grad_out); 132 | cudaFree(grad_points); 133 | return 0; 134 | } 135 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/grouping/test/query_ball_point_block: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/grouping/test/query_ball_point_block -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/grouping/test/query_ball_point_block.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | int index = threadIdx.x; 21 | xyz1 += n*3*index; 22 | xyz2 += m*3*index; 23 | idx += m*nsample*index; 24 | 25 | for (int j=0;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx); 113 | cudaDeviceSynchronize(); 114 | printf("query_ball_point gpu time %f\n",get_time()-t0); 115 | 116 | t0=get_time(); 117 | group_point_gpu<<<1,b>>>(b,n,c,m,nsample,points,idx,out); 118 | cudaDeviceSynchronize(); 119 | printf("grou_point gpu time %f\n",get_time()-t0); 120 | 121 | t0=get_time(); 122 | group_point_grad_gpu<<<1,b>>>(b,n,c,m,nsample,grad_out,idx,grad_points); 123 | cudaDeviceSynchronize(); 124 | printf("grou_point_grad gpu time %f\n",get_time()-t0); 125 | 126 | cudaFree(xyz1); 127 | cudaFree(xyz2); 128 | cudaFree(points); 129 | cudaFree(idx); 130 | cudaFree(out); 131 | cudaFree(grad_out); 132 | cudaFree(grad_points); 133 | return 0; 134 | } 135 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/grouping/test/query_ball_point_cuda: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/grouping/test/query_ball_point_cuda -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/grouping/test/query_ball_point_grid: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/grouping/test/query_ball_point_grid -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/grouping/test/query_ball_point_grid.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | int batch_index = blockIdx.x; 21 | xyz1 += n*3*batch_index; 22 | xyz2 += m*3*batch_index; 23 | idx += m*nsample*batch_index; 24 | 25 | int index = threadIdx.x; 26 | int stride = blockDim.x; 27 | 28 | for (int j=index;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx); 123 | cudaDeviceSynchronize(); 124 | printf("query_ball_point gpu time %f\n",get_time()-t0); 125 | 126 | t0=get_time(); 127 | group_point_gpu<<>>(b,n,c,m,nsample,points,idx,out); 128 | cudaDeviceSynchronize(); 129 | printf("grou_point gpu time %f\n",get_time()-t0); 130 | 131 | t0=get_time(); 132 | group_point_grad_gpu<<>>(b,n,c,m,nsample,grad_out,idx,grad_points); 133 | cudaDeviceSynchronize(); 134 | printf("grou_point_grad gpu time %f\n",get_time()-t0); 135 | 136 | cudaFree(xyz1); 137 | cudaFree(xyz2); 138 | cudaFree(points); 139 | cudaFree(idx); 140 | cudaFree(out); 141 | cudaFree(grad_out); 142 | cudaFree(grad_points); 143 | return 0; 144 | } 145 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/grouping/test/selection_sort: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/grouping/test/selection_sort -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/grouping/test/selection_sort.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // input: k (1), distance matrix dist (b,m,n) 19 | // output: idx (b,m,n), val (b,m,n) 20 | void selection_sort_cpu(int b, int n, int m, int k, const float *dist, int *idx, float *val) { 21 | float *p_dist; 22 | float tmp; 23 | int tmpi; 24 | for (int i=0;i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // input: k (1), distance matrix dist (b,m,n) 19 | // output: idx (b,m,k), val (b,m,k) 20 | __global__ void selection_sort_gpu(int b, int n, int m, int k, float *dist, int *idx, float *val) { 21 | int batch_index = blockIdx.x; 22 | dist+=m*n*batch_index; 23 | idx+=m*k*batch_index; 24 | val+=m*k*batch_index; 25 | 26 | int index = threadIdx.x; 27 | int stride = blockDim.x; 28 | 29 | float *p_dist; 30 | for (int j=index;j>>(b,n,m,k,dist,idx,val); 68 | cudaDeviceSynchronize(); 69 | printf("selection sort cpu time %f\n",get_time()-t0); 70 | 71 | return 0; 72 | } 73 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/grouping/test/selection_sort_const.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // input: k (1), distance matrix dist (b,m,n) 19 | // output: idx (b,m,n), dist_out (b,m,n) 20 | __global__ void selection_sort_gpu(int b, int n, int m, int k, const float *dist, int *outi, float *out) { 21 | int batch_index = blockIdx.x; 22 | dist+=m*n*batch_index; 23 | outi+=m*n*batch_index; 24 | out+=m*n*batch_index; 25 | 26 | int index = threadIdx.x; 27 | int stride = blockDim.x; 28 | 29 | // copy from dist to dist_out 30 | for (int j=index;j>>(b,n,m,k,dist,idx,dist_out); 84 | cudaDeviceSynchronize(); 85 | printf("selection sort cpu time %f\n",get_time()-t0); 86 | 87 | //for (int i=0;i>>(b,n,m,radius,nsample,xyz1,xyz2,idx,pts_cnt); 127 | //cudaDeviceSynchronize(); 128 | } 129 | void selectionSortLauncher(int b, int n, int m, int k, const float *dist, int *outi, float *out) { 130 | selection_sort_gpu<<>>(b,n,m,k,dist,outi,out); 131 | //cudaDeviceSynchronize(); 132 | } 133 | void groupPointLauncher(int b, int n, int c, int m, int nsample, const float *points, const int *idx, float *out){ 134 | group_point_gpu<<>>(b,n,c,m,nsample,points,idx,out); 135 | //cudaDeviceSynchronize(); 136 | } 137 | void groupPointGradLauncher(int b, int n, int c, int m, int nsample, const float *grad_out, const int *idx, float *grad_points){ 138 | group_point_grad_gpu<<>>(b,n,c,m,nsample,grad_out,idx,grad_points); 139 | //group_point_grad_gpu<<<1,1>>>(b,n,c,m,nsample,grad_out,idx,grad_points); 140 | //cudaDeviceSynchronize(); 141 | } 142 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/grouping/tf_grouping_g.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/grouping/tf_grouping_g.cu.o -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/grouping/tf_grouping_op_test.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | from tf_grouping import query_ball_point, group_point 4 | 5 | class GroupPointTest(tf.test.TestCase): 6 | def test(self): 7 | pass 8 | 9 | def test_grad(self): 10 | with tf.device('/gpu:0'): 11 | points = tf.constant(np.random.random((1,128,16)).astype('float32')) 12 | print points 13 | xyz1 = tf.constant(np.random.random((1,128,3)).astype('float32')) 14 | xyz2 = tf.constant(np.random.random((1,8,3)).astype('float32')) 15 | radius = 0.3 16 | nsample = 32 17 | idx, pts_cnt = query_ball_point(radius, nsample, xyz1, xyz2) 18 | grouped_points = group_point(points, idx) 19 | print grouped_points 20 | 21 | with self.test_session(): 22 | print "---- Going to compute gradient error" 23 | err = tf.test.compute_gradient_error(points, (1,128,16), grouped_points, (1,8,32,16)) 24 | print err 25 | self.assertLess(err, 1e-4) 26 | 27 | if __name__=='__main__': 28 | tf.test.main() 29 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/grouping/tf_grouping_so.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/grouping/tf_grouping_so.so -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/pointSIFT_op/__pycache__/pointSIFT_op.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/pointSIFT_op/__pycache__/pointSIFT_op.cpython-36.pyc -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/pointSIFT_op/pointSIFT.cu: -------------------------------------------------------------------------------- 1 | __global__ void selection_k_radius_gpu(int b, int m, int k, float radius, const int* idx, const float* val, int* idx_out, float* val_out){ 2 | int batch_index = blockIdx.x; 3 | int stride = batch_index * m * k; 4 | idx += stride; 5 | val += stride; 6 | idx_out += stride; 7 | val_out += stride; 8 | for(int i = threadIdx.x; i < m;i += blockDim.x) { 9 | 10 | for(int j = 0;j < k;j ++) { 11 | if(val[i * k + j] < radius) { 12 | idx_out[i * k + j] = idx[i * k + j]; 13 | val_out[i * k + j] = val[i * k + j]; 14 | } else { 15 | idx_out[i * k + j] = idx[i * k ]; 16 | val_out[i * k + j] = val[i * k ]; 17 | } 18 | } 19 | } 20 | } 21 | 22 | __global__ void cube_select(int b, int n,float radius, const float* xyz, int* idx_out) { 23 | int batch_idx = blockIdx.x; 24 | xyz += batch_idx * n * 3; 25 | idx_out += batch_idx * n * 8; 26 | float temp_dist[8]; 27 | float judge_dist = radius * radius; 28 | for(int i = threadIdx.x; i < n;i += blockDim.x) { 29 | float x = xyz[i * 3]; 30 | float y = xyz[i * 3 + 1]; 31 | float z = xyz[i * 3 + 2]; 32 | for(int j = 0;j < 8;j ++) { 33 | temp_dist[j] = 1e8; 34 | idx_out[i * 8 + j] = i; // if not found, just return itself.. 35 | } 36 | for(int j = 0;j < n;j ++) { 37 | if(i == j) continue; 38 | float tx = xyz[j * 3]; 39 | float ty = xyz[j * 3 + 1]; 40 | float tz = xyz[j * 3 + 2]; 41 | float dist = (x - tx) * (x - tx) + (y - ty) * (y - ty) + (z - tz) * (z - tz); 42 | if(dist > judge_dist) continue; 43 | int _x = (tx > x); 44 | int _y = (ty > y); 45 | int _z = (tz > z); 46 | int temp_idx = _x * 4 + _y * 2 + _z; 47 | if(dist < temp_dist[temp_idx]) { 48 | idx_out[i * 8 + temp_idx] = j; 49 | temp_dist[temp_idx] = dist; 50 | } 51 | } 52 | } 53 | } 54 | 55 | __global__ void cube_select_two(int b, int n,float radius, const float* xyz, int* idx_out) { 56 | int batch_idx = blockIdx.x; 57 | xyz += batch_idx * n * 3; 58 | idx_out += batch_idx * n * 16; 59 | float temp_dist[16]; 60 | float judge_dist = radius * radius; 61 | for(int i = threadIdx.x; i < n;i += blockDim.x) { 62 | float x = xyz[i * 3]; 63 | float y = xyz[i * 3 + 1]; 64 | float z = xyz[i * 3 + 2]; 65 | for(int j = 0;j < 16;j ++) { 66 | temp_dist[j] = judge_dist; 67 | idx_out[i * 16 + j] = i; // if not found, just return itself.. 68 | } 69 | for(int j = 0;j < n;j ++) { 70 | if(i == j) continue; 71 | float tx = xyz[j * 3]; 72 | float ty = xyz[j * 3 + 1]; 73 | float tz = xyz[j * 3 + 2]; 74 | float dist = (x - tx) * (x - tx) + (y - ty) * (y - ty) + (z - tz) * (z - tz); 75 | if(dist > judge_dist) continue; 76 | int _x = (tx > x); 77 | int _y = (ty > y); 78 | int _z = (tz > z); 79 | int temp_idx = _x * 8 + _y * 4 + _z * 2; 80 | bool flag = false; 81 | for(int k = 0;k < 2;k ++) { 82 | if (dist < temp_dist[temp_idx + k]) { 83 | flag = true; 84 | } 85 | if (flag) { 86 | for (int kk = 1; kk >= k + 1; kk --) { 87 | idx_out[i * 16 + temp_idx + kk] = idx_out[i * 16 + temp_idx + kk - 1]; 88 | temp_dist[temp_idx + kk] = temp_dist[temp_idx + kk - 1]; 89 | } 90 | idx_out[i * 16 + temp_idx + k] = j; 91 | temp_dist[temp_idx + k] = dist; 92 | break; 93 | } 94 | } 95 | 96 | } 97 | } 98 | } 99 | 100 | __global__ void cube_select_four(int b, int n,float radius, const float* xyz, int* idx_out) { 101 | int batch_idx = blockIdx.x; 102 | xyz += batch_idx * n * 3; 103 | idx_out += batch_idx * n * 32; 104 | float temp_dist[32]; 105 | float judge_dist = radius * radius; 106 | for(int i = threadIdx.x; i < n;i += blockDim.x) { 107 | float x = xyz[i * 3]; 108 | float y = xyz[i * 3 + 1]; 109 | float z = xyz[i * 3 + 2]; 110 | for(int j = 0;j < 32;j ++) { 111 | temp_dist[j] = judge_dist; 112 | idx_out[i * 32 + j] = i; // if not found, just return itself.. 113 | } 114 | for(int j = 0;j < n;j ++) { 115 | if(i == j) continue; 116 | float tx = xyz[j * 3]; 117 | float ty = xyz[j * 3 + 1]; 118 | float tz = xyz[j * 3 + 2]; 119 | float dist = (x - tx) * (x - tx) + (y - ty) * (y - ty) + (z - tz) * (z - tz); 120 | if(dist > judge_dist) continue; 121 | int _x = (tx > x); 122 | int _y = (ty > y); 123 | int _z = (tz > z); 124 | int temp_idx = _x * 16 + _y * 8 + _z * 4; 125 | bool flag = false; 126 | for(int k = 0;k < 4;k ++) { 127 | if (dist < temp_dist[temp_idx + k]) { 128 | flag = true; 129 | } 130 | if (flag) { 131 | for (int kk = 3; kk >= k + 1; kk --) { 132 | idx_out[i * 32 + temp_idx + kk] = idx_out[i * 32 + temp_idx + kk - 1]; 133 | temp_dist[temp_idx + kk] = temp_dist[temp_idx + kk - 1]; 134 | } 135 | idx_out[i * 32 + temp_idx + k] = j; 136 | temp_dist[temp_idx + k] = dist; 137 | break; 138 | } 139 | } 140 | 141 | } 142 | } 143 | } 144 | 145 | 146 | 147 | void selectionKRadiusLauncher(int b, int m, int k, float radius, const int* idx, const float* val, int* idx_out, float* val_out){ 148 | selection_k_radius_gpu<<>>(b, m, k, radius, idx, val, idx_out, val_out); 149 | } 150 | void cubeSelectLauncher(int b, int n, float radius, const float* xyz, int* idx_out) { 151 | cube_select<<>>(b, n, radius, xyz, idx_out); 152 | } 153 | void cubeSelectTwoLauncher(int b, int n, float radius, const float* xyz, int* idx_out) { 154 | cube_select_two<<>>(b, n, radius, xyz, idx_out); 155 | } 156 | void cubeSelectFourLauncher(int b, int n, float radius, const float* xyz, int* idx_out) { 157 | cube_select_four<<>>(b, n, radius, xyz, idx_out); 158 | } 159 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/pointSIFT_op/pointSIFT_g.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/pointSIFT_op/pointSIFT_g.cu.o -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/pointSIFT_op/pointSIFT_op.py: -------------------------------------------------------------------------------- 1 | """ 2 | Author: Jiang Mingyang 3 | email: jmydurant@sjtu.edu.cn 4 | pointSIFT module op, do not modify it !!! 5 | """ 6 | 7 | import tensorflow as tf 8 | from tensorflow.python.framework import ops 9 | import sys 10 | import os 11 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 12 | sys.path.append(BASE_DIR) 13 | 14 | pointSIFT_module = tf.load_op_library(os.path.join(BASE_DIR, 'tf_pointSIFT_so.so')) 15 | 16 | def pointSIFT_select(xyz, radius): 17 | """ 18 | :param xyz: (b, n, 3) float 19 | :param radius: float 20 | :return: (b, n, 8) int 21 | """ 22 | idx = pointSIFT_module.cube_select(xyz, radius) 23 | return idx 24 | 25 | 26 | ops.NoGradient('CubeSelect') 27 | 28 | def pointSIFT_select_two(xyz, radius): 29 | """ 30 | :param xyz: (b, n, 3) float 31 | :param radius: float 32 | :return: idx: (b, n, 16) int 33 | """ 34 | idx = pointSIFT_module.cube_select_two(xyz, radius) 35 | return idx 36 | 37 | 38 | ops.NoGradient('CubeSelectTwo') 39 | 40 | def pointSIFT_select_four(xyz, radius): 41 | """ 42 | :param xyz: (b, n, 3) float 43 | :param radius: float 44 | :return: idx: (b, n, 32) int 45 | """ 46 | idx = pointSIFT_module.cube_select_four(xyz, radius) 47 | return idx 48 | 49 | 50 | ops.NoGradient('CubeSelectFour') 51 | 52 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/pointSIFT_op/pointSIFT_op.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/pointSIFT_op/pointSIFT_op.pyc -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/pointSIFT_op/tf_pointSIFT_compile.sh: -------------------------------------------------------------------------------- 1 | #/bin/bash 2 | /usr/local/cuda-9.0/bin/nvcc pointSIFT.cu -o pointSIFT_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 3 | # TF1.2 4 | #g++ -std=c++11 main.cpp pointSIFT_g.cu.o -o tf_pointSIFT_so.so -shared -fPIC -I /home/likewise-open/SENSETIME/dongzhikai/tf_env/tf_env/local/lib/python2.7/site-packages/tensorflow/include -I/usr/local/cuda/include -lcudart -L /usr/local/cuda/lib64/ -O2 -D_GLIBCXX_USE_CXX11_ABI=0 5 | 6 | # TF1.4 7 | g++ -std=c++11 main.cpp pointSIFT_g.cu.o -o tf_pointSIFT_so.so -shared -fPIC -I /home/idesignlab/anaconda3/envs/pointnet2/lib/python3.6/site-packages/tensorflow/include -I /usr/local/cuda-9.0/include -I /home/idesignlab/anaconda3/envs/pointnet2/lib/python3.6/site-packages/tensorflow/include/external/nsync/public -lcudart -L /usr/local/cuda-9.0/lib64/ -L/home/idesignlab/anaconda3/envs/pointnet2/lib/python3.6/site-packages/tensorflow -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0 8 | 9 | 10 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/pointSIFT_op/tf_pointSIFT_so.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/pointSIFT_op/tf_pointSIFT_so.so -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/sampling/tf_sampling.py: -------------------------------------------------------------------------------- 1 | ''' Furthest point sampling 2 | Original author: Haoqiang Fan 3 | Modified by Charles R. Qi 4 | All Rights Reserved. 2017. 5 | ''' 6 | import tensorflow as tf 7 | from tensorflow.python.framework import ops 8 | import sys 9 | import os 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | sys.path.append(BASE_DIR) 12 | sampling_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_sampling_so.so')) 13 | def prob_sample(inp,inpr): 14 | ''' 15 | input: 16 | batch_size * ncategory float32 17 | batch_size * npoints float32 18 | returns: 19 | batch_size * npoints int32 20 | ''' 21 | return sampling_module.prob_sample(inp,inpr) 22 | ops.NoGradient('ProbSample') 23 | # TF1.0 API requires set shape in C++ 24 | #@tf.RegisterShape('ProbSample') 25 | #def _prob_sample_shape(op): 26 | # shape1=op.inputs[0].get_shape().with_rank(2) 27 | # shape2=op.inputs[1].get_shape().with_rank(2) 28 | # return [tf.TensorShape([shape2.dims[0],shape2.dims[1]])] 29 | def gather_point(inp,idx): 30 | ''' 31 | input: 32 | batch_size * ndataset * 3 float32 33 | batch_size * npoints int32 34 | returns: 35 | batch_size * npoints * 3 float32 36 | ''' 37 | return sampling_module.gather_point(inp,idx) 38 | #@tf.RegisterShape('GatherPoint') 39 | #def _gather_point_shape(op): 40 | # shape1=op.inputs[0].get_shape().with_rank(3) 41 | # shape2=op.inputs[1].get_shape().with_rank(2) 42 | # return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[2]])] 43 | @tf.RegisterGradient('GatherPoint') 44 | def _gather_point_grad(op,out_g): 45 | inp=op.inputs[0] 46 | idx=op.inputs[1] 47 | return [sampling_module.gather_point_grad(inp,idx,out_g),None] 48 | def farthest_point_sample(npoint,inp): 49 | ''' 50 | input: 51 | int32 52 | batch_size * ndataset * 3 float32 53 | returns: 54 | batch_size * npoint int32 55 | ''' 56 | return sampling_module.farthest_point_sample(inp, npoint) 57 | ops.NoGradient('FarthestPointSample') 58 | 59 | 60 | if __name__=='__main__': 61 | import numpy as np 62 | np.random.seed(100) 63 | triangles=np.random.rand(1,5,3,3).astype('float32') 64 | with tf.device('/gpu:1'): 65 | inp=tf.constant(triangles) 66 | tria=inp[:,:,0,:] 67 | trib=inp[:,:,1,:] 68 | tric=inp[:,:,2,:] 69 | areas=tf.sqrt(tf.reduce_sum(tf.cross(trib-tria,tric-tria)**2,2)+1e-9) 70 | randomnumbers=tf.random_uniform((1,8192)) 71 | triids=prob_sample(areas,randomnumbers) 72 | tria_sample=gather_point(tria,triids) 73 | trib_sample=gather_point(trib,triids) 74 | tric_sample=gather_point(tric,triids) 75 | us=tf.random_uniform((1,8192)) 76 | vs=tf.random_uniform((1,8192)) 77 | uplusv=1-tf.abs(us+vs-1) 78 | uminusv=us-vs 79 | us=(uplusv+uminusv)*0.5 80 | vs=(uplusv-uminusv)*0.5 81 | pt_sample=tria_sample+(trib_sample-tria_sample)*tf.expand_dims(us,-1)+(tric_sample-tria_sample)*tf.expand_dims(vs,-1) 82 | print ('pt_sample: ', pt_sample) 83 | reduced_sample=gather_point(pt_sample,farthest_point_sample(1024,pt_sample)) 84 | print (reduced_sample) 85 | with tf.Session('') as sess: 86 | ret=sess.run(reduced_sample) 87 | print (ret.shape,ret.dtype) 88 | import cPickle as pickle 89 | pickle.dump(ret,open('1.pkl','wb'),-1) 90 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/sampling/tf_sampling.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/sampling/tf_sampling.pyc -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/sampling/tf_sampling_compile.sh: -------------------------------------------------------------------------------- 1 | #/bin/bash 2 | /usr/local/cuda-9.0/bin/nvcc tf_sampling_g.cu -o tf_sampling_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 3 | 4 | # TF1.2 5 | #g++ -std=c++11 tf_sampling.cpp tf_sampling_g.cu.o -o tf_sampling_so.so -shared -fPIC -I /home/idesignlab/anaconda3/envs/pointnet2/lib/python3.6/site-packages/tensorflow/include -I /usr/local/cuda-8.0/include -lcudart -L /usr/local/cuda-8.0/lib64/ -O2 -D_GLIBCXX_USE_CXX11_ABI=0 6 | 7 | 8 | 9 | # TF1.4 10 | g++ -std=c++11 tf_sampling.cpp tf_sampling_g.cu.o -o tf_sampling_so.so -shared -fPIC -I /home/idesignlab/anaconda3/envs/pointnet2/lib/python3.6/site-packages/tensorflow/include -I /usr/local/cuda-9.0/include -I /home/idesignlab/anaconda3/envs/pointnet2/lib/python3.6/site-packages/tensorflow/include/external/nsync/public -lcudart -L /usr/local/cuda-9.0/lib64/ -L/home/idesignlab/anaconda3/envs/pointnet2/lib/python3.6/site-packages/tensorflow -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0 11 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/sampling/tf_sampling_g.cu: -------------------------------------------------------------------------------- 1 | /* Furthest point sampling GPU implementation 2 | * Original author: Haoqiang Fan 3 | * Modified by Charles R. Qi 4 | * All Rights Reserved. 2017. 5 | */ 6 | 7 | __global__ void cumsumKernel(int b,int n,const float * __restrict__ inp,float * __restrict__ out){ 8 | const int BlockSize=2048; 9 | const int paddingLevel=5; 10 | __shared__ float buffer4[BlockSize*4]; 11 | __shared__ float buffer[BlockSize+(BlockSize>>paddingLevel)]; 12 | for (int i=blockIdx.x;i>2; 18 | for (int k=threadIdx.x*4;k>2)+(k>>(2+paddingLevel))]=v4; 33 | }else{ 34 | float v=0; 35 | for (int k2=k;k2>2)+(k>>(2+paddingLevel))]=v; 43 | } 44 | } 45 | int u=0; 46 | for (;(2<>(u+1));k+=blockDim.x){ 49 | int i1=(((k<<1)+2)<>paddingLevel; 52 | i2+=i2>>paddingLevel; 53 | buffer[i1]+=buffer[i2]; 54 | } 55 | } 56 | u--; 57 | for (;u>=0;u--){ 58 | __syncthreads(); 59 | for (int k=threadIdx.x;k>(u+1));k+=blockDim.x){ 60 | int i1=(((k<<1)+3)<>paddingLevel; 63 | i2+=i2>>paddingLevel; 64 | buffer[i1]+=buffer[i2]; 65 | } 66 | } 67 | __syncthreads(); 68 | for (int k=threadIdx.x*4;k>2)-1)+(((k>>2)-1)>>paddingLevel); 71 | buffer4[k]+=buffer[k2]; 72 | buffer4[k+1]+=buffer[k2]; 73 | buffer4[k+2]+=buffer[k2]; 74 | buffer4[k+3]+=buffer[k2]; 75 | } 76 | } 77 | __syncthreads(); 78 | for (int k=threadIdx.x;k>paddingLevel)]+runningsum2; 82 | float r2=runningsum+t; 83 | runningsum2=t-(r2-runningsum); 84 | runningsum=r2; 85 | __syncthreads(); 86 | } 87 | } 88 | } 89 | 90 | __global__ void binarysearchKernel(int b,int n,int m,const float * __restrict__ dataset,const float * __restrict__ query, int * __restrict__ result){ 91 | int base=1; 92 | while (base=1;k>>=1) 99 | if (r>=k && dataset[i*n+r-k]>=q) 100 | r-=k; 101 | result[i*m+j]=r; 102 | } 103 | } 104 | } 105 | __global__ void farthestpointsamplingKernel(int b,int n,int m,const float * __restrict__ dataset,float * __restrict__ temp,int * __restrict__ idxs){ 106 | if (m<=0) 107 | return; 108 | const int BlockSize=512; 109 | __shared__ float dists[BlockSize]; 110 | __shared__ int dists_i[BlockSize]; 111 | const int BufferSize=3072; 112 | __shared__ float buf[BufferSize*3]; 113 | for (int i=blockIdx.x;ibest){ 147 | best=d2; 148 | besti=k; 149 | } 150 | } 151 | dists[threadIdx.x]=best; 152 | dists_i[threadIdx.x]=besti; 153 | for (int u=0;(1<>(u+1))){ 156 | int i1=(threadIdx.x*2)<>>(b,n,inp,out); 196 | } 197 | //require b*n working space 198 | void probsampleLauncher(int b,int n,int m,const float * inp_p,const float * inp_r,float * temp,int * out){ 199 | cumsumKernel<<<32,512>>>(b,n,inp_p,temp); 200 | binarysearchKernel<<>>(b,n,m,temp,inp_r,out); 201 | } 202 | //require 32*n working space 203 | void farthestpointsamplingLauncher(int b,int n,int m,const float * inp,float * temp,int * out){ 204 | farthestpointsamplingKernel<<<32,512>>>(b,n,m,inp,temp,out); 205 | } 206 | void gatherpointLauncher(int b,int n,int m,const float * inp,const int * idx,float * out){ 207 | gatherpointKernel<<>>(b,n,m,inp,idx,out); 208 | } 209 | void scatteraddpointLauncher(int b,int n,int m,const float * out_g,const int * idx,float * inp_g){ 210 | scatteraddpointKernel<<>>(b,n,m,out_g,idx,inp_g); 211 | } 212 | 213 | -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/sampling/tf_sampling_g.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/sampling/tf_sampling_g.cu.o -------------------------------------------------------------------------------- /convert_to_pointcloud/bunny/tf_ops/sampling/tf_sampling_so.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/convert_to_pointcloud/bunny/tf_ops/sampling/tf_sampling_so.so -------------------------------------------------------------------------------- /h5_dataset/bunny/README: -------------------------------------------------------------------------------- 1 | This is an example of point cloud storage format 2 | "bunny_pprplus.tar" is a pre-trained model for bunny object of IPA datasets for a quick evaluation. 3 | -------------------------------------------------------------------------------- /h5_dataset/bunny/bunny_pprplus.tar: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/h5_dataset/bunny/bunny_pprplus.tar -------------------------------------------------------------------------------- /h5_dataset/bunny/cycle_0000.zip: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/h5_dataset/bunny/cycle_0000.zip -------------------------------------------------------------------------------- /logs/README: -------------------------------------------------------------------------------- 1 | Store training results and logs 2 | -------------------------------------------------------------------------------- /models/calc_lambda.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | # ---------------------utils functions------------------------------ 4 | def load_obj(file_name): 5 | with open(file_name, 'r') as f: 6 | vertexes = [] 7 | faces = [] 8 | for line in f.readlines(): 9 | line = line.strip() 10 | if line.startswith('v'): 11 | words = line.split()[1:] 12 | xyz = [float(w) for w in words] 13 | vertexes.append(xyz) 14 | if line.startswith('f'): 15 | words = line.split()[1:] 16 | indices = [float(w) for w in words] 17 | faces.append(indices) 18 | vertexes = np.array(vertexes) 19 | faces = np.array(faces, dtype=np.int32) 20 | return vertexes, faces 21 | 22 | def triangle_surface_area(a, b, c): 23 | s = np.linalg.norm(np.cross((b-a), (c-a)), ord=2)/2.0 24 | return s 25 | 26 | def triangle_surface_center(a, b, c): 27 | return (a+b+c)/3.0 28 | 29 | 30 | # ---------------------utils functions------------------------------ 31 | 32 | OBJ_PATH = './SileaneBunny.obj' 33 | # the model unit is m 34 | 35 | if __name__ == "__main__": 36 | vertexes, faces = load_obj(OBJ_PATH) 37 | print(vertexes.shape, faces.shape) 38 | # calculate center of mass 39 | center = np.zeros([3]) 40 | total_surface_area = 0 41 | for face in faces: 42 | idx_a, idx_b, idx_c = face[0]-1, face[1]-1, face[2]-1 43 | a, b, c = vertexes[idx_a], vertexes[idx_b], vertexes[idx_c] 44 | s = triangle_surface_area(a,b,c) 45 | total_surface_area += s 46 | center += s * triangle_surface_center(a,b,c) 47 | # print(triangle_surface_center(a,b,c), s) 48 | # exit(0) 49 | center /= total_surface_area 50 | print('Center of mass: ', center) 51 | print('S', total_surface_area) 52 | 53 | # calcuate lambda 54 | vertexes = vertexes - center 55 | L_sq = np.zeros([3,3]) 56 | for face in faces: 57 | idx_a, idx_b, idx_c = face[0]-1, face[1]-1, face[2]-1 58 | a, b, c = vertexes[idx_a], vertexes[idx_b], vertexes[idx_c] 59 | c_tmp, s_tmp = triangle_surface_center(a,b,c), triangle_surface_area(a,b,c) 60 | a = a.reshape([3,1]) 61 | b = b.reshape([3 ,1]) 62 | c = c.reshape([3,1]) 63 | c_tmp = c_tmp.reshape([3,1]) 64 | # aaa = np.dot(c_tmp, c_tmp.T) 65 | theta = s_tmp/12.0*(9.0*np.dot(c_tmp, c_tmp.T) 66 | +np.dot(a,a.T)+np.dot(b,b.T)+np.dot(c,c.T)) 67 | L_sq += theta 68 | L = np.sqrt(L_sq/total_surface_area) 69 | print('lambda: ', L) 70 | print('lambda_rz: ', np.sqrt(L[0,0]*L[0,0]+L[2,2]*L[2,2])) 71 | -------------------------------------------------------------------------------- /pprnet/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/__init__.py -------------------------------------------------------------------------------- /pprnet/__pycache__/__init__.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/__pycache__/__init__.cpython-35.pyc -------------------------------------------------------------------------------- /pprnet/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /pprnet/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /pprnet/__pycache__/object_type.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/__pycache__/object_type.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/__pycache__/pose_loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/__pycache__/pose_loss.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/__pycache__/pprnet.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/__pycache__/pprnet.cpython-35.pyc -------------------------------------------------------------------------------- /pprnet/__pycache__/pprnet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/__pycache__/pprnet.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/__pycache__/pprnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/__pycache__/pprnet.cpython-37.pyc -------------------------------------------------------------------------------- /pprnet/__pycache__/pprnet_test_1.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/__pycache__/pprnet_test_1.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/__pycache__/pprnet_test_1.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/__pycache__/pprnet_test_1.cpython-38.pyc -------------------------------------------------------------------------------- /pprnet/backbone/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | FILE_DIR = os.path.dirname(os.path.abspath(__file__)) 4 | sys.path.append(FILE_DIR) 5 | from pointnet2.pointnet2_backbone import Pointnet2Backbone 6 | from pointnet2.pointnet2_backbone import Pointnet2MSGBackbone 7 | 8 | #from rscnn.rscnn_backbone import RSCNNBackbone 9 | -------------------------------------------------------------------------------- /pprnet/backbone/__pycache__/__init__.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/__pycache__/__init__.cpython-35.pyc -------------------------------------------------------------------------------- /pprnet/backbone/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/pointnet2_backbone.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/pointnet2_backbone.cpython-35.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/pointnet2_backbone.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/pointnet2_backbone.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/pointnet2_backbone.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/pointnet2_backbone.cpython-37.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/pointnet2_modules.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/pointnet2_modules.cpython-35.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/pointnet2_modules.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/pointnet2_modules.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/pointnet2_modules.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/pointnet2_modules.cpython-37.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/pointnet2_utils.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/pointnet2_utils.cpython-35.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/pointnet2_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/pointnet2_utils.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/pointnet2_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/pointnet2_utils.cpython-37.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/pytorch_utils.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/pytorch_utils.cpython-35.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/pytorch_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/pytorch_utils.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/pytorch_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/pytorch_utils.cpython-37.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/__pycache__/test_backbone.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/__pycache__/test_backbone.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/pointnet2/_ext.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/pointnet2/_ext.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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_grad_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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/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 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/__pycache__/rscnn_backbone.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/rscnn/__pycache__/rscnn_backbone.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/rscnn_backbone.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 3 | sys.path.append(BASE_DIR) 4 | sys.path.append(os.path.join(BASE_DIR, "utils")) 5 | import torch 6 | import torch.nn as nn 7 | from torch.autograd import Variable 8 | import rscnn_pytorch_utils as pt_utils 9 | from rscnn_pointnet2_modules import PointnetSAModule, PointnetFPModule, PointnetSAModuleMSG 10 | import numpy as np 11 | 12 | class RSCNNBackbone(nn.Module): 13 | r""" 14 | PointNet2 with multi-scale grouping 15 | Semantic segmentation network that uses feature propogation layers 16 | 17 | Parameters 18 | ---------- 19 | num_classes: int 20 | Number of semantics classes to predict over -- size of softmax classifier that run for each point 21 | input_channels: int = 6 22 | Number of input channels in the feature descriptor for each point. If the point cloud is Nx9, this 23 | value should be 6 as in an Nx9 point cloud, 3 of the channels are xyz, and 6 are feature descriptors 24 | use_xyz: bool = True 25 | Whether or not to use the xyz position of a point as a feature 26 | """ 27 | 28 | def __init__(self, input_channels=0, relation_prior=1, use_xyz=True): 29 | super().__init__() 30 | 31 | self.SA_modules = nn.ModuleList() 32 | c_in = input_channels 33 | self.SA_modules.append( # 0 34 | PointnetSAModuleMSG( 35 | npoint=2048, 36 | radii=[20, 30, 40], 37 | nsamples=[16, 32, 48], 38 | mlps=[[c_in, 64], [c_in, 64], [c_in, 64]], 39 | first_layer=True, 40 | use_xyz=use_xyz, 41 | relation_prior=relation_prior 42 | ) 43 | ) 44 | c_out_0 = 64*3 45 | 46 | c_in = c_out_0 47 | self.SA_modules.append( # 1 48 | PointnetSAModuleMSG( 49 | npoint=1024, 50 | radii=[40, 60, 80], 51 | nsamples=[16, 48, 64], 52 | mlps=[[c_in, 128], [c_in, 128], [c_in, 128]], 53 | use_xyz=use_xyz, 54 | relation_prior=relation_prior 55 | ) 56 | ) 57 | c_out_1 = 128*3 58 | 59 | c_in = c_out_1 60 | self.SA_modules.append( # 2 61 | PointnetSAModuleMSG( 62 | npoint=256, 63 | radii=[80, 120, 160], 64 | nsamples=[16, 32, 48], 65 | mlps=[[c_in, 256], [c_in, 256], [c_in, 256]], 66 | use_xyz=use_xyz, 67 | relation_prior=relation_prior 68 | ) 69 | ) 70 | c_out_2 = 256*3 71 | 72 | c_in = c_out_2 73 | self.SA_modules.append( # 3 74 | PointnetSAModuleMSG( 75 | npoint=64, 76 | radii=[200, 240, 280], 77 | nsamples=[16, 24, 32], 78 | mlps=[[c_in, 512], [c_in, 512], [c_in, 512]], 79 | use_xyz=use_xyz, 80 | relation_prior=relation_prior 81 | ) 82 | ) 83 | c_out_3 = 512*3 84 | 85 | self.SA_modules.append( # 4 global pooling 86 | PointnetSAModule( 87 | nsample = 64, 88 | mlp=[c_out_3, 128], use_xyz=use_xyz 89 | ) 90 | ) 91 | global_out = 128 92 | 93 | self.SA_modules.append( # 5 global pooling 94 | PointnetSAModule( 95 | nsample = 256, 96 | mlp=[c_out_2, 128], use_xyz=use_xyz 97 | ) 98 | ) 99 | global_out2 = 128 100 | 101 | self.FP_modules = nn.ModuleList() 102 | self.FP_modules.append( 103 | PointnetFPModule(mlp=[256 + input_channels, 128, 128]) 104 | ) 105 | self.FP_modules.append(PointnetFPModule(mlp=[512 + c_out_0, 256, 256])) 106 | self.FP_modules.append(PointnetFPModule(mlp=[512 + c_out_1, 512, 512])) 107 | self.FP_modules.append( 108 | PointnetFPModule(mlp=[c_out_3 + c_out_2, 512, 512]) 109 | ) 110 | 111 | self.FC_layer = pt_utils.Conv1d(128+global_out+global_out2, 128, bn=True, activation=None) 112 | 113 | def _break_up_pc(self, pc): 114 | xyz = pc[..., 0:3].contiguous() 115 | features = ( 116 | pc[..., 3:].transpose(1, 2).contiguous() 117 | if pc.size(-1) > 3 else None 118 | ) 119 | 120 | return xyz, features 121 | 122 | def forward(self, pointcloud: torch.cuda.FloatTensor): 123 | r""" 124 | Forward pass of the network 125 | 126 | Parameters 127 | ---------- 128 | pointcloud: Variable(torch.cuda.FloatTensor) 129 | (B, N, 3 + input_channels) tensor 130 | Point cloud to run predicts on 131 | Each point in the point-cloud MUST 132 | be formated as (x, y, z, features...) 133 | """ 134 | xyz, features = self._break_up_pc(pointcloud) 135 | 136 | l_xyz, l_features = [xyz], [features] 137 | for i in range(len(self.SA_modules)): 138 | if i < 5: 139 | li_xyz, li_features = self.SA_modules[i](l_xyz[i], l_features[i]) 140 | if li_xyz is not None: 141 | random_index = np.arange(li_xyz.size()[1]) 142 | np.random.shuffle(random_index) 143 | li_xyz = li_xyz[:, random_index, :] 144 | li_features = li_features[:, :, random_index] 145 | l_xyz.append(li_xyz) 146 | l_features.append(li_features) 147 | 148 | _, global_out2_feat = self.SA_modules[5](l_xyz[3], l_features[3]) 149 | 150 | for i in range(-1, -(len(self.FP_modules) + 1), -1): 151 | l_features[i - 1 - 1] = self.FP_modules[i]( 152 | l_xyz[i - 1 - 1], l_xyz[i - 1], l_features[i - 1 - 1], l_features[i - 1] 153 | ) 154 | 155 | l_features[0] = torch.cat((l_features[0], l_features[-1].repeat(1, 1, l_features[0].size()[2]), global_out2_feat.repeat(1, 1, l_features[0].size()[2])), 1) 156 | return self.FC_layer(l_features[0]) -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/rscnn/utils/__init__.py -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/__pycache__/linalg_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/rscnn/utils/__pycache__/linalg_utils.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/__pycache__/rscnn_pointnet2_modules.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/rscnn/utils/__pycache__/rscnn_pointnet2_modules.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/__pycache__/rscnn_pointnet2_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/rscnn/utils/__pycache__/rscnn_pointnet2_utils.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/_ext/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/rscnn/utils/_ext/__init__.py -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/_ext/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/rscnn/utils/_ext/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/_ext/pointnet2/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from torch.utils.ffi import _wrap_function 3 | from ._pointnet2 import lib as _lib, ffi as _ffi 4 | 5 | __all__ = [] 6 | def _import_symbols(locals): 7 | for symbol in dir(_lib): 8 | fn = getattr(_lib, symbol) 9 | if callable(fn): 10 | locals[symbol] = _wrap_function(fn, _ffi) 11 | else: 12 | locals[symbol] = fn 13 | __all__.append(symbol) 14 | 15 | _import_symbols(locals()) 16 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/_ext/pointnet2/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/rscnn/utils/_ext/pointnet2/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/_ext/pointnet2/_pointnet2.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/rscnn/utils/_ext/pointnet2/_pointnet2.so -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/build_ffi.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import torch 3 | import os.path as osp 4 | from torch.utils.ffi import create_extension 5 | import sys, argparse, shutil 6 | 7 | base_dir = osp.dirname(osp.abspath(__file__)) 8 | 9 | 10 | def parse_args(): 11 | parser = argparse.ArgumentParser( 12 | description="Arguments for building pointnet2 ffi extension" 13 | ) 14 | parser.add_argument("--objs", nargs="*") 15 | clean_arg = parser.add_mutually_exclusive_group() 16 | clean_arg.add_argument("--build", dest='build', action="store_true") 17 | clean_arg.add_argument("--clean", dest='clean', action="store_true") 18 | parser.set_defaults(build=False, clean=False) 19 | 20 | args = parser.parse_args() 21 | assert args.build or args.clean 22 | 23 | return args 24 | 25 | 26 | def build(args): 27 | extra_objects = args.objs 28 | extra_objects += [a for a in glob.glob('/usr/local/cuda/lib64/*.a')] 29 | 30 | ffi = create_extension( 31 | '_ext.pointnet2', 32 | headers=[a for a in glob.glob("cinclude/*_wrapper.h")], 33 | sources=[a for a in glob.glob("csrc/*.c")], 34 | define_macros=[('WITH_CUDA', None)], 35 | relative_to=__file__, 36 | with_cuda=True, 37 | extra_objects=extra_objects, 38 | include_dirs=[osp.join(base_dir, 'cinclude')], 39 | verbose=False, 40 | package=False 41 | ) 42 | ffi.build() 43 | 44 | 45 | def clean(args): 46 | shutil.rmtree(osp.join(base_dir, "_ext")) 47 | 48 | 49 | if __name__ == "__main__": 50 | args = parse_args() 51 | if args.clean: 52 | clean(args) 53 | else: 54 | build(args) 55 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/cinclude/ball_query_gpu.h: -------------------------------------------------------------------------------- 1 | #ifndef _BALL_QUERY_GPU 2 | #define _BALL_QUERY_GPU 3 | 4 | #ifdef __cplusplus 5 | extern "C" { 6 | #endif 7 | 8 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 9 | int nsample, const float *xyz, 10 | const float *new_xyz, const int *fps_idx, int *idx, 11 | cudaStream_t stream); 12 | 13 | #ifdef __cplusplus 14 | } 15 | #endif 16 | #endif 17 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/cinclude/ball_query_wrapper.h: -------------------------------------------------------------------------------- 1 | 2 | int ball_query_wrapper(int b, int n, int m, float radius, int nsample, 3 | THCudaTensor *new_xyz_tensor, THCudaTensor *xyz_tensor, THCudaIntTensor *fps_idx_tensor, 4 | THCudaIntTensor *idx_tensor); 5 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/cinclude/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef _CUDA_UTILS_H 2 | #define _CUDA_UTILS_H 3 | 4 | #include 5 | 6 | #define TOTAL_THREADS 512 7 | 8 | inline int opt_n_threads(int work_size) { 9 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 10 | 11 | return max(min(1 << pow_2, TOTAL_THREADS), 1); 12 | } 13 | 14 | inline dim3 opt_block_config(int x, int y) { 15 | const int x_threads = opt_n_threads(x); 16 | const int y_threads = 17 | max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 18 | dim3 block_config(x_threads, y_threads, 1); 19 | 20 | return block_config; 21 | } 22 | 23 | #endif 24 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/cinclude/group_points_gpu.h: -------------------------------------------------------------------------------- 1 | #ifndef _BALL_QUERY_GPU 2 | #define _BALL_QUERY_GPU 3 | 4 | #ifdef __cplusplus 5 | extern "C" { 6 | #endif 7 | 8 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 9 | const float *points, const int *idx, 10 | float *out, cudaStream_t stream); 11 | 12 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 13 | int nsample, const float *grad_out, 14 | const int *idx, float *grad_points, 15 | cudaStream_t stream); 16 | #ifdef __cplusplus 17 | } 18 | #endif 19 | #endif 20 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/cinclude/group_points_wrapper.h: -------------------------------------------------------------------------------- 1 | int group_points_wrapper(int b, int c, int n, int npoints, int nsample, 2 | THCudaTensor *points_tensor, 3 | THCudaIntTensor *idx_tensor, THCudaTensor *out); 4 | int group_points_grad_wrapper(int b, int c, int n, int npoints, int nsample, 5 | THCudaTensor *grad_out_tensor, 6 | THCudaIntTensor *idx_tensor, 7 | THCudaTensor *grad_points_tensor); 8 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/cinclude/interpolate_gpu.h: -------------------------------------------------------------------------------- 1 | #ifndef _INTERPOLATE_GPU_H 2 | #define _INTERPOLATE_GPU_H 3 | 4 | #ifdef __cplusplus 5 | extern "C" { 6 | #endif 7 | 8 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 9 | const float *known, float *dist2, int *idx, 10 | cudaStream_t stream); 11 | 12 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 13 | const float *points, const int *idx, 14 | const float *weight, float *out, 15 | cudaStream_t stream); 16 | 17 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 18 | const float *grad_out, 19 | const int *idx, const float *weight, 20 | float *grad_points, 21 | cudaStream_t stream); 22 | 23 | #ifdef __cplusplus 24 | } 25 | #endif 26 | 27 | #endif 28 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/cinclude/interpolate_wrapper.h: -------------------------------------------------------------------------------- 1 | 2 | 3 | void three_nn_wrapper(int b, int n, int m, THCudaTensor *unknown_tensor, 4 | THCudaTensor *known_tensor, THCudaTensor *dist2_tensor, 5 | THCudaIntTensor *idx_tensor); 6 | void three_interpolate_wrapper(int b, int c, int m, int n, 7 | THCudaTensor *points_tensor, 8 | THCudaIntTensor *idx_tensor, 9 | THCudaTensor *weight_tensor, 10 | THCudaTensor *out_tensor); 11 | 12 | void three_interpolate_grad_wrapper(int b, int c, int n, int m, 13 | THCudaTensor *grad_out_tensor, 14 | THCudaIntTensor *idx_tensor, 15 | THCudaTensor *weight_tensor, 16 | THCudaTensor *grad_points_tensor); 17 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/cinclude/sampling_gpu.h: -------------------------------------------------------------------------------- 1 | #ifndef _SAMPLING_GPU_H 2 | #define _SAMPLING_GPU_H 3 | 4 | #ifdef __cplusplus 5 | extern "C" { 6 | #endif 7 | 8 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 9 | const float *points, const int *idx, 10 | float *out, cudaStream_t stream); 11 | 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, cudaStream_t stream); 15 | 16 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 17 | const float *dataset, float *temp, 18 | int *idxs, cudaStream_t stream); 19 | 20 | #ifdef __cplusplus 21 | } 22 | #endif 23 | #endif 24 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/cinclude/sampling_wrapper.h: -------------------------------------------------------------------------------- 1 | 2 | int gather_points_wrapper(int b, int c, int n, int npoints, 3 | THCudaTensor *points_tensor, 4 | THCudaIntTensor *idx_tensor, 5 | THCudaTensor *out_tensor); 6 | int gather_points_grad_wrapper(int b, int c, int n, int npoints, 7 | THCudaTensor *grad_out_tensor, 8 | THCudaIntTensor *idx_tensor, 9 | THCudaTensor *grad_points_tensor); 10 | 11 | int furthest_point_sampling_wrapper(int b, int n, int m, 12 | THCudaTensor *points_tensor, 13 | THCudaTensor *temp_tensor, 14 | THCudaIntTensor *idx_tensor); 15 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/csrc/ball_query.c: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include "ball_query_gpu.h" 4 | 5 | extern THCState *state; 6 | 7 | int ball_query_wrapper(int b, int n, int m, float radius, int nsample, 8 | THCudaTensor *new_xyz_tensor, THCudaTensor *xyz_tensor, THCudaIntTensor *fps_idx_tensor, 9 | THCudaIntTensor *idx_tensor) { 10 | 11 | const float *new_xyz = THCudaTensor_data(state, new_xyz_tensor); 12 | const float *xyz = THCudaTensor_data(state, xyz_tensor); 13 | const int *fps_idx = THCudaIntTensor_data(state, fps_idx_tensor); 14 | int *idx = THCudaIntTensor_data(state, idx_tensor); 15 | 16 | cudaStream_t stream = THCState_getCurrentStream(state); 17 | 18 | query_ball_point_kernel_wrapper(b, n, m, radius, nsample, new_xyz, xyz, fps_idx, idx, 19 | stream); 20 | return 1; 21 | } 22 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/csrc/ball_query_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "ball_query_gpu.h" 6 | #include "cuda_utils.h" 7 | 8 | // input: new_xyz(b, m, 3) xyz(b, n, 3) 9 | // output: idx(b, m, nsample) 10 | __global__ void query_ball_point_kernel(int b, int n, int m, float radius, 11 | int nsample, 12 | const float *__restrict__ new_xyz, 13 | const float *__restrict__ xyz, 14 | const int *__restrict__ fps_idx, 15 | int *__restrict__ idx) { 16 | int batch_index = blockIdx.x; 17 | xyz += batch_index * n * 3; 18 | new_xyz += batch_index * m * 3; 19 | fps_idx += batch_index * m; 20 | idx += m * nsample * batch_index; 21 | 22 | int index = threadIdx.x; 23 | int stride = blockDim.x; 24 | 25 | float radius2 = radius * radius; 26 | for (int j = index; j < m; j += stride) { 27 | float new_x = new_xyz[j * 3 + 0]; 28 | float new_y = new_xyz[j * 3 + 1]; 29 | float new_z = new_xyz[j * 3 + 2]; 30 | for (int l = 0; l < nsample; ++l) { 31 | idx[j * nsample + l] = fps_idx[j]; 32 | } 33 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { 34 | float x = xyz[k * 3 + 0]; 35 | float y = xyz[k * 3 + 1]; 36 | float z = xyz[k * 3 + 2]; 37 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 38 | (new_z - z) * (new_z - z); 39 | if (d2 < radius2 && d2 > 0) { 40 | idx[j * nsample + cnt] = k; 41 | ++cnt; 42 | } 43 | } 44 | } 45 | } 46 | 47 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 48 | int nsample, const float *new_xyz, 49 | const float *xyz, const int *fps_idx, int *idx, 50 | cudaStream_t stream) { 51 | 52 | cudaError_t err; 53 | query_ball_point_kernel<<>>( 54 | b, n, m, radius, nsample, new_xyz, xyz, fps_idx, idx); 55 | 56 | err = cudaGetLastError(); 57 | if (cudaSuccess != err) { 58 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 59 | exit(-1); 60 | } 61 | } 62 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/csrc/group_points.c: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include "group_points_gpu.h" 4 | 5 | extern THCState *state; 6 | 7 | int group_points_wrapper(int b, int c, int n, int npoints, int nsample, 8 | THCudaTensor *points_tensor, 9 | THCudaIntTensor *idx_tensor, 10 | THCudaTensor *out_tensor) { 11 | 12 | const float *points = THCudaTensor_data(state, points_tensor); 13 | const int *idx = THCudaIntTensor_data(state, idx_tensor); 14 | float *out = THCudaTensor_data(state, out_tensor); 15 | 16 | cudaStream_t stream = THCState_getCurrentStream(state); 17 | 18 | group_points_kernel_wrapper(b, c, n, npoints, nsample, points, idx, out, 19 | stream); 20 | return 1; 21 | } 22 | 23 | int group_points_grad_wrapper(int b, int c, int n, int npoints, int nsample, 24 | THCudaTensor *grad_out_tensor, 25 | THCudaIntTensor *idx_tensor, 26 | THCudaTensor *grad_points_tensor) { 27 | 28 | float *grad_points = THCudaTensor_data(state, grad_points_tensor); 29 | const int *idx = THCudaIntTensor_data(state, idx_tensor); 30 | const float *grad_out = THCudaTensor_data(state, grad_out_tensor); 31 | 32 | cudaStream_t stream = THCState_getCurrentStream(state); 33 | 34 | group_points_grad_kernel_wrapper(b, c, n, npoints, nsample, grad_out, idx, 35 | grad_points, stream); 36 | return 1; 37 | } 38 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/csrc/group_points_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "cuda_utils.h" 5 | #include "group_points_gpu.h" 6 | 7 | // input: points(b, c, n) idx(b, npoints, nsample) 8 | // output: out(b, c, npoints, nsample) 9 | __global__ void group_points_kernel(int b, int c, int n, int npoints, 10 | int nsample, 11 | const float *__restrict__ points, 12 | const int *__restrict__ idx, 13 | float *__restrict__ out) { 14 | int batch_index = blockIdx.x; 15 | points += batch_index * n * c; 16 | idx += batch_index * npoints * nsample; 17 | out += batch_index * npoints * nsample * c; 18 | 19 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 20 | const int stride = blockDim.y * blockDim.x; 21 | for (int i = index; i < c * npoints; i += stride) { 22 | const int l = i / npoints; 23 | const int j = i % npoints; 24 | for (int k = 0; k < nsample; ++k) { 25 | int ii = idx[j * nsample + k]; 26 | out[(l * npoints + j) * nsample + k] = points[l * n + ii]; 27 | } 28 | } 29 | } 30 | 31 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 32 | const float *points, const int *idx, 33 | float *out, cudaStream_t stream) { 34 | 35 | cudaError_t err; 36 | group_points_kernel<<>>( 37 | b, c, n, npoints, nsample, points, idx, out); 38 | 39 | err = cudaGetLastError(); 40 | if (cudaSuccess != err) { 41 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 42 | exit(-1); 43 | } 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) { 75 | cudaError_t err; 76 | group_points_grad_kernel<<>>( 77 | b, c, n, npoints, nsample, grad_out, idx, grad_points); 78 | 79 | err = cudaGetLastError(); 80 | if (cudaSuccess != err) { 81 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 82 | exit(-1); 83 | } 84 | } 85 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/csrc/interpolate.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #include "interpolate_gpu.h" 7 | 8 | extern THCState *state; 9 | 10 | void three_nn_wrapper(int b, int n, int m, THCudaTensor *unknown_tensor, 11 | THCudaTensor *known_tensor, THCudaTensor *dist2_tensor, 12 | THCudaIntTensor *idx_tensor) { 13 | const float *unknown = THCudaTensor_data(state, unknown_tensor); 14 | const float *known = THCudaTensor_data(state, known_tensor); 15 | float *dist2 = THCudaTensor_data(state, dist2_tensor); 16 | int *idx = THCudaIntTensor_data(state, idx_tensor); 17 | 18 | cudaStream_t stream = THCState_getCurrentStream(state); 19 | three_nn_kernel_wrapper(b, n, m, unknown, known, dist2, idx, stream); 20 | } 21 | 22 | void three_interpolate_wrapper(int b, int c, int m, int n, 23 | THCudaTensor *points_tensor, 24 | THCudaIntTensor *idx_tensor, 25 | THCudaTensor *weight_tensor, 26 | THCudaTensor *out_tensor) { 27 | 28 | const float *points = THCudaTensor_data(state, points_tensor); 29 | const float *weight = THCudaTensor_data(state, weight_tensor); 30 | float *out = THCudaTensor_data(state, out_tensor); 31 | const int *idx = THCudaIntTensor_data(state, idx_tensor); 32 | 33 | cudaStream_t stream = THCState_getCurrentStream(state); 34 | three_interpolate_kernel_wrapper(b, c, m, n, points, idx, weight, out, 35 | stream); 36 | } 37 | 38 | void three_interpolate_grad_wrapper(int b, int c, int n, int m, 39 | THCudaTensor *grad_out_tensor, 40 | THCudaIntTensor *idx_tensor, 41 | THCudaTensor *weight_tensor, 42 | THCudaTensor *grad_points_tensor) { 43 | 44 | const float *grad_out = THCudaTensor_data(state, grad_out_tensor); 45 | const float *weight = THCudaTensor_data(state, weight_tensor); 46 | float *grad_points = THCudaTensor_data(state, grad_points_tensor); 47 | const int *idx = THCudaIntTensor_data(state, idx_tensor); 48 | 49 | cudaStream_t stream = THCState_getCurrentStream(state); 50 | three_interpolate_grad_kernel_wrapper(b, c, n, m, grad_out, idx, weight, 51 | grad_points, stream); 52 | } 53 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/csrc/interpolate_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "cuda_utils.h" 6 | #include "interpolate_gpu.h" 7 | 8 | // input: unknown(b, n, 3) known(b, m, 3) 9 | // output: dist2(b, n, 3), idx(b, n, 3) 10 | __global__ void three_nn_kernel(int b, int n, int m, 11 | const float *__restrict__ unknown, 12 | const float *__restrict__ known, 13 | float *__restrict__ dist2, 14 | int *__restrict__ idx) { 15 | int batch_index = blockIdx.x; 16 | unknown += batch_index * n * 3; 17 | known += batch_index * m * 3; 18 | dist2 += batch_index * n * 3; 19 | idx += batch_index * n * 3; 20 | 21 | int index = threadIdx.x; 22 | int stride = blockDim.x; 23 | for (int j = index; j < n; j += stride) { 24 | float ux = unknown[j * 3 + 0]; 25 | float uy = unknown[j * 3 + 1]; 26 | float uz = unknown[j * 3 + 2]; 27 | 28 | double best1 = 1e40, best2 = 1e40, best3 = 1e40; 29 | int besti1 = 0, besti2 = 0, besti3 = 0; 30 | for (int k = 0; k < m; ++k) { 31 | float x = known[k * 3 + 0]; 32 | float y = known[k * 3 + 1]; 33 | float z = known[k * 3 + 2]; 34 | float d = 35 | (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z); 36 | if (d < best1) { 37 | best3 = best2; 38 | besti3 = besti2; 39 | best2 = best1; 40 | besti2 = besti1; 41 | best1 = d; 42 | besti1 = k; 43 | } else if (d < best2) { 44 | best3 = best2; 45 | besti3 = besti2; 46 | best2 = d; 47 | besti2 = k; 48 | } else if (d < best3) { 49 | best3 = d; 50 | besti3 = k; 51 | } 52 | } 53 | dist2[j * 3 + 0] = best1; 54 | dist2[j * 3 + 1] = best2; 55 | dist2[j * 3 + 2] = best3; 56 | 57 | idx[j * 3 + 0] = besti1; 58 | idx[j * 3 + 1] = besti2; 59 | idx[j * 3 + 2] = besti3; 60 | } 61 | } 62 | 63 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 64 | const float *known, float *dist2, int *idx, 65 | cudaStream_t stream) { 66 | 67 | cudaError_t err; 68 | three_nn_kernel<<>>(b, n, m, unknown, known, 69 | dist2, idx); 70 | 71 | err = cudaGetLastError(); 72 | if (cudaSuccess != err) { 73 | fprintf(stderr, "CUDA kernel " 74 | "failed : %s\n", 75 | cudaGetErrorString(err)); 76 | exit(-1); 77 | } 78 | } 79 | 80 | // input: points(b, c, m), idx(b, n, 3), weight(b, n, 3) 81 | // output: out(b, c, n) 82 | __global__ void three_interpolate_kernel(int b, int c, int m, int n, 83 | const float *__restrict__ points, 84 | const int *__restrict__ idx, 85 | const float *__restrict__ weight, 86 | float *__restrict__ out) { 87 | int batch_index = blockIdx.x; 88 | points += batch_index * m * c; 89 | 90 | idx += batch_index * n * 3; 91 | weight += batch_index * n * 3; 92 | 93 | out += batch_index * n * c; 94 | 95 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 96 | const int stride = blockDim.y * blockDim.x; 97 | for (int i = index; i < c * n; i += stride) { 98 | const int l = i / n; 99 | const int j = i % n; 100 | float w1 = weight[j * 3 + 0]; 101 | float w2 = weight[j * 3 + 1]; 102 | float w3 = weight[j * 3 + 2]; 103 | 104 | int i1 = idx[j * 3 + 0]; 105 | int i2 = idx[j * 3 + 1]; 106 | int i3 = idx[j * 3 + 2]; 107 | 108 | out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 + 109 | points[l * m + i3] * w3; 110 | } 111 | } 112 | 113 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 114 | const float *points, const int *idx, 115 | const float *weight, float *out, 116 | cudaStream_t stream) { 117 | 118 | cudaError_t err; 119 | three_interpolate_kernel<<>>( 120 | b, c, m, n, points, idx, weight, out); 121 | 122 | err = cudaGetLastError(); 123 | if (cudaSuccess != err) { 124 | fprintf(stderr, "CUDA kernel " 125 | "failed : %s\n", 126 | cudaGetErrorString(err)); 127 | exit(-1); 128 | } 129 | } 130 | 131 | // input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3) 132 | // output: grad_points(b, c, m) 133 | 134 | __global__ void three_interpolate_grad_kernel( 135 | int b, int c, int n, int m, const float *__restrict__ grad_out, 136 | const int *__restrict__ idx, const float *__restrict__ weight, 137 | float *__restrict__ grad_points) { 138 | int batch_index = blockIdx.x; 139 | grad_out += batch_index * n * c; 140 | idx += batch_index * n * 3; 141 | weight += batch_index * n * 3; 142 | grad_points += batch_index * m * c; 143 | 144 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 145 | const int stride = blockDim.y * blockDim.x; 146 | for (int i = index; i < c * n; i += stride) { 147 | const int l = i / n; 148 | const int j = i % n; 149 | float w1 = weight[j * 3 + 0]; 150 | float w2 = weight[j * 3 + 1]; 151 | float w3 = weight[j * 3 + 2]; 152 | 153 | int i1 = idx[j * 3 + 0]; 154 | int i2 = idx[j * 3 + 1]; 155 | int i3 = idx[j * 3 + 2]; 156 | 157 | atomicAdd(grad_points + l * m + i1, grad_out[i] * w1); 158 | atomicAdd(grad_points + l * m + i2, grad_out[i] * w2); 159 | atomicAdd(grad_points + l * m + i3, grad_out[i] * w3); 160 | } 161 | } 162 | 163 | void three_interpolate_grad_kernel_wrapper(int b, int n, int c, int m, 164 | const float *grad_out, 165 | const int *idx, const float *weight, 166 | float *grad_points, 167 | cudaStream_t stream) { 168 | 169 | cudaError_t err; 170 | three_interpolate_grad_kernel<<>>( 171 | b, n, c, m, grad_out, idx, weight, grad_points); 172 | 173 | err = cudaGetLastError(); 174 | if (cudaSuccess != err) { 175 | fprintf(stderr, "CUDA kernel " 176 | "failed : %s\n", 177 | cudaGetErrorString(err)); 178 | exit(-1); 179 | } 180 | } 181 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/csrc/sampling.c: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include "sampling_gpu.h" 4 | 5 | extern THCState *state; 6 | 7 | int gather_points_wrapper(int b, int c, int n, int npoints, 8 | THCudaTensor *points_tensor, 9 | THCudaIntTensor *idx_tensor, 10 | THCudaTensor *out_tensor) { 11 | 12 | const float *points = THCudaTensor_data(state, points_tensor); 13 | const int *idx = THCudaIntTensor_data(state, idx_tensor); 14 | float *out = THCudaTensor_data(state, out_tensor); 15 | 16 | cudaStream_t stream = THCState_getCurrentStream(state); 17 | 18 | gather_points_kernel_wrapper(b, c, n, npoints, points, idx, out, stream); 19 | return 1; 20 | } 21 | 22 | int gather_points_grad_wrapper(int b, int c, int n, int npoints, 23 | THCudaTensor *grad_out_tensor, 24 | THCudaIntTensor *idx_tensor, 25 | THCudaTensor *grad_points_tensor) { 26 | 27 | const float *grad_out = THCudaTensor_data(state, grad_out_tensor); 28 | const int *idx = THCudaIntTensor_data(state, idx_tensor); 29 | float *grad_points = THCudaTensor_data(state, grad_points_tensor); 30 | 31 | cudaStream_t stream = THCState_getCurrentStream(state); 32 | 33 | gather_points_grad_kernel_wrapper(b, c, n, npoints, grad_out, idx, 34 | grad_points, stream); 35 | return 1; 36 | } 37 | 38 | int furthest_point_sampling_wrapper(int b, int n, int m, 39 | THCudaTensor *points_tensor, 40 | THCudaTensor *temp_tensor, 41 | THCudaIntTensor *idx_tensor) { 42 | 43 | const float *points = THCudaTensor_data(state, points_tensor); 44 | float *temp = THCudaTensor_data(state, temp_tensor); 45 | int *idx = THCudaIntTensor_data(state, idx_tensor); 46 | 47 | cudaStream_t stream = THCState_getCurrentStream(state); 48 | 49 | furthest_point_sampling_kernel_wrapper(b, n, m, points, temp, idx, stream); 50 | return 1; 51 | } 52 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/csrc/sampling_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "cuda_utils.h" 5 | #include "sampling_gpu.h" 6 | 7 | // input: points(b, c, n) idx(b, m) 8 | // output: out(b, c, m) 9 | __global__ void gather_points_kernel(int b, int c, int n, int m, 10 | const float *__restrict__ points, 11 | const int *__restrict__ idx, 12 | float *__restrict__ out) { 13 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 14 | for (int l = blockIdx.y; l < c; l += gridDim.y) { 15 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 16 | int a = idx[i * m + j]; 17 | out[(i * c + l) * m + j] = points[(i * c + l) * n + a]; 18 | } 19 | } 20 | } 21 | } 22 | 23 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 24 | const float *points, const int *idx, 25 | float *out, cudaStream_t stream) { 26 | 27 | cudaError_t err; 28 | gather_points_kernel<<>>( 29 | b, c, n, npoints, points, idx, out); 30 | 31 | err = cudaGetLastError(); 32 | if (cudaSuccess != err) { 33 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 34 | exit(-1); 35 | } 36 | } 37 | 38 | // input: grad_out(b, c, m) idx(b, m) 39 | // output: grad_points(b, c, n) 40 | __global__ void gather_points_grad_kernel(int b, int c, int n, int m, 41 | const float *__restrict__ grad_out, 42 | const int *__restrict__ idx, 43 | float *__restrict__ grad_points) { 44 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 45 | for (int l = blockIdx.y; l < c; l += gridDim.y) { 46 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 47 | int a = idx[i * m + j]; 48 | atomicAdd(grad_points + (i * c + l) * n + a, 49 | grad_out[(i * c + l) * m + j]); 50 | } 51 | } 52 | } 53 | } 54 | 55 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 56 | const float *grad_out, const int *idx, 57 | float *grad_points, 58 | cudaStream_t stream) { 59 | 60 | cudaError_t err; 61 | gather_points_grad_kernel<<>>(b, c, n, npoints, grad_out, idx, 63 | grad_points); 64 | 65 | err = cudaGetLastError(); 66 | if (cudaSuccess != err) { 67 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 68 | exit(-1); 69 | } 70 | } 71 | 72 | __device__ void __update(float *__restrict__ dists, int *__restrict__ dists_i, 73 | int idx1, int idx2) { 74 | const float v1 = dists[idx1], v2 = dists[idx2]; 75 | const int i1 = dists_i[idx1], i2 = dists_i[idx2]; 76 | dists[idx1] = max(v1, v2); 77 | dists_i[idx1] = v2 > v1 ? i2 : i1; 78 | } 79 | 80 | // Input dataset: (b, n, 3), tmp: (b, n) 81 | // Ouput idxs (b, m) 82 | template 83 | __global__ void furthest_point_sampling_kernel( 84 | int b, int n, int m, const float *__restrict__ dataset, 85 | float *__restrict__ temp, int *__restrict__ idxs) { 86 | if (m <= 0) 87 | return; 88 | __shared__ float dists[block_size]; 89 | __shared__ int dists_i[block_size]; 90 | 91 | int batch_index = blockIdx.x; 92 | dataset += batch_index * n * 3; 93 | temp += batch_index * n; 94 | idxs += batch_index * m; 95 | 96 | int tid = threadIdx.x; 97 | const int stride = block_size; 98 | 99 | int old = 0; 100 | if (threadIdx.x == 0) 101 | idxs[0] = old; 102 | 103 | __syncthreads(); 104 | for (int j = 1; j < m; j++) { 105 | int besti = 0; 106 | float best = -1; 107 | float x1 = dataset[old * 3 + 0]; 108 | float y1 = dataset[old * 3 + 1]; 109 | float z1 = dataset[old * 3 + 2]; 110 | for (int k = tid; k < n; k += stride) { 111 | float x2, y2, z2; 112 | x2 = dataset[k * 3 + 0]; 113 | y2 = dataset[k * 3 + 1]; 114 | z2 = dataset[k * 3 + 2]; 115 | float mag = (x2 * x2) + (y2 * y2) + (z2 * z2); 116 | if (mag <= 1e-3) 117 | continue; 118 | 119 | float d = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + 120 | (z2 - z1) * (z2 - z1); 121 | 122 | float d2 = min(d, temp[k]); 123 | temp[k] = d2; 124 | besti = d2 > best ? k : besti; 125 | best = d2 > best ? d2 : best; 126 | } 127 | dists[tid] = best; 128 | dists_i[tid] = besti; 129 | __syncthreads(); 130 | 131 | if (block_size >= 512) { 132 | if (tid < 256) { 133 | __update(dists, dists_i, tid, tid + 256); 134 | } 135 | __syncthreads(); 136 | } 137 | if (block_size >= 256) { 138 | if (tid < 128) { 139 | __update(dists, dists_i, tid, tid + 128); 140 | } 141 | __syncthreads(); 142 | } 143 | if (block_size >= 128) { 144 | if (tid < 64) { 145 | __update(dists, dists_i, tid, tid + 64); 146 | } 147 | __syncthreads(); 148 | } 149 | if (block_size >= 64) { 150 | if (tid < 32) { 151 | __update(dists, dists_i, tid, tid + 32); 152 | } 153 | __syncthreads(); 154 | } 155 | if (block_size >= 32) { 156 | if (tid < 16) { 157 | __update(dists, dists_i, tid, tid + 16); 158 | } 159 | __syncthreads(); 160 | } 161 | if (block_size >= 16) { 162 | if (tid < 8) { 163 | __update(dists, dists_i, tid, tid + 8); 164 | } 165 | __syncthreads(); 166 | } 167 | if (block_size >= 8) { 168 | if (tid < 4) { 169 | __update(dists, dists_i, tid, tid + 4); 170 | } 171 | __syncthreads(); 172 | } 173 | if (block_size >= 4) { 174 | if (tid < 2) { 175 | __update(dists, dists_i, tid, tid + 2); 176 | } 177 | __syncthreads(); 178 | } 179 | if (block_size >= 2) { 180 | if (tid < 1) { 181 | __update(dists, dists_i, tid, tid + 1); 182 | } 183 | __syncthreads(); 184 | } 185 | 186 | old = dists_i[0]; 187 | if (tid == 0) 188 | idxs[j] = old; 189 | } 190 | } 191 | 192 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 193 | const float *dataset, float *temp, 194 | int *idxs, cudaStream_t stream) { 195 | 196 | cudaError_t err; 197 | unsigned int n_threads = opt_n_threads(n); 198 | 199 | switch (n_threads) { 200 | case 512: 201 | furthest_point_sampling_kernel<512><<>>( 202 | b, n, m, dataset, temp, idxs); 203 | break; 204 | case 256: 205 | furthest_point_sampling_kernel<256><<>>( 206 | b, n, m, dataset, temp, idxs); 207 | break; 208 | case 128: 209 | furthest_point_sampling_kernel<128><<>>( 210 | b, n, m, dataset, temp, idxs); 211 | break; 212 | case 64: 213 | furthest_point_sampling_kernel<64><<>>( 214 | b, n, m, dataset, temp, idxs); 215 | break; 216 | case 32: 217 | furthest_point_sampling_kernel<32><<>>( 218 | b, n, m, dataset, temp, idxs); 219 | break; 220 | case 16: 221 | furthest_point_sampling_kernel<16><<>>( 222 | b, n, m, dataset, temp, idxs); 223 | break; 224 | case 8: 225 | furthest_point_sampling_kernel<8><<>>( 226 | b, n, m, dataset, temp, idxs); 227 | break; 228 | case 4: 229 | furthest_point_sampling_kernel<4><<>>( 230 | b, n, m, dataset, temp, idxs); 231 | break; 232 | case 2: 233 | furthest_point_sampling_kernel<2><<>>( 234 | b, n, m, dataset, temp, idxs); 235 | break; 236 | case 1: 237 | furthest_point_sampling_kernel<1><<>>( 238 | b, n, m, dataset, temp, idxs); 239 | break; 240 | default: 241 | furthest_point_sampling_kernel<512><<>>( 242 | b, n, m, dataset, temp, idxs); 243 | } 244 | 245 | err = cudaGetLastError(); 246 | if (cudaSuccess != err) { 247 | fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err)); 248 | exit(-1); 249 | } 250 | } 251 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/linalg_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from enum import Enum 3 | 4 | PDist2Order = Enum('PDist2Order', 'd_first d_second') 5 | 6 | 7 | def pdist2( 8 | X: torch.Tensor, 9 | Z: torch.Tensor = None, 10 | order: PDist2Order = PDist2Order.d_second 11 | ) -> torch.Tensor: 12 | r""" Calculates the pairwise distance between X and Z 13 | 14 | D[b, i, j] = l2 distance X[b, i] and Z[b, j] 15 | 16 | Parameters 17 | --------- 18 | X : torch.Tensor 19 | X is a (B, N, d) tensor. There are B batches, and N vectors of dimension d 20 | Z: torch.Tensor 21 | Z is a (B, M, d) tensor. If Z is None, then Z = X 22 | 23 | Returns 24 | ------- 25 | torch.Tensor 26 | Distance matrix is size (B, N, M) 27 | """ 28 | 29 | if order == PDist2Order.d_second: 30 | if X.dim() == 2: 31 | X = X.unsqueeze(0) 32 | if Z is None: 33 | Z = X 34 | G = X @ Z.transpose(-2, -1) 35 | S = (X * X).sum(-1, keepdim=True) 36 | R = S.transpose(-2, -1) 37 | else: 38 | if Z.dim() == 2: 39 | Z = Z.unsqueeze(0) 40 | G = X @ Z.transpose(-2, -1) 41 | S = (X * X).sum(-1, keepdim=True) 42 | R = (Z * Z).sum(-1, keepdim=True).transpose(-2, -1) 43 | else: 44 | if X.dim() == 2: 45 | X = X.unsqueeze(0) 46 | if Z is None: 47 | Z = X 48 | G = X.transpose(-2, -1) @ Z 49 | R = (X * X).sum(-2, keepdim=True) 50 | S = R.transpose(-2, -1) 51 | else: 52 | if Z.dim() == 2: 53 | Z = Z.unsqueeze(0) 54 | G = X.transpose(-2, -1) @ Z 55 | S = (X * X).sum(-2, keepdim=True).transpose(-2, -1) 56 | R = (Z * Z).sum(-2, keepdim=True) 57 | 58 | return torch.abs(R + S - 2 * G).squeeze(0) 59 | 60 | 61 | def pdist2_slow(X, Z=None): 62 | if Z is None: Z = X 63 | D = torch.zeros(X.size(0), X.size(2), Z.size(2)) 64 | 65 | for b in range(D.size(0)): 66 | for i in range(D.size(1)): 67 | for j in range(D.size(2)): 68 | D[b, i, j] = torch.dist(X[b, :, i], Z[b, :, j]) 69 | return D 70 | 71 | 72 | if __name__ == "__main__": 73 | X = torch.randn(2, 3, 5) 74 | Z = torch.randn(2, 3, 3) 75 | 76 | print(pdist2(X, order=PDist2Order.d_first)) 77 | print(pdist2_slow(X)) 78 | print(torch.dist(pdist2(X, order=PDist2Order.d_first), pdist2_slow(X))) 79 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/rscnn_pytorch_utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .pytorch_utils import * 2 | -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/rscnn_pytorch_utils/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/rscnn/utils/rscnn_pytorch_utils/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/backbone/rscnn/utils/rscnn_pytorch_utils/__pycache__/pytorch_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/backbone/rscnn/utils/rscnn_pytorch_utils/__pycache__/pytorch_utils.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/data/IDL_pose_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | FILE_PATH = os.path.abspath(__file__) 4 | FILE_DIR = os.path.dirname(FILE_PATH) 5 | ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(FILE_PATH))) 6 | sys.path.append(ROOT_DIR) 7 | 8 | import pprnet.utils.dataset_util as dataset_util 9 | import numpy as np 10 | import torch 11 | import torch.utils.data as data 12 | 13 | class IDLPoseDataset(data.Dataset): 14 | def __init__(self, data_dir, load_ratio=1.0, mode='train', 15 | transforms=None, scale=1000.0): 16 | self.num_point = 16384 17 | self.transforms = transforms 18 | if mode=='train': 19 | self.dataset, _ = dataset_util.load_dataset( \ 20 | data_dir, load_ratio,\ 21 | load_train_set=True, load_test_set=False) 22 | else: 23 | _, self.dataset = dataset_util.load_dataset( \ 24 | data_dir, load_ratio,\ 25 | load_train_set=False, load_test_set=True) 26 | # convert to mm 27 | self.dataset['data'] *= scale 28 | self.dataset['trans_label'] *= scale 29 | 30 | def __len__(self): 31 | return self.dataset['data'].shape[0] 32 | 33 | def __getitem__(self, idx): 34 | sample = { 35 | 'point_clouds': self.dataset['data'][idx].copy().astype(np.float32), 36 | 'rot_label': self.dataset['rot_label'][idx].copy().astype(np.float32), 37 | 'trans_label':self.dataset['trans_label'][idx].copy().astype(np.float32), 38 | 'cls_label':self.dataset['cls_label'][idx].copy().astype(np.int64), 39 | # 'vis_label':self.dataset['vs_label'][idx].copy().astype(np.float32) 40 | } 41 | 42 | if self.transforms is not None: 43 | sample = self.transforms(sample) 44 | 45 | return sample -------------------------------------------------------------------------------- /pprnet/data/IPA_pose_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | FILE_PATH = os.path.abspath(__file__) 4 | FILE_DIR = os.path.dirname(FILE_PATH) 5 | ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(FILE_PATH))) 6 | # print(FILE_PATH) 7 | # print(ROOT_DIR) 8 | # exit() 9 | sys.path.append(ROOT_DIR) 10 | 11 | import pprnet.utils.dataset_util as dataset_util 12 | import numpy as np 13 | import torch 14 | import torch.utils.data as data 15 | 16 | class IPAPoseDataset(data.Dataset): 17 | def __init__(self, data_dir, cycle_range, scene_range, mode='train', 18 | transforms=None, collect_names=False, collect_error_names=False, scale=1000.0): 19 | self.num_point = 16384 20 | self.transforms = transforms 21 | self.collect_names = collect_names 22 | self.dataset = dataset_util.load_dataset_by_cycle( \ 23 | data_dir, range(cycle_range[0], cycle_range[1]), range(scene_range[0], scene_range[1]),\ 24 | mode, collect_names, collect_error_names) 25 | # convert to mm 26 | self.dataset['data'] *= scale 27 | self.dataset['trans_label'] *= scale 28 | 29 | def __len__(self): 30 | return self.dataset['data'].shape[0] 31 | 32 | def __getitem__(self, idx): 33 | sample = { 34 | 'point_clouds': self.dataset['data'][idx].copy().astype(np.float32), 35 | 'rot_label': self.dataset['rot_label'][idx].copy().astype(np.float32), 36 | 'trans_label':self.dataset['trans_label'][idx].copy().astype(np.float32), 37 | 'cls_label':self.dataset['cls_label'][idx].copy().astype(np.int64), 38 | 'vis_label':self.dataset['vs_label'][idx].copy().astype(np.float32) 39 | } 40 | if self.collect_names: 41 | sample['name'] = self.dataset['name'][idx] 42 | 43 | if self.transforms is not None: 44 | sample = self.transforms(sample) 45 | 46 | return sample 47 | -------------------------------------------------------------------------------- /pprnet/data/__pycache__/IDL_pose_dataset.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/data/__pycache__/IDL_pose_dataset.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/data/__pycache__/IPA_pose_dataset.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/data/__pycache__/IPA_pose_dataset.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/data/__pycache__/IPE_pose_dataset.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/data/__pycache__/IPE_pose_dataset.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/data/__pycache__/pointcloud_transforms.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/data/__pycache__/pointcloud_transforms.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/data/pointcloud_transforms.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | class PointCloudShuffle(object): 5 | def __init__(self, num_point): 6 | self.num_point = num_point 7 | 8 | def __call__(self, sample): 9 | pt_idxs = np.arange(0, self.num_point) 10 | np.random.shuffle(pt_idxs) 11 | 12 | sample['point_clouds'] = sample['point_clouds'][pt_idxs] 13 | sample['rot_label'] = sample['rot_label'][pt_idxs] 14 | sample['trans_label'] = sample['trans_label'][pt_idxs] 15 | sample['cls_label'] = sample['cls_label'][pt_idxs] 16 | if 'vis_label' in sample: 17 | sample['vis_label'] = sample['vis_label'][pt_idxs] 18 | 19 | return sample 20 | 21 | class PointCloudJitter(object): 22 | def __init__(self, scale): 23 | self.scale = scale 24 | 25 | def __call__(self, sample): 26 | all_noise = np.random.standard_normal(sample['point_clouds'].shape) * self.scale 27 | sample['point_clouds'] = sample['point_clouds'] + all_noise 28 | sample['point_clouds'] = sample['point_clouds'].astype(np.float32) 29 | return sample 30 | 31 | 32 | class ToTensor(object): 33 | def __call__(self, sample): 34 | sample['point_clouds'] = torch.from_numpy(sample['point_clouds']) 35 | sample['rot_label'] = torch.from_numpy(sample['rot_label']) 36 | sample['trans_label'] = torch.from_numpy(sample['trans_label']) 37 | sample['cls_label'] = torch.from_numpy(sample['cls_label']) 38 | if 'vis_label' in sample: 39 | sample['vis_label'] = torch.from_numpy(sample['vis_label']) 40 | 41 | return sample 42 | 43 | 44 | 45 | 46 | -------------------------------------------------------------------------------- /pprnet/object_type.py: -------------------------------------------------------------------------------- 1 | """ 2 | Class ObjectType holds properties of each type of object. 3 | 4 | Author: Zhikai Dong 5 | """ 6 | 7 | class ObjectType(): 8 | r""" 9 | ObjectType holds properties of each type of object. 10 | 11 | Init args 12 | ---------- 13 | type_name: str, name of the type 14 | class_idx: int, class label of the type 15 | symmetry_type: str, 'revolution' or 'finite' 16 | kwarg: 17 | if symmetry_type == 'revolution': 18 | lambda_p: scalar 19 | retoreflection: bool 20 | if symmetry_type == 'finite': 21 | lambda_p: List[List[float]] (3, 3) 22 | G: List[ List[List[float]] (3, 3) ], len(G)==K, objects with K equal poses(finite) 23 | """ 24 | def __init__(self, type_name, class_idx, symmetry_type, **kwarg): 25 | assert symmetry_type in ['revolution', 'finite'] 26 | self.type_name = type_name 27 | self.class_idx = class_idx 28 | self.symmetry_type = symmetry_type 29 | if symmetry_type == 'revolution': 30 | self.lambda_p = kwarg['lambda_p'] 31 | self.retoreflection = kwarg['retoreflection'] 32 | elif symmetry_type == 'finite': 33 | self.lambda_p = kwarg['lambda_p'] 34 | self.G = kwarg['G'] 35 | 36 | def get_properties(self): 37 | ''' 38 | get args for building PoseLossCalculator 39 | Returns: 40 | args: dict 41 | ''' 42 | args = { 43 | 'symmetry_type': self.symmetry_type, 44 | 'lambda_p': self.lambda_p 45 | } 46 | if self.symmetry_type == 'revolution': 47 | args['retoreflection'] = self.retoreflection 48 | elif self.symmetry_type == 'finite': 49 | args['G'] = self.G 50 | return args 51 | 52 | -------------------------------------------------------------------------------- /pprnet/utils/__pycache__/dataset_util.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/utils/__pycache__/dataset_util.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/utils/__pycache__/dataset_util_ipe.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/utils/__pycache__/dataset_util_ipe.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/utils/__pycache__/eulerangles.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/utils/__pycache__/eulerangles.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/utils/__pycache__/eval_util.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/utils/__pycache__/eval_util.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/utils/__pycache__/show3d_balls.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/utils/__pycache__/show3d_balls.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/utils/__pycache__/train_helper.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/utils/__pycache__/train_helper.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/utils/__pycache__/visualize_util.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/utils/__pycache__/visualize_util.cpython-36.pyc -------------------------------------------------------------------------------- /pprnet/utils/dataset_util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import random 4 | import h5py 5 | import sys 6 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 7 | ROOT_DIR = BASE_DIR 8 | sys.path.append(BASE_DIR) 9 | sys.path.append(os.path.join(ROOT_DIR, '..')) 10 | import eulerangles 11 | 12 | def initialize_dataset(data_path, split_ratio): 13 | file_names = os.listdir(data_path) 14 | random.shuffle(file_names) 15 | num_file = len(file_names) 16 | train_file = open(os.path.join(data_path, 'train_file_list.txt'), 'w') 17 | test_file = open(os.path.join(data_path, 'test_file_list.txt'), 'w') 18 | for i in range(num_file): 19 | if i < num_file*split_ratio: 20 | train_file.write(file_names[i]+'\n') 21 | else: 22 | test_file.write(file_names[i]+'\n') 23 | train_file.close() 24 | test_file.close() 25 | print('Dataset initialized') 26 | 27 | def load_dataset(data_path, load_ratio=1.0, load_train_set=True, load_test_set=True): 28 | train_dataset = {'data':None, 'trans_label':None, 'rot_label':None, 'cls_label':None} 29 | test_dataset = {'data':None, 'trans_label':None, 'rot_label':None, 'cls_label':None} 30 | 31 | def _load(data_path, mode, load_ratio=1.0): 32 | data_list, trans_label_list, rot_label_list, cls_label_list = [], [], [], [] 33 | num_point_in_h5 = 16384 34 | if mode =='train': 35 | split = 'train_file_list.txt' 36 | if mode =='test': 37 | split = 'test_file_list.txt' 38 | with open(os.path.join(data_path,split),'r') as train_file: 39 | train_names = train_file.readlines() 40 | for i,fn in enumerate(train_names): 41 | if i > len(train_names)*load_ratio: 42 | break 43 | fn = fn.strip() 44 | if fn == '': 45 | continue 46 | h5_filename = os.path.join(data_path, fn) 47 | f = h5py.File(h5_filename) 48 | data_list.append(f['data'][:].reshape(1, num_point_in_h5, 3)) 49 | label = f['labels'][:] 50 | trans_label_list.append(label[:,:3].reshape(1, num_point_in_h5, 3)) 51 | rot_quat = label[:,3:7] 52 | rot_mat = np.zeros([1,num_point_in_h5,3,3]) 53 | for i in range(num_point_in_h5): 54 | rot_mat[0,i] = eulerangles.quat2mat(rot_quat[i]) 55 | rot_label_list.append(rot_mat) 56 | cls_label_list.append(label[:,8].reshape(1, num_point_in_h5)) 57 | dataset={ 'data': np.concatenate(data_list, axis=0), # shape: #scene,#point,3 58 | 'trans_label': np.concatenate(trans_label_list, axis=0), # shape: #scene,#point,3 59 | 'rot_label': np.concatenate(rot_label_list, axis=0), # shape: #scene,#point,3，3 60 | 'cls_label': np.concatenate(cls_label_list, axis=0) # shape: #scene,#point 61 | } 62 | return dataset 63 | 64 | if load_train_set: 65 | train_dataset = _load(data_path, 'train', load_ratio) 66 | print('train data loaded') 67 | if load_test_set: 68 | test_dataset = _load(data_path, 'test', load_ratio) 69 | print('test data loaded') 70 | return train_dataset, test_dataset 71 | 72 | def load_dataset_by_cycle(data_path, cycle_idx_list, scene_idx_list, mode='train', collect_names=False, collect_error_names=False): 73 | num_point_in_h5 = 16384 74 | if mode =='train': 75 | data_list, trans_label_list, rot_label_list, vs_label_list, cls_label_list= [], [], [], [],[] 76 | else: 77 | data_list = [] 78 | if collect_names: 79 | name_list = [] 80 | if collect_error_names: 81 | error_name_list = [] 82 | 83 | for cycle_id in cycle_idx_list: 84 | # print('Loading cycle: %d'%cycle_id) 85 | for scene_id in scene_idx_list: 86 | try: 87 | h5_file_name = os.path.join(data_path, 'cycle_{:0>4}'.format(cycle_id), '{:0>3}.h5'.format(scene_id)) 88 | f = h5py.File(h5_file_name) 89 | data_list.append(f['data'][:].reshape(1, num_point_in_h5, 3)) 90 | if mode == 'train': 91 | label = f['labels'][:] 92 | trans_label_list.append(label[:,:3].reshape(1, num_point_in_h5, 3)) 93 | rot_mat = label[:,3:12].reshape(1, num_point_in_h5, 3, 3) 94 | rot_label_list.append(rot_mat) 95 | vs = label[:, 12].reshape(1, num_point_in_h5) 96 | vs_label_list.append(vs) 97 | cls=label[:,-1].reshape(1, num_point_in_h5) 98 | cls_label_list.append(cls) 99 | if collect_names: 100 | name = 'cycle_{:0>4}'.format(cycle_id) + '_scene_{:0>3}'.format(scene_id) 101 | name_list.append(name) 102 | except: 103 | print('Cycle %d scene %d error, please check' % (cycle_id, scene_id)) 104 | if collect_error_names: 105 | error_name = 'cycle_{:0>4}'.format(cycle_id) + '_scene_{:0>3}'.format(scene_id) 106 | error_name_list.append(error_name) 107 | continue 108 | 109 | 110 | if mode == 'train': 111 | dataset={ 'data': np.concatenate(data_list, axis=0), # shape: #scene,#point,3 112 | 'trans_label': np.concatenate(trans_label_list, axis=0), # shape: #scene,#point,3 113 | 'rot_label': np.concatenate(rot_label_list, axis=0), # shape: #scene,#point,3�? 114 | 'vs_label': np.concatenate(vs_label_list, axis=0), # shape: #scene,#point 115 | 'cls_label':np.concatenate(cls_label_list, axis=0), # shape: #scene,#point 116 | } 117 | else: 118 | dataset={ 'data': np.concatenate(data_list, axis=0), # shape: #scene,#point,3 119 | } 120 | if collect_names: 121 | dataset['name'] = name_list 122 | if collect_error_names: 123 | dataset['error_name'] = error_name_list 124 | return dataset 125 | 126 | if __name__ == "__main__": 127 | # data_path = '/home/idesignlab/Desktop/synthetic_data_generation/cylinder_short_dataset/dataset_result' 128 | # initialize_dataset(data_path, 0.8) 129 | data_path = '/home/idesignlab/dongzhikai/Fraunhofer_IPA_Bin-Picking_dataset/h5_dataset/bunny/test' 130 | dataset = load_dataset_by_cycle(data_path, range(20, 21), range(1, 81), 'test') 131 | print('loaded') 132 | -------------------------------------------------------------------------------- /pprnet/utils/dataset_util_ipe.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import random 4 | import h5py 5 | import sys 6 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 7 | ROOT_DIR = BASE_DIR 8 | sys.path.append(BASE_DIR) 9 | sys.path.append(os.path.join(ROOT_DIR, '..')) 10 | import eulerangles 11 | 12 | def initialize_dataset(data_path, split_ratio): 13 | file_names = os.listdir(data_path) 14 | random.shuffle(file_names) 15 | num_file = len(file_names) 16 | train_file = open(os.path.join(data_path, 'train_file_list.txt'), 'w') 17 | test_file = open(os.path.join(data_path, 'test_file_list.txt'), 'w') 18 | for i in range(num_file): 19 | if i < num_file*split_ratio: 20 | train_file.write(file_names[i]+'\n') 21 | else: 22 | test_file.write(file_names[i]+'\n') 23 | train_file.close() 24 | test_file.close() 25 | print('Dataset initialized') 26 | 27 | def load_dataset(data_path, load_ratio=1.0, load_train_set=True, load_test_set=True): 28 | train_dataset = {'data':None, 'trans_label':None, 'rot_label':None, 'cls_label':None} 29 | test_dataset = {'data':None, 'trans_label':None, 'rot_label':None, 'cls_label':None} 30 | 31 | def _load(data_path, mode, load_ratio=1.0): 32 | data_list, trans_label_list, rot_label_list, cls_label_list = [], [], [], [] 33 | num_point_in_h5 = 16384 34 | if mode =='train': 35 | split = 'train_file_list.txt' 36 | if mode =='test': 37 | split = 'test_file_list.txt' 38 | with open(os.path.join(data_path,split),'r') as train_file: 39 | train_names = train_file.readlines() 40 | for i,fn in enumerate(train_names): 41 | if i > len(train_names)*load_ratio: 42 | break 43 | fn = fn.strip() 44 | if fn == '': 45 | continue 46 | h5_filename = os.path.join(data_path, fn) 47 | f = h5py.File(h5_filename) 48 | data_list.append(f['data'][:].reshape(1, num_point_in_h5, 3)) 49 | label = f['labels'][:] 50 | trans_label_list.append(label[:,:3].reshape(1, num_point_in_h5, 3)) 51 | rot_quat = label[:,3:7] 52 | rot_mat = np.zeros([1,num_point_in_h5,3,3]) 53 | for i in range(num_point_in_h5): 54 | rot_mat[0,i] = eulerangles.quat2mat(rot_quat[i]) 55 | rot_label_list.append(rot_mat) 56 | cls_label_list.append(label[:,8].reshape(1, num_point_in_h5)) 57 | dataset={ 'data': np.concatenate(data_list, axis=0), # shape: #scene,#point,3 58 | 'trans_label': np.concatenate(trans_label_list, axis=0), # shape: #scene,#point,3 59 | 'rot_label': np.concatenate(rot_label_list, axis=0), # shape: #scene,#point,3，3 60 | 'cls_label': np.concatenate(cls_label_list, axis=0) # shape: #scene,#point 61 | } 62 | return dataset 63 | 64 | if load_train_set: 65 | train_dataset = _load(data_path, 'train', load_ratio) 66 | print('train data loaded') 67 | if load_test_set: 68 | test_dataset = _load(data_path, 'test', load_ratio) 69 | print('test data loaded') 70 | return train_dataset, test_dataset 71 | 72 | def load_dataset_by_cycle(data_path, cycle_idx_list, scene_idx_list, mode='train', collect_names=False, collect_error_names=False): 73 | num_point_in_h5 = 16384 74 | if mode =='train': 75 | data_list, trans_label_list, rot_label_list, vs_label_list, cls_label_list= [], [], [], [],[] 76 | else: 77 | data_list = [] 78 | if collect_names: 79 | name_list = [] 80 | if collect_error_names: 81 | error_name_list = [] 82 | 83 | for cycle_id in cycle_idx_list: 84 | # print('Loading cycle: %d'%cycle_id) 85 | for scene_id in scene_idx_list: 86 | try: 87 | h5_file_name = os.path.join(data_path, 'cycle_{:0>4}'.format(cycle_id), '{:0>3}.h5'.format(scene_id)) 88 | f = h5py.File(h5_file_name) 89 | data_list.append(f['data'][:].reshape(1, num_point_in_h5, 3)) 90 | if mode == 'train': 91 | label = f['labels'][:] 92 | trans_label_list.append(label[:,:3].reshape(1, num_point_in_h5, 3)) 93 | rot_mat = label[:,3:12].reshape(1, num_point_in_h5, 3, 3) 94 | rot_label_list.append(rot_mat) 95 | vs = label[:, 13].reshape(1, num_point_in_h5) 96 | vs_label_list.append(vs) 97 | cls=label[:,12].reshape(1, num_point_in_h5) 98 | cls_label_list.append(cls) 99 | if collect_names: 100 | name = 'cycle_{:0>4}'.format(cycle_id) + '_scene_{:0>3}'.format(scene_id) 101 | name_list.append(name) 102 | except: 103 | print('Cycle %d scene %d error, please check' % (cycle_id, scene_id)) 104 | if collect_error_names: 105 | error_name = 'cycle_{:0>4}'.format(cycle_id) + '_scene_{:0>3}'.format(scene_id) 106 | error_name_list.append(error_name) 107 | continue 108 | 109 | 110 | if mode == 'train': 111 | dataset={ 'data': np.concatenate(data_list, axis=0), # shape: #scene,#point,3 112 | 'trans_label': np.concatenate(trans_label_list, axis=0), # shape: #scene,#point,3 113 | 'rot_label': np.concatenate(rot_label_list, axis=0), # shape: #scene,#point,3�? 114 | 'vs_label': np.concatenate(vs_label_list, axis=0), # shape: #scene,#point 115 | 'cls_label':np.concatenate(cls_label_list, axis=0), # shape: #scene,#point 116 | } 117 | else: 118 | dataset={ 'data': np.concatenate(data_list, axis=0), # shape: #scene,#point,3 119 | } 120 | if collect_names: 121 | dataset['name'] = name_list 122 | if collect_error_names: 123 | dataset['error_name'] = error_name_list 124 | return dataset 125 | 126 | if __name__ == "__main__": 127 | # data_path = '/home/idesignlab/Desktop/synthetic_data_generation/cylinder_short_dataset/dataset_result' 128 | # initialize_dataset(data_path, 0.8) 129 | data_path = '/home/idesignlab/dongzhikai/Fraunhofer_IPA_Bin-Picking_dataset/h5_dataset/bunny/test' 130 | dataset = load_dataset_by_cycle(data_path, range(20, 21), range(1, 81), 'test') 131 | print('loaded') 132 | -------------------------------------------------------------------------------- /pprnet/utils/eval_util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import random 4 | import sys 5 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 6 | ROOT_DIR = BASE_DIR 7 | sys.path.append(BASE_DIR) 8 | sys.path.append(os.path.join(ROOT_DIR, '..')) 9 | import eulerangles 10 | 11 | import sklearn 12 | from sklearn.cluster import MeanShift 13 | 14 | def read_pcd(file_name, to_mm=True): 15 | with open(file_name, 'r') as f: 16 | begin = False 17 | points = [] 18 | for line in f.readlines(): 19 | if begin: 20 | xyz = list(map(float, line.strip().split())) 21 | if to_mm: 22 | xyz[:3] = [ 1000*t for t in xyz[:3] ] 23 | points.append(xyz[:3]) 24 | if line.startswith('DATA'): 25 | begin = True 26 | return np.array(points) 27 | 28 | def extract_vertexes_from_obj(file_name): 29 | with open(file_name, 'r') as f: 30 | vertexes = [] 31 | for line in f.readlines(): 32 | line = line.strip() 33 | if line.startswith('v'): 34 | words = line.split()[1:] 35 | xyz = [float(w) for w in words] 36 | vertexes.append(xyz) 37 | ori_model_pc = np.array(vertexes) 38 | center = ( np.max(ori_model_pc, axis=0) + np.min(ori_model_pc, axis=0) ) / 2.0 39 | ori_model_pc = ori_model_pc - center 40 | return ori_model_pc 41 | 42 | def cluster_and_average(input_points,pred_vs, pred_trans, pred_mat, meanshift_kwargs): 43 | # ms = MeanShift(bandwidth=10, bin_seeding=True, cluster_all=False, min_bin_freq=100) 44 | ms = MeanShift(**meanshift_kwargs) 45 | ms.fit(pred_trans) 46 | labels = ms.labels_ 47 | # cluster_centers = ms.cluster_centers_ 48 | 49 | # # Number of clusters in labels, ignoring noise if present. 50 | n_clusters = len(set(labels)) - (1 if -1 in labels else 0) 51 | 52 | pred_trans_cluster = [[] for _ in range(n_clusters)] 53 | pred_mat_cluster = [[] for _ in range(n_clusters)] 54 | pred_vs_cluster = [[] for _ in range(n_clusters)] 55 | for idx in range(input_points.shape[0]): 56 | if labels[idx] != -1: 57 | pred_trans_cluster[labels[idx]].append(np.reshape(pred_trans[idx], [1, 3])) 58 | pred_mat_cluster[labels[idx]].append(np.reshape(pred_mat[idx], [1, 3, 3])) 59 | pred_vs_cluster[labels[idx]].append(pred_vs[idx]) 60 | pred_trans_cluster = [np.concatenate(cluster, axis=0) for cluster in pred_trans_cluster] 61 | pred_mat_cluster = [np.concatenate(cluster, axis=0) for cluster in pred_mat_cluster] 62 | pred_vs_cluster = [ np.mean(l) for l in pred_vs_cluster] 63 | 64 | cluster_center_pred = [ np.mean(cluster, axis=0) for cluster in pred_trans_cluster] 65 | #vs_threshold=0.45 66 | #cluster_center_pred=cluster_center_pred[vs_cls_k>vs_threshold] 67 | #pred_mat_cluster=pred_mat_cluster[vs_cls_k>vs_threshold] 68 | cluster_mat_pred = [] 69 | for mat_cluster in pred_mat_cluster: 70 | all_quat = np.zeros([mat_cluster.shape[0], 4]) 71 | for idx in range(mat_cluster.shape[0]): 72 | all_quat[idx] = eulerangles.mat2quat(mat_cluster[idx]) 73 | quat = eulerangles.average_quat(all_quat) 74 | cluster_mat_pred.append( eulerangles.quat2mat(quat) ) 75 | 76 | pc_segments = [] 77 | for k in range(n_clusters): 78 | ii = np.where(labels==k)[0] 79 | pc_tmp = input_points[ii] 80 | pc_segments.append(pc_tmp) 81 | 82 | return n_clusters, pred_vs_cluster, labels, cluster_center_pred, cluster_mat_pred, pc_segments 83 | -------------------------------------------------------------------------------- /pprnet/utils/point_visualize_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import random 4 | import h5py 5 | import sys 6 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 7 | ROOT_DIR = BASE_DIR 8 | sys.path.append(BASE_DIR) 9 | sys.path.append(os.path.join(ROOT_DIR, '..')) 10 | import show3d_balls 11 | 12 | def show_points(point_array, color_array=None, radius=3): 13 | assert isinstance(point_array, list) 14 | all_color = None 15 | if color_array is not None: 16 | assert len(point_array) == len(color_array) 17 | all_color = [ np.zeros( [ pnts.shape[0] ,3] ) for pnts in point_array] 18 | 19 | for i, c in enumerate(color_array): 20 | c=c.tolist()[0] 21 | all_color[i][:] = [c[1],c[0],c[2]] 22 | all_color = np.concatenate(all_color, axis=0) 23 | all_points = np.concatenate(point_array, axis=0) 24 | show3d_balls.showpoints(all_points, c_gt=all_color, ballradius=radius) 25 | 26 | -------------------------------------------------------------------------------- /pprnet/utils/render_balls_so.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lvwj19/PPR-Net-plus/1295d2313fcc3560eb2e61f37edcdccb7d19f128/pprnet/utils/render_balls_so.so -------------------------------------------------------------------------------- /pprnet/utils/show3d_balls.py: -------------------------------------------------------------------------------- 1 | """ Original Author: Haoqiang Fan """ 2 | import numpy as np 3 | import ctypes as ct 4 | import cv2 5 | import sys 6 | import os 7 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 8 | showsz=800 9 | mousex,mousey=0.5,0.5 10 | zoom=1.0 11 | changed=True 12 | def onmouse(*args): 13 | global mousex,mousey,changed 14 | y=args[1] 15 | x=args[2] 16 | mousex=x/float(showsz) 17 | mousey=y/float(showsz) 18 | changed=True 19 | cv2.namedWindow('show3d') 20 | cv2.moveWindow('show3d',0,0) 21 | cv2.setMouseCallback('show3d',onmouse) 22 | 23 | dll=np.ctypeslib.load_library(os.path.join(BASE_DIR, 'render_balls_so'),'.') 24 | 25 | def showpoints(xyz,c_gt=None, c_pred = None ,waittime=0,showrot=False,magnifyBlue=0,freezerot=True,background=(0,0,0),normalizecolor=True,ballradius=10): 26 | global showsz,mousex,mousey,zoom,changed 27 | xyz=xyz-xyz.mean(axis=0) 28 | radius=((xyz**2).sum(axis=-1)**0.5).max() 29 | xyz/=(radius*2.2)/showsz 30 | if c_gt is None: 31 | c0=np.zeros((len(xyz),),dtype='float32')+255 32 | c1=np.zeros((len(xyz),),dtype='float32')+255 33 | c2=np.zeros((len(xyz),),dtype='float32')+255 34 | else: 35 | c0=c_gt[:,0] 36 | c1=c_gt[:,1] 37 | c2=c_gt[:,2] 38 | 39 | 40 | if normalizecolor: 41 | c0/=(c0.max()+1e-14)/255.0 42 | c1/=(c1.max()+1e-14)/255.0 43 | c2/=(c2.max()+1e-14)/255.0 44 | 45 | 46 | c0=np.require(c0,'float32','C') 47 | c1=np.require(c1,'float32','C') 48 | c2=np.require(c2,'float32','C') 49 | 50 | show=np.zeros((showsz*2,showsz*2,3),dtype='uint8') 51 | print(show.shape) 52 | def render(): 53 | rotmat=np.eye(3) 54 | if not freezerot: 55 | xangle=(mousey-0.5)*np.pi*1.2 56 | else: 57 | xangle=180*np.pi/180.0 58 | rotmat=rotmat.dot(np.array([ 59 | [1.0,0.0,0.0], 60 | [0.0,np.cos(xangle),-np.sin(xangle)], 61 | [0.0,np.sin(xangle),np.cos(xangle)], 62 | ])) 63 | if not freezerot: 64 | yangle=(mousex-0.5)*np.pi*1.2 65 | else: 66 | yangle=0 67 | rotmat=-rotmat.dot(np.array([ 68 | [np.cos(yangle),0.0,-np.sin(yangle)], 69 | [0.0,1.0,0.0], 70 | [np.sin(yangle),0.0,np.cos(yangle)], 71 | ])) 72 | rotmat*=zoom 73 | nxyz=xyz.dot(rotmat)+[showsz/2,showsz/2,0] 74 | 75 | ixyz=nxyz.astype('int32') 76 | show[:]=background 77 | dll.render_ball( 78 | ct.c_int(show.shape[0]), 79 | ct.c_int(show.shape[1]), 80 | show.ctypes.data_as(ct.c_void_p), 81 | ct.c_int(ixyz.shape[0]), 82 | ixyz.ctypes.data_as(ct.c_void_p), 83 | c0.ctypes.data_as(ct.c_void_p), 84 | c1.ctypes.data_as(ct.c_void_p), 85 | c2.ctypes.data_as(ct.c_void_p), 86 | ct.c_int(ballradius) 87 | ) 88 | 89 | if magnifyBlue>0: 90 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],1,axis=0)) 91 | if magnifyBlue>=2: 92 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],-1,axis=0)) 93 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],1,axis=1)) 94 | if magnifyBlue>=2: 95 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],-1,axis=1)) 96 | if showrot: 97 | cv2.putText(show,'xangle %d'%(int(xangle/np.pi*180)),(30,showsz-30),0,0.5,cv2.cv.CV_RGB(255,0,0)) 98 | cv2.putText(show,'yangle %d'%(int(yangle/np.pi*180)),(30,showsz-50),0,0.5,cv2.cv.CV_RGB(255,0,0)) 99 | cv2.putText(show,'zoom %d%%'%(int(zoom*100)),(30,showsz-70),0,0.5,cv2.cv.CV_RGB(255,0,0)) 100 | changed=True 101 | while True: 102 | if changed: 103 | render() 104 | changed=False 105 | cv2.imshow('show3d',show) 106 | if waittime==0: 107 | cmd=cv2.waitKey(10)%256 108 | else: 109 | cmd=cv2.waitKey(waittime)%256 110 | if cmd==ord('q'): 111 | break 112 | elif cmd==ord('Q'): 113 | sys.exit(0) 114 | 115 | if cmd==ord('t') or cmd == ord('p'): 116 | if cmd == ord('t'): 117 | if c_gt is None: 118 | c0=np.zeros((len(xyz),),dtype='float32')+255 119 | c1=np.zeros((len(xyz),),dtype='float32')+255 120 | c2=np.zeros((len(xyz),),dtype='float32')+255 121 | else: 122 | c0=c_gt[:,0] 123 | c1=c_gt[:,1] 124 | c2=c_gt[:,2] 125 | else: 126 | if c_pred is None: 127 | c0=np.zeros((len(xyz),),dtype='float32')+255 128 | c1=np.zeros((len(xyz),),dtype='float32')+255 129 | c2=np.zeros((len(xyz),),dtype='float32')+255 130 | else: 131 | c0=c_pred[:,0] 132 | c1=c_pred[:,1] 133 | c2=c_pred[:,2] 134 | if normalizecolor: 135 | c0/=(c0.max()+1e-14)/255.0 136 | c1/=(c1.max()+1e-14)/255.0 137 | c2/=(c2.max()+1e-14)/255.0 138 | c0=np.require(c0,'float32','C') 139 | c1=np.require(c1,'float32','C') 140 | c2=np.require(c2,'float32','C') 141 | changed = True 142 | 143 | 144 | 145 | if cmd==ord('n'): 146 | zoom*=1.1 147 | changed=True 148 | elif cmd==ord('m'): 149 | zoom/=1.1 150 | changed=True 151 | elif cmd==ord('r'): 152 | zoom=1.0 153 | changed=True 154 | elif cmd==ord('s'): 155 | cv2.imwrite('show3d.png',show) 156 | if waittime!=0: 157 | break 158 | return cmd 159 | if __name__=='__main__': 160 | np.random.seed(100) 161 | showpoints(np.random.randn(2500,3)) 162 | 163 | -------------------------------------------------------------------------------- /pprnet/utils/train_helper.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helper classes and functions for pytorch training. 3 | 4 | Author: Zhikai Dong 5 | """ 6 | 7 | import os 8 | import torch 9 | import torch.nn as nn 10 | from collections import OrderedDict 11 | 12 | def set_bn_momentum_default(bn_momentum): 13 | def fn(m): 14 | if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)): 15 | m.momentum = bn_momentum 16 | return fn 17 | 18 | class BNMomentumScheduler(object): 19 | def __init__( 20 | self, model, bn_lambda, last_epoch=-1, 21 | setter=set_bn_momentum_default 22 | ): 23 | if not isinstance(model, nn.Module): 24 | raise RuntimeError( 25 | "Class '{}' is not a PyTorch nn Module".format( 26 | type(model).__name__ 27 | ) 28 | ) 29 | self.model = model 30 | self.setter = setter 31 | self.lmbd = bn_lambda 32 | self.step(last_epoch + 1) 33 | self.last_epoch = last_epoch 34 | 35 | def step(self, epoch=None): 36 | if epoch is None: 37 | epoch = self.last_epoch + 1 38 | self.last_epoch = epoch 39 | self.model.apply(self.setter(self.lmbd(epoch))) 40 | 41 | def get_bn_momentum(self, epoch): 42 | if epoch is None: 43 | epoch = self.last_epoch 44 | bn_momentum = self.lmbd(epoch) 45 | return bn_momentum 46 | 47 | 48 | class OptimizerLRScheduler(object): 49 | """ 50 | Helper class for steping learning rate. Pytorch lr_scheduler behaves differently before and after version 1.10. 51 | So for compatibility of all pytorch version, we use this class instead. 52 | It is a very simple scheduler which sets all learing rate to a given value. 53 | """ 54 | def __init__(self, optimizer, lr_lambda, last_epoch=-1): 55 | if not isinstance(optimizer, torch.optim.Optimizer): 56 | raise RuntimeError( 57 | "Class '{}' is not a PyTorch torch.optim.Optimizer".format( 58 | type(optimizer).__name__ 59 | ) 60 | ) 61 | self.optimizer = optimizer 62 | self.lr_lambda = lr_lambda 63 | self.step(last_epoch + 1) # initialize lr, this will update self.last_epoch 64 | self.last_epoch = last_epoch # so we need to reset self.last_epoch 65 | 66 | def step(self, epoch=None): 67 | if epoch is None: 68 | epoch = self.last_epoch + 1 69 | self.last_epoch = epoch 70 | for param_group in self.optimizer.param_groups: 71 | param_group['lr'] = self.lr_lambda(epoch) 72 | # for param_group, lr in zip(self.optimizer.param_groups, self.lr_lambda(epoch)): 73 | # param_group['lr'] = lr 74 | 75 | def get_optimizer_lr(self): 76 | lrs = [ g['lr'] for g in self.optimizer.param_groups ] 77 | return lrs 78 | 79 | # def get_lr(self, epoch): 80 | # if epoch is None: 81 | # epoch = self.last_epoch 82 | # lr = self.lr_lambda(epoch) 83 | # return lr 84 | 85 | class SimpleLogger(): 86 | def __init__(self, log_dir, file_path): 87 | if os.path.exists(log_dir): 88 | # raise RuntimeError('Dir:%s alreadly exist! Check if you really want to overwrite it.' % log_dir) 89 | pass 90 | else: 91 | os.makedirs(log_dir) 92 | os.system('cp %s %s' % (file_path, log_dir)) # bkp of train procedure 93 | self.log_file = open(os.path.join(log_dir, 'log_train.txt'), 'w') 94 | self.log_file.write('\n') 95 | self.cnt = 0 96 | self.state_dict = OrderedDict() 97 | 98 | def log_string(self, out_str): 99 | self.log_file.write(out_str+'\n') 100 | self.log_file.flush() 101 | print(out_str) 102 | 103 | def reset_state_dict(self, *args): 104 | self.cnt = 0 105 | self.state_dict = OrderedDict() 106 | for k in args: 107 | assert isinstance(k, str) 108 | self.state_dict[k] = 0.0 109 | 110 | def update_state_dict(self, state_dict): 111 | self.cnt += 1 112 | assert set(state_dict.keys()) == set(self.state_dict.keys()) 113 | for k in state_dict.keys(): 114 | self.state_dict[k] += state_dict[k] 115 | 116 | def print_state_dict(self, log=True, one_line=True, line_len=None): 117 | log_fn = self.log_string if log==True else print 118 | out_str = '' 119 | for i, (k,v) in enumerate(self.state_dict.items()): 120 | out_str += '%s: %f' % (k, 1.0*v/self.cnt) 121 | if i != len(self.state_dict.keys()): 122 | if line_len is not None and (i+1)%line_len==0: 123 | out_str += '\n' 124 | else: 125 | out_str += '\t' if one_line else '\n' 126 | log_fn(out_str) 127 | 128 | -------------------------------------------------------------------------------- /pprnet/utils/visualize_util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import random 4 | import h5py 5 | import sys 6 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 7 | ROOT_DIR = BASE_DIR 8 | sys.path.append(BASE_DIR) 9 | sys.path.append(os.path.join(ROOT_DIR, '..')) 10 | import show3d_balls 11 | 12 | def show_points(point_array, color_array=None, radius=3): 13 | assert isinstance(point_array, list) 14 | all_color = None 15 | if color_array is not None: 16 | if color_array == 'random': 17 | color_array = [np.array([random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)]) for i in range(len(point_array))] 18 | assert len(point_array) == len(color_array) 19 | all_color = [ np.zeros( [ pnts.shape[0] ,3] ) for pnts in point_array] 20 | 21 | for i, c in enumerate(color_array): 22 | c=c.tolist() 23 | all_color[i][:] = [c[1],c[0],c[2]] 24 | all_color = np.concatenate(all_color, axis=0) 25 | all_points = np.concatenate(point_array, axis=0) 26 | show3d_balls.showpoints(all_points, c_gt=all_color, ballradius=radius) 27 | 28 | def show_models(model_pc, trans, rot_mat, cls_idx, color_array=None, radius=3): 29 | assert len(trans) == len(rot_mat) == len(cls_idx) 30 | all_points = [] 31 | all_color = [] if color_array is not None else None 32 | if color_array == 'random': 33 | color_array = [ [random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)] for i in range(len(cls_idx))] 34 | for i in range(len(cls_idx)): 35 | model_pc_transformed = np.dot(model_pc[cls_idx[i]], rot_mat[i].T) + \ 36 | np.tile(np.reshape(trans[i], [1, 3]), [model_pc[cls_idx[i]].shape[0], 1]) 37 | all_points.append(model_pc_transformed) 38 | colors_tmp = np.tile(np.array(color_array[i]).reshape(1,3).astype(np.float32), [model_pc_transformed.shape[0], 1]) 39 | if all_color is not None: 40 | all_color.append(colors_tmp) 41 | 42 | all_points = np.concatenate(all_points, axis=0) 43 | if all_color is not None: 44 | all_color = np.concatenate(all_color, axis=0) 45 | show3d_balls.showpoints(all_points, c_gt=all_color, ballradius=radius) 46 | 47 | def get_models_points(model_pc, trans, rot_mat, cls_idx): 48 | assert len(trans) == len(rot_mat) == len(cls_idx) 49 | all_points = [] 50 | for i in range(len(cls_idx)): 51 | model_pc_transformed = np.dot(model_pc[cls_idx[i]], rot_mat[i].T) + \ 52 | np.tile(np.reshape(trans[i], [1, 3]), [model_pc[cls_idx[i]].shape[0], 1]) 53 | all_points.append(model_pc_transformed) 54 | all_points = np.concatenate(all_points, axis=0) 55 | return all_points 56 | 57 | 58 | --------------------------------------------------------------------------------