├── .gitignore ├── README.md ├── cpp_wrappers ├── compile_wrappers.sh ├── cpp_subsampling │ ├── grid_subsampling │ │ ├── grid_subsampling.cpp │ │ └── grid_subsampling.h │ ├── setup.py │ └── wrapper.cpp └── cpp_utils │ ├── cloud │ ├── cloud.cpp │ └── cloud.h │ └── nanoflann │ └── nanoflann.hpp ├── datasets ├── ModelNet.py ├── ShapeNet.py ├── common.py └── dataloader.py ├── environment.yml ├── kernels ├── convolution_ops.py └── kernel_points.py ├── misc └── num_seg_classes.txt ├── models ├── KPCNN.py ├── KPFCNN.py └── network_blocks.py ├── pytorch_ops ├── batch_find_neighbors.cpp ├── cpp_utils │ ├── cloud │ │ ├── cloud.cpp │ │ └── cloud.h │ └── nanoflann │ │ └── nanoflann.hpp ├── neighbors │ ├── neighbors.cpp │ └── neighbors.h └── setup.py ├── trainer.py ├── trainer_cls.py ├── training_ModelNet.py ├── training_ShapeNetCls.py ├── training_ShapeNetPart.py └── utils ├── config.py ├── metrics.py ├── ply.py ├── pointcloud.py └── timer.py /.gitignore: -------------------------------------------------------------------------------- 1 | data/* 2 | kernels/*.ply 3 | snapshot/* 4 | tensorboard/* 5 | */__pycache__/* 6 | __pycache__/* 7 | */pyc 8 | .idea/ 9 | *.o 10 | *.so 11 | pytorch_ops/* -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## KPConv.pytorch 2 | 3 | This repo is implementation for KPConv(https://arxiv.org/abs/1904.08889) in pytorch. 4 | 5 | ## TODO 6 | There are still some works to be done: 7 | - [x] Deformable KPConv. Currently I have only implemented the rigid KPConv. 8 | - [ ] Regularization loss for the deformable convolution needs to be implemented. I have tried using the deformable convolution layer in part segmention on shapenet without the regularization term, the performance is similar with the rigid convolution counterparts. 9 | - [x] Speed up. For current implementation, the `collate_fn` where the neighbor indices and pooling indices are calculated, is too slow. In the tf version, the author implement 2 tensroflow C++ wrapper which is quite efficient. I am planing to write C++ extention using pytorch. 10 | - [ ] But after I implemented the C++ extention, the evaluation time reduces significantly while the model forward and backward pass still cost about 0.8s per iteration. 11 | - [ ] Maybe other datasets. 12 | 13 | 14 | ## Installation 15 | 16 | 1. Create an environment from the environment.yml file, 17 | ``` 18 | conda env create -f environment.yml 19 | ``` 20 | 2. Compile the customized Tensorflow operators and C++ extension module following the [installation instructions](https://github.com/HuguesTHOMAS/KPConv/blob/master/INSTALL.md) provided by the authors. 21 | 3. Go to `pytorch_ops` dictionary and run `python setup.py install` to build and install the C++ extension for `batch_find_neighbors` function. 22 | 23 | 24 | ## Experiments 25 | 26 | Due to the time limitation, I have just implemented the experiments on ShapeNet(classification and part segmentation) and ModelNet40. 27 | 28 | - Shape Classification on ModelNet40 or ShapeNet. 29 | ``` 30 | python training_ModelNet.py[training_ShapeNetCls.py] 31 | ``` 32 | 33 | - Part Segmentation on ShapeNet. (I have only implemented the single class part segmentation.) 34 | ``` 35 | python training_ShapeNetPart.py 36 | ``` 37 | 38 | ## Acknowledgment 39 | 40 | Thank @HuguesTHOMAS for sharing the tensorflow version and valuable explainations. 41 | 42 | -------------------------------------------------------------------------------- /cpp_wrappers/compile_wrappers.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Compile cpp subsampling 4 | cd cpp_subsampling 5 | python3 setup.py build_ext --inplace 6 | cd .. 7 | 8 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "grid_subsampling.h" 3 | 4 | 5 | void grid_subsampling(vector& original_points, 6 | vector& subsampled_points, 7 | vector& original_features, 8 | vector& subsampled_features, 9 | vector& original_classes, 10 | vector& subsampled_classes, 11 | float sampleDl, 12 | int verbose) { 13 | 14 | // Initiate variables 15 | // ****************** 16 | 17 | // Number of points in the cloud 18 | size_t N = original_points.size(); 19 | 20 | // Dimension of the features 21 | size_t fdim = original_features.size() / N; 22 | size_t ldim = original_classes.size() / N; 23 | 24 | // Limits of the cloud 25 | PointXYZ minCorner = min_point(original_points); 26 | PointXYZ maxCorner = max_point(original_points); 27 | PointXYZ originCorner = floor(minCorner * (1/sampleDl)) * sampleDl; 28 | 29 | // Dimensions of the grid 30 | size_t sampleNX = (size_t)floor((maxCorner.x - originCorner.x) / sampleDl) + 1; 31 | size_t sampleNY = (size_t)floor((maxCorner.y - originCorner.y) / sampleDl) + 1; 32 | //size_t sampleNZ = (size_t)floor((maxCorner.z - originCorner.z) / sampleDl) + 1; 33 | 34 | // Check if features and classes need to be processed 35 | bool use_feature = original_features.size() > 0; 36 | bool use_classes = original_classes.size() > 0; 37 | 38 | 39 | // Create the sampled map 40 | // ********************** 41 | 42 | // Verbose parameters 43 | int i = 0; 44 | int nDisp = N / 100; 45 | 46 | // Initiate variables 47 | size_t iX, iY, iZ, mapIdx; 48 | unordered_map data; 49 | 50 | for (auto& p : original_points) 51 | { 52 | // Position of point in sample map 53 | iX = (size_t)floor((p.x - originCorner.x) / sampleDl); 54 | iY = (size_t)floor((p.y - originCorner.y) / sampleDl); 55 | iZ = (size_t)floor((p.z - originCorner.z) / sampleDl); 56 | mapIdx = iX + sampleNX*iY + sampleNX*sampleNY*iZ; 57 | 58 | // If not already created, create key 59 | if (data.count(mapIdx) < 1) 60 | data.emplace(mapIdx, SampledData(fdim, ldim)); 61 | 62 | // Fill the sample map 63 | if (use_feature && use_classes) 64 | data[mapIdx].update_all(p, original_features.begin() + i * fdim, original_classes.begin() + i * ldim); 65 | else if (use_feature) 66 | data[mapIdx].update_features(p, original_features.begin() + i * fdim); 67 | else if (use_classes) 68 | data[mapIdx].update_classes(p, original_classes.begin() + i * ldim); 69 | else 70 | data[mapIdx].update_points(p); 71 | 72 | // Display 73 | i++; 74 | if (verbose > 1 && i%nDisp == 0) 75 | std::cout << "\rSampled Map : " << std::setw(3) << i / nDisp << "%"; 76 | 77 | } 78 | 79 | // Divide for barycentre and transfer to a vector 80 | subsampled_points.reserve(data.size()); 81 | if (use_feature) 82 | subsampled_features.reserve(data.size() * fdim); 83 | if (use_classes) 84 | subsampled_classes.reserve(data.size() * ldim); 85 | for (auto& v : data) 86 | { 87 | subsampled_points.push_back(v.second.point * (1.0 / v.second.count)); 88 | if (use_feature) 89 | { 90 | float count = (float)v.second.count; 91 | transform(v.second.features.begin(), 92 | v.second.features.end(), 93 | v.second.features.begin(), 94 | [count](float f) { return f / count;}); 95 | subsampled_features.insert(subsampled_features.end(),v.second.features.begin(),v.second.features.end()); 96 | } 97 | if (use_classes) 98 | { 99 | for (int i = 0; i < ldim; i++) 100 | subsampled_classes.push_back(max_element(v.second.labels[i].begin(), v.second.labels[i].end())->first); 101 | } 102 | } 103 | 104 | return; 105 | } 106 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h: -------------------------------------------------------------------------------- 1 | 2 | 3 | #include "../../cpp_utils/cloud/cloud.h" 4 | 5 | #include 6 | #include 7 | 8 | using namespace std; 9 | 10 | class SampledData 11 | { 12 | public: 13 | 14 | // Elements 15 | // ******** 16 | 17 | int count; 18 | PointXYZ point; 19 | vector features; 20 | vector> labels; 21 | 22 | 23 | // Methods 24 | // ******* 25 | 26 | // Constructor 27 | SampledData() 28 | { 29 | count = 0; 30 | point = PointXYZ(); 31 | } 32 | 33 | SampledData(const size_t fdim, const size_t ldim) 34 | { 35 | count = 0; 36 | point = PointXYZ(); 37 | features = vector(fdim); 38 | labels = vector>(ldim); 39 | } 40 | 41 | // Method Update 42 | void update_all(const PointXYZ p, vector::iterator f_begin, vector::iterator l_begin) 43 | { 44 | count += 1; 45 | point += p; 46 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 47 | int i = 0; 48 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 49 | { 50 | labels[i][*it] += 1; 51 | i++; 52 | } 53 | return; 54 | } 55 | void update_features(const PointXYZ p, vector::iterator f_begin) 56 | { 57 | count += 1; 58 | point += p; 59 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 60 | return; 61 | } 62 | void update_classes(const PointXYZ p, vector::iterator l_begin) 63 | { 64 | count += 1; 65 | point += p; 66 | int i = 0; 67 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 68 | { 69 | labels[i][*it] += 1; 70 | i++; 71 | } 72 | return; 73 | } 74 | void update_points(const PointXYZ p) 75 | { 76 | count += 1; 77 | point += p; 78 | return; 79 | } 80 | }; 81 | 82 | 83 | 84 | void grid_subsampling(vector& original_points, 85 | vector& subsampled_points, 86 | vector& original_features, 87 | vector& subsampled_features, 88 | vector& original_classes, 89 | vector& subsampled_classes, 90 | float sampleDl, 91 | int verbose); 92 | 93 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup, Extension 2 | import numpy.distutils.misc_util 3 | 4 | # Adding OpenCV to project 5 | # ************************ 6 | 7 | # Adding sources of the project 8 | # ***************************** 9 | 10 | m_name = "grid_subsampling" 11 | 12 | SOURCES = ["../cpp_utils/cloud/cloud.cpp", 13 | "grid_subsampling/grid_subsampling.cpp", 14 | "wrapper.cpp"] 15 | 16 | module = Extension(m_name, 17 | sources=SOURCES, 18 | extra_compile_args=['-std=c++11', 19 | '-D_GLIBCXX_USE_CXX11_ABI=0']) 20 | 21 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | -------------------------------------------------------------------------------- /cpp_wrappers/cpp_subsampling/wrapper.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include "grid_subsampling/grid_subsampling.h" 4 | #include 5 | 6 | 7 | 8 | // docstrings for our module 9 | // ************************* 10 | 11 | static char module_docstring[] = "This module provides an interface for the subsampling of a pointcloud"; 12 | 13 | static char compute_docstring[] = "function subsampling a pointcloud"; 14 | 15 | 16 | // Declare the functions 17 | // ********************* 18 | 19 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds); 20 | 21 | 22 | // Specify the members of the module 23 | // ********************************* 24 | 25 | static PyMethodDef module_methods[] = 26 | { 27 | { "compute", (PyCFunction)grid_subsampling_compute, METH_VARARGS | METH_KEYWORDS, compute_docstring }, 28 | {NULL, NULL, 0, NULL} 29 | }; 30 | 31 | 32 | // Initialize the module 33 | // ********************* 34 | 35 | static struct PyModuleDef moduledef = 36 | { 37 | PyModuleDef_HEAD_INIT, 38 | "grid_subsampling", // m_name 39 | module_docstring, // m_doc 40 | -1, // m_size 41 | module_methods, // m_methods 42 | NULL, // m_reload 43 | NULL, // m_traverse 44 | NULL, // m_clear 45 | NULL, // m_free 46 | }; 47 | 48 | PyMODINIT_FUNC PyInit_grid_subsampling(void) 49 | { 50 | import_array(); 51 | return PyModule_Create(&moduledef); 52 | } 53 | 54 | 55 | // Actual wrapper 56 | // ************** 57 | 58 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds) 59 | { 60 | 61 | // Manage inputs 62 | // ************* 63 | 64 | // Args containers 65 | PyObject *points_obj = NULL; 66 | PyObject *features_obj = NULL; 67 | PyObject *classes_obj = NULL; 68 | 69 | // Keywords containers 70 | static char *kwlist[] = {"points", "features", "classes", "sampleDl", "method", "verbose", NULL }; 71 | float sampleDl = 0.1; 72 | const char *method_buffer = "barycenters"; 73 | int verbose = 0; 74 | 75 | // Parse the input 76 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "O|$OOfsi", kwlist, &points_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &verbose)) 77 | { 78 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); 79 | return NULL; 80 | } 81 | 82 | // Get the method argument 83 | string method(method_buffer); 84 | 85 | // Interpret method 86 | if (method.compare("barycenters") && method.compare("voxelcenters")) 87 | { 88 | PyErr_SetString(PyExc_RuntimeError, "Error parsing method. Valid method names are \"barycenters\" and \"voxelcenters\" "); 89 | return NULL; 90 | } 91 | 92 | // Check if using features or classes 93 | bool use_feature = true, use_classes = true; 94 | if (features_obj == NULL) 95 | use_feature = false; 96 | if (classes_obj == NULL) 97 | use_classes = false; 98 | 99 | // Interpret the input objects as numpy arrays. 100 | PyObject *points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY); 101 | PyObject *features_array = NULL; 102 | PyObject *classes_array = NULL; 103 | if (use_feature) 104 | features_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY); 105 | if (use_classes) 106 | classes_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY); 107 | 108 | // Verify data was load correctly. 109 | if (points_array == NULL) 110 | { 111 | Py_XDECREF(points_array); 112 | Py_XDECREF(classes_array); 113 | Py_XDECREF(features_array); 114 | PyErr_SetString(PyExc_RuntimeError, "Error converting input points to numpy arrays of type float32"); 115 | return NULL; 116 | } 117 | if (use_feature && features_array == NULL) 118 | { 119 | Py_XDECREF(points_array); 120 | Py_XDECREF(classes_array); 121 | Py_XDECREF(features_array); 122 | PyErr_SetString(PyExc_RuntimeError, "Error converting input features to numpy arrays of type float32"); 123 | return NULL; 124 | } 125 | if (use_classes && classes_array == NULL) 126 | { 127 | Py_XDECREF(points_array); 128 | Py_XDECREF(classes_array); 129 | Py_XDECREF(features_array); 130 | PyErr_SetString(PyExc_RuntimeError, "Error converting input classes to numpy arrays of type int32"); 131 | return NULL; 132 | } 133 | 134 | // Check that the input array respect the dims 135 | if ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3) 136 | { 137 | Py_XDECREF(points_array); 138 | Py_XDECREF(classes_array); 139 | Py_XDECREF(features_array); 140 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : points.shape is not (N, 3)"); 141 | return NULL; 142 | } 143 | if (use_feature && ((int)PyArray_NDIM(features_array) != 2)) 144 | { 145 | Py_XDECREF(points_array); 146 | Py_XDECREF(classes_array); 147 | Py_XDECREF(features_array); 148 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 149 | return NULL; 150 | } 151 | 152 | if (use_classes && (int)PyArray_NDIM(classes_array) > 2) 153 | { 154 | Py_XDECREF(points_array); 155 | Py_XDECREF(classes_array); 156 | Py_XDECREF(features_array); 157 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 158 | return NULL; 159 | } 160 | 161 | // Number of points 162 | int N = (int)PyArray_DIM(points_array, 0); 163 | 164 | // Dimension of the features 165 | int fdim = 0; 166 | if (use_feature) 167 | fdim = (int)PyArray_DIM(features_array, 1); 168 | 169 | //Dimension of labels 170 | int ldim = 1; 171 | if (use_classes && (int)PyArray_NDIM(classes_array) == 2) 172 | ldim = (int)PyArray_DIM(classes_array, 1); 173 | 174 | // Check that the input array respect the number of points 175 | if (use_feature && (int)PyArray_DIM(features_array, 0) != N) 176 | { 177 | Py_XDECREF(points_array); 178 | Py_XDECREF(classes_array); 179 | Py_XDECREF(features_array); 180 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 181 | return NULL; 182 | } 183 | if (use_classes && (int)PyArray_DIM(classes_array, 0) != N) 184 | { 185 | Py_XDECREF(points_array); 186 | Py_XDECREF(classes_array); 187 | Py_XDECREF(features_array); 188 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 189 | return NULL; 190 | } 191 | 192 | 193 | // Call the C++ function 194 | // ********************* 195 | 196 | // Create pyramid 197 | if (verbose > 0) 198 | cout << "Computing cloud pyramid with support points: " << endl; 199 | 200 | 201 | // Convert PyArray to Cloud C++ class 202 | vector original_points; 203 | vector original_features; 204 | vector original_classes; 205 | original_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N); 206 | if (use_feature) 207 | original_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N*fdim); 208 | if (use_classes) 209 | original_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N*ldim); 210 | 211 | // Subsample 212 | vector subsampled_points; 213 | vector subsampled_features; 214 | vector subsampled_classes; 215 | grid_subsampling(original_points, 216 | subsampled_points, 217 | original_features, 218 | subsampled_features, 219 | original_classes, 220 | subsampled_classes, 221 | sampleDl, 222 | verbose); 223 | 224 | // Check result 225 | if (subsampled_points.size() < 1) 226 | { 227 | PyErr_SetString(PyExc_RuntimeError, "Error"); 228 | return NULL; 229 | } 230 | 231 | // Manage outputs 232 | // ************** 233 | 234 | // Dimension of input containers 235 | npy_intp* point_dims = new npy_intp[2]; 236 | point_dims[0] = subsampled_points.size(); 237 | point_dims[1] = 3; 238 | npy_intp* feature_dims = new npy_intp[2]; 239 | feature_dims[0] = subsampled_points.size(); 240 | feature_dims[1] = fdim; 241 | npy_intp* classes_dims = new npy_intp[2]; 242 | classes_dims[0] = subsampled_points.size(); 243 | classes_dims[1] = ldim; 244 | 245 | // Create output array 246 | PyObject *res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT); 247 | PyObject *res_features_obj = NULL; 248 | PyObject *res_classes_obj = NULL; 249 | PyObject *ret = NULL; 250 | 251 | // Fill output array with values 252 | size_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float); 253 | memcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes); 254 | if (use_feature) 255 | { 256 | size_in_bytes = subsampled_points.size() * fdim * sizeof(float); 257 | res_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT); 258 | memcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes); 259 | } 260 | if (use_classes) 261 | { 262 | size_in_bytes = subsampled_points.size() * ldim * sizeof(int); 263 | res_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT); 264 | memcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes); 265 | } 266 | 267 | 268 | // Merge results 269 | if (use_feature && use_classes) 270 | ret = Py_BuildValue("NNN", res_points_obj, res_features_obj, res_classes_obj); 271 | else if (use_feature) 272 | ret = Py_BuildValue("NN", res_points_obj, res_features_obj); 273 | else if (use_classes) 274 | ret = Py_BuildValue("NN", res_points_obj, res_classes_obj); 275 | else 276 | ret = Py_BuildValue("N", res_points_obj); 277 | 278 | // Clean up 279 | // ******** 280 | 281 | Py_DECREF(points_array); 282 | Py_XDECREF(features_array); 283 | Py_XDECREF(classes_array); 284 | 285 | return ret; 286 | } -------------------------------------------------------------------------------- /cpp_wrappers/cpp_utils/cloud/cloud.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud source : 13 | // Define usefull Functions/Methods 14 | // 15 | //---------------------------------------------------- 16 | // 17 | // Hugues THOMAS - 10/02/2017 18 | // 19 | 20 | 21 | #include "cloud.h" 22 | 23 | 24 | // Getters 25 | // ******* 26 | 27 | PointXYZ max_point(std::vector points) 28 | { 29 | // Initiate limits 30 | PointXYZ maxP(points[0]); 31 | 32 | // Loop over all points 33 | for (auto p : points) 34 | { 35 | if (p.x > maxP.x) 36 | maxP.x = p.x; 37 | 38 | if (p.y > maxP.y) 39 | maxP.y = p.y; 40 | 41 | if (p.z > maxP.z) 42 | maxP.z = p.z; 43 | } 44 | 45 | return maxP; 46 | } 47 | 48 | PointXYZ min_point(std::vector points) 49 | { 50 | // Initiate limits 51 | PointXYZ minP(points[0]); 52 | 53 | // Loop over all points 54 | for (auto p : points) 55 | { 56 | if (p.x < minP.x) 57 | minP.x = p.x; 58 | 59 | if (p.y < minP.y) 60 | minP.y = p.y; 61 | 62 | if (p.z < minP.z) 63 | minP.z = p.z; 64 | } 65 | 66 | return minP; 67 | } -------------------------------------------------------------------------------- /cpp_wrappers/cpp_utils/cloud/cloud.h: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud header 13 | // 14 | //---------------------------------------------------- 15 | // 16 | // Hugues THOMAS - 10/02/2017 17 | // 18 | 19 | 20 | # pragma once 21 | 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | #include 31 | 32 | 33 | 34 | 35 | // Point class 36 | // *********** 37 | 38 | 39 | class PointXYZ 40 | { 41 | public: 42 | 43 | // Elements 44 | // ******** 45 | 46 | float x, y, z; 47 | 48 | 49 | // Methods 50 | // ******* 51 | 52 | // Constructor 53 | PointXYZ() { x = 0; y = 0; z = 0; } 54 | PointXYZ(float x0, float y0, float z0) { x = x0; y = y0; z = z0; } 55 | 56 | // array type accessor 57 | float operator [] (int i) const 58 | { 59 | if (i == 0) return x; 60 | else if (i == 1) return y; 61 | else return z; 62 | } 63 | 64 | // opperations 65 | float dot(const PointXYZ P) const 66 | { 67 | return x * P.x + y * P.y + z * P.z; 68 | } 69 | 70 | float sq_norm() 71 | { 72 | return x*x + y*y + z*z; 73 | } 74 | 75 | PointXYZ cross(const PointXYZ P) const 76 | { 77 | return PointXYZ(y*P.z - z*P.y, z*P.x - x*P.z, x*P.y - y*P.x); 78 | } 79 | 80 | PointXYZ& operator+=(const PointXYZ& P) 81 | { 82 | x += P.x; 83 | y += P.y; 84 | z += P.z; 85 | return *this; 86 | } 87 | 88 | PointXYZ& operator-=(const PointXYZ& P) 89 | { 90 | x -= P.x; 91 | y -= P.y; 92 | z -= P.z; 93 | return *this; 94 | } 95 | 96 | PointXYZ& operator*=(const float& a) 97 | { 98 | x *= a; 99 | y *= a; 100 | z *= a; 101 | return *this; 102 | } 103 | }; 104 | 105 | 106 | // Point Opperations 107 | // ***************** 108 | 109 | inline PointXYZ operator + (const PointXYZ A, const PointXYZ B) 110 | { 111 | return PointXYZ(A.x + B.x, A.y + B.y, A.z + B.z); 112 | } 113 | 114 | inline PointXYZ operator - (const PointXYZ A, const PointXYZ B) 115 | { 116 | return PointXYZ(A.x - B.x, A.y - B.y, A.z - B.z); 117 | } 118 | 119 | inline PointXYZ operator * (const PointXYZ P, const float a) 120 | { 121 | return PointXYZ(P.x * a, P.y * a, P.z * a); 122 | } 123 | 124 | inline PointXYZ operator * (const float a, const PointXYZ P) 125 | { 126 | return PointXYZ(P.x * a, P.y * a, P.z * a); 127 | } 128 | 129 | inline std::ostream& operator << (std::ostream& os, const PointXYZ P) 130 | { 131 | return os << "[" << P.x << ", " << P.y << ", " << P.z << "]"; 132 | } 133 | 134 | inline bool operator == (const PointXYZ A, const PointXYZ B) 135 | { 136 | return A.x == B.x && A.y == B.y && A.z == B.z; 137 | } 138 | 139 | inline PointXYZ floor(const PointXYZ P) 140 | { 141 | return PointXYZ(floor(P.x), floor(P.y), floor(P.z)); 142 | } 143 | 144 | 145 | PointXYZ max_point(std::vector points); 146 | PointXYZ min_point(std::vector points); 147 | 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | -------------------------------------------------------------------------------- /datasets/ModelNet.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path 3 | import numpy as np 4 | import time 5 | import pickle 6 | from datasets.common import grid_subsampling 7 | import torch.utils.data as data 8 | 9 | 10 | class ModelNetDataset(data.Dataset): 11 | classification = True 12 | 13 | def __init__(self, 14 | root, 15 | split='train', 16 | first_subsampling_dl=0.03, 17 | config=None, 18 | data_augmentation=True): 19 | self.config = config 20 | self.first_subsampling_dl = first_subsampling_dl 21 | self.root = root 22 | self.split = split 23 | self.data_augmentation = data_augmentation 24 | self.points, self.normals, self.labels = [], [], [] 25 | 26 | # Dict from labels to names 27 | self.label_to_names = {0: 'airplane', 28 | 1: 'bathtub', 29 | 2: 'bed', 30 | 3: 'bench', 31 | 4: 'bookshelf', 32 | 5: 'bottle', 33 | 6: 'bowl', 34 | 7: 'car', 35 | 8: 'chair', 36 | 9: 'cone', 37 | 10: 'cup', 38 | 11: 'curtain', 39 | 12: 'desk', 40 | 13: 'door', 41 | 14: 'dresser', 42 | 15: 'flower_pot', 43 | 16: 'glass_box', 44 | 17: 'guitar', 45 | 18: 'keyboard', 46 | 19: 'lamp', 47 | 20: 'laptop', 48 | 21: 'mantel', 49 | 22: 'monitor', 50 | 23: 'night_stand', 51 | 24: 'person', 52 | 25: 'piano', 53 | 26: 'plant', 54 | 27: 'radio', 55 | 28: 'range_hood', 56 | 29: 'sink', 57 | 30: 'sofa', 58 | 31: 'stairs', 59 | 32: 'stool', 60 | 33: 'table', 61 | 34: 'tent', 62 | 35: 'toilet', 63 | 36: 'tv_stand', 64 | 37: 'vase', 65 | 38: 'wardrobe', 66 | 39: 'xbox'} 67 | self.name_to_label = {v: k for k, v in self.label_to_names.items()} 68 | 69 | t0 = time.time() 70 | # Load wanted points if possible 71 | print(f'\nLoading {split} points') 72 | filename = os.path.join(self.root, f'{split}_{first_subsampling_dl:.3f}_record.pkl') 73 | if os.path.exists(filename): 74 | with open(filename, 'rb') as file: 75 | self.points, self.normals, self.labels = pickle.load(file) 76 | else: 77 | # Collect training file names 78 | names = np.loadtxt(os.path.join(self.root, f'modelnet40_{split}.txt'), dtype=np.str) 79 | 80 | # Collect point clouds 81 | for i, cloud_name in enumerate(names): 82 | 83 | # Read points 84 | class_folder = '_'.join(cloud_name.split('_')[:-1]) 85 | txt_file = os.path.join(self.root, class_folder, cloud_name) + '.txt' 86 | data = np.loadtxt(txt_file, delimiter=',', dtype=np.float32) 87 | 88 | # Subsample them 89 | if first_subsampling_dl > 0: 90 | points, normals = grid_subsampling(data[:, :3], 91 | features=data[:, 3:], 92 | sampleDl=first_subsampling_dl) 93 | else: 94 | points = data[:, :3] 95 | normals = data[:, 3:] 96 | 97 | # Add to list 98 | self.points += [points] 99 | self.normals += [normals] 100 | 101 | # Get labels 102 | label_names = ['_'.join(name.split('_')[:-1]) for name in names] 103 | self.labels = np.array([self.name_to_label[name] for name in label_names]) 104 | 105 | # Save for later use 106 | with open(filename, 'wb') as file: 107 | pickle.dump((self.points, self.normals, self.labels), file) 108 | 109 | lengths = [p.shape[0] for p in self.points] 110 | sizes = [l * 4 * 6 for l in lengths] 111 | print('{:.1f} MB loaded in {:.1f}s'.format(np.sum(sizes) * 1e-6, time.time() - t0)) 112 | 113 | def __getitem__(self, index): 114 | points, normals, labels = self.points[index], self.normals[index], self.labels[index] 115 | 116 | if self.data_augmentation and self.split == 'train': 117 | theta = np.random.uniform(0, np.pi * 2) 118 | rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]) 119 | # TODO: why only rotate the x and z axis?? 120 | points[:, [0, 2]] = points[:, [0, 2]].dot(rotation_matrix) # random rotation 121 | points += np.random.normal(0, 0.001, size=points.shape) # random jitter 122 | 123 | if self.config.in_features_dim == 1: 124 | features = np.ones([points.shape[0], 1]) 125 | elif self.config.in_features_dim == 4: 126 | features = np.ones([points.shape[0], 1]) 127 | features = np.concatenate([features, points], axis=1) 128 | 129 | return points, features, labels 130 | 131 | def __len__(self): 132 | return len(self.points) 133 | 134 | 135 | if __name__ == '__main__': 136 | datapath = "./data/modelnet40_normal_resampled/" 137 | from training_ModelNet import ModelNetConfig 138 | 139 | config = ModelNetConfig() 140 | 141 | print("Segmentation task:") 142 | dset = ModelNetDataset(root=datapath, config=config, first_subsampling_dl=0.01) 143 | input = dset[0] 144 | 145 | from datasets.dataloader import get_dataloader 146 | 147 | dataloader = get_dataloader(dset, batch_size=2) 148 | for iter, input in enumerate(dataloader): 149 | print(input) 150 | break 151 | -------------------------------------------------------------------------------- /datasets/ShapeNet.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path 3 | import numpy as np 4 | import json 5 | import open3d 6 | from utils.pointcloud import make_point_cloud 7 | from datasets.common import grid_subsampling 8 | import torch.utils.data as data 9 | 10 | 11 | class ShapeNetDataset(data.Dataset): 12 | # Borrow from https://github.com/fxia22/pointnet.pytorch 13 | def __init__(self, 14 | root, 15 | split='train', 16 | first_subsampling_dl=0.03, 17 | config=None, 18 | classification=False, 19 | class_choice=None, 20 | data_augmentation=True): 21 | self.config = config 22 | self.first_subsampling_dl = first_subsampling_dl 23 | self.root = root 24 | self.split = split 25 | self.cat2id = {} 26 | self.data_augmentation = data_augmentation 27 | self.classification = classification 28 | self.seg_classes = {} 29 | 30 | # parse category file. 31 | with open(os.path.join(self.root, 'synsetoffset2category.txt'), 'r') as f: 32 | for line in f: 33 | ls = line.strip().split() 34 | self.cat2id[ls[0]] = ls[1] 35 | 36 | # parse segment num file. 37 | with open('misc/num_seg_classes.txt', 'r') as f: 38 | for line in f: 39 | ls = line.strip().split() 40 | self.seg_classes[ls[0]] = int(ls[1]) 41 | 42 | # if a subset of classes is specified. 43 | if class_choice is not None: 44 | self.cat2id = {k: v for k, v in self.cat2id.items() if k in class_choice} 45 | self.id2cat = {v: k for k, v in self.cat2id.items()} 46 | 47 | self.datapath = [] 48 | splitfile = os.path.join(self.root, 'train_test_split', 'shuffled_{}_file_list.json'.format(split)) 49 | filelist = json.load(open(splitfile, 'r')) 50 | for file in filelist: 51 | _, category, uuid = file.split('/') 52 | if category in self.cat2id.values(): 53 | self.datapath.append([ 54 | self.id2cat[category], 55 | os.path.join(self.root, category, 'points', uuid + '.pts'), 56 | os.path.join(self.root, category, 'points_label', uuid + '.seg') 57 | ]) 58 | # if split == 'train': 59 | # self.datapath = self.datapath[0:5000] 60 | # else: 61 | # self.datapath = self.datapath[0:500] 62 | self.classes = dict(zip(sorted(self.cat2id), range(len(self.cat2id)))) 63 | # print("classes:", self.classes) 64 | 65 | def __getitem__(self, index): 66 | fn = self.datapath[index] 67 | cls = self.classes[self.datapath[index][0]] 68 | point_set = np.loadtxt(fn[1]).astype(np.float32) 69 | # print("Origin Point size:", len(point_set)) 70 | seg = np.loadtxt(fn[2]).astype(np.int32) 71 | 72 | point_set, seg = grid_subsampling(point_set, labels=seg, sampleDl=self.first_subsampling_dl) 73 | 74 | # Center and rescale point for 1m radius 75 | pmin = np.min(point_set, axis=0) 76 | pmax = np.max(point_set, axis=0) 77 | point_set -= (pmin + pmax) / 2 78 | scale = np.max(np.linalg.norm(point_set, axis=1)) 79 | point_set *= 1.0 / scale 80 | 81 | if self.data_augmentation and self.split == 'train': 82 | theta = np.random.uniform(0, np.pi * 2) 83 | rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]) 84 | # TODO: why only rotate the x and z axis?? 85 | point_set[:, [0, 2]] = point_set[:, [0, 2]].dot(rotation_matrix) # random rotation 86 | point_set += np.random.normal(0, 0.001, size=point_set.shape) # random jitter 87 | 88 | pcd = make_point_cloud(point_set) 89 | open3d.estimate_normals(pcd) 90 | normals = np.array(pcd.normals) 91 | 92 | if self.config.in_features_dim == 1: 93 | features = np.ones([point_set.shape[0], 1]) 94 | elif self.config.in_features_dim == 4: 95 | features = np.ones([point_set.shape[0], 1]) 96 | features = np.concatenate([features, point_set], axis=1) 97 | elif self.config.in_features_dim == 7: 98 | features = np.ones([point_set.shape[0], 1]) 99 | features = np.concatenate([features, point_set, normals], axis=1) 100 | 101 | if self.classification: 102 | # manually convert numpy array to Tensor. 103 | # cls = torch.from_numpy(cls) - 1 # change to 0-based labels 104 | # cls = torch.from_numpy(np.array([cls])) 105 | # dict_inputs = segmentation_inputs(point_set, features, cls, self.config) 106 | # return dict_inputs 107 | return point_set, features, cls 108 | else: 109 | # manually convert numpy array to Tensor. 110 | # seg = torch.from_numpy(seg) - 1 # change to 0-based labels 111 | # dict_inputs = segmentation_inputs(point_set, features, seg, self.config) 112 | # return dict_inputs 113 | seg = seg - 1 114 | return point_set, features, seg 115 | 116 | def __len__(self): 117 | return len(self.datapath) 118 | 119 | 120 | if __name__ == '__main__': 121 | datapath = "./data/shapenetcore_partanno_segmentation_benchmark_v0" 122 | from training_ShapeNetCls import ShapeNetPartConfig 123 | 124 | config = ShapeNetPartConfig() 125 | 126 | print("Segmentation task:") 127 | dset = ShapeNetDataset(root=datapath, config=config, first_subsampling_dl=0.01, classification=False) 128 | input = dset[0] 129 | 130 | from datasets.dataloader import get_dataloader 131 | 132 | dataloader = get_dataloader(dset, batch_size=2) 133 | for iter, input in enumerate(dataloader): 134 | print(input) 135 | break 136 | -------------------------------------------------------------------------------- /datasets/common.py: -------------------------------------------------------------------------------- 1 | import open3d 2 | import numpy as np 3 | from datasets.dataloader import find_neighbors, grid_subsampling 4 | import torch 5 | 6 | 7 | def segmentation_inputs(points, features, labels, config): 8 | # TODO: originally I use this function to prepare the segmentation inputs for one single point cloud inputs and the next function collate_fn 9 | # to aggregate the inputs to a batch. 10 | 11 | # Starting radius of convolutions 12 | r_normal = config.first_subsampling_dl * config.KP_extent * 2.5 13 | 14 | # Starting layer 15 | layer_blocks = [] 16 | 17 | # Lists of inputs 18 | input_points = [] 19 | input_neighbors = [] 20 | input_pools = [] 21 | input_upsamples = [] 22 | input_batches_len = [] 23 | for block_i, block in enumerate(config.architecture): 24 | 25 | # TODO: stop early ? 26 | # Stop when meeting a global pooling or upsampling 27 | if 'upsample' in block: 28 | # if 'global' in block or 'upsample' in block: 29 | break 30 | 31 | # Get all blocks of the layer 32 | if not ('pool' in block or 'strided' in block): 33 | layer_blocks += [block] 34 | if block_i < len(config.architecture) - 1 and not ('upsample' in config.architecture[block_i + 1]): 35 | continue 36 | 37 | # Convolution neighbors indices 38 | # ***************************** 39 | 40 | if layer_blocks: 41 | # Convolutions are done in this layer, compute the neighbors with the good radius 42 | if np.any(['deformable' in blck for blck in layer_blocks[:-1]]): 43 | r = r_normal * config.density_parameter / (config.KP_extent * 2.5) 44 | else: 45 | r = r_normal 46 | conv_i = find_neighbors(points, points, r, max_neighbor=40) 47 | else: 48 | # This layer only perform pooling, no neighbors required 49 | conv_i = torch.zeros((0, 1), dtype=torch.int32) 50 | 51 | # Pooling neighbors indices 52 | # ************************* 53 | # If end of layer is a pooling operation 54 | if 'pool' in block or 'strided' in block: 55 | 56 | # New subsampling length 57 | dl = 2 * r_normal / (config.KP_extent * 2.5) 58 | 59 | # Subsampled points 60 | pool_p = grid_subsampling(points, sampleDl=dl) 61 | pool_b = torch.from_numpy(np.array([len(pool_p)])) 62 | 63 | # Radius of pooled neighbors 64 | if 'deformable' in block: 65 | r = r_normal * config.density_parameter / (config.KP_extent * 2.5) 66 | else: 67 | r = r_normal 68 | 69 | # Subsample indices 70 | pool_i = find_neighbors(pool_p, points, r, max_neighbor=40) 71 | 72 | # Upsample indices (with the radius of the next layer to keep wanted density) 73 | up_i = find_neighbors(points, pool_p, 2 * r, max_neighbor=40) 74 | 75 | else: 76 | # No pooling in the end of this layer, no pooling indices required 77 | pool_i = torch.zeros((0, 1), dtype=torch.int32) 78 | pool_p = torch.zeros((0, 3), dtype=torch.float32) 79 | pool_b = torch.zeros((0,), dtype=torch.int32) 80 | up_i = torch.zeros((0, 1), dtype=torch.int32) 81 | 82 | # TODO: Instead of eliminating the furthest point here, I select as most "max_neighbor" neighbors in find_neighbors function 83 | # Reduce size of neighbors matrices by eliminating furthest point 84 | 85 | # Updating input lists 86 | input_points += [torch.from_numpy(points)] 87 | input_neighbors += [conv_i] 88 | input_pools += [pool_i] 89 | input_upsamples += [up_i] 90 | # input_batches_len += [stacked_lengths] 91 | 92 | # New points for next layer 93 | points = pool_p 94 | # stacked_lengths = pool_b 95 | 96 | # Update radius and reset blocks 97 | r_normal *= 2 98 | layer_blocks = [] 99 | 100 | ############### 101 | # Return inputs 102 | ############### 103 | 104 | dict_inputs = { 105 | 'points': input_points, 106 | 'neighbors': input_neighbors, 107 | 'pools': input_pools, 108 | 'upsamples': input_upsamples, 109 | 'features': features, 110 | 'labels': labels 111 | } 112 | 113 | return dict_inputs 114 | 115 | 116 | def collate_fn(list_data): 117 | # TODO: the problem is, for 'neighbors', 'pools' and 'upsamples', there are dustbin indices for points with less 118 | # neighboring points than max_neighbors, but when we simply do "list_data[i_pcd]['points'][i_layer] + start_ind" 119 | # the dustbin index will be the first point in the next point cloud. 120 | batched_input = { 121 | 'points': [], 122 | 'neighbors': [], 123 | 'pools': [], 124 | 'upsamples': [], 125 | 'features': [], 126 | 'labels': [], 127 | 'batches_len': [], 128 | } 129 | num_layers = len(list_data[0]['points']) 130 | num_pcd = len(list_data) 131 | # process inputs['points'] 132 | for i_layer in range(num_layers): 133 | layer_points_list = [] 134 | start_ind = 0 135 | for i_pcd in range(num_pcd): 136 | layer_points_list.append(list_data[i_pcd]['points'][i_layer] + start_ind) 137 | start_ind += len(list_data[i_pcd]['points'][i_layer]) 138 | layer_points = torch.cat(layer_points_list, dim=0) 139 | batched_input['points'].append(layer_points) 140 | 141 | # process inputs['neighbors'] 142 | for i_layer in range(num_layers): 143 | layer_neighbors_list = [] 144 | start_ind = 0 145 | for i_pcd in range(num_pcd): 146 | layer_neighbors_list.append(list_data[i_pcd]['neighbors'][i_layer] + start_ind) 147 | start_ind += len(list_data[i_pcd]['neighbors'][i_layer]) 148 | layer_neighbors = torch.cat(layer_neighbors_list, dim=0) 149 | batched_input['neighbors'].append(layer_neighbors) 150 | 151 | return batched_input 152 | -------------------------------------------------------------------------------- /datasets/dataloader.py: -------------------------------------------------------------------------------- 1 | import open3d 2 | import numpy as np 3 | from utils.pointcloud import make_point_cloud 4 | from functools import partial 5 | import torch 6 | import cpp_wrappers.cpp_subsampling.grid_subsampling as cpp_subsampling 7 | from utils.timer import Timer 8 | import batch_find_neighbors 9 | 10 | 11 | def find_neighbors(query_points, support_points, radius, max_neighbor): 12 | pcd = make_point_cloud(support_points) 13 | kdtree = open3d.KDTreeFlann(pcd) 14 | neighbor_indices_list = [] 15 | for i, point in enumerate(query_points): 16 | [k, idx, dis] = kdtree.search_radius_vector_3d(point, radius) 17 | if k > max_neighbor: 18 | idx = np.random.choice(idx, max_neighbor, replace=False) 19 | else: 20 | # if not enough neighbor points, then add the dustbin point. 21 | idx = list(idx) + [len(support_points)] * (max_neighbor - k) 22 | neighbor_indices_list.append([idx]) 23 | neighbors = np.concatenate(neighbor_indices_list, axis=0) 24 | return torch.from_numpy(neighbors) 25 | 26 | def batch_find_neighbors_wrapper(query_points, support_points, query_batches, support_batches, radius, max_neighbors): 27 | if True: 28 | cpp = batch_find_neighbors_cpp(query_points, support_points, query_batches, support_batches, radius, max_neighbors) 29 | cpp = cpp.reshape([query_points.shape[0], -1]) 30 | cpp = cpp[:, :max_neighbors] 31 | return cpp 32 | else: 33 | py = batch_find_neighbors_py(query_points, support_points, query_batches, support_batches, radius, max_neighbors) 34 | py = py[:, :max_neighbors] 35 | return py 36 | 37 | def batch_find_neighbors_cpp(query_points, support_points, query_batches, support_batches, radius, max_neighbors): 38 | outputs = batch_find_neighbors.compute(query_points, support_points, query_batches, support_batches, radius) 39 | outputs = outputs.long() 40 | return outputs 41 | 42 | def batch_find_neighbors_py(query_points, support_points, query_batches, support_batches, radius, max_neighbors): 43 | num_batches = len(support_batches) 44 | # Create kdtree for each pcd in support_points 45 | kdtrees = [] 46 | start_ind = 0 47 | for length in support_batches: 48 | pcd = make_point_cloud(support_points[start_ind:start_ind + length]) 49 | kdtrees.append(open3d.KDTreeFlann(pcd)) 50 | start_ind += length 51 | assert len(kdtrees) == num_batches 52 | # Search neigbors indices 53 | neighbors_indices_list = [] 54 | start_ind = 0 55 | support_start_ind = 0 56 | dustbin_ind = len(support_points) 57 | for i_batch, length in enumerate(query_batches): 58 | for i_pts, pts in enumerate(query_points[start_ind:start_ind + length]): 59 | [k, idx, dis] = kdtrees[i_batch].search_radius_vector_3d(pts, radius) 60 | if k > max_neighbors: 61 | # idx = np.random.choice(idx, max_neighbors, replace=False) 62 | # If i use random, the closest_pool will not work as expected. 63 | idx = list(idx[0:max_neighbors]) 64 | else: 65 | # if not enough neighbor points, then add dustbin index. Careful !!! 66 | idx = list(idx) + [dustbin_ind - support_start_ind] * (max_neighbors - k) 67 | idx = np.array(idx) + support_start_ind 68 | neighbors_indices_list.append(idx) 69 | # finish one query_points, update the start_ind 70 | start_ind += int(query_batches[i_batch]) 71 | support_start_ind += int(support_batches[i_batch]) 72 | return torch.from_numpy(np.array(neighbors_indices_list)).long() 73 | 74 | def grid_subsampling(points, features=None, labels=None, sampleDl=0.1, verbose=0): 75 | """ 76 | CPP wrapper for a grid subsampling (method = barycenter for points and features 77 | :param points: (N, 3) matrix of input points 78 | :param features: optional (N, d) matrix of features (floating number) 79 | :param labels: optional (N,) matrix of integer labels 80 | :param sampleDl: parameter defining the size of grid voxels 81 | :param verbose: 1 to display 82 | :return: subsampled points, with features and/or labels depending of the input 83 | """ 84 | 85 | if (features is None) and (labels is None): 86 | return cpp_subsampling.compute(points, sampleDl=sampleDl, verbose=verbose) 87 | elif (labels is None): 88 | return cpp_subsampling.compute(points, features=features, sampleDl=sampleDl, verbose=verbose) 89 | elif (features is None): 90 | return cpp_subsampling.compute(points, classes=labels, sampleDl=sampleDl, verbose=verbose) 91 | else: 92 | return cpp_subsampling.compute(points, features=features, classes=labels, sampleDl=sampleDl, verbose=verbose) 93 | 94 | 95 | def batch_grid_subsampling(points, batches_len, sampleDl=0.1): 96 | """ 97 | CPP wrapper for a batch grid subsampling (method = barycenter for points and features 98 | :param points: (N, 3) matrix of input points 99 | :param batches_len: lengths of batched input points 100 | :param sampleDl: parameter defining the size of grid voxels 101 | :return: 102 | """ 103 | subsampled_points_list = [] 104 | subsampled_batches_len_list = [] 105 | start_ind = 0 106 | for length in batches_len: 107 | b_origin_points = points[start_ind:start_ind + length] 108 | b_subsampled_points = grid_subsampling(b_origin_points, sampleDl=sampleDl) 109 | start_ind += length 110 | subsampled_points_list.append(b_subsampled_points) 111 | subsampled_batches_len_list.append(len(b_subsampled_points)) 112 | subsampled_points = torch.from_numpy(np.concatenate(subsampled_points_list, axis=0)) 113 | subsampled_batches_len = torch.from_numpy(np.array(subsampled_batches_len_list)).int() 114 | return subsampled_points, subsampled_batches_len 115 | 116 | 117 | def collate_fn_segmentation(list_data, config, neighborhood_limits): 118 | batched_points_list = [] 119 | batched_features_list = [] 120 | batched_labels_list = [] 121 | batched_lengths_list = [] 122 | 123 | for ind, (pts, features, labels) in enumerate(list_data): 124 | batched_points_list.append(pts) 125 | batched_features_list.append(features) 126 | batched_labels_list.append(labels) 127 | batched_lengths_list.append(len(pts)) 128 | 129 | batched_features = torch.from_numpy(np.concatenate(batched_features_list, axis=0)) 130 | batched_labels = torch.from_numpy(np.concatenate(batched_labels_list, axis=0)) 131 | batched_points = torch.from_numpy(np.concatenate(batched_points_list, axis=0)) 132 | batched_lengths = torch.from_numpy(np.array(batched_lengths_list)).int() 133 | 134 | # Starting radius of convolutions 135 | r_normal = config.first_subsampling_dl * config.KP_extent * 2.5 136 | 137 | # Starting layer 138 | layer_blocks = [] 139 | layer = 0 140 | 141 | # Lists of inputs 142 | input_points = [] 143 | input_neighbors = [] 144 | input_pools = [] 145 | input_upsamples = [] 146 | input_batches_len = [] 147 | 148 | for block_i, block in enumerate(config.architecture): 149 | 150 | # Stop when meeting a global pooling or upsampling 151 | if 'global' in block or 'upsample' in block: 152 | break 153 | 154 | # Get all blocks of the layer 155 | if not ('pool' in block or 'strided' in block): 156 | layer_blocks += [block] 157 | if block_i < len(config.architecture) - 1 and not ('upsample' in config.architecture[block_i + 1]): 158 | continue 159 | 160 | # Convolution neighbors indices 161 | # ***************************** 162 | 163 | if layer_blocks: 164 | # Convolutions are done in this layer, compute the neighbors with the good radius 165 | if np.any(['deformable' in blck for blck in layer_blocks[:-1]]): 166 | r = r_normal * config.density_parameter / (config.KP_extent * 2.5) 167 | else: 168 | r = r_normal 169 | conv_i = batch_find_neighbors_wrapper(batched_points, batched_points, batched_lengths, batched_lengths, r, neighborhood_limits[layer]) 170 | 171 | else: 172 | # This layer only perform pooling, no neighbors required 173 | conv_i = torch.zeros((0, 1), dtype=torch.int64) 174 | 175 | # Pooling neighbors indices 176 | # ************************* 177 | 178 | # If end of layer is a pooling operation 179 | if 'pool' in block or 'strided' in block: 180 | 181 | # New subsampling length 182 | dl = 2 * r_normal / (config.KP_extent * 2.5) 183 | 184 | # Subsampled points 185 | pool_p, pool_b = batch_grid_subsampling(batched_points, batched_lengths, sampleDl=dl) 186 | 187 | # Radius of pooled neighbors 188 | if 'deformable' in block: 189 | r = r_normal * config.density_parameter / (config.KP_extent * 2.5) 190 | else: 191 | r = r_normal 192 | 193 | # Subsample indices 194 | pool_i = batch_find_neighbors_wrapper(pool_p, batched_points, pool_b, batched_lengths, r, neighborhood_limits[layer]) 195 | 196 | # Upsample indices (with the radius of the next layer to keep wanted density) 197 | up_i = batch_find_neighbors_wrapper(batched_points, pool_p, batched_lengths, pool_b, 2 * r, neighborhood_limits[layer]) 198 | 199 | else: 200 | # No pooling in the end of this layer, no pooling indices required 201 | pool_i = torch.zeros((0, 1), dtype=torch.int64) 202 | pool_p = torch.zeros((0, 3), dtype=torch.float32) 203 | pool_b = torch.zeros((0,), dtype=torch.int64) 204 | up_i = torch.zeros((0, 1), dtype=torch.int64) 205 | 206 | # Updating input lists 207 | input_points += [batched_points.float()] 208 | input_neighbors += [conv_i.long()] 209 | input_pools += [pool_i.long()] 210 | input_upsamples += [up_i.long()] 211 | input_batches_len += [batched_lengths] 212 | 213 | # New points for next layer 214 | batched_points = pool_p 215 | batched_lengths = pool_b 216 | 217 | # Update radius and reset blocks 218 | r_normal *= 2 219 | layer += 1 220 | layer_blocks = [] 221 | 222 | ############### 223 | # Return inputs 224 | ############### 225 | dict_inputs = { 226 | 'points': input_points, 227 | 'neighbors': input_neighbors, 228 | 'pools': input_pools, 229 | 'upsamples': input_upsamples, 230 | 'features': batched_features.float(), 231 | 'labels': batched_labels.long(), 232 | 'stack_lengths': input_batches_len 233 | } 234 | 235 | return dict_inputs 236 | 237 | 238 | def collate_fn_classification(list_data, config, neighborhood_limits): 239 | # The difference with 'collate_fn_classification' is no need to save upsample indices && stop when meeting a global pooling (not including upsample) 240 | batched_points_list = [] 241 | batched_features_list = [] 242 | batched_labels_list = [] 243 | batched_lengths_list = [] 244 | 245 | for ind, (pts, features, labels) in enumerate(list_data): 246 | batched_points_list.append(pts) 247 | batched_features_list.append(features) 248 | batched_labels_list.append(labels) 249 | batched_lengths_list.append(len(pts)) 250 | 251 | batched_features = torch.from_numpy(np.concatenate(batched_features_list, axis=0)) 252 | batched_labels = torch.from_numpy(np.array(batched_labels_list)) 253 | batched_points = torch.from_numpy(np.concatenate(batched_points_list, axis=0)) 254 | batched_lengths = torch.from_numpy(np.array(batched_lengths_list)).int() 255 | 256 | # Starting radius of convolutions 257 | r_normal = config.first_subsampling_dl * config.KP_extent * 2.5 258 | 259 | # Starting layer 260 | layer_blocks = [] 261 | layer = 0 262 | 263 | # Lists of inputs 264 | input_points = [] 265 | input_neighbors = [] 266 | input_pools = [] 267 | input_batches_len = [] 268 | 269 | for block_i, block in enumerate(config.architecture): 270 | 271 | # Stop when meeting a global pooling (not include upsampling) 272 | # TODO: this condition is different with tensorflow implementation. 273 | if 'global' in block and layer_blocks == []: 274 | break 275 | 276 | # Get all blocks of the layer 277 | if not ('pool' in block or 'strided' in block): 278 | layer_blocks += [block] 279 | if block_i < len(config.architecture) - 1 and not ('upsample' in config.architecture[block_i + 1]): 280 | continue 281 | 282 | # Convolution neighbors indices 283 | # ***************************** 284 | 285 | if layer_blocks: 286 | # Convolutions are done in this layer, compute the neighbors with the good radius 287 | if np.any(['deformable' in blck for blck in layer_blocks[:-1]]): 288 | r = r_normal * config.density_parameter / (config.KP_extent * 2.5) 289 | else: 290 | r = r_normal 291 | conv_i = batch_find_neighbors_wrapper(batched_points, batched_points, batched_lengths, batched_lengths, r, neighborhood_limits[layer]) 292 | 293 | else: 294 | # This layer only perform pooling, no neighbors required 295 | conv_i = torch.zeros((0, 1), dtype=torch.int64) 296 | 297 | # Pooling neighbors indices 298 | # ************************* 299 | 300 | # If end of layer is a pooling operation 301 | if 'pool' in block or 'strided' in block: 302 | 303 | # New subsampling length 304 | dl = 2 * r_normal / (config.KP_extent * 2.5) 305 | 306 | # Subsampled points 307 | pool_p, pool_b = batch_grid_subsampling(batched_points, batched_lengths, sampleDl=dl) 308 | 309 | # Radius of pooled neighbors 310 | if 'deformable' in block: 311 | r = r_normal * config.density_parameter / (config.KP_extent * 2.5) 312 | else: 313 | r = r_normal 314 | 315 | # Subsample indices 316 | pool_i = batch_find_neighbors_wrapper(pool_p, batched_points, pool_b, batched_lengths, r, neighborhood_limits[layer]) 317 | 318 | 319 | else: 320 | # No pooling in the end of this layer, no pooling indices required 321 | pool_i = torch.zeros((0, 1), dtype=torch.int64) 322 | pool_p = torch.zeros((0, 3), dtype=torch.float64) 323 | pool_b = torch.zeros((0,), dtype=torch.int64) 324 | 325 | # Updating input lists 326 | input_points += [batched_points.float()] 327 | input_neighbors += [conv_i.long()] 328 | input_pools += [pool_i.long()] 329 | input_batches_len += [batched_lengths] 330 | 331 | # New points for next layer 332 | batched_points = pool_p 333 | batched_lengths = pool_b 334 | 335 | # Update radius and reset blocks 336 | r_normal *= 2 337 | layer += 1 338 | layer_blocks = [] 339 | 340 | ############### 341 | # Return inputs 342 | ############### 343 | dict_inputs = { 344 | 'points': input_points, 345 | 'neighbors': input_neighbors, 346 | 'pools': input_pools, 347 | 'features': batched_features.float(), 348 | 'labels': batched_labels.long(), 349 | 'stack_lengths': input_batches_len 350 | } 351 | 352 | return dict_inputs 353 | 354 | 355 | def calibrate_neighbors(dataset, config, collate_fn, keep_ratio=0.8, samples_threshold=5000): 356 | timer = Timer() 357 | last_display = timer.total_time 358 | 359 | # From config parameter, compute higher bound of neighbors number in a neighborhood 360 | hist_n = int(np.ceil(4 / 3 * np.pi * (config.density_parameter + 1) ** 3)) 361 | neighb_hists = np.zeros((config.num_layers, hist_n), dtype=np.int32) 362 | 363 | # Get histogram of neighborhood sizes i in 1 epoch max. 364 | for i in range(len(dataset)): 365 | timer.tic() 366 | batched_input = collate_fn([dataset[i]], config, neighborhood_limits=[hist_n] * 5) 367 | 368 | # update histogram 369 | counts = [torch.sum(neighb_mat < neighb_mat.shape[0], dim=1).numpy() for neighb_mat in batched_input['neighbors']] 370 | hists = [np.bincount(c, minlength=hist_n)[:hist_n] for c in counts] 371 | neighb_hists += np.vstack(hists) 372 | timer.toc() 373 | 374 | if timer.total_time - last_display > 0.1: 375 | last_display = timer.total_time 376 | print(f"Calib Neighbors {i:08d}: timings {timer.total_time:4.2f}s") 377 | 378 | if np.min(np.sum(neighb_hists, axis=1)) > samples_threshold: 379 | break 380 | 381 | cumsum = np.cumsum(neighb_hists.T, axis=0) 382 | percentiles = np.sum(cumsum < (keep_ratio * cumsum[hist_n - 1, :]), axis=0) 383 | 384 | neighborhood_limits = percentiles 385 | print('\n') 386 | 387 | return neighborhood_limits 388 | 389 | 390 | def get_dataloader(dataset, batch_size=2, num_workers=4, shuffle=True): 391 | if dataset.classification is False: 392 | neighborhood_limits = calibrate_neighbors(dataset, dataset.config, collate_fn=collate_fn_segmentation) 393 | print("neighborhood:", neighborhood_limits) 394 | dataloader = torch.utils.data.DataLoader( 395 | dataset, 396 | batch_size=batch_size, 397 | shuffle=shuffle, 398 | num_workers=num_workers, 399 | # https://discuss.pytorch.org/t/supplying-arguments-to-collate-fn/25754/4 400 | collate_fn=partial(collate_fn_segmentation, config=dataset.config, neighborhood_limits=neighborhood_limits), 401 | drop_last=True 402 | ) 403 | else: 404 | neighborhood_limits = calibrate_neighbors(dataset, dataset.config, collate_fn=collate_fn_classification) 405 | print("neighborhood:", neighborhood_limits) 406 | dataloader = torch.utils.data.DataLoader( 407 | dataset, 408 | batch_size=batch_size, 409 | shuffle=shuffle, 410 | num_workers=num_workers, 411 | collate_fn=partial(collate_fn_classification, config=dataset.config, neighborhood_limits=neighborhood_limits), 412 | drop_last=True 413 | ) 414 | 415 | return dataloader 416 | 417 | 418 | if __name__ == '__main__': 419 | # from training_ShapeNetCls import ShapeNetPartConfig 420 | from training_ShapeNetPart import ShapeNetPartConfig 421 | from datasets.ShapeNet import ShapeNetDataset 422 | 423 | config = ShapeNetPartConfig() 424 | datapath = "./data/shapenetcore_partanno_segmentation_benchmark_v0" 425 | dset = ShapeNetDataset(root=datapath, config=config, first_subsampling_dl=0.01, classification=False) 426 | dataloader = get_dataloader(dset, batch_size=2, num_workers=1) 427 | for iter, inputs in enumerate(dataloader): 428 | print(iter) 429 | print(inputs['labels']) 430 | break 431 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: kpconv_pytorch 2 | channels: 3 | - pytorch 4 | - open3d-admin 5 | - conda-forge 6 | - anaconda 7 | - defaults 8 | dependencies: 9 | - _libgcc_mutex=0.1=main 10 | - attrs=19.3.0=py_0 11 | - backcall=0.1.0=py36_0 12 | - blas=1.0=mkl 13 | - bleach=3.1.0=py36_0 14 | - ca-certificates=2019.11.27=0 15 | - certifi=2019.11.28=py36_0 16 | - cffi=1.13.2=py36h2e261b9_0 17 | - cudatoolkit=10.0.130=0 18 | - cycler=0.10.0=py_2 19 | - dbus=1.13.6=he372182_0 20 | - decorator=4.4.1=py_0 21 | - defusedxml=0.6.0=py_0 22 | - entrypoints=0.3=py36_0 23 | - expat=2.2.6=he6710b0_0 24 | - fontconfig=2.13.1=he4413a7_1000 25 | - freetype=2.9.1=h8a8886c_1 26 | - gettext=0.19.8.1=hc5be6a0_1002 27 | - glib=2.58.3=py36h6f030ca_1002 28 | - gmp=6.1.2=h6c8ec71_1 29 | - gst-plugins-base=1.14.5=h0935bb2_0 30 | - gstreamer=1.14.5=h36ae1b5_0 31 | - icu=58.2=hf484d3e_1000 32 | - importlib_metadata=1.2.0=py36_0 33 | - intel-openmp=2019.4=243 34 | - ipykernel=5.1.3=py36h39e3cac_0 35 | - ipython=7.10.1=py36h39e3cac_0 36 | - ipython_genutils=0.2.0=py36_0 37 | - ipywidgets=7.5.1=py_0 38 | - jedi=0.15.1=py36_0 39 | - jinja2=2.10.3=py_0 40 | - joblib=0.14.0=py_0 41 | - jpeg=9b=h024ee3a_2 42 | - jsonschema=3.2.0=py36_0 43 | - jupyter_client=5.3.4=py36_0 44 | - jupyter_core=4.6.1=py36_0 45 | - kiwisolver=1.1.0=py36hc9558a2_0 46 | - libedit=3.1.20181209=hc058e9b_0 47 | - libffi=3.2.1=hd88cf55_4 48 | - libgcc-ng=9.1.0=hdf63c60_0 49 | - libgfortran-ng=7.3.0=hdf63c60_0 50 | - libiconv=1.15=h516909a_1005 51 | - libpng=1.6.37=hbc83047_0 52 | - libprotobuf=3.11.1=h8b12597_0 53 | - libsodium=1.0.16=h1bed415_0 54 | - libstdcxx-ng=9.1.0=hdf63c60_0 55 | - libtiff=4.1.0=h2733197_0 56 | - libuuid=2.32.1=h14c3975_1000 57 | - libxcb=1.13=h14c3975_1002 58 | - libxml2=2.9.9=hea5a465_1 59 | - markupsafe=1.1.1=py36h7b6447c_0 60 | - matplotlib=3.1.1=py36h5429711_0 61 | - mistune=0.8.4=py36h7b6447c_0 62 | - mkl=2019.4=243 63 | - mkl-service=2.3.0=py36he904b0f_0 64 | - mkl_fft=1.0.15=py36ha843d7b_0 65 | - mkl_random=1.1.0=py36hd6b4f25_0 66 | - more-itertools=7.2.0=py36_0 67 | - nbconvert=5.6.1=py36_0 68 | - nbformat=4.4.0=py36_0 69 | - ncurses=6.1=he6710b0_1 70 | - ninja=1.9.0=py36hfd86e86_0 71 | - notebook=6.0.2=py36_0 72 | - numpy=1.17.4=py36hc1035e2_0 73 | - numpy-base=1.17.4=py36hde5b4d6_0 74 | - olefile=0.46=py36_0 75 | - open3d=0.7.0.0=py36_0 76 | - openssl=1.1.1=h7b6447c_0 77 | - pandoc=2.2.3.2=0 78 | - pandocfilters=1.4.2=py36_1 79 | - parso=0.5.1=py_0 80 | - pcre=8.43=he1b5a44_0 81 | - pexpect=4.7.0=py36_0 82 | - pickleshare=0.7.5=py36_0 83 | - pillow=6.2.1=py36h34e0f95_0 84 | - pip=19.3.1=py36_0 85 | - prometheus_client=0.7.1=py_0 86 | - prompt_toolkit=3.0.2=py_0 87 | - protobuf=3.11.1=py36he1b5a44_0 88 | - pthread-stubs=0.4=h14c3975_1001 89 | - ptyprocess=0.6.0=py36_0 90 | - pycparser=2.19=py36_0 91 | - pygments=2.5.2=py_0 92 | - pyparsing=2.4.5=py_0 93 | - pyqt=5.9.2=py36hcca6a23_4 94 | - pyrsistent=0.15.6=py36h7b6447c_0 95 | - python=3.6.8=h0371630_0 96 | - python-dateutil=2.8.1=py_0 97 | - pytorch=1.1.0=py3.6_cuda10.0.130_cudnn7.5.1_0 98 | - pytz=2019.3=py_0 99 | - pyzmq=18.1.0=py36he6710b0_0 100 | - qt=5.9.7=h5867ecd_1 101 | - readline=7.0=h7b6447c_5 102 | - scikit-learn=0.21.3=py36hd81dba3_0 103 | - scipy=1.3.2=py36h7c811a0_0 104 | - send2trash=1.5.0=py36_0 105 | - setuptools=42.0.2=py36_0 106 | - sip=4.19.8=py36hf484d3e_0 107 | - six=1.13.0=py36_0 108 | - sqlite=3.30.1=h7b6447c_0 109 | - tensorboardx=1.9=py_0 110 | - terminado=0.8.3=py36_0 111 | - testpath=0.4.4=py_0 112 | - tk=8.6.8=hbc83047_0 113 | - torchvision=0.3.0=py36_cu10.0.130_1 114 | - tornado=6.0.3=py36h7b6447c_0 115 | - traitlets=4.3.3=py36_0 116 | - wcwidth=0.1.7=py36_0 117 | - webencodings=0.5.1=py36_1 118 | - wheel=0.33.6=py36_0 119 | - widgetsnbextension=3.5.1=py36_0 120 | - xorg-libxau=1.0.9=h14c3975_0 121 | - xorg-libxdmcp=1.1.3=h516909a_0 122 | - xz=5.2.4=h14c3975_4 123 | - zeromq=4.3.1=he6710b0_3 124 | - zipp=0.6.0=py_0 125 | - zlib=1.2.11=h7b6447c_3 126 | - zstd=1.3.7=h0b5b093_0 127 | prefix: /home/xybai/anaconda3/envs/kpconv_pytorch 128 | 129 | -------------------------------------------------------------------------------- /kernels/convolution_ops.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from kernels.kernel_points import load_kernels as create_kernel_points 3 | import torch 4 | 5 | 6 | def unary_convolution(features, 7 | K_values): 8 | """ 9 | Simple unary convolution in tensorflow. Equivalent to matrix multiplication (space projection) for each features 10 | :param features: float32[n_points, in_fdim] - input features 11 | :param K_values: float32[in_fdim, out_fdim] - weights of the kernel 12 | :return: output_features float32[n_points, out_fdim] 13 | """ 14 | 15 | return torch.matmul(features, K_values) 16 | 17 | 18 | def radius_gaussian(sq_r, sig, eps=1e-9): 19 | """ 20 | Compute a radius gaussian (gaussian of distance) 21 | :param sq_r: input radiuses [dn, ..., d1, d0] 22 | :param sig: extents of gaussians [d1, d0] or [d0] or float 23 | :return: gaussian of sq_r [dn, ..., d1, d0] 24 | """ 25 | return torch.exp(-sq_r / (2 * torch.pow(sig, 2) + eps)) 26 | 27 | 28 | def KPConv(query_points, 29 | support_points, 30 | neighbors_indices, 31 | features, 32 | K_values, 33 | fixed='center', 34 | KP_extent=1.0, 35 | KP_influence='linear', 36 | aggregation_mode='sum'): 37 | """ 38 | This function initiates the kernel point disposition before building KPConv graph ops 39 | :param query_points: float32[n_points, dim] - input query points (center of neighborhoods) 40 | :param support_points: float32[n0_points, dim] - input support points (from which neighbors are taken) 41 | :param neighbors_indices: int32[n_points, n_neighbors] - indices of neighbors of each point 42 | :param features: float32[n_points, in_fdim] - input features 43 | :param K_values: float32[n_kpoints, in_fdim, out_fdim] - weights of the kernel 44 | :param fixed: string in ('none', 'center' or 'verticals') - fix position of certain kernel points 45 | :param KP_extent: float32 - influence radius of each kernel point 46 | :param KP_influence: string in ('constant', 'linear', 'gaussian') - influence function of the kernel points 47 | :param aggregation_mode: string in ('closest', 'sum') - whether to sum influences, or only keep the closest 48 | :return: output_features float32[n_points, out_fdim] 49 | """ 50 | 51 | # Initial kernel extent for this layer 52 | K_radius = 1.5 * KP_extent 53 | 54 | # Number of kernel points 55 | num_kpoints = int(K_values.shape[0]) 56 | 57 | # Check point dimension (currently only 3D is supported) 58 | points_dim = int(query_points.shape[1]) 59 | 60 | # Create one kernel disposition (as numpy array). Choose the KP distance to center thanks to the KP extent 61 | K_points_numpy = create_kernel_points(K_radius, 62 | num_kpoints, 63 | num_kernels=1, 64 | dimension=points_dim, 65 | fixed=fixed) 66 | K_points_numpy = K_points_numpy.reshape((num_kpoints, points_dim)) 67 | 68 | # Create the tensorflow variable 69 | K_points = torch.from_numpy(K_points_numpy.astype(np.float32)) 70 | if K_values.is_cuda: 71 | K_points = K_points.to(K_values.device) 72 | 73 | return KPConv_ops(query_points, 74 | support_points, 75 | neighbors_indices, 76 | features, 77 | K_points, 78 | K_values, 79 | KP_extent, 80 | KP_influence, 81 | aggregation_mode) 82 | 83 | 84 | def KPConv_ops(query_points, 85 | support_points, 86 | neighbors_indices, 87 | features, 88 | K_points, 89 | K_values, 90 | KP_extent, 91 | KP_influence, 92 | aggregation_mode): 93 | """ 94 | This function creates a graph of operations to define Kernel Point Convolution in tensorflow. See KPConv function 95 | above for a description of each parameter 96 | 97 | :param query_points: [n_points, dim] 98 | :param support_points: [n0_points, dim] 99 | :param neighbors_indices: [n_points, n_neighbors] 100 | :param features: [n_points, in_fdim] 101 | :param K_points: [n_kpoints, dim] 102 | :param K_values: [n_kpoints, in_fdim, out_fdim] 103 | :param KP_extent: float32 104 | :param KP_influence: string 105 | :param aggregation_mode: string 106 | :return: [n_points, out_fdim] 107 | """ 108 | 109 | # Get variables 110 | n_kp = int(K_points.shape[0]) 111 | 112 | # Add a fake point in the last row for shadow neighbors 113 | shadow_point = torch.ones_like(support_points[:1, :]) * 1e6 114 | support_points = torch.cat([support_points, shadow_point], dim=0) 115 | 116 | # Get neighbor points [n_points, n_neighbors, dim] 117 | neighbors = support_points[neighbors_indices, :] 118 | 119 | # Center every neighborhood 120 | neighbors = neighbors - query_points.unsqueeze(1) 121 | 122 | # Get all difference matrices [n_points, n_neighbors, n_kpoints, dim] 123 | differences = neighbors.unsqueeze(2) - K_points 124 | 125 | # Get the square distances [n_points, n_neighbors, n_kpoints] 126 | sq_distances = torch.sum(torch.mul(differences, differences), dim=3) 127 | 128 | # Get Kernel point influences [n_points, n_kpoints, n_neighbors] 129 | if KP_influence == 'constant': 130 | # Every point get an influence of 1. 131 | all_weights = torch.ones_like(sq_distances) 132 | all_weights = all_weights.transpose(1, 2) 133 | 134 | elif KP_influence == 'linear': 135 | # Influence decrease linearly with the distance, and get to zero when d = KP_extent. 136 | corr = 1 - torch.sqrt(sq_distances + 1e-10) / KP_extent 137 | all_weights = torch.max(corr, torch.zeros_like(sq_distances)) 138 | all_weights = all_weights.transpose(1, 2) 139 | 140 | elif KP_influence == 'gaussian': 141 | # Influence in gaussian of the distance. 142 | sigma = KP_extent * 0.3 143 | all_weights = radius_gaussian(sq_distances, sigma) 144 | all_weights = all_weights.transpose(1, 2) 145 | else: 146 | raise ValueError('Unknown influence function type (config.KP_influence)') 147 | 148 | # In case of closest mode, only the closest KP can influence each point 149 | if aggregation_mode == 'closest': 150 | pass 151 | # neighbors_1nn = tf.argmin(sq_distances, axis=2, output_type=tf.int32) 152 | # all_weights *= tf.one_hot(neighbors_1nn, n_kp, axis=1, dtype=tf.float32) 153 | 154 | elif aggregation_mode != 'sum': 155 | raise ValueError("Unknown convolution mode. Should be 'closest' or 'sum'") 156 | 157 | features = torch.cat([features, torch.zeros_like(features[:1, :])], dim=0) 158 | 159 | # Get the features of each neighborhood [n_points, n_neighbors, in_fdim] 160 | neighborhood_features = features[neighbors_indices, :] 161 | 162 | # Apply distance weights [n_points, n_kpoints, in_fdim] 163 | weighted_features = torch.matmul(all_weights, neighborhood_features) 164 | 165 | # Apply network weights [n_kpoints, n_points, out_fdim] 166 | weighted_features = weighted_features.transpose(0, 1) 167 | kernel_outputs = torch.matmul(weighted_features, K_values) 168 | 169 | # Convolution sum to get [n_points, out_fdim] 170 | output_features = torch.sum(kernel_outputs, dim=0, keepdim=False) 171 | 172 | # normalization term. 173 | # neighbor_features_sum = torch.sum(neighborhood_features, dim=-1) 174 | # neighbor_num = torch.sum(torch.gt(neighbor_features_sum, 0.0), dim=-1) 175 | # neighbor_num = torch.max(neighbor_num, torch.ones_like(neighbor_num)) 176 | # output_features = output_features / neighbor_num.unsqueeze(1) 177 | 178 | return output_features 179 | 180 | 181 | def KPConv_deformable(query_points, 182 | support_points, 183 | neighbors_indices, 184 | features, 185 | K_values, 186 | w0, b0, 187 | fixed='center', 188 | KP_extent=1.0, 189 | KP_influence='linear', 190 | aggregation_mode='sum', 191 | modulated=False): 192 | """ 193 | This function initiates the kernel point disposition before building deformable KPConv graph ops 194 | 195 | :param query_points: float32[n_points, dim] - input query points (center of neighborhoods) 196 | :param support_points: float32[n0_points, dim] - input support points (from which neighbors are taken) 197 | :param neighbors_indices: int32[n_points, n_neighbors] - indices of neighbors of each point 198 | :param features: float32[n_points, in_fdim] - input features 199 | :param K_values: float32[n_kpoints, in_fdim, out_fdim] - weights of the kernel 200 | :param w0: float32[n_points, dim * n_kpoints] - weights of the rigid KPConv for offsets 201 | :param b0: float32[dim * n_kpoints] - bias of the rigid KPConv for offsets. 202 | :param fixed: string in ('none', 'center' or 'verticals') - fix position of certain kernel points 203 | :param KP_extent: float32 - influence radius of each kernel point 204 | :param KP_influence: string in ('constant', 'linear', 'gaussian') - influence function of the kernel points 205 | :param aggregation_mode: string in ('closest', 'sum') - behavior of the convolution 206 | :param modulated: bool - If deformable conv should be modulated 207 | 208 | :return: output_features float32[n_points, out_fdim] 209 | """ 210 | 211 | # Initial kernel extent for this layer 212 | K_radius = 1.5 * KP_extent 213 | 214 | # Number of kernel points 215 | num_kpoints = int(K_values.shape[0]) 216 | 217 | # Check point dimension (currently only 3D is supported) 218 | points_dim = int(query_points.shape[1]) 219 | 220 | # Create one kernel disposition (as numpy array). Choose the KP distance to center thanks to the KP extent 221 | K_points_numpy = create_kernel_points(K_radius, 222 | num_kpoints, 223 | num_kernels=1, 224 | dimension=points_dim, 225 | fixed=fixed) 226 | K_points_numpy = K_points_numpy.reshape((num_kpoints, points_dim)) 227 | 228 | # Create the pytorch variable 229 | K_points = torch.from_numpy(K_points_numpy.astype(np.float32)) 230 | if K_values.is_cuda: 231 | K_points = K_points.to(K_values.device) 232 | 233 | ############################# 234 | # Standard KPConv for offsets 235 | ############################# 236 | 237 | # Create independant weight for the first convolution and a bias term as no batch normalization happen 238 | # if modulated: 239 | # offset_dim = (points_dim + 1) * num_kpoints 240 | # else: 241 | # offset_dim = points_dim * num_kpoints 242 | # shape0 = list(K_values.shape) 243 | # shape0[-1] = offset_dim # [n_kpoints, in_fdim, offset_dim] 244 | # w0 = torch.zeros(shape0, dtype=torch.float32) # offset_conv_weights 245 | # b0 = torch.zeros(offset_dim, dtype=torch.float32) # offset_conv_bias 246 | 247 | # Get features from standard convolution 248 | features0 = KPConv_ops(query_points, 249 | support_points, 250 | neighbors_indices, 251 | features, 252 | K_points, 253 | w0, 254 | KP_extent, 255 | KP_influence, 256 | aggregation_mode) + b0 257 | if modulated: 258 | 259 | # Get offset (in normalized scale) from features 260 | offsets = features0[:, :points_dim * num_kpoints] 261 | offsets = offsets.reshape([-1, num_kpoints, points_dim]) 262 | 263 | # Get modulations 264 | modulations = 2 * torch.sigmoid(features0[:, points_dim * num_kpoints:]) 265 | 266 | else: 267 | 268 | # Get offset (in normalized scale) from features 269 | offsets = features0.reshape([-1, num_kpoints, points_dim]) 270 | 271 | # No modulations 272 | modulations = None 273 | 274 | # Rescale offset for this layer 275 | offsets *= KP_extent 276 | ############################### 277 | # Build deformable KPConv graph 278 | ############################### 279 | 280 | # Apply deformed convolution 281 | return KPConv_deform_ops(query_points, 282 | support_points, 283 | neighbors_indices, 284 | features, 285 | K_points, 286 | offsets, 287 | modulations, 288 | K_values, 289 | KP_extent, 290 | KP_influence, 291 | aggregation_mode) 292 | 293 | 294 | def KPConv_deformable_v2(query_points, 295 | support_points, 296 | neighbors_indices, 297 | features, 298 | K_values, 299 | w0, b0, 300 | fixed='center', 301 | KP_extent=1.0, 302 | KP_influence='linear', 303 | aggregation_mode='sum', 304 | modulated=False): 305 | """ 306 | This alternate version uses a pointwise MLP instead of KPConv to get the offset. It has thus less parameters. 307 | It also fixes the center point to remain in the center in any case. This definition offers similar performances 308 | 309 | :param query_points: float32[n_points, dim] - input query points (center of neighborhoods) 310 | :param support_points: float32[n0_points, dim] - input support points (from which neighbors are taken) 311 | :param neighbors_indices: int32[n_points, n_neighbors] - indices of neighbors of each point 312 | :param features: float32[n_points, in_fdim] - input features 313 | :param K_values: float32[n_kpoints, in_fdim, out_fdim] - weights of the kernel 314 | :param w0: float32[n_points, dim * n_kpoints] - weights of the unary for offsets 315 | :param b0: float32[dim * n_kpoints] - bias of the unary for offsets. 316 | :param fixed: string in ('none', 'center' or 'verticals') - fix position of certain kernel points 317 | :param KP_extent: float32 - influence radius of each kernel point 318 | :param KP_influence: string in ('constant', 'linear', 'gaussian') - influence function of the kernel points 319 | :param aggregation_mode: string in ('closest', 'sum') - behavior of the convolution 320 | :param modulated: bool - If deformable conv should be modulated 321 | 322 | :return: output_features float32[n_points, out_fdim] 323 | """ 324 | # Initial kernel extent for this layer 325 | K_radius = 1.5 * KP_extent 326 | 327 | # Number of kernel points 328 | num_kpoints = int(K_values.shape[0]) 329 | 330 | # Check point dimension (currently only 3D is supported) 331 | points_dim = int(query_points.shape[1]) 332 | 333 | # Create one kernel disposition (as numpy array). Choose the KP distance to center thanks to the KP extent 334 | K_points_numpy = create_kernel_points(K_radius, 335 | num_kpoints, 336 | num_kernels=1, 337 | dimension=points_dim, 338 | fixed=fixed) 339 | K_points_numpy = K_points_numpy.reshape((num_kpoints, points_dim)) 340 | 341 | # Create the pytorch variable 342 | K_points = torch.from_numpy(K_points_numpy.astype(np.float32)) 343 | if K_values.is_cuda: 344 | K_points = K_points.to(K_values.device) 345 | 346 | ############################# 347 | # Pointwise MLP for offsets 348 | ############################# 349 | # Create independant weight for the first convolution and a bias term as no batch normalization happen 350 | # if modulated: 351 | # offset_dim = (points_dim + 1) * num_kpoints 352 | # else: 353 | # offset_dim = points_dim * num_kpoints 354 | # shape0 = K_values.shape.as_list() 355 | # w0 = torch.zeros([shape0[1], offset_dim], dtype=torch.float32) # offset_mlp_weights 356 | # b0 = torch.zeros(offset_dim, dtype=torch.float32) # offset_mlp_bias 357 | 358 | # Get features from mlp 359 | features0 = unary_convolution(features, w0) + b0 360 | # TODO: need to do something to reduce the point size from len(support_points) to len(query_points). 361 | 362 | if modulated: 363 | 364 | # Get offset (in normalized scale) from features 365 | offsets = features0[:, :points_dim * (num_kpoints - 1)] 366 | offsets = offsets.reshape([-1, (num_kpoints - 1), points_dim]) 367 | 368 | # Get modulations 369 | modulations = 2 * torch.sigmoid(features0[:, points_dim * (num_kpoints-1):]) 370 | 371 | # No offset for the first Kernel points 372 | offsets = torch.cat([torch.zeros_like(offsets[:, :1, :]), offsets], dim=1) 373 | modulations = torch.cat([torch.zeros_like(modulations[:, :1]), modulations], dim=1) 374 | 375 | else: 376 | 377 | # Get offset (in normalized scale) from features 378 | offsets = features0.reshape([-1, (num_kpoints-1), points_dim]) 379 | 380 | # No offset for the first kernle points 381 | offsets = torch.cat([torch.zeros_like(offsets[:, :1, :]), offsets], dim=1) 382 | 383 | # No modulations 384 | modulations = None 385 | 386 | # Rescale offset for this layer 387 | offsets *= KP_extent 388 | 389 | ############################### 390 | # Build deformable KPConv graph 391 | ############################### 392 | 393 | # Apply deformed convolution 394 | return KPConv_deform_ops(query_points, 395 | support_points, 396 | neighbors_indices, 397 | features, 398 | K_points, 399 | offsets, 400 | modulations, 401 | K_values, 402 | KP_extent, 403 | KP_influence, 404 | aggregation_mode) 405 | 406 | 407 | def KPConv_deform_ops(query_points, 408 | support_points, 409 | neighbors_indices, 410 | features, 411 | K_points, 412 | offsets, 413 | modulations, 414 | K_values, 415 | KP_extent, 416 | KP_influence, 417 | mode): 418 | """ 419 | This function creates a graph of operations to define Deformable Kernel Point Convolution in tensorflow. See 420 | KPConv_deformable function above for a description of each parameter 421 | 422 | :param query_points: [n_points, dim] 423 | :param support_points: [n0_points, dim] 424 | :param neighbors_indices: [n_points, n_neighbors] 425 | :param features: [n_points, in_fdim] 426 | :param K_points: [n_kpoints, dim] 427 | :param offsets: [n_points, n_kpoints, dim] 428 | :param modulations: [n_points, n_kpoints] or None 429 | :param K_values: [n_kpoints, in_fdim, out_fdim] 430 | :param KP_extent: float32 431 | :param KP_influence: string 432 | :param mode: string 433 | 434 | :return: [n_points, out_fdim] 435 | """ 436 | 437 | # Get variables 438 | n_kp = int(K_points.shape[0]) 439 | shadow_ind = support_points.shape[0] 440 | 441 | # Add a fake point in the last row for shadow neighbors 442 | shadow_point = torch.ones_like(support_points[:1, :]) * 1e6 443 | support_points = torch.cat([support_points, shadow_point], dim=0) 444 | 445 | # Get neighbor points [n_points, n_neighbors, dim] 446 | neighbors = support_points[neighbors_indices, :] 447 | 448 | # Center every neighborhood 449 | neighbors = neighbors - query_points.unsqueeze(1) 450 | 451 | # Apply offsets to kernel points [n_points, n_kpoints, dim] 452 | deformed_K_points = torch.add(offsets, K_points) 453 | 454 | # Get all difference matrices [n_points, n_neighbors, n_kpoints, dim] 455 | differences = neighbors.unsqueeze(2) - deformed_K_points.unsqueeze(1) 456 | 457 | # Get the square distances [n_points, n_neighbors, n_kpoints] 458 | sq_distances = torch.sum(torch.mul(differences, differences), dim=3) 459 | 460 | # Boolean of the neighbors in range of a kernel point [n_points, n_neighbors] 461 | in_range = torch.any((sq_distances < KP_extent ** 2), dim=2).int() 462 | 463 | # New value of max neighbors 464 | new_max_neighb = torch.max(torch.sum(in_range, dim=1)) 465 | 466 | # For each row of neighbors, indices of the ones that are in range [n_points, new_max_neighb] 467 | new_neighb_bool, new_neighb_inds = in_range.topk(k=int(new_max_neighb)) 468 | 469 | # Gather new neighbor indices [n_points, new_max_neighb] 470 | # new_neighbors_indices = tf.batch_gather(neighbors_indices, new_neigh_inds) 471 | new_neighbors_indices = torch.gather(neighbors_indices, dim=1, index=new_neighb_inds) 472 | 473 | # Gather new distances to KP [n_points, new_max_neighb, n_kpoints] 474 | # new_sq_distances = tf.batch_gather(sq_distances, new_neighb_inds) 475 | # https://pytorch.org/docs/stable/torch.html#torch.gather 476 | # https://discuss.pytorch.org/t/question-about-torch-gather-with-3-dimensions/19891/2 477 | new_sq_distances = sq_distances.gather(dim=1, index=new_neighb_inds.unsqueeze(-1).repeat(1,1,sq_distances.shape[-1])) 478 | 479 | # New shadow neighbors have to point to the last shadow point 480 | new_neighbors_indices *= new_neighb_bool.long() 481 | new_neighbors_indices += (1 - new_neighb_bool.long()) * shadow_ind 482 | 483 | # Get Kernel point influences [n_points, n_kpoints, n_neighbors] 484 | if KP_influence == 'constant': 485 | # Every point get an influence of 1. 486 | all_weights = (new_sq_distances < KP_extent ** 2).float32() 487 | all_weights = all_weights.transpose(1, 2) 488 | 489 | elif KP_influence == 'linear': 490 | # Influence decrease linearly with the distance, and get to zero when d = KP_extent. 491 | corr = 1 - torch.sqrt(new_sq_distances + 1e-10) / KP_extent 492 | all_weights = torch.max(corr, torch.zeros_like(new_sq_distances)) 493 | all_weights = all_weights.transpose(1, 2) 494 | 495 | elif KP_influence == 'gaussian': 496 | # Influence in gaussian of the distance. 497 | sigma = KP_extent * 0.3 498 | all_weights = radius_gaussian(sq_distances, sigma) 499 | all_weights = all_weights.transpose(1, 2) 500 | else: 501 | raise ValueError('Unknown influence function type (config.KP_influence)') 502 | 503 | # In case of closest mode, only the closest KP can influence each point 504 | if mode == 'closest': 505 | pass 506 | elif mode != 'sum': 507 | raise ValueError("Unknown convolution mode. Should be 'closest' or 'sum'") 508 | 509 | features = torch.cat([features, torch.zeros_like(features[:1, :])], dim=0) 510 | 511 | # Get the features of each neighborhood [n_points, n_neighbors, in_fdim] 512 | neighborhood_features = features[new_neighbors_indices, :] 513 | 514 | # Apply distance weights [n_points, n_kpoints, in_fdim] 515 | weighted_features = torch.matmul(all_weights, neighborhood_features) 516 | 517 | # Apply modulations 518 | if modulations is not None: 519 | weighted_features *= modulations.unsqueeze(2) 520 | 521 | # Apply network weights [n_kpoints, n_points, out_fdim] 522 | weighted_features = weighted_features.transpose(0, 1) 523 | kernel_outputs = torch.matmul(weighted_features, K_values) 524 | 525 | # Convolution sum to get [n_points, out_fdim] 526 | output_features = torch.sum(kernel_outputs, dim=0, keepdim=False) 527 | 528 | # normalization term. 529 | # neighbor_features_sum = torch.sum(neighborhood_features, dim=-1) 530 | # neighbor_num = torch.sum(torch.gt(neighbor_features_sum, 0.0), dim=-1) 531 | # neighbor_num = torch.max(neighbor_num, torch.ones_like(neighbor_num)) 532 | # output_features = output_features / neighbor_num.unsqueeze(1) 533 | 534 | return output_features 535 | -------------------------------------------------------------------------------- /kernels/kernel_points.py: -------------------------------------------------------------------------------- 1 | # 2 | # 3 | # 0=================================0 4 | # | Kernel Point Convolutions | 5 | # 0=================================0 6 | # 7 | # 8 | # ---------------------------------------------------------------------------------------------------------------------- 9 | # 10 | # Functions handling the disposition of kernel points. 11 | # 12 | # ---------------------------------------------------------------------------------------------------------------------- 13 | # 14 | # Hugues THOMAS - 11/06/2018 15 | # 16 | 17 | 18 | # ------------------------------------------------------------------------------------------ 19 | # 20 | # Imports and global variables 21 | # \**********************************/ 22 | # 23 | 24 | 25 | # Import numpy package and name it "np" 26 | import numpy as np 27 | import matplotlib.pyplot as plt 28 | from os import makedirs 29 | from os.path import join, exists 30 | 31 | from utils.ply import read_ply, write_ply 32 | 33 | 34 | # ------------------------------------------------------------------------------------------ 35 | # 36 | # Functions 37 | # \***************/ 38 | # 39 | # 40 | 41 | def kernel_point_optimization_debug(radius, num_points, num_kernels=1, dimension=3, fixed='center', ratio=1.0, verbose=0): 42 | """ 43 | Creation of kernel point via optimization of potentials. 44 | :param radius: Radius of the kernels 45 | :param num_points: points composing kernels 46 | :param num_kernels: number of wanted kernels 47 | :param dimension: dimension of the space 48 | :param fixed: fix position of certain kernel points ('none', 'center' or 'verticals') 49 | :param ratio: ratio of the radius where you want the kernels points to be placed 50 | :param verbose: display option 51 | :return: points [num_kernels, num_points, dimension] 52 | """ 53 | 54 | ####################### 55 | # Parameters definition 56 | ####################### 57 | 58 | # Radius used for optimization (points are rescaled afterwards) 59 | radius0 = 1 60 | diameter0 = 2 61 | 62 | # Factor multiplicating gradients for moving points (~learning rate) 63 | moving_factor = 1e-2 64 | continuous_moving_decay = 0.9995 65 | 66 | # Gradient threshold to stop optimization 67 | thresh = 1e-5 68 | 69 | # Gradient clipping value 70 | clip = 0.05 * radius0 71 | 72 | ####################### 73 | # Kernel initialization 74 | ####################### 75 | 76 | # Random kernel points 77 | kernel_points = np.random.rand(num_kernels * num_points - 1, dimension) * diameter0 - radius0 78 | while (kernel_points.shape[0] < num_kernels * num_points): 79 | new_points = np.random.rand(num_kernels * num_points - 1, dimension) * diameter0 - radius0 80 | kernel_points = np.vstack((kernel_points, new_points)) 81 | d2 = np.sum(np.power(kernel_points, 2), axis=1) 82 | kernel_points = kernel_points[d2 < 0.5 * radius0 * radius0, :] 83 | kernel_points = kernel_points[:num_kernels * num_points, :].reshape((num_kernels, num_points, -1)) 84 | 85 | # Optionnal fixing 86 | if fixed == 'center': 87 | kernel_points[:, 0, :] *= 0 88 | if fixed == 'verticals': 89 | kernel_points[:, :3, :] *= 0 90 | kernel_points[:, 1, -1] += 2 * radius0 / 3 91 | kernel_points[:, 2, -1] -= 2 * radius0 / 3 92 | 93 | ##################### 94 | # Kernel optimization 95 | ##################### 96 | 97 | # Initiate figure 98 | if verbose>1: 99 | fig = plt.figure() 100 | 101 | saved_gradient_norms = np.zeros((10000, num_kernels)) 102 | old_gradient_norms = np.zeros((num_kernels, num_points)) 103 | for iter in range(10000): 104 | 105 | # Compute gradients 106 | # ***************** 107 | 108 | # Derivative of the sum of potentials of all points 109 | A = np.expand_dims(kernel_points, axis=2) 110 | B = np.expand_dims(kernel_points, axis=1) 111 | interd2 = np.sum(np.power(A - B, 2), axis=-1) 112 | inter_grads = (A - B) / (np.power(np.expand_dims(interd2, -1), 3/2) + 1e-6) 113 | inter_grads = np.sum(inter_grads, axis=1) 114 | 115 | # Derivative of the radius potential 116 | circle_grads = 10*kernel_points 117 | 118 | # All gradients 119 | gradients = inter_grads + circle_grads 120 | 121 | if fixed == 'verticals': 122 | gradients[:, 1:3, :-1] = 0 123 | 124 | # Stop condition 125 | # ************** 126 | 127 | # Compute norm of gradients 128 | gradients_norms = np.sqrt(np.sum(np.power(gradients, 2), axis=-1)) 129 | saved_gradient_norms[iter, :] = np.max(gradients_norms, axis=1) 130 | 131 | # Stop if all moving points are gradients fixed (low gradients diff) 132 | 133 | if fixed == 'center' and np.max(np.abs(old_gradient_norms[:, 1:] - gradients_norms[:, 1:])) < thresh: 134 | break 135 | elif fixed == 'verticals' and np.max(np.abs(old_gradient_norms[:, 3:] - gradients_norms[:, 3:])) < thresh: 136 | break 137 | elif np.max(np.abs(old_gradient_norms - gradients_norms)) < thresh: 138 | break 139 | old_gradient_norms = gradients_norms 140 | 141 | # Move points 142 | # *********** 143 | 144 | # Clip gradient to get moving dists 145 | moving_dists = np.minimum(moving_factor * gradients_norms, clip) 146 | 147 | # Fix central point 148 | if fixed == 'center': 149 | moving_dists[:, 0] = 0 150 | if fixed == 'verticals': 151 | moving_dists[:, 0] = 0 152 | 153 | # Move points 154 | kernel_points -= np.expand_dims(moving_dists, -1) * gradients / np.expand_dims(gradients_norms + 1e-6, -1) 155 | 156 | if verbose: 157 | print('iter {:5d} / max grad = {:f}'.format(iter, np.max(gradients_norms[:, 3:]))) 158 | if verbose > 1: 159 | plt.clf() 160 | plt.plot(kernel_points[0, :, 0], kernel_points[0, :, 1], '.') 161 | circle = plt.Circle((0, 0), radius, color='r', fill=False) 162 | fig.axes[0].add_artist(circle) 163 | fig.axes[0].set_xlim((-radius*1.1, radius*1.1)) 164 | fig.axes[0].set_ylim((-radius*1.1, radius*1.1)) 165 | fig.axes[0].set_aspect('equal') 166 | plt.draw() 167 | plt.pause(0.001) 168 | plt.show(block=False) 169 | print(moving_factor) 170 | 171 | # moving factor decay 172 | moving_factor *= continuous_moving_decay 173 | 174 | # Rescale radius to fit the wanted ratio of radius 175 | r = np.sqrt(np.sum(np.power(kernel_points, 2), axis=-1)) 176 | kernel_points *= ratio / np.mean(r[:, 1:]) 177 | 178 | # Rescale kernels with real radius 179 | return kernel_points * radius, saved_gradient_norms 180 | 181 | 182 | def load_kernels(radius, num_kpoints, num_kernels, dimension, fixed): 183 | 184 | # Number of tries in the optimization process, to ensure we get the most stable disposition 185 | num_tries = 100 186 | 187 | # Kernel directory 188 | kernel_dir = 'kernels/' 189 | if not exists(kernel_dir): 190 | makedirs(kernel_dir) 191 | 192 | # Kernel_file 193 | if dimension == 3: 194 | kernel_file = join(kernel_dir, 'k_{:03d}_{:s}.ply'.format(num_kpoints, fixed)) 195 | elif dimension == 2: 196 | kernel_file = join(kernel_dir, 'k_{:03d}_{:s}_2D.ply'.format(num_kpoints, fixed)) 197 | else: 198 | raise ValueError('Unsupported dimpension of kernel : ' + str(dimension)) 199 | 200 | # Check if already done 201 | if not exists(kernel_file): 202 | 203 | # Create kernels 204 | kernel_points, grad_norms = kernel_point_optimization_debug(1.0, 205 | num_kpoints, 206 | num_kernels=num_tries, 207 | dimension=dimension, 208 | fixed=fixed, 209 | verbose=0) 210 | 211 | # Find best candidate 212 | best_k = np.argmin(grad_norms[-1, :]) 213 | 214 | # Save points 215 | original_kernel = kernel_points[best_k, :, :] 216 | write_ply(kernel_file, original_kernel, ['x', 'y', 'z']) 217 | 218 | else: 219 | data = read_ply(kernel_file) 220 | original_kernel = np.vstack((data['x'], data['y'], data['z'])).T 221 | 222 | # N.B. 2D kernels are not supported yet 223 | if dimension == 2: 224 | return original_kernel 225 | 226 | # Random rotations depending of the fixed points 227 | if fixed == 'verticals': 228 | 229 | # Create random rotations 230 | thetas = np.random.rand(num_kernels) * 2 * np.pi 231 | c, s = np.cos(thetas), np.sin(thetas) 232 | R = np.zeros((num_kernels, 3, 3), dtype=np.float32) 233 | R[:, 0, 0] = c 234 | R[:, 1, 1] = c 235 | R[:, 2, 2] = 1 236 | R[:, 0, 1] = s 237 | R[:, 1, 0] = -s 238 | 239 | # Scale kernels 240 | original_kernel = radius * np.expand_dims(original_kernel, 0) 241 | 242 | # Rotate kernels 243 | kernels = np.matmul(original_kernel, R) 244 | 245 | else: 246 | 247 | # Create random rotations 248 | u = np.ones((num_kernels, 3)) 249 | v = np.ones((num_kernels, 3)) 250 | wrongs = np.abs(np.sum(u * v, axis=1)) > 0.99 251 | while np.any(wrongs): 252 | new_u = np.random.rand(num_kernels, 3) * 2 - 1 253 | new_u = new_u / np.expand_dims(np.linalg.norm(new_u, axis=1) + 1e-9, -1) 254 | u[wrongs, :] = new_u[wrongs, :] 255 | new_v = np.random.rand(num_kernels, 3) * 2 - 1 256 | new_v = new_v / np.expand_dims(np.linalg.norm(new_v, axis=1) + 1e-9, -1) 257 | v[wrongs, :] = new_v[wrongs, :] 258 | wrongs = np.abs(np.sum(u * v, axis=1)) > 0.99 259 | 260 | # Make v perpendicular to u 261 | v -= np.expand_dims(np.sum(u * v, axis=1), -1) * u 262 | v = v / np.expand_dims(np.linalg.norm(v, axis=1) + 1e-9, -1) 263 | 264 | # Last rotation vector 265 | w = np.cross(u, v) 266 | R = np.stack((u, v, w), axis=-1) 267 | 268 | # Scale kernels 269 | original_kernel = radius * np.expand_dims(original_kernel, 0) 270 | 271 | # Rotate kernels 272 | kernels = np.matmul(original_kernel, R) 273 | 274 | # Add a small noise 275 | kernels = kernels 276 | kernels = kernels + np.random.normal(scale=radius*0.01, size=kernels.shape) 277 | 278 | return kernels 279 | 280 | 281 | 282 | 283 | -------------------------------------------------------------------------------- /misc/num_seg_classes.txt: -------------------------------------------------------------------------------- 1 | Airplane 4 2 | Bag 2 3 | Cap 2 4 | Car 4 5 | Chair 4 6 | Earphone 3 7 | Guitar 3 8 | Knife 2 9 | Lamp 4 10 | Laptop 2 11 | Motorbike 6 12 | Mug 2 13 | Pistol 3 14 | Rocket 3 15 | Skateboard 3 16 | Table 3 17 | -------------------------------------------------------------------------------- /models/KPCNN.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from models.network_blocks import get_block, weight_variable 3 | import torch 4 | import torch.nn as nn 5 | 6 | 7 | class KPCNN(nn.Module): 8 | def __init__(self, config): 9 | # TODO: One big difference is in tensorflow version, each block receives same parameters (layer_ind, inputs, features, radius, fdim, config) 10 | # which in my opinion is because it will be easy to write a unified code for creating models of different architecture, but the readability 11 | # of each block is not satisfactory since the block functions are coupling with the data prepartion pipelines. 12 | # Instead, in pytorch version, I save the fixed parameter(like radius, in_fdim, out_fdim) for each layer (or module) when creating the layer 13 | # and only provide the necessary inputs (query_points, supporting_points, neighbors_indices, features) during forward, so that the function 14 | # of each layer is more clear. And when we want to modify the data preparation pipeline, there is no need to change to code for KPConv layers. 15 | # We only need to modify the code for parsing the architecture list and feed each layer with the correct inputs. 16 | 17 | super(KPCNN, self).__init__() 18 | self.config = config 19 | self.blocks = nn.ModuleDict() 20 | 21 | # Feature Extraction Module 22 | r = config.first_subsampling_dl * config.density_parameter 23 | in_fdim = config.in_features_dim 24 | out_fdim = config.first_features_dim 25 | layer = 0 26 | block_in_layer = 0 27 | for block_i, block in enumerate(config.architecture): 28 | # Detect upsampling block to stop 29 | if 'upsample' in block: 30 | break 31 | 32 | is_strided = 'strided' in block 33 | self.blocks[f'layer{layer}/{block}'] = get_block(block, config, in_fdim, out_fdim, radius=r, strided=is_strided) 34 | 35 | # update feature dimension 36 | in_fdim = out_fdim 37 | block_in_layer += 1 38 | 39 | if 'pool' in block or 'strided' in block: 40 | # Update radius and feature dimension for next layer 41 | out_fdim *= 2 42 | r *= 2 43 | layer += 1 44 | block_in_layer = 0 45 | 46 | # Classification Head 47 | self.blocks['classification_head'] = nn.Sequential( 48 | nn.Linear(out_fdim, 1024), 49 | # nn.BatchNorm1d(1024, momentum=config.batch_norm_momentum, eps=1e-6), 50 | nn.LeakyReLU(negative_slope=0.2), 51 | nn.Dropout(p=config.dropout), 52 | nn.Linear(1024, config.num_classes) 53 | ) 54 | 55 | # print(list(self.parameters())) 56 | 57 | def forward(self, inputs): 58 | F = self.feature_extraction(inputs) 59 | logits = self.classification_head(F[-1]) 60 | return logits 61 | 62 | def feature_extraction(self, inputs): 63 | # Current radius of convolution and feature dimension 64 | r = self.config.first_subsampling_dl * self.config.density_parameter 65 | layer = 0 66 | fdim = self.config.first_features_dim 67 | 68 | # Input features 69 | features = inputs['features'] 70 | F = [] 71 | 72 | # Loop over consecutive blocks 73 | block_in_layer = 0 74 | for block_i, block in enumerate(self.config.architecture): 75 | 76 | # Detect change to next layer 77 | if np.any([tmp in block for tmp in ['pool', 'strided', 'upsample', 'global']]): 78 | # Save this layer features 79 | F += [features] 80 | 81 | # Detect upsampling block to stop 82 | if 'upsample' in block: 83 | break 84 | 85 | # Get the function for this layer 86 | block_ops = self.blocks[f'layer{layer}/{block}'] 87 | 88 | # Apply the layer function defining tf ops 89 | if block == 'global_average': 90 | stack_lengths = inputs['stack_lengths'] 91 | features = block_ops(stack_lengths, features) 92 | else: 93 | if block in ['unary', 'simple', 'resnet', 'resnetb', 'resnetb_deformable']: 94 | query_points = inputs['points'][layer] 95 | support_points = inputs['points'][layer] 96 | neighbors_indices = inputs['neighbors'][layer] 97 | elif block in ['simple_strided', 'resnetb_strided', 'resnetb_deformable_strided']: 98 | query_points = inputs['points'][layer + 1] 99 | support_points = inputs['points'][layer] 100 | neighbors_indices = inputs['pools'][layer] 101 | else: 102 | raise ValueError("Unknown block type.") 103 | features = block_ops(query_points, support_points, neighbors_indices, features) 104 | 105 | # Index of block in this layer 106 | block_in_layer += 1 107 | 108 | # Detect change to a subsampled layer 109 | if 'pool' in block or 'strided' in block: 110 | # Update radius and feature dimension for next layer 111 | layer += 1 112 | r *= 2 113 | fdim *= 2 114 | block_in_layer = 0 115 | 116 | # Save feature vector after global pooling 117 | if 'global' in block: 118 | # Save this layer features 119 | F += [features] 120 | return F 121 | 122 | def classification_head(self, features): 123 | logits = self.blocks['classification_head'](features) 124 | return logits 125 | 126 | 127 | if __name__ == '__main__': 128 | from training_ShapeNetCls import ShapeNetPartConfig 129 | from datasets.ShapeNet import ShapeNetDataset 130 | from datasets.dataloader import get_dataloader 131 | 132 | config = ShapeNetPartConfig() 133 | datapath = "./data/shapenetcore_partanno_segmentation_benchmark_v0" 134 | dset = ShapeNetDataset(root=datapath, config=config, first_subsampling_dl=0.01, classification=True) 135 | dataloader = get_dataloader(dset, batch_size=4) 136 | model = KPCNN(config) 137 | 138 | for iter, input in enumerate(dataloader): 139 | output = model(input) 140 | print("Predict:", output) 141 | print("GT:", input['labels']) 142 | break 143 | -------------------------------------------------------------------------------- /models/KPFCNN.py: -------------------------------------------------------------------------------- 1 | from models.network_blocks import get_block 2 | from models.KPCNN import KPCNN 3 | import torch 4 | import torch.nn as nn 5 | 6 | 7 | class KPFCNN(nn.Module): 8 | def __init__(self, config): 9 | super(KPFCNN, self).__init__() 10 | self.encoder = KPCNN(config) 11 | self.config = config 12 | self.blocks = nn.ModuleDict() 13 | 14 | # Feature Extraction Module 15 | # Find first upsampling block 16 | start_i = 0 17 | for block_i, block in enumerate(config.architecture): 18 | if 'upsample' in block: 19 | start_i = block_i 20 | break 21 | 22 | layer = config.num_layers - 1 23 | r = config.first_subsampling_dl * config.density_parameter * 2 ** layer 24 | in_fdim = config.first_features_dim * 2 ** layer 25 | out_fdim = in_fdim 26 | block_in_layer = 0 27 | for block_i, block in enumerate(config.architecture[start_i:]): 28 | 29 | is_strided = 'strided' in block 30 | self.blocks[f'layer{layer}/{block}'] = get_block(block, config, int(1.5 * in_fdim), out_fdim, radius=r, strided=is_strided) 31 | 32 | # update feature dimension 33 | in_fdim = out_fdim 34 | block_in_layer += 1 35 | 36 | # Detect change to a subsampled layer 37 | if 'upsample' in block: 38 | # Update radius and feature dimension for next layer 39 | out_fdim = out_fdim // 2 40 | r *= 0.5 41 | layer -= 1 42 | block_in_layer = 0 43 | 44 | # Segmentation Head 45 | self.blocks['segmentation_head'] = nn.Sequential( 46 | nn.Linear(out_fdim, config.first_features_dim), 47 | nn.BatchNorm1d(config.first_features_dim, momentum=config.batch_norm_momentum, eps=1e-6), 48 | nn.LeakyReLU(negative_slope=0.2), 49 | # nn.Dropout(p=config.dropout), 50 | nn.Linear(config.first_features_dim, config.num_classes) 51 | ) 52 | # print(list(self.named_parameters())) 53 | 54 | def forward(self, inputs): 55 | features = self.feature_extraction(inputs) 56 | logits = self.segmentation_head(features) 57 | return logits 58 | 59 | def feature_extraction(self, inputs): 60 | F = self.encoder.feature_extraction(inputs) 61 | features = F[-1] 62 | 63 | # Current radius of convolution and feature dimension 64 | layer = self.config.num_layers - 1 65 | r = self.config.first_subsampling_dl * self.config.density_parameter * 2 ** layer 66 | fdim = self.config.first_features_dim * 2 ** layer 67 | 68 | # Find first upsampling block 69 | start_i = 0 70 | for block_i, block in enumerate(self.config.architecture): 71 | if 'upsample' in block: 72 | start_i = block_i 73 | break 74 | 75 | # Loop over upsampling blocks 76 | for block_i, block in enumerate(self.config.architecture[start_i:]): 77 | 78 | # Get the function for this layer 79 | block_ops = self.blocks[f'layer{layer}/{block}'] 80 | 81 | # Apply the layer function defining tf ops 82 | if 'upsample' in block: 83 | if block == 'nearest_upsample': 84 | upsample_indices = inputs['upsamples'][layer - 1] 85 | else: 86 | raise ValueError(f"Unknown block type. {block}") 87 | features = block_ops(upsample_indices, features) 88 | else: 89 | if block in ['unary', 'simple', 'resnet', 'resnetb']: 90 | query_points = inputs['points'][layer] 91 | support_points = inputs['points'][layer] 92 | neighbors_indices = inputs['neighbors'][layer] 93 | elif block in ['simple_strided', 'resnetb_strided', 'resnetb_deformable_strided']: 94 | query_points = inputs['points'][layer + 1] 95 | support_points = inputs['points'][layer] 96 | neighbors_indices = inputs['pools'][layer] 97 | else: 98 | raise ValueError(f"Unknown block type. {block}") 99 | features = block_ops(query_points, support_points, neighbors_indices, features) 100 | 101 | # Detect change to a subsampled layer 102 | if 'upsample' in block: 103 | # Update radius and feature dimension for next layer 104 | layer -= 1 105 | r *= 0.5 106 | fdim = fdim // 2 107 | 108 | # Concatenate with CNN feature map 109 | features = torch.cat((features, F[layer]), dim=1) 110 | 111 | return features 112 | 113 | def segmentation_head(self, features): 114 | logits = self.blocks['segmentation_head'](features) 115 | return logits 116 | 117 | 118 | if __name__ == '__main__': 119 | from training_ShapeNetPart import ShapeNetPartConfig 120 | from datasets.ShapeNet import ShapeNetDataset 121 | from datasets.dataloader import get_dataloader 122 | 123 | config = ShapeNetPartConfig() 124 | datapath = "./data/shapenetcore_partanno_segmentation_benchmark_v0" 125 | dset = ShapeNetDataset(root=datapath, config=config, first_subsampling_dl=0.01, classification=False) 126 | dataloader = get_dataloader(dset, batch_size=1) 127 | model = KPFCNN(config) 128 | 129 | for iter, input in enumerate(dataloader): 130 | output = model(input) 131 | print(output.shape) 132 | break 133 | -------------------------------------------------------------------------------- /models/network_blocks.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import kernels.convolution_ops as conv_ops 3 | import torch 4 | from torch import nn 5 | from torch.nn import functional as F 6 | 7 | 8 | # ---------------------------------------------------------------------------------------------------------------------- 9 | # 10 | # Utilities 11 | # \***************/ 12 | # 13 | 14 | def weight_variable(size): 15 | # https://discuss.pytorch.org/t/implementing-truncated-normal-initializer/4778/21 16 | initial = np.random.normal(scale=np.sqrt(2 / size[-1]), size=size) 17 | initial[initial > 2 * np.sqrt(2 / size[-1])] = 0 # truncated 18 | initial[initial < -2 * np.sqrt(2 / size[-1])] = 0 # truncated 19 | weight = nn.Parameter(torch.from_numpy(initial).float(), requires_grad=True) 20 | return weight 21 | 22 | 23 | def bias_variable(size): 24 | initial = torch.zeros(size=size, dtype=torch.float32) 25 | bias = nn.Parameter(initial, requires_grad=True) 26 | return bias 27 | 28 | 29 | def leaky_relu_layer(negative_slope=0.2): 30 | return nn.LeakyReLU(negative_slope=negative_slope) 31 | 32 | 33 | def ind_max_pool(features, inds): 34 | """ 35 | This pytorch operation compute a maxpooling according to the list of indices 'inds'. 36 | > x = [n1, d] features matrix 37 | > inds = [n2, max_num] each row of this tensor is a list of indices of features to be pooled together 38 | >> output = [n2, d] pooled features matrix 39 | """ 40 | 41 | # Add a last row with minimum features for shadow pools 42 | features = torch.cat([features, torch.min(features, dim=0, keepdim=True)[0]], dim=0) 43 | 44 | # Get features for each pooling cell [n2, max_num, d] 45 | pool_features = features[inds, :] 46 | 47 | # Pool the maximum 48 | return torch.max(pool_features, dim=1)[0] 49 | 50 | 51 | def closest_pool(features, upsample_indices): 52 | """ 53 | This tensorflow operation compute a pooling according to the list of indices 'inds'. 54 | > features = [n1, d] features matrix 55 | > upsample_indices = [n2, max_num] We only use the first column of this which should be the closest points too pooled positions 56 | >> output = [n2, d] pooled features matrix 57 | """ 58 | 59 | # Add a last row with minimum features for shadow pools 60 | features = torch.cat([features, torch.zeros_like(features[:1, :])], dim=0) 61 | 62 | # Get features for each pooling cell [n2, d] 63 | pool_features = features[upsample_indices[:, 0], :] 64 | 65 | return pool_features 66 | 67 | 68 | # def KPConv(query_points, support_points, neighbors_indices, features, K_values, radius, config): 69 | # """ 70 | # Returns the output features of a KPConv 71 | # """ 72 | # 73 | # # Get KP extent from current radius and config density 74 | # extent = config.KP_extent * radius / config.density_parameter 75 | # 76 | # # Convolution 77 | # return conv_ops.KPConv(query_points, 78 | # support_points, 79 | # neighbors_indices, 80 | # features, 81 | # K_values, 82 | # fixed=config.fixed_kernel_points, 83 | # KP_extent=extent, 84 | # KP_influence=config.KP_influence, 85 | # aggregation_mode=config.convolution_mode, ) 86 | 87 | 88 | class unary_block(nn.Module): 89 | def __init__(self, config, in_fdim, out_fdim): 90 | super(unary_block, self).__init__() 91 | self.config = config 92 | self.in_fdim, self.out_fdim = in_fdim, out_fdim 93 | self.weight = weight_variable([in_fdim, out_fdim]) 94 | if config.use_batch_norm: 95 | self.bn = nn.BatchNorm1d(out_fdim, momentum=config.momentum, eps=1e-6) 96 | self.relu = leaky_relu_layer() 97 | 98 | def forward(self, query_points, support_points, neighbors_indices, features): 99 | """ 100 | This module performs a unary 1x1 convolution (same with MLP) 101 | :param features: float32[n_points, in_fdim] - input features 102 | :return: output_features float32[n_points, out_fdim] 103 | """ 104 | x = conv_ops.unary_convolution(features, self.weight) 105 | if self.config.use_batch_norm: 106 | x = self.relu(self.bn(x)) 107 | else: 108 | x = self.relu(x) 109 | return x 110 | 111 | def __repr__(self): 112 | return f'unary(in_fdim={self.in_fdim}, out_fdim={self.out_fdim})' 113 | 114 | 115 | class simple_block(nn.Module): 116 | def __init__(self, config, in_fdim, out_fdim, radius, strided=False): 117 | super(simple_block, self).__init__() 118 | self.config = config 119 | self.radius = radius 120 | self.strided = strided 121 | self.in_fdim, self.out_fdim = in_fdim, out_fdim 122 | 123 | # kernel points weight 124 | self.weight = weight_variable([config.num_kernel_points, in_fdim, out_fdim]) 125 | if config.use_batch_norm: 126 | self.bn = nn.BatchNorm1d(out_fdim, momentum=config.momentum, eps=1e-6) 127 | self.relu = leaky_relu_layer() 128 | 129 | def forward(self, query_points, support_points, neighbors_indices, features): 130 | """ 131 | This module performs a Kernel Point Convolution. (both normal and strided version) 132 | :param query_points: float32[n_points, dim] - input query points (center of neighborhoods) 133 | :param support_points: float32[n0_points, dim] - input support points (from which neighbors are taken) 134 | :param neighbors_indices: int32[n_points, n_neighbors] - indices of neighbors of each point 135 | :param features: float32[n_points, in_fdim] - input features 136 | :return: output_features float32[n_points, out_fdim] 137 | """ 138 | x = conv_ops.KPConv(query_points, 139 | support_points, 140 | neighbors_indices, 141 | features, 142 | K_values=self.weight, 143 | fixed=self.config.fixed_kernel_points, 144 | KP_extent=self.config.KP_extent * self.radius / self.config.density_parameter, 145 | KP_influence=self.config.KP_influence, 146 | aggregation_mode=self.config.convolution_mode, ) 147 | if self.config.use_batch_norm: 148 | x = self.relu(self.bn(x)) 149 | else: 150 | x = self.relu(x) 151 | return x 152 | 153 | def __repr__(self): 154 | return f'simple(in_fdim={self.in_fdim}, out_fdim={self.out_fdim})' 155 | 156 | 157 | class simple_deformable_block(nn.Module): 158 | def __init__(self, config, in_fdim, out_fdim, radius, strided=False): 159 | super(simple_deformable_block, self).__init__() 160 | self.config = config 161 | self.radius = radius 162 | self.strided = strided 163 | self.in_fdim, self.out_fdim = in_fdim, out_fdim 164 | self.deformable_v2 = False 165 | 166 | # kernel points weight 167 | self.weight = weight_variable([config.num_kernel_points, in_fdim, out_fdim]) 168 | if config.use_batch_norm: 169 | self.bn = nn.BatchNorm1d(out_fdim, momentum=config.momentum, eps=1e-6) 170 | self.relu = leaky_relu_layer() 171 | 172 | point_dim = 4 if self.config.modulated else 3 173 | if self.deformable_v2: 174 | # for deformable_v2, there is no offset for the first kernel point. 175 | offset_dim = point_dim * (config.num_kernel_points - 1) 176 | self.offset_weight = weight_variable([in_fdim, offset_dim]) 177 | self.offset_bias = bias_variable([offset_dim]) 178 | else: 179 | offset_dim = point_dim * config.num_kernel_points 180 | self.offset_weight = weight_variable([config.num_kernel_points, in_fdim, offset_dim]) 181 | self.offset_bias = bias_variable([offset_dim]) 182 | 183 | def forward(self, query_points, support_points, neighbors_indices, features): 184 | """ 185 | This module performs a Kernel Point Convolution. (both normal and strided version) 186 | :param query_points: float32[n_points, dim] - input query points (center of neighborhoods) 187 | :param support_points: float32[n0_points, dim] - input support points (from which neighbors are taken) 188 | :param neighbors_indices: int32[n_points, n_neighbors] - indices of neighbors of each point 189 | :param features: float32[n_points, in_fdim] - input features 190 | :return: output_features float32[n_points, out_fdim] 191 | """ 192 | if self.deformable_v2: 193 | x = conv_ops.KPConv_deformable_v2(query_points, 194 | support_points, 195 | neighbors_indices, 196 | features, 197 | K_values=self.weight, 198 | w0=self.offset_weight, 199 | b0=self.offset_bias, 200 | fixed=self.config.fixed_kernel_points, 201 | KP_extent=self.config.KP_extent * self.radius / self.config.density_parameter, 202 | KP_influence=self.config.KP_influence, 203 | aggregation_mode=self.config.convolution_mode, 204 | modulated=self.config.modulated) 205 | else: 206 | x = conv_ops.KPConv_deformable(query_points, 207 | support_points, 208 | neighbors_indices, 209 | features, 210 | K_values=self.weight, 211 | w0=self.offset_weight, 212 | b0=self.offset_bias, 213 | fixed=self.config.fixed_kernel_points, 214 | KP_extent=self.config.KP_extent * self.radius / self.config.density_parameter, 215 | KP_influence=self.config.KP_influence, 216 | aggregation_mode=self.config.convolution_mode, 217 | modulated=self.config.modulated) 218 | if self.config.use_batch_norm: 219 | x = self.relu(self.bn(x)) 220 | else: 221 | x = self.relu(x) 222 | return x 223 | 224 | def __repr__(self): 225 | return f'simple_deformable(in_fdim={self.in_fdim}, out_fdim={self.out_fdim})' 226 | 227 | 228 | class resnet_block(nn.Module): 229 | def __init__(self, config, in_fdim, out_fdim, radius): 230 | super(resnet_block, self).__init__() 231 | self.config = config 232 | self.radius = radius 233 | self.in_fdim, self.out_fdim = in_fdim, out_fdim 234 | self.conv1 = simple_block(config, in_fdim, out_fdim, radius) 235 | self.conv2 = simple_block(config, out_fdim, out_fdim, radius) 236 | self.shortcut = unary_block(config, in_fdim, out_fdim) 237 | self.relu = leaky_relu_layer() 238 | 239 | def forward(self, query_points, support_points, neighbors_indices, features): 240 | """ 241 | This module performs a resnet double convolution (two convolution vgglike and a shortcut) 242 | :param query_points: float32[n_points, dim] - input query points (center of neighborhoods) 243 | :param support_points: float32[n0_points, dim] - input support points (from which neighbors are taken) 244 | :param neighbors_indices: int32[n_points, n_neighbors] - indices of neighbors of each point 245 | :param features: float32[n_points, in_fdim] - input features 246 | :return: output_features float32[n_points, out_fdim] 247 | """ 248 | shortcut = self.shortcut(query_points, support_points, neighbors_indices, features) 249 | features = self.conv1(query_points, support_points, neighbors_indices, features) 250 | features = self.conv2(query_points, support_points, neighbors_indices, features) 251 | return self.relu(shortcut + features) 252 | 253 | 254 | class resnetb_block(nn.Module): 255 | def __init__(self, config, in_fdim, out_fdim, radius, strided): 256 | super(resnetb_block, self).__init__() 257 | self.config = config 258 | self.radius = radius 259 | self.strided = strided 260 | self.in_fdim, self.out_fdim = in_fdim, out_fdim 261 | self.conv1 = unary_block(config, in_fdim, out_fdim // 2) 262 | self.conv2 = simple_block(config, out_fdim // 2, out_fdim // 2, radius, strided=strided) 263 | # TODO: origin implementation this last conv change feature dim to out_fdim * 2 264 | self.conv3 = unary_block(config, out_fdim // 2, out_fdim) 265 | self.shortcut = unary_block(config, in_fdim, out_fdim) 266 | self.relu = leaky_relu_layer() 267 | 268 | def forward(self, query_points, support_points, neighbors_indices, features): 269 | """ 270 | This module performs a resnet bottleneck convolution (1conv > KPconv > 1conv + shortcut) 271 | Both resnetb and resnetb_strided use the module. 272 | :param query_points: float32[n_points, dim] - input query points (center of neighborhoods) 273 | :param support_points: float32[n0_points, dim] - input support points (from which neighbors are taken) 274 | :param neighbors_indices: int32[n_points, n_neighbors] - indices of neighbors of each point 275 | :param features: float32[n_points, in_fdim] - input features 276 | :return: output_features float32[n_points, out_fdim] 277 | """ 278 | origin_features = features # save for shortcut 279 | features = self.conv1(query_points, support_points, neighbors_indices, features) 280 | features = self.conv2(query_points, support_points, neighbors_indices, features) 281 | features = self.conv3(query_points, support_points, neighbors_indices, features) 282 | # TODO: origin implementation has two kinds of shortcut. 283 | if self.strided is False: # for resnetb 284 | shortcut = self.shortcut(query_points, support_points, neighbors_indices, origin_features) 285 | else: # for resnetb_strided 286 | pool_features = ind_max_pool(origin_features, neighbors_indices) 287 | shortcut = self.shortcut(query_points, support_points, neighbors_indices, pool_features) 288 | return self.relu(shortcut + features) 289 | 290 | 291 | class resnetb_deformable_block(nn.Module): 292 | def __init__(self, config, in_fdim, out_fdim, radius, strided): 293 | super(resnetb_deformable_block, self).__init__() 294 | self.config = config 295 | self.radius = radius 296 | self.strided = strided 297 | self.in_fdim, self.out_fdim = in_fdim, out_fdim 298 | self.conv1 = unary_block(config, in_fdim, out_fdim // 2) 299 | self.conv2 = simple_deformable_block(config, out_fdim // 2, out_fdim // 2, radius, strided=strided) 300 | # TODO: origin implementation this last conv change feature dim to out_fdim * 2 301 | self.conv3 = unary_block(config, out_fdim // 2, out_fdim) 302 | self.shortcut = unary_block(config, in_fdim, out_fdim) 303 | self.relu = leaky_relu_layer() 304 | 305 | def forward(self, query_points, support_points, neighbors_indices, features): 306 | """ 307 | This module performs a resnet deformable bottleneck convolution (1conv > KPconv > 1conv + shortcut) 308 | Both resnetb_deformable and resnetb_deformable_strided use the module. 309 | :param query_points: float32[n_points, dim] - input query points (center of neighborhoods) 310 | :param support_points: float32[n0_points, dim] - input support points (from which neighbors are taken) 311 | :param neighbors_indices: int32[n_points, n_neighbors] - indices of neighbors of each point 312 | :param features: float32[n_points, in_fdim] - input features 313 | :return: output_features float32[n_points, out_fdim] 314 | """ 315 | origin_features = features # save for shortcut 316 | features = self.conv1(query_points, support_points, neighbors_indices, features) 317 | features = self.conv2(query_points, support_points, neighbors_indices, features) 318 | features = self.conv3(query_points, support_points, neighbors_indices, features) 319 | # TODO: origin implementation has two kinds of shortcut. 320 | if self.strided is False: # for resnetb 321 | shortcut = self.shortcut(query_points, support_points, neighbors_indices, origin_features) 322 | else: # for resnetb_strided 323 | pool_features = ind_max_pool(origin_features, neighbors_indices) 324 | shortcut = self.shortcut(query_points, support_points, neighbors_indices, pool_features) 325 | return self.relu(shortcut + features) 326 | 327 | 328 | class nearest_upsample_block(nn.Module): 329 | def __init__(self, config): 330 | super(nearest_upsample_block, self).__init__() 331 | self.config = config 332 | 333 | def forward(self, upsample_indices, features): 334 | """ 335 | This module performs an upsampling by nearest interpolation 336 | :param TODO: upsample_indices 337 | :param features: float32[n_points, in_fdim] - input features 338 | :return: 339 | """ 340 | pool_features = closest_pool(features, upsample_indices) 341 | return pool_features 342 | 343 | 344 | class global_average_block(nn.Module): 345 | def __init__(self, config): 346 | super(global_average_block, self).__init__() 347 | self.config = config 348 | 349 | def forward(self, stack_lengths, features): 350 | """ 351 | This module performs a global average over batch pooling 352 | :param features: float32[n_points, in_fdim] - input features 353 | :return: output_features: float32[batch_size, in_fdim] 354 | """ 355 | start_ind = 0 356 | average_feature_list = [] 357 | for length in stack_lengths[-1]: 358 | tmp = torch.mean(features[start_ind:start_ind + length], dim=0, keepdim=True) 359 | average_feature_list.append(tmp) 360 | start_ind += length 361 | return torch.cat(average_feature_list, dim=0) 362 | 363 | 364 | def get_block(block_name, config, in_fdim, out_fdim, radius, strided): 365 | if block_name == 'unary': 366 | return unary_block(config, in_fdim, out_fdim) 367 | if block_name == 'simple' or block_name == 'simple_strided': 368 | return simple_block(config, in_fdim, out_fdim, radius=radius, strided=strided) 369 | if block_name == 'nearest_upsample': 370 | return nearest_upsample_block(config) 371 | if block_name == 'global_average': 372 | return global_average_block(config) 373 | if block_name == 'resnet': # or block_name == 'resnet_strided': 374 | return resnet_block(config, in_fdim, out_fdim, radius=radius) 375 | if block_name == 'resnetb' or block_name == 'resnetb_strided': 376 | return resnetb_block(config, in_fdim, out_fdim, radius=radius, strided=strided) 377 | if block_name == 'resnetb_deformable' or block_name == 'resnetb_deformable_strided': 378 | return resnetb_deformable_block(config, in_fdim, out_fdim, radius=radius, strided=strided) 379 | -------------------------------------------------------------------------------- /pytorch_ops/batch_find_neighbors.cpp: -------------------------------------------------------------------------------- 1 | #include "torch/extension.h" 2 | #include "neighbors/neighbors.h" 3 | #include "vector" 4 | 5 | torch::Tensor batch_find_neighbors( 6 | torch::Tensor query_points, 7 | torch::Tensor support_points, 8 | torch::Tensor query_batches, 9 | torch::Tensor support_batches, 10 | float radius 11 | ){ 12 | 13 | // Points Dimensions 14 | int Nq = (int)query_points.size(0); 15 | int Ns = (int)support_points.size(0); 16 | 17 | // Number of batches 18 | int Nb = (int)query_batches.size(0); 19 | 20 | // get the data as std vector of points 21 | vector queries = vector((PointXYZ*)query_points.data(), 22 | (PointXYZ*)query_points.data() + Nq); 23 | vector supports = vector((PointXYZ*)support_points.data(), 24 | (PointXYZ*)support_points.data()+ Ns); 25 | 26 | // Batches lengths 27 | vector q_batches = vector((int*)query_batches.data(), 28 | (int*)query_batches.data() + Nb); 29 | vector s_batches = vector((int*)support_batches.data(), 30 | (int*)support_batches.data() + Nb); 31 | 32 | // Create result containers 33 | vector neighbors_indices; 34 | 35 | // Compute results 36 | //batch_ordered_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius); 37 | batch_nanoflann_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius); 38 | 39 | // Maximal number of neighbors 40 | int max_neighbors = neighbors_indices.size() / Nq; 41 | 42 | // create output shape 43 | auto output_tensor = torch::zeros({Nq, max_neighbors}); 44 | 45 | // Fill output tensor 46 | // output_tensor = torch::from_blob(&neighbors_indices); 47 | // https://discuss.pytorch.org/t/can-i-initialize-tensor-from-std-vector-in-libtorch/33236 48 | output_tensor = torch::tensor(neighbors_indices); 49 | // for (int i = 0; i < Nq; i++) 50 | // { 51 | // for (int j = 0; j < max_neighbors; j++) 52 | // { 53 | // // std::cout << neighbors_indices[max_neighbors * i + j] << std::endl; 54 | // output_tensor[i][j] = neighbors_indices[max_neighbors * i + j]; 55 | // } 56 | // } 57 | 58 | return output_tensor; 59 | } 60 | 61 | 62 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 63 | m.def("compute", &batch_find_neighbors, "batch find neighbors"); 64 | } -------------------------------------------------------------------------------- /pytorch_ops/cpp_utils/cloud/cloud.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud source : 13 | // Define usefull Functions/Methods 14 | // 15 | //---------------------------------------------------- 16 | // 17 | // Hugues THOMAS - 10/02/2017 18 | // 19 | 20 | 21 | #include "cloud.h" 22 | 23 | 24 | // Getters 25 | // ******* 26 | 27 | PointXYZ max_point(std::vector points) 28 | { 29 | // Initiate limits 30 | PointXYZ maxP(points[0]); 31 | 32 | // Loop over all points 33 | for (auto p : points) 34 | { 35 | if (p.x > maxP.x) 36 | maxP.x = p.x; 37 | 38 | if (p.y > maxP.y) 39 | maxP.y = p.y; 40 | 41 | if (p.z > maxP.z) 42 | maxP.z = p.z; 43 | } 44 | 45 | return maxP; 46 | } 47 | 48 | PointXYZ min_point(std::vector points) 49 | { 50 | // Initiate limits 51 | PointXYZ minP(points[0]); 52 | 53 | // Loop over all points 54 | for (auto p : points) 55 | { 56 | if (p.x < minP.x) 57 | minP.x = p.x; 58 | 59 | if (p.y < minP.y) 60 | minP.y = p.y; 61 | 62 | if (p.z < minP.z) 63 | minP.z = p.z; 64 | } 65 | 66 | return minP; 67 | } -------------------------------------------------------------------------------- /pytorch_ops/cpp_utils/cloud/cloud.h: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud header 13 | // 14 | //---------------------------------------------------- 15 | // 16 | // Hugues THOMAS - 10/02/2017 17 | // 18 | 19 | 20 | # pragma once 21 | 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | #include 31 | 32 | 33 | 34 | 35 | // Point class 36 | // *********** 37 | 38 | 39 | class PointXYZ 40 | { 41 | public: 42 | 43 | // Elements 44 | // ******** 45 | 46 | float x, y, z; 47 | 48 | 49 | // Methods 50 | // ******* 51 | 52 | // Constructor 53 | PointXYZ() { x = 0; y = 0; z = 0; } 54 | PointXYZ(float x0, float y0, float z0) { x = x0; y = y0; z = z0; } 55 | 56 | // array type accessor 57 | float operator [] (int i) const 58 | { 59 | if (i == 0) return x; 60 | else if (i == 1) return y; 61 | else return z; 62 | } 63 | 64 | // opperations 65 | float dot(const PointXYZ P) const 66 | { 67 | return x * P.x + y * P.y + z * P.z; 68 | } 69 | 70 | float sq_norm() 71 | { 72 | return x*x + y*y + z*z; 73 | } 74 | 75 | PointXYZ cross(const PointXYZ P) const 76 | { 77 | return PointXYZ(y*P.z - z*P.y, z*P.x - x*P.z, x*P.y - y*P.x); 78 | } 79 | 80 | PointXYZ& operator+=(const PointXYZ& P) 81 | { 82 | x += P.x; 83 | y += P.y; 84 | z += P.z; 85 | return *this; 86 | } 87 | 88 | PointXYZ& operator-=(const PointXYZ& P) 89 | { 90 | x -= P.x; 91 | y -= P.y; 92 | z -= P.z; 93 | return *this; 94 | } 95 | 96 | PointXYZ& operator*=(const float& a) 97 | { 98 | x *= a; 99 | y *= a; 100 | z *= a; 101 | return *this; 102 | } 103 | }; 104 | 105 | 106 | // Point Opperations 107 | // ***************** 108 | 109 | inline PointXYZ operator + (const PointXYZ A, const PointXYZ B) 110 | { 111 | return PointXYZ(A.x + B.x, A.y + B.y, A.z + B.z); 112 | } 113 | 114 | inline PointXYZ operator - (const PointXYZ A, const PointXYZ B) 115 | { 116 | return PointXYZ(A.x - B.x, A.y - B.y, A.z - B.z); 117 | } 118 | 119 | inline PointXYZ operator * (const PointXYZ P, const float a) 120 | { 121 | return PointXYZ(P.x * a, P.y * a, P.z * a); 122 | } 123 | 124 | inline PointXYZ operator * (const float a, const PointXYZ P) 125 | { 126 | return PointXYZ(P.x * a, P.y * a, P.z * a); 127 | } 128 | 129 | inline std::ostream& operator << (std::ostream& os, const PointXYZ P) 130 | { 131 | return os << "[" << P.x << ", " << P.y << ", " << P.z << "]"; 132 | } 133 | 134 | inline bool operator == (const PointXYZ A, const PointXYZ B) 135 | { 136 | return A.x == B.x && A.y == B.y && A.z == B.z; 137 | } 138 | 139 | inline PointXYZ floor(const PointXYZ P) 140 | { 141 | return PointXYZ(floor(P.x), floor(P.y), floor(P.z)); 142 | } 143 | 144 | 145 | PointXYZ max_point(std::vector points); 146 | PointXYZ min_point(std::vector points); 147 | 148 | struct PointCloud 149 | { 150 | 151 | std::vector pts; 152 | 153 | // Must return the number of data points 154 | inline size_t kdtree_get_point_count() const { return pts.size(); } 155 | 156 | // Returns the dim'th component of the idx'th point in the class: 157 | // Since this is inlined and the "dim" argument is typically an immediate value, the 158 | // "if/else's" are actually solved at compile time. 159 | inline float kdtree_get_pt(const size_t idx, const size_t dim) const 160 | { 161 | if (dim == 0) return pts[idx].x; 162 | else if (dim == 1) return pts[idx].y; 163 | else return pts[idx].z; 164 | } 165 | 166 | // Optional bounding-box computation: return false to default to a standard bbox computation loop. 167 | // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 168 | // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 169 | template 170 | bool kdtree_get_bbox(BBOX& /* bb */) const { return false; } 171 | 172 | }; 173 | -------------------------------------------------------------------------------- /pytorch_ops/neighbors/neighbors.cpp: -------------------------------------------------------------------------------- 1 | #include "neighbors.h" 2 | 3 | 4 | void brute_neighbors(vector& queries, vector& supports, vector& neighbors_indices, float radius, int verbose) 5 | { 6 | 7 | // Initiate variables 8 | // ****************** 9 | 10 | // square radius 11 | float r2 = radius * radius; 12 | 13 | // indices 14 | int i0 = 0; 15 | 16 | // Counting vector 17 | int max_count = 0; 18 | vector> tmp(queries.size()); 19 | 20 | // Search neigbors indices 21 | // *********************** 22 | 23 | for (auto& p0 : queries) 24 | { 25 | int i = 0; 26 | for (auto& p : supports) 27 | { 28 | if ((p0 - p).sq_norm() < r2) 29 | { 30 | tmp[i0].push_back(i); 31 | if (tmp[i0].size() > max_count) 32 | max_count = tmp[i0].size(); 33 | } 34 | i++; 35 | } 36 | i0++; 37 | } 38 | 39 | // Reserve the memory 40 | neighbors_indices.resize(queries.size() * max_count); 41 | i0 = 0; 42 | for (auto& inds : tmp) 43 | { 44 | for (int j = 0; j < max_count; j++) 45 | { 46 | if (j < inds.size()) 47 | neighbors_indices[i0 * max_count + j] = inds[j]; 48 | else 49 | neighbors_indices[i0 * max_count + j] = -1; 50 | } 51 | i0++; 52 | } 53 | 54 | return; 55 | } 56 | 57 | void ordered_neighbors(vector& queries, 58 | vector& supports, 59 | vector& neighbors_indices, 60 | float radius) 61 | { 62 | 63 | // Initiate variables 64 | // ****************** 65 | 66 | // square radius 67 | float r2 = radius * radius; 68 | 69 | // indices 70 | int i0 = 0; 71 | 72 | // Counting vector 73 | int max_count = 0; 74 | float d2; 75 | vector> tmp(queries.size()); 76 | vector> dists(queries.size()); 77 | 78 | // Search neigbors indices 79 | // *********************** 80 | 81 | for (auto& p0 : queries) 82 | { 83 | int i = 0; 84 | for (auto& p : supports) 85 | { 86 | d2 = (p0 - p).sq_norm(); 87 | if (d2 < r2) 88 | { 89 | // Find order of the new point 90 | auto it = std::upper_bound(dists[i0].begin(), dists[i0].end(), d2); 91 | int index = std::distance(dists[i0].begin(), it); 92 | 93 | // Insert element 94 | dists[i0].insert(it, d2); 95 | tmp[i0].insert(tmp[i0].begin() + index, i); 96 | 97 | // Update max count 98 | if (tmp[i0].size() > max_count) 99 | max_count = tmp[i0].size(); 100 | } 101 | i++; 102 | } 103 | i0++; 104 | } 105 | 106 | // Reserve the memory 107 | neighbors_indices.resize(queries.size() * max_count); 108 | i0 = 0; 109 | for (auto& inds : tmp) 110 | { 111 | for (int j = 0; j < max_count; j++) 112 | { 113 | if (j < inds.size()) 114 | neighbors_indices[i0 * max_count + j] = inds[j]; 115 | else 116 | neighbors_indices[i0 * max_count + j] = -1; 117 | } 118 | i0++; 119 | } 120 | 121 | return; 122 | } 123 | 124 | void batch_ordered_neighbors(vector& queries, 125 | vector& supports, 126 | vector& q_batches, 127 | vector& s_batches, 128 | vector& neighbors_indices, 129 | float radius) 130 | { 131 | 132 | // Initiate variables 133 | // ****************** 134 | 135 | // square radius 136 | float r2 = radius * radius; 137 | 138 | // indices 139 | int i0 = 0; 140 | 141 | // Counting vector 142 | int max_count = 0; 143 | float d2; 144 | vector> tmp(queries.size()); 145 | vector> dists(queries.size()); 146 | 147 | // batch index 148 | int b = 0; 149 | int sum_qb = 0; 150 | int sum_sb = 0; 151 | 152 | 153 | // Search neigbors indices 154 | // *********************** 155 | 156 | for (auto& p0 : queries) 157 | { 158 | // Check if we changed batch 159 | if (i0 == sum_qb + q_batches[b]) 160 | { 161 | sum_qb += q_batches[b]; 162 | sum_sb += s_batches[b]; 163 | b++; 164 | } 165 | 166 | // Loop only over the supports of current batch 167 | vector::iterator p_it; 168 | int i = 0; 169 | for(p_it = supports.begin() + sum_sb; p_it < supports.begin() + sum_sb + s_batches[b]; p_it++ ) 170 | { 171 | d2 = (p0 - *p_it).sq_norm(); 172 | if (d2 < r2) 173 | { 174 | // Find order of the new point 175 | auto it = std::upper_bound(dists[i0].begin(), dists[i0].end(), d2); 176 | int index = std::distance(dists[i0].begin(), it); 177 | 178 | // Insert element 179 | dists[i0].insert(it, d2); 180 | tmp[i0].insert(tmp[i0].begin() + index, sum_sb + i); 181 | 182 | // Update max count 183 | if (tmp[i0].size() > max_count) 184 | max_count = tmp[i0].size(); 185 | } 186 | i++; 187 | } 188 | i0++; 189 | } 190 | 191 | // Reserve the memory 192 | neighbors_indices.resize(queries.size() * max_count); 193 | i0 = 0; 194 | for (auto& inds : tmp) 195 | { 196 | for (int j = 0; j < max_count; j++) 197 | { 198 | if (j < inds.size()) 199 | neighbors_indices[i0 * max_count + j] = inds[j]; 200 | else 201 | neighbors_indices[i0 * max_count + j] = supports.size(); 202 | } 203 | i0++; 204 | } 205 | 206 | return; 207 | } 208 | 209 | 210 | void batch_nanoflann_neighbors(vector& queries, 211 | vector& supports, 212 | vector& q_batches, 213 | vector& s_batches, 214 | vector& neighbors_indices, 215 | float radius) 216 | { 217 | 218 | // Initiate variables 219 | // ****************** 220 | 221 | // indices 222 | int i0 = 0; 223 | 224 | // Square radius 225 | float r2 = radius * radius; 226 | 227 | // Counting vector 228 | int max_count = 0; 229 | float d2; 230 | vector>> all_inds_dists(queries.size()); 231 | 232 | // batch index 233 | int b = 0; 234 | int sum_qb = 0; 235 | int sum_sb = 0; 236 | 237 | // Nanoflann related variables 238 | // *************************** 239 | 240 | // CLoud variable 241 | PointCloud current_cloud; 242 | 243 | // Tree parameters 244 | nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */); 245 | 246 | // KDTree type definition 247 | typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Simple_Adaptor , 248 | PointCloud, 249 | 3 > my_kd_tree_t; 250 | 251 | // Pointer to trees 252 | my_kd_tree_t* index; 253 | 254 | // Build KDTree for the first batch element 255 | current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]); 256 | index = new my_kd_tree_t(3, current_cloud, tree_params); 257 | index->buildIndex(); 258 | 259 | 260 | // Search neigbors indices 261 | // *********************** 262 | 263 | // Search params 264 | nanoflann::SearchParams search_params; 265 | search_params.sorted = true; 266 | 267 | for (auto& p0 : queries) 268 | { 269 | 270 | // Check if we changed batch 271 | if (i0 == sum_qb + q_batches[b]) 272 | { 273 | sum_qb += q_batches[b]; 274 | sum_sb += s_batches[b]; 275 | b++; 276 | 277 | // Change the points 278 | current_cloud.pts.clear(); 279 | current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]); 280 | 281 | // Build KDTree of the current element of the batch 282 | delete index; 283 | index = new my_kd_tree_t(3, current_cloud, tree_params); 284 | index->buildIndex(); 285 | } 286 | 287 | // Initial guess of neighbors size 288 | all_inds_dists[i0].reserve(max_count); 289 | 290 | // Find neighbors 291 | float query_pt[3] = { p0.x, p0.y, p0.z}; 292 | size_t nMatches = index->radiusSearch(query_pt, r2, all_inds_dists[i0], search_params); 293 | 294 | // Update max count 295 | if (nMatches > max_count) 296 | max_count = nMatches; 297 | 298 | // Increment query idx 299 | i0++; 300 | } 301 | 302 | // Reserve the memory 303 | neighbors_indices.resize(queries.size() * max_count); 304 | i0 = 0; 305 | sum_sb = 0; 306 | sum_qb = 0; 307 | b = 0; 308 | for (auto& inds_dists : all_inds_dists) 309 | { 310 | // Check if we changed batch 311 | if (i0 == sum_qb + q_batches[b]) 312 | { 313 | sum_qb += q_batches[b]; 314 | sum_sb += s_batches[b]; 315 | b++; 316 | } 317 | 318 | for (int j = 0; j < max_count; j++) 319 | { 320 | if (j < inds_dists.size()) 321 | neighbors_indices[i0 * max_count + j] = inds_dists[j].first + sum_sb; 322 | else 323 | neighbors_indices[i0 * max_count + j] = supports.size(); 324 | } 325 | i0++; 326 | } 327 | 328 | delete index; 329 | 330 | return; 331 | } 332 | 333 | -------------------------------------------------------------------------------- /pytorch_ops/neighbors/neighbors.h: -------------------------------------------------------------------------------- 1 | #include "../cpp_utils/cloud/cloud.h" 2 | #include "../cpp_utils/nanoflann/nanoflann.hpp" 3 | 4 | #include 5 | #include 6 | 7 | using namespace std; 8 | 9 | 10 | void ordered_neighbors(vector& queries, 11 | vector& supports, 12 | vector& neighbors_indices, 13 | float radius); 14 | 15 | void batch_ordered_neighbors(vector& queries, 16 | vector& supports, 17 | vector& q_batches, 18 | vector& s_batches, 19 | vector& neighbors_indices, 20 | float radius); 21 | 22 | void batch_nanoflann_neighbors(vector& queries, 23 | vector& supports, 24 | vector& q_batches, 25 | vector& s_batches, 26 | vector& neighbors_indices, 27 | float radius); 28 | -------------------------------------------------------------------------------- /pytorch_ops/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, Extension 2 | from torch.utils import cpp_extension 3 | 4 | # use the following for macOS 5 | # CC=clang CXX=clang++ CFLAGS='-stdlib=libc++' python setup.py install 6 | 7 | m_name = 'batch_find_neighbors' 8 | setup(name=m_name, 9 | ext_modules=[ 10 | cpp_extension.CppExtension( 11 | name='batch_find_neighbors', 12 | sources=['batch_find_neighbors.cpp', 'neighbors/neighbors.cpp', 'cpp_utils/cloud/cloud.cpp']) 13 | ], 14 | cmdclass={'build_ext': cpp_extension.BuildExtension} 15 | ) 16 | -------------------------------------------------------------------------------- /trainer.py: -------------------------------------------------------------------------------- 1 | import time, os 2 | import numpy as np 3 | from tensorboardX import SummaryWriter 4 | import torch 5 | from utils.timer import Timer, AverageMeter 6 | from utils.metrics import calculate_acc, calculate_iou 7 | 8 | 9 | class Trainer(object): 10 | def __init__(self, args): 11 | self.config = args.config 12 | # parameters 13 | self.start_epoch = args.start_epoch 14 | self.epoch = args.epoch 15 | self.save_dir = args.save_dir 16 | self.result_dir = args.result_dir 17 | self.device = args.device 18 | self.verbose = args.verbose 19 | self.best_iou = 0 20 | self.best_acc = 0 21 | self.best_loss = 10000000 22 | 23 | self.model = args.model.to(self.device) 24 | self.optimizer = args.optimizer 25 | self.scheduler = args.scheduler 26 | self.scheduler_interval = args.scheduler_interval 27 | self.snapshot_interval = args.snapshot_interval 28 | self.evaluate_interval = args.evaluate_interval 29 | self.evaluate_metric = args.evaluate_metric 30 | self.writer = SummaryWriter(log_dir=args.tboard_dir) 31 | # self.writer = SummaryWriter(logdir=args.tboard_dir) 32 | 33 | if args.resume != None: 34 | self._load_pretrain(args.resume) 35 | 36 | self.train_loader = args.train_loader 37 | self.test_loader = args.test_loader 38 | 39 | def train(self): 40 | self.train_hist = { 41 | 'iou': [], 42 | 'loss': [], 43 | 'accuracy': [], 44 | 'per_epoch_time': [], 45 | 'total_time': [] 46 | } 47 | print('training start!!') 48 | start_time = time.time() 49 | 50 | self.model.train() 51 | res = self.evaluate(self.start_epoch) 52 | if self.writer: 53 | self.writer.add_scalar('val/IOU', res['iou'], 0) 54 | self.writer.add_scalar('val/Loss', res['loss'], 0) 55 | self.writer.add_scalar('val/Accuracy', res['accuracy'], 0) 56 | print(f'Evaluation: Epoch 0: Loss {res["loss"]}, Accuracy {res["accuracy"]}, IOU {res["iou"]}') 57 | for epoch in range(self.start_epoch, self.epoch): 58 | self.train_epoch(epoch) 59 | 60 | if (epoch + 1) % self.evaluate_interval == 0 or epoch == 0: 61 | res = self.evaluate(epoch + 1) 62 | print(f'Evaluation: Epoch {epoch + 1}: Loss {res["loss"]}, Accuracy {res["accuracy"]}, IOU {res["iou"]}') 63 | if res['loss'] < self.best_loss: 64 | self.best_loss = res['loss'] 65 | self._snapshot(epoch + 1, 'best_loss') 66 | if res['iou'] > self.best_iou: 67 | self.best_iou = res['iou'] 68 | self._snapshot(epoch + 1, 'best_iou') 69 | 70 | if self.writer: 71 | self.writer.add_scalar('val/IOU', res['iou'], epoch + 1) 72 | self.writer.add_scalar('val/Loss', res['loss'], epoch + 1) 73 | self.writer.add_scalar('val/Accuracy', res['accuracy'], epoch + 1) 74 | 75 | if (epoch + 1) % self.scheduler_interval == 0: 76 | self.scheduler.step() 77 | 78 | if (epoch + 1) % self.snapshot_interval == 0: 79 | self._snapshot(epoch + 1) 80 | 81 | if self.writer: 82 | self.writer.add_scalar('train/Learning Rate', self._get_lr(), epoch + 1) 83 | self.writer.add_scalar('train/IOU', self.train_hist['iou'][-1], epoch + 1) 84 | self.writer.add_scalar('train/Loss', self.train_hist['loss'][-1], epoch + 1) 85 | self.writer.add_scalar('train/Accuracy', self.train_hist['accuracy'][-1], epoch + 1) 86 | 87 | # finish all epoch 88 | self.train_hist['total_time'].append(time.time() - start_time) 89 | print("Avg one epoch time: %.2f, total %d epochs time: %.2f" % (np.mean(self.train_hist['per_epoch_time']), 90 | self.epoch, self.train_hist['total_time'][0])) 91 | print("Training finish!... save training results") 92 | 93 | def train_epoch(self, epoch): 94 | data_timer, model_timer = Timer(), Timer() 95 | loss_meter, acc_meter, iou_meter = AverageMeter(), AverageMeter(), AverageMeter() 96 | num_iter = int(len(self.train_loader.dataset) // self.train_loader.batch_size) 97 | train_loader_iter = self.train_loader.__iter__() 98 | # for iter, inputs in enumerate(self.train_loader): 99 | for iter in range(num_iter): 100 | data_timer.tic() 101 | inputs = train_loader_iter.next() 102 | for k, v in inputs.items(): # load inputs to device. 103 | if type(v) == list: 104 | inputs[k] = [item.to(self.device) for item in v] 105 | else: 106 | inputs[k] = v.to(self.device) 107 | data_timer.toc() 108 | 109 | model_timer.tic() 110 | # forward 111 | self.optimizer.zero_grad() 112 | predict = self.model(inputs) 113 | labels = inputs['labels'].long().squeeze() 114 | loss = self.evaluate_metric(predict, labels) 115 | # acc = torch.sum(torch.max(predict, dim=1)[1].int() == labels.int()) * 100 / predict.shape[0] 116 | acc = calculate_acc(predict, labels) 117 | part_iou = calculate_iou(predict, labels, stack_lengths=inputs['stack_lengths'][0], n_parts=self.config.num_classes) 118 | iou = np.mean(part_iou) 119 | 120 | # backward 121 | loss.backward() 122 | self.optimizer.step() 123 | model_timer.toc() 124 | loss_meter.update(float(loss)) 125 | iou_meter.update(float(iou)) 126 | acc_meter.update(float(acc)) 127 | 128 | if (iter + 1) % 1 == 0 and self.verbose: 129 | print(f"Epoch: {epoch+1} [{iter+1:4d}/{num_iter}] " 130 | f"loss: {loss_meter.avg:.2f} " 131 | f"iou: {iou_meter.avg:.2f} " 132 | f"data time: {data_timer.avg:.2f}s " 133 | f"model time: {model_timer.avg:.2f}s") 134 | # finish one epoch 135 | epoch_time = model_timer.total_time + data_timer.total_time 136 | self.train_hist['per_epoch_time'].append(epoch_time) 137 | self.train_hist['loss'].append(loss_meter.avg) 138 | self.train_hist['accuracy'].append(acc_meter.avg) 139 | self.train_hist['iou'].append(iou_meter.avg) 140 | print(f'Epoch {epoch+1}: Loss : {loss_meter.avg:.2f}, Accuracy: {acc_meter.avg:.2f}, IOU: {iou_meter.avg:.2f}, time {epoch_time:.2f}s') 141 | 142 | def evaluate(self, epoch): 143 | data_timer, model_timer = Timer(), Timer() 144 | loss_meter, acc_meter, iou_meter = AverageMeter(), AverageMeter(), AverageMeter() 145 | num_iter = int(len(self.test_loader.dataset) // self.train_loader.batch_size) 146 | test_loader_iter = self.test_loader.__iter__() 147 | for iter in range(num_iter): 148 | data_timer.tic() 149 | inputs = test_loader_iter.next() 150 | for k, v in inputs.items(): # load inputs to device. 151 | if type(v) == list: 152 | inputs[k] = [item.to(self.device) for item in v] 153 | else: 154 | inputs[k] = v.to(self.device) 155 | data_timer.toc() 156 | 157 | model_timer.tic() 158 | predict = self.model(inputs) 159 | labels = inputs['labels'].long().squeeze() 160 | loss = self.evaluate_metric(predict, labels) 161 | # acc = torch.sum(torch.max(predict, dim=1)[1].int() == labels.int()) * 100 / predict.shape[0] 162 | acc = calculate_acc(predict, labels) 163 | part_iou = calculate_iou(predict, labels, stack_lengths=inputs['stack_lengths'][0], n_parts=self.config.num_classes) 164 | iou = np.mean(part_iou) 165 | 166 | model_timer.toc() 167 | loss_meter.update(float(loss)) 168 | iou_meter.update(float(iou)) 169 | acc_meter.update(float(acc)) 170 | 171 | if (iter + 1) % 1 == 0 and self.verbose: 172 | print(f"Eval epoch: {epoch+1} [{iter+1:4d}/{num_iter}] " 173 | f"loss: {loss_meter.avg:.2f} " 174 | f"iou: {iou_meter.avg:.2f} " 175 | f"data time: {data_timer.avg:.2f}s " 176 | f"model time: {model_timer.avg:.2f}s") 177 | self.model.train() 178 | res = { 179 | 'iou': iou_meter.avg, 180 | 'loss': loss_meter.avg, 181 | 'accuracy': acc_meter.avg 182 | } 183 | return res 184 | 185 | def _snapshot(self, epoch, name=None): 186 | state = { 187 | 'epoch': epoch, 188 | 'state_dict': self.model.state_dict(), 189 | 'optimizer': self.optimizer.state_dict(), 190 | 'scheduler': self.scheduler.state_dict(), 191 | 'best_loss': self.best_loss, 192 | 'best_iou': self.best_iou, 193 | } 194 | if name is None: 195 | filename = os.path.join(self.save_dir, f'model_{epoch}.pth') 196 | else: 197 | filename = os.path.join(self.save_dir, f'model_{name}.pth') 198 | print(f"Save model to {filename}") 199 | torch.save(state, filename) 200 | 201 | def _load_pretrain(self, resume): 202 | if os.path.isfile(resume): 203 | print(f"=> loading checkpoint {resume}") 204 | state = torch.load(resume) 205 | self.start_epoch = state['epoch'] 206 | self.model.load_state_dict(state['state_dict']) 207 | self.scheduler.load_state_dict(state['scheduler']) 208 | self.optimizer.load_state_dict(state['optimizer']) 209 | self.best_loss = state['best_loss'] 210 | self.best_acc = state['best_acc'] 211 | else: 212 | raise ValueError(f"=> no checkpoint found at '{resume}'") 213 | 214 | def _get_lr(self, group=0): 215 | return self.optimizer.param_groups[group]['lr'] 216 | -------------------------------------------------------------------------------- /trainer_cls.py: -------------------------------------------------------------------------------- 1 | import time, os 2 | import numpy as np 3 | from tensorboardX import SummaryWriter 4 | import torch 5 | from utils.timer import Timer, AverageMeter 6 | 7 | 8 | class Trainer(object): 9 | def __init__(self, args): 10 | # parameters 11 | self.start_epoch = args.start_epoch 12 | self.epoch = args.epoch 13 | self.save_dir = args.save_dir 14 | self.result_dir = args.result_dir 15 | self.device = args.device 16 | self.verbose = args.verbose 17 | self.best_acc = 0 18 | self.best_loss = 10000000 19 | 20 | self.model = args.model.to(self.device) 21 | self.optimizer = args.optimizer 22 | self.scheduler = args.scheduler 23 | self.scheduler_interval = args.scheduler_interval 24 | self.snapshot_interval = args.snapshot_interval 25 | self.evaluate_interval = args.evaluate_interval 26 | self.evaluate_metric = args.evaluate_metric 27 | self.writer = SummaryWriter(log_dir=args.tboard_dir) 28 | # self.writer = SummaryWriter(logdir=args.tboard_dir) 29 | 30 | if args.resume != None: 31 | self._load_pretrain(args.resume) 32 | 33 | self.train_loader = args.train_loader 34 | self.test_loader = args.test_loader 35 | 36 | def train(self): 37 | self.train_hist = { 38 | 'loss': [], 39 | 'accuracy': [], 40 | 'per_epoch_time': [], 41 | 'total_time': [] 42 | } 43 | print('training start!!') 44 | start_time = time.time() 45 | 46 | self.model.train() 47 | res = self.evaluate(self.start_epoch) 48 | if self.writer: 49 | self.writer.add_scalar('val/Loss', res['loss'], 0) 50 | self.writer.add_scalar('val/Accuracy', res['accuracy'], 0) 51 | print(f'Evaluation: Epoch 0: Loss {res["loss"]}, Accuracy {res["accuracy"]}') 52 | for epoch in range(self.start_epoch, self.epoch): 53 | self.train_epoch(epoch) 54 | 55 | if (epoch + 1) % self.evaluate_interval == 0 or epoch == 0: 56 | res = self.evaluate(epoch + 1) 57 | print(f'Evaluation: Epoch {epoch + 1}: Loss {res["loss"]}, Accuracy {res["accuracy"]}') 58 | if res['loss'] < self.best_loss: 59 | self.best_loss = res['loss'] 60 | self._snapshot(epoch + 1, 'best') 61 | if res['accuracy'] > self.best_acc: 62 | self.best_acc = res['accuracy'] 63 | self._snapshot(epoch + 1, 'best_acc') 64 | 65 | if self.writer: 66 | self.writer.add_scalar('val/Loss', res['loss'], epoch + 1) 67 | self.writer.add_scalar('val/Accuracy', res['accuracy'], epoch + 1) 68 | 69 | if (epoch + 1) % self.scheduler_interval == 0: 70 | self.scheduler.step() 71 | 72 | if (epoch + 1) % self.snapshot_interval == 0: 73 | self._snapshot(epoch + 1) 74 | 75 | if self.writer: 76 | self.writer.add_scalar('train/Learning Rate', self._get_lr(), epoch + 1) 77 | self.writer.add_scalar('train/Loss', self.train_hist['loss'][-1], epoch + 1) 78 | self.writer.add_scalar('train/Accuracy', self.train_hist['accuracy'][-1], epoch + 1) 79 | 80 | # finish all epoch 81 | self.train_hist['total_time'].append(time.time() - start_time) 82 | print("Avg one epoch time: %.2f, total %d epochs time: %.2f" % (np.mean(self.train_hist['per_epoch_time']), 83 | self.epoch, self.train_hist['total_time'][0])) 84 | print("Training finish!... save training results") 85 | 86 | def train_epoch(self, epoch): 87 | data_timer, model_timer = Timer(), Timer() 88 | loss_meter, acc_meter = AverageMeter(), AverageMeter() 89 | num_iter = int(len(self.train_loader.dataset) // self.train_loader.batch_size) 90 | train_loader_iter = self.train_loader.__iter__() 91 | # for iter, inputs in enumerate(self.train_loader): 92 | for iter in range(num_iter): 93 | data_timer.tic() 94 | inputs = train_loader_iter.next() 95 | for k, v in inputs.items(): # load inputs to device. 96 | if type(v) == list: 97 | inputs[k] = [item.to(self.device) for item in v] 98 | else: 99 | inputs[k] = v.to(self.device) 100 | data_timer.toc() 101 | 102 | model_timer.tic() 103 | # forward 104 | self.optimizer.zero_grad() 105 | predict = self.model(inputs) 106 | labels = inputs['labels'].long() 107 | loss = self.evaluate_metric(predict, labels) 108 | acc = torch.sum(torch.max(predict, dim=1)[1].int() == labels.int()) * 100 / predict.shape[0] 109 | 110 | # backward 111 | loss.backward() 112 | self.optimizer.step() 113 | model_timer.toc() 114 | loss_meter.update(float(loss)) 115 | acc_meter.update(float(acc)) 116 | 117 | if (iter + 1) % 1 == 0 and self.verbose: 118 | print(f"Epoch: {epoch+1} [{iter+1:4d}/{num_iter}] " 119 | f"loss: {loss_meter.avg:.2f} " 120 | f"acc: {acc_meter.avg:.2f} " 121 | f"data time: {data_timer.avg:.2f}s " 122 | f"model time: {model_timer.avg:.2f}s") 123 | # finish one epoch 124 | epoch_time = model_timer.total_time + data_timer.total_time 125 | self.train_hist['per_epoch_time'].append(epoch_time) 126 | self.train_hist['loss'].append(loss_meter.avg) 127 | self.train_hist['accuracy'].append(acc_meter.avg) 128 | print(f'Epoch {epoch+1}: Loss : {loss_meter.avg:.2f}, Accuracy: {acc_meter.avg:.2f} , Total time {epoch_time:.2f}s') 129 | 130 | def evaluate(self, epoch): 131 | self.model.eval() 132 | data_timer, model_timer = Timer(), Timer() 133 | loss_meter, acc_meter = AverageMeter(), AverageMeter() 134 | num_iter = int(len(self.test_loader.dataset) / self.test_loader.batch_size) 135 | test_loader_iter = self.test_loader.__iter__() 136 | for iter in range(num_iter): 137 | data_timer.tic() 138 | inputs = test_loader_iter.next() 139 | for k, v in inputs.items(): # load inputs to device. 140 | if type(v) == list: 141 | inputs[k] = [item.to(self.device) for item in v] 142 | else: 143 | inputs[k] = v.to(self.device) 144 | data_timer.toc() 145 | 146 | model_timer.tic() 147 | predict = self.model(inputs) 148 | labels = inputs['labels'].long() 149 | loss = self.evaluate_metric(predict, labels) 150 | acc = torch.sum(torch.max(predict, dim=1)[1].int() == labels.int()) * 100 / predict.shape[0] 151 | model_timer.toc() 152 | 153 | loss_meter.update(float(loss)) 154 | acc_meter.update(float(acc)) 155 | if (iter + 1) % 1 == 0 and self.verbose: 156 | print(f"Eval epoch {epoch+1}: [{iter+1:3d}/{num_iter}] " 157 | f"loss: {loss_meter.avg:.2f} " 158 | f"acc: {acc_meter.avg:.2f} " 159 | f"data time: {data_timer.avg:.2f}s " 160 | f"model time: {model_timer.avg:.2f}s") 161 | self.model.train() 162 | res = { 163 | 'loss': loss_meter.avg, 164 | 'accuracy': acc_meter.avg 165 | } 166 | return res 167 | 168 | def _snapshot(self, epoch, name=None): 169 | state = { 170 | 'epoch': epoch, 171 | 'state_dict': self.model.state_dict(), 172 | 'optimizer': self.optimizer.state_dict(), 173 | 'scheduler': self.scheduler.state_dict(), 174 | 'best_loss': self.best_loss, 175 | 'best_acc': self.best_acc, 176 | } 177 | if name is None: 178 | filename = os.path.join(self.save_dir, f'model_{epoch}.pth') 179 | else: 180 | filename = os.path.join(self.save_dir, f'model_{name}.pth') 181 | print(f"Save model to {filename}") 182 | torch.save(state, filename) 183 | 184 | def _load_pretrain(self, resume): 185 | if os.path.isfile(resume): 186 | print(f"=> loading checkpoint {resume}") 187 | state = torch.load(resume) 188 | self.start_epoch = state['epoch'] 189 | self.model.load_state_dict(state['state_dict']) 190 | self.scheduler.load_state_dict(state['scheduler']) 191 | self.optimizer.load_state_dict(state['optimizer']) 192 | self.best_loss = state['best_loss'] 193 | self.best_acc = state['best_acc'] 194 | else: 195 | raise ValueError(f"=> no checkpoint found at '{resume}'") 196 | 197 | def _get_lr(self, group=0): 198 | return self.optimizer.param_groups[group]['lr'] 199 | -------------------------------------------------------------------------------- /training_ModelNet.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import shutil 4 | from datasets.ModelNet import ModelNetDataset 5 | from utils.config import Config 6 | from trainer_cls import Trainer 7 | from models.KPCNN import KPCNN 8 | from datasets.dataloader import get_dataloader 9 | from torch import optim 10 | from torch import nn 11 | import torch 12 | 13 | 14 | class ModelNetConfig(Config): 15 | # dataset 16 | dataset = 'ModelNet' 17 | num_classes = 40 18 | first_subsampling_dl = 0.02 19 | in_features_dim = 4 20 | data_train_dir = "./data/modelnet40_normal_resampled/" 21 | data_test_dir = "./data/modelnet40_normal_resampled/" 22 | train_batch_size = 8 23 | test_batch_size = 8 24 | 25 | # model 26 | architecture = ['simple', 27 | 'resnetb', 28 | 'resnetb_strided', 29 | 'resnetb', 30 | 'resnetb_strided', 31 | 'resnetb', 32 | 'resnetb_strided', 33 | 'resnetb', 34 | 'resnetb_strided', 35 | 'resnetb', 36 | 'global_average' 37 | ] 38 | dropout = 0.5 39 | resume = None 40 | use_batch_norm = True 41 | batch_norm_momentum = 0.02 42 | # https://github.com/pytorch/examples/issues/289 pytorch bn momentum 0.02 == tensorflow bn momentum 0.98 43 | 44 | # kernel point convolution 45 | KP_influence = 'linear' 46 | KP_extent = 1.0 47 | convolution_mode = 'sum' 48 | 49 | # training 50 | max_epoch = 200 51 | learning_rate = 5e-3 52 | momentum = 0.98 53 | exp_gamma = 0.1 ** (1 / 80) 54 | exp_interval = 1 55 | 56 | 57 | class Args(object): 58 | def __init__(self, config): 59 | is_test = False 60 | if is_test: 61 | self.experiment_id = "KPConvNet" + time.strftime('%m%d%H%M') + 'Test' 62 | else: 63 | self.experiment_id = "KPConvNet" + time.strftime('%m%d%H%M') 64 | 65 | self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 66 | self.verbose = True 67 | 68 | # snapshot 69 | self.snapshot_interval = 5 70 | snapshot_root = f'snapshot/{config.dataset}_{self.experiment_id}' 71 | tensorboard_root = f'tensorboard/{config.dataset}_{self.experiment_id}' 72 | os.makedirs(snapshot_root, exist_ok=True) 73 | os.makedirs(tensorboard_root, exist_ok=True) 74 | shutil.copy2(os.path.join('.', 'training_ModelNet.py'), os.path.join(snapshot_root, 'train.py')) 75 | shutil.copy2(os.path.join('datasets', 'ModelNet.py'), os.path.join(snapshot_root, 'dataset.py')) 76 | shutil.copy2(os.path.join('datasets', 'dataloader.py'), os.path.join(snapshot_root, 'dataloader.py')) 77 | self.save_dir = os.path.join(snapshot_root, 'models/') 78 | self.result_dir = os.path.join(snapshot_root, 'results/') 79 | self.tboard_dir = tensorboard_root 80 | 81 | # dataset & dataloader 82 | self.train_set = ModelNetDataset(root=config.data_train_dir, 83 | split='train', 84 | first_subsampling_dl=config.first_subsampling_dl, 85 | config=config, 86 | ) 87 | self.test_set = ModelNetDataset(root=config.data_test_dir, 88 | split='test', 89 | first_subsampling_dl=config.first_subsampling_dl, 90 | config=config, 91 | ) 92 | self.train_loader = get_dataloader(dataset=self.train_set, 93 | batch_size=config.train_batch_size, 94 | shuffle=True, 95 | num_workers=config.train_batch_size, 96 | ) 97 | self.test_loader = get_dataloader(dataset=self.test_set, 98 | batch_size=config.test_batch_size, 99 | shuffle=False, 100 | num_workers=config.test_batch_size, 101 | ) 102 | print("Training set size:", self.train_loader.dataset.__len__()) 103 | print("Test set size:", self.test_loader.dataset.__len__()) 104 | 105 | # model 106 | self.model = KPCNN(config) 107 | self.resume = config.resume 108 | # optimizer 109 | self.start_epoch = 0 110 | self.epoch = config.max_epoch 111 | self.optimizer = optim.SGD(self.model.parameters(), lr=config.learning_rate, momentum=config.momentum, weight_decay=1e-6) 112 | self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=config.exp_gamma) 113 | self.scheduler_interval = config.exp_interval 114 | 115 | # evaluate 116 | self.evaluate_interval = 1 117 | self.evaluate_metric = nn.CrossEntropyLoss(reduction='mean') 118 | 119 | self.check_args() 120 | 121 | def check_args(self): 122 | """checking arguments""" 123 | if not os.path.exists(self.save_dir): 124 | os.makedirs(self.save_dir) 125 | if not os.path.exists(self.result_dir): 126 | os.makedirs(self.result_dir) 127 | if not os.path.exists(self.tboard_dir): 128 | os.makedirs(self.tboard_dir) 129 | return self 130 | 131 | 132 | if __name__ == '__main__': 133 | config = ModelNetConfig() 134 | args = Args(config) 135 | trainer = Trainer(args) 136 | trainer.train() 137 | -------------------------------------------------------------------------------- /training_ShapeNetCls.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import shutil 4 | from datasets.ShapeNet import ShapeNetDataset 5 | from utils.config import Config 6 | from trainer_cls import Trainer 7 | from models.KPCNN import KPCNN 8 | from datasets.dataloader import get_dataloader 9 | from torch import optim 10 | from torch import nn 11 | import torch 12 | 13 | 14 | class ShapeNetPartConfig(Config): 15 | # dataset 16 | dataset = 'ShapeNetCls' 17 | num_classes = 16 18 | first_subsampling_dl = 0.02 19 | in_features_dim = 4 20 | data_train_dir = "./data/shapenetcore_partanno_segmentation_benchmark_v0" 21 | data_test_dir = "./data/shapenetcore_partanno_segmentation_benchmark_v0" 22 | train_batch_size = 16 23 | test_batch_size = 16 24 | 25 | # model 26 | architecture = ['simple', 27 | 'resnetb', 28 | 'resnetb_strided', 29 | 'resnetb', 30 | 'resnetb_strided', 31 | 'resnetb', 32 | 'resnetb_strided', 33 | 'resnetb', 34 | 'resnetb_strided', 35 | 'resnetb', 36 | 'global_average' 37 | ] 38 | dropout = 0.5 39 | resume = None 40 | use_batch_norm = True 41 | batch_norm_momentum = 0.02 42 | # https://github.com/pytorch/examples/issues/289 pytorch bn momentum 0.02 == tensorflow bn momentum 0.98 43 | 44 | # kernel point convolution 45 | KP_influence = 'linear' 46 | KP_extent = 1.0 47 | convolution_mode = 'sum' 48 | modulated = False 49 | 50 | # training 51 | max_epoch = 200 52 | learning_rate = 1e-2 53 | momentum = 0.98 54 | exp_gamma = 0.1 ** (1 / 80) 55 | exp_interval = 1 56 | 57 | 58 | class Args(object): 59 | def __init__(self, config): 60 | is_test = False 61 | if is_test: 62 | self.experiment_id = "KPConvNet" + time.strftime('%m%d%H%M') + 'Test' 63 | else: 64 | self.experiment_id = "KPConvNet" + time.strftime('%m%d%H%M') 65 | 66 | self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 67 | self.verbose = True 68 | 69 | # snapshot 70 | self.snapshot_interval = 5 71 | snapshot_root = f'snapshot/{config.dataset}_{self.experiment_id}' 72 | tensorboard_root = f'tensorboard/{config.dataset}_{self.experiment_id}' 73 | os.makedirs(snapshot_root, exist_ok=True) 74 | os.makedirs(tensorboard_root, exist_ok=True) 75 | shutil.copy2(os.path.join('.', 'training_ShapeNetCls.py'), os.path.join(snapshot_root, 'train.py')) 76 | shutil.copy2(os.path.join('datasets', 'ShapeNet.py'), os.path.join(snapshot_root, 'dataset.py')) 77 | shutil.copy2(os.path.join('datasets', 'dataloader.py'), os.path.join(snapshot_root, 'dataloader.py')) 78 | self.save_dir = os.path.join(snapshot_root, 'models/') 79 | self.result_dir = os.path.join(snapshot_root, 'results/') 80 | self.tboard_dir = tensorboard_root 81 | 82 | # dataset & dataloader 83 | self.train_set = ShapeNetDataset(root=config.data_train_dir, 84 | split='train', 85 | first_subsampling_dl=config.first_subsampling_dl, 86 | classification=True, 87 | config=config, 88 | ) 89 | self.test_set = ShapeNetDataset(root=config.data_test_dir, 90 | split='test', 91 | first_subsampling_dl=config.first_subsampling_dl, 92 | classification=True, 93 | config=config, 94 | ) 95 | self.train_loader = get_dataloader(dataset=self.train_set, 96 | batch_size=config.train_batch_size, 97 | shuffle=True, 98 | num_workers=4, 99 | ) 100 | self.test_loader = get_dataloader(dataset=self.test_set, 101 | batch_size=config.test_batch_size, 102 | shuffle=False, 103 | num_workers=4, 104 | ) 105 | print("Training set size:", self.train_loader.dataset.__len__()) 106 | print("Test set size:", self.test_loader.dataset.__len__()) 107 | 108 | # model 109 | self.model = KPCNN(config) 110 | self.resume = config.resume 111 | # optimizer 112 | self.start_epoch = 0 113 | self.epoch = config.max_epoch 114 | self.optimizer = optim.SGD(self.model.parameters(), lr=config.learning_rate, momentum=config.momentum, weight_decay=1e-6) 115 | self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=config.exp_gamma) 116 | self.scheduler_interval = config.exp_interval 117 | 118 | # evaluate 119 | self.evaluate_interval = 1 120 | self.evaluate_metric = nn.CrossEntropyLoss(reduction='mean') 121 | 122 | self.check_args() 123 | 124 | def check_args(self): 125 | """checking arguments""" 126 | if not os.path.exists(self.save_dir): 127 | os.makedirs(self.save_dir) 128 | if not os.path.exists(self.result_dir): 129 | os.makedirs(self.result_dir) 130 | if not os.path.exists(self.tboard_dir): 131 | os.makedirs(self.tboard_dir) 132 | return self 133 | 134 | 135 | if __name__ == '__main__': 136 | config = ShapeNetPartConfig() 137 | args = Args(config) 138 | trainer = Trainer(args) 139 | trainer.train() 140 | -------------------------------------------------------------------------------- /training_ShapeNetPart.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import shutil 4 | from datasets.ShapeNet import ShapeNetDataset 5 | from utils.config import Config 6 | from trainer import Trainer 7 | from models.KPFCNN import KPFCNN 8 | from datasets.dataloader import get_dataloader 9 | from torch import optim 10 | from torch import nn 11 | import torch 12 | 13 | 14 | class ShapeNetPartConfig(Config): 15 | # dataset 16 | dataset = 'ShapeNetPart' 17 | num_classes = 4 18 | first_subsampling_dl = 0.02 19 | in_features_dim = 4 20 | data_train_dir = "./data/shapenetcore_partanno_segmentation_benchmark_v0" 21 | data_test_dir = "./data/shapenetcore_partanno_segmentation_benchmark_v0" 22 | train_batch_size = 16 23 | test_batch_size = 16 24 | 25 | # model 26 | architecture = ['simple', 27 | 'resnetb', 28 | 'resnetb_strided', 29 | 'resnetb', 30 | 'resnetb_strided', 31 | 'resnetb_deformable', 32 | 'resnetb_deformable_strided', 33 | 'resnetb_deformable', 34 | 'resnetb_deformable_strided', 35 | 'resnetb_deformable', 36 | 'nearest_upsample', 37 | 'unary', 38 | 'nearest_upsample', 39 | 'unary', 40 | 'nearest_upsample', 41 | 'unary', 42 | 'nearest_upsample', 43 | 'unary'] 44 | dropout = 0.5 45 | resume = None 46 | use_batch_norm = True 47 | batch_norm_momentum = 0.02 48 | # https://github.com/pytorch/examples/issues/289 pytorch bn momentum 0.02 == tensorflow bn momentum 0.98 49 | 50 | # kernel point convolution 51 | KP_influence = 'linear' 52 | KP_extent = 1.0 53 | convolution_mode = 'sum' 54 | 55 | # training 56 | max_epoch = 200 57 | learning_rate = 1e-2 58 | momentum = 0.98 59 | exp_gamma = 0.1 ** (1 / 80) 60 | exp_interval = 1 61 | 62 | 63 | class Args(object): 64 | def __init__(self, config): 65 | is_test = False 66 | if is_test: 67 | self.experiment_id = "KPConvNet" + time.strftime('%m%d%H%M') + 'Test' 68 | else: 69 | self.experiment_id = "KPConvNet" + time.strftime('%m%d%H%M') 70 | 71 | self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 72 | self.verbose = True 73 | self.config = config 74 | 75 | # snapshot 76 | self.snapshot_interval = 5 77 | snapshot_root = f'snapshot/{config.dataset}_{self.experiment_id}' 78 | tensorboard_root = f'tensorboard/{config.dataset}_{self.experiment_id}' 79 | os.makedirs(snapshot_root, exist_ok=True) 80 | os.makedirs(tensorboard_root, exist_ok=True) 81 | shutil.copy2(os.path.join('.', 'training_ShapeNetPart.py'), os.path.join(snapshot_root, 'train.py')) 82 | shutil.copy2(os.path.join('datasets', 'ShapeNet.py'), os.path.join(snapshot_root, 'dataset.py')) 83 | shutil.copy2(os.path.join('datasets', 'dataloader.py'), os.path.join(snapshot_root, 'dataloader.py')) 84 | self.save_dir = os.path.join(snapshot_root, 'models/') 85 | self.result_dir = os.path.join(snapshot_root, 'results/') 86 | self.tboard_dir = tensorboard_root 87 | 88 | # dataset & dataloader 89 | self.train_set = ShapeNetDataset(root=config.data_train_dir, 90 | split='train', 91 | first_subsampling_dl=config.first_subsampling_dl, 92 | classification=False, 93 | class_choice=['Chair'], 94 | config=config, 95 | ) 96 | self.test_set = ShapeNetDataset(root=config.data_test_dir, 97 | split='test', 98 | first_subsampling_dl=config.first_subsampling_dl, 99 | classification=False, 100 | class_choice=['Chair'], 101 | config=config, 102 | ) 103 | self.train_loader = get_dataloader(dataset=self.train_set, 104 | batch_size=config.train_batch_size, 105 | shuffle=True, 106 | num_workers=config.train_batch_size, 107 | ) 108 | self.test_loader = get_dataloader(dataset=self.test_set, 109 | batch_size=config.test_batch_size, 110 | shuffle=False, 111 | num_workers=config.test_batch_size, 112 | ) 113 | print("Training set size:", self.train_loader.dataset.__len__()) 114 | print("Test set size:", self.test_loader.dataset.__len__()) 115 | 116 | # model 117 | self.model = KPFCNN(config) 118 | self.resume = config.resume 119 | # optimizer 120 | self.start_epoch = 0 121 | self.epoch = config.max_epoch 122 | self.optimizer = optim.SGD(self.model.parameters(), lr=config.learning_rate, momentum=config.momentum, weight_decay=1e-6) 123 | self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=config.exp_gamma) 124 | self.scheduler_interval = config.exp_interval 125 | 126 | # evaluate 127 | self.evaluate_interval = 1 128 | self.evaluate_metric = nn.CrossEntropyLoss(reduction='mean') 129 | 130 | self.check_args() 131 | 132 | def check_args(self): 133 | """checking arguments""" 134 | if not os.path.exists(self.save_dir): 135 | os.makedirs(self.save_dir) 136 | if not os.path.exists(self.result_dir): 137 | os.makedirs(self.result_dir) 138 | if not os.path.exists(self.tboard_dir): 139 | os.makedirs(self.tboard_dir) 140 | return self 141 | 142 | 143 | if __name__ == '__main__': 144 | config = ShapeNetPartConfig() 145 | args = Args(config) 146 | trainer = Trainer(args) 147 | trainer.train() 148 | -------------------------------------------------------------------------------- /utils/config.py: -------------------------------------------------------------------------------- 1 | # 2 | # 3 | # 0=================================0 4 | # | Kernel Point Convolutions | 5 | # 0=================================0 6 | # 7 | # 8 | # ---------------------------------------------------------------------------------------------------------------------- 9 | # 10 | # Configuration class 11 | # 12 | # ---------------------------------------------------------------------------------------------------------------------- 13 | # 14 | # Hugues THOMAS - 11/06/2018 15 | # 16 | 17 | 18 | from os.path import join 19 | 20 | 21 | class Config: 22 | """ 23 | Class containing the parameters you want to modify for this dataset 24 | """ 25 | 26 | ################## 27 | # Input parameters 28 | ################## 29 | # Dataset name 30 | dataset = '' 31 | 32 | # Type of network model 33 | network_model = '' 34 | 35 | # Number of classes in the dataset 36 | num_classes = 0 37 | 38 | # Dimension of input points 39 | in_points_dim = 3 40 | 41 | # Dimension of input features 42 | in_features_dim = 1 43 | 44 | # Radius of the input sphere (ignored for models, only used for point clouds) 45 | in_radius = 1.0 46 | 47 | # Number of CPU threads for the input pipeline 48 | input_threads = 8 49 | 50 | ################## 51 | # Model parameters 52 | ################## 53 | 54 | # Architecture definition. List of blocks 55 | architecture = [] 56 | 57 | # Dimension of the first feature maps 58 | first_features_dim = 64 59 | 60 | # Batch normalization parameters 61 | use_batch_norm = True 62 | batch_norm_momentum = 0.99 63 | 64 | # For segmentation models : ratio between the segmented area and the input area 65 | segmentation_ratio = 1.0 66 | 67 | ################### 68 | # KPConv parameters 69 | ################### 70 | 71 | # First size of grid used for subsampling 72 | first_subsampling_dl = 0.02 73 | 74 | # Radius of the kernels in the first layer (deprecated) 75 | first_kernel_radius = 0.1 76 | 77 | # Number of points in the kernels 78 | num_kernel_points = 15 79 | 80 | # density of neighbors in kernel range 81 | # For each layer, support points are subsampled on a grid with dl = kernel_radius / density_parameter 82 | density_parameter = 3.0 83 | 84 | # Kernel point influence radius 85 | KP_extent = 1.0 86 | 87 | # Influence function when d < KP_extent. ('constant', 'linear', 'gaussian') When d > KP_extent, always zero 88 | KP_influence = 'gaussian' 89 | 90 | # Behavior of convolutions in ('closest', 'sum') 91 | # Decide if you sum all kernel point influences, or if you only take the influence of the closest KP 92 | convolution_mode = 'closest' 93 | 94 | # Fixed points in the kernel : 'none', 'center' or 'verticals' 95 | fixed_kernel_points = 'center' 96 | 97 | # Can the network learn kernel dispositions (deprecated) 98 | trainable_positions = False 99 | 100 | # Use modulateion in deformable convolutions 101 | modulated = False 102 | 103 | ##################### 104 | # Training parameters 105 | ##################### 106 | 107 | # Network optimizer parameters (learning rate and momentum) 108 | learning_rate = 1e-4 109 | momentum = 0.9 110 | 111 | # Learning rate decays. Dictionary of all decay values with their epoch {epoch: decay}. 112 | lr_decays = {200: 0.2, 300: 0.2} 113 | 114 | # Gradient clipping value (negative means no clipping) 115 | grad_clip_norm = 100.0 116 | 117 | # Augmentation parameters 118 | augment_scale_anisotropic = True 119 | augment_scale_min = 0.9 120 | augment_scale_max = 1.1 121 | augment_symmetries = [False, False, False] 122 | augment_rotation = 'vertical' 123 | augment_noise = 0.005 124 | augment_occlusion = 'planar' 125 | augment_occlusion_ratio = 0.2 126 | augment_occlusion_num = 1 127 | augment_color = 0.7 128 | augment_shift_range = 0 129 | 130 | # Regularization loss importance 131 | weights_decay = 1e-6 132 | 133 | # Gaussian loss 134 | gaussian_decay = 1e-3 135 | 136 | # Type of output loss with regard to batches when segmentation 137 | batch_averaged_loss = False 138 | 139 | # Point loss DPRECATED 140 | points_loss = '' 141 | points_decay = 1e-2 142 | 143 | # Offset regularization loss 144 | offsets_loss = 'permissive' 145 | offsets_decay = 1e-2 146 | 147 | # Number of batch 148 | batch_num = 10 149 | 150 | # Maximal number of epochs 151 | max_epoch = 1000 152 | 153 | # Number of steps per epochs 154 | epoch_steps = 1000 155 | 156 | # Number of validation examples per epoch 157 | validation_size = 100 158 | 159 | # Number of epoch between each snapshot 160 | snapshot_gap = 50 161 | 162 | # Do we nee to save convergence 163 | saving = True 164 | saving_path = None 165 | 166 | def __init__(self): 167 | """ 168 | Class Initialyser 169 | """ 170 | 171 | # Number of layers 172 | self.num_layers = len([block for block in self.architecture if 'pool' in block or 'strided' in block]) + 1 173 | 174 | def load(self, path): 175 | 176 | filename = join(path, 'parameters.txt') 177 | with open(filename, 'r') as f: 178 | lines = f.readlines() 179 | 180 | # Class variable dictionary 181 | for line in lines: 182 | line_info = line.split() 183 | if len(line_info) > 1 and line_info[0] != '#': 184 | 185 | if line_info[2] == 'None': 186 | setattr(self, line_info[0], None) 187 | 188 | elif line_info[0] == 'lr_decay_epochs': 189 | self.lr_decays = {int(b.split(':')[0]): float(b.split(':')[1]) for b in line_info[2:]} 190 | 191 | elif line_info[0] == 'architecture': 192 | self.architecture = [b for b in line_info[2:]] 193 | 194 | elif line_info[0] == 'augment_symmetries': 195 | self.augment_symmetries = [bool(int(b)) for b in line_info[2:]] 196 | 197 | elif line_info[0] == 'num_classes': 198 | if len(line_info) > 3: 199 | self.num_classes = [int(c) for c in line_info[2:]] 200 | else: 201 | self.num_classes = int(line_info[2]) 202 | 203 | else: 204 | 205 | attr_type = type(getattr(self, line_info[0])) 206 | if attr_type == bool: 207 | setattr(self, line_info[0], attr_type(int(line_info[2]))) 208 | else: 209 | setattr(self, line_info[0], attr_type(line_info[2])) 210 | 211 | self.saving = True 212 | self.saving_path = path 213 | self.__init__() 214 | 215 | def save(self, path): 216 | 217 | with open(join(path, 'parameters.txt'), "w") as text_file: 218 | 219 | text_file.write('# -----------------------------------#\n') 220 | text_file.write('# Parameters of the training session #\n') 221 | text_file.write('# -----------------------------------#\n\n') 222 | 223 | # Input parameters 224 | text_file.write('# Input parameters\n') 225 | text_file.write('# ****************\n\n') 226 | text_file.write('dataset = {:s}\n'.format(self.dataset)) 227 | text_file.write('network_model = {:s}\n'.format(self.network_model)) 228 | if type(self.num_classes) is list: 229 | text_file.write('num_classes =') 230 | for n in self.num_classes: 231 | text_file.write(' {:d}'.format(n)) 232 | text_file.write('\n') 233 | else: 234 | text_file.write('num_classes = {:d}\n'.format(self.num_classes)) 235 | text_file.write('in_points_dim = {:d}\n'.format(self.in_points_dim)) 236 | text_file.write('in_features_dim = {:d}\n'.format(self.in_features_dim)) 237 | text_file.write('in_radius = {:.3f}\n'.format(self.in_radius)) 238 | text_file.write('input_threads = {:d}\n\n'.format(self.input_threads)) 239 | 240 | # Model parameters 241 | text_file.write('# Model parameters\n') 242 | text_file.write('# ****************\n\n') 243 | 244 | text_file.write('architecture =') 245 | for a in self.architecture: 246 | text_file.write(' {:s}'.format(a)) 247 | text_file.write('\n') 248 | text_file.write('num_layers = {:d}\n'.format(self.num_layers)) 249 | text_file.write('first_features_dim = {:d}\n'.format(self.first_features_dim)) 250 | text_file.write('use_batch_norm = {:d}\n'.format(int(self.use_batch_norm))) 251 | text_file.write('batch_norm_momentum = {:.3f}\n\n'.format(self.batch_norm_momentum)) 252 | text_file.write('segmentation_ratio = {:.3f}\n\n'.format(self.segmentation_ratio)) 253 | 254 | # KPConv parameters 255 | text_file.write('# KPConv parameters\n') 256 | text_file.write('# *****************\n\n') 257 | 258 | text_file.write('first_subsampling_dl = {:.3f}\n'.format(self.first_subsampling_dl)) 259 | text_file.write('num_kernel_points = {:d}\n'.format(self.num_kernel_points)) 260 | text_file.write('density_parameter = {:.3f}\n'.format(self.density_parameter)) 261 | text_file.write('fixed_kernel_points = {:s}\n'.format(self.fixed_kernel_points)) 262 | text_file.write('KP_extent = {:.3f}\n'.format(self.KP_extent)) 263 | text_file.write('KP_influence = {:s}\n'.format(self.KP_influence)) 264 | text_file.write('convolution_mode = {:s}\n'.format(self.convolution_mode)) 265 | text_file.write('trainable_positions = {:d}\n\n'.format(int(self.trainable_positions))) 266 | text_file.write('modulated = {:d}\n\n'.format(int(self.modulated))) 267 | 268 | # Training parameters 269 | text_file.write('# Training parameters\n') 270 | text_file.write('# *******************\n\n') 271 | 272 | text_file.write('learning_rate = {:f}\n'.format(self.learning_rate)) 273 | text_file.write('momentum = {:f}\n'.format(self.momentum)) 274 | text_file.write('lr_decay_epochs =') 275 | for e, d in self.lr_decays.items(): 276 | text_file.write(' {:d}:{:f}'.format(e, d)) 277 | text_file.write('\n') 278 | text_file.write('grad_clip_norm = {:f}\n\n'.format(self.grad_clip_norm)) 279 | 280 | text_file.write('augment_symmetries =') 281 | for a in self.augment_symmetries: 282 | text_file.write(' {:d}'.format(int(a))) 283 | text_file.write('\n') 284 | text_file.write('augment_rotation = {:d}\n'.format(self.augment_rotation)) 285 | text_file.write('augment_noise = {:f}\n'.format(self.augment_noise)) 286 | text_file.write('augment_occlusion = {:s}\n'.format(self.augment_occlusion)) 287 | text_file.write('augment_occlusion_ratio = {:.3f}\n'.format(self.augment_occlusion_ratio)) 288 | text_file.write('augment_occlusion_num = {:d}\n'.format(self.augment_occlusion_num)) 289 | text_file.write('augment_scale_anisotropic = {:d}\n'.format(int(self.augment_scale_anisotropic))) 290 | text_file.write('augment_scale_min = {:.3f}\n'.format(self.augment_scale_min)) 291 | text_file.write('augment_scale_max = {:.3f}\n'.format(self.augment_scale_max)) 292 | text_file.write('augment_color = {:.3f}\n\n'.format(self.augment_color)) 293 | 294 | text_file.write('weights_decay = {:f}\n'.format(self.weights_decay)) 295 | text_file.write('gaussian_decay = {:f}\n'.format(self.gaussian_decay)) 296 | text_file.write('batch_averaged_loss = {:d}\n'.format(int(self.batch_averaged_loss))) 297 | text_file.write('offsets_loss = {:s}\n'.format(self.offsets_loss)) 298 | text_file.write('offsets_decay = {:f}\n'.format(self.offsets_decay)) 299 | text_file.write('batch_num = {:d}\n'.format(self.batch_num)) 300 | text_file.write('max_epoch = {:d}\n'.format(self.max_epoch)) 301 | if self.epoch_steps is None: 302 | text_file.write('epoch_steps = None\n') 303 | else: 304 | text_file.write('epoch_steps = {:d}\n'.format(self.epoch_steps)) 305 | text_file.write('validation_size = {:d}\n'.format(self.validation_size)) 306 | text_file.write('snapshot_gap = {:d}\n'.format(self.snapshot_gap)) 307 | -------------------------------------------------------------------------------- /utils/metrics.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from sklearn.metrics import confusion_matrix 4 | 5 | 6 | def calculate_acc(predict, labels): 7 | pred_labels = torch.max(predict, dim=1)[1].int() 8 | return torch.sum(pred_labels == labels.int()) * 100 / predict.shape[0] 9 | 10 | 11 | def calculate_iou_single_shape(predict, labels, n_parts): 12 | pred_labels = torch.max(predict, dim=1)[1] 13 | Confs = confusion_matrix(labels.detach().cpu().numpy(), pred_labels.detach().cpu().numpy(), np.arange(n_parts)) 14 | 15 | # Objects IoU 16 | IoUs = IoU_from_confusions(Confs) 17 | return IoUs 18 | 19 | 20 | def calculate_iou(predict, labels, stack_lengths, n_parts): 21 | start_ind = 0 22 | iou_list = [] 23 | for length in stack_lengths: 24 | iou = calculate_iou_single_shape(predict[start_ind:start_ind + length], labels[start_ind:start_ind + length], n_parts) 25 | iou_list.append(iou) 26 | start_ind += length 27 | iou_list = np.array(iou_list) 28 | return np.array(iou_list).mean(axis=0) * 100 29 | 30 | 31 | def IoU_from_confusions(confusions): 32 | """ 33 | Computes IoU from confusion matrices. 34 | :param confusions: ([..., n_c, n_c] np.int32). Can be any dimension, the confusion matrices should be described by 35 | the last axes. n_c = number of classes 36 | :param ignore_unclassified: (bool). True if the the first class should be ignored in the results 37 | :return: ([..., n_c] np.float32) IoU score 38 | """ 39 | 40 | # Compute TP, FP, FN. This assume that the second to last axis counts the truths (like the first axis of a 41 | # confusion matrix), and that the last axis counts the predictions (like the second axis of a confusion matrix) 42 | TP = np.diagonal(confusions, axis1=-2, axis2=-1) 43 | TP_plus_FN = np.sum(confusions, axis=-1) 44 | TP_plus_FP = np.sum(confusions, axis=-2) 45 | 46 | # Compute IoU 47 | IoU = TP / (TP_plus_FP + TP_plus_FN - TP + 1e-6) 48 | 49 | # Compute mIoU with only the actual classes 50 | mask = TP_plus_FN < 1e-3 51 | counts = np.sum(1 - mask, axis=-1, keepdims=True) 52 | mIoU = np.sum(IoU, axis=-1, keepdims=True) / (counts + 1e-6) 53 | 54 | # If class is absent, place mIoU in place of 0 IoU to get the actual mean later 55 | IoU += mask * mIoU 56 | 57 | return IoU 58 | -------------------------------------------------------------------------------- /utils/ply.py: -------------------------------------------------------------------------------- 1 | # 2 | # 3 | # 0===============================0 4 | # | PLY files reader/writer | 5 | # 0===============================0 6 | # 7 | # 8 | # ---------------------------------------------------------------------------------------------------------------------- 9 | # 10 | # function to read/write .ply files 11 | # 12 | # ---------------------------------------------------------------------------------------------------------------------- 13 | # 14 | # Hugues THOMAS - 10/02/2017 15 | # 16 | 17 | 18 | # ---------------------------------------------------------------------------------------------------------------------- 19 | # 20 | # Imports and global variables 21 | # \**********************************/ 22 | # 23 | 24 | 25 | # Basic libs 26 | import numpy as np 27 | import sys 28 | 29 | # Define PLY types 30 | ply_dtypes = dict([ 31 | (b'int8', 'i1'), 32 | (b'char', 'i1'), 33 | (b'uint8', 'u1'), 34 | (b'uchar', 'u1'), 35 | (b'int16', 'i2'), 36 | (b'short', 'i2'), 37 | (b'uint16', 'u2'), 38 | (b'ushort', 'u2'), 39 | (b'int32', 'i4'), 40 | (b'int', 'i4'), 41 | (b'uint32', 'u4'), 42 | (b'uint', 'u4'), 43 | (b'float32', 'f4'), 44 | (b'float', 'f4'), 45 | (b'float64', 'f8'), 46 | (b'double', 'f8') 47 | ]) 48 | 49 | # Numpy reader format 50 | valid_formats = {'ascii': '', 'binary_big_endian': '>', 51 | 'binary_little_endian': '<'} 52 | 53 | 54 | # ---------------------------------------------------------------------------------------------------------------------- 55 | # 56 | # Functions 57 | # \***************/ 58 | # 59 | 60 | 61 | def parse_header(plyfile, ext): 62 | # Variables 63 | line = [] 64 | properties = [] 65 | num_points = None 66 | 67 | while b'end_header' not in line and line != b'': 68 | line = plyfile.readline() 69 | 70 | if b'element' in line: 71 | line = line.split() 72 | num_points = int(line[2]) 73 | 74 | elif b'property' in line: 75 | line = line.split() 76 | properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) 77 | 78 | return num_points, properties 79 | 80 | 81 | def parse_mesh_header(plyfile, ext): 82 | # Variables 83 | line = [] 84 | vertex_properties = [] 85 | num_points = None 86 | num_faces = None 87 | current_element = None 88 | 89 | while b'end_header' not in line and line != b'': 90 | line = plyfile.readline() 91 | 92 | # Find point element 93 | if b'element vertex' in line: 94 | current_element = 'vertex' 95 | line = line.split() 96 | num_points = int(line[2]) 97 | 98 | elif b'element face' in line: 99 | current_element = 'face' 100 | line = line.split() 101 | num_faces = int(line[2]) 102 | 103 | elif b'property' in line: 104 | if current_element == 'vertex': 105 | line = line.split() 106 | vertex_properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) 107 | elif current_element == 'vertex': 108 | if not line.startswith('property list uchar int'): 109 | raise ValueError('Unsupported faces property : ' + line) 110 | 111 | return num_points, num_faces, vertex_properties 112 | 113 | 114 | def read_ply(filename, triangular_mesh=False): 115 | """ 116 | Read ".ply" files 117 | 118 | Parameters 119 | ---------- 120 | filename : string 121 | the name of the file to read. 122 | 123 | Returns 124 | ------- 125 | result : array 126 | data stored in the file 127 | 128 | Examples 129 | -------- 130 | Store data in file 131 | 132 | >>> points = np.random.rand(5, 3) 133 | >>> values = np.random.randint(2, size=10) 134 | >>> write_ply('example.ply', [points, values], ['x', 'y', 'z', 'values']) 135 | 136 | Read the file 137 | 138 | >>> data = read_ply('example.ply') 139 | >>> values = data['values'] 140 | array([0, 0, 1, 1, 0]) 141 | 142 | >>> points = np.vstack((data['x'], data['y'], data['z'])).T 143 | array([[ 0.466 0.595 0.324] 144 | [ 0.538 0.407 0.654] 145 | [ 0.850 0.018 0.988] 146 | [ 0.395 0.394 0.363] 147 | [ 0.873 0.996 0.092]]) 148 | 149 | """ 150 | 151 | with open(filename, 'rb') as plyfile: 152 | 153 | # Check if the file start with ply 154 | if b'ply' not in plyfile.readline(): 155 | raise ValueError('The file does not start whith the word ply') 156 | 157 | # get binary_little/big or ascii 158 | fmt = plyfile.readline().split()[1].decode() 159 | if fmt == "ascii": 160 | raise ValueError('The file is not binary') 161 | 162 | # get extension for building the numpy dtypes 163 | ext = valid_formats[fmt] 164 | 165 | # PointCloud reader vs mesh reader 166 | if triangular_mesh: 167 | 168 | # Parse header 169 | num_points, num_faces, properties = parse_mesh_header(plyfile, ext) 170 | 171 | # Get point data 172 | vertex_data = np.fromfile(plyfile, dtype=properties, count=num_points) 173 | 174 | # Get face data 175 | face_properties = [('k', ext + 'u1'), 176 | ('v1', ext + 'i4'), 177 | ('v2', ext + 'i4'), 178 | ('v3', ext + 'i4')] 179 | faces_data = np.fromfile(plyfile, dtype=face_properties, count=num_faces) 180 | 181 | # Return vertex data and concatenated faces 182 | faces = np.vstack((faces_data['v1'], faces_data['v2'], faces_data['v3'])).T 183 | data = [vertex_data, faces] 184 | 185 | else: 186 | 187 | # Parse header 188 | num_points, properties = parse_header(plyfile, ext) 189 | 190 | # Get data 191 | data = np.fromfile(plyfile, dtype=properties, count=num_points) 192 | 193 | return data 194 | 195 | 196 | def header_properties(field_list, field_names): 197 | # List of lines to write 198 | lines = [] 199 | 200 | # First line describing element vertex 201 | lines.append('element vertex %d' % field_list[0].shape[0]) 202 | 203 | # Properties lines 204 | i = 0 205 | for fields in field_list: 206 | for field in fields.T: 207 | lines.append('property %s %s' % (field.dtype.name, field_names[i])) 208 | i += 1 209 | 210 | return lines 211 | 212 | 213 | def write_ply(filename, field_list, field_names, triangular_faces=None): 214 | """ 215 | Write ".ply" files 216 | 217 | Parameters 218 | ---------- 219 | filename : string 220 | the name of the file to which the data is saved. A '.ply' extension will be appended to the 221 | file name if it does no already have one. 222 | 223 | field_list : list, tuple, numpy array 224 | the fields to be saved in the ply file. Either a numpy array, a list of numpy arrays or a 225 | tuple of numpy arrays. Each 1D numpy array and each column of 2D numpy arrays are considered 226 | as one field. 227 | 228 | field_names : list 229 | the name of each fields as a list of strings. Has to be the same length as the number of 230 | fields. 231 | 232 | Examples 233 | -------- 234 | >>> points = np.random.rand(10, 3) 235 | >>> write_ply('example1.ply', points, ['x', 'y', 'z']) 236 | 237 | >>> values = np.random.randint(2, size=10) 238 | >>> write_ply('example2.ply', [points, values], ['x', 'y', 'z', 'values']) 239 | 240 | >>> colors = np.random.randint(255, size=(10,3), dtype=np.uint8) 241 | >>> field_names = ['x', 'y', 'z', 'red', 'green', 'blue', values'] 242 | >>> write_ply('example3.ply', [points, colors, values], field_names) 243 | 244 | """ 245 | 246 | # Format list input to the right form 247 | field_list = list(field_list) if (type(field_list) == list or type(field_list) == tuple) else list((field_list,)) 248 | for i, field in enumerate(field_list): 249 | if field.ndim < 2: 250 | field_list[i] = field.reshape(-1, 1) 251 | if field.ndim > 2: 252 | print('fields have more than 2 dimensions') 253 | return False 254 | 255 | # check all fields have the same number of data 256 | n_points = [field.shape[0] for field in field_list] 257 | if not np.all(np.equal(n_points, n_points[0])): 258 | print('wrong field dimensions') 259 | return False 260 | 261 | # Check if field_names and field_list have same nb of column 262 | n_fields = np.sum([field.shape[1] for field in field_list]) 263 | if (n_fields != len(field_names)): 264 | print('wrong number of field names') 265 | return False 266 | 267 | # Add extension if not there 268 | if not filename.endswith('.ply'): 269 | filename += '.ply' 270 | 271 | # open in text mode to write the header 272 | with open(filename, 'w') as plyfile: 273 | 274 | # First magical word 275 | header = ['ply'] 276 | 277 | # Encoding format 278 | header.append('format binary_' + sys.byteorder + '_endian 1.0') 279 | 280 | # Points properties description 281 | header.extend(header_properties(field_list, field_names)) 282 | 283 | # Add faces if needded 284 | if triangular_faces is not None: 285 | header.append('element face {:d}'.format(triangular_faces.shape[0])) 286 | header.append('property list uchar int vertex_indices') 287 | 288 | # End of header 289 | header.append('end_header') 290 | 291 | # Write all lines 292 | for line in header: 293 | plyfile.write("%s\n" % line) 294 | 295 | # open in binary/append to use tofile 296 | with open(filename, 'ab') as plyfile: 297 | 298 | # Create a structured array 299 | i = 0 300 | type_list = [] 301 | for fields in field_list: 302 | for field in fields.T: 303 | type_list += [(field_names[i], field.dtype.str)] 304 | i += 1 305 | data = np.empty(field_list[0].shape[0], dtype=type_list) 306 | i = 0 307 | for fields in field_list: 308 | for field in fields.T: 309 | data[field_names[i]] = field 310 | i += 1 311 | 312 | data.tofile(plyfile) 313 | 314 | if triangular_faces is not None: 315 | triangular_faces = triangular_faces.astype(np.int32) 316 | type_list = [('k', 'uint8')] + [(str(ind), 'int32') for ind in range(3)] 317 | data = np.empty(triangular_faces.shape[0], dtype=type_list) 318 | data['k'] = np.full((triangular_faces.shape[0],), 3, dtype=np.uint8) 319 | data['0'] = triangular_faces[:, 0] 320 | data['1'] = triangular_faces[:, 1] 321 | data['2'] = triangular_faces[:, 2] 322 | data.tofile(plyfile) 323 | 324 | return True 325 | 326 | 327 | def describe_element(name, df): 328 | """ Takes the columns of the dataframe and builds a ply-like description 329 | 330 | Parameters 331 | ---------- 332 | name: str 333 | df: pandas DataFrame 334 | 335 | Returns 336 | ------- 337 | element: list[str] 338 | """ 339 | property_formats = {'f': 'float', 'u': 'uchar', 'i': 'int'} 340 | element = ['element ' + name + ' ' + str(len(df))] 341 | 342 | if name == 'face': 343 | element.append("property list uchar int points_indices") 344 | 345 | else: 346 | for i in range(len(df.columns)): 347 | # get first letter of dtype to infer format 348 | f = property_formats[str(df.dtypes[i])[0]] 349 | element.append('property ' + f + ' ' + df.columns.values[i]) 350 | 351 | return element 352 | 353 | -------------------------------------------------------------------------------- /utils/pointcloud.py: -------------------------------------------------------------------------------- 1 | import open3d 2 | 3 | 4 | def make_point_cloud(pts): 5 | pcd = open3d.PointCloud() 6 | pcd.points = open3d.Vector3dVector(pts) 7 | return pcd 8 | -------------------------------------------------------------------------------- /utils/timer.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | 4 | class AverageMeter(object): 5 | """Computes and stores the average and current value""" 6 | 7 | def __init__(self): 8 | self.reset() 9 | 10 | def reset(self): 11 | self.val = 0 12 | self.avg = 0 13 | self.sum = 0.0 14 | self.sq_sum = 0.0 15 | self.count = 0 16 | 17 | def update(self, val, n=1): 18 | self.val = val 19 | self.sum += val * n 20 | self.count += n 21 | self.avg = self.sum / self.count 22 | self.sq_sum += val ** 2 * n 23 | self.var = self.sq_sum / self.count - self.avg ** 2 24 | 25 | 26 | class Timer(object): 27 | """A simple timer.""" 28 | 29 | def __init__(self): 30 | self.total_time = 0. 31 | self.calls = 0 32 | self.start_time = 0. 33 | self.diff = 0. 34 | self.avg = 0. 35 | 36 | def reset(self): 37 | self.total_time = 0 38 | self.calls = 0 39 | self.start_time = 0 40 | self.diff = 0 41 | self.avg = 0 42 | 43 | def tic(self): 44 | # using time.time instead of time.clock because time time.clock 45 | # does not normalize for multithreading 46 | self.start_time = time.time() 47 | 48 | def toc(self, average=True): 49 | self.diff = time.time() - self.start_time 50 | self.total_time += self.diff 51 | self.calls += 1 52 | self.avg = self.total_time / self.calls 53 | if average: 54 | return self.avg 55 | else: 56 | return self.diff 57 | --------------------------------------------------------------------------------