├── README.md ├── data_preprocessing ├── README ├── eig_compution_op │ ├── CMakeLists.txt │ ├── cmake.sh │ ├── eig_com.cpp │ ├── eig_com_intensity.cpp │ ├── eig_feature.py │ ├── eig_feature.pyc │ ├── isinblock_gpu.cu │ ├── samlping_gpu.cu │ ├── scene2blocks.cpp │ └── test.py └── prep_stanford │ ├── collect_indoor3d_data.py │ ├── gen_indoor3d_h5.py │ ├── indoor3d_util.py │ ├── indoor3d_util.pyc │ ├── meta │ ├── all_data_label.txt │ ├── anno_paths.txt │ ├── area6_data_label.txt │ └── class_names.txt │ └── prep_scene.py ├── net_S3DIS ├── run.py ├── run.sh ├── run_test.py ├── test.py └── train.py ├── test_utils ├── CMakeLists.txt ├── README ├── cmake.sh ├── interpolate_gpu.cu ├── three_interpolate.cpp ├── three_interpolate.py └── three_interpolate.pyc ├── tf_ops ├── cuda_ulits.cu ├── tf_compile.sh ├── tf_gather.cpp ├── tf_grad_op_test.py ├── tf_interpolate.cpp ├── tf_ops.py ├── tf_ops.pyc ├── tf_ops_.pyc ├── tf_sampling.cpp └── tf_test.py └── utils ├── model.py ├── pcnet_util.py ├── pcnet_util.pyc ├── provider.py ├── provider.pyc ├── tf_util.py └── tf_util.pyc /README.md: -------------------------------------------------------------------------------- 1 | # Graph Attention Convolution for Point Cloud Semantic Segmentation 2 | 3 | This is a Tensorflow implementation of [GACNet](http://openaccess.thecvf.com/content_CVPR_2019/html/Wang_Graph_Attention_Convolution_for_Point_Cloud_Semantic_Segmentation_CVPR_2019_paper.html) for semantic segmentation on [S3DIS dataset](https://shapenet.cs.stanford.edu/media/indoor3d_sem_seg_hdf5_data.zip). 4 | 5 | 6 | Installation 7 | ------ 8 | 9 | The code is tested on Ubuntu 16.04 with Python 2.7 and TF1.12. 10 | 11 | For data processing, [PCL](http://www.pointclouds.org/) is needed for neighbor points searching. 12 | 13 | 14 | 15 | Compile Customized TF Operators 16 | ------- 17 | Most parts here are based on [PointNet++](https://github.com/charlesq34/pointnet2). 18 | 19 | The TF operators are included under tf_ops, you need to compile them first. 20 | 21 | Modify the path of your compiler and run 22 | 23 | cd tf_ops 24 | sh tf_compile.sh 25 | 26 | *Update nvcc and python path if necessary. 27 | 28 | 29 | How to use 30 | ----- 31 | First, you need to prepare your own dataset with the code under the folder data_processing. Slice the input scenes into blocks and down-sampling the points into a certain number, e.g., 4096. 32 | 33 | Here, we also calculate the geometric features in advance as it is slow to put this opteration in the traning phase. 34 | 35 | *[PCL](http://www.pointclouds.org/) is needed for neighbor points searching here. For a prepared dataset for S3DIS, you can download it from [here](https://drive.google.com/drive/folders/1CGY6zY0QvUG4r-DtK4axL972mhImN2bY?usp=sharing). 36 | 37 | 38 | After preparing the dataset, you can run 39 | 40 | cd net_S3DIS 41 | python run.py 42 | python run_test.py 43 | for training and testing on S3DIS. Other/Customized dataset can be done in a similar way. 44 | 45 | 46 | Citation 47 | ----- 48 | If you find our work useful in your research, please consider citing: 49 | 50 | @InProceedings{Wang2019_GACNet, 51 | author = {Wang, Lei and Huang, Yuchun and Hou, Yaolin and Zhang, Shenman and Shan, Jie}, 52 | title = {Graph Attention Convolution for Point Cloud Semantic Segmentation}, 53 | booktitle = {The IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, 54 | month = {June}, 55 | year = {2019} 56 | } 57 | 58 | -------------------------------------------------------------------------------- /data_preprocessing/README: -------------------------------------------------------------------------------- 1 | First, compile the c++ and cuda code for data preprocess using commands 2 | cd eig_compution_op 3 | sh cmake.sh 4 | cd .. 5 | * You may need [PCL](http://www.pointclouds.org/) to compile. 6 | 7 | 8 | Then, modify the data path and prepare the training and testing data on S3DIS with commands 9 | cd prep_stanford 10 | python gen_indoor3d_h5.py 11 | python prep_scene.py -------------------------------------------------------------------------------- /data_preprocessing/eig_compution_op/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | 2 | cmake_minimum_required(VERSION 2.8 FATAL_ERROR) 3 | 4 | project(eig_com) 5 | 6 | find_package(PCL 1.8 REQUIRED) 7 | 8 | include_directories(${PCL_INCLUDE_DIRS} ) 9 | link_directories(${PCL_LIBRARY_DIRS}) 10 | add_definitions(${PCL_DEFINITIONS}) 11 | 12 | MESSAGE(STATUS "This is PCL_INCLUDE_DIRS dir " ${PCL_INCLUDE_DIRS}) 13 | MESSAGE(STATUS "This is PCL_LIBRARY_DIRS dir " ${PCL_LIBRARY_DIRS}) 14 | MESSAGE(STATUS "This is PCL_DEFINITIONS dir " ${PCL_DEFINITIONS}) 15 | MESSAGE(STATUS "This is PCL_LIBRARIES dir " ${PCL_LIBRARIES}) 16 | 17 | FIND_PACKAGE( OpenMP REQUIRED) 18 | if(OPENMP_FOUND) 19 | message("OPENMP FOUND") 20 | set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}") 21 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}") 22 | set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} ${OpenMP_EXE_LINKER_FLAGS}") 23 | endif() 24 | 25 | # packages 26 | find_package(CUDA REQUIRED) 27 | # nvcc flags 28 | set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} -gencode arch=compute_50,code=sm_50) 29 | #set(flannlibs "/usr/lib/x86_64-linux-gnu/libflann.so") 30 | 31 | set(src eig_com.cpp scene2blocks.cpp samlping_gpu.cu) 32 | 33 | CUDA_ADD_LIBRARY(eig_com SHARED ${src}) 34 | target_link_libraries (eig_com ${PCL_LIBRARIES}) 35 | -------------------------------------------------------------------------------- /data_preprocessing/eig_compution_op/cmake.sh: -------------------------------------------------------------------------------- 1 | rm -rf build 2 | mkdir build 3 | cd build 4 | cmake .. 5 | make 6 | -------------------------------------------------------------------------------- /data_preprocessing/eig_compution_op/eig_com.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | #include 18 | #include 19 | #include 20 | #include 21 | 22 | using namespace std; 23 | 24 | struct eigen_index{ 25 | float eigen_value; 26 | int index; 27 | }; 28 | bool my_decrease_sort(const eigen_index &first, const eigen_index &second){ 29 | return first.eigen_value > second.eigen_value; 30 | } 31 | 32 | const int eig_num = 6; 33 | 34 | //////////////////////////////////////////////////////////////////////////////////// 35 | void data2cloud(float *dataset, int num, pcl::PointCloud::Ptr Cloud){ 36 | pcl::PointXYZ tmpPoint; 37 | for(int i=0;ipoints.push_back(tmpPoint); 42 | } 43 | } 44 | 45 | //////////////////////////////////////////////////////////////////////////////////// 46 | void compute_eiginf(std::vector pointIdxRadiusSearch, pcl::PointCloud::Ptr Cloud, float *eigen_inf){ 47 | //matrix and cloud load 48 | Eigen::MatrixXf observe = Eigen::MatrixXf::Random(pointIdxRadiusSearch.size(), 3); 49 | for (size_t ri = 0; ri < pointIdxRadiusSearch.size(); ++ri){ 50 | pcl::PointXYZ tmp_point; 51 | observe(ri, 0) = Cloud->points[pointIdxRadiusSearch[ri]].x; 52 | observe(ri, 1) = Cloud->points[pointIdxRadiusSearch[ri]].y; 53 | observe(ri, 2) = Cloud->points[pointIdxRadiusSearch[ri]].z; 54 | } 55 | //caculate eigen value, eigen vector 56 | Eigen::MatrixXf centered = observe.rowwise() - observe.colwise().mean(); 57 | Eigen::Matrix3f covMat = (centered.adjoint() * centered) / double(observe.rows() - 1); 58 | Eigen::EigenSolver es(covMat); 59 | Eigen::MatrixXf D = es.pseudoEigenvalueMatrix(); 60 | Eigen::MatrixXf V = es.pseudoEigenvectors(); 61 | //sort eigens from big to small 62 | std::vector eigen_indexs; 63 | eigen_index tmp_eigen_index; 64 | for (int i = 0; i < 3; i++){ 65 | tmp_eigen_index.eigen_value = D(i, i); 66 | tmp_eigen_index.index = i; 67 | eigen_indexs.push_back(tmp_eigen_index); 68 | } 69 | sort(eigen_indexs.begin(), eigen_indexs.end(), my_decrease_sort); 70 | 71 | //eigen normlization 72 | float sum_deno = max(sqrt(D(0, 0)*D(0, 0) + D(1, 1)*D(1, 1) + D(2, 2)*D(2, 2)), FLT_MIN); 73 | //float nor_eig[3] = { 0 }; 74 | //normlized eigen values 75 | eigen_inf[0] = max(fabs(D(eigen_indexs[0].index, eigen_indexs[0].index) / sum_deno), FLT_MIN); 76 | eigen_inf[1] = max(fabs(D(eigen_indexs[1].index, eigen_indexs[1].index) / sum_deno), FLT_MIN); 77 | eigen_inf[2] = max(fabs(D(eigen_indexs[2].index, eigen_indexs[2].index) / sum_deno), FLT_MIN); 78 | //norm vector 79 | eigen_inf[3] = fabs(V(0, eigen_indexs[2].index)); 80 | eigen_inf[4] = fabs(V(1, eigen_indexs[2].index)); 81 | eigen_inf[5] = fabs(V(2, eigen_indexs[2].index)); 82 | //printf("%f", eigen_inf[0]); 83 | } 84 | 85 | //////////////////////////////////////////////////////////////////////////////////// 86 | extern "C" void compute_eigen(float *nearpoints, int &point_num, float *seachpoints, int &search_num, float &radius, int &k_num, float *eigen_inf){ 87 | //matrix and cloud load 88 | 89 | pcl::PointCloud::Ptr Cloud(new pcl::PointCloud); 90 | data2cloud(nearpoints, point_num, Cloud); 91 | 92 | //main program 93 | pcl::KdTreeFLANN kdtree; 94 | kdtree.setInputCloud(Cloud); 95 | 96 | #pragma omp parallel for 97 | for (int i=0;i pointIdxRadiusSearch(k_num); 100 | std::vector pointRadiusSquaredDistance(k_num); 101 | float tmp[eig_num] ={0}; 102 | 103 | searchPoint.x = Cloud->points[i].x; 104 | searchPoint.y = Cloud->points[i].y; 105 | searchPoint.z = Cloud->points[i].z; 106 | 107 | /* 108 | if (kdtree.nearestKSearch(searchPoint, k_num, pointIdxNKNSearch, pointNKNSquaredDistance)>3){//if the neighbors are less than 3, we cannot use them for norm calculation 109 | compute_eiginf(pointIdxNKNSearch, Cloud, tmp); 110 | }else{ 111 | cout<<"KNN search failed!"< k_num){ 116 | compute_eiginf(pointIdxRadiusSearch, Cloud, tmp); 117 | } 118 | else{ 119 | kdtree.nearestKSearch(searchPoint, k_num, pointIdxRadiusSearch, pointRadiusSquaredDistance); 120 | compute_eiginf(pointIdxRadiusSearch, Cloud, tmp); 121 | } 122 | 123 | for(int j=0;j mapeiginf; 132 | //#pragma omp parallel for 133 | for(int i=0;i pointIdxRadiusSearch, pcl::PointCloud::Ptr Cloud, float *intensity, float *eigen_inf){ 167 | //matrix and cloud load 168 | Eigen::MatrixXf observe = Eigen::MatrixXf::Random(pointIdxRadiusSearch.size(), 3); 169 | float *tmp_intensity = new float[pointIdxRadiusSearch.size()]; 170 | float sum = 0; 171 | for (size_t ri = 0; ri < pointIdxRadiusSearch.size(); ++ri){ 172 | pcl::PointXYZ tmp_point; 173 | observe(ri, 0) = Cloud->points[pointIdxRadiusSearch[ri]].x; 174 | observe(ri, 1) = Cloud->points[pointIdxRadiusSearch[ri]].y; 175 | observe(ri, 2) = Cloud->points[pointIdxRadiusSearch[ri]].z; 176 | tmp_intensity[ri] = intensity[pointIdxRadiusSearch[ri]]; 177 | sum += tmp_intensity[ri]; 178 | } 179 | float intensity_mean = sum/ (pointIdxRadiusSearch.size()+1e-8); 180 | sum = 0; 181 | for (size_t i = 0; i < pointIdxRadiusSearch.size(); ++i){ 182 | sum += (tmp_intensity[i]-intensity_mean) * (tmp_intensity[i]-intensity_mean) ; 183 | } 184 | float intensity_var = sum/ (pointIdxRadiusSearch.size()+1e-8); 185 | 186 | //caculate eigen value, eigen vector 187 | Eigen::MatrixXf centered = observe.rowwise() - observe.colwise().mean(); 188 | Eigen::Matrix3f covMat = (centered.adjoint() * centered) / double(observe.rows() - 1); 189 | Eigen::EigenSolver es(covMat); 190 | Eigen::MatrixXf D = es.pseudoEigenvalueMatrix(); 191 | Eigen::MatrixXf V = es.pseudoEigenvectors(); 192 | //sort eigens from big to small 193 | std::vector eigen_indexs; 194 | eigen_index tmp_eigen_index; 195 | for (int i = 0; i < 3; i++){ 196 | tmp_eigen_index.eigen_value = D(i, i); 197 | tmp_eigen_index.index = i; 198 | eigen_indexs.push_back(tmp_eigen_index); 199 | } 200 | sort(eigen_indexs.begin(), eigen_indexs.end(), my_decrease_sort); 201 | 202 | //eigen normlization 203 | double sum_deno = sqrt(D(0, 0)*D(0, 0) + D(1, 1)*D(1, 1) + D(2, 2)*D(2, 2)); 204 | eigen_inf[0] = fabs(D(eigen_indexs[0].index, eigen_indexs[0].index) / (sum_deno + 0.00000001)); 205 | eigen_inf[1] = fabs(D(eigen_indexs[1].index, eigen_indexs[1].index) / (sum_deno + 0.00000001)); 206 | eigen_inf[2] = fabs(D(eigen_indexs[2].index, eigen_indexs[2].index) / (sum_deno + 0.00000001)); 207 | eigen_inf[3] = intensity_var; 208 | } 209 | 210 | //////////////////////////////////////////////////////////////////////////////////// 211 | extern "C" void compute_eigen_intensity(float *nearpoints, int &point_num, float *intensity, float *seachpoints, int &search_num, float &radius, int &k_num, float *eigen_inf){ 212 | //matrix and cloud load 213 | 214 | pcl::PointCloud::Ptr Cloud(new pcl::PointCloud); 215 | data2cloud(nearpoints, point_num, Cloud); 216 | 217 | //main program 218 | pcl::KdTreeFLANN kdtree; 219 | kdtree.setInputCloud(Cloud); 220 | 221 | #pragma omp parallel for 222 | for (int i=0;i pointIdxRadiusSearch; 225 | std::vector pointRadiusSquaredDistance; 226 | float tmp[eig_intensity_num] ={0}; 227 | 228 | searchPoint.x = seachpoints[i*3]; 229 | searchPoint.y = seachpoints[i*3+1]; 230 | searchPoint.z = seachpoints[i*3+2]; 231 | 232 | if (kdtree.radiusSearch(searchPoint, radius, pointIdxRadiusSearch, pointRadiusSquaredDistance) > k_num){ 233 | compute_eiginf_intensity(pointIdxRadiusSearch, Cloud, intensity, tmp); 234 | } 235 | else{ 236 | kdtree.nearestKSearch(searchPoint, k_num, pointIdxRadiusSearch, pointRadiusSquaredDistance); 237 | compute_eiginf_intensity(pointIdxRadiusSearch, Cloud, intensity, tmp); 238 | } 239 | 240 | for(int j=0;j mapeiginf; 249 | //#pragma omp parallel for 250 | for(int i=0;i 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | #include 18 | #include 19 | #include 20 | #include 21 | 22 | using namespace std; 23 | 24 | struct eigen_index{ 25 | float eigen_value; 26 | int index; 27 | }; 28 | bool my_decrease_sort(const eigen_index &first, const eigen_index &second){ 29 | return first.eigen_value > second.eigen_value; 30 | } 31 | 32 | const int eig_num = 4; 33 | 34 | 35 | //////////////////////////////////////////////////////////////////////////////////// 36 | void data2cloud(float *dataset, int num, pcl::PointCloud::Ptr Cloud){ 37 | pcl::PointXYZ tmpPoint; 38 | for(int i=0;ipoints.push_back(tmpPoint); 43 | } 44 | } 45 | 46 | 47 | //////////////////////////////////////////////////////////////////////////////////// 48 | void compute_eiginf_intensity(std::vector pointIdxRadiusSearch, pcl::PointCloud::Ptr Cloud, float *intensity, float *eigen_inf){ 49 | //matrix and cloud load 50 | Eigen::MatrixXf observe = Eigen::MatrixXf::Random(pointIdxRadiusSearch.size(), 3); 51 | float *tmp_intensity = new float[pointIdxRadiusSearch.size()]; 52 | float sum = 0; 53 | for (size_t ri = 0; ri < pointIdxRadiusSearch.size(); ++ri){ 54 | pcl::PointXYZ tmp_point; 55 | observe(ri, 0) = Cloud->points[pointIdxRadiusSearch[ri]].x; 56 | observe(ri, 1) = Cloud->points[pointIdxRadiusSearch[ri]].y; 57 | observe(ri, 2) = Cloud->points[pointIdxRadiusSearch[ri]].z; 58 | tmp_intensity[ri] = intensity[pointIdxRadiusSearch[ri]]; 59 | sum += tmp_intensity[ri]; 60 | } 61 | float intensity_mean = sum/ (pointIdxRadiusSearch.size()+1e-8); 62 | sum = 0; 63 | for (size_t i = 0; i < pointIdxRadiusSearch.size(); ++i){ 64 | sum += (tmp_intensity[i]-intensity_mean) * (tmp_intensity[i]-intensity_mean) ; 65 | } 66 | float intensity_var = sum/ (pointIdxRadiusSearch.size()+1e-8); 67 | 68 | //caculate eigen value, eigen vector 69 | Eigen::MatrixXf centered = observe.rowwise() - observe.colwise().mean(); 70 | Eigen::Matrix3f covMat = (centered.adjoint() * centered) / double(observe.rows() - 1); 71 | Eigen::EigenSolver es(covMat); 72 | Eigen::MatrixXf D = es.pseudoEigenvalueMatrix(); 73 | Eigen::MatrixXf V = es.pseudoEigenvectors(); 74 | //sort eigens from big to small 75 | std::vector eigen_indexs; 76 | eigen_index tmp_eigen_index; 77 | for (int i = 0; i < 3; i++){ 78 | tmp_eigen_index.eigen_value = D(i, i); 79 | tmp_eigen_index.index = i; 80 | eigen_indexs.push_back(tmp_eigen_index); 81 | } 82 | sort(eigen_indexs.begin(), eigen_indexs.end(), my_decrease_sort); 83 | 84 | //eigen normlization 85 | double sum_deno = sqrt(D(0, 0)*D(0, 0) + D(1, 1)*D(1, 1) + D(2, 2)*D(2, 2)); 86 | eigen_inf[0] = fabs(D(eigen_indexs[0].index, eigen_indexs[0].index) / (sum_deno + 0.00000001)); 87 | eigen_inf[1] = fabs(D(eigen_indexs[1].index, eigen_indexs[1].index) / (sum_deno + 0.00000001)); 88 | eigen_inf[2] = fabs(D(eigen_indexs[2].index, eigen_indexs[2].index) / (sum_deno + 0.00000001)); 89 | eigen_inf[3] = intensity_var; 90 | } 91 | 92 | //////////////////////////////////////////////////////////////////////////////////// 93 | extern "C" void compute_eigen_intensity(float *nearpoints, int &point_num, float *intensity, float *seachpoints, int &search_num, float &radius, int &k_num, float *eigen_inf){ 94 | //matrix and cloud load 95 | 96 | pcl::PointCloud::Ptr Cloud(new pcl::PointCloud); 97 | data2cloud(nearpoints, point_num, Cloud); 98 | 99 | //main program 100 | pcl::KdTreeFLANN kdtree; 101 | kdtree.setInputCloud(Cloud); 102 | 103 | #pragma omp parallel for 104 | for (int i=0;i pointIdxRadiusSearch; 107 | std::vector pointRadiusSquaredDistance; 108 | float tmp[eig_num] ={0}; 109 | 110 | searchPoint.x = seachpoints[i*3]; 111 | searchPoint.y = seachpoints[i*3+1]; 112 | searchPoint.z = seachpoints[i*3+2]; 113 | 114 | if (kdtree.radiusSearch(searchPoint, radius, pointIdxRadiusSearch, pointRadiusSquaredDistance) > k_num){ 115 | compute_eiginf_intensity(pointIdxRadiusSearch, Cloud, intensity, tmp); 116 | } 117 | else{ 118 | kdtree.nearestKSearch(searchPoint, k_num, pointIdxRadiusSearch, pointRadiusSquaredDistance); 119 | compute_eiginf_intensity(pointIdxRadiusSearch, Cloud, intensity, tmp); 120 | } 121 | 122 | for(int j=0;j mapeiginf; 131 | //#pragma omp parallel for 132 | for(int i=0;ix_low && xy_low && y>>(xy, x_low, y_low, size, is_block, N); 36 | 37 | error = cudaDeviceSynchronize(); 38 | if(error != cudaSuccess){ 39 | printf("code: %d, reason: %s\n",error,cudaGetErrorString(error)); 40 | } 41 | 42 | cudaMemcpy(is_block_host, is_block, sizeof(float)*N, cudaMemcpyDeviceToHost); 43 | cudaFree(xy); 44 | cudaFree(is_block); 45 | } 46 | 47 | -------------------------------------------------------------------------------- /data_preprocessing/eig_compution_op/samlping_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | /* we need these includes for CUDA's random number stuff */ 4 | #include 5 | #include 6 | // sample m points from n points 7 | // input (n,3) 8 | // output (m,3) 9 | __global__ void farthestpointsamplingKernel(int n,int m, const int random_init, float *temp, const float *dataset, int *idxs){ 10 | 11 | const int BlockSize=512; 12 | __shared__ float dists[BlockSize]; 13 | __shared__ int dists_i[BlockSize]; 14 | 15 | int old_id = random_init; 16 | idxs[0] = old_id; 17 | 18 | for (int j=threadIdx.x;jbest){ 42 | best=d2; 43 | besti=j; 44 | } 45 | } 46 | 47 | dists[threadIdx.x]=best; 48 | dists_i[threadIdx.x]=besti; 49 | for (int u=0;(1<>(u+1))){ 52 | int i1=(threadIdx.x*2)<>>(n, m, random_init, temp, dataset, idxs); 81 | 82 | error = cudaDeviceSynchronize(); 83 | if(error != cudaSuccess){ 84 | printf("code: %d, reason: %s\n",error,cudaGetErrorString(error)); 85 | } 86 | 87 | cudaMemcpy(idxs_host, idxs, sizeof(int)*m, cudaMemcpyDeviceToHost); 88 | cudaFree(temp); 89 | cudaFree(dataset); 90 | cudaFree(idxs); 91 | } 92 | 93 | 94 | __global__ void fps_multiblocks_Kernel(int num_sample, int num_block, const int *num_pts, const int *start_id, const int *sort_id, const int *seed, float *temp, const float *dataset, int *idxs){ 95 | 96 | const int BlockSize=512; 97 | __shared__ float dists[BlockSize]; 98 | __shared__ int dists_i[BlockSize]; 99 | 100 | 101 | for (int i=blockIdx.x;i=num_pts[i]){ 111 | curandState_t state; 112 | curand_init(0, 0, 0,&state); 113 | tid = curand(&state) % num_pts[i]; 114 | } 115 | idxs[i*num_sample+j]=sort_id[ start_id[i]+tid ]; 116 | } 117 | 118 | } 119 | else if(num_pts[i]==num_sample ){ 120 | for (int j=threadIdx.x;jbest){ 149 | best=d2; 150 | besti=k; 151 | } 152 | } 153 | dists[threadIdx.x]=best; 154 | dists_i[threadIdx.x]=besti; 155 | for (int u=0;(1<>(u+1))){ 158 | int i1=(threadIdx.x*2)<>>(time(0), states); 225 | randoms<<>>(states, seed, num_pts); 226 | 227 | 228 | fps_multiblocks_Kernel<<<64, 512>>>(num_sample, num_block, num_pts, start_id, sort_id,seed, temp, dataset, idxs); 229 | 230 | //int batchsize=16; 231 | //int num_batch=num_block/batchsize+1; 232 | //for (int i=0;i>>(i, num_sample, num_block, num_pts, start_id, sort_id,seed, temp, dataset, idxs); 234 | //} 235 | error = cudaDeviceSynchronize(); 236 | if(error != cudaSuccess){ 237 | printf("code: %d, reason: %s\n",error,cudaGetErrorString(error)); 238 | } 239 | /* 240 | printf("num_block: %d\n", num_block); 241 | print<<<1,1>>>(start_id,num_block); 242 | printf("num_pts: \n"); 243 | print<<<1,1>>>(num_pts,num_block); 244 | */ 245 | 246 | cudaMemcpy(idxs_host, idxs, sizeof(int)*num_sample*num_block, cudaMemcpyDeviceToHost); 247 | 248 | cudaFree(states); 249 | cudaFree(temp); 250 | cudaFree(dataset); 251 | cudaFree(idxs); 252 | cudaFree(start_id); 253 | cudaFree(num_pts); 254 | cudaFree(seed); 255 | cudaFree(sort_id); 256 | } 257 | 258 | 259 | 260 | 261 | -------------------------------------------------------------------------------- /data_preprocessing/eig_compution_op/scene2blocks.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | 15 | #define Randmod(x) rand()%x 16 | using namespace std; 17 | 18 | 19 | struct pts_idx{ 20 | int idx; 21 | int block_id; 22 | float inner; 23 | }; 24 | bool increase_pts_idx(const pts_idx &first, const pts_idx &second){ 25 | return first.block_id < second.block_id; 26 | } 27 | 28 | float is_in_box(float x, float y, float x_center, float y_center, float radius){ 29 | if( fabs(x-x_center)<=radius && fabs(y-y_center)<=radius) 30 | return 1.0; 31 | else 32 | return 0.0; 33 | } 34 | 35 | struct num_idx{ 36 | int num; 37 | int start_idx; 38 | }; 39 | 40 | 41 | void randomsamlping(vector _pts_idx, vector _num_idx, int MAX_NUM, int *sample_idx){ 42 | int block_num = _num_idx.size(); 43 | //cout< _pts_idx, vector _num_idx, float* xyz,int MAX_NUM, int *sample_idx){ 58 | int totalnum_pts = _pts_idx.size(); 59 | int num_block = _num_idx.size(); 60 | int *num_pts = new int[num_block]; 61 | int *start_id = new int[num_block]; 62 | #pragma omp parallel for 63 | for(int i=0;i _pts_idx; 93 | pts_idx tmp_pts_idx; 94 | for (int i=0;i _num_idx; 106 | num_idx tmp_num_idx; 107 | int tmp_num=1; 108 | tmp_num_idx.start_idx = 0; 109 | for(int i=0;i::iterator it; 125 | for (it=_num_idx.begin();it!=_num_idx.end();){ 126 | if(it->num < count_threshold) 127 | it = _num_idx.erase(it); 128 | else 129 | ++it; 130 | } 131 | block_num = _num_idx.size(); 132 | // 133 | fps_multiblocks_warpper( _pts_idx, _num_idx, pts_xyz, MAX_NUM, sample_idx); 134 | 135 | //randomsamlping(_pts_idx, _num_idx, MAX_NUM, sample_idx); 136 | } 137 | 138 | 139 | /////////////////////////////////////////////////////////////////////////////////////////////////////////////////// 140 | extern "C" void scene2blocks_withinner(float *pts_xyz, int *sample_idx, float *inner, int &MAX_NUM, int &pts_num, int &x_num, float &block_size, float &inner_radius, int &count_threshold, int &block_num){ 141 | 142 | vector _pts_idx; 143 | pts_idx tmp_pts_idx; 144 | int id_x, id_y; 145 | float x_center, y_center; 146 | for (int i=0;i _pts_idx_copy(_pts_idx); 158 | sort(_pts_idx.begin(), _pts_idx.end(), increase_pts_idx); 159 | 160 | //calculate the nummber of points in each block and their start id 161 | vector _num_idx; 162 | num_idx tmp_num_idx; 163 | int tmp_num=0; 164 | tmp_num_idx.start_idx = 0; 165 | for(int i=0;i::iterator it; 181 | for (it=_num_idx.begin();it!=_num_idx.end();){ 182 | if(it->num < count_threshold) 183 | it = _num_idx.erase(it); 184 | else 185 | ++it; 186 | } 187 | 188 | // 189 | block_num = _num_idx.size(); 190 | //cout<=0) 56 | if capacity > 0: 57 | h5_batch_data[buffer_size:buffer_size+capacity, ...] = data[0:capacity, ...] 58 | h5_batch_label[buffer_size:buffer_size+capacity, ...] = label[0:capacity, ...] 59 | h5_batch_spw[buffer_size:buffer_size+capacity, ...] = spw[0:capacity, ...] 60 | # Save batch data and label to h5 file, reset buffer_size 61 | h5_filename = output_filename_prefix + '_' + str(h5_index) + '.h5' 62 | indoor3d_util.save_h5(h5_filename, h5_batch_data, h5_batch_label, h5_batch_spw, data_dtype, label_dtype, spw_dtype) 63 | print('Stored {0} with size {1}'.format(h5_filename, h5_batch_data.shape[0])) 64 | h5_index += 1 65 | buffer_size = 0 66 | # recursive call 67 | insert_batch(data[capacity:, ...], label[capacity:, ...], spw[capacity:, ...], last_batch) 68 | if last_batch and buffer_size > 0: 69 | h5_filename = output_filename_prefix + '_' + str(h5_index) + '.h5' 70 | indoor3d_util.save_h5(h5_filename, h5_batch_data[0:buffer_size, ...], h5_batch_label[0:buffer_size, ...], h5_batch_spw[0:buffer_size, ...], data_dtype, label_dtype, spw_dtype) 71 | print('Stored {0} with size {1}'.format(h5_filename, buffer_size)) 72 | h5_index += 1 73 | buffer_size = 0 74 | return 75 | 76 | 77 | sample_cnt = 0 78 | room_filelist = [] 79 | for i, data_label_filename in enumerate(data_label_files): 80 | print(data_label_filename) 81 | data, label, spw = indoor3d_util.scene2blocks_zeromean(data_label_filename, NUM_POINT, block_size=1.2, inner_size=1.0, stride=0.6, radius=0.1) 82 | print('{0}, {1}'.format(data.shape, label.shape)) 83 | for _ in range(data.shape[0]): 84 | fout_room.write(os.path.basename(data_label_filename)[0:-4]+'\n') 85 | 86 | sample_cnt += data.shape[0] 87 | insert_batch(data, label, spw, i == len(data_label_files)-1) 88 | 89 | fout_room.close() 90 | print("Total samples: {0}".format(sample_cnt)) 91 | -------------------------------------------------------------------------------- /data_preprocessing/prep_stanford/indoor3d_util.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import glob 3 | import os 4 | import sys 5 | import pdb 6 | import h5py 7 | import time 8 | from plyfile import PlyData, PlyElement 9 | 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | ROOT_DIR = os.path.dirname(BASE_DIR) 12 | sys.path.append(os.path.join(ROOT_DIR, 'eig_computation_op')) 13 | from eig_feature import eig_cal_map, scene2blocks_withinner 14 | 15 | # ----------------------------------------------------------------------------- 16 | # PREPARE BLOCK DATA FOR DEEPNETS TRAINING/TESTING 17 | # ----------------------------------------------------------------------------- 18 | 19 | def scene2blocks(data, label, num_point, block_size=1.0, inner_size=0.5, stride=0.5, count_threshold=100): 20 | 21 | assert(stride<=block_size) 22 | data[:,0:2] -= np.amin(data, 0)[0:2] 23 | pts_xyz = data[:,0:3] 24 | sample_idx, is_valid, block_num = scene2blocks_withinner(pts_xyz, num_point, block_size, inner_size, count_threshold) 25 | 26 | #horizen direction 27 | pts_xyz[:,0] += stride 28 | sample_idx_h, is_valid_h, block_num_h = scene2blocks_withinner(pts_xyz, num_point, block_size, inner_size, count_threshold) 29 | 30 | #diag direction 31 | pts_xyz[:,1] += stride 32 | sample_idx_d, is_valid_d, block_num_d = scene2blocks_withinner(pts_xyz, num_point, block_size, inner_size, count_threshold) 33 | 34 | #vihical direction 35 | pts_xyz[:,0] -= stride 36 | sample_idx_v, is_valid_v, block_num_v = scene2blocks_withinner(pts_xyz, num_point, block_size, inner_size, count_threshold) 37 | 38 | sample_idx = np.concatenate([sample_idx, sample_idx_h, sample_idx_v, sample_idx_d], 0) 39 | is_valid = np.concatenate([is_valid, is_valid_h, is_valid_d, is_valid_v], 0) 40 | block_num = block_num+block_num_h+block_num_d+block_num_v 41 | 42 | data_batch = data[sample_idx, ...] 43 | label_batch = label[sample_idx, ...] 44 | 45 | data_batch = np.reshape(data_batch, [block_num, num_point, -1]) 46 | label_batch = np.reshape(label_batch, [block_num, num_point, -1]) 47 | is_valid = np.reshape(is_valid, [block_num, num_point]) 48 | 49 | return data_batch, np.squeeze(label_batch), is_valid 50 | 51 | 52 | def scene2blocks_zeromean(data_label_filename, num_point, block_size=1.0, inner_size=0.5, stride=0.5, radius=0.15): 53 | """ room2block, with input filename and RGB preprocessing. 54 | for each block centralize XY, z channel is keeped 55 | RGB to [0,1] 56 | """ 57 | 58 | if data_label_filename[-3:] == 'txt': 59 | data_label = np.loadtxt(data_label_filename) 60 | elif data_label_filename[-3:] == 'ply': 61 | data_label = read_xyzrgbL_ply(data_label_filename) 62 | else: 63 | print('Unknown file type! exiting.') 64 | exit() 65 | 66 | print('data load done') 67 | data = data_label[:,0:6] 68 | data[:,3:6] /= 255.0 69 | label = data_label[:,-1].astype(np.uint8) 70 | idx = np.array(range(data.shape[0])) 71 | lalel_id = np.stack((label, idx), axis=-1) 72 | 73 | time_start = time.time() 74 | data_batch, label_id_batch, mask_batch = scene2blocks(data, lalel_id, num_point, block_size, inner_size, stride) 75 | #pdb.set_trace() 76 | batch_num = data_batch.shape[0] 77 | #zero mean 78 | for b in range(batch_num): 79 | mean_xy = np.mean(data_batch[b, :, 0:2], 0) 80 | data_batch[b, :, 0] -= mean_xy[0] 81 | data_batch[b, :, 1] -= mean_xy[1] 82 | #pdb.set_trace() 83 | #eigenvalues calculaition 84 | print('eigenvalues calculaiting...') 85 | idx = np.reshape(label_id_batch[:,:,-1], [-1]) 86 | unique_idx = np.unique(idx) 87 | eigs_batch = eig_cal_map(data[:,0:3], unique_idx, idx, radius) 88 | eigs_batch = np.reshape(eigs_batch, [batch_num, num_point, -1]) 89 | data_batch = np.concatenate([data_batch, eigs_batch], -1) 90 | 91 | label_batch = label_id_batch[:,:,0] 92 | #pdb.set_trace() 93 | print('time for scene2block: %ds' %(time.time()-time_start)) 94 | return data_batch, label_batch, mask_batch 95 | 96 | 97 | # Write numpy array data and label to h5_filename 98 | def save_h5(h5_filename, data, label, spw, data_dtype='float32', label_dtype='uint8', spw_dtype='float32'): 99 | h5_fout = h5py.File(h5_filename) 100 | h5_fout.create_dataset( 101 | 'data', data=data, 102 | compression='gzip', compression_opts=4, 103 | dtype=data_dtype) 104 | h5_fout.create_dataset( 105 | 'label', data=label, 106 | compression='gzip', compression_opts=1, 107 | dtype=label_dtype) 108 | h5_fout.create_dataset( 109 | 'spw', data=spw, 110 | compression='gzip', compression_opts=4, 111 | dtype=spw_dtype) 112 | h5_fout.close() 113 | 114 | 115 | def read_xyzrgbL_ply(filename): 116 | """ read XYZ point cloud from filename PLY file """ 117 | plydata = PlyData.read(filename) 118 | pc = plydata['vertex'].data 119 | pc_array = np.array([[x, y, z, red, green, blue, label] for x,y,z, red, green, blue, label in pc]) 120 | return pc_array 121 | 122 | -------------------------------------------------------------------------------- /data_preprocessing/prep_stanford/indoor3d_util.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wleigithub/GACNet/bab6c5b45ac74dc2256bba4e5101524ca8086179/data_preprocessing/prep_stanford/indoor3d_util.pyc -------------------------------------------------------------------------------- /data_preprocessing/prep_stanford/meta/all_data_label.txt: -------------------------------------------------------------------------------- 1 | Area_1_conferenceRoom_1.npy 2 | Area_1_conferenceRoom_2.npy 3 | Area_1_copyRoom_1.npy 4 | Area_1_hallway_1.npy 5 | Area_1_hallway_2.npy 6 | Area_1_hallway_3.npy 7 | Area_1_hallway_4.npy 8 | Area_1_hallway_5.npy 9 | Area_1_hallway_6.npy 10 | Area_1_hallway_7.npy 11 | Area_1_hallway_8.npy 12 | Area_1_office_10.npy 13 | Area_1_office_11.npy 14 | Area_1_office_12.npy 15 | Area_1_office_13.npy 16 | Area_1_office_14.npy 17 | Area_1_office_15.npy 18 | Area_1_office_16.npy 19 | Area_1_office_17.npy 20 | Area_1_office_18.npy 21 | Area_1_office_19.npy 22 | Area_1_office_1.npy 23 | Area_1_office_20.npy 24 | Area_1_office_21.npy 25 | Area_1_office_22.npy 26 | Area_1_office_23.npy 27 | Area_1_office_24.npy 28 | Area_1_office_25.npy 29 | Area_1_office_26.npy 30 | Area_1_office_27.npy 31 | Area_1_office_28.npy 32 | Area_1_office_29.npy 33 | Area_1_office_2.npy 34 | Area_1_office_30.npy 35 | Area_1_office_31.npy 36 | Area_1_office_3.npy 37 | Area_1_office_4.npy 38 | Area_1_office_5.npy 39 | Area_1_office_6.npy 40 | Area_1_office_7.npy 41 | Area_1_office_8.npy 42 | Area_1_office_9.npy 43 | Area_1_pantry_1.npy 44 | Area_1_WC_1.npy 45 | Area_2_auditorium_1.npy 46 | Area_2_auditorium_2.npy 47 | Area_2_conferenceRoom_1.npy 48 | Area_2_hallway_10.npy 49 | Area_2_hallway_11.npy 50 | Area_2_hallway_12.npy 51 | Area_2_hallway_1.npy 52 | Area_2_hallway_2.npy 53 | Area_2_hallway_3.npy 54 | Area_2_hallway_4.npy 55 | Area_2_hallway_5.npy 56 | Area_2_hallway_6.npy 57 | Area_2_hallway_7.npy 58 | Area_2_hallway_8.npy 59 | Area_2_hallway_9.npy 60 | Area_2_office_10.npy 61 | Area_2_office_11.npy 62 | Area_2_office_12.npy 63 | Area_2_office_13.npy 64 | Area_2_office_14.npy 65 | Area_2_office_1.npy 66 | Area_2_office_2.npy 67 | Area_2_office_3.npy 68 | Area_2_office_4.npy 69 | Area_2_office_5.npy 70 | Area_2_office_6.npy 71 | Area_2_office_7.npy 72 | Area_2_office_8.npy 73 | Area_2_office_9.npy 74 | Area_2_storage_1.npy 75 | Area_2_storage_2.npy 76 | Area_2_storage_3.npy 77 | Area_2_storage_4.npy 78 | Area_2_storage_5.npy 79 | Area_2_storage_6.npy 80 | Area_2_storage_7.npy 81 | Area_2_storage_8.npy 82 | Area_2_storage_9.npy 83 | Area_2_WC_1.npy 84 | Area_2_WC_2.npy 85 | Area_3_conferenceRoom_1.npy 86 | Area_3_hallway_1.npy 87 | Area_3_hallway_2.npy 88 | Area_3_hallway_3.npy 89 | Area_3_hallway_4.npy 90 | Area_3_hallway_5.npy 91 | Area_3_hallway_6.npy 92 | Area_3_lounge_1.npy 93 | Area_3_lounge_2.npy 94 | Area_3_office_10.npy 95 | Area_3_office_1.npy 96 | Area_3_office_2.npy 97 | Area_3_office_3.npy 98 | Area_3_office_4.npy 99 | Area_3_office_5.npy 100 | Area_3_office_6.npy 101 | Area_3_office_7.npy 102 | Area_3_office_8.npy 103 | Area_3_office_9.npy 104 | Area_3_storage_1.npy 105 | Area_3_storage_2.npy 106 | Area_3_WC_1.npy 107 | Area_3_WC_2.npy 108 | Area_4_conferenceRoom_1.npy 109 | Area_4_conferenceRoom_2.npy 110 | Area_4_conferenceRoom_3.npy 111 | Area_4_hallway_10.npy 112 | Area_4_hallway_11.npy 113 | Area_4_hallway_12.npy 114 | Area_4_hallway_13.npy 115 | Area_4_hallway_14.npy 116 | Area_4_hallway_1.npy 117 | Area_4_hallway_2.npy 118 | Area_4_hallway_3.npy 119 | Area_4_hallway_4.npy 120 | Area_4_hallway_5.npy 121 | Area_4_hallway_6.npy 122 | Area_4_hallway_7.npy 123 | Area_4_hallway_8.npy 124 | Area_4_hallway_9.npy 125 | Area_4_lobby_1.npy 126 | Area_4_lobby_2.npy 127 | Area_4_office_10.npy 128 | Area_4_office_11.npy 129 | Area_4_office_12.npy 130 | Area_4_office_13.npy 131 | Area_4_office_14.npy 132 | Area_4_office_15.npy 133 | Area_4_office_16.npy 134 | Area_4_office_17.npy 135 | Area_4_office_18.npy 136 | Area_4_office_19.npy 137 | Area_4_office_1.npy 138 | Area_4_office_20.npy 139 | Area_4_office_21.npy 140 | Area_4_office_22.npy 141 | Area_4_office_2.npy 142 | Area_4_office_3.npy 143 | Area_4_office_4.npy 144 | Area_4_office_5.npy 145 | Area_4_office_6.npy 146 | Area_4_office_7.npy 147 | Area_4_office_8.npy 148 | Area_4_office_9.npy 149 | Area_4_storage_1.npy 150 | Area_4_storage_2.npy 151 | Area_4_storage_3.npy 152 | Area_4_storage_4.npy 153 | Area_4_WC_1.npy 154 | Area_4_WC_2.npy 155 | Area_4_WC_3.npy 156 | Area_4_WC_4.npy 157 | Area_5_conferenceRoom_1.npy 158 | Area_5_conferenceRoom_2.npy 159 | Area_5_conferenceRoom_3.npy 160 | Area_5_hallway_10.npy 161 | Area_5_hallway_11.npy 162 | Area_5_hallway_12.npy 163 | Area_5_hallway_13.npy 164 | Area_5_hallway_14.npy 165 | Area_5_hallway_15.npy 166 | Area_5_hallway_1.npy 167 | Area_5_hallway_2.npy 168 | Area_5_hallway_3.npy 169 | Area_5_hallway_4.npy 170 | Area_5_hallway_5.npy 171 | Area_5_hallway_7.npy 172 | Area_5_hallway_8.npy 173 | Area_5_hallway_9.npy 174 | Area_5_lobby_1.npy 175 | Area_5_office_10.npy 176 | Area_5_office_11.npy 177 | Area_5_office_12.npy 178 | Area_5_office_13.npy 179 | Area_5_office_14.npy 180 | Area_5_office_15.npy 181 | Area_5_office_16.npy 182 | Area_5_office_17.npy 183 | Area_5_office_18.npy 184 | Area_5_office_19.npy 185 | Area_5_office_1.npy 186 | Area_5_office_20.npy 187 | Area_5_office_21.npy 188 | Area_5_office_22.npy 189 | Area_5_office_23.npy 190 | Area_5_office_24.npy 191 | Area_5_office_25.npy 192 | Area_5_office_26.npy 193 | Area_5_office_27.npy 194 | Area_5_office_28.npy 195 | Area_5_office_29.npy 196 | Area_5_office_2.npy 197 | Area_5_office_30.npy 198 | Area_5_office_31.npy 199 | Area_5_office_32.npy 200 | Area_5_office_33.npy 201 | Area_5_office_34.npy 202 | Area_5_office_35.npy 203 | Area_5_office_36.npy 204 | Area_5_office_37.npy 205 | Area_5_office_38.npy 206 | Area_5_office_39.npy 207 | Area_5_office_3.npy 208 | Area_5_office_40.npy 209 | Area_5_office_41.npy 210 | Area_5_office_42.npy 211 | Area_5_office_4.npy 212 | Area_5_office_5.npy 213 | Area_5_office_6.npy 214 | Area_5_office_7.npy 215 | Area_5_office_8.npy 216 | Area_5_office_9.npy 217 | Area_5_pantry_1.npy 218 | Area_5_storage_1.npy 219 | Area_5_storage_2.npy 220 | Area_5_storage_3.npy 221 | Area_5_storage_4.npy 222 | Area_5_WC_1.npy 223 | Area_5_WC_2.npy 224 | Area_6_conferenceRoom_1.npy 225 | Area_6_copyRoom_1.npy 226 | Area_6_hallway_1.npy 227 | Area_6_hallway_2.npy 228 | Area_6_hallway_3.npy 229 | Area_6_hallway_4.npy 230 | Area_6_hallway_5.npy 231 | Area_6_hallway_6.npy 232 | Area_6_lounge_1.npy 233 | Area_6_office_10.npy 234 | Area_6_office_11.npy 235 | Area_6_office_12.npy 236 | Area_6_office_13.npy 237 | Area_6_office_14.npy 238 | Area_6_office_15.npy 239 | Area_6_office_16.npy 240 | Area_6_office_17.npy 241 | Area_6_office_18.npy 242 | Area_6_office_19.npy 243 | Area_6_office_1.npy 244 | Area_6_office_20.npy 245 | Area_6_office_21.npy 246 | Area_6_office_22.npy 247 | Area_6_office_23.npy 248 | Area_6_office_24.npy 249 | Area_6_office_25.npy 250 | Area_6_office_26.npy 251 | Area_6_office_27.npy 252 | Area_6_office_28.npy 253 | Area_6_office_29.npy 254 | Area_6_office_2.npy 255 | Area_6_office_30.npy 256 | Area_6_office_31.npy 257 | Area_6_office_32.npy 258 | Area_6_office_33.npy 259 | Area_6_office_34.npy 260 | Area_6_office_35.npy 261 | Area_6_office_36.npy 262 | Area_6_office_37.npy 263 | Area_6_office_3.npy 264 | Area_6_office_4.npy 265 | Area_6_office_5.npy 266 | Area_6_office_6.npy 267 | Area_6_office_7.npy 268 | Area_6_office_8.npy 269 | Area_6_office_9.npy 270 | Area_6_openspace_1.npy 271 | Area_6_pantry_1.npy 272 | -------------------------------------------------------------------------------- /data_preprocessing/prep_stanford/meta/anno_paths.txt: -------------------------------------------------------------------------------- 1 | Area_1/conferenceRoom_1/Annotations 2 | Area_1/conferenceRoom_2/Annotations 3 | Area_1/copyRoom_1/Annotations 4 | Area_1/hallway_1/Annotations 5 | Area_1/hallway_2/Annotations 6 | Area_1/hallway_3/Annotations 7 | Area_1/hallway_4/Annotations 8 | Area_1/hallway_5/Annotations 9 | Area_1/hallway_6/Annotations 10 | Area_1/hallway_7/Annotations 11 | Area_1/hallway_8/Annotations 12 | Area_1/office_10/Annotations 13 | Area_1/office_11/Annotations 14 | Area_1/office_12/Annotations 15 | Area_1/office_13/Annotations 16 | Area_1/office_14/Annotations 17 | Area_1/office_15/Annotations 18 | Area_1/office_16/Annotations 19 | Area_1/office_17/Annotations 20 | Area_1/office_18/Annotations 21 | Area_1/office_19/Annotations 22 | Area_1/office_1/Annotations 23 | Area_1/office_20/Annotations 24 | Area_1/office_21/Annotations 25 | Area_1/office_22/Annotations 26 | Area_1/office_23/Annotations 27 | Area_1/office_24/Annotations 28 | Area_1/office_25/Annotations 29 | Area_1/office_26/Annotations 30 | Area_1/office_27/Annotations 31 | Area_1/office_28/Annotations 32 | Area_1/office_29/Annotations 33 | Area_1/office_2/Annotations 34 | Area_1/office_30/Annotations 35 | Area_1/office_31/Annotations 36 | Area_1/office_3/Annotations 37 | Area_1/office_4/Annotations 38 | Area_1/office_5/Annotations 39 | Area_1/office_6/Annotations 40 | Area_1/office_7/Annotations 41 | Area_1/office_8/Annotations 42 | Area_1/office_9/Annotations 43 | Area_1/pantry_1/Annotations 44 | Area_1/WC_1/Annotations 45 | Area_2/auditorium_1/Annotations 46 | Area_2/auditorium_2/Annotations 47 | Area_2/conferenceRoom_1/Annotations 48 | Area_2/hallway_10/Annotations 49 | Area_2/hallway_11/Annotations 50 | Area_2/hallway_12/Annotations 51 | Area_2/hallway_1/Annotations 52 | Area_2/hallway_2/Annotations 53 | Area_2/hallway_3/Annotations 54 | Area_2/hallway_4/Annotations 55 | Area_2/hallway_5/Annotations 56 | Area_2/hallway_6/Annotations 57 | Area_2/hallway_7/Annotations 58 | Area_2/hallway_8/Annotations 59 | Area_2/hallway_9/Annotations 60 | Area_2/office_10/Annotations 61 | Area_2/office_11/Annotations 62 | Area_2/office_12/Annotations 63 | Area_2/office_13/Annotations 64 | Area_2/office_14/Annotations 65 | Area_2/office_1/Annotations 66 | Area_2/office_2/Annotations 67 | Area_2/office_3/Annotations 68 | Area_2/office_4/Annotations 69 | Area_2/office_5/Annotations 70 | Area_2/office_6/Annotations 71 | Area_2/office_7/Annotations 72 | Area_2/office_8/Annotations 73 | Area_2/office_9/Annotations 74 | Area_2/storage_1/Annotations 75 | Area_2/storage_2/Annotations 76 | Area_2/storage_3/Annotations 77 | Area_2/storage_4/Annotations 78 | Area_2/storage_5/Annotations 79 | Area_2/storage_6/Annotations 80 | Area_2/storage_7/Annotations 81 | Area_2/storage_8/Annotations 82 | Area_2/storage_9/Annotations 83 | Area_2/WC_1/Annotations 84 | Area_2/WC_2/Annotations 85 | Area_3/conferenceRoom_1/Annotations 86 | Area_3/hallway_1/Annotations 87 | Area_3/hallway_2/Annotations 88 | Area_3/hallway_3/Annotations 89 | Area_3/hallway_4/Annotations 90 | Area_3/hallway_5/Annotations 91 | Area_3/hallway_6/Annotations 92 | Area_3/lounge_1/Annotations 93 | Area_3/lounge_2/Annotations 94 | Area_3/office_10/Annotations 95 | Area_3/office_1/Annotations 96 | Area_3/office_2/Annotations 97 | Area_3/office_3/Annotations 98 | Area_3/office_4/Annotations 99 | Area_3/office_5/Annotations 100 | Area_3/office_6/Annotations 101 | Area_3/office_7/Annotations 102 | Area_3/office_8/Annotations 103 | Area_3/office_9/Annotations 104 | Area_3/storage_1/Annotations 105 | Area_3/storage_2/Annotations 106 | Area_3/WC_1/Annotations 107 | Area_3/WC_2/Annotations 108 | Area_4/conferenceRoom_1/Annotations 109 | Area_4/conferenceRoom_2/Annotations 110 | Area_4/conferenceRoom_3/Annotations 111 | Area_4/hallway_10/Annotations 112 | Area_4/hallway_11/Annotations 113 | Area_4/hallway_12/Annotations 114 | Area_4/hallway_13/Annotations 115 | Area_4/hallway_14/Annotations 116 | Area_4/hallway_1/Annotations 117 | Area_4/hallway_2/Annotations 118 | Area_4/hallway_3/Annotations 119 | Area_4/hallway_4/Annotations 120 | Area_4/hallway_5/Annotations 121 | Area_4/hallway_6/Annotations 122 | Area_4/hallway_7/Annotations 123 | Area_4/hallway_8/Annotations 124 | Area_4/hallway_9/Annotations 125 | Area_4/lobby_1/Annotations 126 | Area_4/lobby_2/Annotations 127 | Area_4/office_10/Annotations 128 | Area_4/office_11/Annotations 129 | Area_4/office_12/Annotations 130 | Area_4/office_13/Annotations 131 | Area_4/office_14/Annotations 132 | Area_4/office_15/Annotations 133 | Area_4/office_16/Annotations 134 | Area_4/office_17/Annotations 135 | Area_4/office_18/Annotations 136 | Area_4/office_19/Annotations 137 | Area_4/office_1/Annotations 138 | Area_4/office_20/Annotations 139 | Area_4/office_21/Annotations 140 | Area_4/office_22/Annotations 141 | Area_4/office_2/Annotations 142 | Area_4/office_3/Annotations 143 | Area_4/office_4/Annotations 144 | Area_4/office_5/Annotations 145 | Area_4/office_6/Annotations 146 | Area_4/office_7/Annotations 147 | Area_4/office_8/Annotations 148 | Area_4/office_9/Annotations 149 | Area_4/storage_1/Annotations 150 | Area_4/storage_2/Annotations 151 | Area_4/storage_3/Annotations 152 | Area_4/storage_4/Annotations 153 | Area_4/WC_1/Annotations 154 | Area_4/WC_2/Annotations 155 | Area_4/WC_3/Annotations 156 | Area_4/WC_4/Annotations 157 | Area_5/conferenceRoom_1/Annotations 158 | Area_5/conferenceRoom_2/Annotations 159 | Area_5/conferenceRoom_3/Annotations 160 | Area_5/hallway_10/Annotations 161 | Area_5/hallway_11/Annotations 162 | Area_5/hallway_12/Annotations 163 | Area_5/hallway_13/Annotations 164 | Area_5/hallway_14/Annotations 165 | Area_5/hallway_15/Annotations 166 | Area_5/hallway_1/Annotations 167 | Area_5/hallway_2/Annotations 168 | Area_5/hallway_3/Annotations 169 | Area_5/hallway_4/Annotations 170 | Area_5/hallway_5/Annotations 171 | Area_5/hallway_6/Annotations 172 | Area_5/hallway_7/Annotations 173 | Area_5/hallway_8/Annotations 174 | Area_5/hallway_9/Annotations 175 | Area_5/lobby_1/Annotations 176 | Area_5/office_10/Annotations 177 | Area_5/office_11/Annotations 178 | Area_5/office_12/Annotations 179 | Area_5/office_13/Annotations 180 | Area_5/office_14/Annotations 181 | Area_5/office_15/Annotations 182 | Area_5/office_16/Annotations 183 | Area_5/office_17/Annotations 184 | Area_5/office_18/Annotations 185 | Area_5/office_19/Annotations 186 | Area_5/office_1/Annotations 187 | Area_5/office_20/Annotations 188 | Area_5/office_21/Annotations 189 | Area_5/office_22/Annotations 190 | Area_5/office_23/Annotations 191 | Area_5/office_24/Annotations 192 | Area_5/office_25/Annotations 193 | Area_5/office_26/Annotations 194 | Area_5/office_27/Annotations 195 | Area_5/office_28/Annotations 196 | Area_5/office_29/Annotations 197 | Area_5/office_2/Annotations 198 | Area_5/office_30/Annotations 199 | Area_5/office_31/Annotations 200 | Area_5/office_32/Annotations 201 | Area_5/office_33/Annotations 202 | Area_5/office_34/Annotations 203 | Area_5/office_35/Annotations 204 | Area_5/office_36/Annotations 205 | Area_5/office_37/Annotations 206 | Area_5/office_38/Annotations 207 | Area_5/office_39/Annotations 208 | Area_5/office_3/Annotations 209 | Area_5/office_40/Annotations 210 | Area_5/office_41/Annotations 211 | Area_5/office_42/Annotations 212 | Area_5/office_4/Annotations 213 | Area_5/office_5/Annotations 214 | Area_5/office_6/Annotations 215 | Area_5/office_7/Annotations 216 | Area_5/office_8/Annotations 217 | Area_5/office_9/Annotations 218 | Area_5/pantry_1/Annotations 219 | Area_5/storage_1/Annotations 220 | Area_5/storage_2/Annotations 221 | Area_5/storage_3/Annotations 222 | Area_5/storage_4/Annotations 223 | Area_5/WC_1/Annotations 224 | Area_5/WC_2/Annotations 225 | Area_6/conferenceRoom_1/Annotations 226 | Area_6/copyRoom_1/Annotations 227 | Area_6/hallway_1/Annotations 228 | Area_6/hallway_2/Annotations 229 | Area_6/hallway_3/Annotations 230 | Area_6/hallway_4/Annotations 231 | Area_6/hallway_5/Annotations 232 | Area_6/hallway_6/Annotations 233 | Area_6/lounge_1/Annotations 234 | Area_6/office_10/Annotations 235 | Area_6/office_11/Annotations 236 | Area_6/office_12/Annotations 237 | Area_6/office_13/Annotations 238 | Area_6/office_14/Annotations 239 | Area_6/office_15/Annotations 240 | Area_6/office_16/Annotations 241 | Area_6/office_17/Annotations 242 | Area_6/office_18/Annotations 243 | Area_6/office_19/Annotations 244 | Area_6/office_1/Annotations 245 | Area_6/office_20/Annotations 246 | Area_6/office_21/Annotations 247 | Area_6/office_22/Annotations 248 | Area_6/office_23/Annotations 249 | Area_6/office_24/Annotations 250 | Area_6/office_25/Annotations 251 | Area_6/office_26/Annotations 252 | Area_6/office_27/Annotations 253 | Area_6/office_28/Annotations 254 | Area_6/office_29/Annotations 255 | Area_6/office_2/Annotations 256 | Area_6/office_30/Annotations 257 | Area_6/office_31/Annotations 258 | Area_6/office_32/Annotations 259 | Area_6/office_33/Annotations 260 | Area_6/office_34/Annotations 261 | Area_6/office_35/Annotations 262 | Area_6/office_36/Annotations 263 | Area_6/office_37/Annotations 264 | Area_6/office_3/Annotations 265 | Area_6/office_4/Annotations 266 | Area_6/office_5/Annotations 267 | Area_6/office_6/Annotations 268 | Area_6/office_7/Annotations 269 | Area_6/office_8/Annotations 270 | Area_6/office_9/Annotations 271 | Area_6/openspace_1/Annotations 272 | Area_6/pantry_1/Annotations 273 | -------------------------------------------------------------------------------- /data_preprocessing/prep_stanford/meta/area6_data_label.txt: -------------------------------------------------------------------------------- 1 | data/stanford_indoor3d/Area_6_conferenceRoom_1.npy 2 | data/stanford_indoor3d/Area_6_copyRoom_1.npy 3 | data/stanford_indoor3d/Area_6_hallway_1.npy 4 | data/stanford_indoor3d/Area_6_hallway_2.npy 5 | data/stanford_indoor3d/Area_6_hallway_3.npy 6 | data/stanford_indoor3d/Area_6_hallway_4.npy 7 | data/stanford_indoor3d/Area_6_hallway_5.npy 8 | data/stanford_indoor3d/Area_6_hallway_6.npy 9 | data/stanford_indoor3d/Area_6_lounge_1.npy 10 | data/stanford_indoor3d/Area_6_office_10.npy 11 | data/stanford_indoor3d/Area_6_office_11.npy 12 | data/stanford_indoor3d/Area_6_office_12.npy 13 | data/stanford_indoor3d/Area_6_office_13.npy 14 | data/stanford_indoor3d/Area_6_office_14.npy 15 | data/stanford_indoor3d/Area_6_office_15.npy 16 | data/stanford_indoor3d/Area_6_office_16.npy 17 | data/stanford_indoor3d/Area_6_office_17.npy 18 | data/stanford_indoor3d/Area_6_office_18.npy 19 | data/stanford_indoor3d/Area_6_office_19.npy 20 | data/stanford_indoor3d/Area_6_office_1.npy 21 | data/stanford_indoor3d/Area_6_office_20.npy 22 | data/stanford_indoor3d/Area_6_office_21.npy 23 | data/stanford_indoor3d/Area_6_office_22.npy 24 | data/stanford_indoor3d/Area_6_office_23.npy 25 | data/stanford_indoor3d/Area_6_office_24.npy 26 | data/stanford_indoor3d/Area_6_office_25.npy 27 | data/stanford_indoor3d/Area_6_office_26.npy 28 | data/stanford_indoor3d/Area_6_office_27.npy 29 | data/stanford_indoor3d/Area_6_office_28.npy 30 | data/stanford_indoor3d/Area_6_office_29.npy 31 | data/stanford_indoor3d/Area_6_office_2.npy 32 | data/stanford_indoor3d/Area_6_office_30.npy 33 | data/stanford_indoor3d/Area_6_office_31.npy 34 | data/stanford_indoor3d/Area_6_office_32.npy 35 | data/stanford_indoor3d/Area_6_office_33.npy 36 | data/stanford_indoor3d/Area_6_office_34.npy 37 | data/stanford_indoor3d/Area_6_office_35.npy 38 | data/stanford_indoor3d/Area_6_office_36.npy 39 | data/stanford_indoor3d/Area_6_office_37.npy 40 | data/stanford_indoor3d/Area_6_office_3.npy 41 | data/stanford_indoor3d/Area_6_office_4.npy 42 | data/stanford_indoor3d/Area_6_office_5.npy 43 | data/stanford_indoor3d/Area_6_office_6.npy 44 | data/stanford_indoor3d/Area_6_office_7.npy 45 | data/stanford_indoor3d/Area_6_office_8.npy 46 | data/stanford_indoor3d/Area_6_office_9.npy 47 | data/stanford_indoor3d/Area_6_openspace_1.npy 48 | data/stanford_indoor3d/Area_6_pantry_1.npy 49 | -------------------------------------------------------------------------------- /data_preprocessing/prep_stanford/meta/class_names.txt: -------------------------------------------------------------------------------- 1 | ceiling 2 | floor 3 | wall 4 | beam 5 | column 6 | window 7 | door 8 | table 9 | chair 10 | sofa 11 | bookcase 12 | board 13 | clutter 14 | -------------------------------------------------------------------------------- /data_preprocessing/prep_stanford/prep_scene.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import h5py 3 | import os 4 | import sys 5 | import pdb 6 | import glob 7 | from plyfile import PlyData, PlyElement 8 | 9 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 10 | ROOT_DIR = os.path.dirname(os.path.dirname(BASE_DIR)) 11 | sys.path.append(os.path.join(os.path.dirname(BASE_DIR), 'eig_computation_op')) 12 | from eig_feature import eig_cal_map, FPS 13 | 14 | # ----------------------------------------------------------------------------- 15 | # PREPARE BLOCK DATA FOR DEEPNETS TRAINING/TESTING 16 | # ----------------------------------------------------------------------------- 17 | def test_area2blocks_zeromean(data_filename, num_point, block_size=1.0, inner_size=0.5, radius=0.15): 18 | if data_filename[-3:] == 'ply': 19 | data = read_xyzrgbL_ply(data_filename) 20 | else: 21 | print('Unknown file type! exiting.') 22 | exit() 23 | 24 | print('data load done') 25 | data = data[:, 0:6] 26 | data[:,3:6] /= 255.0 27 | idx = np.array(range(data.shape[0])) 28 | 29 | data_batch, idx_batch, isvalid, pts_num = scene2blocks(data, idx, num_point, block_size=block_size, inner_size=inner_size) 30 | 31 | #eigenvalues calculaition 32 | print('eigenvalues calculaiting...') 33 | idx = np.reshape(idx_batch, [-1]) 34 | unique_idx = np.unique(idx) 35 | eigs_batch = eig_cal_map(data[:,0:3], unique_idx, idx, radius) 36 | eigs_batch = np.reshape(eigs_batch, [-1, 6]) 37 | data_batch = np.concatenate([data_batch, eigs_batch], -1) 38 | 39 | return data_batch, idx_batch, isvalid, pts_num 40 | 41 | 42 | 43 | def scene2blocks(data, idx, num_point, block_size=1.0, inner_size=0.5, count_limit=100): 44 | """ Prepare block data for scene testing. 45 | DISCRIBTION: 46 | Each block consits of two part: the inner part with size: block_size*block_size*hight, 47 | and the buffer area around the inner part with side lenght buffer_size 48 | Args: 49 | data: N x 6 numpy array, 012 are XYZ in meters, 345 are RGB in [0,1] 50 | label_id: N size int numpy array 51 | block_size: float, physical size of the block in meters 52 | buffer_size: float, lenght of buffer area 53 | Returns: 54 | block_data_list: list, each element is the points selected from the block 55 | eachblock_pointnum_list: list, each element is the number of points selected from the block, [0, max_num_point] 56 | eachblock_validnum_list: list, each element is the number of points lying in the inner part 57 | block_label_id_list: list, each element is the label_idc of points lying in the inner part 58 | """ 59 | 60 | buffer_size = (block_size-inner_size)/2 61 | stride = inner_size 62 | 63 | data[:,0:2] -= np.amin(data, 0)[0:2] 64 | limit = np.amax(data, 0)[0:3] 65 | 66 | num_block_x = int(np.ceil(limit[0] / stride)) 67 | num_block_y = int(np.ceil(limit[1] / stride)) 68 | 69 | # Collect blocks 70 | block_data_list = [] 71 | block_idx_list = [] 72 | block_isvalid_list = [] 73 | n_point_list = [] 74 | 75 | for i in range(num_block_x): 76 | print('%d/%d'%(i, num_block_x)) 77 | for j in range(num_block_y): 78 | xbeg = i*stride 79 | ybeg = j*stride 80 | 81 | #find points in [x-buffer,x+stride+buffer]*[y-buffer,y+stride+buffer] 82 | xcond = (data[:,0]<=xbeg+stride+buffer_size) & (data[:,0]>=xbeg-buffer_size) 83 | ycond = (data[:,1]<=ybeg+stride+buffer_size) & (data[:,1]>=ybeg-buffer_size) 84 | 85 | cond = xcond & ycond 86 | block_points = data[cond, :] 87 | block_idx = idx[cond] 88 | 89 | #block_points, block_idx, n_point = random_sample_data_label(block_points, block_idx, num_point) 90 | block_points, block_idx, n_point = FPS_data_label(block_points, block_idx, num_point, res=0.02) 91 | 92 | # then find points in [x,x+stride]*[y,y+stride] 93 | xcond = (block_points[:,0]<=xbeg+stride) & (block_points[:,0]>=xbeg) 94 | ycond = (block_points[:,1]<=ybeg+stride) & (block_points[:,1]>=ybeg) 95 | cond = xcond & ycond 96 | if np.sum(cond) < count_limit: # discard block if there are less than 50 pts in the inner part. 97 | continue 98 | 99 | #zero mean 100 | block_points[:,0:2] -= np.mean(block_points[:,0:2], 0) 101 | 102 | block_data_list.append(block_points) 103 | block_idx_list.append(block_idx) 104 | block_isvalid_list.append(cond) 105 | n_point_list.append(n_point) 106 | return np.concatenate(block_data_list, 0), np.concatenate(block_idx_list, 0), np.concatenate(block_isvalid_list, 0), np.array(n_point_list) 107 | 108 | 109 | def FPS_data_label(data, label, num_sample, res=None): 110 | """ data is in N x ... 111 | we want to keep num_samplexC of them. 112 | if N > num_sample, we will randomly keep num_sample of them. 113 | if N < num_sample, we will randomly duplicate samples. 114 | """ 115 | N = data.shape[0] 116 | xyz = data[:,0:3] 117 | if (N > num_sample): 118 | if res is not None: 119 | ridx = downsampling(xyz, res) 120 | data = data[ridx,...] 121 | label = label[ridx,...] 122 | xyz = xyz[ridx,...] 123 | idx = FPS(xyz, num_sample) 124 | return data[idx, ...], label[idx, ...], num_sample 125 | else: 126 | return data, label, N 127 | 128 | def downsampling(xyz, res): 129 | limitmax = np.amax(xyz, 0)[0:3] 130 | limitmin = np.amin(xyz, 0)[0:3] 131 | num_vox = np.ceil((limitmax-limitmin)/res) 132 | idx = np.ceil((xyz-limitmin)/res) 133 | idx = idx[:,0] + idx[:,1]*num_vox[0] + idx[:,2]*num_vox[0]*num_vox[1] 134 | uidx, ridx = np.unique(idx, return_index=True) 135 | return ridx 136 | 137 | 138 | def random_sample_data_label(data, label, num_sample): 139 | """ data is in N x ... 140 | we want to keep num_samplexC of them. 141 | if N > num_sample, we will randomly keep num_sample of them. 142 | if N < num_sample, we will randomly duplicate samples. 143 | """ 144 | N = data.shape[0] 145 | 146 | if (N > num_sample): 147 | sample = np.random.choice(N, num_sample) 148 | return data[sample, ...], label[sample, ...], num_sample 149 | else: 150 | return data, label, N 151 | 152 | 153 | def read_xyzrgbL_ply(filename): 154 | """ read XYZ point cloud from filename PLY file """ 155 | plydata = PlyData.read(filename) 156 | pc = plydata['vertex'].data 157 | pc_array = np.array([[x, y, z, red, green, blue, label] for x,y,z, red, green, blue, label in pc]) 158 | return pc_array 159 | 160 | 161 | # Write numpy array data and incice to h5_filename 162 | def save_h5_data_and_idx(h5_filename, data, idx, mask, pts_num): 163 | h5_fout = h5py.File(h5_filename) 164 | h5_fout.create_dataset('data', data=data, dtype='float32') 165 | h5_fout.create_dataset('idx', data=idx, dtype='int32') 166 | h5_fout.create_dataset('mask', data=mask, dtype='bool') 167 | h5_fout.create_dataset('pts_num', data=pts_num, dtype='int32') 168 | h5_fout.close() 169 | 170 | 171 | if __name__=='__main__': 172 | # Note: the block size for testing is 3.6m*3.6m for speed consideration, which is larger than training 173 | #settings 174 | org_datadir = '/media/wl/user/3D_pcseg_Data/S3DIS' 175 | 176 | indoor3d_data_dir = os.path.join(org_datadir, 'stanford_indoor3d') 177 | output_dir = os.path.join(ROOT_DIR, 'Data/S3DIS/test_dataset') 178 | NUM_POINT = 4096*9 179 | 180 | if not os.path.exists(output_dir): os.mkdir(output_dir) 181 | test_area = 'Area_5' 182 | 183 | Files = glob.glob(indoor3d_data_dir + '/*.ply') 184 | for i in range(len(Files)): 185 | data_filename = Files[i] 186 | if test_area in data_filename: 187 | h5_filename = os.path.join(output_dir, data_filename.split('.')[0].split('/')[-1]+'.h5') 188 | print(h5_filename) 189 | data, idx, mask, pts_num = test_area2blocks_zeromean(data_filename, NUM_POINT, block_size=3.6, inner_size=3.4, radius=0.1) 190 | save_h5_data_and_idx(h5_filename, data, idx, mask, pts_num) 191 | 192 | -------------------------------------------------------------------------------- /net_S3DIS/run.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import train 4 | 5 | parser = argparse.ArgumentParser() 6 | parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]') 7 | parser.add_argument('--log_dir', default='log', help='Log dir [default: log]') 8 | parser.add_argument('--num_class', type=int, default=13, help='number of class') 9 | parser.add_argument('--num_point', type=int, default=4096, help='Point number [default: 4096]') 10 | parser.add_argument('--max_epoch', type=int, default=50, help='Epoch to run [default: 50]') 11 | parser.add_argument('--batch_size', type=int, default=15, help='Batch Size during training [default: 24]') 12 | parser.add_argument('--learning_rate', type=float, default=0.001, help='Initial learning rate [default: 0.001]') 13 | parser.add_argument('--momentum', type=float, default=0.9, help='Initial learning rate [default: 0.9]') 14 | parser.add_argument('--optimizer', default='adam', help='adam or momentum [default: adam]') 15 | parser.add_argument('--decay_step', type=int, default=300000, help='Decay step for lr decay [default: 300000]') 16 | parser.add_argument('--decay_rate', type=float, default=0.5, help='Decay rate for lr decay [default: 0.5]') 17 | parser.add_argument('--test_area', type=int, default=5, help='Which area to use for test, option: 1-6 [default: 5]') 18 | parser.add_argument('--data_path', type=str, default='../Data/S3DIS/train_dataset', help='data path') 19 | parser.add_argument('--roomlist_file', type=str, default='../Data/S3DIS/train_dataset/room_filelist.txt', help='The file which recorded the room name of each train data ') 20 | args = parser.parse_args() 21 | 22 | 23 | #graph_inf contents parameters for grapg building and coarsing 24 | graph_inf = {'stride_list': [4, 4, 4, 2], #can be seen as the downsampling rate 25 | 'radius_list': [0.1, 0.2, 0.4, 0.8, 1.6], # radius for neighbor points searching 26 | 'maxsample_list': [12, 21, 21, 21, 12] #number of neighbor points for each layer 27 | } 28 | 29 | # number of units for each mlp layer 30 | forward_parm = [ 31 | [ [32,32,64], [64] ], 32 | [ [64,64,128], [128] ], 33 | [ [128,128,256], [256] ], 34 | [ [256,256,512], [512] ], 35 | [ [256,256], [256] ] 36 | ] 37 | 38 | # for feature interpolation stage 39 | upsample_parm = [ 40 | [128, 128], 41 | [128, 128], 42 | [256, 256], 43 | [256, 256] 44 | ] 45 | 46 | # parameters for fully connection layer 47 | fullconect_parm = 128 48 | 49 | net_inf = {'forward_parm': forward_parm, 50 | 'upsample_parm': upsample_parm, 51 | 'fullconect_parm': fullconect_parm 52 | } 53 | 54 | 55 | if not os.path.exists(args.log_dir): os.mkdir(args.log_dir) 56 | os.system('cp run.py %s' % (args.log_dir)) # back up parameter inf 57 | os.system('cp train.py %s' % (args.log_dir)) # back up parameter inf 58 | os.system('cp ../utils/pcnet_util.py %s' % (args.log_dir)) # back up parameter inf 59 | os.system('cp ../utils/model.py %s' % (args.log_dir)) # back up parameter inf 60 | train.train(args, graph_inf, net_inf) 61 | -------------------------------------------------------------------------------- /net_S3DIS/run.sh: -------------------------------------------------------------------------------- 1 | python run.py 2 | python run_test.py 3 | -------------------------------------------------------------------------------- /net_S3DIS/run_test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import test 4 | 5 | parser = argparse.ArgumentParser() 6 | parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]') 7 | parser.add_argument('--log_dir', default='log_area5', help='Log dir [default: log]') 8 | parser.add_argument('--num_class', type=int, default=13, help='number of class') 9 | parser.add_argument('--channel', type=int, default=12, help='number of feature channel') 10 | parser.add_argument('--testdataset_dir', type=str, default='../Data/S3DIS/test_dataset', help='testdataset path') 11 | parser.add_argument('--test_dir', type=str, default='/media/wl/user/3D_pcseg_Data/S3DIS/stanford_indoor3d', help='The file contains original test data') 12 | parser.add_argument('--outdir', type=str, default='../Data/S3DIS/labels_pred_true') 13 | 14 | args = parser.parse_args() 15 | 16 | #graph_inf contents parameters for grapg building and coarsing 17 | graph_inf = {'stride_list': [4, 4, 4, 2], 18 | 'radius_list': [0.1, 0.2, 0.4, 0.8, 1.6], 19 | 'maxsample_list': [12, 21, 21, 21, 12] 20 | } 21 | 22 | # number of units for each mlp layer 23 | forward_parm = [ 24 | [ [32,32,64], [64] ], 25 | [ [64,64,128], [128] ], 26 | [ [128,128,256], [256] ], 27 | [ [256,256,512], [512] ], 28 | [ [256,256], [256] ] 29 | ] 30 | 31 | # for feature interpolation stage 32 | upsample_parm = [ 33 | [128, 128], 34 | [128, 128], 35 | [256, 256], 36 | [256, 256] 37 | ] 38 | 39 | # parameters for fully connection layer 40 | fullconect_parm = 128 41 | 42 | net_inf = {'forward_parm': forward_parm, 43 | 'upsample_parm': upsample_parm, 44 | 'fullconect_parm': fullconect_parm 45 | } 46 | 47 | 48 | if not os.path.exists(args.outdir): os.mkdir(args.outdir) 49 | 50 | test.test(args, graph_inf, net_inf) 51 | test.interpolate(args) 52 | test.acc_report(args) 53 | -------------------------------------------------------------------------------- /net_S3DIS/test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import math 3 | import h5py 4 | import numpy as np 5 | import tensorflow as tf 6 | import glob 7 | import pdb 8 | 9 | import os 10 | import sys 11 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 12 | ROOT_DIR = os.path.dirname(BASE_DIR) 13 | sys.path.append(BASE_DIR) 14 | sys.path.append(ROOT_DIR) 15 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 16 | sys.path.append(os.path.join(ROOT_DIR, 'test_utils')) 17 | 18 | from three_interpolate import three_interpolate 19 | import provider 20 | from model import * 21 | 22 | def test(args, graph_inf, net_inf): 23 | with tf.Graph().as_default(): 24 | with tf.device('/gpu:'+str(args.gpu)): 25 | pointclouds_pl, labels_pl, is_training_pl = placeholder_inputs(args.channel) 26 | 27 | # Get model and loss 28 | graph_prd, coarse_map = build_graph_pyramid(pointclouds_pl[...,0:3], graph_inf) 29 | logits = build_network(graph_prd, coarse_map, net_inf, pointclouds_pl, args.num_class, is_training_pl, bn_decay=None) 30 | pred = tf.nn.softmax(logits) 31 | 32 | # Add ops to save and restore all the variables. 33 | saver = tf.train.Saver() 34 | 35 | # Create a session 36 | config = tf.ConfigProto() 37 | config.gpu_options.allow_growth = True 38 | config.allow_soft_placement = True 39 | config.log_device_placement = True 40 | sess = tf.Session(config=config) 41 | 42 | # Add summary writers 43 | merged = tf.summary.merge_all() 44 | 45 | # Init variables 46 | init = tf.global_variables_initializer() 47 | sess.run(init, {is_training_pl:False}) 48 | 49 | last_model = tf.train.latest_checkpoint(args.log_dir) 50 | saver.restore(sess, last_model) 51 | ops = {'pointclouds_pl': pointclouds_pl, 52 | #'labels_pl': labels_pl, 53 | 'is_training_pl': is_training_pl, 54 | 'pred' : pred 55 | } 56 | 57 | TEST_FILES = glob.glob(args.testdataset_dir + '/*.h5') 58 | for i in range(len(TEST_FILES)): 59 | file_name = TEST_FILES[i] 60 | print(file_name+'\n') 61 | testset = h5py.File(file_name) 62 | points = testset['data'][:] # (blocks_num, 4096, 12) 63 | idx = testset['idx'][:] 64 | mask = testset['mask'][:] 65 | pts_num = testset['pts_num'][:] 66 | 67 | probs = [] 68 | start_id =0 69 | end_id =0 70 | for j in range(len(pts_num)): 71 | end_id += pts_num[j] 72 | current_data = points[start_id:end_id,:] 73 | current_data = np.expand_dims(current_data, 0) 74 | feed_dict = {ops['pointclouds_pl']: current_data, ops['is_training_pl']: False} 75 | current_probs = sess.run(ops['pred'], feed_dict=feed_dict) 76 | probs.append(np.squeeze(current_probs, 0)) 77 | start_id += pts_num[j] 78 | probs = np.concatenate(probs, 0) 79 | probs = probs[mask,:] 80 | idx = idx[mask] 81 | 82 | out_filename = file_name.split('/')[-1].split('.')[0] 83 | save_h5_probs_idx(os.path.join(args.outdir, out_filename+'.h5'), probs, idx) 84 | 85 | sess.close() 86 | 87 | 88 | def interpolate(args): 89 | """ 90 | Calculate the label for the orignal point cloud 91 | Since we only predict the label of part of the points sampled from the orginal point cloud, 92 | we have to calculate the label of the orignal point cloud with interpolation algorithm. 93 | Similar as we do for feature interpolation in the network, we also interpolate the label of each point (of the orignal point cloud) according to its three nearest neighbors 94 | """ 95 | tmp_files = glob.glob(os.path.join(args.outdir, '*.h5')) 96 | for i in range(len(tmp_files)): 97 | file_name = tmp_files[i] 98 | filename = file_name.split('/')[-1].split('.')[0] 99 | print(filename) 100 | 101 | probs, idx = load_h5_probs_idx(os.path.join(args.outdir, filename+'.h5')) 102 | org_pts = provider.read_xyzrgbL_ply( os.path.join(args.test_dir, filename+'.ply') ) 103 | true_label = org_pts[:,-1] 104 | probs_whole = three_interpolate(org_pts, idx, probs, args.num_class) 105 | pred_label = np.argmax(probs_whole, axis=-1) 106 | pred_label = np.uint8(pred_label) 107 | 108 | save_labels_pred_true(os.path.join(args.outdir, filename+'_labels.pred_true'), pred_label, true_label) 109 | 110 | # only save the points without ceiling for vision efficiency 111 | idx_without_ceiling = np.where(true_label!=0)[0] 112 | pts_xyz = org_pts[idx_without_ceiling, 0:3] 113 | pred_label = pred_label[idx_without_ceiling] 114 | #provider.write_xyzL_ply(pts_xyz, pred_label, os.path.join(args.outdir, filename+'.ply')) 115 | 116 | os.remove(os.path.join(args.outdir, filename+'.h5')) 117 | 118 | 119 | ############################################################################################################ 120 | #---------------------------- utilized functions 121 | ############################################################################################################ 122 | # Write numpy array data and label to h5_filename 123 | def save_h5_probs_idx(h5_filename, probs, idx, data_dtype='float32', idx_dtype='int32'): 124 | h5_fout = h5py.File(h5_filename) 125 | h5_fout.create_dataset( 126 | 'probs', data=probs, 127 | compression='gzip', compression_opts=4, 128 | dtype=data_dtype) 129 | h5_fout.create_dataset( 130 | 'idx', data=idx, 131 | compression='gzip', compression_opts=1, 132 | dtype=idx_dtype) 133 | h5_fout.close() 134 | 135 | 136 | def load_h5_probs_idx(h5_filename): 137 | f = h5py.File(h5_filename) 138 | probs = f['probs'][:] 139 | idx = f['idx'][:] 140 | return probs, idx 141 | 142 | 143 | def save_labels_pred_true(filename, pred_label, true_label): 144 | f = open(filename, 'w') 145 | for i in range(len(true_label)): 146 | f.write('%f %f\n' % (pred_label[i], true_label[i])) 147 | f.close() 148 | 149 | 150 | def compute_iou2(args, predl, truel): 151 | gt_classes = [0 for _ in range(args.num_class)] 152 | positive_classes = [0 for _ in range(args.num_class)] 153 | true_positive_classes = [0 for _ in range(args.num_class)] 154 | for i in range(len(truel)): 155 | gt_classes[truel[i]] += 1 156 | positive_classes[predl[i]] += 1 157 | true_positive_classes[truel[i]] += int(truel[i]==predl[i]) 158 | 159 | OA = sum(true_positive_classes)/float(len(truel)) 160 | iou_list = [] 161 | for i in range(args.num_class): 162 | iou = true_positive_classes[i]/float(gt_classes[i]+positive_classes[i]-true_positive_classes[i] + 1e-9) 163 | iou_list.append(iou) 164 | 165 | return iou_list, OA 166 | 167 | 168 | def compute_iou(args, predl, gt_label): 169 | per_class_iou = np.zeros(args.num_class) 170 | 171 | eps = 1e-8 172 | for i in range(args.num_class): 173 | per_class_iou[i] = (float(np.sum((predl==i)&(gt_label==i)))) / (float(np.sum((predl==i) | (gt_label==i))) + eps) 174 | 175 | OA = (float(np.sum(predl==gt_label))) / (float(len(gt_label)) + eps) 176 | 177 | return per_class_iou, OA 178 | 179 | 180 | def acc_report(args): 181 | result_files = glob.glob(os.path.join(args.outdir, '*.pred_true')) 182 | pred_true_labels_list = [] 183 | for i in range(len(result_files)): 184 | pred_true_labels_list.append(np.loadtxt(result_files[i])) 185 | print('%d/%d load'%(i, len(result_files))) 186 | pred_true_labels = np.concatenate(pred_true_labels_list, 0) 187 | 188 | print('computing iou...') 189 | iou_list, OA = compute_iou(args, np.uint8(pred_true_labels[:,0]), np.uint8(pred_true_labels[:,1])) 190 | 191 | print('iou: ') 192 | print(iou_list) 193 | print('mean iou: %f'% np.mean(iou_list)) 194 | print('OA: %f' %(OA)) 195 | 196 | 197 | if __name__ == "__main__": 198 | test() 199 | interpolate() 200 | acc_report() 201 | -------------------------------------------------------------------------------- /net_S3DIS/train.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import numpy as np 4 | import tensorflow as tf 5 | import glob 6 | import pdb 7 | 8 | import os 9 | import sys 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | ROOT_DIR = os.path.dirname(BASE_DIR) 12 | sys.path.append(BASE_DIR) 13 | sys.path.append(ROOT_DIR) 14 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 15 | import provider 16 | from model import * 17 | import tf_util 18 | 19 | 20 | def log_string(LOG_FOUT, out_str): 21 | LOG_FOUT.write(out_str+'\n') 22 | LOG_FOUT.flush() 23 | print(out_str) 24 | 25 | 26 | def train(args, graph_inf, net_inf): 27 | 28 | #creat a log_dir for log record 29 | #if not os.path.exists(args.log_dir): os.mkdir(args.log_dir) 30 | LOG_FOUT = open(os.path.join(args.log_dir, 'log_train.txt'), 'a') 31 | LOG_FOUT.write(str(args)+'\n') 32 | 33 | #load data 34 | train_data, train_label, train_spw, test_data, test_label, test_spw = load_traindata(args.data_path, args.roomlist_file, args.test_area) 35 | feature_channel = train_data.shape[-1] 36 | 37 | with tf.Graph().as_default(): 38 | with tf.device('/gpu:'+str(args.gpu)): 39 | pointclouds_pl, labels_pl, spws_pl, is_training_pl = placeholder_inputs(feature_channel, with_spw=True) 40 | 41 | # Note the global_step=batch parameter to minimize. 42 | # That tells the optimizer to helpfully increment the 'batch' parameter for you every time it trains. 43 | batch = tf.Variable(0) 44 | bn_decay = tf_util.get_bn_decay(batch, args.batch_size, float(args.decay_step)) 45 | tf.summary.scalar('bn_decay', bn_decay) 46 | 47 | # Get model and loss 48 | graph_prd, coarse_map = build_graph_pyramid(pointclouds_pl[...,0:3], graph_inf) 49 | logits = build_network(graph_prd, coarse_map, net_inf, pointclouds_pl, args.num_class, is_training_pl, bn_decay=bn_decay) 50 | loss = get_loss(logits, labels_pl) 51 | tf.summary.scalar('loss', loss) 52 | 53 | pred = tf.nn.softmax(logits) 54 | correct = tf.equal(tf.argmax(pred, 2), tf.to_int64(labels_pl)) 55 | accuracy = tf.reduce_sum(tf.cast(correct, tf.float32)) / float(args.batch_size*args.num_point) 56 | tf.summary.scalar('accuracy', accuracy) 57 | 58 | # Get training operator 59 | learning_rate = tf_util.get_learning_rate(batch, args.learning_rate, args.batch_size, args.decay_step, args.decay_rate) 60 | tf.summary.scalar('learning_rate', learning_rate) 61 | if args.optimizer == 'momentum': 62 | optimizer = tf.train.MomentumOptimizer(learning_rate, momentum=args.momentum) 63 | elif args.optimizer == 'adam': 64 | optimizer = tf.train.AdamOptimizer(learning_rate) 65 | extra_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 66 | with tf.control_dependencies(extra_update_ops): 67 | train_op = optimizer.minimize(loss, global_step=batch) 68 | 69 | # Add ops to save and restore all the variables. 70 | saver = tf.train.Saver() 71 | 72 | 73 | # Create a session 74 | config = tf.ConfigProto() 75 | config.gpu_options.allow_growth = True 76 | config.allow_soft_placement = True 77 | config.log_device_placement = True 78 | sess = tf.Session(config=config) 79 | 80 | # Add summary writers 81 | merged = tf.summary.merge_all() 82 | train_writer = tf.summary.FileWriter(os.path.join(args.log_dir, 'train'), sess.graph) 83 | test_writer = tf.summary.FileWriter(os.path.join(args.log_dir, 'test')) 84 | 85 | # Init variables 86 | init = tf.global_variables_initializer() 87 | sess.run(init, {is_training_pl:True}) 88 | 89 | if os.path.exists('log/checkpoint'): 90 | last_model = tf.train.latest_checkpoint('log/') 91 | saver.restore(sess, last_model) 92 | 93 | ops = {'pointclouds_pl': pointclouds_pl, 94 | 'labels_pl': labels_pl, 95 | 'spws_pl': spws_pl, 96 | 'is_training_pl': is_training_pl, 97 | 'pred': pred, 98 | 'loss': loss, 99 | 'train_op': train_op, 100 | 'merged': merged, 101 | 'step': batch} 102 | 103 | best_iou = -1 104 | for epoch in range(args.max_epoch): 105 | log_string(LOG_FOUT, '**** EPOCH %03d ****' % (epoch)) 106 | sys.stdout.flush() 107 | 108 | train_one_epoch(args, sess, ops, train_writer, train_data, train_label, train_spw, LOG_FOUT) 109 | 110 | # Save the variables to disk. 111 | if epoch % 3 == 0: 112 | iou = eval_one_epoch(args, sess, ops, test_writer, test_data, test_label, test_spw, LOG_FOUT) 113 | if iou>best_iou: 114 | best_iou = iou 115 | save_path = saver.save(sess, os.path.join(args.log_dir, "model.ckpt"), epoch) 116 | log_string(LOG_FOUT, "Model saved in file: %s" % save_path) 117 | #save_path = saver.save(sess, os.path.join(args.log_dir, "model.ckpt")) 118 | LOG_FOUT.close() #close log file 119 | 120 | 121 | def train_one_epoch(args, sess, ops, train_writer, train_data, train_label, train_spw, LOG_FOUT): 122 | """ ops: dict mapping from string to tf ops """ 123 | is_training = True 124 | 125 | log_string(LOG_FOUT, '----') 126 | current_data, current_label, current_spw, _ = provider.shuffle_data_with_spw(train_data[:,0:args.num_point,:], train_label, train_spw) 127 | current_data[:,:,0:3] = provider.rotate_point_cloud_along_z(current_data[:,:,0:3]) # rotate along z-axis 128 | 129 | file_size = current_data.shape[0] 130 | num_batches = file_size // args.batch_size 131 | 132 | total_correct = 0 133 | total_seen = 0 134 | loss_sum = 0 135 | 136 | for batch_idx in range(num_batches): 137 | if batch_idx % 100 == 0: 138 | print('Current batch/total batch num: %d/%d'%(batch_idx,num_batches)) 139 | start_idx = batch_idx * args.batch_size 140 | end_idx = (batch_idx+1) * args.batch_size 141 | 142 | feed_dict = {ops['pointclouds_pl']: current_data[start_idx:end_idx, :, :], 143 | ops['labels_pl']: current_label[start_idx:end_idx], 144 | ops['spws_pl']: current_spw[start_idx:end_idx], 145 | ops['is_training_pl']: is_training,} 146 | summary, step, _, loss_val, pred_val = sess.run([ops['merged'], ops['step'], ops['train_op'], ops['loss'], ops['pred']], 147 | feed_dict=feed_dict) 148 | 149 | train_writer.add_summary(summary, step) 150 | pred_val = np.argmax(pred_val, 2) 151 | correct = np.sum(pred_val == current_label[start_idx:end_idx]) 152 | total_correct += correct 153 | total_seen += (args.batch_size*args.num_point) 154 | loss_sum += loss_val 155 | print('loss: %f' % loss_val) 156 | print('accuracy: %f' % (correct/float(args.batch_size*args.num_point))) 157 | 158 | log_string(LOG_FOUT, 'mean loss: %f' % (loss_sum / float(num_batches))) 159 | log_string(LOG_FOUT, 'accuracy: %f' % (total_correct / float(total_seen))) 160 | 161 | 162 | def eval_one_epoch(args, sess, ops, test_writer, test_data, test_label, test_spw, LOG_FOUT): 163 | """ ops: dict mapping from string to tf ops """ 164 | is_training = False 165 | total_correct = 0 166 | total_seen = 0 167 | loss_sum = 0 168 | total_seen_class = [0 for _ in range(args.num_class)] 169 | total_correct_class = [0 for _ in range(args.num_class)] 170 | 171 | log_string(LOG_FOUT, '----') 172 | current_data = test_data[:,0:args.num_point,:] 173 | current_label = np.squeeze(test_label) 174 | 175 | file_size = current_data.shape[0] 176 | num_batches = file_size // args.batch_size 177 | 178 | gt_classes = [0 for _ in range(args.num_class)] 179 | positive_classes = [0 for _ in range(args.num_class)] 180 | true_positive_classes = [0 for _ in range(args.num_class)] 181 | 182 | for batch_idx in range(num_batches): 183 | if (batch_idx%20==0): 184 | print('eval: %d / %d' %(batch_idx, num_batches)) 185 | start_idx = batch_idx * args.batch_size 186 | end_idx = (batch_idx+1) * args.batch_size 187 | 188 | batch_label = current_label[start_idx:end_idx,:] 189 | batch_spw = test_spw[start_idx:end_idx,:] 190 | feed_dict = {ops['pointclouds_pl']: current_data[start_idx:end_idx, :, :], 191 | ops['labels_pl']: batch_label, 192 | ops['spws_pl']: batch_spw, 193 | ops['is_training_pl']: is_training} 194 | summary, step, loss_val, pred_val = sess.run([ops['merged'], ops['step'], ops['loss'], ops['pred']], 195 | feed_dict=feed_dict) 196 | test_writer.add_summary(summary, step) 197 | pred_val = np.argmax(pred_val, 2) 198 | correct = np.sum((pred_val == batch_label) & (batch_spw>0)) 199 | total_correct += correct 200 | total_seen += np.sum(batch_spw>0) 201 | loss_sum += (loss_val*args.batch_size) 202 | for i in range(start_idx, end_idx): 203 | for j in range(args.num_point): 204 | if test_spw[i,j]==0: continue 205 | l = current_label[i, j] 206 | total_seen_class[l] += 1 207 | total_correct_class[l] += (pred_val[i-start_idx, j] == l) 208 | 209 | gt_classes, positive_classes, true_positive_classes = provider.utils_iou(\ 210 | gt_classes, positive_classes, true_positive_classes, pred_val, batch_label, batch_spw) 211 | iou_list = [] 212 | for i in range(args.num_class): 213 | iou = true_positive_classes[i]/float(gt_classes[i]+positive_classes[i]-true_positive_classes[i]) 214 | iou_list.append(iou) 215 | 216 | log_string(LOG_FOUT, 'eval mean loss: %f' % (loss_sum / float(total_seen/args.num_point))) 217 | log_string(LOG_FOUT, 'eval accuracy: %f'% (total_correct / float(total_seen))) 218 | log_string(LOG_FOUT, 'eval avg class acc: %f' % (np.mean(np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float)))) 219 | log_string(LOG_FOUT, 'eval avg iou: %f' % (np.mean(iou_list))) 220 | 221 | return np.mean(iou_list) 222 | 223 | 224 | def load_traindata(data_path, roomlist_file, test_area): 225 | ALL_FILES = glob.glob(data_path + '/*.h5') 226 | room_filelist = [line.rstrip() for line in open(roomlist_file)] 227 | 228 | # Load ALL data 229 | data_batch_list = [] 230 | label_batch_list = [] 231 | spw_batch_list = [] 232 | for h5_filename in ALL_FILES: 233 | data_batch, label_batch, spw_batch = provider.loadDataFile_with_spw(h5_filename) 234 | data_batch_list.append(data_batch) 235 | label_batch_list.append(label_batch) 236 | spw_batch_list.append(spw_batch) 237 | data_batches = np.concatenate(data_batch_list, 0) 238 | label_batches = np.concatenate(label_batch_list, 0) 239 | spw_batches = np.concatenate(spw_batch_list, 0) 240 | 241 | test_area = 'Area_'+str(test_area) 242 | train_idxs = [] 243 | test_idxs = [] 244 | for i,room_name in enumerate(room_filelist): 245 | if test_area in room_name: 246 | test_idxs.append(i) 247 | else: 248 | train_idxs.append(i) 249 | 250 | train_data = data_batches[train_idxs,...] 251 | train_label = label_batches[train_idxs] 252 | train_spw = spw_batches[train_idxs] 253 | 254 | test_data = data_batches[test_idxs,...] 255 | test_label = label_batches[test_idxs] 256 | test_spw = spw_batches[test_idxs] 257 | print(train_data.shape, train_label.shape) 258 | print(test_data.shape, test_label.shape) 259 | 260 | return train_data, train_label, train_spw, test_data, test_label, test_spw 261 | -------------------------------------------------------------------------------- /test_utils/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | 2 | cmake_minimum_required(VERSION 2.8 FATAL_ERROR) 3 | 4 | project(three_interpolate) 5 | 6 | # packages 7 | find_package(CUDA REQUIRED) 8 | 9 | # nvcc flags 10 | set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} -gencode arch=compute_50,code=sm_50) 11 | 12 | #set(CUDA_NVCC_FLAGS -gencode arch=compute_20,code=sm_20;-G;-g) 13 | #set(CUDA_NVCC_FLAGS -gencode arch=compute_52,code=sm_52;-G;-g) 14 | 15 | FIND_PACKAGE( OpenMP REQUIRED) 16 | if(OPENMP_FOUND) 17 | message("OPENMP FOUND") 18 | set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}") 19 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}") 20 | set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} ${OpenMP_EXE_LINKER_FLAGS}") 21 | endif() 22 | 23 | #set(flannlibs "/usr/lib/x86_64-linux-gnu/libflann.so") 24 | include_directories("/usr/local/cuda-9.0/include" "usr/include/flann") 25 | link_directories("/usr/local/cuda-9.0/lib64/") 26 | 27 | set(flannlibs "/usr/lib/x86_64-linux-gnu/libflann.so") 28 | set(src three_interpolate.cpp interpolate_gpu.cu) 29 | 30 | CUDA_ADD_LIBRARY(three_interpolate SHARED ${src}) 31 | target_link_libraries (three_interpolate ${flannlibs}) 32 | -------------------------------------------------------------------------------- /test_utils/README: -------------------------------------------------------------------------------- 1 | This folder includes code for label interpolation. 2 | 3 | Since we only predict the label of part of the points sampled from the orginal point cloud, 4 | we have to calculate the label of the orignal point cloud with interpolation algorithm. 5 | Similar as we do for feature interpolation in the network, we also interpolate the label of each point (of the orignal point cloud) according to its three nearest neighbors 6 | 7 | To compile these codes, run 8 | sh cmake.sh 9 | -------------------------------------------------------------------------------- /test_utils/cmake.sh: -------------------------------------------------------------------------------- 1 | rm -rf build 2 | mkdir build 3 | cd build 4 | cmake .. 5 | make 6 | -------------------------------------------------------------------------------- /test_utils/interpolate_gpu.cu: -------------------------------------------------------------------------------- 1 | #include "cuda_runtime.h" 2 | #include "device_launch_parameters.h" 3 | #include "stdio.h" 4 | 5 | //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// 6 | // input: features (n,c), idx (n,3), weight (n,3) 7 | // output: probs (n,c) 8 | __global__ void interpolate_kernel(int n, int c, int k_num, float *features, float *weight, int *idx, float *probs) { 9 | 10 | for(int i= blockIdx.x; i>>(n_host, c_host, 3, features, weight, idx, probs); 42 | error = cudaDeviceSynchronize(); 43 | if(error != cudaSuccess){ 44 | printf("code: %d, reason: %s\n",error,cudaGetErrorString(error)); 45 | } 46 | 47 | 48 | cudaMemcpy(probs_host, probs, sizeof(float)*n_host*c_host, cudaMemcpyDeviceToHost); 49 | 50 | cudaFree(weight); 51 | cudaFree(features); 52 | cudaFree(probs); 53 | cudaFree(idx); 54 | } 55 | 56 | 57 | extern "C" void filterLauncher(int n_host, int c_host, int k_num, int iter_num, float *features_host, int *idx_host, float *weight_host, float *probs_host){ 58 | //int *n_dev, *c_dev; 59 | float *weight, *features, *probs; 60 | int *idx; 61 | cudaError_t error; 62 | 63 | cudaMalloc((void**)&weight, sizeof(float)* n_host*k_num); 64 | cudaMalloc((void**)&idx, sizeof(int)* n_host*k_num); 65 | cudaMalloc((void**)&features, sizeof(float)* n_host*c_host); 66 | cudaMalloc((void**)&probs, sizeof(float)* n_host*c_host); 67 | 68 | cudaMemcpy(weight, weight_host, sizeof(float)* n_host*k_num, cudaMemcpyHostToDevice); 69 | cudaMemcpy(idx, idx_host, sizeof(int)* n_host*k_num, cudaMemcpyHostToDevice); 70 | cudaMemcpy(features, features_host, sizeof(float)* n_host*c_host, cudaMemcpyHostToDevice); 71 | cudaMemcpy(probs, probs_host, sizeof(float)*n_host*c_host, cudaMemcpyHostToDevice); 72 | 73 | dim3 grid(32768, 1, 1), block(c_host, 1, 1); 74 | for(int i=0;i>>(n_host, c_host, k_num, features, weight, idx, probs); 76 | error = cudaDeviceSynchronize(); 77 | if(error != cudaSuccess){ 78 | printf("code: %d, reason: %s\n",error,cudaGetErrorString(error)); 79 | } 80 | } 81 | 82 | cudaMemcpy(probs_host, probs, sizeof(float)*n_host*c_host, cudaMemcpyDeviceToHost); 83 | 84 | cudaFree(weight); 85 | cudaFree(features); 86 | cudaFree(probs); 87 | cudaFree(idx); 88 | } 89 | -------------------------------------------------------------------------------- /test_utils/three_interpolate.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | #include 15 | #include 16 | 17 | using namespace std; 18 | 19 | 20 | void weight_cal(int n, float *rgb, float *dist_spatial, int *idx, float *weight) { 21 | #pragma omp parallel for 22 | for(int i=0; i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include "tensorflow/core/framework/op.h" 7 | #include "tensorflow/core/framework/op_kernel.h" 8 | #include "tensorflow/core/framework/shape_inference.h" 9 | #include "tensorflow/core/framework/common_shape_fns.h" 10 | #include 11 | using namespace tensorflow; 12 | 13 | REGISTER_OP("GatherPoint") 14 | .Input("inp: float32") 15 | .Input("idx: int32") 16 | .Output("out: float32") 17 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 18 | ::tensorflow::shape_inference::ShapeHandle dims0; // batch_size * ndataset * 3 19 | c->WithRank(c->input(0), 3, &dims0); 20 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * npoints 21 | c->WithRank(c->input(1), 2, &dims1); 22 | // batch_size * npoints * channel 23 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims0, 0), c->Dim(dims1, 1), c->Dim(dims0, 2)}); 24 | c->set_output(0, output); 25 | return Status::OK(); 26 | }); 27 | REGISTER_OP("GatherPointGrad") 28 | .Input("inp: float32") 29 | .Input("idx: int32") 30 | .Input("out_g: float32") 31 | .Output("inp_g: float32") 32 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 33 | c->set_output(0, c->input(0)); 34 | return Status::OK(); 35 | }); 36 | REGISTER_OP("GatherId") 37 | .Input("inp: int32") 38 | .Input("idx: int32") 39 | .Output("out: int32") 40 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 41 | ::tensorflow::shape_inference::ShapeHandle dims0; // batch_size * ndataset * 3 42 | c->WithRank(c->input(0), 3, &dims0); 43 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * npoints 44 | c->WithRank(c->input(1), 2, &dims1); 45 | // batch_size * npoints * channel 46 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims0, 0), c->Dim(dims1, 1), c->Dim(dims0, 2)}); 47 | c->set_output(0, output); 48 | return Status::OK(); 49 | }); 50 | 51 | REGISTER_OP("GroupPoint") 52 | .Input("points: float32") 53 | .Input("idx: int32") 54 | .Output("out: float32") 55 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 56 | ::tensorflow::shape_inference::ShapeHandle dims0; // batch_size * ndataset * channels 57 | c->WithRank(c->input(0), 3, &dims0); 58 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * npoints * nsample 59 | c->WithRank(c->input(1), 3, &dims1); 60 | // batch_size * npoints * nsample * channels 61 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), c->Dim(dims1, 1), c->Dim(dims1, 2), c->Dim(dims0, 2)}); 62 | c->set_output(0, output); 63 | return Status::OK(); 64 | }); 65 | REGISTER_OP("GroupPointGrad") 66 | .Input("points: float32") 67 | .Input("idx: int32") 68 | .Input("grad_out: float32") 69 | .Output("grad_points: float32") 70 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 71 | c->set_output(0, c->input(0)); 72 | return Status::OK(); 73 | }); 74 | 75 | REGISTER_OP("GatherEigvector") 76 | .Input("eigvs: float32") 77 | .Input("idx: int32") 78 | .Output("out: float32") 79 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 80 | ::tensorflow::shape_inference::ShapeHandle dims0; // batch_size * ndataset * 3 *3 81 | c->WithRank(c->input(0), 4, &dims0); 82 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims0, 0), c->Dim(dims0, 1), 3}); 83 | c->set_output(0, output); 84 | return Status::OK(); 85 | }); 86 | REGISTER_OP("GatherEigvectorGrad") 87 | .Input("eigvs: float32") 88 | .Input("idx: int32") 89 | .Input("grad_out: float32") 90 | .Output("grad_eigvs: float32") 91 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 92 | c->set_output(0, c->input(0)); 93 | return Status::OK(); 94 | }); 95 | 96 | void gatherpointLauncher(int b,int n,int m,int channel, const float * inp,const int * idx,float * out); 97 | class GatherPointGpuOp: public OpKernel{ 98 | public: 99 | explicit GatherPointGpuOp(OpKernelConstruction * context):OpKernel(context){} 100 | void Compute(OpKernelContext * context)override{ 101 | const Tensor& inp_tensor=context->input(0); 102 | OP_REQUIRES(context,inp_tensor.dims()==3,errors::InvalidArgument("GatherPoint expects (batch_size,num_points,channel) inp shape")); 103 | int b=inp_tensor.shape().dim_size(0); 104 | int n=inp_tensor.shape().dim_size(1); 105 | int channel=inp_tensor.shape().dim_size(2); 106 | const Tensor& idx_tensor=context->input(1); 107 | OP_REQUIRES(context,idx_tensor.dims()==2 && idx_tensor.shape().dim_size(0)==b,errors::InvalidArgument("GatherPoint expects (batch_size,num_result) idx shape")); 108 | int m=idx_tensor.shape().dim_size(1); 109 | auto inp_flat=inp_tensor.flat(); 110 | const float * inp=&(inp_flat(0)); 111 | auto idx_flat=idx_tensor.flat(); 112 | const int * idx=&(idx_flat(0)); 113 | Tensor * out_tensor=NULL; 114 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m,channel},&out_tensor)); 115 | auto out_flat=out_tensor->flat(); 116 | float * out=&(out_flat(0)); 117 | gatherpointLauncher(b,n,m,channel,inp,idx,out); 118 | } 119 | }; 120 | REGISTER_KERNEL_BUILDER(Name("GatherPoint").Device(DEVICE_GPU),GatherPointGpuOp); 121 | 122 | void scatteraddpointLauncher(int b,int n,int m,int channel, const float * out_g,const int * idx,float * inp_g); 123 | class GatherPointGradGpuOp: public OpKernel{ 124 | public: 125 | explicit GatherPointGradGpuOp(OpKernelConstruction * context):OpKernel(context){} 126 | void Compute(OpKernelContext * context)override{ 127 | const Tensor& inp_tensor=context->input(0); 128 | OP_REQUIRES(context,inp_tensor.dims()==3,errors::InvalidArgument("GatherPointGradGpuOp expects (batch_size,num_points,channel) inp")); 129 | int b=inp_tensor.shape().dim_size(0); 130 | int n=inp_tensor.shape().dim_size(1); 131 | int channel=inp_tensor.shape().dim_size(2); 132 | const Tensor& idx_tensor=context->input(1); 133 | OP_REQUIRES(context,idx_tensor.dims()==2 && idx_tensor.shape().dim_size(0)==b,errors::InvalidArgument("GatherPointGradGpuOp expects (batch_size,num_result) idx shape")); 134 | int m=idx_tensor.shape().dim_size(1); 135 | auto inp_flat=inp_tensor.flat(); 136 | const float * inp=&(inp_flat(0)); 137 | auto idx_flat=idx_tensor.flat(); 138 | const int * idx=&(idx_flat(0)); 139 | const Tensor& out_g_tensor=context->input(2); 140 | OP_REQUIRES(context,out_g_tensor.dims()==3 && out_g_tensor.shape().dim_size(0)==b && out_g_tensor.shape().dim_size(1)==m,errors::InvalidArgument("GatherPointGradGpuOp expects (batch_size,num_result,channel) out_g shape")); 141 | auto out_g_flat=out_g_tensor.flat(); 142 | const float * out_g=&(out_g_flat(0)); 143 | Tensor * inp_g_tensor=nullptr; 144 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,n,channel},&inp_g_tensor)); 145 | auto inp_g_flat=inp_g_tensor->flat(); 146 | float * inp_g=&(inp_g_flat(0)); 147 | cudaMemset(inp_g,0,sizeof(float)*b*n*channel); 148 | scatteraddpointLauncher(b,n,m,channel,out_g,idx,inp_g); 149 | } 150 | }; 151 | REGISTER_KERNEL_BUILDER(Name("GatherPointGrad").Device(DEVICE_GPU),GatherPointGradGpuOp); 152 | 153 | void gatheridLauncher(int b,int n,int m,int channel, const int * inp,const int * idx,int * out); 154 | class GatherIdGpuOp: public OpKernel{ 155 | public: 156 | explicit GatherIdGpuOp(OpKernelConstruction * context):OpKernel(context){} 157 | void Compute(OpKernelContext * context)override{ 158 | const Tensor& inp_tensor=context->input(0); 159 | OP_REQUIRES(context,inp_tensor.dims()==3,errors::InvalidArgument("GatherId expects (batch_size,num_points,num_sampling1) inp shape")); 160 | int b=inp_tensor.shape().dim_size(0); 161 | int n=inp_tensor.shape().dim_size(1); 162 | int channel=inp_tensor.shape().dim_size(2); 163 | const Tensor& idx_tensor=context->input(1); 164 | OP_REQUIRES(context,idx_tensor.dims()==2 && idx_tensor.shape().dim_size(0)==b,errors::InvalidArgument("GatherId expects (batch_size,num_sampling2) idx shape")); 165 | int m=idx_tensor.shape().dim_size(1); 166 | auto inp_flat=inp_tensor.flat(); 167 | const int * inp=&(inp_flat(0)); 168 | auto idx_flat=idx_tensor.flat(); 169 | const int * idx=&(idx_flat(0)); 170 | Tensor * out_tensor=NULL; 171 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m,channel},&out_tensor)); 172 | auto out_flat=out_tensor->flat(); 173 | int * out=&(out_flat(0)); 174 | gatheridLauncher(b,n,m,channel,inp,idx,out); 175 | } 176 | }; 177 | REGISTER_KERNEL_BUILDER(Name("GatherId").Device(DEVICE_GPU),GatherIdGpuOp); 178 | 179 | void groupPointLauncher(int b, int n, int c, int m, int nsample, const float *points, const int *idx, float *out); 180 | class GroupPointGpuOp: public OpKernel{ 181 | public: 182 | explicit GroupPointGpuOp(OpKernelConstruction * context):OpKernel(context){} 183 | 184 | void Compute(OpKernelContext * context) override { 185 | const Tensor& points_tensor=context->input(0); 186 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("GroupPoint expects (batch_size, num_points, channel) points shape")); 187 | int b = points_tensor.shape().dim_size(0); 188 | int n = points_tensor.shape().dim_size(1); 189 | int c = points_tensor.shape().dim_size(2); 190 | 191 | const Tensor& idx_tensor=context->input(1); 192 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b, errors::InvalidArgument("GroupPoint expects (batch_size, npoints, nsample) idx shape")); 193 | int m = idx_tensor.shape().dim_size(1); 194 | int nsample = idx_tensor.shape().dim_size(2); 195 | 196 | Tensor * out_tensor = nullptr; 197 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,m,nsample,c}, &out_tensor)); 198 | 199 | auto points_flat = points_tensor.flat(); 200 | const float *points = &(points_flat(0)); 201 | auto idx_flat = idx_tensor.flat(); 202 | const int *idx = &(idx_flat(0)); 203 | auto out_flat = out_tensor->flat(); 204 | float *out = &(out_flat(0)); 205 | groupPointLauncher(b,n,c,m,nsample,points,idx,out); 206 | } 207 | }; 208 | REGISTER_KERNEL_BUILDER(Name("GroupPoint").Device(DEVICE_GPU),GroupPointGpuOp); 209 | 210 | void groupPointGradLauncher(int b, int n, int c, int m, int nsample, const float *grad_out, const int *idx, float *grad_points); 211 | class GroupPointGradGpuOp: public OpKernel{ 212 | public: 213 | explicit GroupPointGradGpuOp(OpKernelConstruction * context):OpKernel(context){} 214 | 215 | void Compute(OpKernelContext * context) override { 216 | const Tensor& points_tensor=context->input(0); 217 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("GroupPointGrad expects (batch_size, num_points, channel) points shape")); 218 | int b = points_tensor.shape().dim_size(0); 219 | int n = points_tensor.shape().dim_size(1); 220 | int c = points_tensor.shape().dim_size(2); 221 | 222 | const Tensor& idx_tensor=context->input(1); 223 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b, errors::InvalidArgument("GroupPointGrad expects (batch_size, npoints, nsample) idx shape")); 224 | int m = idx_tensor.shape().dim_size(1); 225 | int nsample = idx_tensor.shape().dim_size(2); 226 | 227 | const Tensor& grad_out_tensor=context->input(2); 228 | OP_REQUIRES(context,grad_out_tensor.dims()==4 && grad_out_tensor.shape().dim_size(0)==b && grad_out_tensor.shape().dim_size(1)==m && grad_out_tensor.shape().dim_size(2)==nsample && grad_out_tensor.shape().dim_size(3)==c, errors::InvalidArgument("GroupPointGrad expects (batch_size, npoints, nsample, channel) grad_out shape")); 229 | 230 | Tensor * grad_points_tensor = nullptr; 231 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,n,c}, &grad_points_tensor)); 232 | 233 | auto points_flat = points_tensor.flat(); 234 | const float *points = &(points_flat(0)); 235 | auto idx_flat = idx_tensor.flat(); 236 | const int *idx = &(idx_flat(0)); 237 | auto grad_out_flat = grad_out_tensor.flat(); 238 | const float *grad_out = &(grad_out_flat(0)); 239 | auto grad_points_flat = grad_points_tensor->flat(); 240 | float *grad_points = &(grad_points_flat(0)); 241 | cudaMemset(grad_points, 0, sizeof(float)*b*n*c); 242 | groupPointGradLauncher(b,n,c,m,nsample,grad_out,idx,grad_points); 243 | } 244 | }; 245 | REGISTER_KERNEL_BUILDER(Name("GroupPointGrad").Device(DEVICE_GPU),GroupPointGradGpuOp); 246 | 247 | 248 | void gatherEigvectorLauncher(int b, int n, const float *eigvs, const int *idx, float *out); 249 | class GatherEigvectorGpuOp: public OpKernel{ 250 | public: 251 | explicit GatherEigvectorGpuOp(OpKernelConstruction * context):OpKernel(context){} 252 | 253 | void Compute(OpKernelContext * context) override { 254 | const Tensor& eigvs_tensor=context->input(0); 255 | OP_REQUIRES(context, eigvs_tensor.dims()==4, errors::InvalidArgument("GatherEigvector expects (batch_size, num_points, 3, 3) points shape")); 256 | int b = eigvs_tensor.shape().dim_size(0); 257 | int n = eigvs_tensor.shape().dim_size(1); 258 | const Tensor& idx_tensor=context->input(1); 259 | OP_REQUIRES(context,idx_tensor.dims()==2 && idx_tensor.shape().dim_size(0)==b, errors::InvalidArgument("GatherEigvector expects (batch_size, npoints) idx shape")); 260 | 261 | Tensor * out_tensor = nullptr; 262 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,n,3}, &out_tensor)); 263 | 264 | auto eigvs_flat = eigvs_tensor.flat(); 265 | const float *eigvs = &(eigvs_flat(0)); 266 | auto idx_flat = idx_tensor.flat(); 267 | const int *idx = &(idx_flat(0)); 268 | auto out_flat = out_tensor->flat(); 269 | float *out = &(out_flat(0)); 270 | gatherEigvectorLauncher(b,n,eigvs,idx,out); 271 | } 272 | }; 273 | REGISTER_KERNEL_BUILDER(Name("GatherEigvector").Device(DEVICE_GPU),GatherEigvectorGpuOp); 274 | 275 | 276 | void gatherEigvectorLauncher(int b, int n, const float *grad_out, const int *idx, float *grad_eigvs); 277 | class GatherEigvectorGradGpuOp: public OpKernel{ 278 | public: 279 | explicit GatherEigvectorGradGpuOp(OpKernelConstruction * context):OpKernel(context){} 280 | 281 | void Compute(OpKernelContext * context) override { 282 | const Tensor& eigvs_tensor=context->input(0); 283 | OP_REQUIRES(context, eigvs_tensor.dims()==4, errors::InvalidArgument("GatherEigvectorGrad expects (batch_size, num_points, 3, 3) points shape")); 284 | int b = eigvs_tensor.shape().dim_size(0); 285 | int n = eigvs_tensor.shape().dim_size(1); 286 | 287 | const Tensor& idx_tensor=context->input(1); 288 | OP_REQUIRES(context,idx_tensor.dims()==2 && idx_tensor.shape().dim_size(0)==b, errors::InvalidArgument("GatherEigvectorGrad expects (batch_size, num_points) idx shape")); 289 | 290 | const Tensor& grad_out_tensor=context->input(2); 291 | OP_REQUIRES(context,grad_out_tensor.dims()==3 && grad_out_tensor.shape().dim_size(0)==b && grad_out_tensor.shape().dim_size(1)==n, errors::InvalidArgument("GatherEigvectorGrad expects (batch_size, npoints, nsample, channel) grad_out shape")); 292 | 293 | Tensor * grad_eigvs_tensor = nullptr; 294 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,n,3,3}, &grad_eigvs_tensor)); 295 | 296 | auto eigvs_flat = eigvs_tensor.flat(); 297 | const float *eigvs = &(eigvs_flat(0)); 298 | auto idx_flat = idx_tensor.flat(); 299 | const int *idx = &(idx_flat(0)); 300 | auto grad_out_flat = grad_out_tensor.flat(); 301 | const float *grad_out = &(grad_out_flat(0)); 302 | auto grad_eigvs_flat = grad_eigvs_tensor->flat(); 303 | float *grad_eigvs = &(grad_eigvs_flat(0)); 304 | cudaMemset(grad_eigvs, 0, sizeof(float)*b*n*3*3); 305 | gatherEigvectorLauncher(b,n,grad_out,idx,grad_eigvs); 306 | } 307 | }; 308 | REGISTER_KERNEL_BUILDER(Name("GatherEigvectorGrad").Device(DEVICE_GPU),GatherEigvectorGradGpuOp); -------------------------------------------------------------------------------- /tf_ops/tf_grad_op_test.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import random 4 | from tf_ops import knn_search, three_interpolate, group_point, gather_point, FPS_downsample, query_ball_point, query_circle_point, shuffle_ids 5 | import pdb 6 | 7 | class TfOpsTest(tf.test.TestCase): 8 | def test(self): 9 | pass 10 | 11 | def test_grad(self): 12 | with self.test_session(): 13 | with tf.device('/gpu:0'): 14 | points = tf.constant(np.random.random((1,68,16)).astype('float32')) 15 | print(points) 16 | xyz1 = tf.constant(np.random.random((1,128,3)).astype('float32')) 17 | xyz2 = tf.constant(np.random.random((1,68,3)).astype('float32')) 18 | 19 | print ('---------------------group------------------') 20 | idxs, dist = knn_search(xyz1, xyz2, 5) 21 | #idxs, dist = query_ball_point(0.05, 5, xyz1, xyz2) 22 | idxs, dist = query_circle_point(0.0, 1.1, 5, xyz1, xyz2) 23 | with tf.device('/cpu:0'): 24 | shuffled_ids = shuffle_ids(xyz1) 25 | with tf.Session() as sess: 26 | idx_out, dist_, shuffled_ids_ = sess.run([idxs, dist, shuffled_ids]) 27 | idx= tf.constant(idx_out) 28 | 29 | print(idx) 30 | grouped_points = group_point(xyz1, idxs) 31 | print(grouped_points) 32 | err = tf.test.compute_gradient_error(xyz1, (1,128,3), grouped_points, (1,68,5,3)) 33 | print(err) 34 | self.assertLess(err, 1e-4) 35 | 36 | print ('---------------------interpolate------------------') 37 | idx, dist = knn_search(xyz1, xyz2, 3) 38 | print(idx) 39 | weight = tf.ones_like(dist)/3.0 40 | interpolated_points = three_interpolate(points, idx, weight) 41 | print(interpolated_points) 42 | err = tf.test.compute_gradient_error(points, (1,68,16), interpolated_points, (1,128,16)) 43 | print(err) 44 | self.assertLess(err, 1e-4) 45 | 46 | print ('---------------------gather------------------') 47 | idxs = FPS_downsample(xyz1, 100) 48 | with tf.Session() as sess: 49 | idx_out = sess.run(idxs) 50 | idx= tf.constant(idx_out) 51 | print(idx) 52 | gathered_points = gather_point(xyz1, idx) 53 | print(gathered_points) 54 | err = tf.test.compute_gradient_error(xyz1, (1,128,3), gathered_points, (1,100,3)) 55 | print(err) 56 | self.assertLess(err, 1e-4) 57 | 58 | 59 | if __name__=='__main__': 60 | ''' 61 | xyz1 = tf.constant(np.random.random((1,128,3)).astype('float32')) 62 | idx = FPS_downsample(xyz1, 100) 63 | with tf.device('/gpu:0'): 64 | with tf.Session() as sess: 65 | idx_out = sess.run(idx) 66 | pdb.set_trace() 67 | ''' 68 | tf.test.main() 69 | -------------------------------------------------------------------------------- /tf_ops/tf_interpolate.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include "tensorflow/core/framework/op.h" 7 | #include "tensorflow/core/framework/op_kernel.h" 8 | #include "tensorflow/core/framework/shape_inference.h" 9 | #include "tensorflow/core/framework/common_shape_fns.h" 10 | #include 11 | using namespace tensorflow; 12 | 13 | REGISTER_OP("ThreeInterpolate") 14 | .Input("points: float32") 15 | .Input("idx: int32") 16 | .Input("weight: float32") 17 | .Output("out: float32") 18 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 19 | ::tensorflow::shape_inference::ShapeHandle dims1; // (b,m,c) 20 | c->WithRank(c->input(0), 3, &dims1); 21 | ::tensorflow::shape_inference::ShapeHandle dims2; // (b,n,3) 22 | c->WithRank(c->input(1), 3, &dims2); 23 | // (b,n,c) 24 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), c->Dim(dims2, 1), c->Dim(dims1, 2)}); 25 | c->set_output(0, output); 26 | return Status::OK(); 27 | }); 28 | REGISTER_OP("ThreeInterpolateGrad") 29 | .Input("points: float32") 30 | .Input("idx: int32") 31 | .Input("weight: float32") 32 | .Input("grad_out: float32") 33 | .Output("grad_points: float32") 34 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 35 | c->set_output(0, c->input(0)); 36 | return Status::OK(); 37 | }); 38 | 39 | 40 | void threeinterpolateLauncher(int b, int m, int c, int n, const float *points, const int *idx, const float *weight, float *out); 41 | class ThreeInterpolateGPUOp: public OpKernel{ 42 | public: 43 | explicit ThreeInterpolateGPUOp(OpKernelConstruction * context):OpKernel(context){} 44 | 45 | void Compute(OpKernelContext * context) override { 46 | const Tensor& points_tensor=context->input(0); 47 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("ThreeInterpolate expects (b,m,c) points shape")); 48 | int b = points_tensor.shape().dim_size(0); 49 | int m = points_tensor.shape().dim_size(1); 50 | int c = points_tensor.shape().dim_size(2); 51 | 52 | const Tensor& idx_tensor=context->input(1); 53 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b && idx_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeInterpolate expects (b,n,3) idx shape")); 54 | int n = idx_tensor.shape().dim_size(1); 55 | const Tensor& weight_tensor=context->input(2); 56 | OP_REQUIRES(context,weight_tensor.dims()==3 && weight_tensor.shape().dim_size(0)==b && weight_tensor.shape().dim_size(1)==n && weight_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeInterpolate expects (b,n,3) weight shape")); 57 | 58 | Tensor * out_tensor = nullptr; 59 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,n,c}, &out_tensor)); 60 | 61 | auto points_flat = points_tensor.flat(); 62 | const float *points = &(points_flat(0)); 63 | auto idx_flat = idx_tensor.flat(); 64 | const int *idx = &(idx_flat(0)); 65 | auto weight_flat = weight_tensor.flat(); 66 | const float *weight = &(weight_flat(0)); 67 | auto out_flat = out_tensor->flat(); 68 | float *out = &(out_flat(0)); 69 | threeinterpolateLauncher(b,m,c,n,points,idx,weight,out); 70 | } 71 | }; 72 | REGISTER_KERNEL_BUILDER(Name("ThreeInterpolate").Device(DEVICE_GPU),ThreeInterpolateGPUOp); 73 | 74 | 75 | void threeinterpolategradLauncher(int b, int n, int c, int m, const float *grad_out, const int *idx, const float *weight, float *grad_points); 76 | class ThreeInterpolateGradGPUOp: public OpKernel{ 77 | public: 78 | explicit ThreeInterpolateGradGPUOp(OpKernelConstruction * context):OpKernel(context){} 79 | 80 | void Compute(OpKernelContext * context) override { 81 | const Tensor& points_tensor=context->input(0); 82 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("ThreeInterpolateGrad expects (b,m,c) points shape")); 83 | int b = points_tensor.shape().dim_size(0); 84 | int m = points_tensor.shape().dim_size(1); 85 | int c = points_tensor.shape().dim_size(2); 86 | 87 | const Tensor& idx_tensor=context->input(1); 88 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b, errors::InvalidArgument("ThreeInterpolateGrad expects (b,n,3) idx shape")); 89 | int n = idx_tensor.shape().dim_size(1); 90 | const Tensor& weight_tensor=context->input(2); 91 | OP_REQUIRES(context,weight_tensor.dims()==3 && weight_tensor.shape().dim_size(0)==b && weight_tensor.shape().dim_size(1)==n && weight_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeInterpolateGrad expects (b,n,3) weight shape")); 92 | 93 | const Tensor& grad_out_tensor=context->input(3); 94 | OP_REQUIRES(context,grad_out_tensor.dims()==3 && grad_out_tensor.shape().dim_size(0)==b && grad_out_tensor.shape().dim_size(1)==n && grad_out_tensor.shape().dim_size(2)==c, errors::InvalidArgument("ThreeInterpolateGrad expects (b,n,c) grad_out shape")); 95 | 96 | Tensor * grad_points_tensor = nullptr; 97 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,m,c}, &grad_points_tensor)); 98 | 99 | auto points_flat = points_tensor.flat(); 100 | const float *points = &(points_flat(0)); 101 | auto idx_flat = idx_tensor.flat(); 102 | const int *idx = &(idx_flat(0)); 103 | auto weight_flat = weight_tensor.flat(); 104 | const float *weight = &(weight_flat(0)); 105 | auto grad_out_flat = grad_out_tensor.flat(); 106 | const float *grad_out = &(grad_out_flat(0)); 107 | auto grad_points_flat = grad_points_tensor->flat(); 108 | float *grad_points = &(grad_points_flat(0)); 109 | cudaMemset(grad_points, 0, sizeof(float)*b*m*c); 110 | threeinterpolategradLauncher(b,n,c,m,grad_out,idx,weight,grad_points); 111 | } 112 | }; 113 | REGISTER_KERNEL_BUILDER(Name("ThreeInterpolateGrad").Device(DEVICE_GPU),ThreeInterpolateGradGPUOp); 114 | 115 | 116 | -------------------------------------------------------------------------------- /tf_ops/tf_ops.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.python.framework import ops 3 | import sys 4 | import os 5 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 6 | sys.path.append(BASE_DIR) 7 | tf_op_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_op_so.so')) 8 | 9 | def query_ball_pnet_point(radius, nsample, xyz1, xyz2): 10 | ''' 11 | Input: 12 | radius: float32, ball search radius 13 | nsample: int32, number of points selected in each ball region 14 | xyz1: (batch_size, ndataset, 3) float32 array, input points 15 | xyz2: (batch_size, npoint, 3) float32 array, query points 16 | Output: 17 | idx: (batch_size, npoint, nsample) int32 array, indices to input points 18 | pts_cnt: (batch_size, npoint) int32 array, number of unique points in each local region 19 | ''' 20 | #return grouping_module.query_ball_point(radius, nsample, xyz1, xyz2) 21 | return tf_op_module.query_ball_pnet_point(xyz1, xyz2, radius, nsample) 22 | ops.NoGradient('QueryBallPnetPoint') 23 | 24 | def rand_seeds(inp): 25 | ''' 26 | input: 27 | batch_size * ndataset * 3 float32 28 | returns: 29 | batch_size int32 30 | ''' 31 | return tf_op_module.rand_seeds(inp) 32 | ops.NoGradient('RandSeeds') 33 | 34 | def shuffle_ids(inp): 35 | ''' 36 | input: 37 | batch_size * ndataset * 3 float32 38 | returns: 39 | batch_size * ndataset int32 40 | ''' 41 | return tf_op_module.shuffle_ids(inp) 42 | ops.NoGradient('ShuffleIds') 43 | 44 | def FPS_downsample(inp, stride): 45 | ''' 46 | input: 47 | batch_size * ndataset * 3 float32 48 | int: downsampling factor 49 | returns: 50 | batch_size * npoint int32 51 | 52 | ''' 53 | seed = rand_seeds(inp) 54 | return tf_op_module.farthest_point_sample(inp, seed, stride) 55 | ops.NoGradient('FarthestPointSample') 56 | 57 | 58 | def query_ball_knn(data_xyz, search_xyz, radius, nsample): 59 | ''' 60 | Input: 61 | radius: float32, ball search radius 62 | nsample: int32, number of points selected in each ball region 63 | data_xyz: (batch_size, ndataset, 3) float32 array, input points 64 | search_xyz: (batch_size, npoint, 3) float32 array, query points 65 | Output: 66 | idx: (batch_size, npoint, nsample) int32 array, indices to input points 67 | pts_cnt: (batch_size, npoint) int32 array, number of unique points in each local region 68 | ''' 69 | 70 | with tf.device('/cpu:0'): 71 | shuffled_ids = shuffle_ids(data_xyz) 72 | return tf_op_module.query_ball_knn(data_xyz, search_xyz, shuffled_ids, radius, nsample) 73 | ops.NoGradient('QueryBallKnn') 74 | 75 | 76 | def query_ball(data_xyz, search_xyz, radius, nsample): 77 | ''' 78 | Input: 79 | radius: float32, ball search radius 80 | nsample: int32, number of points selected in each ball region 81 | data_xyz: (batch_size, ndataset, 3) float32 array, input points 82 | search_xyz: (batch_size, npoint, 3) float32 array, query points 83 | Output: 84 | idx: (batch_size, npoint, nsample) int32 array, indices to input points 85 | pts_cnt: (batch_size, npoint) int32 array, number of unique points in each local region 86 | ''' 87 | with tf.device('/cpu:0'): 88 | shuffled_ids = shuffle_ids(data_xyz) 89 | return tf_op_module.query_ball(data_xyz, search_xyz, shuffled_ids, radius, nsample) 90 | ops.NoGradient('QueryBall') 91 | 92 | def knn_search(data_xyz, search_xyz, k_num): 93 | ''' 94 | Input: 95 | k_num: int32, number of k nearst points 96 | data_xyz: (batch_size, nsearchset, 3) float32 array, search points 97 | search_xyz: (batch_size, ndataset, 3) float32 array, dataset points 98 | Output: 99 | idx: (batch_size, ndataset, k_num) int32 array, indices of knn points 100 | dist: (batch_size, ndataset, k_num) float32 array, distance of knn points (without sort) 101 | ''' 102 | return tf_op_module.knn_search(data_xyz, search_xyz, k_num) 103 | ops.NoGradient('KnnSearch') 104 | 105 | def gather_point(inp,idx): 106 | ''' 107 | input: 108 | batch_size * ndataset * 3 float32 109 | batch_size * npoints int32 110 | returns: 111 | batch_size * npoints * 3 float32 112 | ''' 113 | return tf_op_module.gather_point(inp,idx) 114 | 115 | @tf.RegisterGradient('GatherPoint') 116 | def _gather_point_grad(op,out_g): 117 | inp=op.inputs[0] 118 | idx=op.inputs[1] 119 | return [tf_op_module.gather_point_grad(inp,idx,out_g),None] 120 | 121 | def gather_id(inp,idx): 122 | ''' 123 | input: 124 | batch_size * ndataset * 3 int32 125 | batch_size * npoints int32 126 | returns: 127 | batch_size * npoints * 3 int32 128 | ''' 129 | return tf_op_module.gather_id(inp,idx) 130 | ops.NoGradient('GatherId') 131 | 132 | def group_point(points, idx): 133 | ''' 134 | Input: 135 | points: (batch_size, ndataset, channel) float32 array, points to sample from 136 | idx: (batch_size, npoint, nsample) int32 array, indices to points 137 | Output: 138 | out: (batch_size, npoint, nsample, channel) float32 array, values sampled from points 139 | ''' 140 | return tf_op_module.group_point(points, idx) 141 | @tf.RegisterGradient('GroupPoint') 142 | def _group_point_grad(op, grad_out): 143 | points = op.inputs[0] 144 | idx = op.inputs[1] 145 | return [tf_op_module.group_point_grad(points, idx, grad_out), None] 146 | 147 | def gather_eigvector(eigvs, idx): 148 | ''' 149 | Input: 150 | points: (batch_size, ndataset, 3,3) float32 array, points to sample from 151 | idx: (batch_size, ndataset) int32 array, indices to points 152 | Output: 153 | out: (batch_size, ndataset, 3) float32 array, values sampled from points 154 | ''' 155 | return tf_op_module.gather_eigvector(eigvs, idx) 156 | @tf.RegisterGradient('GatherEigvector') 157 | def _gather_eigvector_grad(op, grad_out): 158 | eigvs = op.inputs[0] 159 | idx = op.inputs[1] 160 | return [tf_op_module.gather_eigvector_grad(eigvs, idx, grad_out), None] 161 | 162 | 163 | def three_interpolate(points, idx, weight): 164 | ''' 165 | Input: 166 | points: (b,m,c) float32 array, known points 167 | idx: (b,n,3) int32 array, indices to known points 168 | weight: (b,n,3) float32 array, weights on known points 169 | Output: 170 | out: (b,n,c) float32 array, interpolated point values 171 | ''' 172 | return tf_op_module.three_interpolate(points, idx, weight) 173 | @tf.RegisterGradient('ThreeInterpolate') 174 | def _three_interpolate_grad(op, grad_out): 175 | points = op.inputs[0] 176 | idx = op.inputs[1] 177 | weight = op.inputs[2] 178 | return [tf_op_module.three_interpolate_grad(points, idx, weight, grad_out), None, None] 179 | -------------------------------------------------------------------------------- /tf_ops/tf_ops.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wleigithub/GACNet/bab6c5b45ac74dc2256bba4e5101524ca8086179/tf_ops/tf_ops.pyc -------------------------------------------------------------------------------- /tf_ops/tf_ops_.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wleigithub/GACNet/bab6c5b45ac74dc2256bba4e5101524ca8086179/tf_ops/tf_ops_.pyc -------------------------------------------------------------------------------- /tf_ops/tf_sampling.cpp: -------------------------------------------------------------------------------- 1 | #include "tensorflow/core/framework/op.h" 2 | #include "tensorflow/core/framework/op_kernel.h" 3 | #include "tensorflow/core/framework/shape_inference.h" 4 | #include "tensorflow/core/framework/common_shape_fns.h" 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | using namespace tensorflow; 12 | 13 | 14 | 15 | REGISTER_OP("RandSeeds") 16 | .Input("inp: float32") 17 | .Output("out: int32") 18 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 19 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * npoint * 3 20 | c->WithRank(c->input(0), 3, &dims1); 21 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0)}); 22 | c->set_output(0, output); 23 | return Status::OK(); 24 | }); 25 | 26 | 27 | REGISTER_OP("ShuffleIds") 28 | .Input("inp: float32") 29 | .Output("out: int32") 30 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 31 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * npoint * 3 32 | c->WithRank(c->input(0), 3, &dims1); 33 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), c->Dim(dims1, 1)}); 34 | c->set_output(0, output); 35 | return Status::OK(); 36 | }); 37 | 38 | 39 | Status GetFPSNum( 40 | shape_inference::InferenceContext* c, 41 | shape_inference::DimensionHandle input_size, 42 | int64 stride, 43 | shape_inference::DimensionHandle* output_size) { 44 | if (stride <= 0) { 45 | return errors::InvalidArgument("Stride must be > 0, but got ", stride); 46 | } 47 | TF_RETURN_IF_ERROR(c->Add(input_size, stride - 1, output_size)); 48 | TF_RETURN_IF_ERROR(c->Divide(*output_size, stride, false /* evenly_divisible */, output_size)); 49 | return Status::OK(); 50 | } 51 | 52 | REGISTER_OP("FarthestPointSample") 53 | .Attr("stride: int") 54 | .Input("inp: float32") 55 | .Input("initid: int32") 56 | .Output("out: int32") 57 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 58 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * npoint * 3 59 | c->WithRank(c->input(0), 3, &dims1); 60 | 61 | ::tensorflow::shape_inference::DimensionHandle org_pcnum = c->Dim(dims1, 1); 62 | int stride; 63 | TF_RETURN_IF_ERROR(c->GetAttr("stride", &stride)); 64 | ::tensorflow::shape_inference::DimensionHandle npoint; 65 | GetFPSNum(c, org_pcnum, stride, &npoint); 66 | 67 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), npoint}); 68 | c->set_output(0, output); 69 | return Status::OK(); 70 | }); 71 | 72 | 73 | REGISTER_OP("QueryBallKnn") 74 | .Attr("radius: float") 75 | .Attr("k_num: int") 76 | .Input("data_xyz: float32") 77 | .Input("search_xyz: float32") 78 | .Input("shuffled_ids: int32") 79 | .Output("idx: int32") 80 | .Output("dist: float32") 81 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 82 | ::tensorflow::shape_inference::ShapeHandle dims2; // batch_size * npoint * 3 83 | c->WithRank(c->input(1), 3, &dims2); 84 | int k_num; 85 | TF_RETURN_IF_ERROR(c->GetAttr("k_num", &k_num)); 86 | ::tensorflow::shape_inference::ShapeHandle output1 = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1), k_num}); 87 | c->set_output(0, output1); 88 | ::tensorflow::shape_inference::ShapeHandle output2 = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1), k_num}); 89 | c->set_output(1, output2); 90 | return Status::OK(); 91 | }); 92 | 93 | 94 | REGISTER_OP("QueryBall") 95 | .Attr("radius: float") 96 | .Attr("k_num: int") 97 | .Input("data_xyz: float32") 98 | .Input("search_xyz: float32") 99 | .Input("shuffled_ids: int32") 100 | .Output("idx: int32") 101 | .Output("pts_cnt: int32") 102 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 103 | ::tensorflow::shape_inference::ShapeHandle dims2; // batch_size * npoint * 3 104 | c->WithRank(c->input(1), 3, &dims2); 105 | int k_num; 106 | TF_RETURN_IF_ERROR(c->GetAttr("k_num", &k_num)); 107 | ::tensorflow::shape_inference::ShapeHandle output1 = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1), k_num}); 108 | c->set_output(0, output1); 109 | ::tensorflow::shape_inference::ShapeHandle output2 = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1)}); 110 | c->set_output(1, output2); 111 | return Status::OK(); 112 | }); 113 | 114 | 115 | REGISTER_OP("KnnSearch") 116 | .Attr("k_num: int") 117 | .Input("data_xyz: float32") 118 | .Input("search_xyz: float32") 119 | .Output("idx: int32") 120 | .Output("dist: float32") 121 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 122 | ::tensorflow::shape_inference::ShapeHandle dims2; // batch_size * npoint * 3 123 | c->WithRank(c->input(1), 3, &dims2); 124 | int k_num; 125 | TF_RETURN_IF_ERROR(c->GetAttr("k_num", &k_num)); 126 | ::tensorflow::shape_inference::ShapeHandle output1 = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1), k_num}); 127 | c->set_output(0, output1); 128 | ::tensorflow::shape_inference::ShapeHandle output2 = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1), k_num}); 129 | c->set_output(1, output2); 130 | return Status::OK(); 131 | }); 132 | 133 | void randomLauncher(int b, int *gpu_nums, int upbound); 134 | class RandSeedGpuOp: public OpKernel{ 135 | public: 136 | explicit RandSeedGpuOp(OpKernelConstruction* context) : OpKernel(context) {} 137 | 138 | void Compute(OpKernelContext * context)override{ 139 | 140 | const Tensor& inp_tensor=context->input(0); 141 | OP_REQUIRES(context,inp_tensor.dims()==3 && inp_tensor.shape().dim_size(2)==3,errors::InvalidArgument("RandSeedGpuOp expects (batch_size,num_points,3) inp shape")); 142 | int b=inp_tensor.shape().dim_size(0); 143 | int n=inp_tensor.shape().dim_size(1); 144 | auto inp_flat=inp_tensor.flat(); 145 | const float * inp=&(inp_flat(0)); 146 | Tensor * out_tensor; 147 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b},&out_tensor)); 148 | auto out_flat=out_tensor->flat(); 149 | int * out=&(out_flat(0)); 150 | randomLauncher(b, out, n); 151 | } 152 | }; 153 | REGISTER_KERNEL_BUILDER(Name("RandSeeds").Device(DEVICE_GPU),RandSeedGpuOp); 154 | 155 | 156 | void shuffleidcpu(int b, int n, int *shuffled_ids){ 157 | //initialization 158 | for(int i=0;iinput(0); 171 | OP_REQUIRES(context,inp_tensor.dims()==3 && inp_tensor.shape().dim_size(2)==3,errors::InvalidArgument("ShuffleIdOp expects (batch_size,num_points,3) inp shape")); 172 | int b=inp_tensor.shape().dim_size(0); 173 | int n=inp_tensor.shape().dim_size(1); 174 | auto inp_flat=inp_tensor.flat(); 175 | const float * inp=&(inp_flat(0)); 176 | Tensor * out_tensor; 177 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b, n},&out_tensor)); 178 | auto out_flat=out_tensor->flat(); 179 | int * out=&(out_flat(0)); 180 | shuffleidcpu(b, n, out); 181 | } 182 | }; 183 | REGISTER_KERNEL_BUILDER(Name("ShuffleIds").Device(DEVICE_CPU),ShuffleIdOp); 184 | 185 | void farthestpointsamplingLauncher(int b,int n,int m,const int *init, const float * inp,float * temp,int * out); 186 | class FarthestPointSampleGpuOp: public OpKernel{ 187 | public: 188 | explicit FarthestPointSampleGpuOp(OpKernelConstruction* context):OpKernel(context) { 189 | OP_REQUIRES_OK(context, context->GetAttr("stride", &stride_)); 190 | OP_REQUIRES(context, stride_ > 0, errors::InvalidArgument("FarthestPointSample expects positive stride")); 191 | } 192 | void Compute(OpKernelContext * context)override{ 193 | 194 | const Tensor& inp_tensor=context->input(0); 195 | OP_REQUIRES(context,inp_tensor.dims()==3 && inp_tensor.shape().dim_size(2)==3,errors::InvalidArgument("FarthestPointSample expects (batch_size,num_points,3) inp shape")); 196 | int b=inp_tensor.shape().dim_size(0); 197 | int n=inp_tensor.shape().dim_size(1); 198 | auto inp_flat=inp_tensor.flat(); 199 | const float * inp=&(inp_flat(0)); 200 | int m = (n+stride_-1)/stride_; //point number to sample 201 | Tensor * out_tensor; 202 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m},&out_tensor)); 203 | auto out_flat=out_tensor->flat(); 204 | int * out=&(out_flat(0)); 205 | Tensor temp_tensor; 206 | OP_REQUIRES_OK(context,context->allocate_temp(DataTypeToEnum::value,TensorShape{32,n},&temp_tensor)); 207 | auto temp_flat=temp_tensor.flat(); 208 | float * temp=&(temp_flat(0)); 209 | 210 | const Tensor& initid_tensor=context->input(1); 211 | //Tensor& initid_tensor=context->input(1); 212 | //OP_REQUIRES(context,inp_tensor.dims()==1 ,errors::InvalidArgument("FarthestPointSample expects (batch_size,num_points,3) inp shape")); 213 | auto initid_flat=initid_tensor.flat(); 214 | const int * init=&(initid_flat(0)); 215 | farthestpointsamplingLauncher(b,n,m,init, inp,temp,out); 216 | 217 | } 218 | private: 219 | int stride_; 220 | }; 221 | REGISTER_KERNEL_BUILDER(Name("FarthestPointSample").Device(DEVICE_GPU),FarthestPointSampleGpuOp); 222 | 223 | 224 | void queryBallKnnLauncher(int b, int n, int m, float radius, int k_num, const int *shuffled_ids, const float *data_xyz, const float *search_xyz, int *idx, float *dist); 225 | class QueryBallKnnGpuOp : public OpKernel { 226 | public: 227 | explicit QueryBallKnnGpuOp(OpKernelConstruction* context) : OpKernel(context) { 228 | OP_REQUIRES_OK(context, context->GetAttr("radius", &radius_)); 229 | OP_REQUIRES(context, radius_ > 0, errors::InvalidArgument("QueryBallKnn expects positive radius")); 230 | 231 | OP_REQUIRES_OK(context, context->GetAttr("k_num", &k_num_)); 232 | OP_REQUIRES(context, k_num_ > 0, errors::InvalidArgument("QueryBallKnn expects positive k_num")); 233 | } 234 | 235 | void Compute(OpKernelContext* context) override { 236 | const Tensor& data_xyz_tensor = context->input(0); 237 | OP_REQUIRES(context, data_xyz_tensor.dims()==3 && data_xyz_tensor.shape().dim_size(2)==3, errors::InvalidArgument("QueryBallKnn expects (batch_size, ndataset, 3) data_xyz shape.")); 238 | int b = data_xyz_tensor.shape().dim_size(0); 239 | int n = data_xyz_tensor.shape().dim_size(1); 240 | 241 | const Tensor& search_xyz_tensor = context->input(1); 242 | OP_REQUIRES(context, search_xyz_tensor.dims()==3 && search_xyz_tensor.shape().dim_size(2)==3, errors::InvalidArgument("QueryBallKnn expects (batch_size, npoint, 3) search_xyz shape.")); 243 | int m = search_xyz_tensor.shape().dim_size(1); 244 | 245 | const Tensor& shuffled_ids_tensor = context->input(2); 246 | OP_REQUIRES(context, shuffled_ids_tensor.dims()==2, errors::InvalidArgument("QueryBallKnn expects (batch_size, ndataset) shuffled_ids shape.")); 247 | 248 | Tensor *idx_tensor = nullptr; 249 | OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape{b,m,k_num_}, &idx_tensor)); 250 | Tensor *dist_tensor = nullptr; 251 | OP_REQUIRES_OK(context, context->allocate_output(1, TensorShape{b,m,k_num_}, &dist_tensor)); 252 | 253 | auto shuffled_ids_flat = shuffled_ids_tensor.flat(); 254 | const int *shuffled_ids = &(shuffled_ids_flat(0)); 255 | auto data_xyz_flat = data_xyz_tensor.flat(); 256 | const float *data_xyz = &(data_xyz_flat(0)); 257 | auto search_xyz_flat = search_xyz_tensor.flat(); 258 | const float *search_xyz = &(search_xyz_flat(0)); 259 | auto idx_flat = idx_tensor->flat(); 260 | int *idx = &(idx_flat(0)); 261 | auto dist_flat = dist_tensor->flat(); 262 | float *dist = &(dist_flat(0)); 263 | 264 | queryBallKnnLauncher(b,n,m,radius_,k_num_,shuffled_ids,data_xyz,search_xyz,idx,dist); 265 | } 266 | private: 267 | float radius_; 268 | int k_num_; 269 | }; 270 | REGISTER_KERNEL_BUILDER(Name("QueryBallKnn").Device(DEVICE_GPU), QueryBallKnnGpuOp); 271 | 272 | 273 | void queryBallLauncher(int b, int n, int m, float radius, int k_num, const int *shuffled_ids, const float *data_xyz, const float *search_xyz, int *idx, int *pts_cnt); 274 | class QueryBallGpuOp : public OpKernel { 275 | public: 276 | explicit QueryBallGpuOp(OpKernelConstruction* context) : OpKernel(context) { 277 | OP_REQUIRES_OK(context, context->GetAttr("radius", &radius_)); 278 | OP_REQUIRES(context, radius_ > 0, errors::InvalidArgument("QueryBall expects positive radius")); 279 | 280 | OP_REQUIRES_OK(context, context->GetAttr("k_num", &k_num_)); 281 | OP_REQUIRES(context, k_num_ > 0, errors::InvalidArgument("QueryBall expects positive k_num")); 282 | } 283 | 284 | void Compute(OpKernelContext* context) override { 285 | const Tensor& data_xyz_tensor = context->input(0); 286 | OP_REQUIRES(context, data_xyz_tensor.dims()==3 && data_xyz_tensor.shape().dim_size(2)==3, errors::InvalidArgument("QueryBall expects (batch_size, ndataset, 3) data_xyz shape.")); 287 | int b = data_xyz_tensor.shape().dim_size(0); 288 | int n = data_xyz_tensor.shape().dim_size(1); 289 | 290 | const Tensor& search_xyz_tensor = context->input(1); 291 | OP_REQUIRES(context, search_xyz_tensor.dims()==3 && search_xyz_tensor.shape().dim_size(2)==3, errors::InvalidArgument("QueryBall expects (batch_size, npoint, 3) search_xyz shape.")); 292 | int m = search_xyz_tensor.shape().dim_size(1); 293 | 294 | const Tensor& shuffled_ids_tensor = context->input(2); 295 | OP_REQUIRES(context, shuffled_ids_tensor.dims()==2, errors::InvalidArgument("QueryBall expects (batch_size, ndataset) shuffled_ids shape.")); 296 | 297 | Tensor *idx_tensor = nullptr; 298 | OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape{b,m,k_num_}, &idx_tensor)); 299 | Tensor *pts_cnt_tensor = nullptr; 300 | OP_REQUIRES_OK(context, context->allocate_output(1, TensorShape{b,m}, &pts_cnt_tensor)); 301 | 302 | auto shuffled_ids_flat = shuffled_ids_tensor.flat(); 303 | const int *shuffled_ids = &(shuffled_ids_flat(0)); 304 | auto data_xyz_flat = data_xyz_tensor.flat(); 305 | const float *data_xyz = &(data_xyz_flat(0)); 306 | auto search_xyz_flat = search_xyz_tensor.flat(); 307 | const float *search_xyz = &(search_xyz_flat(0)); 308 | auto idx_flat = idx_tensor->flat(); 309 | int *idx = &(idx_flat(0)); 310 | auto pts_cnt_flat = pts_cnt_tensor->flat(); 311 | int *pts_cnt = &(pts_cnt_flat(0)); 312 | 313 | queryBallLauncher(b,n,m,radius_,k_num_,shuffled_ids,data_xyz,search_xyz,idx,pts_cnt); 314 | } 315 | private: 316 | float radius_; 317 | int k_num_; 318 | }; 319 | REGISTER_KERNEL_BUILDER(Name("QueryBall").Device(DEVICE_GPU), QueryBallGpuOp); 320 | 321 | void knnLauncher(int b, int n, int m, int k_num, const float *data_xyz, const float *search_xyz, int *idx, float *dist); 322 | class KnnSearchGpuOp : public OpKernel { 323 | public: 324 | explicit KnnSearchGpuOp(OpKernelConstruction* context) : OpKernel(context) { 325 | OP_REQUIRES_OK(context, context->GetAttr("k_num", &k_num_)); 326 | OP_REQUIRES(context, k_num_ > 0, errors::InvalidArgument("KnnSearch expects positive k_num")); 327 | } 328 | 329 | void Compute(OpKernelContext* context) override { 330 | const Tensor& data_xyz_tensor = context->input(0); 331 | OP_REQUIRES(context, data_xyz_tensor.dims()==3 && data_xyz_tensor.shape().dim_size(2)==3, errors::InvalidArgument("KnnSearch expects (batch_size, ndataset, 3) xyz1 shape.")); 332 | int b = data_xyz_tensor.shape().dim_size(0); 333 | int n = data_xyz_tensor.shape().dim_size(1); 334 | 335 | const Tensor& search_xyz_tensor = context->input(1); 336 | OP_REQUIRES(context, search_xyz_tensor.dims()==3 && search_xyz_tensor.shape().dim_size(2)==3, errors::InvalidArgument("KnnSearch expects (batch_size, npoint, 3) xyz2 shape.")); 337 | int m = search_xyz_tensor.shape().dim_size(1); 338 | 339 | Tensor *idx_tensor = nullptr; 340 | OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape{b,m,k_num_}, &idx_tensor)); 341 | Tensor *dist_tensor = nullptr; 342 | OP_REQUIRES_OK(context, context->allocate_output(1, TensorShape{b,m,k_num_}, &dist_tensor)); 343 | 344 | auto data_xyz_flat = data_xyz_tensor.flat(); 345 | const float *data_xyz = &(data_xyz_flat(0)); 346 | auto search_xyz_flat = search_xyz_tensor.flat(); 347 | const float *search_xyz = &(search_xyz_flat(0)); 348 | auto idx_flat = idx_tensor->flat(); 349 | int *idx = &(idx_flat(0)); 350 | auto dist_flat = dist_tensor->flat(); 351 | float *dist = &(dist_flat(0)); 352 | 353 | knnLauncher(b, n, m, k_num_, data_xyz, search_xyz, idx, dist); 354 | } 355 | 356 | private: 357 | int k_num_; 358 | }; 359 | REGISTER_KERNEL_BUILDER(Name("KnnSearch").Device(DEVICE_GPU), KnnSearchGpuOp); 360 | 361 | -------------------------------------------------------------------------------- /tf_ops/tf_test.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import random 4 | from tf_ops import knn_search, three_interpolate, group_point, gather_point, FPS_downsample, query_ball_knn, shuffle_ids,query_ball 5 | import pdb 6 | 7 | 8 | if __name__=='__main__': 9 | #arr_xyz1 = np.random.random((1,128,3)).astype('float32') 10 | arr_xyz2 = np.random.random((1,68,3)).astype('float32') 11 | arr_xyz1 = np.loadtxt('/media/wl/data/CVPR/pointnet-master/data/modelnet40_ply_hdf5_2048/test/1.txt') 12 | arr_xyz1 = np.reshape(arr_xyz1[:, 0:3], [1,-1, 3]) 13 | #pdb.set_trace() 14 | xyz1 = tf.constant(arr_xyz1, dtype=tf.float32) 15 | xyz2 = tf.constant(arr_xyz1[:,0:500,:], dtype=tf.float32) 16 | #idxs, dist = knn_search(xyz2, xyz1, 3) 17 | #idxs, dist = query_ball_knn(xyz1, xyz1, 0.2, 15) 18 | #idxs, dist = query_ball(xyz2, xyz1, 0.1, 21) 19 | idxs = FPS_downsample(xyz1, 3) 20 | 21 | #idxs = shuffle_ids(xyz1) 22 | #pts = gather_point(xyz1, idxs) 23 | #with tf.device('/cpu:0'): 24 | # shuffled_ids = shuffle_ids(xyz1) 25 | with tf.Session() as sess: 26 | idx_out = sess.run(idxs) 27 | idx= tf.constant(idx_out) 28 | pdb.set_trace() 29 | -------------------------------------------------------------------------------- /utils/model.py: -------------------------------------------------------------------------------- 1 | 2 | import os, pdb 3 | import sys 4 | BASE_DIR = os.path.dirname(__file__) 5 | sys.path.append(BASE_DIR) 6 | sys.path.append(os.path.join(BASE_DIR, '../utils')) 7 | sys.path.append(os.path.join(BASE_DIR, '../tf_ops')) 8 | import tensorflow as tf 9 | import numpy as np 10 | import tf_util 11 | from pcnet_util import build_graph, graph_coarse, graph_attention_layer, graph_pooling_layer, point_upsample_layer, crf_layer, graph_attention_layer_for_featurerefine 12 | from tf_ops import knn_search 13 | 14 | 15 | def placeholder_inputs(feature_channel, with_spw=False): 16 | pointclouds_pl = tf.placeholder(tf.float32, shape=(None, None, feature_channel)) 17 | labels_pl = tf.placeholder(tf.int32, shape=(None, None)) 18 | is_training_pl = tf.placeholder(tf.bool, shape=()) 19 | if with_spw: 20 | spws_pl = tf.placeholder(tf.float32, shape=(None, None)) 21 | return pointclouds_pl, labels_pl, spws_pl, is_training_pl 22 | else: 23 | return pointclouds_pl, labels_pl, is_training_pl 24 | 25 | 26 | def build_graph_pyramid(xyz, graph_inf): 27 | """ Builds a pyramid of graphs and pooling operations corresponding to progressively coarsened point cloud. 28 | Inputs: 29 | xyz: (batchsize, num_point, nfeature) 30 | graph_inf: parameters for graph building (see run.py) 31 | Outputs: 32 | graph_prd: graph pyramid contains the vertices and their edges at each layer 33 | coarse_map: record the corresponding relation between two close graph layers (for graph coarseing/pooling) 34 | """ 35 | stride_list, radius_list, maxsample_list = graph_inf['stride_list'], graph_inf['radius_list'], graph_inf['maxsample_list'] 36 | 37 | graph_prd = [] 38 | graph = {} #save subsampled points and their neighbor indexes at each level 39 | coarse_map = [] # save index map from previous level to next level 40 | 41 | ids = build_graph(xyz, radius_list[0], maxsample_list[0]) # (batchsize, num_point, maxsample_list[0]) neighbor indexes at current level 42 | graph['vertex'], graph['adjids'] = xyz, ids 43 | graph_prd.append(graph.copy()) 44 | 45 | for stride, radius, maxsample in zip(stride_list, radius_list[1:], maxsample_list[1:]): 46 | xyz, coarse_map_ids = graph_coarse(xyz, ids, stride, radius, maxsample) 47 | coarse_map.append(coarse_map_ids) 48 | ids = build_graph(xyz, radius, maxsample) 49 | graph['vertex'], graph['adjids'] = xyz, ids 50 | graph_prd.append(graph.copy()) 51 | 52 | return graph_prd, coarse_map 53 | 54 | 55 | def build_network(graph_prd, coarse_map, net_inf, features, num_class, is_training, bn_decay=None): 56 | """build GACNet on the graph pyramid. 57 | Inputs: 58 | graph_prd: graph pyramid contains the vertices and their edges at each layer 59 | coarse_map: record the corresponding relation between two close graph layers (for graph coarseing/pooling) 60 | net_inf: parameters for GACNet (see run.py for details) 61 | features: input signal (x,y,z,r,g,b,...) 62 | num_class: number of class to be classified 63 | is_training: training or not 64 | bn_decay: use weight decay 65 | Outputs: 66 | features: learned feature for each point 67 | """ 68 | forward_parm, upsample_parm, fullconect_parm = net_inf['forward_parm'], net_inf['upsample_parm'], net_inf['fullconect_parm'] 69 | 70 | inif = features[...,0:6] # (x,y,z,r,g,b) 71 | features = features[...,2:] #(z, r, g, b, and (initial geofeatures if possible)) 72 | #features = tf.concat( [features[...,2:3],features[...,6:]], axis=-1) #remove the rgb 73 | 74 | #forward layers 75 | feature_prd = [] 76 | for i in range(len(coarse_map)): 77 | features = graph_attention_layer(graph_prd[i], features, forward_parm[i][0], forward_parm[i][1], is_training, bn_decay, scope='attention_%d'%(i), bn=True) 78 | feature_prd.append(features) 79 | features = graph_pooling_layer(features, coarse_map[i], scope='graph_pooling_%d'%(i)) 80 | 81 | features = graph_attention_layer(graph_prd[-1], features, forward_parm[-1][0], forward_parm[-1][1], is_training, bn_decay, scope='attention_%d'%(len(coarse_map)), bn=True) 82 | 83 | #upsample/interpolation layers 84 | for i in range(len(coarse_map)): 85 | j = len(coarse_map) -i-1 86 | features = point_upsample_layer(graph_prd[j]['vertex'], graph_prd[j+1]['vertex'], feature_prd[j], features, 87 | upsampling=True, mlp=upsample_parm[j], is_training=is_training, bn_decay=bn_decay, scope='up%d'%(j)) 88 | 89 | #fully connected layer 90 | features = tf_util.conv1d(features, fullconect_parm, 1, bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay) 91 | features = tf_util.dropout(features, keep_prob=0.5, is_training=is_training, scope='dp1') 92 | features = tf_util.conv1d(features, num_class, 1, activation_fn=None, is_training=is_training, scope='fc2') 93 | 94 | #features = crf_layer(graph_prd[0]['vertex'], inif[...,3:6], features, graph_prd[0]['adjids'], is_training, bn_decay, scope='crflayer', iter_num=5) 95 | features = graph_attention_layer_for_featurerefine(inif, features, graph_prd[0]['adjids'], is_training, bn_decay, scope='refine') # in this layer, the ouput is not through actication function 96 | 97 | return features 98 | 99 | 100 | def get_loss(pred, label, spw=None): 101 | """ 102 | pred: BxNxC, 103 | label: BxN, 104 | smpw: BxN 105 | """ 106 | if spw is not None: 107 | classify_loss = tf.losses.sparse_softmax_cross_entropy(labels=label, logits=pred, weights=spw) 108 | else: 109 | classify_loss = tf.losses.sparse_softmax_cross_entropy(labels=label, logits=pred) 110 | tf.summary.scalar('classify loss', classify_loss) 111 | return classify_loss 112 | 113 | 114 | # test ------------------------------------------------------------------------------------------------ 115 | if __name__ == '__main__': 116 | import pdb 117 | num_point = 4096 118 | data = np.random.rand(10, num_point, 6) 119 | #xyz = tf.constant(data, dtype=np.float32) 120 | xyz = tf.placeholder(tf.float32, shape=(None,None,6)) 121 | num_point = tf.placeholder(tf.int32,shape=()) 122 | 123 | graph_inf = {'stride_list': [4, 4, 4, 2], 124 | 'radius_list': [0.1, 0.2, 0.4, 0.8, 1.6], 125 | 'maxsample_list': [12, 21, 21, 21, 12] 126 | } 127 | 128 | forward_parm = [ 129 | [ [32,32,64], [64] ], 130 | [ [64,64,128], [128] ], 131 | [ [128,128,256], [256] ], 132 | [ [256,256,512], [512] ], 133 | [ [256,256], [256] ] 134 | ] 135 | upsample_parm = [ 136 | [128, 128], 137 | [128, 128], 138 | [256, 256], 139 | [256, 256] 140 | ] 141 | fullconect_parm = 128 142 | 143 | net_inf = {'forward_parm': forward_parm, 144 | 'upsample_parm': upsample_parm, 145 | 'fullconect_parm': fullconect_parm 146 | } 147 | 148 | #with tf.Graph().as_default(): 149 | graph_prd, coarse_map = build_graph_pyramid(xyz, graph_inf) 150 | logits = build_network(graph_prd, coarse_map, net_inf, xyz, 10, tf.constant(True)) 151 | 152 | #pdb.set_trace() -------------------------------------------------------------------------------- /utils/pcnet_util.py: -------------------------------------------------------------------------------- 1 | 2 | import tensorflow as tf 3 | import numpy as np 4 | import os 5 | import sys 6 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 7 | ROOT_DIR = os.path.dirname(BASE_DIR) 8 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 9 | sys.path.append(os.path.join(ROOT_DIR, 'tf_ops')) 10 | 11 | from tf_ops import FPS_downsample, query_ball_knn, knn_search, gather_point, gather_id, group_point, three_interpolate, query_ball, gather_eigvector 12 | 13 | import tf_util 14 | import pdb 15 | 16 | #parameters for graph attention convolution (GAC) 17 | #which combines both the position and feature difference to assign weights to different neighbors 18 | gac_par = [ 19 | [32, 16], #MLP for xyz 20 | [16, 16], #MLP for feature 21 | [64] #hidden node of MLP for mergering 22 | ] 23 | 24 | ############################################################################################################################## 25 | #------------------------------------------graph building modules------------------------------------------------------------ 26 | ############################################################################################################################## 27 | def build_graph(xyz, radius, maxsample): 28 | """ Converts point cloud to graph. 29 | Inputs: 30 | xyz: (batchsize, npoint, nfeature) 31 | radius: radius to search neighbors 32 | maxsample: max number to sample in their neighbors 33 | Outputs: 34 | ids: neighbors' indices 35 | """ 36 | ids, cnts = query_ball(xyz, xyz, radius, maxsample) #(batchsize, npoint, maxsample) (batchsize, npoint, maxsample) 37 | return ids 38 | 39 | 40 | def graph_coarse(xyz_org, ids_full, stride, radius, maxsample): 41 | """ Coarse graph with down sampling, and find their corresponding vertexes at previous (or father) level. """ 42 | if stride>1: 43 | sub_pts_ids = FPS_downsample(xyz_org, stride) #(batch_size, num_point) 44 | sub_xyz = gather_point(xyz_org, sub_pts_ids) #(batch_size, num_point, 3) 45 | ids = gather_id(ids_full, sub_pts_ids)#(batchsize, num_point, maxsample) 46 | return sub_xyz, ids 47 | else: 48 | return xyz_org, ids_full 49 | 50 | 51 | ############################################################################################################################## 52 | #------------------------------------------graph network modules------------------------------------------------------------ 53 | ############################################################################################################################## 54 | 55 | def covInf(grouped_xyz): 56 | #calculate covMat 57 | pts_mean = grouped_xyz - tf.reduce_mean(grouped_xyz, axis=2, keep_dims=True) # (b, n, K, 3) 58 | pts_BNK31 = tf.expand_dims(pts_mean, axis=-1) 59 | covMat = tf.matmul(pts_BNK31, pts_BNK31, transpose_b=True) # (b, n, K, 3, 3) 60 | covMat = tf.reduce_mean(covMat, axis=2, keep_dims=False) # (b, n, 3, 3) 61 | 62 | covMat_flat = tf.concat([covMat[...,0],covMat[...,1],covMat[...,2]], axis=-1) 63 | ''' 64 | #calculate covInf 65 | eigs, eigvs = tf.self_adjoint_eig(covMat) 66 | eigs = tf.nn.l2_normalize(tf.abs(eigs)) 67 | eigs, idx = tf.nn.top_k(eigs, k=3) 68 | # 69 | idx = idx[...,2] 70 | #pdb.set_trace() 71 | eigv = gather_eigvector(eigvs, idx) 72 | eigv = tf.nn.l2_normalize(eigv) 73 | 74 | return tf.concat([eigs, eigv], axis=-1) 75 | ''' 76 | return covMat_flat 77 | 78 | 79 | def coeff_generation(grouped_features, features, grouped_xyz, is_training, bn_decay, scope, bn=True, mode='with_feature'): 80 | with tf.variable_scope(scope) as sc: 81 | if mode == 'with_feature': 82 | coeff = grouped_features - tf.expand_dims(features, axis=2) 83 | coeff = tf_util.MLP_2d(coeff, gac_par[1], is_training, bn_decay, bn, scope='conv_with_feature')## compress to a hiden feature space 84 | coeff = tf.concat([grouped_xyz, coeff], -1) 85 | if mode == 'edge_only': 86 | coeff = grouped_xyz 87 | if mode == 'feature_only': 88 | coeff = grouped_features - tf.expand_dims(features, axis=2) 89 | coeff = tf_util.MLP_2d(coeff, gac_par[1], is_training, bn_decay, bn, scope='conv_feature')## compress to a hiden feature space 90 | 91 | grouped_features = tf.concat([grouped_xyz, grouped_features], axis=-1) #updata feature 92 | out_chal = grouped_features.get_shape()[-1].value 93 | coeff = tf_util.MLP_2d(coeff, gac_par[2], is_training, bn_decay, bn, scope='conv')## map to a hiden feature space 94 | coeff = tf_util.conv2d(coeff, out_chal, [1,1], scope='conv2d', is_training=is_training, bn_decay=bn_decay, activation_fn=None) #output coefficent 95 | coeff = tf.nn.softmax(coeff, axis=2) #coefffient normalization 96 | 97 | grouped_features = tf.multiply(coeff, grouped_features) 98 | grouped_features = tf.reduce_sum(grouped_features, axis=[2], keep_dims=False) 99 | return grouped_features 100 | 101 | 102 | def graph_attention_layer(graph, features, mlp, mlp2, is_training, bn_decay, scope, bn=True): 103 | with tf.variable_scope(scope) as sc: 104 | xyz, ids = graph['vertex'], graph['adjids'] 105 | grouped_xyz = group_point(xyz, ids) # (batch_size, ndataset, maxsample, 3) 106 | ''' 107 | #it's slow to calculate the inital features during training, so we compute it prior 108 | #calculate inital features 109 | if int(scope.split('_')[1])==0: 110 | cov_inf = covInf(grouped_xyz) 111 | features = tf.concat([features, cov_inf], axis=-1) 112 | ''' 113 | grouped_xyz -= tf.expand_dims(xyz, 2) # translation normalization 114 | grouped_xyz = edge_mapping(grouped_xyz, is_training, bn_decay, scope='edge') # map local postion to a feature space with MLP 115 | features = tf_util.MLP_1d(features, mlp, is_training, bn_decay, bn, scope='feature') #feature transform 116 | grouped_features = group_point(features, ids) # (batch_size, ndataset, maxsample, channel) 117 | 118 | new_features = coeff_generation(grouped_features, features,grouped_xyz, is_training, bn_decay, scope='coeff_gen') 119 | #new_features = tf.reduce_max(tf.concat([grouped_xyz, grouped_features], axis=-1), axis=2, keep_dims=False) #for comparison of GAC and MAX 120 | #new_features = tf.reduce_mean(tf.concat([grouped_xyz, grouped_features], axis=-1), axis=2, keep_dims=False) #for comparison of GAC and MEAN 121 | 122 | if mlp2 is not None and features is not None: 123 | new_features = tf.concat([features, new_features], axis=-1) # (batch_size, ndataset, maxsample, 3+channel) #merge the feature with itself 124 | new_features = tf_util.MLP_1d(new_features, mlp2, is_training, bn_decay, bn) 125 | return new_features 126 | 127 | 128 | def graph_attention_layer_for_featurerefine(initf, features, ids, is_training, bn_decay, scope, bn=True): 129 | """ 130 | Graph attention convolution for post-processing (for comparion convenience with CRF) 131 | """ 132 | with tf.variable_scope(scope) as sc: 133 | out_chal = features.get_shape()[-1].value 134 | 135 | grouped_initf = group_point(initf, ids) # (batch_size, ndataset, maxsample, 3) 136 | grouped_initf -= tf.expand_dims(initf, 2) # translation normalization 137 | grouped_initf = edge_mapping(grouped_initf, is_training, bn_decay, scope='edge') #map local postion to a feature space with MLP 138 | grouped_features = group_point(features, ids) # (batch_size, ndataset, maxsample, channel) 139 | 140 | new_features = coeff_generation(grouped_features, features, grouped_initf, is_training, bn_decay, scope='coeff_gen') 141 | #new_features = tf.reduce_max(tf.concat([grouped_initf, grouped_features], axis=-1), axis=2, keep_dims=False) 142 | 143 | new_features = tf.concat([features, new_features], axis=-1) # (batch_size, ndataset, maxsample, 3+channel) #merge the feature with itself 144 | new_features = tf_util.conv1d(new_features, out_chal, 1, activation_fn=None, is_training=is_training, scope=scope) 145 | return new_features 146 | 147 | 148 | def edge_mapping(grouped_xyz, is_training, bn_decay, scope, bn=True): 149 | """ mapping edges (deta.x) into feature space {deep sets} 150 | input: grouped_xyz (b,n,k,3) 151 | output: (b,n,k,output_channel) 152 | """ 153 | with tf.variable_scope(scope) as sc: 154 | #MLP for edge feature transform 155 | grouped_xyz = tf_util.MLP_2d(grouped_xyz, gac_par[0], is_training, bn_decay, bn, scope='pn_conv') 156 | return grouped_xyz 157 | 158 | 159 | ############################################################################################################################## 160 | 161 | def graph_pooling_layer(features, coarse_map, scope, pooling='max'): 162 | ''' 163 | Input: 164 | xyz: (batch_size, ndataset, 3) TF tensor 165 | features: (batch_size, ndataset, nchannel) TF tensor 166 | ball_ids: (batch_size, ndataset, maxsample) TF tensor 167 | num_subsampling: int32 or None, None means no subsampling 168 | mlp: list of int32 -- output size for MLP on each point 169 | Return: 170 | new_xyz: (batch_size, num_subsampling or ndataset, 3) TF tensor 171 | new_features: (batch_size, num_subsampling or ndataset, mlp[-1]) TF tensor 172 | ''' 173 | with tf.variable_scope(scope) as sc: 174 | grouped_features = group_point(features, coarse_map) # (batch_size, ndataset or num_subsampling, nsample, channel) 175 | 176 | if pooling=='max': 177 | new_features = tf.reduce_max(grouped_features, axis=[2]) # (batch_size, ndataset or num_subsampling, channel) 178 | 179 | return new_features 180 | 181 | 182 | def point_upsample_layer(xyz1, xyz2, features1, features2, upsampling, mlp, is_training, bn_decay, scope, bn=True): 183 | ''' 184 | Input: 185 | xyz1: (batch_size, ndataset1, 3) TF tensor 186 | xyz2: (batch_size, ndataset2, 3) TF tensor, sparser than xyz1 187 | features1: (batch_size, ndataset1, nchannel1) TF tensor 188 | features2: (batch_size, ndataset2, nchannel2) TF tensor 189 | mlp: list of int32 -- output size for MLP on each point 190 | Return: 191 | new_features: (batch_size, ndataset1, mlp[-1]) TF tensor 192 | ''' 193 | with tf.variable_scope(scope) as sc: 194 | if upsampling: 195 | idx, dist = knn_search(data_xyz=xyz2, search_xyz=xyz1, k_num=3) 196 | dist = 1.0/tf.maximum(dist, 1e-10) 197 | weight = tf.nn.softmax(dist, axis=2) 198 | #with tf.device('/cpu:0'): 199 | interpolated_features = three_interpolate(features2, idx, weight) 200 | else: 201 | interpolated_features = features2 202 | 203 | new_features = tf.concat(axis=2, values=[interpolated_features, features1]) # B,ndataset1,nchannel1+nchannel2 204 | 205 | if mlp is not None: 206 | new_features = tf_util.MLP_1d(new_features, mlp, is_training, bn_decay, bn) 207 | 208 | return new_features 209 | 210 | ############################################################################################################################## 211 | #------------------------------------------crf modules------------------------------------------------------------ 212 | ############################################################################################################################## 213 | 214 | def crf_layer(xyz, rgb, features, idx, is_training, bn_decay, scope, iter_num=1): 215 | """ 216 | CRF module with Gaussian kernel as post-processing 217 | """ 218 | with tf.variable_scope(scope) as sc: 219 | #idx, dist = knn_search(xyz, xyz, knn_num) 220 | grouped_xyz = group_point(xyz, idx) # (batch_size, num_point, knn_num, 3) 221 | grouped_xyz -= tf.expand_dims(xyz, 2) 222 | grouped_rgb = group_point(rgb, idx) # (batch_size, num_point, knn_num, 3) 223 | grouped_rgb -= tf.expand_dims(rgb, 2) # translation normalization 224 | 225 | #Gaussian weights 226 | theta = tf.Variable(tf.truncated_normal([3], mean=1, stddev=0.01)) #(theta_alpha, theta_beta, theta_gama) 227 | W = tf.Variable(tf.truncated_normal([2], mean=0.5, stddev=0.01)) 228 | 229 | num_class = features.get_shape()[-1].value 230 | compatible_weight = tf.get_variable('compatible_weight', shape=[1, num_class, num_class], initializer=tf.constant_initializer(np.identity(num_class))) 231 | 232 | #initial normlizing 233 | #features = tf.nn.softmax(features) 234 | for i in range(iter_num): 235 | features_normed = tf.nn.softmax(features) 236 | grouped_features = group_point(features_normed, idx) #(batch_size, num_point, knn_num, channels) 237 | 238 | #compute weights with Gaussian kernels 239 | ker_appearance = tf.reduce_sum(tf.square(grouped_xyz*theta[0]), axis=3) + tf.reduce_sum(tf.square(grouped_rgb*theta[1]), axis=3) #(batch_size, num_point, knn_num) 240 | ker_smooth = tf.reduce_sum(tf.square(grouped_xyz*theta[2]), axis=3) 241 | ker_appearance = tf.exp(-ker_appearance) 242 | ker_smooth = tf.exp(-ker_smooth) 243 | 244 | Q_weight = W[0]*ker_appearance + W[1]*ker_smooth #(batch_size, num_point, knn_num) 245 | Q_weight = tf.expand_dims(Q_weight, axis=2) #(batch_size, num_point, 1, knn_num) 246 | 247 | #message passing 248 | Q_til_weighted = tf.matmul(Q_weight, grouped_features) #(batch_size, num_point, 1, channels) 249 | Q_til_weighted = tf.squeeze(Q_til_weighted, axis=2) #(batch_size, num_point, channels) 250 | 251 | #compatibility transform 252 | Q_til_weighted = tf.nn.conv1d(Q_til_weighted, compatible_weight, 1, padding='SAME') 253 | 254 | #adding unary potentials 255 | features += Q_til_weighted 256 | 257 | #normalization 258 | #features = tf.nn.softmax(features) 259 | return features 260 | -------------------------------------------------------------------------------- /utils/pcnet_util.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wleigithub/GACNet/bab6c5b45ac74dc2256bba4e5101524ca8086179/utils/pcnet_util.pyc -------------------------------------------------------------------------------- /utils/provider.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import numpy as np 4 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 5 | sys.path.append(BASE_DIR) 6 | 7 | ################################################################################################## 8 | ######################### data processing utils 9 | ################################################################################################## 10 | def shuffle_data(data, labels): 11 | """ Shuffle data and labels. 12 | Input: 13 | data: B,N,... numpy array 14 | label: B,... numpy array 15 | Return: 16 | shuffled data, label and shuffle indices 17 | """ 18 | idx = np.arange(len(labels)) 19 | np.random.shuffle(idx) 20 | return data[idx, ...], labels[idx,...], idx 21 | 22 | def shuffle_data_with_spw(data, labels, spws): 23 | """ Shuffle data and labels. 24 | Input: 25 | data: B,N,... numpy array 26 | label: B,... numpy array 27 | Return: 28 | shuffled data, label and shuffle indices 29 | """ 30 | idx = np.arange(len(labels)) 31 | np.random.shuffle(idx) 32 | return data[idx, ...], labels[idx,...], spws[idx,...], idx 33 | 34 | 35 | def rotate_point_cloud(batch_data): 36 | """ Randomly rotate the point clouds to augument the dataset 37 | rotation is per shape based along up direction 38 | Input: 39 | BxNx3 array, original batch of point clouds 40 | Return: 41 | BxNx3 array, rotated batch of point clouds 42 | """ 43 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 44 | for k in range(batch_data.shape[0]): 45 | rotation_angle = np.random.uniform() * 2 * np.pi 46 | cosval = np.cos(rotation_angle) 47 | sinval = np.sin(rotation_angle) 48 | rotation_matrix = np.array([[cosval, 0, sinval], 49 | [0, 1, 0], 50 | [-sinval, 0, cosval]]) 51 | shape_pc = batch_data[k, ...] 52 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 53 | return rotated_data 54 | 55 | def rotate_point_cloud_along_z(batch_data): 56 | """ Randomly rotate the point clouds to augument the dataset 57 | rotation is per shape based along up direction 58 | Input: 59 | BxNx3 array, original batch of point clouds 60 | Return: 61 | BxNx3 array, rotated batch of point clouds 62 | """ 63 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 64 | for k in range(batch_data.shape[0]): 65 | rotation_angle = np.random.uniform() * 2 * np.pi 66 | cosval = np.cos(rotation_angle) 67 | sinval = np.sin(rotation_angle) 68 | rotation_matrix = np.array([[cosval, sinval, 0], 69 | [-sinval, cosval, 0], 70 | [0, 0, 1]]) 71 | shape_pc = batch_data[k, ...] 72 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 73 | return rotated_data 74 | 75 | 76 | def rotate_point_cloud_by_angle(batch_data, rotation_angle): 77 | """ Rotate the point cloud along up direction with certain angle. 78 | Input: 79 | BxNx3 array, original batch of point clouds 80 | Return: 81 | BxNx3 array, rotated batch of point clouds 82 | """ 83 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 84 | for k in range(batch_data.shape[0]): 85 | #rotation_angle = np.random.uniform() * 2 * np.pi 86 | cosval = np.cos(rotation_angle) 87 | sinval = np.sin(rotation_angle) 88 | rotation_matrix = np.array([[cosval, 0, sinval], 89 | [0, 1, 0], 90 | [-sinval, 0, cosval]]) 91 | shape_pc = batch_data[k, ...] 92 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 93 | return rotated_data 94 | 95 | 96 | def jitter_point_cloud(batch_data, sigma=0.01): 97 | """ Randomly jitter points. jittering is per point. 98 | Input: 99 | BxNx3 array, original batch of point clouds 100 | Return: 101 | BxNx3 array, jittered batch of point clouds 102 | """ 103 | B, N, C = batch_data.shape 104 | jittered_data = sigma * np.random.randn(B, N, C) 105 | jittered_data += batch_data 106 | return jittered_data 107 | 108 | 109 | def shift_point_cloud(batch_data, shift_range=0.1): 110 | """ Randomly shift point cloud. Shift is per point cloud. 111 | Input: 112 | BxNx3 array, original batch of point clouds 113 | Return: 114 | BxNx3 array, shifted batch of point clouds 115 | """ 116 | B, N, C = batch_data.shape 117 | shifts = np.random.uniform(-shift_range, shift_range, (B,3)) 118 | for batch_index in range(B): 119 | batch_data[batch_index,:,:] += shifts[batch_index,:] 120 | return batch_data 121 | 122 | 123 | def random_scale_point_cloud(batch_data, scale_low=0.8, scale_high=1.2): 124 | """ Randomly scale the point cloud. Scale is per point cloud. 125 | Input:s 126 | BxNx3 array, original batch of point clouds 127 | Return: 128 | BxNx3 array, scaled batch of point clouds 129 | """ 130 | B, N, C = batch_data.shape 131 | scales = np.random.uniform(scale_low, scale_high, B) 132 | for batch_index in range(B): 133 | batch_data[batch_index,:,:] *= scales[batch_index] 134 | return batch_data 135 | 136 | 137 | ################################################################################################## 138 | ######################### h5 io utils 139 | ################################################################################################## 140 | import h5py 141 | 142 | def getDataFiles(list_filename): 143 | return [line.rstrip() for line in open(list_filename)] 144 | 145 | def load_h5(h5_filename): 146 | f = h5py.File(h5_filename) 147 | data = f['data'][:] 148 | label = f['label'][:] 149 | return (data, label) 150 | 151 | def loadDataFile(filename): 152 | return load_h5(filename) 153 | 154 | def loadDataFile_with_spw(h5_filename): 155 | f = h5py.File(h5_filename) 156 | data = f['data'][:] 157 | label = f['label'][:] 158 | spw = f['spw'][:] 159 | #spw = f['sampleweight'][:] 160 | return (data, label, spw) 161 | 162 | def loadDataFile_with_spw_catl(h5_filename): 163 | f = h5py.File(h5_filename) 164 | data = f['data'][:] 165 | label = f['label'][:] 166 | spw = f['spw'][:] 167 | catl = f['catl'][:] 168 | return (data, label, spw, catl) 169 | 170 | def load_h5_eig_curve(h5_filename): 171 | f = h5py.File(h5_filename) 172 | eig_curve = f['eig_curve'][:] 173 | return eig_curve 174 | 175 | def utils_iou(gt_classes, positive_classes, true_positive_classes, pred_label, gt_label, mask=None): 176 | pred_label = np.reshape(pred_label, [-1]) 177 | gt_label = np.reshape(gt_label, [-1]) 178 | if mask is not None: 179 | mask = np.reshape(mask, [-1]) 180 | for j in range(gt_label.shape[0]): 181 | if mask[j]==0: continue 182 | gt_l = int(gt_label[j]) 183 | pred_l = int(pred_label[j]) 184 | gt_classes[gt_l] += 1 185 | positive_classes[pred_l] += 1 186 | true_positive_classes[gt_l] += int(gt_l==pred_l) 187 | else: 188 | for j in range(gt_label.shape[0]): 189 | gt_l = int(gt_label[j]) 190 | pred_l = int(pred_label[j]) 191 | gt_classes[gt_l] += 1 192 | positive_classes[pred_l] += 1 193 | true_positive_classes[gt_l] += int(gt_l==pred_l) 194 | return gt_classes, positive_classes, true_positive_classes 195 | 196 | 197 | ################################################################################################## 198 | ######################### ply io utils 199 | ################################################################################################## 200 | from plyfile import PlyData, PlyElement 201 | 202 | def read_xyzrgbIL_ply(filename): 203 | """ read XYZ point cloud from filename PLY file """ 204 | plydata = PlyData.read(filename) 205 | pc = plydata['vertex'].data 206 | pc_array = np.array([[x, y, z, red, green, blue, scalar_intensity, scalar_label] for x,y,z, red, green, blue, scalar_intensity, scalar_label in pc]) 207 | return pc_array 208 | 209 | def read_xyzrgbI_ply(filename): 210 | """ read XYZ point cloud from filename PLY file """ 211 | plydata = PlyData.read(filename) 212 | pc = plydata['vertex'].data 213 | pc_array = np.array([[x, y, z, red, green, blue, intensity] for x,y,z, red, green, blue, intensity in pc]) 214 | return pc_array 215 | 216 | def read_xyzrgbL_ply(filename): 217 | """ read XYZ point cloud from filename PLY file """ 218 | plydata = PlyData.read(filename) 219 | pc = plydata['vertex'].data 220 | pc_array = np.array([[x, y, z, red, green, blue, label] for x,y,z, red, green, blue, label in pc]) 221 | return pc_array 222 | 223 | def write_xyzrgbL_ply(points, labels, filename, text=False): 224 | """ input: Nx3, write points to filename as PLY format. """ 225 | points = [(points[i,0], points[i,1], points[i,2], points[i,3], points[i,4], points[i,5], labels[i]) for i in range(points.shape[0])] 226 | vertex = np.array(points, dtype=[('x', 'f4'), ('y', 'f4'),('z', 'f4'), ('red', 'u1'), ('green', 'u1'),('blue', 'u1'), ('label', 'u1')]) 227 | el = PlyElement.describe(vertex, 'vertex', comments=['vertices']) 228 | PlyData([el], text=text).write(filename) 229 | 230 | def write_xyzL_ply(points, labels, filename, text=False): 231 | """ input: Nx3, write points to filename as PLY format. """ 232 | points = [(points[i,0], points[i,1], points[i,2], labels[i]) for i in range(points.shape[0])] 233 | vertex = np.array(points, dtype=[('x', 'f4'), ('y', 'f4'),('z', 'f4'), ('label', 'u1')]) 234 | el = PlyElement.describe(vertex, 'vertex', comments=['vertices']) 235 | PlyData([el], text=text).write(filename) 236 | -------------------------------------------------------------------------------- /utils/provider.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wleigithub/GACNet/bab6c5b45ac74dc2256bba4e5101524ca8086179/utils/provider.pyc -------------------------------------------------------------------------------- /utils/tf_util.py: -------------------------------------------------------------------------------- 1 | """ Wrapper functions for TensorFlow layers. 2 | """ 3 | 4 | import numpy as np 5 | import tensorflow as tf 6 | 7 | def variable_with_weight_decay(name, shape, wd, trainable=True): 8 | """Helper to create an initialized Variable with weight decay. 9 | 10 | Note that the Variable is initialized with a truncated normal distribution. 11 | A weight decay is added only if one is specified. 12 | 13 | Args: 14 | name: name of the variable 15 | shape: list of ints 16 | wd: add L2Loss weight decay multiplied by this float. If None, weight 17 | decay is not added for this Variable. 18 | Returns: 19 | Variable Tensor 20 | """ 21 | var = tf.get_variable(name, shape=shape, dtype=tf.float32, initializer=tf.contrib.layers.xavier_initializer(), trainable=trainable) 22 | if wd is not None: 23 | weight_decay = tf.multiply(tf.nn.l2_loss(var), wd, name='weight_loss') 24 | tf.add_to_collection('losses', weight_decay) 25 | return var 26 | 27 | def get_bias(shape, trainable=True): 28 | return tf.get_variable(name='biases', shape=shape, dtype=tf.float32, initializer=tf.constant_initializer(0.0), trainable=trainable) 29 | 30 | def conv1d(inputs,num_output_channels,kernel_size,scope, stride=1,padding='SAME',weight_decay=0.0,bn=False, 31 | bn_decay=None,is_training=None, activation_fn=tf.nn.elu, trainable=True): 32 | """ 1D convolution with non-linear operation. 33 | Args: 34 | inputs: 3-D tensor variable BxLxC 35 | num_output_channels: int 36 | kernel_size: int 37 | scope: string 38 | stride: int 39 | padding: 'SAME' or 'VALID' 40 | weight_decay: float 41 | activation_fn: function 42 | bn: bool, whether to use batch norm 43 | bn_decay: float or float tensor variable in [0,1] 44 | is_training: bool Tensor variable 45 | 46 | Returns: 47 | Variable tensor 48 | """ 49 | with tf.variable_scope(scope) as sc: 50 | num_in_channels = inputs.get_shape()[-1].value 51 | kernel_shape = [kernel_size, num_in_channels, num_output_channels] 52 | kernel = variable_with_weight_decay('weights', shape=kernel_shape, wd=weight_decay, trainable=trainable) 53 | outputs = tf.nn.conv1d(inputs, kernel,stride=stride,padding=padding) 54 | biases = get_bias([num_output_channels], trainable) 55 | outputs = tf.nn.bias_add(outputs, biases) 56 | 57 | if bn: 58 | outputs = batch_norm(outputs, is_training, bn_decay=bn_decay, scope='bn') 59 | 60 | if activation_fn is not None: 61 | outputs = activation_fn(outputs) 62 | return outputs 63 | 64 | 65 | def conv2d(inputs, num_output_channels,kernel_size,scope,stride=[1, 1],padding='SAME',weight_decay=0.0,bn=False, 66 | bn_decay=None,is_training=None, activation_fn=tf.nn.elu, trainable=True): 67 | """ 2D convolution with non-linear operation. 68 | 69 | Args: 70 | inputs: 4-D tensor variable BxHxWxC 71 | num_output_channels: int 72 | kernel_size: a list of 2 ints 73 | scope: string 74 | stride: a list of 2 ints 75 | padding: 'SAME' or 'VALID' 76 | weight_decay: float 77 | activation_fn: function 78 | bn: bool, whether to use batch norm 79 | bn_decay: float or float tensor variable in [0,1] 80 | is_training: bool Tensor variable 81 | 82 | Returns: 83 | Variable tensor 84 | """ 85 | with tf.variable_scope(scope) as sc: 86 | kernel_h, kernel_w = kernel_size 87 | num_in_channels = inputs.get_shape()[-1].value 88 | kernel_shape = [kernel_h, kernel_w, num_in_channels, num_output_channels] 89 | kernel = variable_with_weight_decay('weights',shape=kernel_shape, wd=weight_decay, trainable=trainable) 90 | stride_h, stride_w = stride 91 | outputs = tf.nn.conv2d(inputs, kernel,[1, stride_h, stride_w, 1],padding=padding) 92 | biases = get_bias([num_output_channels], trainable) 93 | outputs = tf.nn.bias_add(outputs, biases) 94 | 95 | if bn: 96 | outputs = batch_norm(outputs, is_training, bn_decay=bn_decay, scope='bn') 97 | 98 | if activation_fn is not None: 99 | outputs = activation_fn(outputs) 100 | return outputs 101 | 102 | 103 | def MLP_2d(features, mlp, is_training, bn_decay, bn, scope='mlp2d', padding='VALID'): 104 | for i, num_out_channel in enumerate(mlp): 105 | features = conv2d(features, num_out_channel, [1,1], 106 | padding=padding, stride=[1,1], 107 | bn=bn, is_training=is_training, 108 | scope= scope+'_%d'%(i), bn_decay=bn_decay) 109 | return features 110 | 111 | 112 | def MLP_1d(features, mlp, is_training, bn_decay, bn, scope='mlp1d', padding='VALID'): 113 | for i, num_out_channel in enumerate(mlp): 114 | features = conv1d(features, num_out_channel, 1, 115 | padding=padding, stride=1, 116 | bn=bn, is_training=is_training, 117 | scope=scope+'_%d'%(i), bn_decay=bn_decay) 118 | return features 119 | 120 | 121 | def fully_connected(inputs,num_outputs,scope,weight_decay=0.0,activation_fn=tf.nn.elu, 122 | bn=False, bn_decay=None,is_training=None, trainable=True): 123 | """ Fully connected layer with non-linear operation. 124 | Args: 125 | inputs: 2-D tensor BxN 126 | num_outputs: int 127 | 128 | Returns: 129 | Variable tensor of size B x num_outputs. 130 | """ 131 | with tf.variable_scope(scope) as sc: 132 | num_input_units = inputs.get_shape()[-1].value 133 | weights = variable_with_weight_decay('weights',shape=[num_input_units, num_outputs],wd=weight_decay, trainable=trainable) 134 | outputs = tf.matmul(inputs, weights) 135 | biases = get_bias([num_outputs], trainable) 136 | outputs = tf.nn.bias_add(outputs, biases) 137 | 138 | if bn: 139 | outputs = batch_norm(outputs, is_training, bn_decay, 'bn') 140 | 141 | if activation_fn is not None: 142 | outputs = activation_fn(outputs) 143 | return outputs 144 | 145 | 146 | def dropout(inputs, 147 | is_training, 148 | scope, 149 | keep_prob=0.5, 150 | noise_shape=None): 151 | """ Dropout layer. 152 | 153 | Args: 154 | inputs: tensor 155 | is_training: boolean tf.Variable 156 | scope: string 157 | keep_prob: float in [0,1] 158 | noise_shape: list of ints 159 | 160 | Returns: 161 | tensor variable 162 | """ 163 | with tf.variable_scope(scope) as sc: 164 | outputs = tf.cond(is_training, 165 | lambda: tf.nn.dropout(inputs, keep_prob, noise_shape), 166 | lambda: inputs) 167 | return outputs 168 | 169 | 170 | def batch_norm(data, is_training, bn_decay, scope): 171 | decay = bn_decay if bn_decay is not None else 0.5 172 | return tf.layers.batch_normalization(data, momentum=decay, training=is_training, 173 | beta_regularizer=tf.contrib.layers.l2_regularizer(scale=1.0), 174 | gamma_regularizer=tf.contrib.layers.l2_regularizer(scale=1.0), 175 | name=scope) 176 | 177 | 178 | def _batch_norm(inputs, is_training, bn_decay, scope): 179 | with tf.variable_scope(scope): 180 | beta = tf.Variable(tf.constant(0.0, shape=[inputs.shape[-1]]), name='beta', trainable=True) 181 | gamma = tf.Variable(tf.constant(1.0, shape=[inputs.shape[-1]]), name='gamma', trainable=True) 182 | batch_mean, batch_var = tf.nn.moments(inputs, range(len(inputs.shape) - 1), name='moments') 183 | 184 | decay = bn_decay if bn_decay is not None else 0.5 185 | ema = tf.train.ExponentialMovingAverage(decay=decay) 186 | 187 | def mean_var_with_update(): 188 | ema_apply_op = ema.apply([batch_mean, batch_var]) 189 | with tf.control_dependencies([ema_apply_op]): 190 | return tf.identity(batch_mean), tf.identity(batch_var) 191 | 192 | mean, var = tf.cond(is_training, mean_var_with_update, 193 | lambda: (ema.average(batch_mean), ema.average(batch_var))) 194 | normed = tf.nn.batch_normalization(inputs, mean, var, beta, gamma, 1e-3) 195 | return normed 196 | 197 | 198 | def get_learning_rate(batch, BASE_LEARNING_RATE, BATCH_SIZE, DECAY_STEP, DECAY_RATE): 199 | """Learning rate decay. 200 | """ 201 | learning_rate = tf.train.exponential_decay( 202 | BASE_LEARNING_RATE, # Base learning rate. 203 | batch * BATCH_SIZE, # Current index into the dataset. 204 | DECAY_STEP, # Decay step. 205 | DECAY_RATE, # Decay rate. 206 | staircase=True) 207 | learing_rate = tf.maximum(learning_rate, 0.00001) # CLIP THE LEARNING RATE!! 208 | return learning_rate 209 | 210 | 211 | def get_bn_decay(batch, BATCH_SIZE, BN_DECAY_DECAY_STEP=300000.0, BN_DECAY_DECAY_RATE=0.5, BN_INIT_DECAY = 0.5, BN_DECAY_CLIP = 0.99): 212 | """ 213 | BN_DECAY_DECAY_STEP: accept BN_DECAY_DECAY_STEP with float type 214 | """ 215 | bn_momentum = tf.train.exponential_decay( 216 | BN_INIT_DECAY, 217 | batch*BATCH_SIZE, 218 | BN_DECAY_DECAY_STEP, 219 | BN_DECAY_DECAY_RATE, 220 | staircase=True) 221 | bn_decay = tf.minimum(BN_DECAY_CLIP, 1 - bn_momentum) 222 | return bn_decay -------------------------------------------------------------------------------- /utils/tf_util.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wleigithub/GACNet/bab6c5b45ac74dc2256bba4e5101524ca8086179/utils/tf_util.pyc --------------------------------------------------------------------------------