commit c13e1ba86344319e2a44519751e8a50c39ee684b Author: HuguesTHOMAS Date: Tue Mar 31 15:42:35 2020 -0400 Initial commit diff --git a/cpp_wrappers/compile_wrappers.sh b/cpp_wrappers/compile_wrappers.sh new file mode 100644 index 0000000..9f41ae1 --- /dev/null +++ b/cpp_wrappers/compile_wrappers.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +# Compile cpp subsampling +cd cpp_subsampling +python3 setup.py build_ext --inplace +cd .. + +# Compile cpp neighbors +cd cpp_neighbors +python3 setup.py build_ext --inplace +cd .. \ No newline at end of file diff --git a/cpp_wrappers/cpp_neighbors/neighbors/neighbors.cpp b/cpp_wrappers/cpp_neighbors/neighbors/neighbors.cpp new file mode 100644 index 0000000..bf22af8 --- /dev/null +++ b/cpp_wrappers/cpp_neighbors/neighbors/neighbors.cpp @@ -0,0 +1,333 @@ + +#include "neighbors.h" + + +void brute_neighbors(vector& queries, vector& supports, vector& neighbors_indices, float radius, int verbose) +{ + + // Initialize variables + // ****************** + + // square radius + float r2 = radius * radius; + + // indices + int i0 = 0; + + // Counting vector + int max_count = 0; + vector> tmp(queries.size()); + + // Search neigbors indices + // *********************** + + for (auto& p0 : queries) + { + int i = 0; + for (auto& p : supports) + { + if ((p0 - p).sq_norm() < r2) + { + tmp[i0].push_back(i); + if (tmp[i0].size() > max_count) + max_count = tmp[i0].size(); + } + i++; + } + i0++; + } + + // Reserve the memory + neighbors_indices.resize(queries.size() * max_count); + i0 = 0; + for (auto& inds : tmp) + { + for (int j = 0; j < max_count; j++) + { + if (j < inds.size()) + neighbors_indices[i0 * max_count + j] = inds[j]; + else + neighbors_indices[i0 * max_count + j] = -1; + } + i0++; + } + + return; +} + +void ordered_neighbors(vector& queries, + vector& supports, + vector& neighbors_indices, + float radius) +{ + + // Initialize variables + // ****************** + + // square radius + float r2 = radius * radius; + + // indices + int i0 = 0; + + // Counting vector + int max_count = 0; + float d2; + vector> tmp(queries.size()); + vector> dists(queries.size()); + + // Search neigbors indices + // *********************** + + for (auto& p0 : queries) + { + int i = 0; + for (auto& p : supports) + { + d2 = (p0 - p).sq_norm(); + if (d2 < r2) + { + // Find order of the new point + auto it = std::upper_bound(dists[i0].begin(), dists[i0].end(), d2); + int index = std::distance(dists[i0].begin(), it); + + // Insert element + dists[i0].insert(it, d2); + tmp[i0].insert(tmp[i0].begin() + index, i); + + // Update max count + if (tmp[i0].size() > max_count) + max_count = tmp[i0].size(); + } + i++; + } + i0++; + } + + // Reserve the memory + neighbors_indices.resize(queries.size() * max_count); + i0 = 0; + for (auto& inds : tmp) + { + for (int j = 0; j < max_count; j++) + { + if (j < inds.size()) + neighbors_indices[i0 * max_count + j] = inds[j]; + else + neighbors_indices[i0 * max_count + j] = -1; + } + i0++; + } + + return; +} + +void batch_ordered_neighbors(vector& queries, + vector& supports, + vector& q_batches, + vector& s_batches, + vector& neighbors_indices, + float radius) +{ + + // Initialize variables + // ****************** + + // square radius + float r2 = radius * radius; + + // indices + int i0 = 0; + + // Counting vector + int max_count = 0; + float d2; + vector> tmp(queries.size()); + vector> dists(queries.size()); + + // batch index + int b = 0; + int sum_qb = 0; + int sum_sb = 0; + + + // Search neigbors indices + // *********************** + + for (auto& p0 : queries) + { + // Check if we changed batch + if (i0 == sum_qb + q_batches[b]) + { + sum_qb += q_batches[b]; + sum_sb += s_batches[b]; + b++; + } + + // Loop only over the supports of current batch + vector::iterator p_it; + int i = 0; + for(p_it = supports.begin() + sum_sb; p_it < supports.begin() + sum_sb + s_batches[b]; p_it++ ) + { + d2 = (p0 - *p_it).sq_norm(); + if (d2 < r2) + { + // Find order of the new point + auto it = std::upper_bound(dists[i0].begin(), dists[i0].end(), d2); + int index = std::distance(dists[i0].begin(), it); + + // Insert element + dists[i0].insert(it, d2); + tmp[i0].insert(tmp[i0].begin() + index, sum_sb + i); + + // Update max count + if (tmp[i0].size() > max_count) + max_count = tmp[i0].size(); + } + i++; + } + i0++; + } + + // Reserve the memory + neighbors_indices.resize(queries.size() * max_count); + i0 = 0; + for (auto& inds : tmp) + { + for (int j = 0; j < max_count; j++) + { + if (j < inds.size()) + neighbors_indices[i0 * max_count + j] = inds[j]; + else + neighbors_indices[i0 * max_count + j] = supports.size(); + } + i0++; + } + + return; +} + + +void batch_nanoflann_neighbors(vector& queries, + vector& supports, + vector& q_batches, + vector& s_batches, + vector& neighbors_indices, + float radius) +{ + + // Initialize variables + // ****************** + + // indices + int i0 = 0; + + // Square radius + float r2 = radius * radius; + + // Counting vector + int max_count = 0; + float d2; + vector>> all_inds_dists(queries.size()); + + // batch index + int b = 0; + int sum_qb = 0; + int sum_sb = 0; + + // Nanoflann related variables + // *************************** + + // CLoud variable + PointCloud current_cloud; + + // Tree parameters + nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */); + + // KDTree type definition + typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Simple_Adaptor , + PointCloud, + 3 > my_kd_tree_t; + + // Pointer to trees + my_kd_tree_t* index; + + // Build KDTree for the first batch element + current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]); + index = new my_kd_tree_t(3, current_cloud, tree_params); + index->buildIndex(); + + + // Search neigbors indices + // *********************** + + // Search params + nanoflann::SearchParams search_params; + search_params.sorted = true; + + for (auto& p0 : queries) + { + + // Check if we changed batch + if (i0 == sum_qb + q_batches[b]) + { + sum_qb += q_batches[b]; + sum_sb += s_batches[b]; + b++; + + // Change the points + current_cloud.pts.clear(); + current_cloud.pts = vector(supports.begin() + sum_sb, supports.begin() + sum_sb + s_batches[b]); + + // Build KDTree of the current element of the batch + delete index; + index = new my_kd_tree_t(3, current_cloud, tree_params); + index->buildIndex(); + } + + // Initial guess of neighbors size + all_inds_dists[i0].reserve(max_count); + + // Find neighbors + float query_pt[3] = { p0.x, p0.y, p0.z}; + size_t nMatches = index->radiusSearch(query_pt, r2, all_inds_dists[i0], search_params); + + // Update max count + if (nMatches > max_count) + max_count = nMatches; + + // Increment query idx + i0++; + } + + // Reserve the memory + neighbors_indices.resize(queries.size() * max_count); + i0 = 0; + sum_sb = 0; + sum_qb = 0; + b = 0; + for (auto& inds_dists : all_inds_dists) + { + // Check if we changed batch + if (i0 == sum_qb + q_batches[b]) + { + sum_qb += q_batches[b]; + sum_sb += s_batches[b]; + b++; + } + + for (int j = 0; j < max_count; j++) + { + if (j < inds_dists.size()) + neighbors_indices[i0 * max_count + j] = inds_dists[j].first + sum_sb; + else + neighbors_indices[i0 * max_count + j] = supports.size(); + } + i0++; + } + + delete index; + + return; +} + diff --git a/cpp_wrappers/cpp_neighbors/neighbors/neighbors.h b/cpp_wrappers/cpp_neighbors/neighbors/neighbors.h new file mode 100644 index 0000000..ff612b0 --- /dev/null +++ b/cpp_wrappers/cpp_neighbors/neighbors/neighbors.h @@ -0,0 +1,29 @@ + + +#include "../../cpp_utils/cloud/cloud.h" +#include "../../cpp_utils/nanoflann/nanoflann.hpp" + +#include +#include + +using namespace std; + + +void ordered_neighbors(vector& queries, + vector& supports, + vector& neighbors_indices, + float radius); + +void batch_ordered_neighbors(vector& queries, + vector& supports, + vector& q_batches, + vector& s_batches, + vector& neighbors_indices, + float radius); + +void batch_nanoflann_neighbors(vector& queries, + vector& supports, + vector& q_batches, + vector& s_batches, + vector& neighbors_indices, + float radius); diff --git a/cpp_wrappers/cpp_neighbors/setup.py b/cpp_wrappers/cpp_neighbors/setup.py new file mode 100644 index 0000000..8f53a9c --- /dev/null +++ b/cpp_wrappers/cpp_neighbors/setup.py @@ -0,0 +1,28 @@ +from distutils.core import setup, Extension +import numpy.distutils.misc_util + +# Adding OpenCV to project +# ************************ + +# Adding sources of the project +# ***************************** + +SOURCES = ["../cpp_utils/cloud/cloud.cpp", + "neighbors/neighbors.cpp", + "wrapper.cpp"] + +module = Extension(name="radius_neighbors", + sources=SOURCES, + extra_compile_args=['-std=c++11', + '-D_GLIBCXX_USE_CXX11_ABI=0']) + + +setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) + + + + + + + + diff --git a/cpp_wrappers/cpp_neighbors/wrapper.cpp b/cpp_wrappers/cpp_neighbors/wrapper.cpp new file mode 100644 index 0000000..a4e2809 --- /dev/null +++ b/cpp_wrappers/cpp_neighbors/wrapper.cpp @@ -0,0 +1,238 @@ +#include +#include +#include "neighbors/neighbors.h" +#include + + + +// docstrings for our module +// ************************* + +static char module_docstring[] = "This module provides two methods to compute radius neighbors from pointclouds or batch of pointclouds"; + +static char batch_query_docstring[] = "Method to get radius neighbors in a batch of stacked pointclouds"; + + +// Declare the functions +// ********************* + +static PyObject *batch_neighbors(PyObject *self, PyObject *args, PyObject *keywds); + + +// Specify the members of the module +// ********************************* + +static PyMethodDef module_methods[] = +{ + { "batch_query", (PyCFunction)batch_neighbors, METH_VARARGS | METH_KEYWORDS, batch_query_docstring }, + {NULL, NULL, 0, NULL} +}; + + +// Initialize the module +// ********************* + +static struct PyModuleDef moduledef = +{ + PyModuleDef_HEAD_INIT, + "radius_neighbors", // m_name + module_docstring, // m_doc + -1, // m_size + module_methods, // m_methods + NULL, // m_reload + NULL, // m_traverse + NULL, // m_clear + NULL, // m_free +}; + +PyMODINIT_FUNC PyInit_radius_neighbors(void) +{ + import_array(); + return PyModule_Create(&moduledef); +} + + +// Definition of the batch_subsample method +// ********************************** + +static PyObject* batch_neighbors(PyObject* self, PyObject* args, PyObject* keywds) +{ + + // Manage inputs + // ************* + + // Args containers + PyObject* queries_obj = NULL; + PyObject* supports_obj = NULL; + PyObject* q_batches_obj = NULL; + PyObject* s_batches_obj = NULL; + + // Keywords containers + static char* kwlist[] = { "queries", "supports", "q_batches", "s_batches", "radius", NULL }; + float radius = 0.1; + + // Parse the input + if (!PyArg_ParseTupleAndKeywords(args, keywds, "OOOO|$f", kwlist, &queries_obj, &supports_obj, &q_batches_obj, &s_batches_obj, &radius)) + { + PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); + return NULL; + } + + + // Interpret the input objects as numpy arrays. + PyObject* queries_array = PyArray_FROM_OTF(queries_obj, NPY_FLOAT, NPY_IN_ARRAY); + PyObject* supports_array = PyArray_FROM_OTF(supports_obj, NPY_FLOAT, NPY_IN_ARRAY); + PyObject* q_batches_array = PyArray_FROM_OTF(q_batches_obj, NPY_INT, NPY_IN_ARRAY); + PyObject* s_batches_array = PyArray_FROM_OTF(s_batches_obj, NPY_INT, NPY_IN_ARRAY); + + // Verify data was load correctly. + if (queries_array == NULL) + { + Py_XDECREF(queries_array); + Py_XDECREF(supports_array); + Py_XDECREF(q_batches_array); + Py_XDECREF(s_batches_array); + PyErr_SetString(PyExc_RuntimeError, "Error converting query points to numpy arrays of type float32"); + return NULL; + } + if (supports_array == NULL) + { + Py_XDECREF(queries_array); + Py_XDECREF(supports_array); + Py_XDECREF(q_batches_array); + Py_XDECREF(s_batches_array); + PyErr_SetString(PyExc_RuntimeError, "Error converting support points to numpy arrays of type float32"); + return NULL; + } + if (q_batches_array == NULL) + { + Py_XDECREF(queries_array); + Py_XDECREF(supports_array); + Py_XDECREF(q_batches_array); + Py_XDECREF(s_batches_array); + PyErr_SetString(PyExc_RuntimeError, "Error converting query batches to numpy arrays of type int32"); + return NULL; + } + if (s_batches_array == NULL) + { + Py_XDECREF(queries_array); + Py_XDECREF(supports_array); + Py_XDECREF(q_batches_array); + Py_XDECREF(s_batches_array); + PyErr_SetString(PyExc_RuntimeError, "Error converting support batches to numpy arrays of type int32"); + return NULL; + } + + // Check that the input array respect the dims + if ((int)PyArray_NDIM(queries_array) != 2 || (int)PyArray_DIM(queries_array, 1) != 3) + { + Py_XDECREF(queries_array); + Py_XDECREF(supports_array); + Py_XDECREF(q_batches_array); + Py_XDECREF(s_batches_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : query.shape is not (N, 3)"); + return NULL; + } + if ((int)PyArray_NDIM(supports_array) != 2 || (int)PyArray_DIM(supports_array, 1) != 3) + { + Py_XDECREF(queries_array); + Py_XDECREF(supports_array); + Py_XDECREF(q_batches_array); + Py_XDECREF(s_batches_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : support.shape is not (N, 3)"); + return NULL; + } + if ((int)PyArray_NDIM(q_batches_array) > 1) + { + Py_XDECREF(queries_array); + Py_XDECREF(supports_array); + Py_XDECREF(q_batches_array); + Py_XDECREF(s_batches_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : queries_batches.shape is not (B,) "); + return NULL; + } + if ((int)PyArray_NDIM(s_batches_array) > 1) + { + Py_XDECREF(queries_array); + Py_XDECREF(supports_array); + Py_XDECREF(q_batches_array); + Py_XDECREF(s_batches_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : supports_batches.shape is not (B,) "); + return NULL; + } + if ((int)PyArray_DIM(q_batches_array, 0) != (int)PyArray_DIM(s_batches_array, 0)) + { + Py_XDECREF(queries_array); + Py_XDECREF(supports_array); + Py_XDECREF(q_batches_array); + Py_XDECREF(s_batches_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong number of batch elements: different for queries and supports "); + return NULL; + } + + // Number of points + int Nq = (int)PyArray_DIM(queries_array, 0); + int Ns= (int)PyArray_DIM(supports_array, 0); + + // Number of batches + int Nb = (int)PyArray_DIM(q_batches_array, 0); + + // Call the C++ function + // ********************* + + // Convert PyArray to Cloud C++ class + vector queries; + vector supports; + vector q_batches; + vector s_batches; + queries = vector((PointXYZ*)PyArray_DATA(queries_array), (PointXYZ*)PyArray_DATA(queries_array) + Nq); + supports = vector((PointXYZ*)PyArray_DATA(supports_array), (PointXYZ*)PyArray_DATA(supports_array) + Ns); + q_batches = vector((int*)PyArray_DATA(q_batches_array), (int*)PyArray_DATA(q_batches_array) + Nb); + s_batches = vector((int*)PyArray_DATA(s_batches_array), (int*)PyArray_DATA(s_batches_array) + Nb); + + // Create result containers + vector neighbors_indices; + + // Compute results + //batch_ordered_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius); + batch_nanoflann_neighbors(queries, supports, q_batches, s_batches, neighbors_indices, radius); + + // Check result + if (neighbors_indices.size() < 1) + { + PyErr_SetString(PyExc_RuntimeError, "Error"); + return NULL; + } + + // Manage outputs + // ************** + + // Maximal number of neighbors + int max_neighbors = neighbors_indices.size() / Nq; + + // Dimension of output containers + npy_intp* neighbors_dims = new npy_intp[2]; + neighbors_dims[0] = Nq; + neighbors_dims[1] = max_neighbors; + + // Create output array + PyObject* res_obj = PyArray_SimpleNew(2, neighbors_dims, NPY_INT); + PyObject* ret = NULL; + + // Fill output array with values + size_t size_in_bytes = Nq * max_neighbors * sizeof(int); + memcpy(PyArray_DATA(res_obj), neighbors_indices.data(), size_in_bytes); + + // Merge results + ret = Py_BuildValue("N", res_obj); + + // Clean up + // ******** + + Py_XDECREF(queries_array); + Py_XDECREF(supports_array); + Py_XDECREF(q_batches_array); + Py_XDECREF(s_batches_array); + + return ret; +} diff --git a/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp b/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp new file mode 100644 index 0000000..24276bb --- /dev/null +++ b/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp @@ -0,0 +1,211 @@ + +#include "grid_subsampling.h" + + +void grid_subsampling(vector& original_points, + vector& subsampled_points, + vector& original_features, + vector& subsampled_features, + vector& original_classes, + vector& subsampled_classes, + float sampleDl, + int verbose) { + + // Initialize variables + // ****************** + + // Number of points in the cloud + size_t N = original_points.size(); + + // Dimension of the features + size_t fdim = original_features.size() / N; + size_t ldim = original_classes.size() / N; + + // Limits of the cloud + PointXYZ minCorner = min_point(original_points); + PointXYZ maxCorner = max_point(original_points); + PointXYZ originCorner = floor(minCorner * (1/sampleDl)) * sampleDl; + + // Dimensions of the grid + size_t sampleNX = (size_t)floor((maxCorner.x - originCorner.x) / sampleDl) + 1; + size_t sampleNY = (size_t)floor((maxCorner.y - originCorner.y) / sampleDl) + 1; + //size_t sampleNZ = (size_t)floor((maxCorner.z - originCorner.z) / sampleDl) + 1; + + // Check if features and classes need to be processed + bool use_feature = original_features.size() > 0; + bool use_classes = original_classes.size() > 0; + + + // Create the sampled map + // ********************** + + // Verbose parameters + int i = 0; + int nDisp = N / 100; + + // Initialize variables + size_t iX, iY, iZ, mapIdx; + unordered_map data; + + for (auto& p : original_points) + { + // Position of point in sample map + iX = (size_t)floor((p.x - originCorner.x) / sampleDl); + iY = (size_t)floor((p.y - originCorner.y) / sampleDl); + iZ = (size_t)floor((p.z - originCorner.z) / sampleDl); + mapIdx = iX + sampleNX*iY + sampleNX*sampleNY*iZ; + + // If not already created, create key + if (data.count(mapIdx) < 1) + data.emplace(mapIdx, SampledData(fdim, ldim)); + + // Fill the sample map + if (use_feature && use_classes) + data[mapIdx].update_all(p, original_features.begin() + i * fdim, original_classes.begin() + i * ldim); + else if (use_feature) + data[mapIdx].update_features(p, original_features.begin() + i * fdim); + else if (use_classes) + data[mapIdx].update_classes(p, original_classes.begin() + i * ldim); + else + data[mapIdx].update_points(p); + + // Display + i++; + if (verbose > 1 && i%nDisp == 0) + std::cout << "\rSampled Map : " << std::setw(3) << i / nDisp << "%"; + + } + + // Divide for barycentre and transfer to a vector + subsampled_points.reserve(data.size()); + if (use_feature) + subsampled_features.reserve(data.size() * fdim); + if (use_classes) + subsampled_classes.reserve(data.size() * ldim); + for (auto& v : data) + { + subsampled_points.push_back(v.second.point * (1.0 / v.second.count)); + if (use_feature) + { + float count = (float)v.second.count; + transform(v.second.features.begin(), + v.second.features.end(), + v.second.features.begin(), + [count](float f) { return f / count;}); + subsampled_features.insert(subsampled_features.end(),v.second.features.begin(),v.second.features.end()); + } + if (use_classes) + { + for (int i = 0; i < ldim; i++) + subsampled_classes.push_back(max_element(v.second.labels[i].begin(), v.second.labels[i].end(), + [](const pair&a, const pair&b){return a.second < b.second;})->first); + } + } + + return; +} + + +void batch_grid_subsampling(vector& original_points, + vector& subsampled_points, + vector& original_features, + vector& subsampled_features, + vector& original_classes, + vector& subsampled_classes, + vector& original_batches, + vector& subsampled_batches, + float sampleDl, + int max_p) +{ + // Initialize variables + // ****************** + + int b = 0; + int sum_b = 0; + + // Number of points in the cloud + size_t N = original_points.size(); + + // Dimension of the features + size_t fdim = original_features.size() / N; + size_t ldim = original_classes.size() / N; + + // Handle max_p = 0 + if (max_p < 1) + max_p = N; + + // Loop over batches + // ***************** + + for (b = 0; b < original_batches.size(); b++) + { + + // Extract batch points features and labels + vector b_o_points = vector(original_points.begin () + sum_b, + original_points.begin () + sum_b + original_batches[b]); + + vector b_o_features; + if (original_features.size() > 0) + { + b_o_features = vector(original_features.begin () + sum_b * fdim, + original_features.begin () + (sum_b + original_batches[b]) * fdim); + } + + vector b_o_classes; + if (original_classes.size() > 0) + { + b_o_classes = vector(original_classes.begin () + sum_b * ldim, + original_classes.begin () + sum_b + original_batches[b] * ldim); + } + + + // Create result containers + vector b_s_points; + vector b_s_features; + vector b_s_classes; + + // Compute subsampling on current batch + grid_subsampling(b_o_points, + b_s_points, + b_o_features, + b_s_features, + b_o_classes, + b_s_classes, + sampleDl, + 0); + + // Stack batches points features and labels + // **************************************** + + // If too many points remove some + if (b_s_points.size() <= max_p) + { + subsampled_points.insert(subsampled_points.end(), b_s_points.begin(), b_s_points.end()); + + if (original_features.size() > 0) + subsampled_features.insert(subsampled_features.end(), b_s_features.begin(), b_s_features.end()); + + if (original_classes.size() > 0) + subsampled_classes.insert(subsampled_classes.end(), b_s_classes.begin(), b_s_classes.end()); + + subsampled_batches.push_back(b_s_points.size()); + } + else + { + subsampled_points.insert(subsampled_points.end(), b_s_points.begin(), b_s_points.begin() + max_p); + + if (original_features.size() > 0) + subsampled_features.insert(subsampled_features.end(), b_s_features.begin(), b_s_features.begin() + max_p * fdim); + + if (original_classes.size() > 0) + subsampled_classes.insert(subsampled_classes.end(), b_s_classes.begin(), b_s_classes.begin() + max_p * ldim); + + subsampled_batches.push_back(max_p); + } + + // Stack new batch lengths + sum_b += original_batches[b]; + } + + return; +} diff --git a/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h b/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h new file mode 100644 index 0000000..37f775d --- /dev/null +++ b/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h @@ -0,0 +1,101 @@ + + +#include "../../cpp_utils/cloud/cloud.h" + +#include +#include + +using namespace std; + +class SampledData +{ +public: + + // Elements + // ******** + + int count; + PointXYZ point; + vector features; + vector> labels; + + + // Methods + // ******* + + // Constructor + SampledData() + { + count = 0; + point = PointXYZ(); + } + + SampledData(const size_t fdim, const size_t ldim) + { + count = 0; + point = PointXYZ(); + features = vector(fdim); + labels = vector>(ldim); + } + + // Method Update + void update_all(const PointXYZ p, vector::iterator f_begin, vector::iterator l_begin) + { + count += 1; + point += p; + transform (features.begin(), features.end(), f_begin, features.begin(), plus()); + int i = 0; + for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) + { + labels[i][*it] += 1; + i++; + } + return; + } + void update_features(const PointXYZ p, vector::iterator f_begin) + { + count += 1; + point += p; + transform (features.begin(), features.end(), f_begin, features.begin(), plus()); + return; + } + void update_classes(const PointXYZ p, vector::iterator l_begin) + { + count += 1; + point += p; + int i = 0; + for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) + { + labels[i][*it] += 1; + i++; + } + return; + } + void update_points(const PointXYZ p) + { + count += 1; + point += p; + return; + } +}; + +void grid_subsampling(vector& original_points, + vector& subsampled_points, + vector& original_features, + vector& subsampled_features, + vector& original_classes, + vector& subsampled_classes, + float sampleDl, + int verbose); + +void batch_grid_subsampling(vector& original_points, + vector& subsampled_points, + vector& original_features, + vector& subsampled_features, + vector& original_classes, + vector& subsampled_classes, + vector& original_batches, + vector& subsampled_batches, + float sampleDl, + int max_p); + diff --git a/cpp_wrappers/cpp_subsampling/setup.py b/cpp_wrappers/cpp_subsampling/setup.py new file mode 100644 index 0000000..3206299 --- /dev/null +++ b/cpp_wrappers/cpp_subsampling/setup.py @@ -0,0 +1,28 @@ +from distutils.core import setup, Extension +import numpy.distutils.misc_util + +# Adding OpenCV to project +# ************************ + +# Adding sources of the project +# ***************************** + +SOURCES = ["../cpp_utils/cloud/cloud.cpp", + "grid_subsampling/grid_subsampling.cpp", + "wrapper.cpp"] + +module = Extension(name="grid_subsampling", + sources=SOURCES, + extra_compile_args=['-std=c++11', + '-D_GLIBCXX_USE_CXX11_ABI=0']) + + +setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) + + + + + + + + diff --git a/cpp_wrappers/cpp_subsampling/wrapper.cpp b/cpp_wrappers/cpp_subsampling/wrapper.cpp new file mode 100644 index 0000000..8a92aaa --- /dev/null +++ b/cpp_wrappers/cpp_subsampling/wrapper.cpp @@ -0,0 +1,566 @@ +#include +#include +#include "grid_subsampling/grid_subsampling.h" +#include + + + +// docstrings for our module +// ************************* + +static char module_docstring[] = "This module provides an interface for the subsampling of a batch of stacked pointclouds"; + +static char subsample_docstring[] = "function subsampling a pointcloud"; + +static char subsample_batch_docstring[] = "function subsampling a batch of stacked pointclouds"; + + +// Declare the functions +// ********************* + +static PyObject *cloud_subsampling(PyObject* self, PyObject* args, PyObject* keywds); +static PyObject *batch_subsampling(PyObject *self, PyObject *args, PyObject *keywds); + + +// Specify the members of the module +// ********************************* + +static PyMethodDef module_methods[] = +{ + { "subsample", (PyCFunction)cloud_subsampling, METH_VARARGS | METH_KEYWORDS, subsample_docstring }, + { "subsample_batch", (PyCFunction)batch_subsampling, METH_VARARGS | METH_KEYWORDS, subsample_batch_docstring }, + {NULL, NULL, 0, NULL} +}; + + +// Initialize the module +// ********************* + +static struct PyModuleDef moduledef = +{ + PyModuleDef_HEAD_INIT, + "grid_subsampling", // m_name + module_docstring, // m_doc + -1, // m_size + module_methods, // m_methods + NULL, // m_reload + NULL, // m_traverse + NULL, // m_clear + NULL, // m_free +}; + +PyMODINIT_FUNC PyInit_grid_subsampling(void) +{ + import_array(); + return PyModule_Create(&moduledef); +} + + +// Definition of the batch_subsample method +// ********************************** + +static PyObject* batch_subsampling(PyObject* self, PyObject* args, PyObject* keywds) +{ + + // Manage inputs + // ************* + + // Args containers + PyObject* points_obj = NULL; + PyObject* features_obj = NULL; + PyObject* classes_obj = NULL; + PyObject* batches_obj = NULL; + + // Keywords containers + static char* kwlist[] = { "points", "batches", "features", "classes", "sampleDl", "method", "max_p", "verbose", NULL }; + float sampleDl = 0.1; + const char* method_buffer = "barycenters"; + int verbose = 0; + int max_p = 0; + + // Parse the input + if (!PyArg_ParseTupleAndKeywords(args, keywds, "OO|$OOfsii", kwlist, &points_obj, &batches_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &max_p, &verbose)) + { + PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); + return NULL; + } + + // Get the method argument + string method(method_buffer); + + // Interpret method + if (method.compare("barycenters") && method.compare("voxelcenters")) + { + PyErr_SetString(PyExc_RuntimeError, "Error parsing method. Valid method names are \"barycenters\" and \"voxelcenters\" "); + return NULL; + } + + // Check if using features or classes + bool use_feature = true, use_classes = true; + if (features_obj == NULL) + use_feature = false; + if (classes_obj == NULL) + use_classes = false; + + // Interpret the input objects as numpy arrays. + PyObject* points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY); + PyObject* batches_array = PyArray_FROM_OTF(batches_obj, NPY_INT, NPY_IN_ARRAY); + PyObject* features_array = NULL; + PyObject* classes_array = NULL; + if (use_feature) + features_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY); + if (use_classes) + classes_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY); + + // Verify data was load correctly. + if (points_array == NULL) + { + Py_XDECREF(points_array); + Py_XDECREF(batches_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Error converting input points to numpy arrays of type float32"); + return NULL; + } + if (batches_array == NULL) + { + Py_XDECREF(points_array); + Py_XDECREF(batches_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Error converting input batches to numpy arrays of type int32"); + return NULL; + } + if (use_feature && features_array == NULL) + { + Py_XDECREF(points_array); + Py_XDECREF(batches_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Error converting input features to numpy arrays of type float32"); + return NULL; + } + if (use_classes && classes_array == NULL) + { + Py_XDECREF(points_array); + Py_XDECREF(batches_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Error converting input classes to numpy arrays of type int32"); + return NULL; + } + + // Check that the input array respect the dims + if ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3) + { + Py_XDECREF(points_array); + Py_XDECREF(batches_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : points.shape is not (N, 3)"); + return NULL; + } + if ((int)PyArray_NDIM(batches_array) > 1) + { + Py_XDECREF(points_array); + Py_XDECREF(batches_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : batches.shape is not (B,) "); + return NULL; + } + if (use_feature && ((int)PyArray_NDIM(features_array) != 2)) + { + Py_XDECREF(points_array); + Py_XDECREF(batches_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); + return NULL; + } + + if (use_classes && (int)PyArray_NDIM(classes_array) > 2) + { + Py_XDECREF(points_array); + Py_XDECREF(batches_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); + return NULL; + } + + // Number of points + int N = (int)PyArray_DIM(points_array, 0); + + // Number of batches + int Nb = (int)PyArray_DIM(batches_array, 0); + + // Dimension of the features + int fdim = 0; + if (use_feature) + fdim = (int)PyArray_DIM(features_array, 1); + + //Dimension of labels + int ldim = 1; + if (use_classes && (int)PyArray_NDIM(classes_array) == 2) + ldim = (int)PyArray_DIM(classes_array, 1); + + // Check that the input array respect the number of points + if (use_feature && (int)PyArray_DIM(features_array, 0) != N) + { + Py_XDECREF(points_array); + Py_XDECREF(batches_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); + return NULL; + } + if (use_classes && (int)PyArray_DIM(classes_array, 0) != N) + { + Py_XDECREF(points_array); + Py_XDECREF(batches_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); + return NULL; + } + + + // Call the C++ function + // ********************* + + // Create pyramid + if (verbose > 0) + cout << "Computing cloud pyramid with support points: " << endl; + + + // Convert PyArray to Cloud C++ class + vector original_points; + vector original_batches; + vector original_features; + vector original_classes; + original_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N); + original_batches = vector((int*)PyArray_DATA(batches_array), (int*)PyArray_DATA(batches_array) + Nb); + if (use_feature) + original_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N * fdim); + if (use_classes) + original_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N * ldim); + + // Subsample + vector subsampled_points; + vector subsampled_features; + vector subsampled_classes; + vector subsampled_batches; + batch_grid_subsampling(original_points, + subsampled_points, + original_features, + subsampled_features, + original_classes, + subsampled_classes, + original_batches, + subsampled_batches, + sampleDl, + max_p); + + // Check result + if (subsampled_points.size() < 1) + { + PyErr_SetString(PyExc_RuntimeError, "Error"); + return NULL; + } + + // Manage outputs + // ************** + + // Dimension of input containers + npy_intp* point_dims = new npy_intp[2]; + point_dims[0] = subsampled_points.size(); + point_dims[1] = 3; + npy_intp* feature_dims = new npy_intp[2]; + feature_dims[0] = subsampled_points.size(); + feature_dims[1] = fdim; + npy_intp* classes_dims = new npy_intp[2]; + classes_dims[0] = subsampled_points.size(); + classes_dims[1] = ldim; + npy_intp* batches_dims = new npy_intp[1]; + batches_dims[0] = Nb; + + // Create output array + PyObject* res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT); + PyObject* res_batches_obj = PyArray_SimpleNew(1, batches_dims, NPY_INT); + PyObject* res_features_obj = NULL; + PyObject* res_classes_obj = NULL; + PyObject* ret = NULL; + + // Fill output array with values + size_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float); + memcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes); + size_in_bytes = Nb * sizeof(int); + memcpy(PyArray_DATA(res_batches_obj), subsampled_batches.data(), size_in_bytes); + if (use_feature) + { + size_in_bytes = subsampled_points.size() * fdim * sizeof(float); + res_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT); + memcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes); + } + if (use_classes) + { + size_in_bytes = subsampled_points.size() * ldim * sizeof(int); + res_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT); + memcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes); + } + + + // Merge results + if (use_feature && use_classes) + ret = Py_BuildValue("NNNN", res_points_obj, res_batches_obj, res_features_obj, res_classes_obj); + else if (use_feature) + ret = Py_BuildValue("NNN", res_points_obj, res_batches_obj, res_features_obj); + else if (use_classes) + ret = Py_BuildValue("NNN", res_points_obj, res_batches_obj, res_classes_obj); + else + ret = Py_BuildValue("NN", res_points_obj, res_batches_obj); + + // Clean up + // ******** + + Py_DECREF(points_array); + Py_DECREF(batches_array); + Py_XDECREF(features_array); + Py_XDECREF(classes_array); + + return ret; +} + +// Definition of the subsample method +// **************************************** + +static PyObject* cloud_subsampling(PyObject* self, PyObject* args, PyObject* keywds) +{ + + // Manage inputs + // ************* + + // Args containers + PyObject* points_obj = NULL; + PyObject* features_obj = NULL; + PyObject* classes_obj = NULL; + + // Keywords containers + static char* kwlist[] = { "points", "features", "classes", "sampleDl", "method", "verbose", NULL }; + float sampleDl = 0.1; + const char* method_buffer = "barycenters"; + int verbose = 0; + + // Parse the input + if (!PyArg_ParseTupleAndKeywords(args, keywds, "O|$OOfsi", kwlist, &points_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &verbose)) + { + PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); + return NULL; + } + + // Get the method argument + string method(method_buffer); + + // Interpret method + if (method.compare("barycenters") && method.compare("voxelcenters")) + { + PyErr_SetString(PyExc_RuntimeError, "Error parsing method. Valid method names are \"barycenters\" and \"voxelcenters\" "); + return NULL; + } + + // Check if using features or classes + bool use_feature = true, use_classes = true; + if (features_obj == NULL) + use_feature = false; + if (classes_obj == NULL) + use_classes = false; + + // Interpret the input objects as numpy arrays. + PyObject* points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY); + PyObject* features_array = NULL; + PyObject* classes_array = NULL; + if (use_feature) + features_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY); + if (use_classes) + classes_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY); + + // Verify data was load correctly. + if (points_array == NULL) + { + Py_XDECREF(points_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Error converting input points to numpy arrays of type float32"); + return NULL; + } + if (use_feature && features_array == NULL) + { + Py_XDECREF(points_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Error converting input features to numpy arrays of type float32"); + return NULL; + } + if (use_classes && classes_array == NULL) + { + Py_XDECREF(points_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Error converting input classes to numpy arrays of type int32"); + return NULL; + } + + // Check that the input array respect the dims + if ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3) + { + Py_XDECREF(points_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : points.shape is not (N, 3)"); + return NULL; + } + if (use_feature && ((int)PyArray_NDIM(features_array) != 2)) + { + Py_XDECREF(points_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); + return NULL; + } + + if (use_classes && (int)PyArray_NDIM(classes_array) > 2) + { + Py_XDECREF(points_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); + return NULL; + } + + // Number of points + int N = (int)PyArray_DIM(points_array, 0); + + // Dimension of the features + int fdim = 0; + if (use_feature) + fdim = (int)PyArray_DIM(features_array, 1); + + //Dimension of labels + int ldim = 1; + if (use_classes && (int)PyArray_NDIM(classes_array) == 2) + ldim = (int)PyArray_DIM(classes_array, 1); + + // Check that the input array respect the number of points + if (use_feature && (int)PyArray_DIM(features_array, 0) != N) + { + Py_XDECREF(points_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); + return NULL; + } + if (use_classes && (int)PyArray_DIM(classes_array, 0) != N) + { + Py_XDECREF(points_array); + Py_XDECREF(classes_array); + Py_XDECREF(features_array); + PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); + return NULL; + } + + + // Call the C++ function + // ********************* + + // Create pyramid + if (verbose > 0) + cout << "Computing cloud pyramid with support points: " << endl; + + + // Convert PyArray to Cloud C++ class + vector original_points; + vector original_features; + vector original_classes; + original_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N); + if (use_feature) + original_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N * fdim); + if (use_classes) + original_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N * ldim); + + // Subsample + vector subsampled_points; + vector subsampled_features; + vector subsampled_classes; + grid_subsampling(original_points, + subsampled_points, + original_features, + subsampled_features, + original_classes, + subsampled_classes, + sampleDl, + verbose); + + // Check result + if (subsampled_points.size() < 1) + { + PyErr_SetString(PyExc_RuntimeError, "Error"); + return NULL; + } + + // Manage outputs + // ************** + + // Dimension of input containers + npy_intp* point_dims = new npy_intp[2]; + point_dims[0] = subsampled_points.size(); + point_dims[1] = 3; + npy_intp* feature_dims = new npy_intp[2]; + feature_dims[0] = subsampled_points.size(); + feature_dims[1] = fdim; + npy_intp* classes_dims = new npy_intp[2]; + classes_dims[0] = subsampled_points.size(); + classes_dims[1] = ldim; + + // Create output array + PyObject* res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT); + PyObject* res_features_obj = NULL; + PyObject* res_classes_obj = NULL; + PyObject* ret = NULL; + + // Fill output array with values + size_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float); + memcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes); + if (use_feature) + { + size_in_bytes = subsampled_points.size() * fdim * sizeof(float); + res_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT); + memcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes); + } + if (use_classes) + { + size_in_bytes = subsampled_points.size() * ldim * sizeof(int); + res_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT); + memcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes); + } + + + // Merge results + if (use_feature && use_classes) + ret = Py_BuildValue("NNN", res_points_obj, res_features_obj, res_classes_obj); + else if (use_feature) + ret = Py_BuildValue("NN", res_points_obj, res_features_obj); + else if (use_classes) + ret = Py_BuildValue("NN", res_points_obj, res_classes_obj); + else + ret = Py_BuildValue("N", res_points_obj); + + // Clean up + // ******** + + Py_DECREF(points_array); + Py_XDECREF(features_array); + Py_XDECREF(classes_array); + + return ret; +} \ No newline at end of file diff --git a/cpp_wrappers/cpp_utils/cloud/cloud.cpp b/cpp_wrappers/cpp_utils/cloud/cloud.cpp new file mode 100644 index 0000000..c285140 --- /dev/null +++ b/cpp_wrappers/cpp_utils/cloud/cloud.cpp @@ -0,0 +1,67 @@ +// +// +// 0==========================0 +// | Local feature test | +// 0==========================0 +// +// version 1.0 : +// > +// +//--------------------------------------------------- +// +// Cloud source : +// Define usefull Functions/Methods +// +//---------------------------------------------------- +// +// Hugues THOMAS - 10/02/2017 +// + + +#include "cloud.h" + + +// Getters +// ******* + +PointXYZ max_point(std::vector points) +{ + // Initialize limits + PointXYZ maxP(points[0]); + + // Loop over all points + for (auto p : points) + { + if (p.x > maxP.x) + maxP.x = p.x; + + if (p.y > maxP.y) + maxP.y = p.y; + + if (p.z > maxP.z) + maxP.z = p.z; + } + + return maxP; +} + +PointXYZ min_point(std::vector points) +{ + // Initialize limits + PointXYZ minP(points[0]); + + // Loop over all points + for (auto p : points) + { + if (p.x < minP.x) + minP.x = p.x; + + if (p.y < minP.y) + minP.y = p.y; + + if (p.z < minP.z) + minP.z = p.z; + } + + return minP; +} \ No newline at end of file diff --git a/cpp_wrappers/cpp_utils/cloud/cloud.h b/cpp_wrappers/cpp_utils/cloud/cloud.h new file mode 100644 index 0000000..99d4e19 --- /dev/null +++ b/cpp_wrappers/cpp_utils/cloud/cloud.h @@ -0,0 +1,185 @@ +// +// +// 0==========================0 +// | Local feature test | +// 0==========================0 +// +// version 1.0 : +// > +// +//--------------------------------------------------- +// +// Cloud header +// +//---------------------------------------------------- +// +// Hugues THOMAS - 10/02/2017 +// + + +# pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + + + + +// Point class +// *********** + + +class PointXYZ +{ +public: + + // Elements + // ******** + + float x, y, z; + + + // Methods + // ******* + + // Constructor + PointXYZ() { x = 0; y = 0; z = 0; } + PointXYZ(float x0, float y0, float z0) { x = x0; y = y0; z = z0; } + + // array type accessor + float operator [] (int i) const + { + if (i == 0) return x; + else if (i == 1) return y; + else return z; + } + + // opperations + float dot(const PointXYZ P) const + { + return x * P.x + y * P.y + z * P.z; + } + + float sq_norm() + { + return x*x + y*y + z*z; + } + + PointXYZ cross(const PointXYZ P) const + { + return PointXYZ(y*P.z - z*P.y, z*P.x - x*P.z, x*P.y - y*P.x); + } + + PointXYZ& operator+=(const PointXYZ& P) + { + x += P.x; + y += P.y; + z += P.z; + return *this; + } + + PointXYZ& operator-=(const PointXYZ& P) + { + x -= P.x; + y -= P.y; + z -= P.z; + return *this; + } + + PointXYZ& operator*=(const float& a) + { + x *= a; + y *= a; + z *= a; + return *this; + } +}; + + +// Point Opperations +// ***************** + +inline PointXYZ operator + (const PointXYZ A, const PointXYZ B) +{ + return PointXYZ(A.x + B.x, A.y + B.y, A.z + B.z); +} + +inline PointXYZ operator - (const PointXYZ A, const PointXYZ B) +{ + return PointXYZ(A.x - B.x, A.y - B.y, A.z - B.z); +} + +inline PointXYZ operator * (const PointXYZ P, const float a) +{ + return PointXYZ(P.x * a, P.y * a, P.z * a); +} + +inline PointXYZ operator * (const float a, const PointXYZ P) +{ + return PointXYZ(P.x * a, P.y * a, P.z * a); +} + +inline std::ostream& operator << (std::ostream& os, const PointXYZ P) +{ + return os << "[" << P.x << ", " << P.y << ", " << P.z << "]"; +} + +inline bool operator == (const PointXYZ A, const PointXYZ B) +{ + return A.x == B.x && A.y == B.y && A.z == B.z; +} + +inline PointXYZ floor(const PointXYZ P) +{ + return PointXYZ(std::floor(P.x), std::floor(P.y), std::floor(P.z)); +} + + +PointXYZ max_point(std::vector points); +PointXYZ min_point(std::vector points); + + +struct PointCloud +{ + + std::vector pts; + + // Must return the number of data points + inline size_t kdtree_get_point_count() const { return pts.size(); } + + // Returns the dim'th component of the idx'th point in the class: + // Since this is inlined and the "dim" argument is typically an immediate value, the + // "if/else's" are actually solved at compile time. + inline float kdtree_get_pt(const size_t idx, const size_t dim) const + { + if (dim == 0) return pts[idx].x; + else if (dim == 1) return pts[idx].y; + else return pts[idx].z; + } + + // Optional bounding-box computation: return false to default to a standard bbox computation loop. + // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. + // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) + template + bool kdtree_get_bbox(BBOX& /* bb */) const { return false; } + +}; + + + + + + + + + + + diff --git a/cpp_wrappers/cpp_utils/nanoflann/nanoflann.hpp b/cpp_wrappers/cpp_utils/nanoflann/nanoflann.hpp new file mode 100644 index 0000000..8d2ab6c --- /dev/null +++ b/cpp_wrappers/cpp_utils/nanoflann/nanoflann.hpp @@ -0,0 +1,2043 @@ +/*********************************************************************** + * Software License Agreement (BSD License) + * + * Copyright 2008-2009 Marius Muja (mariusm@cs.ubc.ca). All rights reserved. + * Copyright 2008-2009 David G. Lowe (lowe@cs.ubc.ca). All rights reserved. + * Copyright 2011-2016 Jose Luis Blanco (joseluisblancoc@gmail.com). + * All rights reserved. + * + * THE BSD LICENSE + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met: + * + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * + * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES + * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. + * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, + * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT + * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF + * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + *************************************************************************/ + +/** \mainpage nanoflann C++ API documentation + * nanoflann is a C++ header-only library for building KD-Trees, mostly + * optimized for 2D or 3D point clouds. + * + * nanoflann does not require compiling or installing, just an + * #include in your code. + * + * See: + * - C++ API organized by modules + * - Online README + * - Doxygen + * documentation + */ + +#ifndef NANOFLANN_HPP_ +#define NANOFLANN_HPP_ + +#include +#include +#include +#include // for abs() +#include // for fwrite() +#include // for abs() +#include +#include // std::reference_wrapper +#include +#include + +/** Library version: 0xMmP (M=Major,m=minor,P=patch) */ +#define NANOFLANN_VERSION 0x130 + +// Avoid conflicting declaration of min/max macros in windows headers +#if !defined(NOMINMAX) && \ + (defined(_WIN32) || defined(_WIN32_) || defined(WIN32) || defined(_WIN64)) +#define NOMINMAX +#ifdef max +#undef max +#undef min +#endif +#endif + +namespace nanoflann { +/** @addtogroup nanoflann_grp nanoflann C++ library for ANN + * @{ */ + +/** the PI constant (required to avoid MSVC missing symbols) */ +template T pi_const() { + return static_cast(3.14159265358979323846); +} + +/** + * Traits if object is resizable and assignable (typically has a resize | assign + * method) + */ +template struct has_resize : std::false_type {}; + +template +struct has_resize().resize(1), 0)> + : std::true_type {}; + +template struct has_assign : std::false_type {}; + +template +struct has_assign().assign(1, 0), 0)> + : std::true_type {}; + +/** + * Free function to resize a resizable object + */ +template +inline typename std::enable_if::value, void>::type +resize(Container &c, const size_t nElements) { + c.resize(nElements); +} + +/** + * Free function that has no effects on non resizable containers (e.g. + * std::array) It raises an exception if the expected size does not match + */ +template +inline typename std::enable_if::value, void>::type +resize(Container &c, const size_t nElements) { + if (nElements != c.size()) + throw std::logic_error("Try to change the size of a std::array."); +} + +/** + * Free function to assign to a container + */ +template +inline typename std::enable_if::value, void>::type +assign(Container &c, const size_t nElements, const T &value) { + c.assign(nElements, value); +} + +/** + * Free function to assign to a std::array + */ +template +inline typename std::enable_if::value, void>::type +assign(Container &c, const size_t nElements, const T &value) { + for (size_t i = 0; i < nElements; i++) + c[i] = value; +} + +/** @addtogroup result_sets_grp Result set classes + * @{ */ +template +class KNNResultSet { +public: + typedef _DistanceType DistanceType; + typedef _IndexType IndexType; + typedef _CountType CountType; + +private: + IndexType *indices; + DistanceType *dists; + CountType capacity; + CountType count; + +public: + inline KNNResultSet(CountType capacity_) + : indices(0), dists(0), capacity(capacity_), count(0) {} + + inline void init(IndexType *indices_, DistanceType *dists_) { + indices = indices_; + dists = dists_; + count = 0; + if (capacity) + dists[capacity - 1] = (std::numeric_limits::max)(); + } + + inline CountType size() const { return count; } + + inline bool full() const { return count == capacity; } + + /** + * Called during search to add an element matching the criteria. + * @return true if the search should be continued, false if the results are + * sufficient + */ + inline bool addPoint(DistanceType dist, IndexType index) { + CountType i; + for (i = count; i > 0; --i) { +#ifdef NANOFLANN_FIRST_MATCH // If defined and two points have the same + // distance, the one with the lowest-index will be + // returned first. + if ((dists[i - 1] > dist) || + ((dist == dists[i - 1]) && (indices[i - 1] > index))) { +#else + if (dists[i - 1] > dist) { +#endif + if (i < capacity) { + dists[i] = dists[i - 1]; + indices[i] = indices[i - 1]; + } + } else + break; + } + if (i < capacity) { + dists[i] = dist; + indices[i] = index; + } + if (count < capacity) + count++; + + // tell caller that the search shall continue + return true; + } + + inline DistanceType worstDist() const { return dists[capacity - 1]; } +}; + +/** operator "<" for std::sort() */ +struct IndexDist_Sorter { + /** PairType will be typically: std::pair */ + template + inline bool operator()(const PairType &p1, const PairType &p2) const { + return p1.second < p2.second; + } +}; + +/** + * A result-set class used when performing a radius based search. + */ +template +class RadiusResultSet { +public: + typedef _DistanceType DistanceType; + typedef _IndexType IndexType; + +public: + const DistanceType radius; + + std::vector> &m_indices_dists; + + inline RadiusResultSet( + DistanceType radius_, + std::vector> &indices_dists) + : radius(radius_), m_indices_dists(indices_dists) { + init(); + } + + inline void init() { clear(); } + inline void clear() { m_indices_dists.clear(); } + + inline size_t size() const { return m_indices_dists.size(); } + + inline bool full() const { return true; } + + /** + * Called during search to add an element matching the criteria. + * @return true if the search should be continued, false if the results are + * sufficient + */ + inline bool addPoint(DistanceType dist, IndexType index) { + if (dist < radius) + m_indices_dists.push_back(std::make_pair(index, dist)); + return true; + } + + inline DistanceType worstDist() const { return radius; } + + /** + * Find the worst result (furtherest neighbor) without copying or sorting + * Pre-conditions: size() > 0 + */ + std::pair worst_item() const { + if (m_indices_dists.empty()) + throw std::runtime_error("Cannot invoke RadiusResultSet::worst_item() on " + "an empty list of results."); + typedef + typename std::vector>::const_iterator + DistIt; + DistIt it = std::max_element(m_indices_dists.begin(), m_indices_dists.end(), + IndexDist_Sorter()); + return *it; + } +}; + +/** @} */ + +/** @addtogroup loadsave_grp Load/save auxiliary functions + * @{ */ +template +void save_value(FILE *stream, const T &value, size_t count = 1) { + fwrite(&value, sizeof(value), count, stream); +} + +template +void save_value(FILE *stream, const std::vector &value) { + size_t size = value.size(); + fwrite(&size, sizeof(size_t), 1, stream); + fwrite(&value[0], sizeof(T), size, stream); +} + +template +void load_value(FILE *stream, T &value, size_t count = 1) { + size_t read_cnt = fread(&value, sizeof(value), count, stream); + if (read_cnt != count) { + throw std::runtime_error("Cannot read from file"); + } +} + +template void load_value(FILE *stream, std::vector &value) { + size_t size; + size_t read_cnt = fread(&size, sizeof(size_t), 1, stream); + if (read_cnt != 1) { + throw std::runtime_error("Cannot read from file"); + } + value.resize(size); + read_cnt = fread(&value[0], sizeof(T), size, stream); + if (read_cnt != size) { + throw std::runtime_error("Cannot read from file"); + } +} +/** @} */ + +/** @addtogroup metric_grp Metric (distance) classes + * @{ */ + +struct Metric {}; + +/** Manhattan distance functor (generic version, optimized for + * high-dimensionality data sets). Corresponding distance traits: + * nanoflann::metric_L1 \tparam T Type of the elements (e.g. double, float, + * uint8_t) \tparam _DistanceType Type of distance variables (must be signed) + * (e.g. float, double, int64_t) + */ +template +struct L1_Adaptor { + typedef T ElementType; + typedef _DistanceType DistanceType; + + const DataSource &data_source; + + L1_Adaptor(const DataSource &_data_source) : data_source(_data_source) {} + + inline DistanceType evalMetric(const T *a, const size_t b_idx, size_t size, + DistanceType worst_dist = -1) const { + DistanceType result = DistanceType(); + const T *last = a + size; + const T *lastgroup = last - 3; + size_t d = 0; + + /* Process 4 items with each loop for efficiency. */ + while (a < lastgroup) { + const DistanceType diff0 = + std::abs(a[0] - data_source.kdtree_get_pt(b_idx, d++)); + const DistanceType diff1 = + std::abs(a[1] - data_source.kdtree_get_pt(b_idx, d++)); + const DistanceType diff2 = + std::abs(a[2] - data_source.kdtree_get_pt(b_idx, d++)); + const DistanceType diff3 = + std::abs(a[3] - data_source.kdtree_get_pt(b_idx, d++)); + result += diff0 + diff1 + diff2 + diff3; + a += 4; + if ((worst_dist > 0) && (result > worst_dist)) { + return result; + } + } + /* Process last 0-3 components. Not needed for standard vector lengths. */ + while (a < last) { + result += std::abs(*a++ - data_source.kdtree_get_pt(b_idx, d++)); + } + return result; + } + + template + inline DistanceType accum_dist(const U a, const V b, const size_t) const { + return std::abs(a - b); + } +}; + +/** Squared Euclidean distance functor (generic version, optimized for + * high-dimensionality data sets). Corresponding distance traits: + * nanoflann::metric_L2 \tparam T Type of the elements (e.g. double, float, + * uint8_t) \tparam _DistanceType Type of distance variables (must be signed) + * (e.g. float, double, int64_t) + */ +template +struct L2_Adaptor { + typedef T ElementType; + typedef _DistanceType DistanceType; + + const DataSource &data_source; + + L2_Adaptor(const DataSource &_data_source) : data_source(_data_source) {} + + inline DistanceType evalMetric(const T *a, const size_t b_idx, size_t size, + DistanceType worst_dist = -1) const { + DistanceType result = DistanceType(); + const T *last = a + size; + const T *lastgroup = last - 3; + size_t d = 0; + + /* Process 4 items with each loop for efficiency. */ + while (a < lastgroup) { + const DistanceType diff0 = a[0] - data_source.kdtree_get_pt(b_idx, d++); + const DistanceType diff1 = a[1] - data_source.kdtree_get_pt(b_idx, d++); + const DistanceType diff2 = a[2] - data_source.kdtree_get_pt(b_idx, d++); + const DistanceType diff3 = a[3] - data_source.kdtree_get_pt(b_idx, d++); + result += diff0 * diff0 + diff1 * diff1 + diff2 * diff2 + diff3 * diff3; + a += 4; + if ((worst_dist > 0) && (result > worst_dist)) { + return result; + } + } + /* Process last 0-3 components. Not needed for standard vector lengths. */ + while (a < last) { + const DistanceType diff0 = *a++ - data_source.kdtree_get_pt(b_idx, d++); + result += diff0 * diff0; + } + return result; + } + + template + inline DistanceType accum_dist(const U a, const V b, const size_t) const { + return (a - b) * (a - b); + } +}; + +/** Squared Euclidean (L2) distance functor (suitable for low-dimensionality + * datasets, like 2D or 3D point clouds) Corresponding distance traits: + * nanoflann::metric_L2_Simple \tparam T Type of the elements (e.g. double, + * float, uint8_t) \tparam _DistanceType Type of distance variables (must be + * signed) (e.g. float, double, int64_t) + */ +template +struct L2_Simple_Adaptor { + typedef T ElementType; + typedef _DistanceType DistanceType; + + const DataSource &data_source; + + L2_Simple_Adaptor(const DataSource &_data_source) + : data_source(_data_source) {} + + inline DistanceType evalMetric(const T *a, const size_t b_idx, + size_t size) const { + DistanceType result = DistanceType(); + for (size_t i = 0; i < size; ++i) { + const DistanceType diff = a[i] - data_source.kdtree_get_pt(b_idx, i); + result += diff * diff; + } + return result; + } + + template + inline DistanceType accum_dist(const U a, const V b, const size_t) const { + return (a - b) * (a - b); + } +}; + +/** SO2 distance functor + * Corresponding distance traits: nanoflann::metric_SO2 + * \tparam T Type of the elements (e.g. double, float) + * \tparam _DistanceType Type of distance variables (must be signed) (e.g. + * float, double) orientation is constrained to be in [-pi, pi] + */ +template +struct SO2_Adaptor { + typedef T ElementType; + typedef _DistanceType DistanceType; + + const DataSource &data_source; + + SO2_Adaptor(const DataSource &_data_source) : data_source(_data_source) {} + + inline DistanceType evalMetric(const T *a, const size_t b_idx, + size_t size) const { + return accum_dist(a[size - 1], data_source.kdtree_get_pt(b_idx, size - 1), + size - 1); + } + + /** Note: this assumes that input angles are already in the range [-pi,pi] */ + template + inline DistanceType accum_dist(const U a, const V b, const size_t) const { + DistanceType result = DistanceType(), PI = pi_const(); + result = b - a; + if (result > PI) + result -= 2 * PI; + else if (result < -PI) + result += 2 * PI; + return result; + } +}; + +/** SO3 distance functor (Uses L2_Simple) + * Corresponding distance traits: nanoflann::metric_SO3 + * \tparam T Type of the elements (e.g. double, float) + * \tparam _DistanceType Type of distance variables (must be signed) (e.g. + * float, double) + */ +template +struct SO3_Adaptor { + typedef T ElementType; + typedef _DistanceType DistanceType; + + L2_Simple_Adaptor distance_L2_Simple; + + SO3_Adaptor(const DataSource &_data_source) + : distance_L2_Simple(_data_source) {} + + inline DistanceType evalMetric(const T *a, const size_t b_idx, + size_t size) const { + return distance_L2_Simple.evalMetric(a, b_idx, size); + } + + template + inline DistanceType accum_dist(const U a, const V b, const size_t idx) const { + return distance_L2_Simple.accum_dist(a, b, idx); + } +}; + +/** Metaprogramming helper traits class for the L1 (Manhattan) metric */ +struct metric_L1 : public Metric { + template struct traits { + typedef L1_Adaptor distance_t; + }; +}; +/** Metaprogramming helper traits class for the L2 (Euclidean) metric */ +struct metric_L2 : public Metric { + template struct traits { + typedef L2_Adaptor distance_t; + }; +}; +/** Metaprogramming helper traits class for the L2_simple (Euclidean) metric */ +struct metric_L2_Simple : public Metric { + template struct traits { + typedef L2_Simple_Adaptor distance_t; + }; +}; +/** Metaprogramming helper traits class for the SO3_InnerProdQuat metric */ +struct metric_SO2 : public Metric { + template struct traits { + typedef SO2_Adaptor distance_t; + }; +}; +/** Metaprogramming helper traits class for the SO3_InnerProdQuat metric */ +struct metric_SO3 : public Metric { + template struct traits { + typedef SO3_Adaptor distance_t; + }; +}; + +/** @} */ + +/** @addtogroup param_grp Parameter structs + * @{ */ + +/** Parameters (see README.md) */ +struct KDTreeSingleIndexAdaptorParams { + KDTreeSingleIndexAdaptorParams(size_t _leaf_max_size = 10) + : leaf_max_size(_leaf_max_size) {} + + size_t leaf_max_size; +}; + +/** Search options for KDTreeSingleIndexAdaptor::findNeighbors() */ +struct SearchParams { + /** Note: The first argument (checks_IGNORED_) is ignored, but kept for + * compatibility with the FLANN interface */ + SearchParams(int checks_IGNORED_ = 32, float eps_ = 0, bool sorted_ = true) + : checks(checks_IGNORED_), eps(eps_), sorted(sorted_) {} + + int checks; //!< Ignored parameter (Kept for compatibility with the FLANN + //!< interface). + float eps; //!< search for eps-approximate neighbours (default: 0) + bool sorted; //!< only for radius search, require neighbours sorted by + //!< distance (default: true) +}; +/** @} */ + +/** @addtogroup memalloc_grp Memory allocation + * @{ */ + +/** + * Allocates (using C's malloc) a generic type T. + * + * Params: + * count = number of instances to allocate. + * Returns: pointer (of type T*) to memory buffer + */ +template inline T *allocate(size_t count = 1) { + T *mem = static_cast(::malloc(sizeof(T) * count)); + return mem; +} + +/** + * Pooled storage allocator + * + * The following routines allow for the efficient allocation of storage in + * small chunks from a specified pool. Rather than allowing each structure + * to be freed individually, an entire pool of storage is freed at once. + * This method has two advantages over just using malloc() and free(). First, + * it is far more efficient for allocating small objects, as there is + * no overhead for remembering all the information needed to free each + * object or consolidating fragmented memory. Second, the decision about + * how long to keep an object is made at the time of allocation, and there + * is no need to track down all the objects to free them. + * + */ + +const size_t WORDSIZE = 16; +const size_t BLOCKSIZE = 8192; + +class PooledAllocator { + /* We maintain memory alignment to word boundaries by requiring that all + allocations be in multiples of the machine wordsize. */ + /* Size of machine word in bytes. Must be power of 2. */ + /* Minimum number of bytes requested at a time from the system. Must be + * multiple of WORDSIZE. */ + + size_t remaining; /* Number of bytes left in current block of storage. */ + void *base; /* Pointer to base of current block of storage. */ + void *loc; /* Current location in block to next allocate memory. */ + + void internal_init() { + remaining = 0; + base = NULL; + usedMemory = 0; + wastedMemory = 0; + } + +public: + size_t usedMemory; + size_t wastedMemory; + + /** + Default constructor. Initializes a new pool. + */ + PooledAllocator() { internal_init(); } + + /** + * Destructor. Frees all the memory allocated in this pool. + */ + ~PooledAllocator() { free_all(); } + + /** Frees all allocated memory chunks */ + void free_all() { + while (base != NULL) { + void *prev = + *(static_cast(base)); /* Get pointer to prev block. */ + ::free(base); + base = prev; + } + internal_init(); + } + + /** + * Returns a pointer to a piece of new memory of the given size in bytes + * allocated from the pool. + */ + void *malloc(const size_t req_size) { + /* Round size up to a multiple of wordsize. The following expression + only works for WORDSIZE that is a power of 2, by masking last bits of + incremented size to zero. + */ + const size_t size = (req_size + (WORDSIZE - 1)) & ~(WORDSIZE - 1); + + /* Check whether a new block must be allocated. Note that the first word + of a block is reserved for a pointer to the previous block. + */ + if (size > remaining) { + + wastedMemory += remaining; + + /* Allocate new storage. */ + const size_t blocksize = + (size + sizeof(void *) + (WORDSIZE - 1) > BLOCKSIZE) + ? size + sizeof(void *) + (WORDSIZE - 1) + : BLOCKSIZE; + + // use the standard C malloc to allocate memory + void *m = ::malloc(blocksize); + if (!m) { + fprintf(stderr, "Failed to allocate memory.\n"); + return NULL; + } + + /* Fill first word of new block with pointer to previous block. */ + static_cast(m)[0] = base; + base = m; + + size_t shift = 0; + // int size_t = (WORDSIZE - ( (((size_t)m) + sizeof(void*)) & + // (WORDSIZE-1))) & (WORDSIZE-1); + + remaining = blocksize - sizeof(void *) - shift; + loc = (static_cast(m) + sizeof(void *) + shift); + } + void *rloc = loc; + loc = static_cast(loc) + size; + remaining -= size; + + usedMemory += size; + + return rloc; + } + + /** + * Allocates (using this pool) a generic type T. + * + * Params: + * count = number of instances to allocate. + * Returns: pointer (of type T*) to memory buffer + */ + template T *allocate(const size_t count = 1) { + T *mem = static_cast(this->malloc(sizeof(T) * count)); + return mem; + } +}; +/** @} */ + +/** @addtogroup nanoflann_metaprog_grp Auxiliary metaprogramming stuff + * @{ */ + +/** Used to declare fixed-size arrays when DIM>0, dynamically-allocated vectors + * when DIM=-1. Fixed size version for a generic DIM: + */ +template struct array_or_vector_selector { + typedef std::array container_t; +}; +/** Dynamic size version */ +template struct array_or_vector_selector<-1, T> { + typedef std::vector container_t; +}; + +/** @} */ + +/** kd-tree base-class + * + * Contains the member functions common to the classes KDTreeSingleIndexAdaptor + * and KDTreeSingleIndexDynamicAdaptor_. + * + * \tparam Derived The name of the class which inherits this class. + * \tparam DatasetAdaptor The user-provided adaptor (see comments above). + * \tparam Distance The distance metric to use, these are all classes derived + * from nanoflann::Metric \tparam DIM Dimensionality of data points (e.g. 3 for + * 3D points) \tparam IndexType Will be typically size_t or int + */ + +template +class KDTreeBaseClass { + +public: + /** Frees the previously-built index. Automatically called within + * buildIndex(). */ + void freeIndex(Derived &obj) { + obj.pool.free_all(); + obj.root_node = NULL; + obj.m_size_at_index_build = 0; + } + + typedef typename Distance::ElementType ElementType; + typedef typename Distance::DistanceType DistanceType; + + /*--------------------- Internal Data Structures --------------------------*/ + struct Node { + /** Union used because a node can be either a LEAF node or a non-leaf node, + * so both data fields are never used simultaneously */ + union { + struct leaf { + IndexType left, right; //!< Indices of points in leaf node + } lr; + struct nonleaf { + int divfeat; //!< Dimension used for subdivision. + DistanceType divlow, divhigh; //!< The values used for subdivision. + } sub; + } node_type; + Node *child1, *child2; //!< Child nodes (both=NULL mean its a leaf node) + }; + + typedef Node *NodePtr; + + struct Interval { + ElementType low, high; + }; + + /** + * Array of indices to vectors in the dataset. + */ + std::vector vind; + + NodePtr root_node; + + size_t m_leaf_max_size; + + size_t m_size; //!< Number of current points in the dataset + size_t m_size_at_index_build; //!< Number of points in the dataset when the + //!< index was built + int dim; //!< Dimensionality of each data point + + /** Define "BoundingBox" as a fixed-size or variable-size container depending + * on "DIM" */ + typedef + typename array_or_vector_selector::container_t BoundingBox; + + /** Define "distance_vector_t" as a fixed-size or variable-size container + * depending on "DIM" */ + typedef typename array_or_vector_selector::container_t + distance_vector_t; + + /** The KD-tree used to find neighbours */ + + BoundingBox root_bbox; + + /** + * Pooled memory allocator. + * + * Using a pooled memory allocator is more efficient + * than allocating memory directly when there is a large + * number small of memory allocations. + */ + PooledAllocator pool; + + /** Returns number of points in dataset */ + size_t size(const Derived &obj) const { return obj.m_size; } + + /** Returns the length of each point in the dataset */ + size_t veclen(const Derived &obj) { + return static_cast(DIM > 0 ? DIM : obj.dim); + } + + /// Helper accessor to the dataset points: + inline ElementType dataset_get(const Derived &obj, size_t idx, + int component) const { + return obj.dataset.kdtree_get_pt(idx, component); + } + + /** + * Computes the inde memory usage + * Returns: memory used by the index + */ + size_t usedMemory(Derived &obj) { + return obj.pool.usedMemory + obj.pool.wastedMemory + + obj.dataset.kdtree_get_point_count() * + sizeof(IndexType); // pool memory and vind array memory + } + + void computeMinMax(const Derived &obj, IndexType *ind, IndexType count, + int element, ElementType &min_elem, + ElementType &max_elem) { + min_elem = dataset_get(obj, ind[0], element); + max_elem = dataset_get(obj, ind[0], element); + for (IndexType i = 1; i < count; ++i) { + ElementType val = dataset_get(obj, ind[i], element); + if (val < min_elem) + min_elem = val; + if (val > max_elem) + max_elem = val; + } + } + + /** + * Create a tree node that subdivides the list of vecs from vind[first] + * to vind[last]. The routine is called recursively on each sublist. + * + * @param left index of the first vector + * @param right index of the last vector + */ + NodePtr divideTree(Derived &obj, const IndexType left, const IndexType right, + BoundingBox &bbox) { + NodePtr node = obj.pool.template allocate(); // allocate memory + + /* If too few exemplars remain, then make this a leaf node. */ + if ((right - left) <= static_cast(obj.m_leaf_max_size)) { + node->child1 = node->child2 = NULL; /* Mark as leaf node. */ + node->node_type.lr.left = left; + node->node_type.lr.right = right; + + // compute bounding-box of leaf points + for (int i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i) { + bbox[i].low = dataset_get(obj, obj.vind[left], i); + bbox[i].high = dataset_get(obj, obj.vind[left], i); + } + for (IndexType k = left + 1; k < right; ++k) { + for (int i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i) { + if (bbox[i].low > dataset_get(obj, obj.vind[k], i)) + bbox[i].low = dataset_get(obj, obj.vind[k], i); + if (bbox[i].high < dataset_get(obj, obj.vind[k], i)) + bbox[i].high = dataset_get(obj, obj.vind[k], i); + } + } + } else { + IndexType idx; + int cutfeat; + DistanceType cutval; + middleSplit_(obj, &obj.vind[0] + left, right - left, idx, cutfeat, cutval, + bbox); + + node->node_type.sub.divfeat = cutfeat; + + BoundingBox left_bbox(bbox); + left_bbox[cutfeat].high = cutval; + node->child1 = divideTree(obj, left, left + idx, left_bbox); + + BoundingBox right_bbox(bbox); + right_bbox[cutfeat].low = cutval; + node->child2 = divideTree(obj, left + idx, right, right_bbox); + + node->node_type.sub.divlow = left_bbox[cutfeat].high; + node->node_type.sub.divhigh = right_bbox[cutfeat].low; + + for (int i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i) { + bbox[i].low = std::min(left_bbox[i].low, right_bbox[i].low); + bbox[i].high = std::max(left_bbox[i].high, right_bbox[i].high); + } + } + + return node; + } + + void middleSplit_(Derived &obj, IndexType *ind, IndexType count, + IndexType &index, int &cutfeat, DistanceType &cutval, + const BoundingBox &bbox) { + const DistanceType EPS = static_cast(0.00001); + ElementType max_span = bbox[0].high - bbox[0].low; + for (int i = 1; i < (DIM > 0 ? DIM : obj.dim); ++i) { + ElementType span = bbox[i].high - bbox[i].low; + if (span > max_span) { + max_span = span; + } + } + ElementType max_spread = -1; + cutfeat = 0; + for (int i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i) { + ElementType span = bbox[i].high - bbox[i].low; + if (span > (1 - EPS) * max_span) { + ElementType min_elem, max_elem; + computeMinMax(obj, ind, count, i, min_elem, max_elem); + ElementType spread = max_elem - min_elem; + ; + if (spread > max_spread) { + cutfeat = i; + max_spread = spread; + } + } + } + // split in the middle + DistanceType split_val = (bbox[cutfeat].low + bbox[cutfeat].high) / 2; + ElementType min_elem, max_elem; + computeMinMax(obj, ind, count, cutfeat, min_elem, max_elem); + + if (split_val < min_elem) + cutval = min_elem; + else if (split_val > max_elem) + cutval = max_elem; + else + cutval = split_val; + + IndexType lim1, lim2; + planeSplit(obj, ind, count, cutfeat, cutval, lim1, lim2); + + if (lim1 > count / 2) + index = lim1; + else if (lim2 < count / 2) + index = lim2; + else + index = count / 2; + } + + /** + * Subdivide the list of points by a plane perpendicular on axe corresponding + * to the 'cutfeat' dimension at 'cutval' position. + * + * On return: + * dataset[ind[0..lim1-1]][cutfeat]cutval + */ + void planeSplit(Derived &obj, IndexType *ind, const IndexType count, + int cutfeat, DistanceType &cutval, IndexType &lim1, + IndexType &lim2) { + /* Move vector indices for left subtree to front of list. */ + IndexType left = 0; + IndexType right = count - 1; + for (;;) { + while (left <= right && dataset_get(obj, ind[left], cutfeat) < cutval) + ++left; + while (right && left <= right && + dataset_get(obj, ind[right], cutfeat) >= cutval) + --right; + if (left > right || !right) + break; // "!right" was added to support unsigned Index types + std::swap(ind[left], ind[right]); + ++left; + --right; + } + /* If either list is empty, it means that all remaining features + * are identical. Split in the middle to maintain a balanced tree. + */ + lim1 = left; + right = count - 1; + for (;;) { + while (left <= right && dataset_get(obj, ind[left], cutfeat) <= cutval) + ++left; + while (right && left <= right && + dataset_get(obj, ind[right], cutfeat) > cutval) + --right; + if (left > right || !right) + break; // "!right" was added to support unsigned Index types + std::swap(ind[left], ind[right]); + ++left; + --right; + } + lim2 = left; + } + + DistanceType computeInitialDistances(const Derived &obj, + const ElementType *vec, + distance_vector_t &dists) const { + assert(vec); + DistanceType distsq = DistanceType(); + + for (int i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i) { + if (vec[i] < obj.root_bbox[i].low) { + dists[i] = obj.distance.accum_dist(vec[i], obj.root_bbox[i].low, i); + distsq += dists[i]; + } + if (vec[i] > obj.root_bbox[i].high) { + dists[i] = obj.distance.accum_dist(vec[i], obj.root_bbox[i].high, i); + distsq += dists[i]; + } + } + return distsq; + } + + void save_tree(Derived &obj, FILE *stream, NodePtr tree) { + save_value(stream, *tree); + if (tree->child1 != NULL) { + save_tree(obj, stream, tree->child1); + } + if (tree->child2 != NULL) { + save_tree(obj, stream, tree->child2); + } + } + + void load_tree(Derived &obj, FILE *stream, NodePtr &tree) { + tree = obj.pool.template allocate(); + load_value(stream, *tree); + if (tree->child1 != NULL) { + load_tree(obj, stream, tree->child1); + } + if (tree->child2 != NULL) { + load_tree(obj, stream, tree->child2); + } + } + + /** Stores the index in a binary file. + * IMPORTANT NOTE: The set of data points is NOT stored in the file, so when + * loading the index object it must be constructed associated to the same + * source of data points used while building it. See the example: + * examples/saveload_example.cpp \sa loadIndex */ + void saveIndex_(Derived &obj, FILE *stream) { + save_value(stream, obj.m_size); + save_value(stream, obj.dim); + save_value(stream, obj.root_bbox); + save_value(stream, obj.m_leaf_max_size); + save_value(stream, obj.vind); + save_tree(obj, stream, obj.root_node); + } + + /** Loads a previous index from a binary file. + * IMPORTANT NOTE: The set of data points is NOT stored in the file, so the + * index object must be constructed associated to the same source of data + * points used while building the index. See the example: + * examples/saveload_example.cpp \sa loadIndex */ + void loadIndex_(Derived &obj, FILE *stream) { + load_value(stream, obj.m_size); + load_value(stream, obj.dim); + load_value(stream, obj.root_bbox); + load_value(stream, obj.m_leaf_max_size); + load_value(stream, obj.vind); + load_tree(obj, stream, obj.root_node); + } +}; + +/** @addtogroup kdtrees_grp KD-tree classes and adaptors + * @{ */ + +/** kd-tree static index + * + * Contains the k-d trees and other information for indexing a set of points + * for nearest-neighbor matching. + * + * The class "DatasetAdaptor" must provide the following interface (can be + * non-virtual, inlined methods): + * + * \code + * // Must return the number of data poins + * inline size_t kdtree_get_point_count() const { ... } + * + * + * // Must return the dim'th component of the idx'th point in the class: + * inline T kdtree_get_pt(const size_t idx, const size_t dim) const { ... } + * + * // Optional bounding-box computation: return false to default to a standard + * bbox computation loop. + * // Return true if the BBOX was already computed by the class and returned + * in "bb" so it can be avoided to redo it again. + * // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 + * for point clouds) template bool kdtree_get_bbox(BBOX &bb) const + * { + * bb[0].low = ...; bb[0].high = ...; // 0th dimension limits + * bb[1].low = ...; bb[1].high = ...; // 1st dimension limits + * ... + * return true; + * } + * + * \endcode + * + * \tparam DatasetAdaptor The user-provided adaptor (see comments above). + * \tparam Distance The distance metric to use: nanoflann::metric_L1, + * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. \tparam DIM + * Dimensionality of data points (e.g. 3 for 3D points) \tparam IndexType Will + * be typically size_t or int + */ +template +class KDTreeSingleIndexAdaptor + : public KDTreeBaseClass< + KDTreeSingleIndexAdaptor, + Distance, DatasetAdaptor, DIM, IndexType> { +public: + /** Deleted copy constructor*/ + KDTreeSingleIndexAdaptor( + const KDTreeSingleIndexAdaptor + &) = delete; + + /** + * The dataset used by this index + */ + const DatasetAdaptor &dataset; //!< The source of our data + + const KDTreeSingleIndexAdaptorParams index_params; + + Distance distance; + + typedef typename nanoflann::KDTreeBaseClass< + nanoflann::KDTreeSingleIndexAdaptor, + Distance, DatasetAdaptor, DIM, IndexType> + BaseClassRef; + + typedef typename BaseClassRef::ElementType ElementType; + typedef typename BaseClassRef::DistanceType DistanceType; + + typedef typename BaseClassRef::Node Node; + typedef Node *NodePtr; + + typedef typename BaseClassRef::Interval Interval; + /** Define "BoundingBox" as a fixed-size or variable-size container depending + * on "DIM" */ + typedef typename BaseClassRef::BoundingBox BoundingBox; + + /** Define "distance_vector_t" as a fixed-size or variable-size container + * depending on "DIM" */ + typedef typename BaseClassRef::distance_vector_t distance_vector_t; + + /** + * KDTree constructor + * + * Refer to docs in README.md or online in + * https://github.com/jlblancoc/nanoflann + * + * The KD-Tree point dimension (the length of each point in the datase, e.g. 3 + * for 3D points) is determined by means of: + * - The \a DIM template parameter if >0 (highest priority) + * - Otherwise, the \a dimensionality parameter of this constructor. + * + * @param inputData Dataset with the input features + * @param params Basically, the maximum leaf node size + */ + KDTreeSingleIndexAdaptor(const int dimensionality, + const DatasetAdaptor &inputData, + const KDTreeSingleIndexAdaptorParams ¶ms = + KDTreeSingleIndexAdaptorParams()) + : dataset(inputData), index_params(params), distance(inputData) { + BaseClassRef::root_node = NULL; + BaseClassRef::m_size = dataset.kdtree_get_point_count(); + BaseClassRef::m_size_at_index_build = BaseClassRef::m_size; + BaseClassRef::dim = dimensionality; + if (DIM > 0) + BaseClassRef::dim = DIM; + BaseClassRef::m_leaf_max_size = params.leaf_max_size; + + // Create a permutable array of indices to the input vectors. + init_vind(); + } + + /** + * Builds the index + */ + void buildIndex() { + BaseClassRef::m_size = dataset.kdtree_get_point_count(); + BaseClassRef::m_size_at_index_build = BaseClassRef::m_size; + init_vind(); + this->freeIndex(*this); + BaseClassRef::m_size_at_index_build = BaseClassRef::m_size; + if (BaseClassRef::m_size == 0) + return; + computeBoundingBox(BaseClassRef::root_bbox); + BaseClassRef::root_node = + this->divideTree(*this, 0, BaseClassRef::m_size, + BaseClassRef::root_bbox); // construct the tree + } + + /** \name Query methods + * @{ */ + + /** + * Find set of nearest neighbors to vec[0:dim-1]. Their indices are stored + * inside the result object. + * + * Params: + * result = the result object in which the indices of the + * nearest-neighbors are stored vec = the vector for which to search the + * nearest neighbors + * + * \tparam RESULTSET Should be any ResultSet + * \return True if the requested neighbors could be found. + * \sa knnSearch, radiusSearch + */ + template + bool findNeighbors(RESULTSET &result, const ElementType *vec, + const SearchParams &searchParams) const { + assert(vec); + if (this->size(*this) == 0) + return false; + if (!BaseClassRef::root_node) + throw std::runtime_error( + "[nanoflann] findNeighbors() called before building the index."); + float epsError = 1 + searchParams.eps; + + distance_vector_t + dists; // fixed or variable-sized container (depending on DIM) + auto zero = static_cast(0); + assign(dists, (DIM > 0 ? DIM : BaseClassRef::dim), + zero); // Fill it with zeros. + DistanceType distsq = this->computeInitialDistances(*this, vec, dists); + + searchLevel(result, vec, BaseClassRef::root_node, distsq, dists, + epsError); // "count_leaf" parameter removed since was neither + // used nor returned to the user. + + return result.full(); + } + + /** + * Find the "num_closest" nearest neighbors to the \a query_point[0:dim-1]. + * Their indices are stored inside the result object. \sa radiusSearch, + * findNeighbors \note nChecks_IGNORED is ignored but kept for compatibility + * with the original FLANN interface. \return Number `N` of valid points in + * the result set. Only the first `N` entries in `out_indices` and + * `out_distances_sq` will be valid. Return may be less than `num_closest` + * only if the number of elements in the tree is less than `num_closest`. + */ + size_t knnSearch(const ElementType *query_point, const size_t num_closest, + IndexType *out_indices, DistanceType *out_distances_sq, + const int /* nChecks_IGNORED */ = 10) const { + nanoflann::KNNResultSet resultSet(num_closest); + resultSet.init(out_indices, out_distances_sq); + this->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); + return resultSet.size(); + } + + /** + * Find all the neighbors to \a query_point[0:dim-1] within a maximum radius. + * The output is given as a vector of pairs, of which the first element is a + * point index and the second the corresponding distance. Previous contents of + * \a IndicesDists are cleared. + * + * If searchParams.sorted==true, the output list is sorted by ascending + * distances. + * + * For a better performance, it is advisable to do a .reserve() on the vector + * if you have any wild guess about the number of expected matches. + * + * \sa knnSearch, findNeighbors, radiusSearchCustomCallback + * \return The number of points within the given radius (i.e. indices.size() + * or dists.size() ) + */ + size_t + radiusSearch(const ElementType *query_point, const DistanceType &radius, + std::vector> &IndicesDists, + const SearchParams &searchParams) const { + RadiusResultSet resultSet(radius, IndicesDists); + const size_t nFound = + radiusSearchCustomCallback(query_point, resultSet, searchParams); + if (searchParams.sorted) + std::sort(IndicesDists.begin(), IndicesDists.end(), IndexDist_Sorter()); + return nFound; + } + + /** + * Just like radiusSearch() but with a custom callback class for each point + * found in the radius of the query. See the source of RadiusResultSet<> as a + * start point for your own classes. \sa radiusSearch + */ + template + size_t radiusSearchCustomCallback( + const ElementType *query_point, SEARCH_CALLBACK &resultSet, + const SearchParams &searchParams = SearchParams()) const { + this->findNeighbors(resultSet, query_point, searchParams); + return resultSet.size(); + } + + /** @} */ + +public: + /** Make sure the auxiliary list \a vind has the same size than the current + * dataset, and re-generate if size has changed. */ + void init_vind() { + // Create a permutable array of indices to the input vectors. + BaseClassRef::m_size = dataset.kdtree_get_point_count(); + if (BaseClassRef::vind.size() != BaseClassRef::m_size) + BaseClassRef::vind.resize(BaseClassRef::m_size); + for (size_t i = 0; i < BaseClassRef::m_size; i++) + BaseClassRef::vind[i] = i; + } + + void computeBoundingBox(BoundingBox &bbox) { + resize(bbox, (DIM > 0 ? DIM : BaseClassRef::dim)); + if (dataset.kdtree_get_bbox(bbox)) { + // Done! It was implemented in derived class + } else { + const size_t N = dataset.kdtree_get_point_count(); + if (!N) + throw std::runtime_error("[nanoflann] computeBoundingBox() called but " + "no data points found."); + for (int i = 0; i < (DIM > 0 ? DIM : BaseClassRef::dim); ++i) { + bbox[i].low = bbox[i].high = this->dataset_get(*this, 0, i); + } + for (size_t k = 1; k < N; ++k) { + for (int i = 0; i < (DIM > 0 ? DIM : BaseClassRef::dim); ++i) { + if (this->dataset_get(*this, k, i) < bbox[i].low) + bbox[i].low = this->dataset_get(*this, k, i); + if (this->dataset_get(*this, k, i) > bbox[i].high) + bbox[i].high = this->dataset_get(*this, k, i); + } + } + } + } + + /** + * Performs an exact search in the tree starting from a node. + * \tparam RESULTSET Should be any ResultSet + * \return true if the search should be continued, false if the results are + * sufficient + */ + template + bool searchLevel(RESULTSET &result_set, const ElementType *vec, + const NodePtr node, DistanceType mindistsq, + distance_vector_t &dists, const float epsError) const { + /* If this is a leaf node, then do check and return. */ + if ((node->child1 == NULL) && (node->child2 == NULL)) { + // count_leaf += (node->lr.right-node->lr.left); // Removed since was + // neither used nor returned to the user. + DistanceType worst_dist = result_set.worstDist(); + for (IndexType i = node->node_type.lr.left; i < node->node_type.lr.right; + ++i) { + const IndexType index = BaseClassRef::vind[i]; // reorder... : i; + DistanceType dist = distance.evalMetric( + vec, index, (DIM > 0 ? DIM : BaseClassRef::dim)); + if (dist < worst_dist) { + if (!result_set.addPoint(dist, BaseClassRef::vind[i])) { + // the resultset doesn't want to receive any more points, we're done + // searching! + return false; + } + } + } + return true; + } + + /* Which child branch should be taken first? */ + int idx = node->node_type.sub.divfeat; + ElementType val = vec[idx]; + DistanceType diff1 = val - node->node_type.sub.divlow; + DistanceType diff2 = val - node->node_type.sub.divhigh; + + NodePtr bestChild; + NodePtr otherChild; + DistanceType cut_dist; + if ((diff1 + diff2) < 0) { + bestChild = node->child1; + otherChild = node->child2; + cut_dist = distance.accum_dist(val, node->node_type.sub.divhigh, idx); + } else { + bestChild = node->child2; + otherChild = node->child1; + cut_dist = distance.accum_dist(val, node->node_type.sub.divlow, idx); + } + + /* Call recursively to search next level down. */ + if (!searchLevel(result_set, vec, bestChild, mindistsq, dists, epsError)) { + // the resultset doesn't want to receive any more points, we're done + // searching! + return false; + } + + DistanceType dst = dists[idx]; + mindistsq = mindistsq + cut_dist - dst; + dists[idx] = cut_dist; + if (mindistsq * epsError <= result_set.worstDist()) { + if (!searchLevel(result_set, vec, otherChild, mindistsq, dists, + epsError)) { + // the resultset doesn't want to receive any more points, we're done + // searching! + return false; + } + } + dists[idx] = dst; + return true; + } + +public: + /** Stores the index in a binary file. + * IMPORTANT NOTE: The set of data points is NOT stored in the file, so when + * loading the index object it must be constructed associated to the same + * source of data points used while building it. See the example: + * examples/saveload_example.cpp \sa loadIndex */ + void saveIndex(FILE *stream) { this->saveIndex_(*this, stream); } + + /** Loads a previous index from a binary file. + * IMPORTANT NOTE: The set of data points is NOT stored in the file, so the + * index object must be constructed associated to the same source of data + * points used while building the index. See the example: + * examples/saveload_example.cpp \sa loadIndex */ + void loadIndex(FILE *stream) { this->loadIndex_(*this, stream); } + +}; // class KDTree + +/** kd-tree dynamic index + * + * Contains the k-d trees and other information for indexing a set of points + * for nearest-neighbor matching. + * + * The class "DatasetAdaptor" must provide the following interface (can be + * non-virtual, inlined methods): + * + * \code + * // Must return the number of data poins + * inline size_t kdtree_get_point_count() const { ... } + * + * // Must return the dim'th component of the idx'th point in the class: + * inline T kdtree_get_pt(const size_t idx, const size_t dim) const { ... } + * + * // Optional bounding-box computation: return false to default to a standard + * bbox computation loop. + * // Return true if the BBOX was already computed by the class and returned + * in "bb" so it can be avoided to redo it again. + * // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 + * for point clouds) template bool kdtree_get_bbox(BBOX &bb) const + * { + * bb[0].low = ...; bb[0].high = ...; // 0th dimension limits + * bb[1].low = ...; bb[1].high = ...; // 1st dimension limits + * ... + * return true; + * } + * + * \endcode + * + * \tparam DatasetAdaptor The user-provided adaptor (see comments above). + * \tparam Distance The distance metric to use: nanoflann::metric_L1, + * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. \tparam DIM + * Dimensionality of data points (e.g. 3 for 3D points) \tparam IndexType Will + * be typically size_t or int + */ +template +class KDTreeSingleIndexDynamicAdaptor_ + : public KDTreeBaseClass, + Distance, DatasetAdaptor, DIM, IndexType> { +public: + /** + * The dataset used by this index + */ + const DatasetAdaptor &dataset; //!< The source of our data + + KDTreeSingleIndexAdaptorParams index_params; + + std::vector &treeIndex; + + Distance distance; + + typedef typename nanoflann::KDTreeBaseClass< + nanoflann::KDTreeSingleIndexDynamicAdaptor_, + Distance, DatasetAdaptor, DIM, IndexType> + BaseClassRef; + + typedef typename BaseClassRef::ElementType ElementType; + typedef typename BaseClassRef::DistanceType DistanceType; + + typedef typename BaseClassRef::Node Node; + typedef Node *NodePtr; + + typedef typename BaseClassRef::Interval Interval; + /** Define "BoundingBox" as a fixed-size or variable-size container depending + * on "DIM" */ + typedef typename BaseClassRef::BoundingBox BoundingBox; + + /** Define "distance_vector_t" as a fixed-size or variable-size container + * depending on "DIM" */ + typedef typename BaseClassRef::distance_vector_t distance_vector_t; + + /** + * KDTree constructor + * + * Refer to docs in README.md or online in + * https://github.com/jlblancoc/nanoflann + * + * The KD-Tree point dimension (the length of each point in the datase, e.g. 3 + * for 3D points) is determined by means of: + * - The \a DIM template parameter if >0 (highest priority) + * - Otherwise, the \a dimensionality parameter of this constructor. + * + * @param inputData Dataset with the input features + * @param params Basically, the maximum leaf node size + */ + KDTreeSingleIndexDynamicAdaptor_( + const int dimensionality, const DatasetAdaptor &inputData, + std::vector &treeIndex_, + const KDTreeSingleIndexAdaptorParams ¶ms = + KDTreeSingleIndexAdaptorParams()) + : dataset(inputData), index_params(params), treeIndex(treeIndex_), + distance(inputData) { + BaseClassRef::root_node = NULL; + BaseClassRef::m_size = 0; + BaseClassRef::m_size_at_index_build = 0; + BaseClassRef::dim = dimensionality; + if (DIM > 0) + BaseClassRef::dim = DIM; + BaseClassRef::m_leaf_max_size = params.leaf_max_size; + } + + /** Assignment operator definiton */ + KDTreeSingleIndexDynamicAdaptor_ + operator=(const KDTreeSingleIndexDynamicAdaptor_ &rhs) { + KDTreeSingleIndexDynamicAdaptor_ tmp(rhs); + std::swap(BaseClassRef::vind, tmp.BaseClassRef::vind); + std::swap(BaseClassRef::m_leaf_max_size, tmp.BaseClassRef::m_leaf_max_size); + std::swap(index_params, tmp.index_params); + std::swap(treeIndex, tmp.treeIndex); + std::swap(BaseClassRef::m_size, tmp.BaseClassRef::m_size); + std::swap(BaseClassRef::m_size_at_index_build, + tmp.BaseClassRef::m_size_at_index_build); + std::swap(BaseClassRef::root_node, tmp.BaseClassRef::root_node); + std::swap(BaseClassRef::root_bbox, tmp.BaseClassRef::root_bbox); + std::swap(BaseClassRef::pool, tmp.BaseClassRef::pool); + return *this; + } + + /** + * Builds the index + */ + void buildIndex() { + BaseClassRef::m_size = BaseClassRef::vind.size(); + this->freeIndex(*this); + BaseClassRef::m_size_at_index_build = BaseClassRef::m_size; + if (BaseClassRef::m_size == 0) + return; + computeBoundingBox(BaseClassRef::root_bbox); + BaseClassRef::root_node = + this->divideTree(*this, 0, BaseClassRef::m_size, + BaseClassRef::root_bbox); // construct the tree + } + + /** \name Query methods + * @{ */ + + /** + * Find set of nearest neighbors to vec[0:dim-1]. Their indices are stored + * inside the result object. + * + * Params: + * result = the result object in which the indices of the + * nearest-neighbors are stored vec = the vector for which to search the + * nearest neighbors + * + * \tparam RESULTSET Should be any ResultSet + * \return True if the requested neighbors could be found. + * \sa knnSearch, radiusSearch + */ + template + bool findNeighbors(RESULTSET &result, const ElementType *vec, + const SearchParams &searchParams) const { + assert(vec); + if (this->size(*this) == 0) + return false; + if (!BaseClassRef::root_node) + return false; + float epsError = 1 + searchParams.eps; + + // fixed or variable-sized container (depending on DIM) + distance_vector_t dists; + // Fill it with zeros. + assign(dists, (DIM > 0 ? DIM : BaseClassRef::dim), + static_cast(0)); + DistanceType distsq = this->computeInitialDistances(*this, vec, dists); + + searchLevel(result, vec, BaseClassRef::root_node, distsq, dists, + epsError); // "count_leaf" parameter removed since was neither + // used nor returned to the user. + + return result.full(); + } + + /** + * Find the "num_closest" nearest neighbors to the \a query_point[0:dim-1]. + * Their indices are stored inside the result object. \sa radiusSearch, + * findNeighbors \note nChecks_IGNORED is ignored but kept for compatibility + * with the original FLANN interface. \return Number `N` of valid points in + * the result set. Only the first `N` entries in `out_indices` and + * `out_distances_sq` will be valid. Return may be less than `num_closest` + * only if the number of elements in the tree is less than `num_closest`. + */ + size_t knnSearch(const ElementType *query_point, const size_t num_closest, + IndexType *out_indices, DistanceType *out_distances_sq, + const int /* nChecks_IGNORED */ = 10) const { + nanoflann::KNNResultSet resultSet(num_closest); + resultSet.init(out_indices, out_distances_sq); + this->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); + return resultSet.size(); + } + + /** + * Find all the neighbors to \a query_point[0:dim-1] within a maximum radius. + * The output is given as a vector of pairs, of which the first element is a + * point index and the second the corresponding distance. Previous contents of + * \a IndicesDists are cleared. + * + * If searchParams.sorted==true, the output list is sorted by ascending + * distances. + * + * For a better performance, it is advisable to do a .reserve() on the vector + * if you have any wild guess about the number of expected matches. + * + * \sa knnSearch, findNeighbors, radiusSearchCustomCallback + * \return The number of points within the given radius (i.e. indices.size() + * or dists.size() ) + */ + size_t + radiusSearch(const ElementType *query_point, const DistanceType &radius, + std::vector> &IndicesDists, + const SearchParams &searchParams) const { + RadiusResultSet resultSet(radius, IndicesDists); + const size_t nFound = + radiusSearchCustomCallback(query_point, resultSet, searchParams); + if (searchParams.sorted) + std::sort(IndicesDists.begin(), IndicesDists.end(), IndexDist_Sorter()); + return nFound; + } + + /** + * Just like radiusSearch() but with a custom callback class for each point + * found in the radius of the query. See the source of RadiusResultSet<> as a + * start point for your own classes. \sa radiusSearch + */ + template + size_t radiusSearchCustomCallback( + const ElementType *query_point, SEARCH_CALLBACK &resultSet, + const SearchParams &searchParams = SearchParams()) const { + this->findNeighbors(resultSet, query_point, searchParams); + return resultSet.size(); + } + + /** @} */ + +public: + void computeBoundingBox(BoundingBox &bbox) { + resize(bbox, (DIM > 0 ? DIM : BaseClassRef::dim)); + + if (dataset.kdtree_get_bbox(bbox)) { + // Done! It was implemented in derived class + } else { + const size_t N = BaseClassRef::m_size; + if (!N) + throw std::runtime_error("[nanoflann] computeBoundingBox() called but " + "no data points found."); + for (int i = 0; i < (DIM > 0 ? DIM : BaseClassRef::dim); ++i) { + bbox[i].low = bbox[i].high = + this->dataset_get(*this, BaseClassRef::vind[0], i); + } + for (size_t k = 1; k < N; ++k) { + for (int i = 0; i < (DIM > 0 ? DIM : BaseClassRef::dim); ++i) { + if (this->dataset_get(*this, BaseClassRef::vind[k], i) < bbox[i].low) + bbox[i].low = this->dataset_get(*this, BaseClassRef::vind[k], i); + if (this->dataset_get(*this, BaseClassRef::vind[k], i) > bbox[i].high) + bbox[i].high = this->dataset_get(*this, BaseClassRef::vind[k], i); + } + } + } + } + + /** + * Performs an exact search in the tree starting from a node. + * \tparam RESULTSET Should be any ResultSet + */ + template + void searchLevel(RESULTSET &result_set, const ElementType *vec, + const NodePtr node, DistanceType mindistsq, + distance_vector_t &dists, const float epsError) const { + /* If this is a leaf node, then do check and return. */ + if ((node->child1 == NULL) && (node->child2 == NULL)) { + // count_leaf += (node->lr.right-node->lr.left); // Removed since was + // neither used nor returned to the user. + DistanceType worst_dist = result_set.worstDist(); + for (IndexType i = node->node_type.lr.left; i < node->node_type.lr.right; + ++i) { + const IndexType index = BaseClassRef::vind[i]; // reorder... : i; + if (treeIndex[index] == -1) + continue; + DistanceType dist = distance.evalMetric( + vec, index, (DIM > 0 ? DIM : BaseClassRef::dim)); + if (dist < worst_dist) { + if (!result_set.addPoint( + static_cast(dist), + static_cast( + BaseClassRef::vind[i]))) { + // the resultset doesn't want to receive any more points, we're done + // searching! + return; // false; + } + } + } + return; + } + + /* Which child branch should be taken first? */ + int idx = node->node_type.sub.divfeat; + ElementType val = vec[idx]; + DistanceType diff1 = val - node->node_type.sub.divlow; + DistanceType diff2 = val - node->node_type.sub.divhigh; + + NodePtr bestChild; + NodePtr otherChild; + DistanceType cut_dist; + if ((diff1 + diff2) < 0) { + bestChild = node->child1; + otherChild = node->child2; + cut_dist = distance.accum_dist(val, node->node_type.sub.divhigh, idx); + } else { + bestChild = node->child2; + otherChild = node->child1; + cut_dist = distance.accum_dist(val, node->node_type.sub.divlow, idx); + } + + /* Call recursively to search next level down. */ + searchLevel(result_set, vec, bestChild, mindistsq, dists, epsError); + + DistanceType dst = dists[idx]; + mindistsq = mindistsq + cut_dist - dst; + dists[idx] = cut_dist; + if (mindistsq * epsError <= result_set.worstDist()) { + searchLevel(result_set, vec, otherChild, mindistsq, dists, epsError); + } + dists[idx] = dst; + } + +public: + /** Stores the index in a binary file. + * IMPORTANT NOTE: The set of data points is NOT stored in the file, so when + * loading the index object it must be constructed associated to the same + * source of data points used while building it. See the example: + * examples/saveload_example.cpp \sa loadIndex */ + void saveIndex(FILE *stream) { this->saveIndex_(*this, stream); } + + /** Loads a previous index from a binary file. + * IMPORTANT NOTE: The set of data points is NOT stored in the file, so the + * index object must be constructed associated to the same source of data + * points used while building the index. See the example: + * examples/saveload_example.cpp \sa loadIndex */ + void loadIndex(FILE *stream) { this->loadIndex_(*this, stream); } +}; + +/** kd-tree dynaimic index + * + * class to create multiple static index and merge their results to behave as + * single dynamic index as proposed in Logarithmic Approach. + * + * Example of usage: + * examples/dynamic_pointcloud_example.cpp + * + * \tparam DatasetAdaptor The user-provided adaptor (see comments above). + * \tparam Distance The distance metric to use: nanoflann::metric_L1, + * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. \tparam DIM + * Dimensionality of data points (e.g. 3 for 3D points) \tparam IndexType Will + * be typically size_t or int + */ +template +class KDTreeSingleIndexDynamicAdaptor { +public: + typedef typename Distance::ElementType ElementType; + typedef typename Distance::DistanceType DistanceType; + +protected: + size_t m_leaf_max_size; + size_t treeCount; + size_t pointCount; + + /** + * The dataset used by this index + */ + const DatasetAdaptor &dataset; //!< The source of our data + + std::vector treeIndex; //!< treeIndex[idx] is the index of tree in which + //!< point at idx is stored. treeIndex[idx]=-1 + //!< means that point has been removed. + + KDTreeSingleIndexAdaptorParams index_params; + + int dim; //!< Dimensionality of each data point + + typedef KDTreeSingleIndexDynamicAdaptor_ + index_container_t; + std::vector index; + +public: + /** Get a const ref to the internal list of indices; the number of indices is + * adapted dynamically as the dataset grows in size. */ + const std::vector &getAllIndices() const { return index; } + +private: + /** finds position of least significant unset bit */ + int First0Bit(IndexType num) { + int pos = 0; + while (num & 1) { + num = num >> 1; + pos++; + } + return pos; + } + + /** Creates multiple empty trees to handle dynamic support */ + void init() { + typedef KDTreeSingleIndexDynamicAdaptor_ + my_kd_tree_t; + std::vector index_( + treeCount, my_kd_tree_t(dim /*dim*/, dataset, treeIndex, index_params)); + index = index_; + } + +public: + Distance distance; + + /** + * KDTree constructor + * + * Refer to docs in README.md or online in + * https://github.com/jlblancoc/nanoflann + * + * The KD-Tree point dimension (the length of each point in the datase, e.g. 3 + * for 3D points) is determined by means of: + * - The \a DIM template parameter if >0 (highest priority) + * - Otherwise, the \a dimensionality parameter of this constructor. + * + * @param inputData Dataset with the input features + * @param params Basically, the maximum leaf node size + */ + KDTreeSingleIndexDynamicAdaptor(const int dimensionality, + const DatasetAdaptor &inputData, + const KDTreeSingleIndexAdaptorParams ¶ms = + KDTreeSingleIndexAdaptorParams(), + const size_t maximumPointCount = 1000000000U) + : dataset(inputData), index_params(params), distance(inputData) { + treeCount = static_cast(std::log2(maximumPointCount)); + pointCount = 0U; + dim = dimensionality; + treeIndex.clear(); + if (DIM > 0) + dim = DIM; + m_leaf_max_size = params.leaf_max_size; + init(); + const size_t num_initial_points = dataset.kdtree_get_point_count(); + if (num_initial_points > 0) { + addPoints(0, num_initial_points - 1); + } + } + + /** Deleted copy constructor*/ + KDTreeSingleIndexDynamicAdaptor( + const KDTreeSingleIndexDynamicAdaptor &) = delete; + + /** Add points to the set, Inserts all points from [start, end] */ + void addPoints(IndexType start, IndexType end) { + size_t count = end - start + 1; + treeIndex.resize(treeIndex.size() + count); + for (IndexType idx = start; idx <= end; idx++) { + int pos = First0Bit(pointCount); + index[pos].vind.clear(); + treeIndex[pointCount] = pos; + for (int i = 0; i < pos; i++) { + for (int j = 0; j < static_cast(index[i].vind.size()); j++) { + index[pos].vind.push_back(index[i].vind[j]); + if (treeIndex[index[i].vind[j]] != -1) + treeIndex[index[i].vind[j]] = pos; + } + index[i].vind.clear(); + index[i].freeIndex(index[i]); + } + index[pos].vind.push_back(idx); + index[pos].buildIndex(); + pointCount++; + } + } + + /** Remove a point from the set (Lazy Deletion) */ + void removePoint(size_t idx) { + if (idx >= pointCount) + return; + treeIndex[idx] = -1; + } + + /** + * Find set of nearest neighbors to vec[0:dim-1]. Their indices are stored + * inside the result object. + * + * Params: + * result = the result object in which the indices of the + * nearest-neighbors are stored vec = the vector for which to search the + * nearest neighbors + * + * \tparam RESULTSET Should be any ResultSet + * \return True if the requested neighbors could be found. + * \sa knnSearch, radiusSearch + */ + template + bool findNeighbors(RESULTSET &result, const ElementType *vec, + const SearchParams &searchParams) const { + for (size_t i = 0; i < treeCount; i++) { + index[i].findNeighbors(result, &vec[0], searchParams); + } + return result.full(); + } +}; + +/** An L2-metric KD-tree adaptor for working with data directly stored in an + * Eigen Matrix, without duplicating the data storage. Each row in the matrix + * represents a point in the state space. + * + * Example of usage: + * \code + * Eigen::Matrix mat; + * // Fill out "mat"... + * + * typedef KDTreeEigenMatrixAdaptor< Eigen::Matrix > + * my_kd_tree_t; const int max_leaf = 10; my_kd_tree_t mat_index(mat, max_leaf + * ); mat_index.index->buildIndex(); mat_index.index->... \endcode + * + * \tparam DIM If set to >0, it specifies a compile-time fixed dimensionality + * for the points in the data set, allowing more compiler optimizations. \tparam + * Distance The distance metric to use: nanoflann::metric_L1, + * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. + */ +template +struct KDTreeEigenMatrixAdaptor { + typedef KDTreeEigenMatrixAdaptor self_t; + typedef typename MatrixType::Scalar num_t; + typedef typename MatrixType::Index IndexType; + typedef + typename Distance::template traits::distance_t metric_t; + typedef KDTreeSingleIndexAdaptor + index_t; + + index_t *index; //! The kd-tree index for the user to call its methods as + //! usual with any other FLANN index. + + /// Constructor: takes a const ref to the matrix object with the data points + KDTreeEigenMatrixAdaptor(const size_t dimensionality, + const std::reference_wrapper &mat, + const int leaf_max_size = 10) + : m_data_matrix(mat) { + const auto dims = mat.get().cols(); + if (size_t(dims) != dimensionality) + throw std::runtime_error( + "Error: 'dimensionality' must match column count in data matrix"); + if (DIM > 0 && int(dims) != DIM) + throw std::runtime_error( + "Data set dimensionality does not match the 'DIM' template argument"); + index = + new index_t(static_cast(dims), *this /* adaptor */, + nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size)); + index->buildIndex(); + } + +public: + /** Deleted copy constructor */ + KDTreeEigenMatrixAdaptor(const self_t &) = delete; + + ~KDTreeEigenMatrixAdaptor() { delete index; } + + const std::reference_wrapper m_data_matrix; + + /** Query for the \a num_closest closest points to a given point (entered as + * query_point[0:dim-1]). Note that this is a short-cut method for + * index->findNeighbors(). The user can also call index->... methods as + * desired. \note nChecks_IGNORED is ignored but kept for compatibility with + * the original FLANN interface. + */ + inline void query(const num_t *query_point, const size_t num_closest, + IndexType *out_indices, num_t *out_distances_sq, + const int /* nChecks_IGNORED */ = 10) const { + nanoflann::KNNResultSet resultSet(num_closest); + resultSet.init(out_indices, out_distances_sq); + index->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); + } + + /** @name Interface expected by KDTreeSingleIndexAdaptor + * @{ */ + + const self_t &derived() const { return *this; } + self_t &derived() { return *this; } + + // Must return the number of data points + inline size_t kdtree_get_point_count() const { + return m_data_matrix.get().rows(); + } + + // Returns the dim'th component of the idx'th point in the class: + inline num_t kdtree_get_pt(const IndexType idx, size_t dim) const { + return m_data_matrix.get().coeff(idx, IndexType(dim)); + } + + // Optional bounding-box computation: return false to default to a standard + // bbox computation loop. + // Return true if the BBOX was already computed by the class and returned in + // "bb" so it can be avoided to redo it again. Look at bb.size() to find out + // the expected dimensionality (e.g. 2 or 3 for point clouds) + template bool kdtree_get_bbox(BBOX & /*bb*/) const { + return false; + } + + /** @} */ + +}; // end of KDTreeEigenMatrixAdaptor + /** @} */ + +/** @} */ // end of grouping +} // namespace nanoflann + +#endif /* NANOFLANN_HPP_ */ diff --git a/datasets/ModelNet40.py b/datasets/ModelNet40.py new file mode 100644 index 0000000..90c1868 --- /dev/null +++ b/datasets/ModelNet40.py @@ -0,0 +1,996 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Class handling ModelNet40 dataset. +# Implements a Dataset, a Sampler, and a collate_fn +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 11/06/2018 +# + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Imports and global variables +# \**********************************/ +# + +# Common libs +import time +import numpy as np +import pickle +import torch +import math + + +# OS functions +from os import listdir +from os.path import exists, join + +# Dataset parent class +from datasets.common import PointCloudDataset +from torch.utils.data import Sampler, get_worker_info +from utils.mayavi_visu import * + +from datasets.common import grid_subsampling +from utils.config import bcolors + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Dataset class definition +# \******************************/ + + +class ModelNet40Dataset(PointCloudDataset): + """Class to handle Modelnet 40 dataset.""" + + def __init__(self, config, train=True, orient_correction=True): + """ + This dataset is small enough to be stored in-memory, so load all point clouds here + """ + PointCloudDataset.__init__(self, 'ModelNet40') + + ############ + # Parameters + ############ + + # Dict from labels to names + self.label_to_names = {0: 'airplane', + 1: 'bathtub', + 2: 'bed', + 3: 'bench', + 4: 'bookshelf', + 5: 'bottle', + 6: 'bowl', + 7: 'car', + 8: 'chair', + 9: 'cone', + 10: 'cup', + 11: 'curtain', + 12: 'desk', + 13: 'door', + 14: 'dresser', + 15: 'flower_pot', + 16: 'glass_box', + 17: 'guitar', + 18: 'keyboard', + 19: 'lamp', + 20: 'laptop', + 21: 'mantel', + 22: 'monitor', + 23: 'night_stand', + 24: 'person', + 25: 'piano', + 26: 'plant', + 27: 'radio', + 28: 'range_hood', + 29: 'sink', + 30: 'sofa', + 31: 'stairs', + 32: 'stool', + 33: 'table', + 34: 'tent', + 35: 'toilet', + 36: 'tv_stand', + 37: 'vase', + 38: 'wardrobe', + 39: 'xbox'} + + # Initialize a bunch of variables concerning class labels + self.init_labels() + + # List of classes ignored during training (can be empty) + self.ignored_labels = np.array([]) + + # Dataset folder + self.path = '../../Data/ModelNet40' + + # Type of task conducted on this dataset + self.dataset_task = 'classification' + + # Update number of class and data task in configuration + config.num_classes = self.num_classes + config.dataset_task = self.dataset_task + + # Parameters from config + self.config = config + + # Training or test set + self.train = train + + # Number of models and models used per epoch + if self.train: + self.num_models = 9843 + if config.epoch_steps and config.epoch_steps * config.batch_num < self.num_models: + self.epoch_n = config.epoch_steps * config.batch_num + else: + self.epoch_n = self.num_models + else: + self.num_models = 2468 + self.epoch_n = min(self.num_models, config.validation_size * config.batch_num) + + ############# + # Load models + ############# + + if 0 < self.config.first_subsampling_dl <= 0.01: + raise ValueError('subsampling_parameter too low (should be over 1 cm') + + self.input_points, self.input_normals, self.input_labels = self.load_subsampled_clouds(orient_correction) + + return + + def __len__(self): + """ + Return the length of data here + """ + return self.num_models + + def __getitem__(self, idx_list): + """ + The main thread gives a list of indices to load a batch. Each worker is going to work in parallel to load a + different list of indices. + """ + + ################### + # Gather batch data + ################### + + tp_list = [] + tn_list = [] + tl_list = [] + ti_list = [] + s_list = [] + R_list = [] + + for p_i in idx_list: + + # Get points and labels + points = self.input_points[p_i].astype(np.float32) + normals = self.input_normals[p_i].astype(np.float32) + label = self.label_to_idx[self.input_labels[p_i]] + + # Data augmentation + points, normals, scale, R = self.augmentation_transform(points, normals) + + # Stack batch + tp_list += [points] + tn_list += [normals] + tl_list += [label] + ti_list += [p_i] + s_list += [scale] + R_list += [R] + + ################### + # Concatenate batch + ################### + + #show_ModelNet_examples(tp_list, cloud_normals=tn_list) + + stacked_points = np.concatenate(tp_list, axis=0) + stacked_normals = np.concatenate(tn_list, axis=0) + labels = np.array(tl_list, dtype=np.int64) + model_inds = np.array(ti_list, dtype=np.int32) + stack_lengths = np.array([tp.shape[0] for tp in tp_list], dtype=np.int32) + scales = np.array(s_list, dtype=np.float32) + rots = np.stack(R_list, axis=0) + + # Input features + stacked_features = np.ones_like(stacked_points[:, :1], dtype=np.float32) + if self.config.in_features_dim == 1: + pass + elif self.config.in_features_dim == 4: + stacked_features = np.hstack((stacked_features, stacked_normals)) + else: + raise ValueError('Only accepted input dimensions are 1, 4 and 7 (without and with XYZ)') + + ####################### + # Create network inputs + ####################### + # + # Points, neighbors, pooling indices for each layers + # + + # Get the whole input list + input_list = self.classification_inputs(stacked_points, + stacked_features, + labels, + stack_lengths) + + # Add scale and rotation for testing + input_list += [scales, rots, model_inds] + + return input_list + + def load_subsampled_clouds(self, orient_correction): + + # Restart timer + t0 = time.time() + + # Load wanted points if possible + if self.train: + split ='training' + else: + split = 'test' + + print('\nLoading {:s} points subsampled at {:.3f}'.format(split, self.config.first_subsampling_dl)) + filename = join(self.path, '{:s}_{:.3f}_record.pkl'.format(split, self.config.first_subsampling_dl)) + + if exists(filename): + with open(filename, 'rb') as file: + input_points, input_normals, input_labels = pickle.load(file) + + # Else compute them from original points + else: + + # Collect training file names + if self.train: + names = np.loadtxt(join(self.path, 'modelnet40_train.txt'), dtype=np.str) + else: + names = np.loadtxt(join(self.path, 'modelnet40_test.txt'), dtype=np.str) + + # Initialize containers + input_points = [] + input_normals = [] + + # Advanced display + N = len(names) + progress_n = 30 + fmt_str = '[{:<' + str(progress_n) + '}] {:5.1f}%' + + # Collect point clouds + for i, cloud_name in enumerate(names): + + # Read points + class_folder = '_'.join(cloud_name.split('_')[:-1]) + txt_file = join(self.path, class_folder, cloud_name) + '.txt' + data = np.loadtxt(txt_file, delimiter=',', dtype=np.float32) + + # Subsample them + if self.config.first_subsampling_dl > 0: + points, normals = grid_subsampling(data[:, :3], + features=data[:, 3:], + sampleDl=self.config.first_subsampling_dl) + else: + points = data[:, :3] + normals = data[:, 3:] + + print('', end='\r') + print(fmt_str.format('#' * ((i * progress_n) // N), 100 * i / N), end='', flush=True) + + # Add to list + input_points += [points] + input_normals += [normals] + + print('', end='\r') + print(fmt_str.format('#' * progress_n, 100), end='', flush=True) + print() + + # Get labels + label_names = ['_'.join(name.split('_')[:-1]) for name in names] + input_labels = np.array([self.name_to_label[name] for name in label_names]) + + # Save for later use + with open(filename, 'wb') as file: + pickle.dump((input_points, + input_normals, + input_labels), file) + + lengths = [p.shape[0] for p in input_points] + sizes = [l * 4 * 6 for l in lengths] + print('{:.1f} MB loaded in {:.1f}s'.format(np.sum(sizes) * 1e-6, time.time() - t0)) + + if orient_correction: + input_points = [pp[:, [0, 2, 1]] for pp in input_points] + input_normals = [nn[:, [0, 2, 1]] for nn in input_normals] + + return input_points, input_normals, input_labels + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Utility classes definition +# \********************************/ + + +class ModelNet40Sampler(Sampler): + """Sampler for ModelNet40""" + + def __init__(self, dataset: ModelNet40Dataset, use_potential=True, balance_labels=False): + Sampler.__init__(self, dataset) + + # Does the sampler use potential for regular sampling + self.use_potential = use_potential + + # Should be balance the classes when sampling + self.balance_labels = balance_labels + + # Dataset used by the sampler (no copy is made in memory) + self.dataset = dataset + + # Create potentials + if self.use_potential: + self.potentials = np.random.rand(len(dataset.input_labels)) * 0.1 + 0.1 + else: + self.potentials = None + + # Initialize value for batch limit (max number of points per batch). + self.batch_limit = 10000 + + return + + def __iter__(self): + """ + Yield next batch indices here + """ + + ########################################## + # Initialize the list of generated indices + ########################################## + + if self.use_potential: + if self.balance_labels: + + gen_indices = [] + pick_n = self.dataset.epoch_n // self.dataset.num_classes + 1 + for i, l in enumerate(self.dataset.label_values): + + # Get the potentials of the objects of this class + label_inds = np.where(np.equal(self.dataset.input_labels, l))[0] + class_potentials = self.potentials[label_inds] + + # Get the indices to generate thanks to potentials + if pick_n < class_potentials.shape[0]: + pick_indices = np.argpartition(class_potentials, pick_n)[:pick_n] + else: + pick_indices = np.random.permutation(class_potentials.shape[0]) + class_indices = label_inds[pick_indices] + gen_indices.append(class_indices) + + # Stack the chosen indices of all classes + gen_indices = np.random.permutation(np.hstack(gen_indices)) + + else: + + # Get indices with the minimum potential + if self.dataset.epoch_n < self.potentials.shape[0]: + gen_indices = np.argpartition(self.potentials, self.dataset.epoch_n)[:self.dataset.epoch_n] + else: + gen_indices = np.random.permutation(self.potentials.shape[0]) + gen_indices = np.random.permutation(gen_indices) + + # Update potentials (Change the order for the next epoch) + self.potentials[gen_indices] = np.ceil(self.potentials[gen_indices]) + self.potentials[gen_indices] += np.random.rand(gen_indices.shape[0]) * 0.1 + 0.1 + + else: + if self.balance_labels: + pick_n = self.dataset.epoch_n // self.dataset.num_classes + 1 + gen_indices = [] + for l in self.dataset.label_values: + label_inds = np.where(np.equal(self.dataset.input_labels, l))[0] + rand_inds = np.random.choice(label_inds, size=pick_n, replace=True) + gen_indices += [rand_inds] + gen_indices = np.random.permutation(np.hstack(gen_indices)) + else: + gen_indices = np.random.permutation(self.dataset.num_models)[:self.dataset.epoch_n] + + ################ + # Generator loop + ################ + + # Initialize concatenation lists + ti_list = [] + batch_n = 0 + + # Generator loop + for p_i in gen_indices: + + # Size of picked cloud + n = self.dataset.input_points[p_i].shape[0] + + # In case batch is full, yield it and reset it + if batch_n + n > self.batch_limit and batch_n > 0: + yield np.array(ti_list, dtype=np.int32) + ti_list = [] + batch_n = 0 + + # Add data to current batch + ti_list += [p_i] + + # Update batch size + batch_n += n + + yield np.array(ti_list, dtype=np.int32) + + return 0 + + def __len__(self): + """ + The number of yielded samples is variable + """ + return None + + def calibration(self, dataloader, untouched_ratio=0.9, verbose=False): + """ + Method performing batch and neighbors calibration. + Batch calibration: Set "batch_limit" (the maximum number of points allowed in every batch) so that the + average batch size (number of stacked pointclouds) is the one asked. + Neighbors calibration: Set the "neighborhood_limits" (the maximum number of neighbors allowed in convolutions) + so that 90% of the neighborhoods remain untouched. There is a limit for each layer. + """ + + ############################## + # Previously saved calibration + ############################## + + print('\nStarting Calibration (use verbose=True for more details)') + t0 = time.time() + + redo = False + + # Batch limit + # *********** + + # Load batch_limit dictionary + batch_lim_file = join(self.dataset.path, 'batch_limits.pkl') + if exists(batch_lim_file): + with open(batch_lim_file, 'rb') as file: + batch_lim_dict = pickle.load(file) + else: + batch_lim_dict = {} + + # Check if the batch limit associated with current parameters exists + key = '{:.3f}_{:d}'.format(self.dataset.config.first_subsampling_dl, + self.dataset.config.batch_num) + if key in batch_lim_dict: + self.batch_limit = batch_lim_dict[key] + else: + redo = True + + if verbose: + print('\nPrevious calibration found:') + print('Check batch limit dictionary') + if key in batch_lim_dict: + color = bcolors.OKGREEN + v = str(int(batch_lim_dict[key])) + else: + color = bcolors.FAIL + v = '?' + print('{:}\"{:s}\": {:s}{:}'.format(color, key, v, bcolors.ENDC)) + + # Neighbors limit + # *************** + + # Load neighb_limits dictionary + neighb_lim_file = join(self.dataset.path, 'neighbors_limits.pkl') + if exists(neighb_lim_file): + with open(neighb_lim_file, 'rb') as file: + neighb_lim_dict = pickle.load(file) + else: + neighb_lim_dict = {} + + # Check if the limit associated with current parameters exists (for each layer) + neighb_limits = [] + for layer_ind in range(self.dataset.config.num_layers): + + dl = self.dataset.config.first_subsampling_dl * (2**layer_ind) + if self.dataset.config.deform_layers[layer_ind]: + r = dl * self.dataset.config.deform_radius + else: + r = dl * self.dataset.config.conv_radius + + key = '{:.3f}_{:.3f}'.format(dl, r) + if key in neighb_lim_dict: + neighb_limits += [neighb_lim_dict[key]] + + if len(neighb_limits) == self.dataset.config.num_layers: + self.dataset.neighborhood_limits = neighb_limits + else: + redo = True + + if verbose: + print('Check neighbors limit dictionary') + for layer_ind in range(self.dataset.config.num_layers): + dl = self.dataset.config.first_subsampling_dl * (2**layer_ind) + if self.dataset.config.deform_layers[layer_ind]: + r = dl * self.dataset.config.deform_radius + else: + r = dl * self.dataset.config.conv_radius + key = '{:.3f}_{:.3f}'.format(dl, r) + + if key in neighb_lim_dict: + color = bcolors.OKGREEN + v = str(neighb_lim_dict[key]) + else: + color = bcolors.FAIL + v = '?' + print('{:}\"{:s}\": {:s}{:}'.format(color, key, v, bcolors.ENDC)) + + if redo: + + ############################ + # Neighbors calib parameters + ############################ + + # From config parameter, compute higher bound of neighbors number in a neighborhood + hist_n = int(np.ceil(4 / 3 * np.pi * (self.dataset.config.conv_radius + 1) ** 3)) + + # Histogram of neighborhood sizes + neighb_hists = np.zeros((self.dataset.config.num_layers, hist_n), dtype=np.int32) + + ######################## + # Batch calib parameters + ######################## + + # Estimated average batch size and target value + estim_b = 0 + target_b = self.dataset.config.batch_num + + # Calibration parameters + low_pass_T = 10 + Kp = 100.0 + finer = False + + # Convergence parameters + smooth_errors = [] + converge_threshold = 0.1 + + # Loop parameters + last_display = time.time() + i = 0 + breaking = False + + ##################### + # Perform calibration + ##################### + + for epoch in range(10): + for batch_i, batch in enumerate(dataloader): + + # Update neighborhood histogram + counts = [np.sum(neighb_mat.numpy() < neighb_mat.shape[0], axis=1) for neighb_mat in batch.neighbors] + hists = [np.bincount(c, minlength=hist_n)[:hist_n] for c in counts] + neighb_hists += np.vstack(hists) + + # batch length + b = len(batch.labels) + + # Update estim_b (low pass filter) + estim_b += (b - estim_b) / low_pass_T + + # Estimate error (noisy) + error = target_b - b + + # Save smooth errors for convergene check + smooth_errors.append(target_b - estim_b) + if len(smooth_errors) > 10: + smooth_errors = smooth_errors[1:] + + # Update batch limit with P controller + self.batch_limit += Kp * error + + # finer low pass filter when closing in + if not finer and np.abs(estim_b - target_b) < 1: + low_pass_T = 100 + finer = True + + # Convergence + if finer and np.max(np.abs(smooth_errors)) < converge_threshold: + breaking = True + break + + i += 1 + t = time.time() + + # Console display (only one per second) + if verbose and (t - last_display) > 1.0: + last_display = t + message = 'Step {:5d} estim_b ={:5.2f} batch_limit ={:7d}' + print(message.format(i, + estim_b, + int(self.batch_limit))) + + if breaking: + break + + # Use collected neighbor histogram to get neighbors limit + cumsum = np.cumsum(neighb_hists.T, axis=0) + percentiles = np.sum(cumsum < (untouched_ratio * cumsum[hist_n - 1, :]), axis=0) + self.dataset.neighborhood_limits = percentiles + + if verbose: + + # Crop histogram + while np.sum(neighb_hists[:, -1]) == 0: + neighb_hists = neighb_hists[:, :-1] + hist_n = neighb_hists.shape[1] + + print('\n**************************************************\n') + line0 = 'neighbors_num ' + for layer in range(neighb_hists.shape[0]): + line0 += '| layer {:2d} '.format(layer) + print(line0) + for neighb_size in range(hist_n): + line0 = ' {:4d} '.format(neighb_size) + for layer in range(neighb_hists.shape[0]): + if neighb_size > percentiles[layer]: + color = bcolors.FAIL + else: + color = bcolors.OKGREEN + line0 += '|{:}{:10d}{:} '.format(color, + neighb_hists[layer, neighb_size], + bcolors.ENDC) + + print(line0) + + print('\n**************************************************\n') + print('\nchosen neighbors limits: ', percentiles) + print() + + # Save batch_limit dictionary + key = '{:.3f}_{:d}'.format(self.dataset.config.first_subsampling_dl, + self.dataset.config.batch_num) + batch_lim_dict[key] = self.batch_limit + with open(batch_lim_file, 'wb') as file: + pickle.dump(batch_lim_dict, file) + + # Save neighb_limit dictionary + for layer_ind in range(self.dataset.config.num_layers): + dl = self.dataset.config.first_subsampling_dl * (2 ** layer_ind) + if self.dataset.config.deform_layers[layer_ind]: + r = dl * self.dataset.config.deform_radius + else: + r = dl * self.dataset.config.conv_radius + key = '{:.3f}_{:.3f}'.format(dl, r) + neighb_lim_dict[key] = self.dataset.neighborhood_limits[layer_ind] + with open(neighb_lim_file, 'wb') as file: + pickle.dump(neighb_lim_dict, file) + + + print('Calibration done in {:.1f}s\n'.format(time.time() - t0)) + return + + +class ModelNet40CustomBatch: + """Custom batch definition with memory pinning for ModelNet40""" + + def __init__(self, input_list): + + # Get rid of batch dimension + input_list = input_list[0] + + # Number of layers + L = (len(input_list) - 5) // 4 + + # Extract input tensors from the list of numpy array + ind = 0 + self.points = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]] + ind += L + self.neighbors = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]] + ind += L + self.pools = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]] + ind += L + self.lengths = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]] + ind += L + self.features = torch.from_numpy(input_list[ind]) + ind += 1 + self.labels = torch.from_numpy(input_list[ind]) + ind += 1 + self.scales = torch.from_numpy(input_list[ind]) + ind += 1 + self.rots = torch.from_numpy(input_list[ind]) + ind += 1 + self.model_inds = torch.from_numpy(input_list[ind]) + + return + + def pin_memory(self): + """ + Manual pinning of the memory + """ + + self.points = [in_tensor.pin_memory() for in_tensor in self.points] + self.neighbors = [in_tensor.pin_memory() for in_tensor in self.neighbors] + self.pools = [in_tensor.pin_memory() for in_tensor in self.pools] + self.lengths = [in_tensor.pin_memory() for in_tensor in self.lengths] + self.features = self.features.pin_memory() + self.labels = self.labels.pin_memory() + self.scales = self.scales.pin_memory() + self.rots = self.rots.pin_memory() + self.model_inds = self.model_inds.pin_memory() + + return self + + def to(self, device): + + self.points = [in_tensor.to(device) for in_tensor in self.points] + self.neighbors = [in_tensor.to(device) for in_tensor in self.neighbors] + self.pools = [in_tensor.to(device) for in_tensor in self.pools] + self.lengths = [in_tensor.to(device) for in_tensor in self.lengths] + self.features = self.features.to(device) + self.labels = self.labels.to(device) + self.scales = self.scales.to(device) + self.rots = self.rots.to(device) + self.model_inds = self.model_inds.to(device) + + return self + + def unstack_points(self, layer=None): + """Unstack the points""" + return self.unstack_elements('points', layer) + + def unstack_neighbors(self, layer=None): + """Unstack the neighbors indices""" + return self.unstack_elements('neighbors', layer) + + def unstack_pools(self, layer=None): + """Unstack the pooling indices""" + return self.unstack_elements('pools', layer) + + def unstack_elements(self, element_name, layer=None, to_numpy=True): + """ + Return a list of the stacked elements in the batch at a certain layer. If no layer is given, then return all + layers + """ + + if element_name == 'points': + elements = self.points + elif element_name == 'neighbors': + elements = self.neighbors + elif element_name == 'pools': + elements = self.pools[:-1] + else: + raise ValueError('Unknown element name: {:s}'.format(element_name)) + + all_p_list = [] + for layer_i, layer_elems in enumerate(elements): + + if layer is None or layer == layer_i: + + i0 = 0 + p_list = [] + if element_name == 'pools': + lengths = self.lengths[layer_i+1] + else: + lengths = self.lengths[layer_i] + + for b_i, length in enumerate(lengths): + + elem = layer_elems[i0:i0 + length] + if element_name == 'neighbors': + elem[elem >= self.points[layer_i].shape[0]] = -1 + elem[elem >= 0] -= i0 + elif element_name == 'pools': + elem[elem >= self.points[layer_i].shape[0]] = -1 + elem[elem >= 0] -= torch.sum(self.lengths[layer_i][:b_i]) + i0 += length + + if to_numpy: + p_list.append(elem.numpy()) + else: + p_list.append(elem) + + if layer == layer_i: + return p_list + + all_p_list.append(p_list) + + return all_p_list + + + + +def ModelNet40Collate(batch_data): + return ModelNet40CustomBatch(batch_data) + + +class ModelNet40WorkerInitDebug(): + """Callable class that Initializes workers.""" + + def __init__(self, dataset): + self.dataset = dataset + return + + def __call__(self, worker_id): + + # Print workers info + worker_info = get_worker_info() + print(worker_info) + + # Get associated dataset + dataset = worker_info.dataset # the dataset copy in this worker process + + # In windows, each worker has its own copy of the dataset. In Linux, this is shared in memory + print(dataset.input_labels.__array_interface__['data']) + print(worker_info.dataset.input_labels.__array_interface__['data']) + print(self.dataset.input_labels.__array_interface__['data']) + + # configure the dataset to only process the split workload + + return + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Debug functions +# \*********************/ + + +def debug_sampling(dataset, sampler, loader): + """Shows which labels are sampled according to strategy chosen""" + label_sum = np.zeros((dataset.num_classes), dtype=np.int32) + for epoch in range(10): + + for batch_i, (points, normals, labels, indices, in_sizes) in enumerate(loader): + # print(batch_i, tuple(points.shape), tuple(normals.shape), labels, indices, in_sizes) + + label_sum += np.bincount(labels.numpy(), minlength=dataset.num_classes) + print(label_sum) + #print(sampler.potentials[:6]) + + print('******************') + print('*******************************************') + + _, counts = np.unique(dataset.input_labels, return_counts=True) + print(counts) + + +def debug_timing(dataset, sampler, loader): + """Timing of generator function""" + + t = [time.time()] + last_display = time.time() + mean_dt = np.zeros(2) + estim_b = dataset.config.batch_num + + for epoch in range(10): + + for batch_i, batch in enumerate(loader): + # print(batch_i, tuple(points.shape), tuple(normals.shape), labels, indices, in_sizes) + + # New time + t = t[-1:] + t += [time.time()] + + # Update estim_b (low pass filter) + estim_b += (len(batch.labels) - estim_b) / 100 + + # Pause simulating computations + time.sleep(0.050) + t += [time.time()] + + # Average timing + mean_dt = 0.9 * mean_dt + 0.1 * (np.array(t[1:]) - np.array(t[:-1])) + + # Console display (only one per second) + if (t[-1] - last_display) > -1.0: + last_display = t[-1] + message = 'Step {:08d} -> (ms/batch) {:8.2f} {:8.2f} / batch = {:.2f}' + print(message.format(batch_i, + 1000 * mean_dt[0], + 1000 * mean_dt[1], + estim_b)) + + print('************* Epoch ended *************') + + _, counts = np.unique(dataset.input_labels, return_counts=True) + print(counts) + + +def debug_show_clouds(dataset, sampler, loader): + + + for epoch in range(10): + + clouds = [] + cloud_normals = [] + cloud_labels = [] + + L = dataset.config.num_layers + + for batch_i, batch in enumerate(loader): + + # Print characteristics of input tensors + print('\nPoints tensors') + for i in range(L): + print(batch.points[i].dtype, batch.points[i].shape) + print('\nNeigbors tensors') + for i in range(L): + print(batch.neighbors[i].dtype, batch.neighbors[i].shape) + print('\nPools tensors') + for i in range(L): + print(batch.pools[i].dtype, batch.pools[i].shape) + print('\nStack lengths') + for i in range(L): + print(batch.lengths[i].dtype, batch.lengths[i].shape) + print('\nFeatures') + print(batch.features.dtype, batch.features.shape) + print('\nLabels') + print(batch.labels.dtype, batch.labels.shape) + print('\nAugment Scales') + print(batch.scales.dtype, batch.scales.shape) + print('\nAugment Rotations') + print(batch.rots.dtype, batch.rots.shape) + print('\nModel indices') + print(batch.model_inds.dtype, batch.model_inds.shape) + + print('\nAre input tensors pinned') + print(batch.neighbors[0].is_pinned()) + print(batch.neighbors[-1].is_pinned()) + print(batch.points[0].is_pinned()) + print(batch.points[-1].is_pinned()) + print(batch.labels.is_pinned()) + print(batch.scales.is_pinned()) + print(batch.rots.is_pinned()) + print(batch.model_inds.is_pinned()) + + show_input_batch(batch) + + print('*******************************************') + + _, counts = np.unique(dataset.input_labels, return_counts=True) + print(counts) + + +def debug_batch_and_neighbors_calib(dataset, sampler, loader): + """Timing of generator function""" + + t = [time.time()] + last_display = time.time() + mean_dt = np.zeros(2) + + for epoch in range(10): + + for batch_i, input_list in enumerate(loader): + # print(batch_i, tuple(points.shape), tuple(normals.shape), labels, indices, in_sizes) + + # New time + t = t[-1:] + t += [time.time()] + + # Pause simulating computations + time.sleep(0.01) + t += [time.time()] + + # Average timing + mean_dt = 0.9 * mean_dt + 0.1 * (np.array(t[1:]) - np.array(t[:-1])) + + # Console display (only one per second) + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + message = 'Step {:08d} -> Average timings (ms/batch) {:8.2f} {:8.2f} ' + print(message.format(batch_i, + 1000 * mean_dt[0], + 1000 * mean_dt[1])) + + print('************* Epoch ended *************') + + _, counts = np.unique(dataset.input_labels, return_counts=True) + print(counts) diff --git a/datasets/S3DIS.py b/datasets/S3DIS.py new file mode 100644 index 0000000..89fd350 --- /dev/null +++ b/datasets/S3DIS.py @@ -0,0 +1,1344 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Class handling S3DIS dataset. +# Implements a Dataset, a Sampler, and a collate_fn +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 11/06/2018 +# + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Imports and global variables +# \**********************************/ +# + +# Common libs +import time +import numpy as np +import pickle +import torch +import math +#from mayavi import mlab +from multiprocessing import Lock + + +# OS functions +from os import listdir +from os.path import exists, join, isdir + +# Dataset parent class +from datasets.common import PointCloudDataset +from torch.utils.data import Sampler, get_worker_info +from utils.mayavi_visu import * + +from datasets.common import grid_subsampling +from utils.config import bcolors + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Dataset class definition +# \******************************/ + + +class S3DISDataset(PointCloudDataset): + """Class to handle Modelnet 40 dataset.""" + + def __init__(self, config, set='training', use_potentials=True, load_data=True): + """ + This dataset is small enough to be stored in-memory, so load all point clouds here + """ + PointCloudDataset.__init__(self, 'S3DIS') + + ############ + # Parameters + ############ + + # Dict from labels to names + self.label_to_names = {0: 'ceiling', + 1: 'floor', + 2: 'wall', + 3: 'beam', + 4: 'column', + 5: 'window', + 6: 'door', + 7: 'chair', + 8: 'table', + 9: 'bookcase', + 10: 'sofa', + 11: 'board', + 12: 'clutter'} + + # Initialize a bunch of variables concerning class labels + self.init_labels() + + # List of classes ignored during training (can be empty) + self.ignored_labels = np.array([]) + + # Dataset folder + self.path = '../../Data/S3DIS' + + # Type of task conducted on this dataset + self.dataset_task = 'cloud_segmentation' + + # Update number of class and data task in configuration + config.num_classes = self.num_classes + config.dataset_task = self.dataset_task + + # Parameters from config + self.config = config + + # Training or test set + self.set = set + + # Using potential or random epoch generation + self.use_potentials = use_potentials + + # Path of the training files + self.train_path = 'original_ply' + + # List of files to process + ply_path = join(self.path, self.train_path) + + # Proportion of validation scenes + self.cloud_names = ['Area_1', 'Area_2', 'Area_3', 'Area_4', 'Area_5', 'Area_6'] + self.all_splits = [0, 1, 2, 3, 4, 5] + self.validation_split = 4 + + # Number of models used per epoch + if self.set == 'training': + self.epoch_n = config.epoch_steps * config.batch_num + elif self.set in ['validation', 'test', 'ERF']: + self.epoch_n = config.validation_size * config.batch_num + else: + raise ValueError('Unknown set for S3DIS data: ', self.set) + + # Stop data is not needed + if not load_data: + return + + ################### + # Prepare ply files + ################### + + self.prepare_S3DIS_ply() + + ################ + # Load ply files + ################ + + # List of training files + self.train_files = [join(ply_path, f + '.ply') for f in self.cloud_names] + + if 0 < self.config.first_subsampling_dl <= 0.01: + raise ValueError('subsampling_parameter too low (should be over 1 cm') + + # Initiate containers + self.input_trees = [] + self.input_colors = [] + self.input_labels = [] + self.pot_trees = [] + self.num_clouds = 0 + self.validation_proj = [] + self.validation_labels = [] + + # Start loading + self.load_subsampled_clouds() + + ############################ + # Batch selection parameters + ############################ + + # Initialize value for batch limit (max number of points per batch). + self.batch_limit = torch.tensor([1], dtype=torch.float32) + self.batch_limit.share_memory_() + + # Initialize potentials + if use_potentials: + self.potentials = [] + self.min_potentials = [] + self.argmin_potentials = [] + for i, tree in enumerate(self.pot_trees): + self.potentials += [torch.from_numpy(np.random.rand(tree.data.shape[0]) * 1e-3)] + min_ind = int(torch.argmin(self.potentials[-1])) + self.argmin_potentials += [min_ind] + self.min_potentials += [float(self.potentials[-1][min_ind])] + + # Share potential memory + self.pot_lock = Lock() + self.argmin_potentials = torch.from_numpy(np.array(self.argmin_potentials, dtype=np.int64)) + self.min_potentials = torch.from_numpy(np.array(self.min_potentials, dtype=np.float64)) + self.argmin_potentials.share_memory_() + self.min_potentials.share_memory_() + for i, _ in enumerate(self.pot_trees): + self.potentials[i].share_memory_() + + self.worker_waiting = torch.tensor([0 for _ in range(config.input_threads)], dtype=torch.int32) + self.worker_waiting.share_memory_() + + else: + self.pot_lock = None + self.potentials = None + self.min_potentials = None + self.argmin_potentials = None + + # For ERF visualization, we want only one cloud per batch and no randomness + if self.set == 'ERF': + self.batch_limit = torch.tensor([1], dtype=torch.float32) + self.batch_limit.share_memory_() + np.random.seed(42) + + return + + def __len__(self): + """ + Return the length of data here + """ + return len(self.cloud_names) + + def __getitem__(self, batch_i, debug_workers=False): + """ + The main thread gives a list of indices to load a batch. Each worker is going to work in parallel to load a + different list of indices. + """ + + # Initiate concatanation lists + p_list = [] + f_list = [] + l_list = [] + i_list = [] + pi_list = [] + ci_list = [] + s_list = [] + R_list = [] + batch_n = 0 + + info = get_worker_info() + wid = info.id + + + while True: + + if debug_workers: + message = '' + for wi in range(info.num_workers): + if wi == wid: + message += ' {:}X{:} '.format(bcolors.FAIL, bcolors.ENDC) + elif self.worker_waiting[wi] == 0: + message += ' ' + elif self.worker_waiting[wi] == 1: + message += ' | ' + elif self.worker_waiting[wi] == 2: + message += ' o ' + print(message) + self.worker_waiting[wid] = 0 + + with self.pot_lock: + + if debug_workers: + message = '' + for wi in range(info.num_workers): + if wi == wid: + message += ' {:}v{:} '.format(bcolors.OKGREEN, bcolors.ENDC) + elif self.worker_waiting[wi] == 0: + message += ' ' + elif self.worker_waiting[wi] == 1: + message += ' | ' + elif self.worker_waiting[wi] == 2: + message += ' o ' + print(message) + self.worker_waiting[wid] = 1 + + # Get potential minimum + cloud_ind = int(torch.argmin(self.min_potentials)) + point_ind = int(self.argmin_potentials[cloud_ind]) + + # Get potential points from tree structure + pot_points = np.array(self.pot_trees[cloud_ind].data, copy=False) + + # Center point of input region + center_point = pot_points[point_ind, :].reshape(1, -1) + + # Add a small noise to center point + if self.set != 'ERF': + center_point += np.random.normal(scale=self.config.in_radius/10, size=center_point.shape) + + # Indices of points in input region + pot_inds, dists = self.pot_trees[cloud_ind].query_radius(center_point, + r=self.config.in_radius, + return_distance=True) + + d2s = np.square(dists[0]) + pot_inds = pot_inds[0] + + # Update potentials (Tukey weights) + if self.set != 'ERF': + tukeys = np.square(1 - d2s / np.square(self.config.in_radius)) + tukeys[d2s > np.square(self.config.in_radius)] = 0 + self.potentials[cloud_ind][pot_inds] += tukeys + min_ind = torch.argmin(self.potentials[cloud_ind]) + self.min_potentials[[cloud_ind]] = self.potentials[cloud_ind][min_ind] + self.argmin_potentials[[cloud_ind]] = min_ind + + # Get points from tree structure + points = np.array(self.input_trees[cloud_ind].data, copy=False) + + # Indices of points in input region + input_inds = self.input_trees[cloud_ind].query_radius(center_point, + r=self.config.in_radius)[0] + + # Number collected + n = input_inds.shape[0] + + # Collect labels and colors + input_points = (points[input_inds] - center_point).astype(np.float32) + input_colors = self.input_colors[cloud_ind][input_inds] + if self.set in ['test', 'ERF']: + input_labels = np.zeros(input_points.shape[0]) + else: + input_labels = self.input_labels[cloud_ind][input_inds] + input_labels = np.array([self.label_to_idx[l] for l in input_labels]) + + # Data augmentation + input_points, scale, R = self.augmentation_transform(input_points) + + # Color augmentation + if np.random.rand() > self.config.augment_color: + input_colors *= 0 + + # Get original height as additional feature + input_features = np.hstack((input_colors, input_points[:, 2:] + center_point[:, 2:])).astype(np.float32) + + # Stack batch + p_list += [input_points] + f_list += [input_features] + l_list += [input_labels] + pi_list += [input_inds] + i_list += [point_ind] + ci_list += [cloud_ind] + s_list += [scale] + R_list += [R] + + # Update batch size + batch_n += n + + # In case batch is full, stop + if batch_n > int(self.batch_limit): + break + + # Randomly drop some points (act as an augmentation process and a safety for GPU memory consumption) + #if n > int(self.batch_limit): + # input_inds = np.random.choice(input_inds, size=int(self.batch_limit) - 1, replace=False) + # n = input_inds.shape[0] + + ################### + # Concatenate batch + ################### + + stacked_points = np.concatenate(p_list, axis=0) + features = np.concatenate(f_list, axis=0) + labels = np.concatenate(l_list, axis=0) + point_inds = np.array(i_list, dtype=np.int32) + cloud_inds = np.array(ci_list, dtype=np.int32) + input_inds = np.concatenate(pi_list, axis=0) + stack_lengths = np.array([pp.shape[0] for pp in p_list], dtype=np.int32) + scales = np.array(s_list, dtype=np.float32) + rots = np.stack(R_list, axis=0) + + # Input features + stacked_features = np.ones_like(stacked_points[:, :1], dtype=np.float32) + if self.config.in_features_dim == 1: + pass + elif self.config.in_features_dim == 4: + stacked_features = np.hstack((stacked_features, features[:, :3])) + elif self.config.in_features_dim == 5: + stacked_features = np.hstack((stacked_features, features)) + else: + raise ValueError('Only accepted input dimensions are 1, 4 and 7 (without and with XYZ)') + + ####################### + # Create network inputs + ####################### + # + # Points, neighbors, pooling indices for each layers + # + + # Get the whole input list + input_list = self.segmentation_inputs(stacked_points, + stacked_features, + labels, + stack_lengths) + + # Add scale and rotation for testing + input_list += [scales, rots, cloud_inds, point_inds, input_inds] + + if debug_workers: + message = '' + for wi in range(info.num_workers): + if wi == wid: + message += ' {:}0{:} '.format(bcolors.OKBLUE, bcolors.ENDC) + elif self.worker_waiting[wi] == 0: + message += ' ' + elif self.worker_waiting[wi] == 1: + message += ' | ' + elif self.worker_waiting[wi] == 2: + message += ' o ' + print(message) + self.worker_waiting[wid] = 2 + + return input_list + + def prepare_S3DIS_ply(self): + + print('\nPreparing ply files') + t0 = time.time() + + # Folder for the ply files + ply_path = join(self.path, self.train_path) + if not exists(ply_path): + makedirs(ply_path) + + for cloud_name in self.cloud_names: + + # Pass if the cloud has already been computed + cloud_file = join(ply_path, cloud_name + '.ply') + if exists(cloud_file): + continue + + # Get rooms of the current cloud + cloud_folder = join(self.path, cloud_name) + room_folders = [join(cloud_folder, room) for room in listdir(cloud_folder) if isdir(join(cloud_folder, room))] + + # Initiate containers + cloud_points = np.empty((0, 3), dtype=np.float32) + cloud_colors = np.empty((0, 3), dtype=np.uint8) + cloud_classes = np.empty((0, 1), dtype=np.int32) + + # Loop over rooms + for i, room_folder in enumerate(room_folders): + + print('Cloud %s - Room %d/%d : %s' % (cloud_name, i+1, len(room_folders), room_folder.split('/')[-1])) + + for object_name in listdir(join(room_folder, 'Annotations')): + + if object_name[-4:] == '.txt': + + # Text file containing point of the object + object_file = join(room_folder, 'Annotations', object_name) + + # Object class and ID + tmp = object_name[:-4].split('_')[0] + if tmp in self.name_to_label: + object_class = self.name_to_label[tmp] + elif tmp in ['stairs']: + object_class = self.name_to_label['clutter'] + else: + raise ValueError('Unknown object name: ' + str(tmp)) + + # Correct bug in S3DIS dataset + if object_name == 'ceiling_1.txt': + with open(object_file, 'r') as f: + lines = f.readlines() + for l_i, line in enumerate(lines): + if '103.0\x100000' in line: + lines[l_i] = line.replace('103.0\x100000', '103.000000') + with open(object_file, 'w') as f: + f.writelines(lines) + + # Read object points and colors + object_data = np.loadtxt(object_file, dtype=np.float32) + + # Stack all data + cloud_points = np.vstack((cloud_points, object_data[:, 0:3].astype(np.float32))) + cloud_colors = np.vstack((cloud_colors, object_data[:, 3:6].astype(np.uint8))) + object_classes = np.full((object_data.shape[0], 1), object_class, dtype=np.int32) + cloud_classes = np.vstack((cloud_classes, object_classes)) + + # Save as ply + write_ply(cloud_file, + (cloud_points, cloud_colors, cloud_classes), + ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) + + print('Done in {:.1f}s'.format(time.time() - t0)) + return + + def load_subsampled_clouds(self): + + # Parameter + dl = self.config.first_subsampling_dl + + # Create path for files + tree_path = join(self.path, 'input_{:.3f}'.format(dl)) + if not exists(tree_path): + makedirs(tree_path) + + ############## + # Load KDTrees + ############## + + for i, file_path in enumerate(self.train_files): + + # Restart timer + t0 = time.time() + + # Skip split that is not in current set + if self.set == 'training': + if self.all_splits[i] == self.validation_split: + continue + elif self.set in ['validation', 'test', 'ERF']: + if self.all_splits[i] != self.validation_split: + continue + else: + raise ValueError('Unknown set for S3DIS data: ', self.set) + + # Get cloud name + cloud_name = self.cloud_names[i] + + # Name of the input files + KDTree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) + sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) + + # Check if inputs have already been computed + if exists(KDTree_file): + print('\nFound KDTree for cloud {:s}, subsampled at {:.3f}'.format(cloud_name, dl)) + + # read ply with data + data = read_ply(sub_ply_file) + sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T + sub_labels = data['class'] + + # Read pkl with search tree + with open(KDTree_file, 'rb') as f: + search_tree = pickle.load(f) + + else: + print('\nPreparing KDTree for cloud {:s}, subsampled at {:.3f}'.format(cloud_name, dl)) + + # Read ply file + data = read_ply(file_path) + points = np.vstack((data['x'], data['y'], data['z'])).T + colors = np.vstack((data['red'], data['green'], data['blue'])).T + labels = data['class'] + + # Subsample cloud + sub_points, sub_colors, sub_labels = grid_subsampling(points, + features=colors, + labels=labels, + sampleDl=dl) + + # Rescale float color and squeeze label + sub_colors = sub_colors / 255 + sub_labels = np.squeeze(sub_labels) + + # Get chosen neighborhoods + search_tree = KDTree(sub_points, leaf_size=10) + #search_tree = nnfln.KDTree(n_neighbors=1, metric='L2', leaf_size=10) + #search_tree.fit(sub_points) + + # Save KDTree + with open(KDTree_file, 'wb') as f: + pickle.dump(search_tree, f) + + # Save ply + write_ply(sub_ply_file, + [sub_points, sub_colors, sub_labels], + ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) + + # Fill data containers + self.input_trees += [search_tree] + self.input_colors += [sub_colors] + self.input_labels += [sub_labels] + + size = sub_colors.shape[0] * 4 * 7 + print('{:.1f} MB loaded in {:.1f}s'.format(size * 1e-6, time.time() - t0)) + + ############################ + # Coarse potential locations + ############################ + + # Only necessary for validation and test sets + if self.use_potentials: + print('\nPreparing potentials') + + # Restart timer + t0 = time.time() + + pot_dl = self.config.in_radius / 10 + cloud_ind = 0 + + for i, file_path in enumerate(self.train_files): + + # Skip split that is not in current set + if self.set == 'training': + if self.all_splits[i] == self.validation_split: + continue + elif self.set in ['validation', 'test', 'ERF']: + if self.all_splits[i] != self.validation_split: + continue + else: + raise ValueError('Unknown set for S3DIS data: ', self.set) + + # Get cloud name + cloud_name = self.cloud_names[i] + + # Name of the input files + coarse_KDTree_file = join(tree_path, '{:s}_coarse_KDTree.pkl'.format(cloud_name)) + + # Check if inputs have already been computed + if exists(coarse_KDTree_file): + # Read pkl with search tree + with open(coarse_KDTree_file, 'rb') as f: + search_tree = pickle.load(f) + + else: + # Subsample cloud + sub_points = np.array(self.input_trees[cloud_ind].data, copy=False) + coarse_points = grid_subsampling(sub_points.astype(np.float32), sampleDl=pot_dl) + + # Get chosen neighborhoods + search_tree = KDTree(coarse_points, leaf_size=10) + + # Save KDTree + with open(coarse_KDTree_file, 'wb') as f: + pickle.dump(search_tree, f) + + # Fill data containers + self.pot_trees += [search_tree] + cloud_ind += 1 + + print('Done in {:.1f}s'.format(time.time() - t0)) + + ###################### + # Reprojection indices + ###################### + + # Only necessary for validation and test sets + if self.set in ['validation', 'test']: + + print('\nPreparing reprojection indices for testing') + + # Get number of clouds + self.num_clouds = len(self.input_trees) + + # Get validation/test reprojection indices + i_cloud = 0 + for i, file_path in enumerate(self.train_files): + + # Skip split that is not in current set + if self.all_splits[i] != self.validation_split: + continue + + # Restart timer + t0 = time.time() + + # Get info on this cloud + cloud_name = self.cloud_names[i] + + # File name for saving + proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) + + # Try to load previous indices + if exists(proj_file): + with open(proj_file, 'rb') as f: + proj_inds, labels = pickle.load(f) + else: + data = read_ply(file_path) + points = np.vstack((data['x'], data['y'], data['z'])).T + labels = data['class'] + + # Compute projection inds + idxs = self.input_trees[i_cloud].query(points, return_distance=False) + #dists, idxs = self.input_trees[i_cloud].kneighbors(points) + proj_inds = np.squeeze(idxs).astype(np.int32) + + # Save + with open(proj_file, 'wb') as f: + pickle.dump([proj_inds, labels], f) + + self.validation_proj += [proj_inds] + self.validation_labels += [labels] + i_cloud += 1 + print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) + + print() + return + + def load_evaluation_points(self, file_path): + """ + Load points (from test or validation split) on which the metrics should be evaluated + """ + + # Get original points + data = read_ply(file_path) + return np.vstack((data['x'], data['y'], data['z'])).T +# ---------------------------------------------------------------------------------------------------------------------- +# +# Utility classes definition +# \********************************/ + + +class S3DISSampler(Sampler): + """Sampler for S3DIS""" + + def __init__(self, dataset: S3DISDataset,): + Sampler.__init__(self, dataset) + + # Dataset used by the sampler (no copy is made in memory) + self.dataset = dataset + + # Number of step per epoch + if dataset.set == 'training': + self.N = dataset.config.epoch_steps + else: + self.N = dataset.config.validation_size + + return + + def __iter__(self): + """ + Yield next batch indices here. In this dataset, this is a dummy sampler that yield the index of batch element + (input sphere) in epoch instead of the list of point indices + """ + + # Generator loop + for i in range(self.N): + yield i + + def __len__(self): + """ + The number of yielded samples is variable + """ + return None + + def fast_calib(self): + """ + This method calibrates the batch sizes while ensuring the potentials are well initialized. Indeed on a dataset + like Semantic3D, before potential have been updated over the dataset, there are cahnces that all the dense area + are picked in the begining and in the end, we will have very large batch of small point clouds + :return: + """ + + # Estimated average batch size and target value + estim_b = 0 + target_b = self.dataset.config.batch_num + + # Calibration parameters + low_pass_T = 10 + Kp = 100.0 + finer = False + breaking = False + + # Convergence parameters + smooth_errors = [] + converge_threshold = 0.1 + + t = [time.time()] + last_display = time.time() + mean_dt = np.zeros(2) + + for epoch in range(10): + for i, test in enumerate(self): + + # New time + t = t[-1:] + t += [time.time()] + + # batch length + b = len(test) + + # Update estim_b (low pass filter) + estim_b += (b - estim_b) / low_pass_T + + # Estimate error (noisy) + error = target_b - b + + # Save smooth errors for convergene check + smooth_errors.append(target_b - estim_b) + if len(smooth_errors) > 10: + smooth_errors = smooth_errors[1:] + + # Update batch limit with P controller + self.dataset.batch_limit += Kp * error + + # finer low pass filter when closing in + if not finer and np.abs(estim_b - target_b) < 1: + low_pass_T = 100 + finer = True + + # Convergence + if finer and np.max(np.abs(smooth_errors)) < converge_threshold: + breaking = True + break + + # Average timing + t += [time.time()] + mean_dt = 0.9 * mean_dt + 0.1 * (np.array(t[1:]) - np.array(t[:-1])) + + # Console display (only one per second) + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + message = 'Step {:5d} estim_b ={:5.2f} batch_limit ={:7d}, // {:.1f}ms {:.1f}ms' + print(message.format(i, + estim_b, + int(self.dataset.batch_limit), + 1000 * mean_dt[0], + 1000 * mean_dt[1])) + + if breaking: + break + + def calibration(self, dataloader, untouched_ratio=0.9, verbose=False): + """ + Method performing batch and neighbors calibration. + Batch calibration: Set "batch_limit" (the maximum number of points allowed in every batch) so that the + average batch size (number of stacked pointclouds) is the one asked. + Neighbors calibration: Set the "neighborhood_limits" (the maximum number of neighbors allowed in convolutions) + so that 90% of the neighborhoods remain untouched. There is a limit for each layer. + """ + + ############################## + # Previously saved calibration + ############################## + + print('\nStarting Calibration (use verbose=True for more details)') + t0 = time.time() + + redo = False + + # Batch limit + # *********** + + # Load batch_limit dictionary + batch_lim_file = join(self.dataset.path, 'batch_limits.pkl') + if exists(batch_lim_file): + with open(batch_lim_file, 'rb') as file: + batch_lim_dict = pickle.load(file) + else: + batch_lim_dict = {} + + # Check if the batch limit associated with current parameters exists + key = '{:.3f}_{:.3f}_{:d}'.format(self.dataset.config.in_radius, + self.dataset.config.first_subsampling_dl, + self.dataset.config.batch_num) + if key in batch_lim_dict: + self.dataset.batch_limit[0] = batch_lim_dict[key] + else: + redo = True + + if verbose: + print('\nPrevious calibration found:') + print('Check batch limit dictionary') + if key in batch_lim_dict: + color = bcolors.OKGREEN + v = str(int(batch_lim_dict[key])) + else: + color = bcolors.FAIL + v = '?' + print('{:}\"{:s}\": {:s}{:}'.format(color, key, v, bcolors.ENDC)) + + # Neighbors limit + # *************** + + # Load neighb_limits dictionary + neighb_lim_file = join(self.dataset.path, 'neighbors_limits.pkl') + if exists(neighb_lim_file): + with open(neighb_lim_file, 'rb') as file: + neighb_lim_dict = pickle.load(file) + else: + neighb_lim_dict = {} + + # Check if the limit associated with current parameters exists (for each layer) + neighb_limits = [] + for layer_ind in range(self.dataset.config.num_layers): + + dl = self.dataset.config.first_subsampling_dl * (2**layer_ind) + if self.dataset.config.deform_layers[layer_ind]: + r = dl * self.dataset.config.deform_radius + else: + r = dl * self.dataset.config.conv_radius + + key = '{:.3f}_{:.3f}'.format(dl, r) + if key in neighb_lim_dict: + neighb_limits += [neighb_lim_dict[key]] + + if len(neighb_limits) == self.dataset.config.num_layers: + self.dataset.neighborhood_limits = neighb_limits + else: + redo = True + + if verbose: + print('Check neighbors limit dictionary') + for layer_ind in range(self.dataset.config.num_layers): + dl = self.dataset.config.first_subsampling_dl * (2**layer_ind) + if self.dataset.config.deform_layers[layer_ind]: + r = dl * self.dataset.config.deform_radius + else: + r = dl * self.dataset.config.conv_radius + key = '{:.3f}_{:.3f}'.format(dl, r) + + if key in neighb_lim_dict: + color = bcolors.OKGREEN + v = str(neighb_lim_dict[key]) + else: + color = bcolors.FAIL + v = '?' + print('{:}\"{:s}\": {:s}{:}'.format(color, key, v, bcolors.ENDC)) + + if redo: + + ############################ + # Neighbors calib parameters + ############################ + + # From config parameter, compute higher bound of neighbors number in a neighborhood + hist_n = int(np.ceil(4 / 3 * np.pi * (self.dataset.config.deform_radius + 1) ** 3)) + + print(hist_n) + + # Histogram of neighborhood sizes + neighb_hists = np.zeros((self.dataset.config.num_layers, hist_n), dtype=np.int32) + + ######################## + # Batch calib parameters + ######################## + + # Estimated average batch size and target value + estim_b = 0 + target_b = self.dataset.config.batch_num + + # Calibration parameters + low_pass_T = 10 + Kp = 100.0 + finer = False + + # Convergence parameters + smooth_errors = [] + converge_threshold = 0.1 + + # Loop parameters + last_display = time.time() + i = 0 + breaking = False + + ##################### + # Perform calibration + ##################### + + for epoch in range(10): + for batch_i, batch in enumerate(dataloader): + + # Update neighborhood histogram + counts = [np.sum(neighb_mat.numpy() < neighb_mat.shape[0], axis=1) for neighb_mat in batch.neighbors] + hists = [np.bincount(c, minlength=hist_n)[:hist_n] for c in counts] + neighb_hists += np.vstack(hists) + + # batch length + b = len(batch.cloud_inds) + + # Update estim_b (low pass filter) + estim_b += (b - estim_b) / low_pass_T + + # Estimate error (noisy) + error = target_b - b + + # Save smooth errors for convergene check + smooth_errors.append(target_b - estim_b) + if len(smooth_errors) > 10: + smooth_errors = smooth_errors[1:] + + # Update batch limit with P controller + self.dataset.batch_limit += Kp * error + + # finer low pass filter when closing in + if not finer and np.abs(estim_b - target_b) < 1: + low_pass_T = 100 + finer = True + + # Convergence + if finer and np.max(np.abs(smooth_errors)) < converge_threshold: + breaking = True + break + + i += 1 + t = time.time() + + # Console display (only one per second) + if verbose and (t - last_display) > 1.0: + last_display = t + message = 'Step {:5d} estim_b ={:5.2f} batch_limit ={:7d}' + print(message.format(i, + estim_b, + int(self.dataset.batch_limit))) + + if breaking: + break + + # Use collected neighbor histogram to get neighbors limit + cumsum = np.cumsum(neighb_hists.T, axis=0) + percentiles = np.sum(cumsum < (untouched_ratio * cumsum[hist_n - 1, :]), axis=0) + self.dataset.neighborhood_limits = percentiles + + if verbose: + + # Crop histogram + while np.sum(neighb_hists[:, -1]) == 0: + neighb_hists = neighb_hists[:, :-1] + hist_n = neighb_hists.shape[1] + + print('\n**************************************************\n') + line0 = 'neighbors_num ' + for layer in range(neighb_hists.shape[0]): + line0 += '| layer {:2d} '.format(layer) + print(line0) + for neighb_size in range(hist_n): + line0 = ' {:4d} '.format(neighb_size) + for layer in range(neighb_hists.shape[0]): + if neighb_size > percentiles[layer]: + color = bcolors.FAIL + else: + color = bcolors.OKGREEN + line0 += '|{:}{:10d}{:} '.format(color, + neighb_hists[layer, neighb_size], + bcolors.ENDC) + + print(line0) + + print('\n**************************************************\n') + print('\nchosen neighbors limits: ', percentiles) + print() + + # Save batch_limit dictionary + key = '{:.3f}_{:.3f}_{:d}'.format(self.dataset.config.in_radius, + self.dataset.config.first_subsampling_dl, + self.dataset.config.batch_num) + batch_lim_dict[key] = float(self.dataset.batch_limit) + with open(batch_lim_file, 'wb') as file: + pickle.dump(batch_lim_dict, file) + + # Save neighb_limit dictionary + for layer_ind in range(self.dataset.config.num_layers): + dl = self.dataset.config.first_subsampling_dl * (2 ** layer_ind) + if self.dataset.config.deform_layers[layer_ind]: + r = dl * self.dataset.config.deform_radius + else: + r = dl * self.dataset.config.conv_radius + key = '{:.3f}_{:.3f}'.format(dl, r) + neighb_lim_dict[key] = self.dataset.neighborhood_limits[layer_ind] + with open(neighb_lim_file, 'wb') as file: + pickle.dump(neighb_lim_dict, file) + + + print('Calibration done in {:.1f}s\n'.format(time.time() - t0)) + return + + +class S3DISCustomBatch: + """Custom batch definition with memory pinning for S3DIS""" + + def __init__(self, input_list): + + # Get rid of batch dimension + input_list = input_list[0] + + # Number of layers + L = (len(input_list) - 7) // 5 + + # Extract input tensors from the list of numpy array + ind = 0 + self.points = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]] + ind += L + self.neighbors = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]] + ind += L + self.pools = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]] + ind += L + self.upsamples = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]] + ind += L + self.lengths = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]] + ind += L + self.features = torch.from_numpy(input_list[ind]) + ind += 1 + self.labels = torch.from_numpy(input_list[ind]) + ind += 1 + self.scales = torch.from_numpy(input_list[ind]) + ind += 1 + self.rots = torch.from_numpy(input_list[ind]) + ind += 1 + self.cloud_inds = torch.from_numpy(input_list[ind]) + ind += 1 + self.center_inds = torch.from_numpy(input_list[ind]) + ind += 1 + self.input_inds = torch.from_numpy(input_list[ind]) + + return + + def pin_memory(self): + """ + Manual pinning of the memory + """ + + self.points = [in_tensor.pin_memory() for in_tensor in self.points] + self.neighbors = [in_tensor.pin_memory() for in_tensor in self.neighbors] + self.pools = [in_tensor.pin_memory() for in_tensor in self.pools] + self.upsamples = [in_tensor.pin_memory() for in_tensor in self.upsamples] + self.lengths = [in_tensor.pin_memory() for in_tensor in self.lengths] + self.features = self.features.pin_memory() + self.labels = self.labels.pin_memory() + self.scales = self.scales.pin_memory() + self.rots = self.rots.pin_memory() + self.cloud_inds = self.cloud_inds.pin_memory() + self.center_inds = self.center_inds.pin_memory() + self.input_inds = self.input_inds.pin_memory() + + return self + + def to(self, device): + + self.points = [in_tensor.to(device) for in_tensor in self.points] + self.neighbors = [in_tensor.to(device) for in_tensor in self.neighbors] + self.pools = [in_tensor.to(device) for in_tensor in self.pools] + self.upsamples = [in_tensor.to(device) for in_tensor in self.upsamples] + self.lengths = [in_tensor.to(device) for in_tensor in self.lengths] + self.features = self.features.to(device) + self.labels = self.labels.to(device) + self.scales = self.scales.to(device) + self.rots = self.rots.to(device) + self.cloud_inds = self.cloud_inds.to(device) + self.center_inds = self.center_inds.to(device) + self.input_inds = self.input_inds.to(device) + + return self + + def unstack_points(self, layer=None): + """Unstack the points""" + return self.unstack_elements('points', layer) + + def unstack_neighbors(self, layer=None): + """Unstack the neighbors indices""" + return self.unstack_elements('neighbors', layer) + + def unstack_pools(self, layer=None): + """Unstack the pooling indices""" + return self.unstack_elements('pools', layer) + + def unstack_elements(self, element_name, layer=None, to_numpy=True): + """ + Return a list of the stacked elements in the batch at a certain layer. If no layer is given, then return all + layers + """ + + if element_name == 'points': + elements = self.points + elif element_name == 'neighbors': + elements = self.neighbors + elif element_name == 'pools': + elements = self.pools[:-1] + else: + raise ValueError('Unknown element name: {:s}'.format(element_name)) + + all_p_list = [] + for layer_i, layer_elems in enumerate(elements): + + if layer is None or layer == layer_i: + + i0 = 0 + p_list = [] + if element_name == 'pools': + lengths = self.lengths[layer_i+1] + else: + lengths = self.lengths[layer_i] + + for b_i, length in enumerate(lengths): + + elem = layer_elems[i0:i0 + length] + if element_name == 'neighbors': + elem[elem >= self.points[layer_i].shape[0]] = -1 + elem[elem >= 0] -= i0 + elif element_name == 'pools': + elem[elem >= self.points[layer_i].shape[0]] = -1 + elem[elem >= 0] -= torch.sum(self.lengths[layer_i][:b_i]) + i0 += length + + if to_numpy: + p_list.append(elem.numpy()) + else: + p_list.append(elem) + + if layer == layer_i: + return p_list + + all_p_list.append(p_list) + + return all_p_list + + +def S3DISCollate(batch_data): + return S3DISCustomBatch(batch_data) + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Debug functions +# \*********************/ + + +def debug_sampling(dataset, sampler, loader): + """Shows which labels are sampled according to strategy chosen""" + label_sum = np.zeros((dataset.num_classes), dtype=np.int32) + for epoch in range(10): + + for batch_i, (points, normals, labels, indices, in_sizes) in enumerate(loader): + # print(batch_i, tuple(points.shape), tuple(normals.shape), labels, indices, in_sizes) + + label_sum += np.bincount(labels.numpy(), minlength=dataset.num_classes) + print(label_sum) + #print(sampler.potentials[:6]) + + print('******************') + print('*******************************************') + + _, counts = np.unique(dataset.input_labels, return_counts=True) + print(counts) + + +def debug_timing(dataset, sampler, loader): + """Timing of generator function""" + + t = [time.time()] + last_display = time.time() + mean_dt = np.zeros(2) + estim_b = dataset.config.batch_num + + for epoch in range(10): + + for batch_i, batch in enumerate(loader): + # print(batch_i, tuple(points.shape), tuple(normals.shape), labels, indices, in_sizes) + + # New time + t = t[-1:] + t += [time.time()] + + # Update estim_b (low pass filter) + estim_b += (len(batch.cloud_inds) - estim_b) / 100 + + # Pause simulating computations + time.sleep(0.05) + t += [time.time()] + + # Average timing + mean_dt = 0.9 * mean_dt + 0.1 * (np.array(t[1:]) - np.array(t[:-1])) + + # Console display (only one per second) + if (t[-1] - last_display) > -1.0: + last_display = t[-1] + message = 'Step {:08d} -> (ms/batch) {:8.2f} {:8.2f} / batch = {:.2f}' + print(message.format(batch_i, + 1000 * mean_dt[0], + 1000 * mean_dt[1], + estim_b)) + + print('************* Epoch ended *************') + + _, counts = np.unique(dataset.input_labels, return_counts=True) + print(counts) + + +def debug_show_clouds(dataset, sampler, loader): + + + for epoch in range(10): + + clouds = [] + cloud_normals = [] + cloud_labels = [] + + L = dataset.config.num_layers + + for batch_i, batch in enumerate(loader): + + # Print characteristics of input tensors + print('\nPoints tensors') + for i in range(L): + print(batch.points[i].dtype, batch.points[i].shape) + print('\nNeigbors tensors') + for i in range(L): + print(batch.neighbors[i].dtype, batch.neighbors[i].shape) + print('\nPools tensors') + for i in range(L): + print(batch.pools[i].dtype, batch.pools[i].shape) + print('\nStack lengths') + for i in range(L): + print(batch.lengths[i].dtype, batch.lengths[i].shape) + print('\nFeatures') + print(batch.features.dtype, batch.features.shape) + print('\nLabels') + print(batch.labels.dtype, batch.labels.shape) + print('\nAugment Scales') + print(batch.scales.dtype, batch.scales.shape) + print('\nAugment Rotations') + print(batch.rots.dtype, batch.rots.shape) + print('\nModel indices') + print(batch.model_inds.dtype, batch.model_inds.shape) + + print('\nAre input tensors pinned') + print(batch.neighbors[0].is_pinned()) + print(batch.neighbors[-1].is_pinned()) + print(batch.points[0].is_pinned()) + print(batch.points[-1].is_pinned()) + print(batch.labels.is_pinned()) + print(batch.scales.is_pinned()) + print(batch.rots.is_pinned()) + print(batch.model_inds.is_pinned()) + + show_input_batch(batch) + + print('*******************************************') + + _, counts = np.unique(dataset.input_labels, return_counts=True) + print(counts) + + +def debug_batch_and_neighbors_calib(dataset, sampler, loader): + """Timing of generator function""" + + t = [time.time()] + last_display = time.time() + mean_dt = np.zeros(2) + + for epoch in range(10): + + for batch_i, input_list in enumerate(loader): + # print(batch_i, tuple(points.shape), tuple(normals.shape), labels, indices, in_sizes) + + # New time + t = t[-1:] + t += [time.time()] + + # Pause simulating computations + time.sleep(0.01) + t += [time.time()] + + # Average timing + mean_dt = 0.9 * mean_dt + 0.1 * (np.array(t[1:]) - np.array(t[:-1])) + + # Console display (only one per second) + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + message = 'Step {:08d} -> Average timings (ms/batch) {:8.2f} {:8.2f} ' + print(message.format(batch_i, + 1000 * mean_dt[0], + 1000 * mean_dt[1])) + + print('************* Epoch ended *************') + + _, counts = np.unique(dataset.input_labels, return_counts=True) + print(counts) diff --git a/datasets/common.py b/datasets/common.py new file mode 100644 index 0000000..845698d --- /dev/null +++ b/datasets/common.py @@ -0,0 +1,517 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Class handling datasets +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 11/06/2018 +# + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Imports and global variables +# \**********************************/ +# + +# Common libs +import time +import os +import numpy as np +import sys +import torch +from torch.utils.data import DataLoader, Dataset +from utils.config import Config +from utils.mayavi_visu import * +from kernels.kernel_points import create_3D_rotations + +# Subsampling extension +import cpp_wrappers.cpp_subsampling.grid_subsampling as cpp_subsampling +import cpp_wrappers.cpp_neighbors.radius_neighbors as cpp_neighbors + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Utility functions +# \***********************/ +# + +def grid_subsampling(points, features=None, labels=None, sampleDl=0.1, verbose=0): + """ + CPP wrapper for a grid subsampling (method = barycenter for points and features) + :param points: (N, 3) matrix of input points + :param features: optional (N, d) matrix of features (floating number) + :param labels: optional (N,) matrix of integer labels + :param sampleDl: parameter defining the size of grid voxels + :param verbose: 1 to display + :return: subsampled points, with features and/or labels depending of the input + """ + + if (features is None) and (labels is None): + return cpp_subsampling.subsample(points, + sampleDl=sampleDl, + verbose=verbose) + elif (labels is None): + return cpp_subsampling.subsample(points, + features=features, + sampleDl=sampleDl, + verbose=verbose) + elif (features is None): + return cpp_subsampling.subsample(points, + classes=labels, + sampleDl=sampleDl, + verbose=verbose) + else: + return cpp_subsampling.subsample(points, + features=features, + classes=labels, + sampleDl=sampleDl, + verbose=verbose) + + +def batch_grid_subsampling(points, batches_len, features=None, labels=None, sampleDl=0.1, max_p=0, verbose=0): + """ + CPP wrapper for a grid subsampling (method = barycenter for points and features) + :param points: (N, 3) matrix of input points + :param features: optional (N, d) matrix of features (floating number) + :param labels: optional (N,) matrix of integer labels + :param sampleDl: parameter defining the size of grid voxels + :param verbose: 1 to display + :return: subsampled points, with features and/or labels depending of the input + """ + + if (features is None) and (labels is None): + return cpp_subsampling.subsample_batch(points, + batches_len, + sampleDl=sampleDl, + max_p=max_p, + verbose=verbose) + elif (labels is None): + return cpp_subsampling.subsample_batch(points, + batches_len, + features=features, + sampleDl=sampleDl, + max_p=max_p, + verbose=verbose) + elif (features is None): + return cpp_subsampling.subsample_batch(points, + batches_len, + classes=labels, + sampleDl=sampleDl, + max_p=max_p, + verbose=verbose) + else: + return cpp_subsampling.subsample_batch(points, + batches_len, + features=features, + classes=labels, + sampleDl=sampleDl, + max_p=max_p, + verbose=verbose) + + +def batch_neighbors(queries, supports, q_batches, s_batches, radius): + """ + Computes neighbors for a batch of queries and supports + :param queries: (N1, 3) the query points + :param supports: (N2, 3) the support points + :param q_batches: (B) the list of lengths of batch elements in queries + :param s_batches: (B)the list of lengths of batch elements in supports + :param radius: float32 + :return: neighbors indices + """ + + return cpp_neighbors.batch_query(queries, supports, q_batches, s_batches, radius=radius) + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Class definition +# \**********************/ + + +class PointCloudDataset(Dataset): + """Parent class for Point Cloud Datasets.""" + + def __init__(self, name): + """ + Initialize parameters of the dataset here. + """ + + self.name = name + self.path = '' + self.label_to_names = {} + self.num_classes = 0 + self.label_values = np.zeros((0,), dtype=np.int32) + self.label_names = [] + self.label_to_idx = {} + self.name_to_label = {} + self.config = Config() + self.neighborhood_limits = [] + + return + + def __len__(self): + """ + Return the length of data here + """ + return 0 + + def __getitem__(self, idx): + """ + Return the item at the given index + """ + + return 0 + + def init_labels(self): + + # Initialize all label parameters given the label_to_names dict + self.num_classes = len(self.label_to_names) + self.label_values = np.sort([k for k, v in self.label_to_names.items()]) + self.label_names = [self.label_to_names[k] for k in self.label_values] + self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} + self.name_to_label = {v: k for k, v in self.label_to_names.items()} + + def augmentation_transform(self, points, normals=None, verbose=False): + """Implementation of an augmentation transform for point clouds.""" + + ########## + # Rotation + ########## + + # Initialize rotation matrix + R = np.eye(points.shape[1]) + + if points.shape[1] == 3: + if self.config.augment_rotation == 'vertical': + + # Create random rotations + theta = np.random.rand() * 2 * np.pi + c, s = np.cos(theta), np.sin(theta) + R = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]], dtype=np.float32) + + elif self.config.augment_rotation == 'all': + + # Choose two random angles for the first vector in polar coordinates + theta = np.random.rand() * 2 * np.pi + phi = (np.random.rand() - 0.5) * np.pi + + # Create the first vector in carthesian coordinates + u = np.array([np.cos(theta) * np.cos(phi), np.sin(theta) * np.cos(phi), np.sin(phi)]) + + # Choose a random rotation angle + alpha = np.random.rand() * 2 * np.pi + + # Create the rotation matrix with this vector and angle + R = create_3D_rotations(np.reshape(u, (1, -1)), np.reshape(alpha, (1, -1)))[0] + + R = R.astype(np.float32) + + ####### + # Scale + ####### + + # Choose random scales for each example + min_s = self.config.augment_scale_min + max_s = self.config.augment_scale_max + if self.config.augment_scale_anisotropic: + scale = np.random.rand(points.shape[1]) * (max_s - min_s) + min_s + else: + scale = np.random.rand() * (max_s - min_s) - min_s + + # Add random symmetries to the scale factor + symmetries = np.array(self.config.augment_symmetries).astype(np.int32) + symmetries *= np.random.randint(2, size=points.shape[1]) + scale = (scale * symmetries * 2 - 1).astype(np.float32) + + ####### + # Noise + ####### + + noise = (np.random.randn(points.shape[0], points.shape[1]) * self.config.augment_noise).astype(np.float32) + + ################## + # Apply transforms + ################## + + augmented_points = np.dot(points, R) * scale + noise + + if normals is None: + return augmented_points, scale, R + else: + # Anisotropic scale of the normals thanks to cross product formula + normal_scale = scale[[1, 2, 0]] * scale[[2, 0, 1]] + augmented_normals = np.dot(normals, R) * normal_scale + # Renormalise + augmented_normals *= 1 / (np.linalg.norm(augmented_normals, axis=1, keepdims=True) + 1e-6) + + if verbose: + test_p = [np.vstack([points, augmented_points])] + test_n = [np.vstack([normals, augmented_normals])] + test_l = [np.hstack([points[:, 2]*0, augmented_points[:, 2]*0+1])] + show_ModelNet_examples(test_p, test_n, test_l) + + return augmented_points, augmented_normals, scale, R + + def big_neighborhood_filter(self, neighbors, layer): + """ + Filter neighborhoods with max number of neighbors. Limit is set to keep XX% of the neighborhoods untouched. + Limit is computed at initialization + """ + + # crop neighbors matrix + if len(self.neighborhood_limits) > 0: + return neighbors[:, :self.neighborhood_limits[layer]] + else: + return neighbors + + def classification_inputs(self, + stacked_points, + stacked_features, + labels, + stack_lengths): + + # Starting radius of convolutions + r_normal = self.config.first_subsampling_dl * self.config.conv_radius + + # Starting layer + layer_blocks = [] + + # Lists of inputs + input_points = [] + input_neighbors = [] + input_pools = [] + input_stack_lengths = [] + deform_layers = [] + + ###################### + # Loop over the blocks + ###################### + + arch = self.config.architecture + + for block_i, block in enumerate(arch): + + # Get all blocks of the layer + if not ('pool' in block or 'strided' in block or 'global' in block or 'upsample' in block): + layer_blocks += [block] + continue + + # Convolution neighbors indices + # ***************************** + + deform_layer = False + if layer_blocks: + # Convolutions are done in this layer, compute the neighbors with the good radius + if np.any(['deformable' in blck for blck in layer_blocks]): + r = r_normal * self.config.deform_radius / self.config.conv_radius + deform_layer = True + else: + r = r_normal + conv_i = batch_neighbors(stacked_points, stacked_points, stack_lengths, stack_lengths, r) + + else: + # This layer only perform pooling, no neighbors required + conv_i = np.zeros((0, 1), dtype=np.int32) + + # Pooling neighbors indices + # ************************* + + # If end of layer is a pooling operation + if 'pool' in block or 'strided' in block: + + # New subsampling length + dl = 2 * r_normal / self.config.conv_radius + + # Subsampled points + pool_p, pool_b = batch_grid_subsampling(stacked_points, stack_lengths, sampleDl=dl) + + # Radius of pooled neighbors + if 'deformable' in block: + r = r_normal * self.config.deform_radius / self.config.conv_radius + deform_layer = True + else: + r = r_normal + + # Subsample indices + pool_i = batch_neighbors(pool_p, stacked_points, pool_b, stack_lengths, r) + + else: + # No pooling in the end of this layer, no pooling indices required + pool_i = np.zeros((0, 1), dtype=np.int32) + pool_p = np.zeros((0, 3), dtype=np.float32) + pool_b = np.zeros((0,), dtype=np.int32) + + # Reduce size of neighbors matrices by eliminating furthest point + conv_i = self.big_neighborhood_filter(conv_i, len(input_points)) + pool_i = self.big_neighborhood_filter(pool_i, len(input_points)) + + # Updating input lists + input_points += [stacked_points] + input_neighbors += [conv_i.astype(np.int64)] + input_pools += [pool_i.astype(np.int64)] + input_stack_lengths += [stack_lengths] + deform_layers += [deform_layer] + + # New points for next layer + stacked_points = pool_p + stack_lengths = pool_b + + # Update radius and reset blocks + r_normal *= 2 + layer_blocks = [] + + # Stop when meeting a global pooling or upsampling + if 'global' in block or 'upsample' in block: + break + + ############### + # Return inputs + ############### + + # Save deform layers + + # list of network inputs + li = input_points + input_neighbors + input_pools + input_stack_lengths + li += [stacked_features, labels] + + return li + + + def segmentation_inputs(self, + stacked_points, + stacked_features, + labels, + stack_lengths): + + # Starting radius of convolutions + r_normal = self.config.first_subsampling_dl * self.config.conv_radius + + # Starting layer + layer_blocks = [] + + # Lists of inputs + input_points = [] + input_neighbors = [] + input_pools = [] + input_upsamples = [] + input_stack_lengths = [] + deform_layers = [] + + ###################### + # Loop over the blocks + ###################### + + arch = self.config.architecture + + for block_i, block in enumerate(arch): + + # Get all blocks of the layer + if not ('pool' in block or 'strided' in block or 'global' in block or 'upsample' in block): + layer_blocks += [block] + continue + + # Convolution neighbors indices + # ***************************** + + deform_layer = False + if layer_blocks: + # Convolutions are done in this layer, compute the neighbors with the good radius + if np.any(['deformable' in blck for blck in layer_blocks]): + r = r_normal * self.config.deform_radius / self.config.conv_radius + deform_layer = True + else: + r = r_normal + conv_i = batch_neighbors(stacked_points, stacked_points, stack_lengths, stack_lengths, r) + + else: + # This layer only perform pooling, no neighbors required + conv_i = np.zeros((0, 1), dtype=np.int32) + + # Pooling neighbors indices + # ************************* + + # If end of layer is a pooling operation + if 'pool' in block or 'strided' in block: + + # New subsampling length + dl = 2 * r_normal / self.config.conv_radius + + # Subsampled points + pool_p, pool_b = batch_grid_subsampling(stacked_points, stack_lengths, sampleDl=dl) + + # Radius of pooled neighbors + if 'deformable' in block: + r = r_normal * self.config.deform_radius / self.config.conv_radius + deform_layer = True + else: + r = r_normal + + # Subsample indices + pool_i = batch_neighbors(pool_p, stacked_points, pool_b, stack_lengths, r) + + # Upsample indices (with the radius of the next layer to keep wanted density) + up_i = batch_neighbors(stacked_points, pool_p, stack_lengths, pool_b, 2 * r) + + else: + # No pooling in the end of this layer, no pooling indices required + pool_i = np.zeros((0, 1), dtype=np.int32) + pool_p = np.zeros((0, 3), dtype=np.float32) + pool_b = np.zeros((0,), dtype=np.int32) + up_i = np.zeros((0, 1), dtype=np.int32) + + # Reduce size of neighbors matrices by eliminating furthest point + conv_i = self.big_neighborhood_filter(conv_i, len(input_points)) + pool_i = self.big_neighborhood_filter(pool_i, len(input_points)) + up_i = self.big_neighborhood_filter(up_i, len(input_points)) + + # Updating input lists + input_points += [stacked_points] + input_neighbors += [conv_i.astype(np.int64)] + input_pools += [pool_i.astype(np.int64)] + input_upsamples += [up_i.astype(np.int64)] + input_stack_lengths += [stack_lengths] + deform_layers += [deform_layer] + + # New points for next layer + stacked_points = pool_p + stack_lengths = pool_b + + # Update radius and reset blocks + r_normal *= 2 + layer_blocks = [] + + # Stop when meeting a global pooling or upsampling + if 'global' in block or 'upsample' in block: + break + + ############### + # Return inputs + ############### + + # Save deform layers + + # list of network inputs + li = input_points + input_neighbors + input_pools + input_upsamples + input_stack_lengths + li += [stacked_features, labels] + + return li + + + + + + + + + + + + + diff --git a/kernels/kernel_points.py b/kernels/kernel_points.py new file mode 100644 index 0000000..9eb6d16 --- /dev/null +++ b/kernels/kernel_points.py @@ -0,0 +1,481 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Functions handling the disposition of kernel points. +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 11/06/2018 +# + + +# ------------------------------------------------------------------------------------------ +# +# Imports and global variables +# \**********************************/ +# + + +# Import numpy package and name it "np" +import time +import numpy as np +import matplotlib.pyplot as plt +from matplotlib import cm +from os import makedirs +from os.path import join, exists + +from utils.ply import read_ply, write_ply +from utils.config import bcolors + + +# ------------------------------------------------------------------------------------------ +# +# Functions +# \***************/ +# +# + +def create_3D_rotations(axis, angle): + """ + Create rotation matrices from a list of axes and angles. Code from wikipedia on quaternions + :param axis: float32[N, 3] + :param angle: float32[N,] + :return: float32[N, 3, 3] + """ + + t1 = np.cos(angle) + t2 = 1 - t1 + t3 = axis[:, 0] * axis[:, 0] + t6 = t2 * axis[:, 0] + t7 = t6 * axis[:, 1] + t8 = np.sin(angle) + t9 = t8 * axis[:, 2] + t11 = t6 * axis[:, 2] + t12 = t8 * axis[:, 1] + t15 = axis[:, 1] * axis[:, 1] + t19 = t2 * axis[:, 1] * axis[:, 2] + t20 = t8 * axis[:, 0] + t24 = axis[:, 2] * axis[:, 2] + R = np.stack([t1 + t2 * t3, + t7 - t9, + t11 + t12, + t7 + t9, + t1 + t2 * t15, + t19 - t20, + t11 - t12, + t19 + t20, + t1 + t2 * t24], axis=1) + + return np.reshape(R, (-1, 3, 3)) + +def spherical_Lloyd(radius, num_cells, dimension=3, fixed='center', approximation='monte-carlo', + approx_n=5000, max_iter=500, momentum=0.9, verbose=0): + """ + Creation of kernel point via Lloyd algorithm. We use an approximation of the algorithm, and compute the Voronoi + cell centers with discretization of space. The exact formula is not trivial with part of the sphere as sides. + :param radius: Radius of the kernels + :param num_cells: Number of cell (kernel points) in the Voronoi diagram. + :param dimension: dimension of the space + :param fixed: fix position of certain kernel points ('none', 'center' or 'verticals') + :param approximation: Approximation method for Lloyd's algorithm ('discretization', 'monte-carlo') + :param approx_n: Number of point used for approximation. + :param max_iter: Maximum nu;ber of iteration for the algorithm. + :param momentum: Momentum of the low pass filter smoothing kernel point positions + :param verbose: display option + :return: points [num_kernels, num_points, dimension] + """ + + ####################### + # Parameters definition + ####################### + + # Radius used for optimization (points are rescaled afterwards) + radius0 = 1.0 + + ####################### + # Kernel initialization + ####################### + + # Random kernel points (Uniform distribution in a sphere) + kernel_points = np.zeros((0, dimension)) + while kernel_points.shape[0] < num_cells: + new_points = np.random.rand(num_cells, dimension) * 2 * radius0 - radius0 + kernel_points = np.vstack((kernel_points, new_points)) + d2 = np.sum(np.power(kernel_points, 2), axis=1) + kernel_points = kernel_points[np.logical_and(d2 < radius0 ** 2, (0.9 * radius0) ** 2 < d2), :] + kernel_points = kernel_points[:num_cells, :].reshape((num_cells, -1)) + + # Optional fixing + if fixed == 'center': + kernel_points[0, :] *= 0 + if fixed == 'verticals': + kernel_points[:3, :] *= 0 + kernel_points[1, -1] += 2 * radius0 / 3 + kernel_points[2, -1] -= 2 * radius0 / 3 + + ############################## + # Approximation initialization + ############################## + + # Initialize figure + if verbose > 1: + fig = plt.figure() + + # Initialize discretization in this method is chosen + if approximation == 'discretization': + side_n = int(np.floor(approx_n ** (1. / dimension))) + dl = 2 * radius0 / side_n + coords = np.arange(-radius0 + dl/2, radius0, dl) + if dimension == 2: + x, y = np.meshgrid(coords, coords) + X = np.vstack((np.ravel(x), np.ravel(y))).T + elif dimension == 3: + x, y, z = np.meshgrid(coords, coords, coords) + X = np.vstack((np.ravel(x), np.ravel(y), np.ravel(z))).T + elif dimension == 4: + x, y, z, t = np.meshgrid(coords, coords, coords, coords) + X = np.vstack((np.ravel(x), np.ravel(y), np.ravel(z), np.ravel(t))).T + else: + raise ValueError('Unsupported dimension (max is 4)') + elif approximation == 'monte-carlo': + X = np.zeros((0, dimension)) + else: + raise ValueError('Wrong approximation method chosen: "{:s}"'.format(approximation)) + + # Only points inside the sphere are used + d2 = np.sum(np.power(X, 2), axis=1) + X = X[d2 < radius0 * radius0, :] + + ##################### + # Kernel optimization + ##################### + + # Warning if at least one kernel point has no cell + warning = False + + # moving vectors of kernel points saved to detect convergence + max_moves = np.zeros((0,)) + + for iter in range(max_iter): + + # In the case of monte-carlo, renew the sampled points + if approximation == 'monte-carlo': + X = np.random.rand(approx_n, dimension) * 2 * radius0 - radius0 + d2 = np.sum(np.power(X, 2), axis=1) + X = X[d2 < radius0 * radius0, :] + + # Get the distances matrix [n_approx, K, dim] + differences = np.expand_dims(X, 1) - kernel_points + sq_distances = np.sum(np.square(differences), axis=2) + + # Compute cell centers + cell_inds = np.argmin(sq_distances, axis=1) + centers = [] + for c in range(num_cells): + bool_c = (cell_inds == c) + num_c = np.sum(bool_c.astype(np.int32)) + if num_c > 0: + centers.append(np.sum(X[bool_c, :], axis=0) / num_c) + else: + warning = True + centers.append(kernel_points[c]) + + # Update kernel points with low pass filter to smooth mote carlo + centers = np.vstack(centers) + moves = (1 - momentum) * (centers - kernel_points) + kernel_points += moves + + # Check moves for convergence + max_moves = np.append(max_moves, np.max(np.linalg.norm(moves, axis=1))) + + # Optional fixing + if fixed == 'center': + kernel_points[0, :] *= 0 + if fixed == 'verticals': + kernel_points[0, :] *= 0 + kernel_points[:3, :-1] *= 0 + + if verbose: + print('iter {:5d} / max move = {:f}'.format(iter, np.max(np.linalg.norm(moves, axis=1)))) + if warning: + print('{:}WARNING: at least one point has no cell{:}'.format(bcolors.WARNING, bcolors.ENDC)) + if verbose > 1: + plt.clf() + plt.scatter(X[:, 0], X[:, 1], c=cell_inds, s=20.0, + marker='.', cmap=plt.get_cmap('tab20')) + #plt.scatter(kernel_points[:, 0], kernel_points[:, 1], c=np.arange(num_cells), s=100.0, + # marker='+', cmap=plt.get_cmap('tab20')) + plt.plot(kernel_points[:, 0], kernel_points[:, 1], 'k+') + circle = plt.Circle((0, 0), radius0, color='r', fill=False) + fig.axes[0].add_artist(circle) + fig.axes[0].set_xlim((-radius0 * 1.1, radius0 * 1.1)) + fig.axes[0].set_ylim((-radius0 * 1.1, radius0 * 1.1)) + fig.axes[0].set_aspect('equal') + plt.draw() + plt.pause(0.001) + plt.show(block=False) + + ################### + # User verification + ################### + + # Show the convergence to ask user if this kernel is correct + if verbose: + if dimension == 2: + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=[10.4, 4.8]) + ax1.plot(max_moves) + ax2.scatter(X[:, 0], X[:, 1], c=cell_inds, s=20.0, + marker='.', cmap=plt.get_cmap('tab20')) + # plt.scatter(kernel_points[:, 0], kernel_points[:, 1], c=np.arange(num_cells), s=100.0, + # marker='+', cmap=plt.get_cmap('tab20')) + ax2.plot(kernel_points[:, 0], kernel_points[:, 1], 'k+') + circle = plt.Circle((0, 0), radius0, color='r', fill=False) + ax2.add_artist(circle) + ax2.set_xlim((-radius0 * 1.1, radius0 * 1.1)) + ax2.set_ylim((-radius0 * 1.1, radius0 * 1.1)) + ax2.set_aspect('equal') + plt.title('Check if kernel is correct.') + plt.draw() + plt.show() + + if dimension > 2: + plt.figure() + plt.plot(max_moves) + plt.title('Check if kernel is correct.') + plt.show() + + # Rescale kernels with real radius + return kernel_points * radius + + +def kernel_point_optimization_debug(radius, num_points, num_kernels=1, dimension=3, + fixed='center', ratio=0.66, verbose=0): + """ + Creation of kernel point via optimization of potentials. + :param radius: Radius of the kernels + :param num_points: points composing kernels + :param num_kernels: number of wanted kernels + :param dimension: dimension of the space + :param fixed: fix position of certain kernel points ('none', 'center' or 'verticals') + :param ratio: ratio of the radius where you want the kernels points to be placed + :param verbose: display option + :return: points [num_kernels, num_points, dimension] + """ + + ####################### + # Parameters definition + ####################### + + # Radius used for optimization (points are rescaled afterwards) + radius0 = 1 + diameter0 = 2 + + # Factor multiplicating gradients for moving points (~learning rate) + moving_factor = 1e-2 + continuous_moving_decay = 0.9995 + + # Gradient threshold to stop optimization + thresh = 1e-5 + + # Gradient clipping value + clip = 0.05 * radius0 + + ####################### + # Kernel initialization + ####################### + + # Random kernel points + kernel_points = np.random.rand(num_kernels * num_points - 1, dimension) * diameter0 - radius0 + while (kernel_points.shape[0] < num_kernels * num_points): + new_points = np.random.rand(num_kernels * num_points - 1, dimension) * diameter0 - radius0 + kernel_points = np.vstack((kernel_points, new_points)) + d2 = np.sum(np.power(kernel_points, 2), axis=1) + kernel_points = kernel_points[d2 < 0.5 * radius0 * radius0, :] + kernel_points = kernel_points[:num_kernels * num_points, :].reshape((num_kernels, num_points, -1)) + + # Optionnal fixing + if fixed == 'center': + kernel_points[:, 0, :] *= 0 + if fixed == 'verticals': + kernel_points[:, :3, :] *= 0 + kernel_points[:, 1, -1] += 2 * radius0 / 3 + kernel_points[:, 2, -1] -= 2 * radius0 / 3 + + ##################### + # Kernel optimization + ##################### + + # Initialize figure + if verbose>1: + fig = plt.figure() + + saved_gradient_norms = np.zeros((10000, num_kernels)) + old_gradient_norms = np.zeros((num_kernels, num_points)) + for iter in range(10000): + + # Compute gradients + # ***************** + + # Derivative of the sum of potentials of all points + A = np.expand_dims(kernel_points, axis=2) + B = np.expand_dims(kernel_points, axis=1) + interd2 = np.sum(np.power(A - B, 2), axis=-1) + inter_grads = (A - B) / (np.power(np.expand_dims(interd2, -1), 3/2) + 1e-6) + inter_grads = np.sum(inter_grads, axis=1) + + # Derivative of the radius potential + circle_grads = 10*kernel_points + + # All gradients + gradients = inter_grads + circle_grads + + if fixed == 'verticals': + gradients[:, 1:3, :-1] = 0 + + # Stop condition + # ************** + + # Compute norm of gradients + gradients_norms = np.sqrt(np.sum(np.power(gradients, 2), axis=-1)) + saved_gradient_norms[iter, :] = np.max(gradients_norms, axis=1) + + # Stop if all moving points are gradients fixed (low gradients diff) + + if fixed == 'center' and np.max(np.abs(old_gradient_norms[:, 1:] - gradients_norms[:, 1:])) < thresh: + break + elif fixed == 'verticals' and np.max(np.abs(old_gradient_norms[:, 3:] - gradients_norms[:, 3:])) < thresh: + break + elif np.max(np.abs(old_gradient_norms - gradients_norms)) < thresh: + break + old_gradient_norms = gradients_norms + + # Move points + # *********** + + # Clip gradient to get moving dists + moving_dists = np.minimum(moving_factor * gradients_norms, clip) + + # Fix central point + if fixed == 'center': + moving_dists[:, 0] = 0 + if fixed == 'verticals': + moving_dists[:, 0] = 0 + + # Move points + kernel_points -= np.expand_dims(moving_dists, -1) * gradients / np.expand_dims(gradients_norms + 1e-6, -1) + + if verbose: + print('iter {:5d} / max grad = {:f}'.format(iter, np.max(gradients_norms[:, 3:]))) + if verbose > 1: + plt.clf() + plt.plot(kernel_points[0, :, 0], kernel_points[0, :, 1], '.') + circle = plt.Circle((0, 0), radius, color='r', fill=False) + fig.axes[0].add_artist(circle) + fig.axes[0].set_xlim((-radius*1.1, radius*1.1)) + fig.axes[0].set_ylim((-radius*1.1, radius*1.1)) + fig.axes[0].set_aspect('equal') + plt.draw() + plt.pause(0.001) + plt.show(block=False) + print(moving_factor) + + # moving factor decay + moving_factor *= continuous_moving_decay + + # Rescale radius to fit the wanted ratio of radius + r = np.sqrt(np.sum(np.power(kernel_points, 2), axis=-1)) + kernel_points *= ratio / np.mean(r[:, 1:]) + + # Rescale kernels with real radius + return kernel_points * radius, saved_gradient_norms + + +def load_kernels(radius, num_kpoints, dimension, fixed, lloyd=False): + + # Kernel directory + kernel_dir = 'kernels/dispositions' + if not exists(kernel_dir): + makedirs(kernel_dir) + + # To many points switch to Lloyds + if num_kpoints > 30: + lloyd = True + + # Kernel_file + kernel_file = join(kernel_dir, 'k_{:03d}_{:s}_{:d}D.ply'.format(num_kpoints, fixed, dimension)) + + # Check if already done + if not exists(kernel_file): + if lloyd: + # Create kernels + kernel_points = spherical_Lloyd(1.0, + num_kpoints, + dimension=dimension, + fixed=fixed, + verbose=0) + + else: + # Create kernels + kernel_points, grad_norms = kernel_point_optimization_debug(1.0, + num_kpoints, + num_kernels=100, + dimension=dimension, + fixed=fixed, + verbose=0) + + # Find best candidate + best_k = np.argmin(grad_norms[-1, :]) + + # Save points + kernel_points = kernel_points[best_k, :, :] + + write_ply(kernel_file, kernel_points, ['x', 'y', 'z']) + + else: + data = read_ply(kernel_file) + kernel_points = np.vstack((data['x'], data['y'], data['z'])).T + + # Random roations for the kernel + # N.B. 4D random rotations not supported yet + R = np.eye(dimension) + theta = np.random.rand() * 2 * np.pi + if dimension == 2: + if fixed != 'vertical': + c, s = np.cos(theta), np.sin(theta) + R = np.array([[c, -s], [s, c]], dtype=np.float32) + + elif dimension == 3: + if fixed != 'vertical': + c, s = np.cos(theta), np.sin(theta) + R = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]], dtype=np.float32) + + else: + phi = (np.random.rand() - 0.5) * np.pi + + # Create the first vector in carthesian coordinates + u = np.array([np.cos(theta) * np.cos(phi), np.sin(theta) * np.cos(phi), np.sin(phi)]) + + # Choose a random rotation angle + alpha = np.random.rand() * 2 * np.pi + + # Create the rotation matrix with this vector and angle + R = create_3D_rotations(np.reshape(u, (1, -1)), np.reshape(alpha, (1, -1)))[0] + + R = R.astype(np.float32) + + # Add a small noise + kernel_points = kernel_points + np.random.normal(scale=0.01, size=kernel_points.shape) + + # Scale kernels + kernel_points = radius * kernel_points + + # Rotate kernels + kernel_points = np.matmul(kernel_points, R) + + return kernel_points.astype(np.float32) \ No newline at end of file diff --git a/models/architectures.py b/models/architectures.py new file mode 100644 index 0000000..941af3f --- /dev/null +++ b/models/architectures.py @@ -0,0 +1,451 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Define network architectures +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 06/03/2020 +# + + +from models.blocks import * +import numpy as np + +class KPCNN(nn.Module): + """ + Class defining KPCNN + """ + + def __init__(self, config): + super(KPCNN, self).__init__() + + ##################### + # Network opperations + ##################### + + # Current radius of convolution and feature dimension + layer = 0 + r = config.first_subsampling_dl * config.conv_radius + in_dim = config.in_features_dim + out_dim = config.first_features_dim + self.K = config.num_kernel_points + + # Save all block operations in a list of modules + self.block_ops = nn.ModuleList() + + # Loop over consecutive blocks + block_in_layer = 0 + for block_i, block in enumerate(config.architecture): + + # Check equivariance + if ('equivariant' in block) and (not out_dim % 3 == 0): + raise ValueError('Equivariant block but features dimension is not a factor of 3') + + # Detect upsampling block to stop + if 'upsample' in block: + break + + # Apply the good block function defining tf ops + self.block_ops.append(block_decider(block, + r, + in_dim, + out_dim, + layer, + config)) + + # Index of block in this layer + block_in_layer += 1 + + # Update dimension of input from output + in_dim = out_dim + + # Detect change to a subsampled layer + if 'pool' in block or 'strided' in block: + # Update radius and feature dimension for next layer + layer += 1 + r *= 2 + out_dim *= 2 + block_in_layer = 0 + + self.head_mlp = UnaryBlock(out_dim, 1024, False, 0) + self.head_softmax = UnaryBlock(1024, config.num_classes, False, 0) + + ################ + # Network Losses + ################ + + self.criterion = torch.nn.CrossEntropyLoss() + self.offset_loss = config.offsets_loss + self.offset_decay = config.offsets_decay + self.output_loss = 0 + self.reg_loss = 0 + self.l1 = nn.L1Loss() + + return + + def forward(self, batch, config): + + # Save all block operations in a list of modules + x = batch.features.clone().detach() + + # Loop over consecutive blocks + for block_op in self.block_ops: + x = block_op(x, batch) + + # Head of network + x = self.head_mlp(x, batch) + x = self.head_softmax(x, batch) + + return x + + def loss(self, outputs, labels): + """ + Runs the loss on outputs of the model + :param outputs: logits + :param labels: labels + :return: loss + """ + + # TODO: Ignore unclassified points in loss for segmentation architecture + + # Cross entropy loss + self.output_loss = self.criterion(outputs, labels) + + # Regularization of deformable offsets + self.reg_loss = self.offset_regularizer() + + # Combined loss + return self.output_loss + self.reg_loss + + @staticmethod + def accuracy(outputs, labels): + """ + Computes accuracy of the current batch + :param outputs: logits predicted by the network + :param labels: labels + :return: accuracy value + """ + + predicted = torch.argmax(outputs.data, dim=1) + total = labels.size(0) + correct = (predicted == labels).sum().item() + + return correct / total + + def offset_regularizer(self): + + fitting_loss = 0 + repulsive_loss = 0 + + for m in self.modules(): + + if isinstance(m, KPConv) and m.deformable: + + ############################## + # divide offset gradient by 10 + ############################## + + m.unscaled_offsets.register_hook(lambda grad: grad * 0.1) + #m.unscaled_offsets.register_hook(lambda grad: print('GRAD2', grad[10, 5, :])) + + ############## + # Fitting loss + ############## + + # Get the distance to closest input point + KP_min_d2, _ = torch.min(m.deformed_d2, dim=1) + + # Normalize KP locations to be independant from layers + KP_min_d2 = KP_min_d2 / (m.KP_extent ** 2) + + # Loss will be the square distance to closest input point. We use L1 because dist is already squared + fitting_loss += self.l1(KP_min_d2, torch.zeros_like(KP_min_d2)) + + ################ + # Repulsive loss + ################ + + # Normalized KP locations + KP_locs = m.deformed_KP / m.KP_extent + + # Point should not be close to each other + for i in range(self.K): + + other_KP = torch.cat([KP_locs[:, :i, :], KP_locs[:, i + 1:, :]], dim=1).detach() + distances = torch.sqrt(torch.sum((other_KP - KP_locs[:, i:i + 1, :]) ** 2, dim=2)) + rep_loss = torch.sum(torch.clamp_max(distances - 1.5, max=0.0) ** 2, dim=1) + repulsive_loss += self.l1(rep_loss, torch.zeros_like(rep_loss)) + + + + return self.offset_decay * (fitting_loss + repulsive_loss) + + + + + +class KPFCNN(nn.Module): + """ + Class defining KPFCNN + """ + + def __init__(self, config): + super(KPFCNN, self).__init__() + + ############ + # Parameters + ############ + + # Current radius of convolution and feature dimension + layer = 0 + r = config.first_subsampling_dl * config.conv_radius + in_dim = config.in_features_dim + out_dim = config.first_features_dim + self.K = config.num_kernel_points + + ##################### + # List Encoder blocks + ##################### + + # Save all block operations in a list of modules + self.encoder_blocs = nn.ModuleList() + self.encoder_skip_dims = [] + self.encoder_skips = [] + + # Loop over consecutive blocks + for block_i, block in enumerate(config.architecture): + + # Check equivariance + if ('equivariant' in block) and (not out_dim % 3 == 0): + raise ValueError('Equivariant block but features dimension is not a factor of 3') + + # Detect change to next layer for skip connection + if np.any([tmp in block for tmp in ['pool', 'strided', 'upsample', 'global']]): + self.encoder_skips.append(block_i) + self.encoder_skip_dims.append(in_dim) + + # Detect upsampling block to stop + if 'upsample' in block: + break + + # Apply the good block function defining tf ops + self.encoder_blocs.append(block_decider(block, + r, + in_dim, + out_dim, + layer, + config)) + + # Update dimension of input from output + in_dim = out_dim + + # Detect change to a subsampled layer + if 'pool' in block or 'strided' in block: + # Update radius and feature dimension for next layer + layer += 1 + r *= 2 + out_dim *= 2 + + ##################### + # List Decoder blocks + ##################### + + # Save all block operations in a list of modules + self.decoder_blocs = nn.ModuleList() + self.decoder_concats = [] + + # Find first upsampling block + start_i = 0 + for block_i, block in enumerate(config.architecture): + if 'upsample' in block: + start_i = block_i + break + + # Loop over consecutive blocks + for block_i, block in enumerate(config.architecture[start_i:]): + + # Add dimension of skip connection concat + if block_i > 0 and 'upsample' in config.architecture[start_i + block_i - 1]: + in_dim += self.encoder_skip_dims[layer] + self.decoder_concats.append(block_i) + + # Apply the good block function defining tf ops + self.decoder_blocs.append(block_decider(block, + r, + in_dim, + out_dim, + layer, + config)) + + # Update dimension of input from output + in_dim = out_dim + + # Detect change to a subsampled layer + if 'upsample' in block: + # Update radius and feature dimension for next layer + layer -= 1 + r *= 0.5 + out_dim = out_dim // 2 + + self.head_mlp = UnaryBlock(out_dim, config.first_features_dim, False, 0) + self.head_softmax = UnaryBlock(config.first_features_dim, config.num_classes, False, 0) + + ################ + # Network Losses + ################ + + # Choose segmentation loss + if config.segloss_balance == 'none': + self.criterion = torch.nn.CrossEntropyLoss() + elif config.segloss_balance == 'class': + self.criterion = torch.nn.CrossEntropyLoss() + elif config.segloss_balance == 'batch': + self.criterion = torch.nn.CrossEntropyLoss() + else: + raise ValueError('Unknown segloss_balance:', config.segloss_balance) + self.offset_loss = config.offsets_loss + self.offset_decay = config.offsets_decay + self.output_loss = 0 + self.reg_loss = 0 + self.l1 = nn.L1Loss() + + return + + def forward(self, batch, config): + + # Get input features + x = batch.features.clone().detach() + + # Loop over consecutive blocks + skip_x = [] + for block_i, block_op in enumerate(self.encoder_blocs): + if block_i in self.encoder_skips: + skip_x.append(x) + x = block_op(x, batch) + + for block_i, block_op in enumerate(self.decoder_blocs): + if block_i in self.decoder_concats: + x = torch.cat([x, skip_x.pop()], dim=1) + x = block_op(x, batch) + + # Head of network + x = self.head_mlp(x, batch) + x = self.head_softmax(x, batch) + + return x + + def loss(self, outputs, labels): + """ + Runs the loss on outputs of the model + :param outputs: logits + :param labels: labels + :return: loss + """ + + outputs = torch.transpose(outputs, 0, 1) + outputs = outputs.unsqueeze(0) + labels = labels.unsqueeze(0) + + # Cross entropy loss + self.output_loss = self.criterion(outputs, labels) + + # Regularization of deformable offsets + self.reg_loss = self.offset_regularizer() + + # Combined loss + return self.output_loss + self.reg_loss + + @staticmethod + def accuracy(outputs, labels): + """ + Computes accuracy of the current batch + :param outputs: logits predicted by the network + :param labels: labels + :return: accuracy value + """ + + predicted = torch.argmax(outputs.data, dim=1) + total = labels.size(0) + correct = (predicted == labels).sum().item() + + return correct / total + + def offset_regularizer(self): + + fitting_loss = 0 + repulsive_loss = 0 + + for m in self.modules(): + + if isinstance(m, KPConv) and m.deformable: + + ############################## + # divide offset gradient by 10 + ############################## + + m.unscaled_offsets.register_hook(lambda grad: grad * 0.1) + #m.unscaled_offsets.register_hook(lambda grad: print('GRAD2', grad[10, 5, :])) + + ############## + # Fitting loss + ############## + + # Get the distance to closest input point + KP_min_d2, _ = torch.min(m.deformed_d2, dim=1) + + # Normalize KP locations to be independant from layers + KP_min_d2 = KP_min_d2 / (m.KP_extent ** 2) + + # Loss will be the square distance to closest input point. We use L1 because dist is already squared + fitting_loss += self.l1(KP_min_d2, torch.zeros_like(KP_min_d2)) + + ################ + # Repulsive loss + ################ + + # Normalized KP locations + KP_locs = m.deformed_KP / m.KP_extent + + # Point should not be close to each other + for i in range(self.K): + + other_KP = torch.cat([KP_locs[:, :i, :], KP_locs[:, i + 1:, :]], dim=1).detach() + distances = torch.sqrt(torch.sum((other_KP - KP_locs[:, i:i + 1, :]) ** 2, dim=2)) + rep_loss = torch.sum(torch.clamp_max(distances - 1.5, max=0.0) ** 2, dim=1) + repulsive_loss += self.l1(rep_loss, torch.zeros_like(rep_loss)) + + + + return self.offset_decay * (fitting_loss + repulsive_loss) + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/models/blocks.py b/models/blocks.py new file mode 100644 index 0000000..c8332e3 --- /dev/null +++ b/models/blocks.py @@ -0,0 +1,651 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Define network blocks +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 06/03/2020 +# + + +import time +import math +import torch +import torch.nn as nn +from torch.nn.parameter import Parameter +from torch.nn.init import kaiming_uniform_ +from kernels.kernel_points import load_kernels + +from utils.ply import write_ply + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Simple functions +# \**********************/ +# + + +def gather(x, idx, method=2): + """ + implementation of a custom gather operation for faster backwards. + :param x: input with shape [N, D_1, ... D_d] + :param idx: indexing with shape [n_1, ..., n_m] + :param method: Choice of the method + :return: x[idx] with shape [n_1, ..., n_m, D_1, ... D_d] + """ + + if method == 0: + return x[idx] + elif method == 1: + x = x.unsqueeze(1) + x = x.expand((-1, idx.shape[-1], -1)) + idx = idx.unsqueeze(2) + idx = idx.expand((-1, -1, x.shape[-1])) + return x.gather(0, idx) + elif method == 2: + for i, ni in enumerate(idx.size()[1:]): + x = x.unsqueeze(i+1) + new_s = list(x.size()) + new_s[i+1] = ni + x = x.expand(new_s) + n = len(idx.size()) + for i, di in enumerate(x.size()[n:]): + idx = idx.unsqueeze(i+n) + new_s = list(idx.size()) + new_s[i+n] = di + idx = idx.expand(new_s) + return x.gather(0, idx) + else: + raise ValueError('Unkown method') + + +def radius_gaussian(sq_r, sig, eps=1e-9): + """ + Compute a radius gaussian (gaussian of distance) + :param sq_r: input radiuses [dn, ..., d1, d0] + :param sig: extents of gaussians [d1, d0] or [d0] or float + :return: gaussian of sq_r [dn, ..., d1, d0] + """ + return torch.exp(-sq_r / (2 * sig**2 + eps)) + + +def closest_pool(x, inds): + """ + Pools features from the closest neighbors. WARNING: this function assumes the neighbors are ordered. + :param x: [n1, d] features matrix + :param inds: [n2, max_num] Only the first column is used for pooling + :return: [n2, d] pooled features matrix + """ + + # Add a last row with minimum features for shadow pools + x = torch.cat((x, torch.zeros_like(x[:1, :])), 0) + + # Get features for each pooling location [n2, d] + return gather(x, inds[:, 0]) + + +def max_pool(x, inds): + """ + Pools features with the maximum values. + :param x: [n1, d] features matrix + :param inds: [n2, max_num] pooling indices + :return: [n2, d] pooled features matrix + """ + + # Add a last row with minimum features for shadow pools + x = torch.cat((x, torch.zeros_like(x[:1, :])), 0) + + # Get all features for each pooling location [n2, max_num, d] + pool_features = gather(x, inds) + + # Pool the maximum [n2, d] + max_features, _ = torch.max(pool_features, 1) + return max_features + + +def global_average(x, batch_lengths): + """ + Block performing a global average over batch pooling + :param x: [N, D] input features + :param batch_lengths: [B] list of batch lengths + :return: [B, D] averaged features + """ + + # Loop over the clouds of the batch + averaged_features = [] + i0 = 0 + for b_i, length in enumerate(batch_lengths): + + # Average features for each batch cloud + averaged_features.append(torch.mean(x[i0:i0 + length], dim=0)) + + # Increment for next cloud + i0 += length + + # Average features in each batch + return torch.stack(averaged_features) + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# KPConv class +# \******************/ +# + + +class KPConv(nn.Module): + + def __init__(self, kernel_size, p_dim, in_channels, out_channels, KP_extent, radius, + fixed_kernel_points='center', KP_influence='linear', aggregation_mode='sum', + deformable=False, modulated=False): + """ + Initialize parameters for KPConvDeformable. + :param kernel_size: Number of kernel points. + :param p_dim: dimension of the point space. + :param in_channels: dimension of input features. + :param out_channels: dimension of output features. + :param KP_extent: influence radius of each kernel point. + :param radius: radius used for kernel point init. Even for deformable, use the config.conv_radius + :param fixed_kernel_points: fix position of certain kernel points ('none', 'center' or 'verticals'). + :param KP_influence: influence function of the kernel points ('constant', 'linear', 'gaussian'). + :param aggregation_mode: choose to sum influences, or only keep the closest ('closest', 'sum'). + :param deformable: choose deformable or not + :param modulated: choose if kernel weights are modulated in addition to deformed + """ + super(KPConv, self).__init__() + + # Save parameters + self.K = kernel_size + self.p_dim = p_dim + self.in_channels = in_channels + self.out_channels = out_channels + self.radius = radius + self.KP_extent = KP_extent + self.fixed_kernel_points = fixed_kernel_points + self.KP_influence = KP_influence + self.aggregation_mode = aggregation_mode + self.deformable = deformable + self.modulated = modulated + + # Running variable containing deformed KP distance to input points. (used in regularization loss) + self.deformed_d2 = None + self.deformed_KP = None + self.unscaled_offsets = None + + # Initialize weights + self.weights = Parameter(torch.zeros((self.K, in_channels, out_channels), dtype=torch.float32), + requires_grad=True) + + # Initiate weights for offsets + if deformable: + if modulated: + self.offset_dim = (self.p_dim + 1) * self.K + else: + self.offset_dim = self.p_dim * self.K + self.offset_conv = KPConv(self.K, + self.p_dim, + in_channels, + self.offset_dim, + KP_extent, + radius, + fixed_kernel_points=fixed_kernel_points, + KP_influence=KP_influence, + aggregation_mode=aggregation_mode) + self.offset_bias = Parameter(torch.zeros(self.offset_dim, dtype=torch.float32), requires_grad=True) + + else: + self.offset_dim = None + self.offset_conv = None + self.offset_bias = None + + # Reset parameters + self.reset_parameters() + + # Initialize kernel points + self.kernel_points = self.init_KP() + + return + + def reset_parameters(self): + kaiming_uniform_(self.weights, a=math.sqrt(5)) + if self.deformable: + nn.init.zeros_(self.offset_bias) + return + + def init_KP(self): + """ + Initialize the kernel point positions in a sphere + :return: the tensor of kernel points + """ + + # Create one kernel disposition (as numpy array). Choose the KP distance to center thanks to the KP extent + K_points_numpy = load_kernels(self.radius, + self.K, + dimension=self.p_dim, + fixed=self.fixed_kernel_points) + + return Parameter(torch.tensor(K_points_numpy, dtype=torch.float32), + requires_grad=False) + + def forward(self, q_pts, s_pts, neighb_inds, x): + + ################### + # Offset generation + ################### + + if self.deformable: + offset_features = self.offset_conv(q_pts, s_pts, neighb_inds, x) + self.offset_bias + + if self.modulated: + + # Get offset (in normalized scale) from features + offsets = offset_features[:, :self.p_dim * self.K] + self.unscaled_offsets = offsets.view(-1, self.K, self.p_dim) + + # Get modulations + modulations = 2 * torch.sigmoid(offset_features[:, self.p_dim * self.K:]) + + else: + + # Get offset (in normalized scale) from features + self.unscaled_offsets = offset_features.view(-1, self.K, self.p_dim) + + # No modulations + modulations = None + + # Rescale offset for this layer + offsets = self.unscaled_offsets * self.KP_extent + + else: + offsets = None + modulations = None + + ###################### + # Deformed convolution + ###################### + + # Add a fake point in the last row for shadow neighbors + s_pts = torch.cat((s_pts, torch.zeros_like(s_pts[:1, :]) + 1e6), 0) + + # Get neighbor points [n_points, n_neighbors, dim] + neighbors = s_pts[neighb_inds, :] + + # Center every neighborhood + neighbors = neighbors - q_pts.unsqueeze(1) + + # Apply offsets to kernel points [n_points, n_kpoints, dim] + if self.deformable: + self.deformed_KP = offsets + self.kernel_points + deformed_K_points = self.deformed_KP.unsqueeze(1) + else: + deformed_K_points = self.kernel_points + + # Get all difference matrices [n_points, n_neighbors, n_kpoints, dim] + neighbors.unsqueeze_(2) + differences = neighbors - deformed_K_points + + # Get the square distances [n_points, n_neighbors, n_kpoints] + sq_distances = torch.sum(differences ** 2, dim=3) + + # Optimization by ignoring points outside a deformed KP range + if False and self.deformable: + # Boolean of the neighbors in range of a kernel point [n_points, n_neighbors] + in_range = torch.any(sq_distances < self.KP_extent ** 2, dim=2) + + # New value of max neighbors + new_max_neighb = torch.max(torch.sum(in_range, dim=1)) + + print(sq_distances.shape[1], '=>', new_max_neighb.item()) + + # Save distances for loss + if self.deformable: + self.deformed_d2 = sq_distances + + # Get Kernel point influences [n_points, n_kpoints, n_neighbors] + if self.KP_influence == 'constant': + # Every point get an influence of 1. + all_weights = torch.ones_like(sq_distances) + all_weights = torch.transpose(all_weights, 1, 2) + + elif self.KP_influence == 'linear': + # Influence decrease linearly with the distance, and get to zero when d = KP_extent. + all_weights = torch.clamp(1 - torch.sqrt(sq_distances) / self.KP_extent, min=0.0) + all_weights = torch.transpose(all_weights, 1, 2) + + elif self.KP_influence == 'gaussian': + # Influence in gaussian of the distance. + sigma = self.KP_extent * 0.3 + all_weights = radius_gaussian(sq_distances, sigma) + all_weights = torch.transpose(all_weights, 1, 2) + else: + raise ValueError('Unknown influence function type (config.KP_influence)') + + # In case of closest mode, only the closest KP can influence each point + if self.aggregation_mode == 'closest': + neighbors_1nn = torch.argmin(sq_distances, dim=2) + all_weights *= torch.transpose(nn.functional.one_hot(neighbors_1nn, self.K), 1, 2) + + elif self.aggregation_mode != 'sum': + raise ValueError("Unknown convolution mode. Should be 'closest' or 'sum'") + + # Add a zero feature for shadow neighbors + x = torch.cat((x, torch.zeros_like(x[:1, :])), 0) + + # Get the features of each neighborhood [n_points, n_neighbors, in_fdim] + neighb_x = gather(x, neighb_inds) + + # Apply distance weights [n_points, n_kpoints, in_fdim] + weighted_features = torch.matmul(all_weights, neighb_x) + + # Apply modulations + if self.deformable and self.modulated: + weighted_features *= modulations.unsqueeze(2) + + # Apply network weights [n_kpoints, n_points, out_fdim] + weighted_features = weighted_features.permute((1, 0, 2)) + kernel_outputs = torch.matmul(weighted_features, self.weights) + + # Convolution sum [n_points, out_fdim] + return torch.sum(kernel_outputs, dim=0) + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Complex blocks +# \********************/ +# + +def block_decider(block_name, + radius, + in_dim, + out_dim, + layer_ind, + config): + + if block_name == 'unary': + return UnaryBlock(in_dim, out_dim, config.use_batch_norm, config.batch_norm_momentum) + + elif block_name in ['simple', + 'simple_deformable', + 'simple_invariant', + 'simple_equivariant', + 'simple_strided', + 'simple_deformable_strided', + 'simple_invariant_strided', + 'simple_equivariant_strided']: + return SimpleBlock(block_name, in_dim, out_dim, radius, layer_ind, config) + + elif block_name in ['resnetb', + 'resnetb_invariant', + 'resnetb_equivariant', + 'resnetb_deformable', + 'resnetb_strided', + 'resnetb_deformable_strided', + 'resnetb_equivariant_strided', + 'resnetb_invariant_strided']: + return ResnetBottleneckBlock(block_name, in_dim, out_dim, radius, layer_ind, config) + + elif block_name == 'max_pool' or block_name == 'max_pool_wide': + return MaxPoolBlock(layer_ind) + + elif block_name == 'global_average': + return GlobalAverageBlock() + + elif block_name == 'nearest_upsample': + return NearestUpsampleBlock(layer_ind) + + else: + raise ValueError('Unknown block name in the architecture definition : ' + block_name) + + +class BatchNormBlock(nn.Module): + + def __init__(self, in_dim, use_bn, bn_momentum): + """ + Initialize a batch normalization block. If network does not use batch normalization, replace with biases. + :param in_dim: dimension input features + :param use_bn: boolean indicating if we use Batch Norm + :param bn_momentum: Batch norm momentum + """ + super(BatchNormBlock, self).__init__() + self.bn_momentum = bn_momentum + self.use_bn = use_bn + if self.use_bn: + self.batch_norm = nn.BatchNorm1d(in_dim, momentum=bn_momentum) + #self.batch_norm = nn.InstanceNorm1d(in_dim, momentum=bn_momentum) + else: + self.bias = Parameter(torch.zeros(in_dim, dtype=torch.float32), requires_grad=True) + return + + def reset_parameters(self): + nn.init.zeros_(self.bias) + + def forward(self, x): + if self.use_bn: + + x = x.unsqueeze(2) + x = x.transpose(0, 2) + x = self.batch_norm(x) + x = x.transpose(0, 2) + return x.squeeze() + else: + return x + self.bias + + +class UnaryBlock(nn.Module): + + def __init__(self, in_dim, out_dim, use_bn, bn_momentum, no_relu=False): + """ + Initialize a standard unary block with its ReLU and BatchNorm. + :param in_dim: dimension input features + :param out_dim: dimension input features + :param use_bn: boolean indicating if we use Batch Norm + :param bn_momentum: Batch norm momentum + """ + + super(UnaryBlock, self).__init__() + self.bn_momentum = bn_momentum + self.use_bn = use_bn + self.no_relu = no_relu + self.mlp = nn.Linear(in_dim, out_dim, bias=False) + self.batch_norm = BatchNormBlock(out_dim, self.use_bn, self.bn_momentum) + if not no_relu: + self.leaky_relu = nn.LeakyReLU(0.1) + return + + def forward(self, x, batch=None): + x = self.mlp(x) + x = self.batch_norm(x) + if not self.no_relu: + x = self.leaky_relu(x) + return x + + +class SimpleBlock(nn.Module): + + def __init__(self, block_name, in_dim, out_dim, radius, layer_ind, config): + """ + Initialize a simple convolution block with its ReLU and BatchNorm. + :param in_dim: dimension input features + :param out_dim: dimension input features + :param radius: current radius of convolution + :param config: parameters + """ + super(SimpleBlock, self).__init__() + + # get KP_extent from current radius + current_extent = radius * config.KP_extent / config.conv_radius + + # Get other parameters + self.bn_momentum = config.batch_norm_momentum + self.use_bn = config.use_batch_norm + self.layer_ind = layer_ind + self.block_name = block_name + + # Define the KPConv class + self.KPConv = KPConv(config.num_kernel_points, + config.in_points_dim, + in_dim, + out_dim, + current_extent, + radius, + fixed_kernel_points=config.fixed_kernel_points, + KP_influence=config.KP_influence, + aggregation_mode=config.aggregation_mode, + deformable='deform' in block_name, + modulated=config.modulated) + + # Other opperations + self.batch_norm = BatchNormBlock(out_dim, self.use_bn, self.bn_momentum) + self.leaky_relu = nn.LeakyReLU(0.1) + + return + + def forward(self, x, batch): + + if 'strided' in self.block_name: + q_pts = batch.points[self.layer_ind + 1] + s_pts = batch.points[self.layer_ind] + neighb_inds = batch.pools[self.layer_ind] + else: + q_pts = batch.points[self.layer_ind] + s_pts = batch.points[self.layer_ind] + neighb_inds = batch.neighbors[self.layer_ind] + + x = self.KPConv(q_pts, s_pts, neighb_inds, x) + return self.leaky_relu(self.batch_norm(x)) + + +class ResnetBottleneckBlock(nn.Module): + + def __init__(self, block_name, in_dim, out_dim, radius, layer_ind, config): + """ + Initialize a resnet bottleneck block. + :param in_dim: dimension input features + :param out_dim: dimension input features + :param radius: current radius of convolution + :param config: parameters + """ + super(ResnetBottleneckBlock, self).__init__() + + # get KP_extent from current radius + current_extent = radius * config.KP_extent / config.conv_radius + + # Get other parameters + self.bn_momentum = config.batch_norm_momentum + self.use_bn = config.use_batch_norm + self.block_name = block_name + self.layer_ind = layer_ind + + # First downscaling mlp + if in_dim != out_dim // 2: + self.unary1 = UnaryBlock(in_dim, out_dim // 2, self.use_bn, self.bn_momentum) + else: + self.unary1 = nn.Identity() + + # KPConv block + self.KPConv = KPConv(config.num_kernel_points, + config.in_points_dim, + out_dim // 2, + out_dim // 2, + current_extent, + radius, + fixed_kernel_points=config.fixed_kernel_points, + KP_influence=config.KP_influence, + aggregation_mode=config.aggregation_mode, + deformable='deform' in block_name, + modulated=config.modulated) + self.batch_norm_conv = BatchNormBlock(out_dim // 2, self.use_bn, self.bn_momentum) + + # Second upscaling mlp + self.unary2 = UnaryBlock(out_dim // 2, out_dim, self.use_bn, self.bn_momentum, no_relu=True) + + # Shortcut optional mpl + if in_dim != out_dim: + self.unary_shortcut = UnaryBlock(in_dim, out_dim, self.use_bn, self.bn_momentum, no_relu=True) + else: + self.unary_shortcut = nn.Identity() + + # Other operations + self.leaky_relu = nn.LeakyReLU(0.1) + + return + + def forward(self, features, batch): + + if 'strided' in self.block_name: + q_pts = batch.points[self.layer_ind + 1] + s_pts = batch.points[self.layer_ind] + neighb_inds = batch.pools[self.layer_ind] + else: + q_pts = batch.points[self.layer_ind] + s_pts = batch.points[self.layer_ind] + neighb_inds = batch.neighbors[self.layer_ind] + + # First downscaling mlp + x = self.unary1(features) + + # Convolution + x = self.KPConv(q_pts, s_pts, neighb_inds, x) + x = self.leaky_relu(self.batch_norm_conv(x)) + + # Second upscaling mlp + x = self.unary2(x) + + # Shortcut + if 'strided' in self.block_name: + shortcut = max_pool(features, neighb_inds) + else: + shortcut = features + shortcut = self.unary_shortcut(shortcut) + + return self.leaky_relu(x + shortcut) + + +class GlobalAverageBlock(nn.Module): + + def __init__(self): + """ + Initialize a global average block with its ReLU and BatchNorm. + """ + super(GlobalAverageBlock, self).__init__() + return + + def forward(self, x, batch): + return global_average(x, batch.lengths[-1]) + + +class NearestUpsampleBlock(nn.Module): + + def __init__(self, layer_ind): + """ + Initialize a nearest upsampling block with its ReLU and BatchNorm. + """ + super(NearestUpsampleBlock, self).__init__() + self.layer_ind = layer_ind + return + + def forward(self, x, batch): + return closest_pool(x, batch.upsamples[self.layer_ind - 1]) + + +class MaxPoolBlock(nn.Module): + + def __init__(self, layer_ind): + """ + Initialize a max pooling block with its ReLU and BatchNorm. + """ + super(MaxPoolBlock, self).__init__() + self.layer_ind = layer_ind + return + + def forward(self, x, batch): + return max_pool(x, batch.pools[self.layer_ind + 1]) + diff --git a/plot_convergence.py b/plot_convergence.py new file mode 100644 index 0000000..5371220 --- /dev/null +++ b/plot_convergence.py @@ -0,0 +1,1455 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Callable script to test any model on any dataset +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 11/06/2018 +# + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Imports and global variables +# \**********************************/ +# + +# Common libs +import os +import torch +import numpy as np +import matplotlib.pyplot as plt +from os.path import isfile, join, exists +from os import listdir, remove, getcwd +from sklearn.metrics import confusion_matrix +import time + +# My libs +from utils.config import Config +from utils.metrics import IoU_from_confusions, smooth_metrics, fast_confusion +from utils.ply import read_ply + +# Datasets +from datasets.ModelNet40 import ModelNet40Dataset +from datasets.S3DIS import S3DISDataset + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Utility functions +# \***********************/ +# + + +def running_mean(signal, n, axis=0, stride=1): + signal = np.array(signal) + torch_conv = torch.nn.Conv1d(1, 1, kernel_size=2*n+1, stride=stride, bias=False) + torch_conv.weight.requires_grad_(False) + torch_conv.weight *= 0 + torch_conv.weight += 1 / (2*n+1) + if signal.ndim == 1: + torch_signal = torch.from_numpy(signal.reshape([1, 1, -1]).astype(np.float32)) + return torch_conv(torch_signal).squeeze().numpy() + + elif signal.ndim == 2: + print('TODO implement with torch and stride here') + smoothed = np.empty(signal.shape) + if axis == 0: + for i, sig in enumerate(signal): + sig_sum = np.convolve(sig, np.ones((2*n+1,)), mode='same') + sig_num = np.convolve(sig*0+1, np.ones((2*n+1,)), mode='same') + smoothed[i, :] = sig_sum / sig_num + elif axis == 1: + for i, sig in enumerate(signal.T): + sig_sum = np.convolve(sig, np.ones((2*n+1,)), mode='same') + sig_num = np.convolve(sig*0+1, np.ones((2*n+1,)), mode='same') + smoothed[:, i] = sig_sum / sig_num + else: + print('wrong axis') + return smoothed + + else: + print('wrong dimensions') + return None + + +def IoU_multi_metrics(all_IoUs, smooth_n): + + # Get mean IoU for consecutive epochs to directly get a mean + all_mIoUs = [np.hstack([np.mean(obj_IoUs, axis=1) for obj_IoUs in epoch_IoUs]) for epoch_IoUs in all_IoUs] + smoothed_mIoUs = [] + for epoch in range(len(all_mIoUs)): + i0 = max(epoch - smooth_n, 0) + i1 = min(epoch + smooth_n + 1, len(all_mIoUs)) + smoothed_mIoUs += [np.mean(np.hstack(all_mIoUs[i0:i1]))] + + # Get mean for each class + all_objs_mIoUs = [[np.mean(obj_IoUs, axis=1) for obj_IoUs in epoch_IoUs] for epoch_IoUs in all_IoUs] + smoothed_obj_mIoUs = [] + for epoch in range(len(all_objs_mIoUs)): + i0 = max(epoch - smooth_n, 0) + i1 = min(epoch + smooth_n + 1, len(all_objs_mIoUs)) + + epoch_obj_mIoUs = [] + for obj in range(len(all_objs_mIoUs[0])): + epoch_obj_mIoUs += [np.mean(np.hstack([objs_mIoUs[obj] for objs_mIoUs in all_objs_mIoUs[i0:i1]]))] + + smoothed_obj_mIoUs += [epoch_obj_mIoUs] + + return np.array(smoothed_mIoUs), np.array(smoothed_obj_mIoUs) + + +def IoU_class_metrics(all_IoUs, smooth_n): + + # Get mean IoU per class for consecutive epochs to directly get a mean without further smoothing + smoothed_IoUs = [] + for epoch in range(len(all_IoUs)): + i0 = max(epoch - smooth_n, 0) + i1 = min(epoch + smooth_n + 1, len(all_IoUs)) + smoothed_IoUs += [np.mean(np.vstack(all_IoUs[i0:i1]), axis=0)] + smoothed_IoUs = np.vstack(smoothed_IoUs) + smoothed_mIoUs = np.mean(smoothed_IoUs, axis=1) + + return smoothed_IoUs, smoothed_mIoUs + + +def load_confusions(filename, n_class): + + with open(filename, 'r') as f: + lines = f.readlines() + + confs = np.zeros((len(lines), n_class, n_class)) + for i, line in enumerate(lines): + C = np.array([int(value) for value in line.split()]) + confs[i, :, :] = C.reshape((n_class, n_class)) + + return confs + + +def load_training_results(path): + + filename = join(path, 'training.txt') + with open(filename, 'r') as f: + lines = f.readlines() + + epochs = [] + steps = [] + L_out = [] + L_p = [] + acc = [] + t = [] + for line in lines[1:]: + line_info = line.split() + if (len(line) > 0): + epochs += [int(line_info[0])] + steps += [int(line_info[1])] + L_out += [float(line_info[2])] + L_p += [float(line_info[3])] + acc += [float(line_info[4])] + t += [float(line_info[5])] + else: + break + + return epochs, steps, L_out, L_p, acc, t + + +def load_single_IoU(filename, n_parts): + + with open(filename, 'r') as f: + lines = f.readlines() + + # Load all IoUs + all_IoUs = [] + for i, line in enumerate(lines): + all_IoUs += [np.reshape([float(IoU) for IoU in line.split()], [-1, n_parts])] + return all_IoUs + + +def load_snap_clouds(path, dataset, only_last=False): + + cloud_folders = np.array([join(path, f) for f in listdir(path) if f.startswith('val_preds')]) + cloud_epochs = np.array([int(f.split('_')[-1]) for f in cloud_folders]) + epoch_order = np.argsort(cloud_epochs) + cloud_epochs = cloud_epochs[epoch_order] + cloud_folders = cloud_folders[epoch_order] + + Confs = np.zeros((len(cloud_epochs), dataset.num_classes, dataset.num_classes), dtype=np.int32) + for c_i, cloud_folder in enumerate(cloud_folders): + if only_last and c_i < len(cloud_epochs) - 1: + continue + + # Load confusion if previously saved + conf_file = join(cloud_folder, 'conf.txt') + if isfile(conf_file): + Confs[c_i] += np.loadtxt(conf_file, dtype=np.int32) + + else: + for f in listdir(cloud_folder): + if f.endswith('.ply') and not f.endswith('sub.ply'): + data = read_ply(join(cloud_folder, f)) + labels = data['class'] + preds = data['preds'] + Confs[c_i] += fast_confusion(labels, preds, dataset.label_values).astype(np.int32) + + np.savetxt(conf_file, Confs[c_i], '%12d') + + # Erase ply to save disk memory + if c_i < len(cloud_folders) - 1: + for f in listdir(cloud_folder): + if f.endswith('.ply'): + remove(join(cloud_folder, f)) + + # Remove ignored labels from confusions + for l_ind, label_value in reversed(list(enumerate(dataset.label_values))): + if label_value in dataset.ignored_labels: + Confs = np.delete(Confs, l_ind, axis=1) + Confs = np.delete(Confs, l_ind, axis=2) + + return cloud_epochs, IoU_from_confusions(Confs) + + +def load_multi_snap_clouds(path, dataset, file_i, only_last=False): + + cloud_folders = np.array([join(path, f) for f in listdir(path) if f.startswith('val_preds')]) + cloud_epochs = np.array([int(f.split('_')[-1]) for f in cloud_folders]) + epoch_order = np.argsort(cloud_epochs) + cloud_epochs = cloud_epochs[epoch_order] + cloud_folders = cloud_folders[epoch_order] + + if len(cloud_folders) > 0: + dataset_folders = [f for f in listdir(cloud_folders[0]) if dataset.name in f] + cloud_folders = [join(f, dataset_folders[file_i]) for f in cloud_folders] + + Confs = np.zeros((len(cloud_epochs), dataset.num_classes, dataset.num_classes), dtype=np.int32) + for c_i, cloud_folder in enumerate(cloud_folders): + if only_last and c_i < len(cloud_epochs) - 1: + continue + + # Load confusion if previously saved + conf_file = join(cloud_folder, 'conf_{:s}.txt'.format(dataset.name)) + if isfile(conf_file): + Confs[c_i] += np.loadtxt(conf_file, dtype=np.int32) + + else: + for f in listdir(cloud_folder): + if f.endswith('.ply') and not f.endswith('sub.ply'): + if np.any([cloud_path.endswith(f) for cloud_path in dataset.train_files]): + data = read_ply(join(cloud_folder, f)) + labels = data['class'] + preds = data['preds'] + Confs[c_i] += confusion_matrix(labels, preds, dataset.label_values).astype(np.int32) + + np.savetxt(conf_file, Confs[c_i], '%12d') + + # Erase ply to save disk memory + if c_i < len(cloud_folders) - 1: + for f in listdir(cloud_folder): + if f.endswith('.ply'): + remove(join(cloud_folder, f)) + + # Remove ignored labels from confusions + for l_ind, label_value in reversed(list(enumerate(dataset.label_values))): + if label_value in dataset.ignored_labels: + Confs = np.delete(Confs, l_ind, axis=1) + Confs = np.delete(Confs, l_ind, axis=2) + + return cloud_epochs, IoU_from_confusions(Confs) + + +def load_multi_IoU(filename, n_parts): + + with open(filename, 'r') as f: + lines = f.readlines() + + # Load all IoUs + all_IoUs = [] + for i, line in enumerate(lines): + obj_IoUs = [[float(IoU) for IoU in s.split()] for s in line.split('/')] + obj_IoUs = [np.reshape(IoUs, [-1, n_parts[obj]]) for obj, IoUs in enumerate(obj_IoUs)] + all_IoUs += [obj_IoUs] + return all_IoUs + + +def compare_trainings(list_of_paths, list_of_labels=None): + + # Parameters + # ********** + + plot_lr = False + smooth_epochs = 0.5 + stride = 2 + + if list_of_labels is None: + list_of_labels = [str(i) for i in range(len(list_of_paths))] + + # Read Training Logs + # ****************** + + all_epochs = [] + all_loss = [] + all_lr = [] + all_times = [] + all_RAMs = [] + + for path in list_of_paths: + + print(path) + + if ('val_IoUs.txt' in [f for f in listdir(path)]) or ('val_confs.txt' in [f for f in listdir(path)]): + config = Config() + config.load(path) + else: + continue + + # Load results + epochs, steps, L_out, L_p, acc, t = load_training_results(path) + epochs = np.array(epochs, dtype=np.int32) + epochs_d = np.array(epochs, dtype=np.float32) + steps = np.array(steps, dtype=np.float32) + + # Compute number of steps per epoch + max_e = np.max(epochs) + first_e = np.min(epochs) + epoch_n = [] + for i in range(first_e, max_e): + bool0 = epochs == i + e_n = np.sum(bool0) + epoch_n.append(e_n) + epochs_d[bool0] += steps[bool0] / e_n + smooth_n = int(np.mean(epoch_n) * smooth_epochs) + smooth_loss = running_mean(L_out, smooth_n, stride=stride) + all_loss += [smooth_loss] + all_epochs += [epochs_d[smooth_n:-smooth_n:stride]] + all_times += [t[smooth_n:-smooth_n:stride]] + + # Learning rate + if plot_lr: + lr_decay_v = np.array([lr_d for ep, lr_d in config.lr_decays.items()]) + lr_decay_e = np.array([ep for ep, lr_d in config.lr_decays.items()]) + max_e = max(np.max(all_epochs[-1]) + 1, np.max(lr_decay_e) + 1) + lr_decays = np.ones(int(np.ceil(max_e)), dtype=np.float32) + lr_decays[0] = float(config.learning_rate) + lr_decays[lr_decay_e] = lr_decay_v + lr = np.cumprod(lr_decays) + all_lr += [lr[np.floor(all_epochs[-1]).astype(np.int32)]] + + # Plots learning rate + # ******************* + + + if plot_lr: + # Figure + fig = plt.figure('lr') + for i, label in enumerate(list_of_labels): + plt.plot(all_epochs[i], all_lr[i], linewidth=1, label=label) + + # Set names for axes + plt.xlabel('epochs') + plt.ylabel('lr') + plt.yscale('log') + + # Display legends and title + plt.legend(loc=1) + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + # ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + # Plots loss + # ********** + + # Figure + fig = plt.figure('loss') + for i, label in enumerate(list_of_labels): + plt.plot(all_epochs[i], all_loss[i], linewidth=1, label=label) + + # Set names for axes + plt.xlabel('epochs') + plt.ylabel('loss') + plt.yscale('log') + + # Display legends and title + plt.legend(loc=1) + plt.title('Losses compare') + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + # ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + # Plot Times + # ********** + + # Figure + fig = plt.figure('time') + for i, label in enumerate(list_of_labels): + plt.plot(all_epochs[i], np.array(all_times[i]) / 3600, linewidth=1, label=label) + + # Set names for axes + plt.xlabel('epochs') + plt.ylabel('time') + # plt.yscale('log') + + # Display legends and title + plt.legend(loc=0) + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + # ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + # Show all + plt.show() + + +def compare_convergences_multisegment(list_of_paths, list_of_labels=None): + + # Parameters + # ********** + + steps_per_epoch = 0 + smooth_n = 5 + + if list_of_labels is None: + list_of_labels = [str(i) for i in range(len(list_of_paths))] + + # Read Logs + # ********* + + all_pred_epochs = [] + all_instances_mIoUs = [] + all_objs_mIoUs = [] + all_objs_IoUs = [] + all_parts = [] + + obj_list = ['Air', 'Bag', 'Cap', 'Car', 'Cha', 'Ear', 'Gui', 'Kni', + 'Lam', 'Lap', 'Mot', 'Mug', 'Pis', 'Roc', 'Ska', 'Tab'] + print('Objs | Inst | Air Bag Cap Car Cha Ear Gui Kni Lam Lap Mot Mug Pis Roc Ska Tab') + print('-----|------|--------------------------------------------------------------------------------') + for path in list_of_paths: + + # Load parameters + config = Config() + config.load(path) + + # Get the number of classes + n_parts = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3] + part = config.dataset.split('_')[-1] + + # Get validation confusions + file = join(path, 'val_IoUs.txt') + val_IoUs = load_multi_IoU(file, n_parts) + + file = join(path, 'vote_IoUs.txt') + vote_IoUs = load_multi_IoU(file, n_parts) + + #print(len(val_IoUs[0])) + #print(val_IoUs[0][0].shape) + + # Get mean IoU + #instances_mIoUs, objs_mIoUs = IoU_multi_metrics(val_IoUs, smooth_n) + + # Get mean IoU + instances_mIoUs, objs_mIoUs = IoU_multi_metrics(vote_IoUs, smooth_n) + + # Aggregate results + all_pred_epochs += [np.array([i for i in range(len(val_IoUs))])] + all_instances_mIoUs += [instances_mIoUs] + all_objs_IoUs += [objs_mIoUs] + all_objs_mIoUs += [np.mean(objs_mIoUs, axis=1)] + + if part == 'multi': + s = '{:4.1f} | {:4.1f} | '.format(100 * np.mean(objs_mIoUs[-1]), 100 * instances_mIoUs[-1]) + for obj_mIoU in objs_mIoUs[-1]: + s += '{:4.1f} '.format(100 * obj_mIoU) + print(s) + else: + s = ' -- | -- | ' + for obj_name in obj_list: + if part.startswith(obj_name): + s += '{:4.1f} '.format(100 * instances_mIoUs[-1]) + else: + s += ' -- '.format(100 * instances_mIoUs[-1]) + print(s) + all_parts += [part] + + # Plots + # ***** + + if 'multi' in all_parts: + + # Figure + fig = plt.figure('Instances mIoU') + for i, label in enumerate(list_of_labels): + if all_parts[i] == 'multi': + plt.plot(all_pred_epochs[i], all_instances_mIoUs[i], linewidth=1, label=label) + plt.xlabel('epochs') + plt.ylabel('IoU') + + # Set limits for y axis + #plt.ylim(0.55, 0.95) + + # Display legends and title + plt.legend(loc=4) + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + #ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + # Figure + fig = plt.figure('mean of categories mIoU') + for i, label in enumerate(list_of_labels): + if all_parts[i] == 'multi': + plt.plot(all_pred_epochs[i], all_objs_mIoUs[i], linewidth=1, label=label) + plt.xlabel('epochs') + plt.ylabel('IoU') + + # Set limits for y axis + #plt.ylim(0.8, 1) + + # Display legends and title + plt.legend(loc=4) + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + #ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + for obj_i, obj_name in enumerate(obj_list): + if np.any([part.startswith(obj_name) for part in all_parts]): + # Figure + fig = plt.figure(obj_name + ' mIoU') + for i, label in enumerate(list_of_labels): + if all_parts[i] == 'multi': + plt.plot(all_pred_epochs[i], all_objs_IoUs[i][:, obj_i], linewidth=1, label=label) + elif all_parts[i].startswith(obj_name): + plt.plot(all_pred_epochs[i], all_objs_mIoUs[i], linewidth=1, label=label) + plt.xlabel('epochs') + plt.ylabel('IoU') + + # Set limits for y axis + #plt.ylim(0.8, 1) + + # Display legends and title + plt.legend(loc=4) + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + #ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + + + # Show all + plt.show() + + +def compare_convergences_segment(dataset, list_of_paths, list_of_names=None): + + # Parameters + # ********** + + smooth_n = 10 + + if list_of_names is None: + list_of_names = [str(i) for i in range(len(list_of_paths))] + + # Read Logs + # ********* + + all_pred_epochs = [] + all_mIoUs = [] + all_class_IoUs = [] + all_snap_epochs = [] + all_snap_IoUs = [] + + # Load parameters + config = Config() + config.load(list_of_paths[0]) + + class_list = [dataset.label_to_names[label] for label in dataset.label_values + if label not in dataset.ignored_labels] + + s = '{:^10}|'.format('mean') + for c in class_list: + s += '{:^10}'.format(c) + print(s) + print(10*'-' + '|' + 10*config.num_classes*'-') + for path in list_of_paths: + + # Get validation IoUs + file = join(path, 'val_IoUs.txt') + val_IoUs = load_single_IoU(file, config.num_classes) + + # Get mean IoU + class_IoUs, mIoUs = IoU_class_metrics(val_IoUs, smooth_n) + + # Aggregate results + all_pred_epochs += [np.array([i for i in range(len(val_IoUs))])] + all_mIoUs += [mIoUs] + all_class_IoUs += [class_IoUs] + + s = '{:^10.1f}|'.format(100*mIoUs[-1]) + for IoU in class_IoUs[-1]: + s += '{:^10.1f}'.format(100*IoU) + print(s) + + # Get optional full validation on clouds + snap_epochs, snap_IoUs = load_snap_clouds(path, dataset) + all_snap_epochs += [snap_epochs] + all_snap_IoUs += [snap_IoUs] + + print(10*'-' + '|' + 10*config.num_classes*'-') + for snap_IoUs in all_snap_IoUs: + if len(snap_IoUs) > 0: + s = '{:^10.1f}|'.format(100*np.mean(snap_IoUs[-1])) + for IoU in snap_IoUs[-1]: + s += '{:^10.1f}'.format(100*IoU) + else: + s = '{:^10s}'.format('-') + for _ in range(config.num_classes): + s += '{:^10s}'.format('-') + print(s) + + # Plots + # ***** + + # Figure + fig = plt.figure('mIoUs') + for i, name in enumerate(list_of_names): + p = plt.plot(all_pred_epochs[i], all_mIoUs[i], '--', linewidth=1, label=name) + plt.plot(all_snap_epochs[i], np.mean(all_snap_IoUs[i], axis=1), linewidth=1, color=p[-1].get_color()) + plt.xlabel('epochs') + plt.ylabel('IoU') + + # Set limits for y axis + #plt.ylim(0.55, 0.95) + + # Display legends and title + plt.legend(loc=4) + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + #ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + displayed_classes = [0, 1, 2, 3, 4, 5, 6, 7] + displayed_classes = [] + for c_i, c_name in enumerate(class_list): + if c_i in displayed_classes: + + # Figure + fig = plt.figure(c_name + ' IoU') + for i, name in enumerate(list_of_names): + plt.plot(all_pred_epochs[i], all_class_IoUs[i][:, c_i], linewidth=1, label=name) + plt.xlabel('epochs') + plt.ylabel('IoU') + + # Set limits for y axis + #plt.ylim(0.8, 1) + + # Display legends and title + plt.legend(loc=4) + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + #ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + # Show all + plt.show() + + +def compare_convergences_classif(list_of_paths, list_of_labels=None): + + # Parameters + # ********** + + steps_per_epoch = 0 + smooth_n = 12 + + if list_of_labels is None: + list_of_labels = [str(i) for i in range(len(list_of_paths))] + + # Read Logs + # ********* + + all_pred_epochs = [] + all_val_OA = [] + all_train_OA = [] + all_vote_OA = [] + all_vote_confs = [] + + + for path in list_of_paths: + + # Load parameters + config = Config() + config.load(list_of_paths[0]) + + # Get the number of classes + n_class = config.num_classes + + # Load epochs + epochs, _, _, _, _, _ = load_training_results(path) + first_e = np.min(epochs) + + # Get validation confusions + file = join(path, 'val_confs.txt') + val_C1 = load_confusions(file, n_class) + val_PRE, val_REC, val_F1, val_IoU, val_ACC = smooth_metrics(val_C1, smooth_n=smooth_n) + + # Get vote confusions + file = join(path, 'vote_confs.txt') + if exists(file): + vote_C2 = load_confusions(file, n_class) + vote_PRE, vote_REC, vote_F1, vote_IoU, vote_ACC = smooth_metrics(vote_C2, smooth_n=2) + else: + vote_C2 = val_C1 + vote_PRE, vote_REC, vote_F1, vote_IoU, vote_ACC = (val_PRE, val_REC, val_F1, val_IoU, val_ACC) + + # Aggregate results + all_pred_epochs += [np.array([i+first_e for i in range(len(val_ACC))])] + all_val_OA += [val_ACC] + all_vote_OA += [vote_ACC] + all_vote_confs += [vote_C2] + + print() + + # Best scores + # *********** + + for i, label in enumerate(list_of_labels): + + print('\n' + label + '\n' + '*' * len(label) + '\n') + print(list_of_paths[i]) + + best_epoch = np.argmax(all_vote_OA[i]) + print('Best Accuracy : {:.1f} % (epoch {:d})'.format(100 * all_vote_OA[i][best_epoch], best_epoch)) + + confs = all_vote_confs[i] + + """ + s = '' + for cc in confs[best_epoch]: + for c in cc: + s += '{:.0f} '.format(c) + s += '\n' + print(s) + """ + + TP_plus_FN = np.sum(confs, axis=-1, keepdims=True) + class_avg_confs = confs.astype(np.float32) / TP_plus_FN.astype(np.float32) + diags = np.diagonal(class_avg_confs, axis1=-2, axis2=-1) + class_avg_ACC = np.sum(diags, axis=-1) / np.sum(class_avg_confs, axis=(-1, -2)) + + print('Corresponding mAcc : {:.1f} %'.format(100 * class_avg_ACC[best_epoch])) + + # Plots + # ***** + + for fig_name, OA in zip(['Validation', 'Vote'], [all_val_OA, all_vote_OA]): + + # Figure + fig = plt.figure(fig_name) + for i, label in enumerate(list_of_labels): + plt.plot(all_pred_epochs[i], OA[i], linewidth=1, label=label) + plt.xlabel('epochs') + plt.ylabel(fig_name + ' Accuracy') + + # Set limits for y axis + #plt.ylim(0.55, 0.95) + + # Display legends and title + plt.legend(loc=4) + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + #ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + #for i, label in enumerate(list_of_labels): + # print(label, np.max(all_train_OA[i]), np.max(all_val_OA[i])) + + # Show all + plt.show() + + +def compare_convergences_multicloud(list_of_paths, multi, multi_datasets, list_of_names=None): + + # Parameters + # ********** + + smooth_n = 10 + + if list_of_names is None: + list_of_names = [str(i) for i in range(len(list_of_paths))] + + + # Loop on all datasets: + for plot_dataset in multi_datasets: + print('\n') + print(plot_dataset) + print('*'*len(plot_dataset)) + print() + + # Load dataset parameters + if plot_dataset.startswith('S3DIS'): + dataset = S3DISDataset() + elif plot_dataset.startswith('Scann'): + dataset = ScannetDataset() + elif plot_dataset.startswith('Seman'): + dataset = Semantic3DDataset() + elif plot_dataset.startswith('NPM3D'): + dataset = NPM3DDataset() + else: + raise ValueError('Unsupported dataset : ' + plot_dataset) + + # Read Logs + # ********* + + all_pred_epochs = [] + all_mIoUs = [] + all_class_IoUs = [] + all_snap_epochs = [] + all_snap_IoUs = [] + all_names = [] + + class_list = [dataset.label_to_names[label] for label in dataset.label_values + if label not in dataset.ignored_labels] + + s = '{:^10}|'.format('mean') + for c in class_list: + s += '{:^10}'.format(c) + print(s) + print(10*'-' + '|' + 10*dataset.num_classes*'-') + for log_i, (path, is_multi) in enumerate(zip(list_of_paths, multi)): + + n_c = None + if is_multi: + config = MultiConfig() + config.load(path) + if plot_dataset in config.datasets: + val_IoU_files = [] + for d_i in np.where(np.array(config.datasets) == plot_dataset)[0]: + n_c = config.num_classes[d_i] + val_IoU_files.append(join(path, 'val_IoUs_{:d}_{:s}.txt'.format(d_i, plot_dataset))) + else: + continue + else: + config = Config() + config.load(path) + if plot_dataset == config.dataset: + n_c = config.num_classes + val_IoU_files = [join(path, 'val_IoUs.txt')] + else: + continue + + for file_i, file in enumerate(val_IoU_files): + + # Load validation IoUs + val_IoUs = load_single_IoU(file, n_c) + + # Get mean IoU + class_IoUs, mIoUs = IoU_class_metrics(val_IoUs, smooth_n) + + # Aggregate results + all_pred_epochs += [np.array([i for i in range(len(val_IoUs))])] + all_mIoUs += [mIoUs] + all_class_IoUs += [class_IoUs] + all_names += [list_of_names[log_i]+'_{:d}'.format(file_i+1)] + + s = '{:^10.1f}|'.format(100*mIoUs[-1]) + for IoU in class_IoUs[-1]: + s += '{:^10.1f}'.format(100*IoU) + print(s) + + # Get optional full validation on clouds + if is_multi: + snap_epochs, snap_IoUs = load_multi_snap_clouds(path, dataset, file_i) + else: + snap_epochs, snap_IoUs = load_snap_clouds(path, dataset) + all_snap_epochs += [snap_epochs] + all_snap_IoUs += [snap_IoUs] + + print(10*'-' + '|' + 10*dataset.num_classes*'-') + for snap_IoUs in all_snap_IoUs: + if len(snap_IoUs) > 0: + s = '{:^10.1f}|'.format(100*np.mean(snap_IoUs[-1])) + for IoU in snap_IoUs[-1]: + s += '{:^10.1f}'.format(100*IoU) + else: + s = '{:^10s}'.format('-') + for _ in range(dataset.num_classes): + s += '{:^10s}'.format('-') + print(s) + + # Plots + # ***** + + # Figure + fig = plt.figure('mIoUs') + for i, name in enumerate(all_names): + p = plt.plot(all_pred_epochs[i], all_mIoUs[i], '--', linewidth=1, label=name) + plt.plot(all_snap_epochs[i], np.mean(all_snap_IoUs[i], axis=1), linewidth=1, color=p[-1].get_color()) + + plt.title(plot_dataset) + plt.xlabel('epochs') + plt.ylabel('IoU') + + # Set limits for y axis + #plt.ylim(0.55, 0.95) + + # Display legends and title + plt.legend(loc=4) + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + #ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + displayed_classes = [0, 1, 2, 3, 4, 5, 6, 7] + displayed_classes = [] + for c_i, c_name in enumerate(class_list): + if c_i in displayed_classes: + + # Figure + fig = plt.figure(c_name + ' IoU') + for i, name in enumerate(list_of_names): + plt.plot(all_pred_epochs[i], all_class_IoUs[i][:, c_i], linewidth=1, label=name) + plt.xlabel('epochs') + plt.ylabel('IoU') + + # Set limits for y axis + #plt.ylim(0.8, 1) + + # Display legends and title + plt.legend(loc=4) + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + #ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + + + # Show all + plt.show() + + +def compare_convergences_SLAM(dataset, list_of_paths, list_of_names=None): + + # Parameters + # ********** + + smooth_n = 10 + + if list_of_names is None: + list_of_names = [str(i) for i in range(len(list_of_paths))] + + # Read Logs + # ********* + + all_pred_epochs = [] + all_val_mIoUs = [] + all_val_class_IoUs = [] + all_subpart_mIoUs = [] + all_subpart_class_IoUs = [] + + # Load parameters + config = Config() + config.load(list_of_paths[0]) + + class_list = [dataset.label_to_names[label] for label in dataset.label_values + if label not in dataset.ignored_labels] + + s = '{:^10}|'.format('mean') + for c in class_list: + s += '{:^10}'.format(c) + print(s) + print(10*'-' + '|' + 10*config.num_classes*'-') + for path in list_of_paths: + + # Get validation IoUs + file = join(path, 'val_IoUs.txt') + val_IoUs = load_single_IoU(file, config.num_classes) + + # Get Subpart IoUs + file = join(path, 'subpart_IoUs.txt') + subpart_IoUs = load_single_IoU(file, config.num_classes) + + # Get mean IoU + val_class_IoUs, val_mIoUs = IoU_class_metrics(val_IoUs, smooth_n) + subpart_class_IoUs, subpart_mIoUs = IoU_class_metrics(subpart_IoUs, smooth_n) + + # Aggregate results + all_pred_epochs += [np.array([i for i in range(len(val_IoUs))])] + all_val_mIoUs += [val_mIoUs] + all_val_class_IoUs += [val_class_IoUs] + all_subpart_mIoUs += [subpart_mIoUs] + all_subpart_class_IoUs += [subpart_class_IoUs] + + s = '{:^10.1f}|'.format(100*subpart_mIoUs[-1]) + for IoU in subpart_class_IoUs[-1]: + s += '{:^10.1f}'.format(100*IoU) + print(s) + + + print(10*'-' + '|' + 10*config.num_classes*'-') + for snap_IoUs in all_val_class_IoUs: + if len(snap_IoUs) > 0: + s = '{:^10.1f}|'.format(100*np.mean(snap_IoUs[-1])) + for IoU in snap_IoUs[-1]: + s += '{:^10.1f}'.format(100*IoU) + else: + s = '{:^10s}'.format('-') + for _ in range(config.num_classes): + s += '{:^10s}'.format('-') + print(s) + + # Plots + # ***** + + # Figure + fig = plt.figure('mIoUs') + for i, name in enumerate(list_of_names): + p = plt.plot(all_pred_epochs[i], all_subpart_mIoUs[i], '--', linewidth=1, label=name) + plt.plot(all_pred_epochs[i], all_val_mIoUs[i], linewidth=1, color=p[-1].get_color()) + plt.xlabel('epochs') + plt.ylabel('IoU') + + # Set limits for y axis + #plt.ylim(0.55, 0.95) + + # Display legends and title + plt.legend(loc=4) + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + #ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + displayed_classes = [0, 1, 2, 3, 4, 5, 6, 7] + displayed_classes = [] + for c_i, c_name in enumerate(class_list): + if c_i in displayed_classes: + + # Figure + fig = plt.figure(c_name + ' IoU') + for i, name in enumerate(list_of_names): + plt.plot(all_pred_epochs[i], all_val_class_IoUs[i][:, c_i], linewidth=1, label=name) + plt.xlabel('epochs') + plt.ylabel('IoU') + + # Set limits for y axis + #plt.ylim(0.8, 1) + + # Display legends and title + plt.legend(loc=4) + + # Customize the graph + ax = fig.gca() + ax.grid(linestyle='-.', which='both') + #ax.set_yticks(np.arange(0.8, 1.02, 0.02)) + + + + # Show all + plt.show() + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Main Call +# \***************/ +# + + +def ModelNet40_first_test(): + """ + First tries with ModelNet40 + First we compare convergence of a very very deep network on ModelNet40, with our without bn + Then, We try the resuming of previous trainings. Which works quite well. + However in the mean time, we change how validation worked by calling net.eval()/net.train() before/after + validation. It seems that the network perform strange when calling net.eval()/net.train() although it should be the + right way to do it. + Then we try to change BatchNorm1D with InstanceNorm1D and compare with and without calling eval/train at validation. + (Also with a faster lr decay). + --- MISTAKE FOUND --- the batch norm momentum was inverted 0.98 instead of 0.02. + See next experiment for correct convergences. Instance norm seems not as good + """ + + # Using the dates of the logs, you can easily gather consecutive ones. All logs should be of the same dataset. + start = 'Log_2020-03-18_16-04-20' + end = 'Log_2020-03-20_16-59-40' + + if end < 'Log_2020-03-22_19-30-19': + res_path = 'old_results' + else: + res_path = 'results' + + logs = np.sort([join(res_path, l) for l in listdir(res_path) if start <= l <= end]) + + # Give names to the logs (for legends) + logs_names = ['with_bn', + 'without_bn', + 'with_bn2', + 'without_bn2', + 'lrd_80_Inorm_eval_train', + 'lrd_80_Inorm_always_train', + 'test'] + + logs_names = np.array(logs_names[:len(logs)]) + + return logs, logs_names + + +def ModelNet40_batch_norm(): + """ + Compare different type of batch norm now that it has been fixed. Batch norm seems the best easily. Instance norm + crewated a NAN loss so avoid this one. + Now try fast experiments. First reduce network size. Reducing the number of convolution per layer does not affect + results (maybe because dataset is too simple???). 5 small layers is way better that 4 big layers. + Now reduce number of step per epoch and maybe try balanced sampler. Balanced sampler with fewer steps per epoch is + way faster for convergence and gets nearly the same scores. so good for experimenting. However we cant really + conclude between parameters which will get the same score (like the more layers) because the dataset my be + limitating. We can only conclude if something is not good and reduce score. + """ + + # Using the dates of the logs, you can easily gather consecutive ones. All logs should be of the same dataset. + start = 'Log_2020-03-20_16-59-41' + end = 'Log_2020-04-13_18-14-44' + + if end < 'Log_2020-03-22_19-30-19': + res_path = 'old_results' + else: + res_path = 'results' + + logs = np.sort([join(res_path, l) for l in listdir(res_path) if start <= l <= end]) + + # Give names to the logs (for legends) + logs_names = ['no_norm', + 'IN', + 'BN', + '5_small_layer-d0=0.02', + '3_big_layer-d0=0.02', + '3_big_layer-d0=0.04', + 'small-e_n=300', + 'small-e_n=300-balanced_train', + 'small-e_n=300-balanced_traintest', + 'test'] + + logs_names = np.array(logs_names[:len(logs)]) + + return logs, logs_names + + +def ModelNet40_fast_vs_results(): + """ + Try lr decay with fast convergence (epoch_n=300 and balanced traintest). 80 is a good value. + """ + + # Using the dates of the logs, you can easily gather consecutive ones. All logs should be of the same dataset. + start = 'Log_2020-03-21_16-09-17' + end = 'Log_2020-03-21_16-09-36' + + if end < 'Log_2020-03-22_19-30-19': + res_path = 'old_results' + else: + res_path = 'results' + + logs = np.sort([join(res_path, l) for l in listdir(res_path) if start <= l <= end]) + logs = np.insert(logs, 1, join(res_path, 'Log_2020-03-21_11-57-45')) + + # Give names to the logs (for legends) + logs_names = ['lrd=120', + 'lrd=80', + 'lrd=40', + 'test'] + + logs_names = np.array(logs_names[:len(logs)]) + + return logs, logs_names + + +def ModelNet40_grad_clipping(): + """ + Test different grad clipping. No difference so we can move on + """ + + # Using the dates of the logs, you can easily gather consecutive ones. All logs should be of the same dataset. + start = 'Log_2020-03-21_18-21-37' + end = 'Log_2020-03-21_18-30-01' + + if end < 'Log_2020-03-22_19-30-19': + res_path = 'old_results' + else: + res_path = 'results' + + logs = np.sort([join(res_path, l) for l in listdir(res_path) if start <= l <= end]) + logs = np.insert(logs, 0, join(res_path, 'Log_2020-03-21_11-57-45')) + + # Give names to the logs (for legends) + logs_names = ['value_clip_100', + 'norm_clip_100', + 'value_clip_10', + 'norm_clip_10', + 'no_clip', + 'test'] + + logs_names = np.array(logs_names[:len(logs)]) + + return logs, logs_names + + +def ModelNet40_KP_extent(): + """ + Test differents mode et kp extent. sum et extent=2.0 definitivement moins bon (trop de recouvrement des kp + influences, noyau moins versatile). les closest semble plutot bon et le sum extent=1.5 pas mal du tout () peut + etre le meilleur. A confirmer sur gros dataset + """ + + # Using the dates of the logs, you can easily gather consecutive ones. All logs should be of the same dataset. + start = 'Log_2020-03-21_18-30-02' + end = 'Log_2020-03-21_23-36-18' + + if end < 'Log_2020-03-22_19-30-19': + res_path = 'old_results' + else: + res_path = 'results' + + logs = np.sort([join(res_path, l) for l in listdir(res_path) if start <= l <= end]) + logs = np.insert(logs, 0, join(res_path, 'Log_2020-03-21_11-57-45')) + + # Give names to the logs (for legends) + logs_names = ['KPe=1.0_sum_linear', + 'KPe=1.5_sum_linear', + 'KPe=2.0_sum_linear', + 'KPe=1.5_closest_linear', + 'KPe=2.0_closest_linear', + 'test'] + + logs_names = np.array(logs_names[:len(logs)]) + + return logs, logs_names + + +def ModelNet40_gaussian(): + """ + Test different extent in gaussian mode. extent=1.5 seems the best. 2.0 is not bad. But in any case, it does not + perform better than 1.5-linear-sum at least on this dataset. + """ + + # Using the dates of the logs, you can easily gather consecutive ones. All logs should be of the same dataset. + start = 'Log_2020-03-21_23-36-19' + end = 'Log_2020-04-13_18-14-44' + + if end < 'Log_2020-03-22_19-30-19': + res_path = 'old_results' + else: + res_path = 'results' + + logs = np.sort([join(res_path, l) for l in listdir(res_path) if start <= l <= end]) + logs = np.insert(logs, 4, join(res_path, 'Log_2020-03-21_19-35-11')) + + # Give names to the logs (for legends) + logs_names = ['KPe=1.0_sum_gaussian', + 'KPe=1.5_sum_gaussian', + 'KPe=2.0_sum_gaussian', + 'KPe=2.5_sum_gaussian', + 'KPe=1.5_sum_linear', + 'test'] + + logs_names = np.array(logs_names[:len(logs)]) + + return logs, logs_names + + +def ModelNet40_normals(): + """ + Test different way to add normals. Seems pretty much the same and we dont care about normals. + """ + + # Using the dates of the logs, you can easily gather consecutive ones. All logs should be of the same dataset. + start = 'Log_2020-03-22_10-18-56' + end = 'Log_2020-03-22_13-32-51' + + if end < 'Log_2020-03-22_19-30-19': + res_path = 'old_results' + else: + res_path = 'results' + + logs = np.sort([join(res_path, l) for l in listdir(res_path) if start <= l <= end]) + logs = np.insert(logs, 0, join(res_path, 'Log_2020-03-21_19-35-11')) + + # Give names to the logs (for legends) + logs_names = ['no_normals', + 'anisotropic_scale_normals', + 'wrong_scale_normals', + 'only_symmetries_normals(cheating)', + 'test'] + + logs_names = np.array(logs_names[:len(logs)]) + + return logs, logs_names + + +def ModelNet40_radius(): + """ + Test different convolution radius. It was expected that larger radius would means slower networks but better + performances. In fact we do not see much difference (again because of the dataset maybe?) + """ + + # Using the dates of the logs, you can easily gather consecutive ones. All logs should be of the same dataset. + start = 'Log_2020-03-22_13-32-52' + end = 'Log_2020-03-22_19-30-17' + + if end < 'Log_2020-03-22_19-30-19': + res_path = 'old_results' + else: + res_path = 'results' + + logs = np.sort([join(res_path, l) for l in listdir(res_path) if start <= l <= end]) + logs = np.insert(logs, 2, join(res_path, 'Log_2020-03-21_19-35-11')) + + # Give names to the logs (for legends) + logs_names = ['KPe=0.9_r=1.5', + 'KPe=1.2_r=2.0', + 'KPe=1.5_r=2.5', + 'KPe=1.8_r=3.0', + 'KPe=2.1_r=3.5', + 'test'] + + logs_names = np.array(logs_names[:len(logs)]) + + return logs, logs_names + + +def ModelNet40_deform(old_result_limit): + """ + Test deformable convolution with different offset decay. Without modulations 0.01 seems the best. With + modulations 0.1 seems the best. In all cases 1.0 is to much. We need to show deformations for verification. + + It seems that deformations are not really fittig the point cloud. They just reach further away. W need to try on + other datasets and with deformations earlier to see if fitting loss works + """ + + # Using the dates of the logs, you can easily gather consecutive ones. All logs should be of the same dataset. + start = 'Log_2020-03-22_19-30-21' + end = 'Log_2020-03-25_19-30-17' + + if end < old_result_limit: + res_path = 'old_results' + else: + res_path = 'results' + + logs = np.sort([join(res_path, l) for l in listdir(res_path) if start <= l <= end]) + logs = logs.astype(' 'none': Each point in the whole batch has the same contribution. + # > 'class': Each class has the same contribution (points are weighted according to class balance) + # > 'batch': Each cloud in the batch has the same contribution (points are weighted according cloud sizes) + segloss_balance = 'none' + + # Do we nee to save convergence + saving = True + saving_path = None + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Main Call +# \***************/ +# + +if __name__ == '__main__': + + ############################ + # Initialize the environment + ############################ + + # Set which gpu is going to be used + GPU_ID = '3' + + # Set GPU visible device + os.environ['CUDA_VISIBLE_DEVICES'] = GPU_ID + + ############### + # Previous chkp + ############### + + # Choose here if you want to start training from a previous snapshot (None for new training) + #previous_training_path = 'Log_2020-03-19_19-53-27' + previous_training_path = '' + + # Choose index of checkpoint to start from. If None, uses the latest chkp + chkp_idx = None + if previous_training_path: + + # Find all snapshot in the chosen training folder + chkp_path = os.path.join('results', previous_training_path, 'checkpoints') + chkps = [f for f in os.listdir(chkp_path) if f[:4] == 'chkp'] + + # Find which snapshot to restore + if chkp_idx is None: + chosen_chkp = 'current_chkp.tar' + else: + chosen_chkp = np.sort(chkps)[chkp_idx] + chosen_chkp = os.path.join('results', previous_training_path, 'checkpoints', chosen_chkp) + + else: + chosen_chkp = None + + ############## + # Prepare Data + ############## + + print() + print('Data Preparation') + print('****************') + + # Initialize configuration class + config = Modelnet40Config() + if previous_training_path: + config.load(os.path.join('results', previous_training_path)) + config.saving_path = None + + # Get path from argument if given + if len(sys.argv) > 1: + config.saving_path = sys.argv[1] + + # Initialize datasets + training_dataset = ModelNet40Dataset(config, train=True) + test_dataset = ModelNet40Dataset(config, train=False) + + # Initialize samplers + training_sampler = ModelNet40Sampler(training_dataset, balance_labels=True) + test_sampler = ModelNet40Sampler(test_dataset, balance_labels=True) + + # Initialize the dataloader + training_loader = DataLoader(training_dataset, + batch_size=1, + sampler=training_sampler, + collate_fn=ModelNet40Collate, + num_workers=config.input_threads, + pin_memory=True) + test_loader = DataLoader(test_dataset, + batch_size=1, + sampler=test_sampler, + collate_fn=ModelNet40Collate, + num_workers=config.input_threads, + pin_memory=True) + + # Calibrate samplers + training_sampler.calibration(training_loader) + test_sampler.calibration(test_loader) + + #debug_timing(test_dataset, test_sampler, test_loader) + #debug_show_clouds(training_dataset, training_sampler, training_loader) + + print('\nModel Preparation') + print('*****************') + + # Define network model + t1 = time.time() + net = KPCNN(config) + + # Define a trainer class + trainer = ModelTrainer(net, config, chkp_path=chosen_chkp) + print('Done in {:.1f}s\n'.format(time.time() - t1)) + + print('\nStart training') + print('**************') + + # Training + try: + trainer.train(net, training_loader, test_loader, config) + except: + print('Caught an error') + os.kill(os.getpid(), signal.SIGINT) + + print('Forcing exit now') + os.kill(os.getpid(), signal.SIGINT) + + + diff --git a/train_S3DIS.py b/train_S3DIS.py new file mode 100644 index 0000000..210f9d0 --- /dev/null +++ b/train_S3DIS.py @@ -0,0 +1,296 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Callable script to start a training on S3DIS dataset +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 06/03/2020 +# + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Imports and global variables +# \**********************************/ +# + +# Common libs +import signal +import os +import numpy as np +import sys +import torch + +# Dataset +from datasets.S3DIS import * +from torch.utils.data import DataLoader + +from utils.config import Config +from utils.trainer import ModelTrainer +from models.architectures import KPFCNN + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Config Class +# \******************/ +# + +class S3DISConfig(Config): + """ + Override the parameters you want to modify for this dataset + """ + + #################### + # Dataset parameters + #################### + + # Dataset name + dataset = 'S3DIS' + + # Number of classes in the dataset (This value is overwritten by dataset class when Initializating dataset). + num_classes = None + + # Type of task performed on this dataset (also overwritten) + dataset_task = '' + + # Number of CPU threads for the input pipeline + input_threads = 10 + + ######################### + # Architecture definition + ######################### + + # Define layers + architecture = ['simple', + 'resnetb_strided', + 'resnetb', + 'resnetb_strided', + 'resnetb', + 'resnetb_deformable_strided', + 'resnetb_deformable', + 'resnetb_deformable_strided', + 'resnetb_deformable', + 'nearest_upsample', + 'unary', + 'nearest_upsample', + 'unary', + 'nearest_upsample', + 'unary', + 'nearest_upsample', + 'unary'] + + ################### + # KPConv parameters + ################### + + # Radius of the input sphere + in_radius = 2.5 + + # Number of kernel points + num_kernel_points = 15 + + # Size of the first subsampling grid in meter + first_subsampling_dl = 0.04 + + # Radius of convolution in "number grid cell". (2.5 is the standard value) + conv_radius = 2.5 + + # Radius of deformable convolution in "number grid cell". Larger so that deformed kernel can spread out + deform_radius = 6.0 + + # Radius of the area of influence of each kernel point in "number grid cell". (1.0 is the standard value) + KP_extent = 1.5 + + # Behavior of convolutions in ('constant', 'linear', 'gaussian') + KP_influence = 'linear' + + # Aggregation function of KPConv in ('closest', 'sum') + aggregation_mode = 'sum' + + # Choice of input features + in_features_dim = 5 + + # Can the network learn modulations + modulated = True + + # Batch normalization parameters + use_batch_norm = True + batch_norm_momentum = 0.05 + + # Offset loss + # 'permissive' only constrains offsets inside the deform radius (NOT implemented yet) + # 'fitting' helps deformed kernels to adapt to the geometry by penalizing distance to input points + offsets_loss = 'fitting' + offsets_decay = 0.01 + + ##################### + # Training parameters + ##################### + + # Maximal number of epochs + max_epoch = 500 + + # Learning rate management + learning_rate = 1e-2 + momentum = 0.98 + lr_decays = {i: 0.1**(1/100) for i in range(1, max_epoch)} + grad_clip_norm = 100.0 + + # Number of batch + batch_num = 8 + + # Number of steps per epochs + epoch_steps = 500 + + # Number of validation examples per epoch + validation_size = 30 + + # Number of epoch between each checkpoint + checkpoint_gap = 50 + + # Augmentations + augment_scale_anisotropic = True + augment_symmetries = [True, False, False] + augment_rotation = 'vertical' + augment_scale_min = 0.9 + augment_scale_max = 1.1 + augment_noise = 0.001 + augment_color = 0.9 + + # The way we balance segmentation loss TODO: implement and test 'class' and 'batch' modes + # > 'none': Each point in the whole batch has the same contribution. + # > 'class': Each class has the same contribution (points are weighted according to class balance) + # > 'batch': Each cloud in the batch has the same contribution (points are weighted according cloud sizes) + segloss_balance = 'none' + + # Do we nee to save convergence + saving = True + saving_path = None + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Main Call +# \***************/ +# + +if __name__ == '__main__': + + ############################ + # Initialize the environment + ############################ + + # Set which gpu is going to be used + GPU_ID = '1' + + # Set GPU visible device + os.environ['CUDA_VISIBLE_DEVICES'] = GPU_ID + + ############### + # Previous chkp + ############### + + # Choose here if you want to start training from a previous snapshot (None for new training) + #previous_training_path = 'Log_2020-03-19_19-53-27' + previous_training_path = '' + + # Choose index of checkpoint to start from. If None, uses the latest chkp + chkp_idx = None + if previous_training_path: + + # Find all snapshot in the chosen training folder + chkp_path = os.path.join('results', previous_training_path, 'checkpoints') + chkps = [f for f in os.listdir(chkp_path) if f[:4] == 'chkp'] + + # Find which snapshot to restore + if chkp_idx is None: + chosen_chkp = 'current_chkp.tar' + else: + chosen_chkp = np.sort(chkps)[chkp_idx] + chosen_chkp = os.path.join('results', previous_training_path, 'checkpoints', chosen_chkp) + + else: + chosen_chkp = None + + ############## + # Prepare Data + ############## + + print() + print('Data Preparation') + print('****************') + + # Initialize configuration class + config = S3DISConfig() + if previous_training_path: + config.load(os.path.join('results', previous_training_path)) + config.saving_path = None + + # Get path from argument if given + if len(sys.argv) > 1: + config.saving_path = sys.argv[1] + + # Initialize datasets + training_dataset = S3DISDataset(config, set='training') + test_dataset = S3DISDataset(config, set='validation') + + # Initialize samplers + training_sampler = S3DISSampler(training_dataset) + test_sampler = S3DISSampler(test_dataset) + + # Initialize the dataloader + training_loader = DataLoader(training_dataset, + batch_size=1, + sampler=training_sampler, + collate_fn=S3DISCollate, + num_workers=config.input_threads, + pin_memory=True) + test_loader = DataLoader(test_dataset, + batch_size=1, + sampler=test_sampler, + collate_fn=S3DISCollate, + num_workers=config.input_threads, + pin_memory=True) + + # Calibrate samplers + training_sampler.calibration(training_loader, verbose=True) + test_sampler.calibration(test_loader, verbose=True) + + #debug_timing(training_dataset, training_sampler, training_loader) + #debug_timing(test_dataset, test_sampler, test_loader) + #debug_show_clouds(training_dataset, training_sampler, training_loader) + + print('\nModel Preparation') + print('*****************') + + # Define network model + t1 = time.time() + net = KPFCNN(config) + + # Define a trainer class + trainer = ModelTrainer(net, config, chkp_path=chosen_chkp) + print('Done in {:.1f}s\n'.format(time.time() - t1)) + + print('\nStart training') + print('**************') + + # Training + try: + trainer.train(net, training_loader, test_loader, config) + except: + print('Caught an error') + os.kill(os.getpid(), signal.SIGINT) + + print('Forcing exit now') + os.kill(os.getpid(), signal.SIGINT) + + + diff --git a/utils/config.py b/utils/config.py new file mode 100644 index 0000000..934dd2c --- /dev/null +++ b/utils/config.py @@ -0,0 +1,363 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Configuration class +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 11/06/2018 +# + + +from os.path import join +import numpy as np + + +# Colors for printing +class bcolors: + HEADER = '\033[95m' + OKBLUE = '\033[94m' + OKGREEN = '\033[92m' + WARNING = '\033[93m' + FAIL = '\033[91m' + ENDC = '\033[0m' + BOLD = '\033[1m' + UNDERLINE = '\033[4m' + + +class Config: + """ + Class containing the parameters you want to modify for this dataset + """ + + ################## + # Input parameters + ################## + + # Dataset name + dataset = '' + + # Type of network model + dataset_task = '' + + # Number of classes in the dataset + num_classes = 0 + + # Dimension of input points + in_points_dim = 3 + + # Dimension of input features + in_features_dim = 1 + + # Radius of the input sphere (ignored for models, only used for point clouds) + in_radius = 1.0 + + # Number of CPU threads for the input pipeline + input_threads = 8 + + ################## + # Model parameters + ################## + + # Architecture definition. List of blocks + architecture = [] + + # Decide the mode of equivariance and invariance + equivar_mode = '' + invar_mode = '' + + # Dimension of the first feature maps + first_features_dim = 64 + + # Batch normalization parameters + use_batch_norm = True + batch_norm_momentum = 0.99 + + # For segmentation models : ratio between the segmented area and the input area + segmentation_ratio = 1.0 + + ################### + # KPConv parameters + ################### + + # Number of kernel points + num_kernel_points = 15 + + # Size of the first subsampling grid in meter + first_subsampling_dl = 0.02 + + # Radius of convolution in "number grid cell". (2.5 is the standard value) + conv_radius = 2.5 + + # Radius of deformable convolution in "number grid cell". Larger so that deformed kernel can spread out + deform_radius = 5.0 + + # Kernel point influence radius + KP_extent = 1.0 + + # Influence function when d < KP_extent. ('constant', 'linear', 'gaussian') When d > KP_extent, always zero + KP_influence = 'linear' + + # Aggregation function of KPConv in ('closest', 'sum') + # Decide if you sum all kernel point influences, or if you only take the influence of the closest KP + aggregation_mode = 'sum' + + # Fixed points in the kernel : 'none', 'center' or 'verticals' + fixed_kernel_points = 'center' + + # Use modulateion in deformable convolutions + modulated = False + + # For SLAM datasets like SemanticKitti number of frames used (minimum one) + n_frames = 1 + + # For SLAM datasets like SemanticKitti max number of point in input cloud + max_in_points = 0 + + ##################### + # Training parameters + ##################### + + # Network optimizer parameters (learning rate and momentum) + learning_rate = 1e-3 + momentum = 0.9 + + # Learning rate decays. Dictionary of all decay values with their epoch {epoch: decay}. + lr_decays = {200: 0.2, 300: 0.2} + + # Gradient clipping value (negative means no clipping) + grad_clip_norm = 100.0 + + # Augmentation parameters + augment_scale_anisotropic = True + augment_scale_min = 0.9 + augment_scale_max = 1.1 + augment_symmetries = [False, False, False] + augment_rotation = 'vertical' + augment_noise = 0.005 + augment_color = 0.7 + + # Augment with occlusions (not implemented yet) + augment_occlusion = 'planar' + augment_occlusion_ratio = 0.2 + augment_occlusion_num = 1 + + # Regularization loss importance + weight_decay = 1e-3 + + # The way we balance segmentation loss + # > 'none': Each point in the whole batch has the same contribution. + # > 'class': Each class has the same contribution (points are weighted according to class balance) + # > 'batch': Each cloud in the batch has the same contribution (points are weighted according cloud sizes) + segloss_balance = 'none' + + # Offset regularization loss + offsets_loss = 'permissive' + offsets_decay = 1e-2 + + # Number of batch + batch_num = 10 + + # Maximal number of epochs + max_epoch = 1000 + + # Number of steps per epochs + epoch_steps = 1000 + + # Number of validation examples per epoch + validation_size = 100 + + # Number of epoch between each checkpoint + checkpoint_gap = 50 + + # Do we nee to save convergence + saving = True + saving_path = None + + def __init__(self): + """ + Class Initialyser + """ + + # Number of layers + self.num_layers = len([block for block in self.architecture if 'pool' in block or 'strided' in block]) + 1 + + ################### + # Deform layer list + ################### + # + # List of boolean indicating which layer has a deformable convolution + # + + layer_blocks = [] + self.deform_layers = [] + arch = self.architecture + for block_i, block in enumerate(arch): + + # Get all blocks of the layer + if not ('pool' in block or 'strided' in block or 'global' in block or 'upsample' in block): + layer_blocks += [block] + continue + + # Convolution neighbors indices + # ***************************** + + deform_layer = False + if layer_blocks: + if np.any(['deformable' in blck for blck in layer_blocks]): + deform_layer = True + + if 'pool' in block or 'strided' in block: + if 'deformable' in block: + deform_layer = True + + self.deform_layers += [deform_layer] + layer_blocks = [] + + # Stop when meeting a global pooling or upsampling + if 'global' in block or 'upsample' in block: + break + + def load(self, path): + + filename = join(path, 'parameters.txt') + with open(filename, 'r') as f: + lines = f.readlines() + + # Class variable dictionary + for line in lines: + line_info = line.split() + if len(line_info) > 1 and line_info[0] != '#': + + if line_info[2] == 'None': + setattr(self, line_info[0], None) + + elif line_info[0] == 'lr_decay_epochs': + self.lr_decays = {int(b.split(':')[0]): float(b.split(':')[1]) for b in line_info[2:]} + + elif line_info[0] == 'architecture': + self.architecture = [b for b in line_info[2:]] + + elif line_info[0] == 'augment_symmetries': + self.augment_symmetries = [bool(int(b)) for b in line_info[2:]] + + elif line_info[0] == 'num_classes': + if len(line_info) > 3: + self.num_classes = [int(c) for c in line_info[2:]] + else: + self.num_classes = int(line_info[2]) + + else: + attr_type = type(getattr(self, line_info[0])) + if attr_type == bool: + setattr(self, line_info[0], attr_type(int(line_info[2]))) + else: + setattr(self, line_info[0], attr_type(line_info[2])) + + self.saving = True + self.saving_path = path + self.__init__() + + def save(self): + + with open(join(self.saving_path, 'parameters.txt'), "w") as text_file: + + text_file.write('# -----------------------------------#\n') + text_file.write('# Parameters of the training session #\n') + text_file.write('# -----------------------------------#\n\n') + + # Input parameters + text_file.write('# Input parameters\n') + text_file.write('# ****************\n\n') + text_file.write('dataset = {:s}\n'.format(self.dataset)) + text_file.write('dataset_task = {:s}\n'.format(self.dataset_task)) + if type(self.num_classes) is list: + text_file.write('num_classes =') + for n in self.num_classes: + text_file.write(' {:d}'.format(n)) + text_file.write('\n') + else: + text_file.write('num_classes = {:d}\n'.format(self.num_classes)) + text_file.write('in_points_dim = {:d}\n'.format(self.in_points_dim)) + text_file.write('in_features_dim = {:d}\n'.format(self.in_features_dim)) + text_file.write('in_radius = {:.3f}\n'.format(self.in_radius)) + text_file.write('input_threads = {:d}\n\n'.format(self.input_threads)) + + # Model parameters + text_file.write('# Model parameters\n') + text_file.write('# ****************\n\n') + + text_file.write('architecture =') + for a in self.architecture: + text_file.write(' {:s}'.format(a)) + text_file.write('\n') + text_file.write('equivar_mode = {:s}\n'.format(self.equivar_mode)) + text_file.write('invar_mode = {:s}\n'.format(self.invar_mode)) + text_file.write('num_layers = {:d}\n'.format(self.num_layers)) + text_file.write('first_features_dim = {:d}\n'.format(self.first_features_dim)) + text_file.write('use_batch_norm = {:d}\n'.format(int(self.use_batch_norm))) + text_file.write('batch_norm_momentum = {:.3f}\n\n'.format(self.batch_norm_momentum)) + text_file.write('segmentation_ratio = {:.3f}\n\n'.format(self.segmentation_ratio)) + + # KPConv parameters + text_file.write('# KPConv parameters\n') + text_file.write('# *****************\n\n') + + text_file.write('first_subsampling_dl = {:.3f}\n'.format(self.first_subsampling_dl)) + text_file.write('num_kernel_points = {:d}\n'.format(self.num_kernel_points)) + text_file.write('conv_radius = {:.3f}\n'.format(self.conv_radius)) + text_file.write('deform_radius = {:.3f}\n'.format(self.deform_radius)) + text_file.write('fixed_kernel_points = {:s}\n'.format(self.fixed_kernel_points)) + text_file.write('KP_extent = {:.3f}\n'.format(self.KP_extent)) + text_file.write('KP_influence = {:s}\n'.format(self.KP_influence)) + text_file.write('aggregation_mode = {:s}\n'.format(self.aggregation_mode)) + text_file.write('modulated = {:d}\n'.format(int(self.modulated))) + text_file.write('n_frames = {:d}\n'.format(self.n_frames)) + text_file.write('max_in_points = {:d}\n\n'.format(self.max_in_points)) + + # Training parameters + text_file.write('# Training parameters\n') + text_file.write('# *******************\n\n') + + text_file.write('learning_rate = {:f}\n'.format(self.learning_rate)) + text_file.write('momentum = {:f}\n'.format(self.momentum)) + text_file.write('lr_decay_epochs =') + for e, d in self.lr_decays.items(): + text_file.write(' {:d}:{:f}'.format(e, d)) + text_file.write('\n') + text_file.write('grad_clip_norm = {:f}\n\n'.format(self.grad_clip_norm)) + + + text_file.write('augment_symmetries =') + for a in self.augment_symmetries: + text_file.write(' {:d}'.format(int(a))) + text_file.write('\n') + text_file.write('augment_rotation = {:s}\n'.format(self.augment_rotation)) + text_file.write('augment_noise = {:f}\n'.format(self.augment_noise)) + text_file.write('augment_occlusion = {:s}\n'.format(self.augment_occlusion)) + text_file.write('augment_occlusion_ratio = {:.3f}\n'.format(self.augment_occlusion_ratio)) + text_file.write('augment_occlusion_num = {:d}\n'.format(self.augment_occlusion_num)) + text_file.write('augment_scale_anisotropic = {:d}\n'.format(int(self.augment_scale_anisotropic))) + text_file.write('augment_scale_min = {:.3f}\n'.format(self.augment_scale_min)) + text_file.write('augment_scale_max = {:.3f}\n'.format(self.augment_scale_max)) + text_file.write('augment_color = {:.3f}\n\n'.format(self.augment_color)) + + text_file.write('weight_decay = {:f}\n'.format(self.weight_decay)) + text_file.write('segloss_balance = {:s}\n'.format(self.segloss_balance)) + text_file.write('offsets_loss = {:s}\n'.format(self.offsets_loss)) + text_file.write('offsets_decay = {:f}\n'.format(self.offsets_decay)) + text_file.write('batch_num = {:d}\n'.format(self.batch_num)) + text_file.write('max_epoch = {:d}\n'.format(self.max_epoch)) + if self.epoch_steps is None: + text_file.write('epoch_steps = None\n') + else: + text_file.write('epoch_steps = {:d}\n'.format(self.epoch_steps)) + text_file.write('validation_size = {:d}\n'.format(self.validation_size)) + text_file.write('checkpoint_gap = {:d}\n'.format(self.checkpoint_gap)) + diff --git a/utils/mayavi_visu.py b/utils/mayavi_visu.py new file mode 100644 index 0000000..b1c3821 --- /dev/null +++ b/utils/mayavi_visu.py @@ -0,0 +1,436 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Script for various visualization with mayavi +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 11/06/2018 +# + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Imports and global variables +# \**********************************/ +# + + +# Basic libs +import torch +import numpy as np +from sklearn.neighbors import KDTree +from os import makedirs, remove, rename, listdir +from os.path import exists, join +import time + +import sys + +# PLY reader +from utils.ply import write_ply, read_ply + +# Configuration class +from utils.config import Config + + +def show_ModelNet_models(all_points): + from mayavi import mlab + + ########################### + # Interactive visualization + ########################### + + # Create figure for features + fig1 = mlab.figure('Models', bgcolor=(1, 1, 1), size=(1000, 800)) + fig1.scene.parallel_projection = False + + # Indices + global file_i + file_i = 0 + + def update_scene(): + + # clear figure + mlab.clf(fig1) + + # Plot new data feature + points = all_points[file_i] + + # Rescale points for visu + points = (points * 1.5 + np.array([1.0, 1.0, 1.0])) * 50.0 + + # Show point clouds colorized with activations + activations = mlab.points3d(points[:, 0], + points[:, 1], + points[:, 2], + points[:, 2], + scale_factor=3.0, + scale_mode='none', + figure=fig1) + + # New title + mlab.title(str(file_i), color=(0, 0, 0), size=0.3, height=0.01) + text = '<--- (press g for previous)' + 50 * ' ' + '(press h for next) --->' + mlab.text(0.01, 0.01, text, color=(0, 0, 0), width=0.98) + mlab.orientation_axes() + + return + + def keyboard_callback(vtk_obj, event): + global file_i + + if vtk_obj.GetKeyCode() in ['g', 'G']: + + file_i = (file_i - 1) % len(all_points) + update_scene() + + elif vtk_obj.GetKeyCode() in ['h', 'H']: + + file_i = (file_i + 1) % len(all_points) + update_scene() + + return + + # Draw a first plot + update_scene() + fig1.scene.interactor.add_observer('KeyPressEvent', keyboard_callback) + mlab.show() + + +def show_ModelNet_examples(clouds, cloud_normals=None, cloud_labels=None): + from mayavi import mlab + + ########################### + # Interactive visualization + ########################### + + # Create figure for features + fig1 = mlab.figure('Models', bgcolor=(1, 1, 1), size=(1000, 800)) + fig1.scene.parallel_projection = False + + if cloud_labels is None: + cloud_labels = [points[:, 2] for points in clouds] + + # Indices + global file_i, show_normals + file_i = 0 + show_normals = True + + def update_scene(): + + # clear figure + mlab.clf(fig1) + + # Plot new data feature + points = clouds[file_i] + labels = cloud_labels[file_i] + if cloud_normals is not None: + normals = cloud_normals[file_i] + else: + normals = None + + # Rescale points for visu + points = (points * 1.5 + np.array([1.0, 1.0, 1.0])) * 50.0 + + # Show point clouds colorized with activations + activations = mlab.points3d(points[:, 0], + points[:, 1], + points[:, 2], + labels, + scale_factor=3.0, + scale_mode='none', + figure=fig1) + if normals is not None and show_normals: + activations = mlab.quiver3d(points[:, 0], + points[:, 1], + points[:, 2], + normals[:, 0], + normals[:, 1], + normals[:, 2], + scale_factor=10.0, + scale_mode='none', + figure=fig1) + + # New title + mlab.title(str(file_i), color=(0, 0, 0), size=0.3, height=0.01) + text = '<--- (press g for previous)' + 50 * ' ' + '(press h for next) --->' + mlab.text(0.01, 0.01, text, color=(0, 0, 0), width=0.98) + mlab.orientation_axes() + + return + + def keyboard_callback(vtk_obj, event): + global file_i, show_normals + + if vtk_obj.GetKeyCode() in ['g', 'G']: + file_i = (file_i - 1) % len(clouds) + update_scene() + + elif vtk_obj.GetKeyCode() in ['h', 'H']: + file_i = (file_i + 1) % len(clouds) + update_scene() + + elif vtk_obj.GetKeyCode() in ['n', 'N']: + show_normals = not show_normals + update_scene() + + return + + # Draw a first plot + update_scene() + fig1.scene.interactor.add_observer('KeyPressEvent', keyboard_callback) + mlab.show() + + +def show_neighbors(query, supports, neighbors): + from mayavi import mlab + + ########################### + # Interactive visualization + ########################### + + # Create figure for features + fig1 = mlab.figure('Models', bgcolor=(1, 1, 1), size=(1000, 800)) + fig1.scene.parallel_projection = False + + # Indices + global file_i + file_i = 0 + + def update_scene(): + + # clear figure + mlab.clf(fig1) + + # Rescale points for visu + p1 = (query * 1.5 + np.array([1.0, 1.0, 1.0])) * 50.0 + p2 = (supports * 1.5 + np.array([1.0, 1.0, 1.0])) * 50.0 + + l1 = p1[:, 2]*0 + l1[file_i] = 1 + + l2 = p2[:, 2]*0 + 2 + l2[neighbors[file_i]] = 3 + + # Show point clouds colorized with activations + activations = mlab.points3d(p1[:, 0], + p1[:, 1], + p1[:, 2], + l1, + scale_factor=2.0, + scale_mode='none', + vmin=0.0, + vmax=3.0, + figure=fig1) + + activations = mlab.points3d(p2[:, 0], + p2[:, 1], + p2[:, 2], + l2, + scale_factor=3.0, + scale_mode='none', + vmin=0.0, + vmax=3.0, + figure=fig1) + + # New title + mlab.title(str(file_i), color=(0, 0, 0), size=0.3, height=0.01) + text = '<--- (press g for previous)' + 50 * ' ' + '(press h for next) --->' + mlab.text(0.01, 0.01, text, color=(0, 0, 0), width=0.98) + mlab.orientation_axes() + + return + + def keyboard_callback(vtk_obj, event): + global file_i + + if vtk_obj.GetKeyCode() in ['g', 'G']: + + file_i = (file_i - 1) % len(query) + update_scene() + + elif vtk_obj.GetKeyCode() in ['h', 'H']: + + file_i = (file_i + 1) % len(query) + update_scene() + + return + + # Draw a first plot + update_scene() + fig1.scene.interactor.add_observer('KeyPressEvent', keyboard_callback) + mlab.show() + + +def show_input_batch(batch): + from mayavi import mlab + + ########################### + # Interactive visualization + ########################### + + # Create figure for features + fig1 = mlab.figure('Input', bgcolor=(1, 1, 1), size=(1000, 800)) + fig1.scene.parallel_projection = False + + # Unstack batch + all_points = batch.unstack_points() + all_neighbors = batch.unstack_neighbors() + all_pools = batch.unstack_pools() + + # Indices + global b_i, l_i, neighb_i, show_pools + b_i = 0 + l_i = 0 + neighb_i = 0 + show_pools = False + + def update_scene(): + + # clear figure + mlab.clf(fig1) + + # Rescale points for visu + p = (all_points[l_i][b_i] * 1.5 + np.array([1.0, 1.0, 1.0])) * 50.0 + labels = p[:, 2]*0 + + if show_pools: + p2 = (all_points[l_i+1][b_i][neighb_i:neighb_i+1] * 1.5 + np.array([1.0, 1.0, 1.0])) * 50.0 + p = np.vstack((p, p2)) + labels = np.hstack((labels, np.ones((1,), dtype=np.int32)*3)) + pool_inds = all_pools[l_i][b_i][neighb_i] + pool_inds = pool_inds[pool_inds >= 0] + labels[pool_inds] = 2 + else: + neighb_inds = all_neighbors[l_i][b_i][neighb_i] + neighb_inds = neighb_inds[neighb_inds >= 0] + labels[neighb_inds] = 2 + labels[neighb_i] = 3 + + # Show point clouds colorized with activations + mlab.points3d(p[:, 0], + p[:, 1], + p[:, 2], + labels, + scale_factor=2.0, + scale_mode='none', + vmin=0.0, + vmax=3.0, + figure=fig1) + + + """ + mlab.points3d(p[-2:, 0], + p[-2:, 1], + p[-2:, 2], + labels[-2:]*0 + 3, + scale_factor=0.16 * 1.5 * 50, + scale_mode='none', + mode='cube', + vmin=0.0, + vmax=3.0, + figure=fig1) + mlab.points3d(p[-1:, 0], + p[-1:, 1], + p[-1:, 2], + labels[-1:]*0 + 2, + scale_factor=0.16 * 2 * 2.5 * 1.5 * 50, + scale_mode='none', + mode='sphere', + vmin=0.0, + vmax=3.0, + figure=fig1) + + """ + + # New title + title_str = '<([) b_i={:d} (])> <(,) l_i={:d} (.)> <(N) n_i={:d} (M)>'.format(b_i, l_i, neighb_i) + mlab.title(title_str, color=(0, 0, 0), size=0.3, height=0.90) + if show_pools: + text = 'pools (switch with G)' + else: + text = 'neighbors (switch with G)' + mlab.text(0.01, 0.01, text, color=(0, 0, 0), width=0.3) + mlab.orientation_axes() + + return + + def keyboard_callback(vtk_obj, event): + global b_i, l_i, neighb_i, show_pools + + if vtk_obj.GetKeyCode() in ['[', '{']: + b_i = (b_i - 1) % len(all_points[l_i]) + neighb_i = 0 + update_scene() + + elif vtk_obj.GetKeyCode() in [']', '}']: + b_i = (b_i + 1) % len(all_points[l_i]) + neighb_i = 0 + update_scene() + + elif vtk_obj.GetKeyCode() in [',', '<']: + if show_pools: + l_i = (l_i - 1) % (len(all_points) - 1) + else: + l_i = (l_i - 1) % len(all_points) + neighb_i = 0 + update_scene() + + elif vtk_obj.GetKeyCode() in ['.', '>']: + if show_pools: + l_i = (l_i + 1) % (len(all_points) - 1) + else: + l_i = (l_i + 1) % len(all_points) + neighb_i = 0 + update_scene() + + elif vtk_obj.GetKeyCode() in ['n', 'N']: + neighb_i = (neighb_i - 1) % all_points[l_i][b_i].shape[0] + update_scene() + + elif vtk_obj.GetKeyCode() in ['m', 'M']: + neighb_i = (neighb_i + 1) % all_points[l_i][b_i].shape[0] + update_scene() + + elif vtk_obj.GetKeyCode() in ['g', 'G']: + if l_i < len(all_points) - 1: + show_pools = not show_pools + neighb_i = 0 + update_scene() + + return + + # Draw a first plot + update_scene() + fig1.scene.interactor.add_observer('KeyPressEvent', keyboard_callback) + mlab.show() + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/utils/metrics.py b/utils/metrics.py new file mode 100644 index 0000000..22e6bfb --- /dev/null +++ b/utils/metrics.py @@ -0,0 +1,230 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Metric utility functions +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 11/06/2018 +# + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Imports and global variables +# \**********************************/ +# + + +# Basic libs +import numpy as np + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Utilities +# \***************/ +# + +def fast_confusion(true, pred, label_values=None): + """ + Fast confusion matrix (100x faster than Scikit learn). But only works if labels are la + :param true: + :param false: + :param num_classes: + :return: + """ + + # Ensure data is in the right format + true = np.squeeze(true) + pred = np.squeeze(pred) + if len(true.shape) != 1: + raise ValueError('Truth values are stored in a {:d}D array instead of 1D array'. format(len(true.shape))) + if len(pred.shape) != 1: + raise ValueError('Prediction values are stored in a {:d}D array instead of 1D array'. format(len(pred.shape))) + if true.dtype not in [np.int32, np.int64]: + raise ValueError('Truth values are {:s} instead of int32 or int64'.format(true.dtype)) + if pred.dtype not in [np.int32, np.int64]: + raise ValueError('Prediction values are {:s} instead of int32 or int64'.format(pred.dtype)) + true = true.astype(np.int32) + pred = pred.astype(np.int32) + + # Get the label values + if label_values is None: + # From data if they are not given + label_values = np.unique(np.hstack((true, pred))) + else: + # Ensure they are good if given + if label_values.dtype not in [np.int32, np.int64]: + raise ValueError('label values are {:s} instead of int32 or int64'.format(label_values.dtype)) + if len(np.unique(label_values)) < len(label_values): + raise ValueError('Given labels are not unique') + + # Sort labels + label_values = np.sort(label_values) + + # Get the number of classes + num_classes = len(label_values) + + #print(num_classes) + #print(label_values) + #print(np.max(true)) + #print(np.max(pred)) + #print(np.max(true * num_classes + pred)) + + # Start confusion computations + if label_values[0] == 0 and label_values[-1] == num_classes - 1: + + # Vectorized confusion + vec_conf = np.bincount(true * num_classes + pred) + + # Add possible missing values due to classes not being in pred or true + #print(vec_conf.shape) + if vec_conf.shape[0] < num_classes ** 2: + vec_conf = np.pad(vec_conf, (0, num_classes ** 2 - vec_conf.shape[0]), 'constant') + #print(vec_conf.shape) + + # Reshape confusion in a matrix + return vec_conf.reshape((num_classes, num_classes)) + + + else: + + # Ensure no negative classes + if label_values[0] < 0: + raise ValueError('Unsupported negative classes') + + # Get the data in [0,num_classes[ + label_map = np.zeros((label_values[-1] + 1,), dtype=np.int32) + for k, v in enumerate(label_values): + label_map[v] = k + + pred = label_map[pred] + true = label_map[true] + + # Vectorized confusion + vec_conf = np.bincount(true * num_classes + pred) + + # Add possible missing values due to classes not being in pred or true + if vec_conf.shape[0] < num_classes ** 2: + vec_conf = np.pad(vec_conf, (0, num_classes ** 2 - vec_conf.shape[0]), 'constant') + + # Reshape confusion in a matrix + return vec_conf.reshape((num_classes, num_classes)) + +def metrics(confusions, ignore_unclassified=False): + """ + Computes different metrics from confusion matrices. + :param confusions: ([..., n_c, n_c] np.int32). Can be any dimension, the confusion matrices should be described by + the last axes. n_c = number of classes + :param ignore_unclassified: (bool). True if the the first class should be ignored in the results + :return: ([..., n_c] np.float32) precision, recall, F1 score, IoU score + """ + + # If the first class (often "unclassified") should be ignored, erase it from the confusion. + if (ignore_unclassified): + confusions[..., 0, :] = 0 + confusions[..., :, 0] = 0 + + # Compute TP, FP, FN. This assume that the second to last axis counts the truths (like the first axis of a + # confusion matrix), and that the last axis counts the predictions (like the second axis of a confusion matrix) + TP = np.diagonal(confusions, axis1=-2, axis2=-1) + TP_plus_FP = np.sum(confusions, axis=-1) + TP_plus_FN = np.sum(confusions, axis=-2) + + # Compute precision and recall. This assume that the second to last axis counts the truths (like the first axis of + # a confusion matrix), and that the last axis counts the predictions (like the second axis of a confusion matrix) + PRE = TP / (TP_plus_FN + 1e-6) + REC = TP / (TP_plus_FP + 1e-6) + + # Compute Accuracy + ACC = np.sum(TP, axis=-1) / (np.sum(confusions, axis=(-2, -1)) + 1e-6) + + # Compute F1 score + F1 = 2 * TP / (TP_plus_FP + TP_plus_FN + 1e-6) + + # Compute IoU + IoU = F1 / (2 - F1) + + return PRE, REC, F1, IoU, ACC + + +def smooth_metrics(confusions, smooth_n=0, ignore_unclassified=False): + """ + Computes different metrics from confusion matrices. Smoothed over a number of epochs. + :param confusions: ([..., n_c, n_c] np.int32). Can be any dimension, the confusion matrices should be described by + the last axes. n_c = number of classes + :param smooth_n: (int). smooth extent + :param ignore_unclassified: (bool). True if the the first class should be ignored in the results + :return: ([..., n_c] np.float32) precision, recall, F1 score, IoU score + """ + + # If the first class (often "unclassified") should be ignored, erase it from the confusion. + if ignore_unclassified: + confusions[..., 0, :] = 0 + confusions[..., :, 0] = 0 + + # Sum successive confusions for smoothing + smoothed_confusions = confusions.copy() + if confusions.ndim > 2 and smooth_n > 0: + for epoch in range(confusions.shape[-3]): + i0 = max(epoch - smooth_n, 0) + i1 = min(epoch + smooth_n + 1, confusions.shape[-3]) + smoothed_confusions[..., epoch, :, :] = np.sum(confusions[..., i0:i1, :, :], axis=-3) + + # Compute TP, FP, FN. This assume that the second to last axis counts the truths (like the first axis of a + # confusion matrix), and that the last axis counts the predictions (like the second axis of a confusion matrix) + TP = np.diagonal(smoothed_confusions, axis1=-2, axis2=-1) + TP_plus_FP = np.sum(smoothed_confusions, axis=-2) + TP_plus_FN = np.sum(smoothed_confusions, axis=-1) + + # Compute precision and recall. This assume that the second to last axis counts the truths (like the first axis of + # a confusion matrix), and that the last axis counts the predictions (like the second axis of a confusion matrix) + PRE = TP / (TP_plus_FN + 1e-6) + REC = TP / (TP_plus_FP + 1e-6) + + # Compute Accuracy + ACC = np.sum(TP, axis=-1) / (np.sum(smoothed_confusions, axis=(-2, -1)) + 1e-6) + + # Compute F1 score + F1 = 2 * TP / (TP_plus_FP + TP_plus_FN + 1e-6) + + # Compute IoU + IoU = F1 / (2 - F1) + + return PRE, REC, F1, IoU, ACC + + +def IoU_from_confusions(confusions): + """ + Computes IoU from confusion matrices. + :param confusions: ([..., n_c, n_c] np.int32). Can be any dimension, the confusion matrices should be described by + the last axes. n_c = number of classes + :param ignore_unclassified: (bool). True if the the first class should be ignored in the results + :return: ([..., n_c] np.float32) IoU score + """ + + # Compute TP, FP, FN. This assume that the second to last axis counts the truths (like the first axis of a + # confusion matrix), and that the last axis counts the predictions (like the second axis of a confusion matrix) + TP = np.diagonal(confusions, axis1=-2, axis2=-1) + TP_plus_FN = np.sum(confusions, axis=-1) + TP_plus_FP = np.sum(confusions, axis=-2) + + # Compute IoU + IoU = TP / (TP_plus_FP + TP_plus_FN - TP + 1e-6) + + # Compute mIoU with only the actual classes + mask = TP_plus_FN < 1e-3 + counts = np.sum(1 - mask, axis=-1, keepdims=True) + mIoU = np.sum(IoU, axis=-1, keepdims=True) / (counts + 1e-6) + + # If class is absent, place mIoU in place of 0 IoU to get the actual mean later + IoU += mask * mIoU + + return IoU diff --git a/utils/ply.py b/utils/ply.py new file mode 100644 index 0000000..0f5bfd3 --- /dev/null +++ b/utils/ply.py @@ -0,0 +1,355 @@ +# +# +# 0===============================0 +# | PLY files reader/writer | +# 0===============================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# function to read/write .ply files +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 10/02/2017 +# + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Imports and global variables +# \**********************************/ +# + + +# Basic libs +import numpy as np +import sys + + +# Define PLY types +ply_dtypes = dict([ + (b'int8', 'i1'), + (b'char', 'i1'), + (b'uint8', 'u1'), + (b'uchar', 'u1'), + (b'int16', 'i2'), + (b'short', 'i2'), + (b'uint16', 'u2'), + (b'ushort', 'u2'), + (b'int32', 'i4'), + (b'int', 'i4'), + (b'uint32', 'u4'), + (b'uint', 'u4'), + (b'float32', 'f4'), + (b'float', 'f4'), + (b'float64', 'f8'), + (b'double', 'f8') +]) + +# Numpy reader format +valid_formats = {'ascii': '', 'binary_big_endian': '>', + 'binary_little_endian': '<'} + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Functions +# \***************/ +# + + +def parse_header(plyfile, ext): + # Variables + line = [] + properties = [] + num_points = None + + while b'end_header' not in line and line != b'': + line = plyfile.readline() + + if b'element' in line: + line = line.split() + num_points = int(line[2]) + + elif b'property' in line: + line = line.split() + properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) + + return num_points, properties + + +def parse_mesh_header(plyfile, ext): + # Variables + line = [] + vertex_properties = [] + num_points = None + num_faces = None + current_element = None + + + while b'end_header' not in line and line != b'': + line = plyfile.readline() + + # Find point element + if b'element vertex' in line: + current_element = 'vertex' + line = line.split() + num_points = int(line[2]) + + elif b'element face' in line: + current_element = 'face' + line = line.split() + num_faces = int(line[2]) + + elif b'property' in line: + if current_element == 'vertex': + line = line.split() + vertex_properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) + elif current_element == 'vertex': + if not line.startswith('property list uchar int'): + raise ValueError('Unsupported faces property : ' + line) + + return num_points, num_faces, vertex_properties + + +def read_ply(filename, triangular_mesh=False): + """ + Read ".ply" files + + Parameters + ---------- + filename : string + the name of the file to read. + + Returns + ------- + result : array + data stored in the file + + Examples + -------- + Store data in file + + >>> points = np.random.rand(5, 3) + >>> values = np.random.randint(2, size=10) + >>> write_ply('example.ply', [points, values], ['x', 'y', 'z', 'values']) + + Read the file + + >>> data = read_ply('example.ply') + >>> values = data['values'] + array([0, 0, 1, 1, 0]) + + >>> points = np.vstack((data['x'], data['y'], data['z'])).T + array([[ 0.466 0.595 0.324] + [ 0.538 0.407 0.654] + [ 0.850 0.018 0.988] + [ 0.395 0.394 0.363] + [ 0.873 0.996 0.092]]) + + """ + + with open(filename, 'rb') as plyfile: + + + # Check if the file start with ply + if b'ply' not in plyfile.readline(): + raise ValueError('The file does not start whith the word ply') + + # get binary_little/big or ascii + fmt = plyfile.readline().split()[1].decode() + if fmt == "ascii": + raise ValueError('The file is not binary') + + # get extension for building the numpy dtypes + ext = valid_formats[fmt] + + # PointCloud reader vs mesh reader + if triangular_mesh: + + # Parse header + num_points, num_faces, properties = parse_mesh_header(plyfile, ext) + + # Get point data + vertex_data = np.fromfile(plyfile, dtype=properties, count=num_points) + + # Get face data + face_properties = [('k', ext + 'u1'), + ('v1', ext + 'i4'), + ('v2', ext + 'i4'), + ('v3', ext + 'i4')] + faces_data = np.fromfile(plyfile, dtype=face_properties, count=num_faces) + + # Return vertex data and concatenated faces + faces = np.vstack((faces_data['v1'], faces_data['v2'], faces_data['v3'])).T + data = [vertex_data, faces] + + else: + + # Parse header + num_points, properties = parse_header(plyfile, ext) + + # Get data + data = np.fromfile(plyfile, dtype=properties, count=num_points) + + return data + + +def header_properties(field_list, field_names): + + # List of lines to write + lines = [] + + # First line describing element vertex + lines.append('element vertex %d' % field_list[0].shape[0]) + + # Properties lines + i = 0 + for fields in field_list: + for field in fields.T: + lines.append('property %s %s' % (field.dtype.name, field_names[i])) + i += 1 + + return lines + + +def write_ply(filename, field_list, field_names, triangular_faces=None): + """ + Write ".ply" files + + Parameters + ---------- + filename : string + the name of the file to which the data is saved. A '.ply' extension will be appended to the + file name if it does no already have one. + + field_list : list, tuple, numpy array + the fields to be saved in the ply file. Either a numpy array, a list of numpy arrays or a + tuple of numpy arrays. Each 1D numpy array and each column of 2D numpy arrays are considered + as one field. + + field_names : list + the name of each fields as a list of strings. Has to be the same length as the number of + fields. + + Examples + -------- + >>> points = np.random.rand(10, 3) + >>> write_ply('example1.ply', points, ['x', 'y', 'z']) + + >>> values = np.random.randint(2, size=10) + >>> write_ply('example2.ply', [points, values], ['x', 'y', 'z', 'values']) + + >>> colors = np.random.randint(255, size=(10,3), dtype=np.uint8) + >>> field_names = ['x', 'y', 'z', 'red', 'green', 'blue', values'] + >>> write_ply('example3.ply', [points, colors, values], field_names) + + """ + + # Format list input to the right form + field_list = list(field_list) if (type(field_list) == list or type(field_list) == tuple) else list((field_list,)) + for i, field in enumerate(field_list): + if field.ndim < 2: + field_list[i] = field.reshape(-1, 1) + if field.ndim > 2: + print('fields have more than 2 dimensions') + return False + + # check all fields have the same number of data + n_points = [field.shape[0] for field in field_list] + if not np.all(np.equal(n_points, n_points[0])): + print('wrong field dimensions') + return False + + # Check if field_names and field_list have same nb of column + n_fields = np.sum([field.shape[1] for field in field_list]) + if (n_fields != len(field_names)): + print('wrong number of field names') + return False + + # Add extension if not there + if not filename.endswith('.ply'): + filename += '.ply' + + # open in text mode to write the header + with open(filename, 'w') as plyfile: + + # First magical word + header = ['ply'] + + # Encoding format + header.append('format binary_' + sys.byteorder + '_endian 1.0') + + # Points properties description + header.extend(header_properties(field_list, field_names)) + + # Add faces if needded + if triangular_faces is not None: + header.append('element face {:d}'.format(triangular_faces.shape[0])) + header.append('property list uchar int vertex_indices') + + # End of header + header.append('end_header') + + # Write all lines + for line in header: + plyfile.write("%s\n" % line) + + # open in binary/append to use tofile + with open(filename, 'ab') as plyfile: + + # Create a structured array + i = 0 + type_list = [] + for fields in field_list: + for field in fields.T: + type_list += [(field_names[i], field.dtype.str)] + i += 1 + data = np.empty(field_list[0].shape[0], dtype=type_list) + i = 0 + for fields in field_list: + for field in fields.T: + data[field_names[i]] = field + i += 1 + + data.tofile(plyfile) + + if triangular_faces is not None: + triangular_faces = triangular_faces.astype(np.int32) + type_list = [('k', 'uint8')] + [(str(ind), 'int32') for ind in range(3)] + data = np.empty(triangular_faces.shape[0], dtype=type_list) + data['k'] = np.full((triangular_faces.shape[0],), 3, dtype=np.uint8) + data['0'] = triangular_faces[:, 0] + data['1'] = triangular_faces[:, 1] + data['2'] = triangular_faces[:, 2] + data.tofile(plyfile) + + return True + + +def describe_element(name, df): + """ Takes the columns of the dataframe and builds a ply-like description + + Parameters + ---------- + name: str + df: pandas DataFrame + + Returns + ------- + element: list[str] + """ + property_formats = {'f': 'float', 'u': 'uchar', 'i': 'int'} + element = ['element ' + name + ' ' + str(len(df))] + + if name == 'face': + element.append("property list uchar int points_indices") + + else: + for i in range(len(df.columns)): + # get first letter of dtype to infer format + f = property_formats[str(df.dtypes[i])[0]] + element.append('property ' + f + ' ' + df.columns.values[i]) + + return element \ No newline at end of file diff --git a/utils/trainer.py b/utils/trainer.py new file mode 100644 index 0000000..c3acf33 --- /dev/null +++ b/utils/trainer.py @@ -0,0 +1,1910 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Class handling the training of any model +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 11/06/2018 +# + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Imports and global variables +# \**********************************/ +# + + +# Basic libs +import torch +import torch.nn as nn +import numpy as np +import pickle +import os +from os import makedirs, remove +from os.path import exists, join +import time +import sys + +# PLY reader +from utils.ply import read_ply, write_ply + +# Metrics +from utils.metrics import IoU_from_confusions +from utils.config import Config +from sklearn.metrics import confusion_matrix +from sklearn.neighbors import KDTree + +from models.blocks import KPConv + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Trainer Class +# \*******************/ +# + + +class ModelTrainer: + + # Initialization methods + # ------------------------------------------------------------------------------------------------------------------ + + def __init__(self, net, config, chkp_path=None, finetune=False, on_gpu=True): + """ + Initialize training parameters and reload previous model for restore/finetune + :param net: network object + :param config: configuration object + :param chkp_path: path to the checkpoint that needs to be loaded (None for new training) + :param finetune: finetune from checkpoint (True) or restore training from checkpoint (False) + :param on_gpu: Train on GPU or CPU + """ + + ############ + # Parameters + ############ + + # Epoch index + self.epoch = 0 + self.step = 0 + + # Optimizer + self.optimizer = torch.optim.SGD(net.parameters(), + lr=config.learning_rate, + momentum=config.momentum, + weight_decay=config.weight_decay) + + # Choose to train on CPU or GPU + if on_gpu and torch.cuda.is_available(): + self.device = torch.device("cuda:0") + else: + self.device = torch.device("cpu") + net.to(self.device) + + ########################## + # Load previous checkpoint + ########################## + + if (chkp_path is not None): + if finetune: + checkpoint = torch.load(chkp_path) + net.load_state_dict(checkpoint['model_state_dict']) + net.train() + print("Model restored and ready for finetuning.") + else: + checkpoint = torch.load(chkp_path) + net.load_state_dict(checkpoint['model_state_dict']) + self.optimizer.load_state_dict(checkpoint['optimizer_state_dict']) + self.epoch = checkpoint['epoch'] + net.train() + print("Model and training state restored.") + + # Path of the result folder + if config.saving: + if config.saving_path is None: + config.saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) + if not exists(config.saving_path): + makedirs(config.saving_path) + config.save() + + return + + # Training main method + # ------------------------------------------------------------------------------------------------------------------ + + def train(self, net, training_loader, val_loader, config): + """ + Train the model on a particular dataset. + """ + + ################ + # Initialization + ################ + + if config.saving: + # Training log file + with open(join(config.saving_path, 'training.txt'), "w") as file: + file.write('epochs steps out_loss offset_loss train_accuracy time\n') + + # Killing file (simply delete this file when you want to stop the training) + PID_file = join(config.saving_path, 'running_PID.txt') + if not exists(PID_file): + with open(PID_file, "w") as file: + file.write('Launched with PyCharm') + + # Checkpoints directory + checkpoint_directory = join(config.saving_path, 'checkpoints') + if not exists(checkpoint_directory): + makedirs(checkpoint_directory) + else: + checkpoint_directory = None + PID_file = None + + # Loop variables + t0 = time.time() + t = [time.time()] + last_display = time.time() + mean_dt = np.zeros(1) + + # Start training loop + for epoch in range(config.max_epoch): + + # Remove File for kill signal + if epoch == config.max_epoch - 1 and exists(PID_file): + remove(PID_file) + + self.step = 0 + for batch in training_loader: + + # Check kill signal (running_PID.txt deleted) + if config.saving and not exists(PID_file): + continue + + ################## + # Processing batch + ################## + + # New time + t = t[-1:] + t += [time.time()] + + if 'cuda' in self.device.type: + batch.to(self.device) + + # zero the parameter gradients + self.optimizer.zero_grad() + + # Forward pass + outputs = net(batch, config) + loss = net.loss(outputs, batch.labels) + acc = net.accuracy(outputs, batch.labels) + + t += [time.time()] + + # Backward + optimize + loss.backward() + + if config.grad_clip_norm > 0: + #torch.nn.utils.clip_grad_norm_(net.parameters(), config.grad_clip_norm) + torch.nn.utils.clip_grad_value_(net.parameters(), config.grad_clip_norm) + self.optimizer.step() + torch.cuda.synchronize(self.device) + + t += [time.time()] + + # Average timing + if self.step < 2: + mean_dt = np.array(t[1:]) - np.array(t[:-1]) + else: + mean_dt = 0.9 * mean_dt + 0.1 * (np.array(t[1:]) - np.array(t[:-1])) + + # Console display (only one per second) + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + message = 'e{:03d}-i{:04d} => L={:.3f} acc={:3.0f}% / t(ms): {:5.1f} {:5.1f} {:5.1f})' + print(message.format(self.epoch, self.step, + loss.item(), + 100*acc, + 1000 * mean_dt[0], + 1000 * mean_dt[1], + 1000 * mean_dt[2])) + + # Log file + if config.saving: + with open(join(config.saving_path, 'training.txt'), "a") as file: + message = '{:d} {:d} {:.3f} {:.3f} {:.3f} {:.3f}\n' + file.write(message.format(self.epoch, + self.step, + net.output_loss, + net.reg_loss, + acc, + t[-1] - t0)) + + + self.step += 1 + + ############## + # End of epoch + ############## + + # Check kill signal (running_PID.txt deleted) + if config.saving and not exists(PID_file): + break + + # Update learning rate + if self.epoch in config.lr_decays: + for param_group in self.optimizer.param_groups: + param_group['lr'] *= config.lr_decays[self.epoch] + + # Update epoch + self.epoch += 1 + + # Saving + if config.saving: + # Get current state dict + save_dict = {'epoch': self.epoch, + 'model_state_dict': net.state_dict(), + 'optimizer_state_dict': self.optimizer.state_dict(), + 'saving_path': config.saving_path} + + # Save current state of the network (for restoring purposes) + checkpoint_path = join(checkpoint_directory, 'current_chkp.tar') + torch.save(save_dict, checkpoint_path) + + # Save checkpoints occasionally + if (self.epoch + 1) % config.checkpoint_gap == 0: + checkpoint_path = join(checkpoint_directory, 'chkp_{:04d}.tar'.format(self.epoch)) + torch.save(save_dict, checkpoint_path) + + + # Validation + net.eval() + self.validation(net, val_loader, config) + net.train() + + print('Finished Training') + return + + # Validation methods + # ------------------------------------------------------------------------------------------------------------------ + + def validation(self, net, val_loader, config: Config): + + if config.dataset_task == 'classification': + self.object_classification_validation(net, val_loader, config) + elif config.dataset_task == 'segmentation': + self.object_segmentation_validation(net, val_loader, config) + elif config.dataset_task == 'cloud_segmentation': + self.cloud_segmentation_validation(net, val_loader, config) + elif config.dataset_task == 'slam_segmentation': + self.slam_segmentation_validation(net, val_loader, config) + else: + raise ValueError('No validation method implemented for this network type') + + def object_classification_validation(self, net, val_loader, config): + """ + Perform a round of validation and show/save results + :param net: network object + :param val_loader: data loader for validation set + :param config: configuration object + """ + + ############ + # Initialize + ############ + + # Choose validation smoothing parameter (0 for no smothing, 0.99 for big smoothing) + val_smooth = 0.95 + + # Number of classes predicted by the model + nc_model = config.num_classes + softmax = torch.nn.Softmax(1) + + # Initialize global prediction over all models + if not hasattr(self, 'val_probs'): + self.val_probs = np.zeros((val_loader.dataset.num_models, nc_model)) + + ##################### + # Network predictions + ##################### + + probs = [] + targets = [] + obj_inds = [] + + t = [time.time()] + last_display = time.time() + mean_dt = np.zeros(1) + + # Start validation loop + for batch in val_loader: + + # New time + t = t[-1:] + t += [time.time()] + + if 'cuda' in self.device.type: + batch.to(self.device) + + # Forward pass + outputs = net(batch, config) + + # Get probs and labels + probs += [softmax(outputs).cpu().detach().numpy()] + targets += [batch.labels.cpu().numpy()] + obj_inds += [batch.model_inds.cpu().numpy()] + torch.cuda.synchronize(self.device) + + # Average timing + t += [time.time()] + mean_dt = 0.95 * mean_dt + 0.05 * (np.array(t[1:]) - np.array(t[:-1])) + + # Display + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + message = 'Validation : {:.1f}% (timings : {:4.2f} {:4.2f})' + print(message.format(100 * len(obj_inds) / config.validation_size, + 1000 * (mean_dt[0]), + 1000 * (mean_dt[1]))) + + # Stack all validation predictions + probs = np.vstack(probs) + targets = np.hstack(targets) + obj_inds = np.hstack(obj_inds) + + ################### + # Voting validation + ################### + + self.val_probs[obj_inds] = val_smooth * self.val_probs[obj_inds] + (1-val_smooth) * probs + + ############ + # Confusions + ############ + + validation_labels = np.array(val_loader.dataset.label_values) + + # Compute classification results + C1 = confusion_matrix(targets, + np.argmax(probs, axis=1), + validation_labels) + + # Compute votes confusion + C2 = confusion_matrix(val_loader.dataset.input_labels, + np.argmax(self.val_probs, axis=1), + validation_labels) + + + # Saving (optionnal) + if config.saving: + print("Save confusions") + conf_list = [C1, C2] + file_list = ['val_confs.txt', 'vote_confs.txt'] + for conf, conf_file in zip(conf_list, file_list): + test_file = join(config.saving_path, conf_file) + if exists(test_file): + with open(test_file, "a") as text_file: + for line in conf: + for value in line: + text_file.write('%d ' % value) + text_file.write('\n') + else: + with open(test_file, "w") as text_file: + for line in conf: + for value in line: + text_file.write('%d ' % value) + text_file.write('\n') + + val_ACC = 100 * np.sum(np.diag(C1)) / (np.sum(C1) + 1e-6) + vote_ACC = 100 * np.sum(np.diag(C2)) / (np.sum(C2) + 1e-6) + print('Accuracies : val = {:.1f}% / vote = {:.1f}%'.format(val_ACC, vote_ACC)) + + return C1 + + def cloud_segmentation_validation(self, net, val_loader, config): + """ + Validation method for cloud segmentation models + """ + + ############ + # Initialize + ############ + + # Choose validation smoothing parameter (0 for no smothing, 0.99 for big smoothing) + val_smooth = 0.95 + softmax = torch.nn.Softmax(1) + + # Do not validate if dataset has no validation cloud + if val_loader.dataset.validation_split not in val_loader.dataset.all_splits: + return + + # Number of classes including ignored labels + nc_tot = val_loader.dataset.num_classes + + # Number of classes predicted by the model + nc_model = config.num_classes + + #print(nc_tot) + #print(nc_model) + + # Initiate global prediction over validation clouds + if not hasattr(self, 'validation_probs'): + self.validation_probs = [np.zeros((l.shape[0], nc_model)) + for l in val_loader.dataset.input_labels] + self.val_proportions = np.zeros(nc_model, dtype=np.float32) + i = 0 + for label_value in val_loader.dataset.label_values: + if label_value not in val_loader.dataset.ignored_labels: + self.val_proportions[i] = np.sum([np.sum(labels == label_value) + for labels in val_loader.dataset.validation_labels]) + i += 1 + + ##################### + # Network predictions + ##################### + + predictions = [] + targets = [] + + t = [time.time()] + last_display = time.time() + mean_dt = np.zeros(1) + + # Start validation loop + for i, batch in enumerate(val_loader): + + # New time + t = t[-1:] + t += [time.time()] + + if 'cuda' in self.device.type: + batch.to(self.device) + + # Forward pass + outputs = net(batch, config) + + # Get probs and labels + stacked_probs = softmax(outputs).cpu().detach().numpy() + labels = batch.labels.cpu().numpy() + lengths = batch.lengths[0].cpu().numpy() + in_inds = batch.input_inds.cpu().numpy() + cloud_inds = batch.cloud_inds.cpu().numpy() + torch.cuda.synchronize(self.device) + + # Get predictions and labels per instance + # *************************************** + + i0 = 0 + for b_i, length in enumerate(lengths): + + # Get prediction + target = labels[i0:i0 + length] + probs = stacked_probs[i0:i0 + length] + inds = in_inds[i0:i0 + length] + c_i = cloud_inds[b_i] + + # Update current probs in whole cloud + self.validation_probs[c_i][inds] = val_smooth * self.validation_probs[c_i][inds] \ + + (1 - val_smooth) * probs + + # Stack all prediction for this epoch + predictions.append(probs) + targets.append(target) + i0 += length + + # Average timing + t += [time.time()] + mean_dt = 0.95 * mean_dt + 0.05 * (np.array(t[1:]) - np.array(t[:-1])) + + # Display + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + message = 'Validation : {:.1f}% (timings : {:4.2f} {:4.2f})' + print(message.format(100 * i / config.validation_size, + 1000 * (mean_dt[0]), + 1000 * (mean_dt[1]))) + + # Confusions for our subparts of validation set + Confs = np.zeros((len(predictions), nc_tot, nc_tot), dtype=np.int32) + for i, (probs, truth) in enumerate(zip(predictions, targets)): + + # Insert false columns for ignored labels + for l_ind, label_value in enumerate(val_loader.dataset.label_values): + if label_value in val_loader.dataset.ignored_labels: + probs = np.insert(probs, l_ind, 0, axis=1) + + # Predicted labels + preds = val_loader.dataset.label_values[np.argmax(probs, axis=1)] + + # Confusions + Confs[i, :, :] = confusion_matrix(truth, preds, val_loader.dataset.label_values) + + # Sum all confusions + C = np.sum(Confs, axis=0).astype(np.float32) + + # Remove ignored labels from confusions + for l_ind, label_value in reversed(list(enumerate(val_loader.dataset.label_values))): + if label_value in val_loader.dataset.ignored_labels: + C = np.delete(C, l_ind, axis=0) + C = np.delete(C, l_ind, axis=1) + + # Balance with real validation proportions + C *= np.expand_dims(self.val_proportions / (np.sum(C, axis=1) + 1e-6), 1) + + # Objects IoU + IoUs = IoU_from_confusions(C) + + # Saving (optionnal) + if config.saving: + + # Name of saving file + test_file = join(config.saving_path, 'val_IoUs.txt') + + # Line to write: + line = '' + for IoU in IoUs: + line += '{:.3f} '.format(IoU) + line = line + '\n' + + # Write in file + if exists(test_file): + with open(test_file, "a") as text_file: + text_file.write(line) + else: + with open(test_file, "w") as text_file: + text_file.write(line) + + # Save potentials + pot_path = join(config.saving_path, 'potentials') + if not exists(pot_path): + makedirs(pot_path) + files = val_loader.dataset.train_files + i_val = 0 + for i, file_path in enumerate(files): + if val_loader.dataset.all_splits[i] == val_loader.dataset.validation_split: + pot_points = np.array(val_loader.dataset.pot_trees[i_val].data, copy=False) + cloud_name = file_path.split('/')[-1] + pot_name = join(pot_path, cloud_name) + pots = val_loader.dataset.potentials[i_val].numpy().astype(np.float32) + write_ply(pot_name, + [pot_points.astype(np.float32), pots], + ['x', 'y', 'z', 'pots']) + + # Print instance mean + mIoU = 100 * np.mean(IoUs) + print('{:s} mean IoU = {:.1f}%'.format(config.dataset, mIoU)) + + # Save predicted cloud occasionally + if config.saving and (self.epoch + 1) % config.checkpoint_gap == 0: + val_path = join(config.saving_path, 'val_preds_{:d}'.format(self.epoch)) + if not exists(val_path): + makedirs(val_path) + files = val_loader.dataset.train_files + i_val = 0 + for i, file_path in enumerate(files): + if val_loader.dataset.all_splits[i] == val_loader.dataset.validation_split: + + # Get points + points = val_loader.dataset.load_evaluation_points(file_path) + + # Get probs on our own ply points + sub_probs = self.validation_probs[i_val] + + # Insert false columns for ignored labels + for l_ind, label_value in enumerate(val_loader.dataset.label_values): + if label_value in val_loader.dataset.ignored_labels: + sub_probs = np.insert(sub_probs, l_ind, 0, axis=1) + + # Get the predicted labels + sub_preds = val_loader.dataset.label_values[np.argmax(sub_probs, axis=1).astype(np.int32)] + + # Reproject preds on the evaluations points + preds = (sub_preds[val_loader.dataset.validation_proj[i_val]]).astype(np.int32) + + # Path of saved validation file + cloud_name = file_path.split('/')[-1] + val_name = join(val_path, cloud_name) + + # Save file + labels = val_loader.dataset.validation_labels[i_val].astype(np.int32) + write_ply(val_name, + [points, preds, labels], + ['x', 'y', 'z', 'preds', 'class']) + + i_val += 1 + + return + + + + + + def validation_error(self, model, dataset): + """ + Validation method for classification models + """ + + ############ + # Initialize + ############ + + # Choose validation smoothing parameter (0 for no smothing, 0.99 for big smoothing) + val_smooth = 0.95 + + # Initialise iterator with train data + self.sess.run(dataset.val_init_op) + + # Number of classes predicted by the model + nc_model = model.config.num_classes + + # Initialize global prediction over all models + if not hasattr(self, 'val_probs'): + self.val_probs = np.zeros((len(dataset.input_labels['validation']), nc_model)) + + ##################### + # Network predictions + ##################### + + probs = [] + targets = [] + obj_inds = [] + + mean_dt = np.zeros(2) + last_display = time.time() + while True: + try: + # Run one step of the model. + t = [time.time()] + ops = (self.prob_logits, model.labels, model.inputs['object_inds']) + prob, labels, inds = self.sess.run(ops, {model.dropout_prob: 1.0}) + t += [time.time()] + + # Get probs and labels + probs += [prob] + targets += [labels] + obj_inds += [inds] + + # Average timing + t += [time.time()] + mean_dt = 0.95 * mean_dt + 0.05 * (np.array(t[1:]) - np.array(t[:-1])) + + # Display + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + message = 'Validation : {:.1f}% (timings : {:4.2f} {:4.2f})' + print(message.format(100 * len(obj_inds) / model.config.validation_size, + 1000 * (mean_dt[0]), + 1000 * (mean_dt[1]))) + + except tf.errors.OutOfRangeError: + break + + # Stack all validation predictions + probs = np.vstack(probs) + targets = np.hstack(targets) + obj_inds = np.hstack(obj_inds) + + ################### + # Voting validation + ################### + + self.val_probs[obj_inds] = val_smooth * self.val_probs[obj_inds] + (1-val_smooth) * probs + + ############ + # Confusions + ############ + + validation_labels = np.array(dataset.label_values) + + # Compute classification results + C1 = confusion_matrix(targets, + np.argmax(probs, axis=1), + validation_labels) + + # Compute training confusion + C2 = confusion_matrix(self.training_labels, + self.training_preds, + validation_labels) + + # Compute votes confusion + C3 = confusion_matrix(dataset.input_labels['validation'], + np.argmax(self.val_probs, axis=1), + validation_labels) + + + # Saving (optionnal) + if model.config.saving: + print("Save confusions") + conf_list = [C1, C2, C3] + file_list = ['val_confs.txt', 'training_confs.txt', 'vote_confs.txt'] + for conf, conf_file in zip(conf_list, file_list): + test_file = join(model.saving_path, conf_file) + if exists(test_file): + with open(test_file, "a") as text_file: + for line in conf: + for value in line: + text_file.write('%d ' % value) + text_file.write('\n') + else: + with open(test_file, "w") as text_file: + for line in conf: + for value in line: + text_file.write('%d ' % value) + text_file.write('\n') + + train_ACC = 100 * np.sum(np.diag(C2)) / (np.sum(C2) + 1e-6) + val_ACC = 100 * np.sum(np.diag(C1)) / (np.sum(C1) + 1e-6) + vote_ACC = 100 * np.sum(np.diag(C3)) / (np.sum(C3) + 1e-6) + print('Accuracies : train = {:.1f}% / val = {:.1f}% / vote = {:.1f}%'.format(train_ACC, val_ACC, vote_ACC)) + + return C1 + + def segment_validation_error(self, model, dataset): + """ + Validation method for single object segmentation models + """ + + ########## + # Initialize + ########## + + # Choose validation smoothing parameter (0 for no smothing, 0.99 for big smoothing) + val_smooth = 0.95 + + # Initialise iterator with train data + self.sess.run(dataset.val_init_op) + + # Number of classes predicted by the model + nc_model = model.config.num_classes + + # Initialize global prediction over all models + if not hasattr(self, 'val_probs'): + self.val_probs = [np.zeros((len(p_l), nc_model)) for p_l in dataset.input_point_labels['validation']] + + ##################### + # Network predictions + ##################### + + probs = [] + targets = [] + obj_inds = [] + mean_dt = np.zeros(2) + last_display = time.time() + for i0 in range(model.config.validation_size): + try: + # Run one step of the model. + t = [time.time()] + ops = (self.prob_logits, model.labels, model.inputs['in_batches'], model.inputs['object_inds']) + prob, labels, batches, o_inds = self.sess.run(ops, {model.dropout_prob: 1.0}) + t += [time.time()] + + # Get predictions and labels per instance + # *************************************** + + # Stack all validation predictions for each class separately + max_ind = np.max(batches) + for b_i, b in enumerate(batches): + + # Eliminate shadow indices + b = b[b < max_ind-0.5] + + # Stack all results + probs += [prob[b]] + targets += [labels[b]] + obj_inds += [o_inds[b_i]] + + # Average timing + t += [time.time()] + mean_dt = 0.95 * mean_dt + 0.05 * (np.array(t[1:]) - np.array(t[:-1])) + + # Display + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + message = 'Validation : {:.1f}% (timings : {:4.2f} {:4.2f})' + print(message.format(100 * i0 / model.config.validation_size, + 1000 * (mean_dt[0]), + 1000 * (mean_dt[1]))) + + except tf.errors.OutOfRangeError: + break + + ################### + # Voting validation + ################### + + for o_i, o_probs in zip(obj_inds, probs): + self.val_probs[o_i] = val_smooth * self.val_probs[o_i] + (1 - val_smooth) * o_probs + + ############ + # Confusions + ############ + + # Confusion matrix for each instance + n_parts = model.config.num_classes + Confs = np.zeros((len(probs), n_parts, n_parts), dtype=np.int32) + for i, (pred, truth) in enumerate(zip(probs, targets)): + parts = [j for j in range(pred.shape[1])] + Confs[i, :, :] = confusion_matrix(truth, np.argmax(pred, axis=1), parts) + + # Objects IoU + IoUs = IoU_from_confusions(Confs) + + + # Compute votes confusion + Confs = np.zeros((len(self.val_probs), n_parts, n_parts), dtype=np.int32) + for i, (pred, truth) in enumerate(zip(self.val_probs, dataset.input_point_labels['validation'])): + parts = [j for j in range(pred.shape[1])] + Confs[i, :, :] = confusion_matrix(truth, np.argmax(pred, axis=1), parts) + + # Objects IoU + vote_IoUs = IoU_from_confusions(Confs) + + # Saving (optionnal) + if model.config.saving: + + IoU_list = [IoUs, vote_IoUs] + file_list = ['val_IoUs.txt', 'vote_IoUs.txt'] + for IoUs_to_save, IoU_file in zip(IoU_list, file_list): + + # Name of saving file + test_file = join(model.saving_path, IoU_file) + + # Line to write: + line = '' + for instance_IoUs in IoUs_to_save: + for IoU in instance_IoUs: + line += '{:.3f} '.format(IoU) + line = line + '\n' + + # Write in file + if exists(test_file): + with open(test_file, "a") as text_file: + text_file.write(line) + else: + with open(test_file, "w") as text_file: + text_file.write(line) + + # Print instance mean + mIoU = 100 * np.mean(IoUs) + mIoU2 = 100 * np.mean(vote_IoUs) + print('{:s} : mIoU = {:.1f}% / vote mIoU = {:.1f}%'.format(model.config.dataset, mIoU, mIoU2)) + + return + + def cloud_validation_error(self, model, dataset): + """ + Validation method for cloud segmentation models + """ + + ########## + # Initialize + ########## + + # Choose validation smoothing parameter (0 for no smothing, 0.99 for big smoothing) + val_smooth = 0.95 + + # Do not validate if dataset has no validation cloud + if dataset.validation_split not in dataset.all_splits: + return + + # Initialise iterator with train data + self.sess.run(dataset.val_init_op) + + # Number of classes including ignored labels + nc_tot = dataset.num_classes + + # Number of classes predicted by the model + nc_model = model.config.num_classes + + # Initialize global prediction over validation clouds + if not hasattr(self, 'validation_probs'): + self.validation_probs = [np.zeros((l.shape[0], nc_model)) for l in dataset.input_labels['validation']] + self.val_proportions = np.zeros(nc_model, dtype=np.float32) + i = 0 + for label_value in dataset.label_values: + if label_value not in dataset.ignored_labels: + self.val_proportions[i] = np.sum([np.sum(labels == label_value) + for labels in dataset.validation_labels]) + i += 1 + + ##################### + # Network predictions + ##################### + + predictions = [] + targets = [] + mean_dt = np.zeros(2) + last_display = time.time() + for i0 in range(model.config.validation_size): + try: + # Run one step of the model. + t = [time.time()] + ops = (self.prob_logits, + model.labels, + model.inputs['in_batches'], + model.inputs['point_inds'], + model.inputs['cloud_inds']) + stacked_probs, labels, batches, point_inds, cloud_inds = self.sess.run(ops, {model.dropout_prob: 1.0}) + t += [time.time()] + + # Get predictions and labels per instance + # *************************************** + + # Stack all validation predictions for each class separately + max_ind = np.max(batches) + for b_i, b in enumerate(batches): + + # Eliminate shadow indices + b = b[b < max_ind-0.5] + + # Get prediction (only for the concerned parts) + probs = stacked_probs[b] + inds = point_inds[b] + c_i = cloud_inds[b_i] + + # Update current probs in whole cloud + self.validation_probs[c_i][inds] = val_smooth * self.validation_probs[c_i][inds] \ + + (1-val_smooth) * probs + + # Stack all prediction for this epoch + predictions += [probs] + targets += [dataset.input_labels['validation'][c_i][inds]] + + # Average timing + t += [time.time()] + mean_dt = 0.95 * mean_dt + 0.05 * (np.array(t[1:]) - np.array(t[:-1])) + + # Display + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + message = 'Validation : {:.1f}% (timings : {:4.2f} {:4.2f})' + print(message.format(100 * i0 / model.config.validation_size, + 1000 * (mean_dt[0]), + 1000 * (mean_dt[1]))) + + except tf.errors.OutOfRangeError: + break + + # Confusions for our subparts of validation set + Confs = np.zeros((len(predictions), nc_tot, nc_tot), dtype=np.int32) + for i, (probs, truth) in enumerate(zip(predictions, targets)): + + # Insert false columns for ignored labels + for l_ind, label_value in enumerate(dataset.label_values): + if label_value in dataset.ignored_labels: + probs = np.insert(probs, l_ind, 0, axis=1) + + # Predicted labels + preds = dataset.label_values[np.argmax(probs, axis=1)] + + # Confusions + Confs[i, :, :] = confusion_matrix(truth, preds, dataset.label_values) + + # Sum all confusions + C = np.sum(Confs, axis=0).astype(np.float32) + + # Remove ignored labels from confusions + for l_ind, label_value in reversed(list(enumerate(dataset.label_values))): + if label_value in dataset.ignored_labels: + C = np.delete(C, l_ind, axis=0) + C = np.delete(C, l_ind, axis=1) + + # Balance with real validation proportions + C *= np.expand_dims(self.val_proportions / (np.sum(C, axis=1) + 1e-6), 1) + + # Objects IoU + IoUs = IoU_from_confusions(C) + + # Saving (optionnal) + if model.config.saving: + + # Name of saving file + test_file = join(model.saving_path, 'val_IoUs.txt') + + # Line to write: + line = '' + for IoU in IoUs: + line += '{:.3f} '.format(IoU) + line = line + '\n' + + # Write in file + if exists(test_file): + with open(test_file, "a") as text_file: + text_file.write(line) + else: + with open(test_file, "w") as text_file: + text_file.write(line) + + # Print instance mean + mIoU = 100 * np.mean(IoUs) + print('{:s} mean IoU = {:.1f}%'.format(model.config.dataset, mIoU)) + + # Save predicted cloud occasionally + if model.config.saving and (self.training_epoch + 1) % model.config.checkpoint_gap == 0: + val_path = join(model.saving_path, 'val_preds_{:d}'.format(self.training_epoch)) + if not exists(val_path): + makedirs(val_path) + files = dataset.train_files + i_val = 0 + for i, file_path in enumerate(files): + if dataset.all_splits[i] == dataset.validation_split: + + # Get points + points = dataset.load_evaluation_points(file_path) + + # Get probs on our own ply points + sub_probs = self.validation_probs[i_val] + + # Insert false columns for ignored labels + for l_ind, label_value in enumerate(dataset.label_values): + if label_value in dataset.ignored_labels: + sub_probs = np.insert(sub_probs, l_ind, 0, axis=1) + + # Get the predicted labels + sub_preds = dataset.label_values[np.argmax(sub_probs, axis=1).astype(np.int32)] + + # Reproject preds on the evaluations points + preds = (sub_preds[dataset.validation_proj[i_val]]).astype(np.int32) + + # Path of saved validation file + cloud_name = file_path.split('/')[-1] + val_name = join(val_path, cloud_name) + + # Save file + labels = dataset.validation_labels[i_val].astype(np.int32) + write_ply(val_name, + [points, preds, labels], + ['x', 'y', 'z', 'preds', 'class']) + + i_val += 1 + + return + + def multi_cloud_validation_error(self, model, multi_dataset): + """ + Validation method for cloud segmentation models + """ + + ########## + # Initialize + ########## + + # Choose validation smoothing parameter (0 for no smothing, 0.99 for big smoothing) + val_smooth = 0.95 + + # Initialise iterator with train data + self.sess.run(multi_dataset.val_init_op) + + if not hasattr(self, 'validation_probs'): + + self.validation_probs = [] + self.val_proportions = [] + + for d_i, dataset in enumerate(multi_dataset.datasets): + + # Do not validate if dataset has no validation cloud + if dataset.validation_split not in dataset.all_splits: + continue + + # Number of classes including ignored labels + nc_tot = dataset.num_classes + + # Number of classes predicted by the model + nc_model = model.config.num_classes[d_i] + + # Initialize global prediction over validation clouds + self.validation_probs.append([np.zeros((l.shape[0], nc_model)) for l in dataset.input_labels['validation']]) + self.val_proportions.append(np.zeros(nc_model, dtype=np.float32)) + i = 0 + for label_value in dataset.label_values: + if label_value not in dataset.ignored_labels: + self.val_proportions[-1][i] = np.sum([np.sum(labels == label_value) + for labels in dataset.validation_labels]) + i += 1 + + ##################### + # Network predictions + ##################### + + pred_d_inds = [] + predictions = [] + targets = [] + mean_dt = np.zeros(2) + last_display = time.time() + for i0 in range(model.config.validation_size): + try: + # Run one step of the model. + t = [time.time()] + ops = (self.val_logits, + model.labels, + model.inputs['in_batches'], + model.inputs['point_inds'], + model.inputs['cloud_inds'], + model.inputs['dataset_inds']) + stacked_probs, labels, batches, p_inds, c_inds, d_inds = self.sess.run(ops, {model.dropout_prob: 1.0}) + t += [time.time()] + + # Get predictions and labels per instance + # *************************************** + + # Stack all validation predictions for each class separately + max_ind = np.max(batches) + for b_i, b in enumerate(batches): + + # Eliminate shadow indices + b = b[b < max_ind-0.5] + + # Get prediction (only for the concerned parts) + d_i = d_inds[b_i] + probs = stacked_probs[b, :model.config.num_classes[d_i]] + inds = p_inds[b] + c_i = c_inds[b_i] + + # Update current probs in whole cloud + self.validation_probs[d_i][c_i][inds] = val_smooth * self.validation_probs[d_i][c_i][inds] \ + + (1-val_smooth) * probs + + # Stack all prediction for this epoch + pred_d_inds += [d_i] + predictions += [probs] + targets += [multi_dataset.datasets[d_i].input_labels['validation'][c_i][inds]] + + # Average timing + t += [time.time()] + mean_dt = 0.95 * mean_dt + 0.05 * (np.array(t[1:]) - np.array(t[:-1])) + + # Display + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + message = 'Validation : {:.1f}% (timings : {:4.2f} {:4.2f})' + print(message.format(100 * i0 / model.config.validation_size, + 1000 * (mean_dt[0]), + 1000 * (mean_dt[1]))) + + except tf.errors.OutOfRangeError: + break + + # Convert list to np array for indexing + predictions = np.array(predictions) + targets = np.array(targets) + pred_d_inds = np.array(pred_d_inds, np.int32) + + IoUs = [] + for d_i, dataset in enumerate(multi_dataset.datasets): + + # Do not validate if dataset has no validation cloud + if dataset.validation_split not in dataset.all_splits: + continue + + # Number of classes including ignored labels + nc_tot = dataset.num_classes + + # Number of classes predicted by the model + nc_model = model.config.num_classes[d_i] + + # Extract the spheres from this dataset + tmp_inds = np.where(pred_d_inds == d_i)[0] + + # Confusions for our subparts of validation set + Confs = np.zeros((len(tmp_inds), nc_tot, nc_tot), dtype=np.int32) + for i, (probs, truth) in enumerate(zip(predictions[tmp_inds], targets[tmp_inds])): + + # Insert false columns for ignored labels + for l_ind, label_value in enumerate(dataset.label_values): + if label_value in dataset.ignored_labels: + probs = np.insert(probs, l_ind, 0, axis=1) + + # Predicted labels + preds = dataset.label_values[np.argmax(probs, axis=1)] + + # Confusions + Confs[i, :, :] = confusion_matrix(truth, preds, dataset.label_values) + + # Sum all confusions + C = np.sum(Confs, axis=0).astype(np.float32) + + # Remove ignored labels from confusions + for l_ind, label_value in reversed(list(enumerate(dataset.label_values))): + if label_value in dataset.ignored_labels: + C = np.delete(C, l_ind, axis=0) + C = np.delete(C, l_ind, axis=1) + + # Balance with real validation proportions + C *= np.expand_dims(self.val_proportions[d_i] / (np.sum(C, axis=1) + 1e-6), 1) + + # Objects IoU + IoUs += [IoU_from_confusions(C)] + + # Saving (optionnal) + if model.config.saving: + + # Name of saving file + test_file = join(model.saving_path, 'val_IoUs_{:d}_{:s}.txt'.format(d_i, dataset.name)) + + # Line to write: + line = '' + for IoU in IoUs[-1]: + line += '{:.3f} '.format(IoU) + line = line + '\n' + + # Write in file + if exists(test_file): + with open(test_file, "a") as text_file: + text_file.write(line) + else: + with open(test_file, "w") as text_file: + text_file.write(line) + + # Print instance mean + mIoU = 100 * np.mean(IoUs[-1]) + print('{:s} mean IoU = {:.1f}%'.format(dataset.name, mIoU)) + + # Save predicted cloud occasionally + if model.config.saving and (self.training_epoch + 1) % model.config.checkpoint_gap == 0: + val_path = join(model.saving_path, 'val_preds_{:d}'.format(self.training_epoch)) + if not exists(val_path): + makedirs(val_path) + + for d_i, dataset in enumerate(multi_dataset.datasets): + + dataset_val_path = join(val_path, '{:d}_{:s}'.format(d_i, dataset.name)) + if not exists(dataset_val_path): + makedirs(dataset_val_path) + + files = dataset.train_files + i_val = 0 + for i, file_path in enumerate(files): + if dataset.all_splits[i] == dataset.validation_split: + + # Get points + points = dataset.load_evaluation_points(file_path) + + # Get probs on our own ply points + sub_probs = self.validation_probs[d_i][i_val] + + # Insert false columns for ignored labels + for l_ind, label_value in enumerate(dataset.label_values): + if label_value in dataset.ignored_labels: + sub_probs = np.insert(sub_probs, l_ind, 0, axis=1) + + # Get the predicted labels + sub_preds = dataset.label_values[np.argmax(sub_probs, axis=1).astype(np.int32)] + + # Reproject preds on the evaluations points + preds = (sub_preds[dataset.validation_proj[i_val]]).astype(np.int32) + + # Path of saved validation file + cloud_name = file_path.split('/')[-1] + val_name = join(dataset_val_path, cloud_name) + + # Save file + labels = dataset.validation_labels[i_val].astype(np.int32) + write_ply(val_name, + [points, preds, labels], + ['x', 'y', 'z', 'preds', 'class']) + + i_val += 1 + + return + + def multi_validation_error(self, model, dataset): + """ + Validation method for multi object segmentation models + """ + + ########## + # Initialize + ########## + + # Choose validation smoothing parameter (0 for no smothing, 0.99 for big smoothing) + val_smooth = 0.95 + + # Initialise iterator with train data + self.sess.run(dataset.val_init_op) + + # Initialize global prediction over all models + if not hasattr(self, 'val_probs'): + self.val_probs = [] + for p_l, o_l in zip(dataset.input_point_labels['validation'], dataset.input_labels['validation']): + self.val_probs += [np.zeros((len(p_l), dataset.num_parts[o_l]))] + + ##################### + # Network predictions + ##################### + + probs = [] + targets = [] + objects = [] + obj_inds = [] + mean_dt = np.zeros(2) + last_display = time.time() + for i0 in range(model.config.validation_size): + try: + # Run one step of the model. + t = [time.time()] + ops = (model.logits, + model.labels, + model.inputs['super_labels'], + model.inputs['object_inds'], + model.inputs['in_batches']) + prob, labels, object_labels, o_inds, batches = self.sess.run(ops, {model.dropout_prob: 1.0}) + t += [time.time()] + + # Get predictions and labels per instance + # *************************************** + + # Stack all validation predictions for each class separately + max_ind = np.max(batches) + for b_i, b in enumerate(batches): + + # Eliminate shadow indices + b = b[b < max_ind-0.5] + + # Get prediction (only for the concerned parts) + obj = object_labels[b[0]] + pred = prob[b][:, :model.config.num_classes[obj]] + + # Stack all results + objects += [obj] + obj_inds += [o_inds[b_i]] + probs += [prob[b, :model.config.num_classes[obj]]] + targets += [labels[b]] + + # Average timing + t += [time.time()] + mean_dt = 0.95 * mean_dt + 0.05 * (np.array(t[1:]) - np.array(t[:-1])) + + # Display + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + message = 'Validation : {:.1f}% (timings : {:4.2f} {:4.2f})' + print(message.format(100 * i0 / model.config.validation_size, + 1000 * (mean_dt[0]), + 1000 * (mean_dt[1]))) + + except tf.errors.OutOfRangeError: + break + + ################### + # Voting validation + ################### + + for o_i, o_probs in zip(obj_inds, probs): + self.val_probs[o_i] = val_smooth * self.val_probs[o_i] + (1 - val_smooth) * o_probs + + ############ + # Confusions + ############ + + # Confusion matrix for each object + n_objs = [np.sum(np.array(objects) == l) for l in dataset.label_values] + Confs = [np.zeros((n_obj, n_parts, n_parts), dtype=np.int32) for n_parts, n_obj in + zip(dataset.num_parts, n_objs)] + obj_count = [0 for _ in n_objs] + for obj, pred, truth in zip(objects, probs, targets): + parts = [i for i in range(pred.shape[1])] + Confs[obj][obj_count[obj], :, :] = confusion_matrix(truth, np.argmax(pred, axis=1), parts) + obj_count[obj] += 1 + + # Objects mIoU + IoUs = [IoU_from_confusions(C) for C in Confs] + + + # Compute votes confusion + n_objs = [np.sum(np.array(dataset.input_labels['validation']) == l) for l in dataset.label_values] + Confs = [np.zeros((n_obj, n_parts, n_parts), dtype=np.int32) for n_parts, n_obj in + zip(dataset.num_parts, n_objs)] + obj_count = [0 for _ in n_objs] + for obj, pred, truth in zip(dataset.input_labels['validation'], + self.val_probs, + dataset.input_point_labels['validation']): + parts = [i for i in range(pred.shape[1])] + Confs[obj][obj_count[obj], :, :] = confusion_matrix(truth, np.argmax(pred, axis=1), parts) + obj_count[obj] += 1 + + # Objects mIoU + vote_IoUs = [IoU_from_confusions(C) for C in Confs] + + # Saving (optionnal) + if model.config.saving: + + IoU_list = [IoUs, vote_IoUs] + file_list = ['val_IoUs.txt', 'vote_IoUs.txt'] + + for IoUs_to_save, IoU_file in zip(IoU_list, file_list): + + # Name of saving file + test_file = join(model.saving_path, IoU_file) + + # Line to write: + line = '' + for obj_IoUs in IoUs_to_save: + for part_IoUs in obj_IoUs: + for IoU in part_IoUs: + line += '{:.3f} '.format(IoU) + line += '/ ' + line = line[:-2] + '\n' + + # Write in file + if exists(test_file): + with open(test_file, "a") as text_file: + text_file.write(line) + else: + with open(test_file, "w") as text_file: + text_file.write(line) + + # Print instance mean + mIoU = 100 * np.mean(np.hstack([np.mean(obj_IoUs, axis=1) for obj_IoUs in IoUs])) + class_mIoUs = [np.mean(obj_IoUs) for obj_IoUs in IoUs] + mcIoU = 100 * np.mean(class_mIoUs) + print('Val : mIoU = {:.1f}% / mcIoU = {:.1f}% '.format(mIoU, mcIoU)) + mIoU = 100 * np.mean(np.hstack([np.mean(obj_IoUs, axis=1) for obj_IoUs in vote_IoUs])) + class_mIoUs = [np.mean(obj_IoUs) for obj_IoUs in vote_IoUs] + mcIoU = 100 * np.mean(class_mIoUs) + print('Vote : mIoU = {:.1f}% / mcIoU = {:.1f}% '.format(mIoU, mcIoU)) + + return + + def slam_validation_error(self, model, dataset): + """ + Validation method for slam segmentation models + """ + + ########## + # Initialize + ########## + + # Do not validate if dataset has no validation cloud + if dataset.validation_split not in dataset.seq_splits: + return + + # Create folder for validation predictions + if not exists (join(model.saving_path, 'val_preds')): + makedirs(join(model.saving_path, 'val_preds')) + + # Initialize the dataset validation containers + dataset.val_points = [] + dataset.val_labels = [] + + # Initialise iterator with train data + self.sess.run(dataset.val_init_op) + + # Number of classes including ignored labels + nc_tot = dataset.num_classes + + # Number of classes predicted by the model + nc_model = model.config.num_classes + + ##################### + # Network predictions + ##################### + + predictions = [] + targets = [] + inds = [] + mean_dt = np.zeros(2) + last_display = time.time() + val_i = 0 + for i0 in range(model.config.validation_size): + try: + # Run one step of the model. + t = [time.time()] + ops = (self.prob_logits, + model.labels, + model.inputs['points'][0], + model.inputs['in_batches'], + model.inputs['frame_inds'], + model.inputs['frame_centers'], + model.inputs['augment_scales'], + model.inputs['augment_rotations']) + s_probs, s_labels, s_points, batches, f_inds, p0s, S, R = self.sess.run(ops, {model.dropout_prob: 1.0}) + t += [time.time()] + + # Get predictions and labels per instance + # *************************************** + + # Stack all validation predictions for each class separately + max_ind = np.max(batches) + for b_i, b in enumerate(batches): + + # Eliminate shadow indices + b = b[b < max_ind-0.5] + + # Get prediction (only for the concerned parts) + probs = s_probs[b] + labels = s_labels[b] + points = s_points[b, :] + S_i = S[b_i] + R_i = R[b_i] + p0 = p0s[b_i] + + # Get input points in their original positions + points2 = (points * (1/S_i)).dot(R_i.T) + + # get val_points that are in range + radiuses = np.sum(np.square(dataset.val_points[val_i] - p0), axis=1) + mask = radiuses < (0.9 * model.config.in_radius) ** 2 + + # Project predictions on the frame points + search_tree = KDTree(points2, leaf_size=50) + proj_inds = search_tree.query(dataset.val_points[val_i][mask, :], return_distance=False) + proj_inds = np.squeeze(proj_inds).astype(np.int32) + proj_probs = probs[proj_inds] + #proj_labels = labels[proj_inds] + + # Safe check if only one point: + if proj_probs.ndim < 2: + proj_probs = np.expand_dims(proj_probs, 0) + + # Insert false columns for ignored labels + for l_ind, label_value in enumerate(dataset.label_values): + if label_value in dataset.ignored_labels: + proj_probs = np.insert(proj_probs, l_ind, 0, axis=1) + + # Predicted labels + preds = dataset.label_values[np.argmax(proj_probs, axis=1)] + + # Save predictions in a binary file + filename ='{:02d}_{:07d}.npy'.format(f_inds[b_i, 0], f_inds[b_i, 1]) + filepath = join(model.saving_path, 'val_preds', filename) + if exists(filepath): + frame_preds = np.load(filepath) + else: + frame_preds = np.zeros(dataset.val_labels[val_i].shape, dtype=np.uint8) + frame_preds[mask] = preds.astype(np.uint8) + np.save(filepath, frame_preds) + + # Save some of the frame pots + if f_inds[b_i, 1] % 10 == 0: + pots = dataset.f_potentials['validation'][f_inds[b_i, 0]][f_inds[b_i, 1]] + write_ply(filepath[:-4]+'_pots.ply', + [dataset.val_points[val_i], dataset.val_labels[val_i], frame_preds, pots], + ['x', 'y', 'z', 'gt', 'pre', 'pots']) + + # Update validation confusions + frame_C = confusion_matrix(dataset.val_labels[val_i], frame_preds, dataset.label_values) + dataset.val_confs[f_inds[b_i, 0]][f_inds[b_i, 1], :, :] = frame_C + + # Stack all prediction for this epoch + predictions += [preds] + targets += [dataset.val_labels[val_i][mask]] + inds += [f_inds[b_i, :]] + val_i += 1 + + # Average timing + t += [time.time()] + mean_dt = 0.95 * mean_dt + 0.05 * (np.array(t[1:]) - np.array(t[:-1])) + + # Display + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + message = 'Validation : {:.1f}% (timings : {:4.2f} {:4.2f})' + print(message.format(100 * i0 / model.config.validation_size, + 1000 * (mean_dt[0]), + 1000 * (mean_dt[1]))) + + except tf.errors.OutOfRangeError: + break + + # Confusions for our subparts of validation set + Confs = np.zeros((len(predictions), nc_tot, nc_tot), dtype=np.int32) + for i, (preds, truth) in enumerate(zip(predictions, targets)): + + # Confusions + Confs[i, :, :] = confusion_matrix(truth, preds, dataset.label_values) + + ####################################### + # Results on this subpart of validation + ####################################### + + # Sum all confusions + C = np.sum(Confs, axis=0).astype(np.float32) + + # Balance with real validation proportions + C *= np.expand_dims(dataset.class_proportions['validation'] / (np.sum(C, axis=1) + 1e-6), 1) + + # Remove ignored labels from confusions + for l_ind, label_value in reversed(list(enumerate(dataset.label_values))): + if label_value in dataset.ignored_labels: + C = np.delete(C, l_ind, axis=0) + C = np.delete(C, l_ind, axis=1) + + # Objects IoU + IoUs = IoU_from_confusions(C) + + ##################################### + # Results on the whole validation set + ##################################### + + # Sum all validation confusions + C_tot = [np.sum(seq_C, axis=0) for seq_C in dataset.val_confs if len(seq_C) > 0] + C_tot = np.sum(np.stack(C_tot, axis=0), axis=0) + + s = '' + for cc in C_tot: + for c in cc: + s += '{:8.1f} '.format(c) + s += '\n' + print(s) + + # Remove ignored labels from confusions + for l_ind, label_value in reversed(list(enumerate(dataset.label_values))): + if label_value in dataset.ignored_labels: + C_tot = np.delete(C_tot, l_ind, axis=0) + C_tot = np.delete(C_tot, l_ind, axis=1) + + # Objects IoU + val_IoUs = IoU_from_confusions(C_tot) + + # Saving (optionnal) + if model.config.saving: + + IoU_list = [IoUs, val_IoUs] + file_list = ['subpart_IoUs.txt', 'val_IoUs.txt'] + for IoUs_to_save, IoU_file in zip(IoU_list, file_list): + + # Name of saving file + test_file = join(model.saving_path, IoU_file) + + # Line to write: + line = '' + for IoU in IoUs_to_save: + line += '{:.3f} '.format(IoU) + line = line + '\n' + + # Write in file + if exists(test_file): + with open(test_file, "a") as text_file: + text_file.write(line) + else: + with open(test_file, "w") as text_file: + text_file.write(line) + + # Print instance mean + mIoU = 100 * np.mean(IoUs) + print('{:s} : subpart mIoU = {:.1f} %'.format(model.config.dataset, mIoU)) + mIoU = 100 * np.mean(val_IoUs) + print('{:s} : val mIoU = {:.1f} %'.format(model.config.dataset, mIoU)) + + return + + # Saving methods + # ------------------------------------------------------------------------------------------------------------------ + + def save_kernel_points(self, model, epoch): + """ + Method saving kernel point disposition and current model weights for later visualization + """ + + if model.config.saving: + + # Create a directory to save kernels of this epoch + kernels_dir = join(model.saving_path, 'kernel_points', 'epoch{:d}'.format(epoch)) + if not exists(kernels_dir): + makedirs(kernels_dir) + + # Get points + all_kernel_points_tf = [v for v in tf.global_variables() if 'kernel_points' in v.name + and v.name.startswith('KernelPoint')] + all_kernel_points = self.sess.run(all_kernel_points_tf) + + # Get Extents + if False and 'gaussian' in model.config.convolution_mode: + all_kernel_params_tf = [v for v in tf.global_variables() if 'kernel_extents' in v.name + and v.name.startswith('KernelPoint')] + all_kernel_params = self.sess.run(all_kernel_params_tf) + else: + all_kernel_params = [None for p in all_kernel_points] + + # Save in ply file + for kernel_points, kernel_extents, v in zip(all_kernel_points, all_kernel_params, all_kernel_points_tf): + + # Name of saving file + ply_name = '_'.join(v.name[:-2].split('/')[1:-1]) + '.ply' + ply_file = join(kernels_dir, ply_name) + + # Data to save + if kernel_points.ndim > 2: + kernel_points = kernel_points[:, 0, :] + if False and 'gaussian' in model.config.convolution_mode: + data = [kernel_points, kernel_extents] + keys = ['x', 'y', 'z', 'sigma'] + else: + data = kernel_points + keys = ['x', 'y', 'z'] + + # Save + write_ply(ply_file, data, keys) + + # Get Weights + all_kernel_weights_tf = [v for v in tf.global_variables() if 'weights' in v.name + and v.name.startswith('KernelPointNetwork')] + all_kernel_weights = self.sess.run(all_kernel_weights_tf) + + # Save in numpy file + for kernel_weights, v in zip(all_kernel_weights, all_kernel_weights_tf): + np_name = '_'.join(v.name[:-2].split('/')[1:-1]) + '.npy' + np_file = join(kernels_dir, np_name) + np.save(np_file, kernel_weights) + + # Debug methods + # ------------------------------------------------------------------------------------------------------------------ + + def show_memory_usage(self, batch_to_feed): + + for l in range(self.config.num_layers): + neighb_size = list(batch_to_feed[self.in_neighbors_f32[l]].shape) + dist_size = neighb_size + [self.config.num_kernel_points, 3] + dist_memory = np.prod(dist_size) * 4 * 1e-9 + in_feature_size = neighb_size + [self.config.first_features_dim * 2**l] + in_feature_memory = np.prod(in_feature_size) * 4 * 1e-9 + out_feature_size = [neighb_size[0], self.config.num_kernel_points, self.config.first_features_dim * 2**(l+1)] + out_feature_memory = np.prod(out_feature_size) * 4 * 1e-9 + + print('Layer {:d} => {:.1f}GB {:.1f}GB {:.1f}GB'.format(l, + dist_memory, + in_feature_memory, + out_feature_memory)) + print('************************************') + + def debug_nan(self, model, inputs, logits): + """ + NaN happened, find where + """ + + print('\n\n------------------------ NaN DEBUG ------------------------\n') + + # First save everything to reproduce error + file1 = join(model.saving_path, 'all_debug_inputs.pkl') + with open(file1, 'wb') as f1: + pickle.dump(inputs, f1) + + # First save all inputs + file1 = join(model.saving_path, 'all_debug_logits.pkl') + with open(file1, 'wb') as f1: + pickle.dump(logits, f1) + + # Then print a list of the trainable variables and if they have nan + print('List of variables :') + print('*******************\n') + all_vars = self.sess.run(tf.global_variables()) + for v, value in zip(tf.global_variables(), all_vars): + nan_percentage = 100 * np.sum(np.isnan(value)) / np.prod(value.shape) + print(v.name, ' => {:.1f}% of values are NaN'.format(nan_percentage)) + + + print('Inputs :') + print('********') + + #Print inputs + nl = model.config.num_layers + for layer in range(nl): + + print('Layer : {:d}'.format(layer)) + + points = inputs[layer] + neighbors = inputs[nl + layer] + pools = inputs[2*nl + layer] + upsamples = inputs[3*nl + layer] + + nan_percentage = 100 * np.sum(np.isnan(points)) / np.prod(points.shape) + print('Points =>', points.shape, '{:.1f}% NaN'.format(nan_percentage)) + nan_percentage = 100 * np.sum(np.isnan(neighbors)) / np.prod(neighbors.shape) + print('neighbors =>', neighbors.shape, '{:.1f}% NaN'.format(nan_percentage)) + nan_percentage = 100 * np.sum(np.isnan(pools)) / np.prod(pools.shape) + print('pools =>', pools.shape, '{:.1f}% NaN'.format(nan_percentage)) + nan_percentage = 100 * np.sum(np.isnan(upsamples)) / np.prod(upsamples.shape) + print('upsamples =>', upsamples.shape, '{:.1f}% NaN'.format(nan_percentage)) + + ind = 4 * nl + features = inputs[ind] + nan_percentage = 100 * np.sum(np.isnan(features)) / np.prod(features.shape) + print('features =>', features.shape, '{:.1f}% NaN'.format(nan_percentage)) + ind += 1 + batch_weights = inputs[ind] + ind += 1 + in_batches = inputs[ind] + max_b = np.max(in_batches) + print(in_batches.shape) + in_b_sizes = np.sum(in_batches < max_b - 0.5, axis=-1) + print('in_batch_sizes =>', in_b_sizes) + ind += 1 + out_batches = inputs[ind] + max_b = np.max(out_batches) + print(out_batches.shape) + out_b_sizes = np.sum(out_batches < max_b - 0.5, axis=-1) + print('out_batch_sizes =>', out_b_sizes) + ind += 1 + point_labels = inputs[ind] + print('point labels, ', point_labels.shape, ', values : ', np.unique(point_labels)) + print(np.array([int(100 * np.sum(point_labels == l) / len(point_labels)) for l in np.unique(point_labels)])) + + ind += 1 + if model.config.dataset.startswith('ShapeNetPart_multi'): + object_labels = inputs[ind] + nan_percentage = 100 * np.sum(np.isnan(object_labels)) / np.prod(object_labels.shape) + print('object_labels =>', object_labels.shape, '{:.1f}% NaN'.format(nan_percentage)) + ind += 1 + augment_scales = inputs[ind] + ind += 1 + augment_rotations = inputs[ind] + ind += 1 + + print('\npoolings and upsamples nums :\n') + + #Print inputs + for layer in range(nl): + + print('\nLayer : {:d}'.format(layer)) + + neighbors = inputs[nl + layer] + pools = inputs[2*nl + layer] + upsamples = inputs[3*nl + layer] + + max_n = np.max(neighbors) + nums = np.sum(neighbors < max_n - 0.5, axis=-1) + print('min neighbors =>', np.min(nums)) + + if np.prod(pools.shape) > 0: + max_n = np.max(pools) + nums = np.sum(pools < max_n - 0.5, axis=-1) + print('min pools =>', np.min(nums)) + else: + print('pools empty') + + + if np.prod(upsamples.shape) > 0: + max_n = np.max(upsamples) + nums = np.sum(upsamples < max_n - 0.5, axis=-1) + print('min upsamples =>', np.min(nums)) + else: + print('upsamples empty') + + + print('\nFinished\n\n') + time.sleep(0.5) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/utils/visualizer.py b/utils/visualizer.py new file mode 100644 index 0000000..566da39 --- /dev/null +++ b/utils/visualizer.py @@ -0,0 +1,1831 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Class handling the visualization +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 11/06/2018 +# + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Imports and global variables +# \**********************************/ +# + + +# Basic libs +import torch +import numpy as np +from sklearn.neighbors import KDTree +from os import makedirs, remove, rename, listdir +from os.path import exists, join +import time +from mayavi import mlab +import sys + +from models.blocks import KPConv + +# PLY reader +from utils.ply import write_ply, read_ply + +# Configuration class +from utils.config import Config, bcolors + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Trainer Class +# \*******************/ +# + + +class ModelVisualizer: + + # Initialization methods + # ------------------------------------------------------------------------------------------------------------------ + + def __init__(self, net, config, chkp_path, on_gpu=True): + """ + Initialize training parameters and reload previous model for restore/finetune + :param net: network object + :param config: configuration object + :param chkp_path: path to the checkpoint that needs to be loaded (None for new training) + :param finetune: finetune from checkpoint (True) or restore training from checkpoint (False) + :param on_gpu: Train on GPU or CPU + """ + + ############ + # Parameters + ############ + + # Choose to train on CPU or GPU + if on_gpu and torch.cuda.is_available(): + self.device = torch.device("cuda:0") + else: + self.device = torch.device("cpu") + net.to(self.device) + + ########################## + # Load previous checkpoint + ########################## + + checkpoint = torch.load(chkp_path) + net.load_state_dict(checkpoint['model_state_dict']) + self.epoch = checkpoint['epoch'] + net.eval() + print("Model and training state restored.") + + return + + # Main visualization methods + # ------------------------------------------------------------------------------------------------------------------ + + def top_relu_activations(self, model, dataset, relu_idx=0, top_num=5): + """ + Test the model on test dataset to see which points activate the most each neurons in a relu layer + :param model: model used at training + :param dataset: dataset used at training + :param relu_idx: which features are to be visualized + :param top_num: how many top candidates are kept per features + """ + + ##################################### + # First choose the visualized feature + ##################################### + + # List all relu ops + all_ops = [op for op in tf.get_default_graph().get_operations() if op.name.startswith('KernelPointNetwork') + and op.name.endswith('LeakyRelu')] + + # List all possible Relu indices + print('\nPossible Relu indices:') + for i, t in enumerate(all_ops): + print(i, ': ', t.name) + + # Print the chosen one + if relu_idx is not None: + features_tensor = all_ops[relu_idx].outputs[0] + else: + relu_idx = int(input('Choose a Relu index: ')) + features_tensor = all_ops[relu_idx].outputs[0] + + # Get parameters + layer_idx = int(features_tensor.name.split('/')[1][6:]) + if 'strided' in all_ops[relu_idx].name and not ('strided' in all_ops[relu_idx+1].name): + layer_idx += 1 + features_dim = int(features_tensor.shape[1]) + radius = model.config.first_subsampling_dl * model.config.density_parameter * (2 ** layer_idx) + + print('You chose to compute the output of operation named:\n' + all_ops[relu_idx].name) + print('\nIt contains {:d} features.'.format(int(features_tensor.shape[1]))) + + print('\n****************************************************************************') + + ####################### + # Initialize containers + ####################### + + # Initialize containers + self.top_features = -np.ones((top_num, features_dim)) + self.top_classes = -np.ones((top_num, features_dim), dtype=np.int32) + self.saving = model.config.saving + + # Testing parameters + num_votes = 3 + + # Create visu folder + self.visu_path = None + self.fmt_str = None + if model.config.saving: + self.visu_path = join('visu', + 'visu_' + model.saving_path.split('/')[-1], + 'top_activations', + 'Relu{:02d}'.format(relu_idx)) + self.fmt_str = 'f{:04d}_top{:02d}.ply' + if not exists(self.visu_path): + makedirs(self.visu_path) + + # ******************* + # Network predictions + # ******************* + + mean_dt = np.zeros(2) + last_display = time.time() + for v in range(num_votes): + + # Run model on all test examples + # ****************************** + + # Initialise iterator with test data + if model.config.dataset.startswith('S3DIS'): + self.sess.run(dataset.val_init_op) + else: + self.sess.run(dataset.test_init_op) + count = 0 + + while True: + try: + + if model.config.dataset.startswith('ShapeNetPart'): + if model.config.dataset.split('_')[1] == 'multi': + label_op = model.inputs['super_labels'] + else: + label_op = model.inputs['point_labels'] + elif model.config.dataset.startswith('S3DIS'): + label_op = model.inputs['point_labels'] + elif model.config.dataset.startswith('Scannet'): + label_op = model.inputs['point_labels'] + elif model.config.dataset.startswith('ModelNet40'): + label_op = model.inputs['labels'] + else: + raise ValueError('Unsupported dataset') + + # Run one step of the model + t = [time.time()] + ops = (all_ops[-1].outputs[0], + features_tensor, + label_op, + model.inputs['points'], + model.inputs['pools'], + model.inputs['in_batches']) + _, stacked_features, labels, all_points, all_pools, in_batches = self.sess.run(ops, {model.dropout_prob: 1.0}) + t += [time.time()] + count += in_batches.shape[0] + + # Stack all batches + max_ind = np.max(in_batches) + stacked_batches = [] + for b_i, b in enumerate(in_batches): + stacked_batches += [b[b < max_ind - 0.5]*0+b_i] + stacked_batches = np.hstack(stacked_batches) + + # Find batches at wanted layer + for l in range(model.config.num_layers - 1): + if l >= layer_idx: + break + stacked_batches = stacked_batches[all_pools[l][:, 0]] + + # Get each example and update top_activations + for b_i, b in enumerate(in_batches): + b = b[b < max_ind - 0.5] + in_points = all_points[0][b] + features = stacked_features[stacked_batches == b_i] + points = all_points[layer_idx][stacked_batches == b_i] + if model.config.dataset in ['ShapeNetPart_multi', 'ModelNet40_classif']: + l = labels[b_i] + else: + l = np.argmax(np.bincount(labels[b])) + + self.update_top_activations(features, labels[b_i], points, in_points, radius) + + # Average timing + t += [time.time()] + mean_dt = 0.95 * mean_dt + 0.05 * (np.array(t[1:]) - np.array(t[:-1])) + + # Display + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + if model.config.dataset.startswith('S3DIS'): + completed = count / (model.config.validation_size * model.config.batch_num) + else: + completed = count / dataset.num_test + message = 'Vote {:d} : {:.1f}% (timings : {:4.2f} {:4.2f})' + print(message.format(v, + 100 * completed, + 1000 * (mean_dt[0]), + 1000 * (mean_dt[1]))) + #class_names = np.array([dataset.label_to_names[i] for i in range(dataset.num_classes)]) + #print(class_names[self.top_classes[:, :20]].T) + + except tf.errors.OutOfRangeError: + break + + return relu_idx + + def update_top_activations(self, features, label, l_points, input_points, radius, max_computed=60): + + top_num = self.top_features.shape[0] + + # Compute top indice for each feature + max_indices = np.argmax(features, axis=0) + + # get top_point neighborhoods + for features_i, idx in enumerate(max_indices[:max_computed]): + if features[idx, features_i] <= self.top_features[-1, features_i]: + continue + if label in self.top_classes[:, features_i]: + ind0 = np.where(self.top_classes[:, features_i] == label)[0][0] + if features[idx, features_i] <= self.top_features[ind0, features_i]: + continue + elif ind0 < top_num - 1: + self.top_features[ind0:-1, features_i] = self.top_features[ind0+1:, features_i] + self.top_classes[ind0:-1, features_i] = self.top_classes[ind0+1:, features_i] + for next_i in range(ind0 + 1, top_num): + old_f = join(self.visu_path, self.fmt_str.format(features_i, next_i + 1)) + new_f = join(self.visu_path, self.fmt_str.format(features_i, next_i)) + if exists(old_f): + if exists(new_f): + remove(new_f) + rename(old_f, new_f) + + # Find indice where new top should be placed + top_i = np.where(features[idx, features_i] > self.top_features[:, features_i])[0][0] + + # Update top features + if top_i < top_num - 1: + self.top_features[top_i + 1:, features_i] = self.top_features[top_i:-1, features_i] + self.top_features[top_i, features_i] = features[idx, features_i] + self.top_classes[top_i + 1:, features_i] = self.top_classes[top_i:-1, features_i] + self.top_classes[top_i, features_i] = label + + # Find in which batch lays the point + if self.saving: + + # Get inputs + l_features = features[:, features_i] + point = l_points[idx, :] + dist = np.linalg.norm(input_points - point, axis=1) + influence = (radius - dist) / radius + + # Project response on input cloud + if l_points.shape[0] == input_points.shape[0]: + responses = l_features + else: + tree = KDTree(l_points, leaf_size=50) + nn_k = min(l_points.shape[0], 10) + interp_dists, interp_inds = tree.query(input_points, nn_k, return_distance=True) + tukeys = np.square(1 - np.square(interp_dists / radius)) + tukeys[interp_dists > radius] = 0 + responses = np.sum(l_features[interp_inds] * tukeys, axis=1) + + # Handle last examples + for next_i in range(top_num - 1, top_i, -1): + old_f = join(self.visu_path, self.fmt_str.format(features_i, next_i)) + new_f = join(self.visu_path, self.fmt_str.format(features_i, next_i + 1)) + if exists(old_f): + if exists(new_f): + remove(new_f) + rename(old_f, new_f) + + # Save + filename = join(self.visu_path, self.fmt_str.format(features_i, top_i + 1)) + write_ply(filename, + [input_points, influence, responses], + ['x', 'y', 'z', 'influence', 'responses']) + + def show_deformable_kernels_old(self, model, dataset, deform_idx=0): + + ########################################## + # First choose the visualized deformations + ########################################## + + # List all deformation ops + all_ops = [op for op in tf.get_default_graph().get_operations() if op.name.startswith('KernelPointNetwork') + and op.name.endswith('deformed_KP')] + + print('\nPossible deformed indices:') + for i, t in enumerate(all_ops): + print(i, ': ', t.name) + + # Chosen deformations + deformed_KP_tensor = all_ops[deform_idx].outputs[0] + + # Layer index + layer_idx = int(all_ops[deform_idx].name.split('/')[1].split('_')[-1]) + + # Original kernel point positions + KP_vars = [v for v in tf.global_variables() if 'kernel_points' in v.name] + tmp = np.array(all_ops[deform_idx].name.split('/')) + test = [] + for v in KP_vars: + cmp = np.array(v.name.split('/')) + l = min(len(cmp), len(tmp)) + cmp = cmp[:l] + tmp = tmp[:l] + test += [np.sum(cmp == tmp)] + chosen_KP = np.argmax(test) + + print('You chose to visualize the output of operation named: ' + all_ops[deform_idx].name) + + print('\n****************************************************************************') + + # Run model on all test examples + # ****************************** + + # Initialise iterator with test data + if model.config.dataset.startswith('S3DIS'): + self.sess.run(dataset.val_init_op) + else: + self.sess.run(dataset.test_init_op) + count = 0 + + while True: + try: + + # Run one step of the model + t = [time.time()] + ops = (deformed_KP_tensor, + model.inputs['points'], + model.inputs['features'], + model.inputs['pools'], + model.inputs['in_batches'], + KP_vars) + stacked_deformed_KP, \ + all_points, \ + all_colors, \ + all_pools, \ + in_batches, \ + original_KPs = self.sess.run(ops, {model.dropout_prob: 1.0}) + t += [time.time()] + count += in_batches.shape[0] + + # Stack all batches + max_ind = np.max(in_batches) + stacked_batches = [] + for b_i, b in enumerate(in_batches): + stacked_batches += [b[b < max_ind - 0.5] * 0 + b_i] + stacked_batches = np.hstack(stacked_batches) + + # Find batches at wanted layer + for l in range(model.config.num_layers - 1): + if l >= layer_idx: + break + stacked_batches = stacked_batches[all_pools[l][:, 0]] + + # Get each example and update top_activations + in_points = [] + in_colors = [] + deformed_KP = [] + points = [] + lookuptrees = [] + for b_i, b in enumerate(in_batches): + b = b[b < max_ind - 0.5] + in_points += [all_points[0][b]] + deformed_KP += [stacked_deformed_KP[stacked_batches == b_i]] + points += [all_points[layer_idx][stacked_batches == b_i]] + lookuptrees += [KDTree(points[-1])] + if all_colors.shape[1] == 4: + in_colors += [all_colors[b, 1:]] + else: + in_colors += [None] + + print('New batch size : ', len(in_batches)) + + ########################### + # Interactive visualization + ########################### + + # Create figure for features + fig1 = mlab.figure('Features', bgcolor=(1.0, 1.0, 1.0), size=(1280, 920)) + fig1.scene.parallel_projection = False + + # Indices + global obj_i, point_i, plots, offsets, p_scale, show_in_p, aim_point + p_scale = 0.03 + obj_i = 0 + point_i = 0 + plots = {} + offsets = False + show_in_p = 2 + aim_point = np.zeros((1, 3)) + + def picker_callback(picker): + """ Picker callback: this get called when on pick events. + """ + global plots, aim_point + + if 'in_points' in plots: + if plots['in_points'].actor.actor._vtk_obj in [o._vtk_obj for o in picker.actors]: + point_rez = plots['in_points'].glyph.glyph_source.glyph_source.output.points.to_array().shape[0] + new_point_i = int(np.floor(picker.point_id / point_rez)) + if new_point_i < len(plots['in_points'].mlab_source.points): + # Get closest point in the layer we are interested in + aim_point = plots['in_points'].mlab_source.points[new_point_i:new_point_i + 1] + update_scene() + + if 'points' in plots: + if plots['points'].actor.actor._vtk_obj in [o._vtk_obj for o in picker.actors]: + point_rez = plots['points'].glyph.glyph_source.glyph_source.output.points.to_array().shape[0] + new_point_i = int(np.floor(picker.point_id / point_rez)) + if new_point_i < len(plots['points'].mlab_source.points): + # Get closest point in the layer we are interested in + aim_point = plots['points'].mlab_source.points[new_point_i:new_point_i + 1] + update_scene() + + def update_scene(): + global plots, offsets, p_scale, show_in_p, aim_point, point_i + + # Get the current view + v = mlab.view() + roll = mlab.roll() + + # clear figure + for key in plots.keys(): + plots[key].remove() + + plots = {} + + # Plot new data feature + p = points[obj_i] + + # Rescale points for visu + p = (p * 1.5 / model.config.in_radius) + + + # Show point cloud + if show_in_p <= 1: + plots['points'] = mlab.points3d(p[:, 0], + p[:, 1], + p[:, 2], + resolution=8, + scale_factor=p_scale, + scale_mode='none', + color=(0, 1, 1), + figure=fig1) + + if show_in_p >= 1: + + # Get points and colors + in_p = in_points[obj_i] + in_p = (in_p * 1.5 / model.config.in_radius) + + # Color point cloud if possible + in_c = in_colors[obj_i] + if in_c is not None: + + # Primitives + scalars = np.arange(len(in_p)) # Key point: set an integer for each point + + # Define color table (including alpha), which must be uint8 and [0,255] + colors = np.hstack((in_c, np.ones_like(in_c[:, :1]))) + colors = (colors * 255).astype(np.uint8) + + plots['in_points'] = mlab.points3d(in_p[:, 0], + in_p[:, 1], + in_p[:, 2], + scalars, + resolution=8, + scale_factor=p_scale*0.8, + scale_mode='none', + figure=fig1) + plots['in_points'].module_manager.scalar_lut_manager.lut.table = colors + + else: + + plots['in_points'] = mlab.points3d(in_p[:, 0], + in_p[:, 1], + in_p[:, 2], + resolution=8, + scale_factor=p_scale*0.8, + scale_mode='none', + figure=fig1) + + + # Get KP locations + rescaled_aim_point = aim_point * model.config.in_radius / 1.5 + point_i = lookuptrees[obj_i].query(rescaled_aim_point, return_distance=False)[0][0] + if offsets: + KP = points[obj_i][point_i] + deformed_KP[obj_i][point_i] + scals = np.ones_like(KP[:, 0]) + else: + KP = points[obj_i][point_i] + original_KPs[chosen_KP] + scals = np.zeros_like(KP[:, 0]) + + KP = (KP * 1.5 / model.config.in_radius) + + plots['KP'] = mlab.points3d(KP[:, 0], + KP[:, 1], + KP[:, 2], + scals, + colormap='autumn', + resolution=8, + scale_factor=1.2*p_scale, + scale_mode='none', + vmin=0, + vmax=1, + figure=fig1) + + + if True: + plots['center'] = mlab.points3d(p[point_i, 0], + p[point_i, 1], + p[point_i, 2], + scale_factor=1.1*p_scale, + scale_mode='none', + color=(0, 1, 0), + figure=fig1) + + # New title + plots['title'] = mlab.title(str(obj_i), color=(0, 0, 0), size=0.3, height=0.01) + text = '<--- (press g for previous)' + 50 * ' ' + '(press h for next) --->' + plots['text'] = mlab.text(0.01, 0.01, text, color=(0, 0, 0), width=0.98) + plots['orient'] = mlab.orientation_axes() + + # Set the saved view + mlab.view(*v) + mlab.roll(roll) + + return + + def animate_kernel(): + global plots, offsets, p_scale, show_in_p + + # Get KP locations + + KP_def = points[obj_i][point_i] + deformed_KP[obj_i][point_i] + KP_def = (KP_def * 1.5 / model.config.in_radius) + KP_def_color = (1, 0, 0) + + KP_rigid = points[obj_i][point_i] + original_KPs[chosen_KP] + KP_rigid = (KP_rigid * 1.5 / model.config.in_radius) + KP_rigid_color = (1, 0.7, 0) + + if offsets: + t_list = np.linspace(0, 1, 150, dtype=np.float32) + else: + t_list = np.linspace(1, 0, 150, dtype=np.float32) + + @mlab.animate(delay=10) + def anim(): + for t in t_list: + plots['KP'].mlab_source.set(x=t * KP_def[:, 0] + (1 - t) * KP_rigid[:, 0], + y=t * KP_def[:, 1] + (1 - t) * KP_rigid[:, 1], + z=t * KP_def[:, 2] + (1 - t) * KP_rigid[:, 2], + scalars=t * np.ones_like(KP_def[:, 0])) + + yield + + anim() + + return + + def keyboard_callback(vtk_obj, event): + global obj_i, point_i, offsets, p_scale, show_in_p + + if vtk_obj.GetKeyCode() in ['b', 'B']: + p_scale /= 1.5 + update_scene() + + elif vtk_obj.GetKeyCode() in ['n', 'N']: + p_scale *= 1.5 + update_scene() + + if vtk_obj.GetKeyCode() in ['g', 'G']: + obj_i = (obj_i - 1) % len(deformed_KP) + point_i = 0 + update_scene() + + elif vtk_obj.GetKeyCode() in ['h', 'H']: + obj_i = (obj_i + 1) % len(deformed_KP) + point_i = 0 + update_scene() + + elif vtk_obj.GetKeyCode() in ['k', 'K']: + offsets = not offsets + animate_kernel() + + elif vtk_obj.GetKeyCode() in ['z', 'Z']: + show_in_p = (show_in_p + 1) % 3 + update_scene() + + elif vtk_obj.GetKeyCode() in ['0']: + + print('Saving') + + # Find a new name + file_i = 0 + file_name = 'KP_{:03d}.ply'.format(file_i) + files = [f for f in listdir('KP_clouds') if f.endswith('.ply')] + while file_name in files: + file_i += 1 + file_name = 'KP_{:03d}.ply'.format(file_i) + + KP_deform = points[obj_i][point_i] + deformed_KP[obj_i][point_i] + KP_normal = points[obj_i][point_i] + original_KPs[chosen_KP] + + # Save + write_ply(join('KP_clouds', file_name), + [in_points[obj_i], in_colors[obj_i]], + ['x', 'y', 'z', 'red', 'green', 'blue']) + write_ply(join('KP_clouds', 'KP_{:03d}_deform.ply'.format(file_i)), + [KP_deform], + ['x', 'y', 'z']) + write_ply(join('KP_clouds', 'KP_{:03d}_normal.ply'.format(file_i)), + [KP_normal], + ['x', 'y', 'z']) + print('OK') + + return + + # Draw a first plot + pick_func = fig1.on_mouse_pick(picker_callback) + pick_func.tolerance = 0.01 + update_scene() + fig1.scene.interactor.add_observer('KeyPressEvent', keyboard_callback) + mlab.show() + + + + + except tf.errors.OutOfRangeError: + break + + def show_effective_recep_field(self, model, dataset, relu_idx=0): + + ################################################### + # First add a modulation variable on input features + ################################################### + + # Tensorflow random seed + random_seed = 42 + + # Create a modulated input feature op + with tf.variable_scope('input_modulations'): + initial = tf.constant(0., shape=[200000, 1]) + input_modulations_var = tf.Variable(initial, name='alphas') + input_modulations = 2 * tf.sigmoid(input_modulations_var) + assert_op = tf.assert_less(tf.shape(model.inputs['features'])[0], tf.shape(input_modulations)[0]) + with tf.control_dependencies([assert_op]): + modulated_input = model.inputs['features'] * input_modulations[:tf.shape(model.inputs['features'])[0]] + modulated_input = tf.identity(modulated_input, name='modulated_features') + + print('*******************************************') + + # Swap the op with the normal input features + for op in tf.get_default_graph().get_operations(): + + if 'input_modulations' in op.name: + continue + + if model.inputs['features'].name in [in_t.name for in_t in op.inputs]: + input_list = [] + for in_t in op.inputs: + if in_t.name == model.inputs['features'].name: + input_list += [modulated_input] + else: + input_list += [in_t] + print('swapping op ', op.name) + print('old inputs ', [in_t.name for in_t in op.inputs]) + print('new inputs ', [in_t.name for in_t in input_list]) + ge.swap_inputs(op, input_list) + + print('*******************************************') + + ########################## + # Create the ERF optimizer + ########################## + + # This optimizer only computes gradients for the feature modulation variables. We set the ERF loss, which + # consists of modifying the features in one location a the wanted layer + + with tf.variable_scope('ERF_loss'): + + # List all relu ops + all_ops = [op for op in tf.get_default_graph().get_operations() if op.name.startswith('KernelPointNetwork') + and op.name.endswith('LeakyRelu')] + + # Print the chosen one + features_tensor = all_ops[relu_idx].outputs[0] + + # Get parameters + layer_idx = int(features_tensor.name.split('/')[1][6:]) + if 'strided' in all_ops[relu_idx].name and not ('strided' in all_ops[relu_idx + 1].name): + layer_idx += 1 + features_dim = int(features_tensor.shape[1]) + radius = model.config.first_subsampling_dl * model.config.density_parameter * (2 ** layer_idx) + + print('You chose to visualize the output of operation named: ' + all_ops[relu_idx].name) + print('It contains {:d} features.'.format(int(features_tensor.shape[1]))) + + print('\nPossible Relu indices:') + for i, t in enumerate(all_ops): + print(i, ': ', t.name) + + print('\n****************************************************************************') + + # Get the receptive field of a random point + N = tf.shape(features_tensor)[0] + #random_ind = tf.random_uniform([1], minval=0, maxval=N, dtype=np.int32, seed=random_seed)[0] + #chosen_i_holder = tf.placeholder(tf.int32, name='chosen_ind') + aimed_coordinates = tf.placeholder(tf.float32, shape=(1, 3), name='aimed_coordinates') + d2 = tf.reduce_sum(tf.square(model.inputs['points'][layer_idx] - aimed_coordinates), axis=1) + chosen_i_tf = tf.argmin(d2, output_type=tf.int32) + + #test1 = tf.multiply(features_tensor, 2.0, name='test1') + #test2 = tf.multiply(features_tensor, 2.0, name='test2') + + # Gradient scaling operation + @tf.custom_gradient + def scale_grad_layer(x): + def scaled_grad(dy): + p_op = tf.print(x.name, + tf.reduce_mean(tf.abs(x)), + tf.reduce_mean(tf.abs(dy)), + output_stream=sys.stdout) + with tf.control_dependencies([p_op]): + new_dy = 1.0 * dy + return new_dy + return tf.identity(x), scaled_grad + + #test2 = scale_grad_layer(test2) + + # Get the tensor of error for these features (one for the chosen point, zero for the rest) + chosen_f_tf = tf.placeholder(tf.int32, name='feature_ind') + ERF_error = tf.expand_dims(tf.cast(tf.equal(tf.range(N), chosen_i_tf), tf.float32), 1) + ERF_error *= tf.expand_dims(tf.cast(tf.equal(tf.range(features_dim), chosen_f_tf), tf.float32), 0) + + # Get objective for the features (with a stop gradient so that we can get a gradient on the loss) + objective_features = features_tensor + ERF_error + objective_features = tf.stop_gradient(objective_features) + + # Loss is the error but with the features that can be learned to correct it + ERF_loss = tf.reduce_sum(tf.square(objective_features - features_tensor)) + + + with tf.variable_scope('ERF_optimizer'): + + # Create the gradient descent optimizer with a dummy learning rate + optimizer = tf.train.GradientDescentOptimizer(1.0) + + # Get the gradients with respect to the modulation variable + ERF_var_grads = optimizer.compute_gradients(ERF_loss, var_list=[input_modulations_var]) + + # Gradient of the modulations + ERF_train_op = optimizer.apply_gradients(ERF_var_grads) + + ################################ + # Run model on all test examples + ################################ + + # Init our modulation variable + self.sess.run(tf.variables_initializer([input_modulations_var])) + + # Initialise iterator with test data + self.sess.run(dataset.test_init_op) + count = 0 + + global plots, p_scale, show_in_p, remove_h, aim_point + aim_point = np.zeros((1, 3), dtype=np.float32) + remove_h = 1.05 + p_scale = 0.1 + plots = {} + show_in_p = False + + global points, in_points, grad_values, chosen_point, in_colors + points = None + in_points = np.zeros((0, 3)) + grad_values = None + chosen_point = None + in_colors = None + + ########################### + # Interactive visualization + ########################### + + # Create figure for features + fig1 = mlab.figure('Features', bgcolor=(0.5, 0.5, 0.5), size=(640, 480)) + fig1.scene.parallel_projection = False + + # Indices + + def update_ERF(only_points=False): + global points, in_points, grad_values, chosen_point, aim_point, in_colors + + # Generate clouds until we effectively changed + if only_points: + for i in range(50): + all_points = self.sess.run(model.inputs['points']) + if all_points[0].shape[0] != in_points.shape[0]: + break + + sum_grads = 0 + if only_points: + num_tries = 1 + else: + num_tries = 10 + + for test_i in range(num_tries): + + print('Updating ERF {:.0f}%'.format((test_i + 1) * 100 / num_tries)) + rand_f_i = np.random.randint(features_dim) + + # Reset input modulation variable + reset_op = input_modulations_var.assign(tf.zeros_like(input_modulations_var)) + self.sess.run(reset_op) + + # Apply gradient to input modulations + t = [time.time()] + ops = (ERF_train_op, + chosen_i_tf, + input_modulations_var, + model.inputs['points'], + model.inputs['features'], + model.inputs['pools'], + model.inputs['in_batches']) + feed_dict = {aimed_coordinates: aim_point, + chosen_f_tf: rand_f_i, + model.dropout_prob: 1.0} + _, chosen_i, new_mods, all_points, all_colors, all_pools, in_batches = self.sess.run(ops, feed_dict) + t += [time.time()] + + # Get the new value of the modulations + sum_grads += np.abs(self.sess.run(input_modulations_var)) + + grad = sum_grads / num_tries + + # Stack all batches + max_ind = np.max(in_batches) + stacked_batches = [] + for b_i, b in enumerate(in_batches): + stacked_batches += [b[b < max_ind - 0.5] * 0 + b_i] + stacked_batches = np.hstack(stacked_batches) + + # Find batches at wanted layer + for l in range(model.config.num_layers - 1): + if l >= layer_idx: + break + stacked_batches = stacked_batches[all_pools[l][:, 0]] + + # Get each example and update top_activations + for b_i, b in enumerate(in_batches): + b = b[b < max_ind - 0.5] + in_points = all_points[0][b] + in_colors = all_colors[b, 1:] + points = all_points[layer_idx][stacked_batches == b_i] + grad_values = grad[b] + + chosen_point = all_points[layer_idx][chosen_i] + + def update_scene(): + global plots, p_scale, show_in_p, remove_h + global points, in_points, grad_values, chosen_point + + # Get the current view + v = mlab.view() + roll = mlab.roll() + + # clear figure + for key in plots.keys(): + plots[key].remove() + + plots = {} + + # Plot new data feature + in_p = in_points + p = points + p0 = chosen_point + responses = 100 * np.abs(np.ravel(grad_values)) + #xresponses = responses ** (1/2) + + # Remove roof + if 0.0 < remove_h < 1.0: + floor_h = np.min(in_p[:, 2]) + ceil_h = np.max(in_p[:, 2]) + threshold = floor_h + (ceil_h - floor_h) * remove_h + responses = responses[in_p[:, 2] < threshold] + in_p = in_p[in_p[:, 2] < threshold] + p = p[p[:, 2] < threshold] + + # Rescale responses + min_response, max_response = np.min(responses), np.max(responses) + + # Show point cloud + if show_in_p: + plots['points'] = mlab.points3d(p[:, 0], + p[:, 1], + p[:, 2], + resolution=8, + scale_factor=p_scale, + scale_mode='none', + color=(0, 1, 1), + figure=fig1) + + plots['in_points'] = mlab.points3d(in_p[:, 0], + in_p[:, 1], + in_p[:, 2], + responses, + resolution=8, + scale_factor=p_scale * 0.8, + scale_mode='none', + vmin=0.1, + vmax=1.5, + figure=fig1) + + plots['center'] = mlab.points3d(p0[0], + p0[1], + p0[2], + scale_factor=1.5 * p_scale, + scale_mode='none', + color=(0, 0, 0), + figure=fig1) + + # New title + plots['title'] = mlab.title(str(int(100*remove_h)) + '%', color=(0, 0, 0), size=0.3, height=0.01) + text = '<--- (press g to remove ceiling)' + 50 * ' ' + '(press h to add ceiling) --->' + plots['text'] = mlab.text(0.01, 0.01, text, color=(0, 0, 0), width=0.98) + plots['orient'] = mlab.orientation_axes() + + # Set the saved view + mlab.view(*v) + mlab.roll(roll) + + return + + def picker_callback(picker): + """ Picker callback: this get called when on pick events. + """ + global plots, aim_point, in_points + + if plots['in_points'].actor.actor._vtk_obj in [o._vtk_obj for o in picker.actors]: + point_rez = plots['in_points'].glyph.glyph_source.glyph_source.output.points.to_array().shape[0] + new_point_i = int(np.floor(picker.point_id / point_rez)) + if new_point_i < len(plots['in_points'].mlab_source.points): + + # Get closest point in the layer we are interested in + aim_point = plots['in_points'].mlab_source.points[new_point_i:new_point_i + 1] + update_ERF() + update_scene() + + def keyboard_callback(vtk_obj, event): + global remove_h, p_scale, show_in_p + global in_points, grad_values, chosen_point, in_colors + + print(vtk_obj.GetKeyCode()) + + + if vtk_obj.GetKeyCode() in ['b', 'B']: + p_scale /= 1.5 + update_scene() + + elif vtk_obj.GetKeyCode() in ['n', 'N']: + p_scale *= 1.5 + update_scene() + + if vtk_obj.GetKeyCode() in ['g', 'G']: + if remove_h > 0.0: + remove_h -= 0.1 + update_scene() + + elif vtk_obj.GetKeyCode() in ['h', 'H']: + if remove_h < 1.0: + remove_h += 0.1 + update_ERF() + update_scene() + + elif vtk_obj.GetKeyCode() in ['z', 'Z']: + show_in_p = not show_in_p + update_scene() + + elif vtk_obj.GetKeyCode() in ['x', 'X']: + # Reset potentials + dataset.potentials['ERF'] = [] + dataset.min_potentials['ERF'] = [] + for i, tree in enumerate(dataset.input_trees['test']): + dataset.potentials['ERF'] += [np.random.rand(tree.data.shape[0]) * 1e-3] + dataset.min_potentials['ERF'] += [float(np.min(dataset.potentials['ERF'][-1]))] + + # Update figure + update_ERF(only_points=True) + update_scene() + + elif vtk_obj.GetKeyCode() in ['0']: + + print('Saving') + + # Find a new name + file_i = 0 + file_name = 'ERF_{:03d}.ply'.format(file_i) + files = [f for f in listdir('ERF_clouds') if f.endswith('.ply')] + while file_name in files: + file_i += 1 + file_name = 'ERF_{:03d}.ply'.format(file_i) + + # Save + responses = 100 * np.abs(np.ravel(grad_values)) + write_ply(join('ERF_clouds', file_name), + [in_points, in_colors, responses], + ['x', 'y', 'z', 'red', 'green', 'blue', 'erf']) + write_ply(join('ERF_clouds', 'ERF_{:03d}_center.ply'.format(file_i)), + [chosen_point.reshape([1, -1])], + ['x', 'y', 'z']) + print('OK') + + return + + # Draw a first plot + pick_func = fig1.on_mouse_pick(picker_callback) + pick_func.tolerance = 0.01 + update_ERF(only_points=True) + update_scene() + fig1.scene.interactor.add_observer('KeyPressEvent', keyboard_callback) + mlab.show() + + def show_deformable_kernels(self, net, loader, config, deform_idx=0): + """ + Show some inference with deformable kernels + """ + + ########################################## + # First choose the visualized deformations + ########################################## + + print('\nList of the deformable convolution available (chosen one highlighted in green)') + fmt_str = ' {:}{:2d} > KPConv(r={:.3f}, Din={:d}, Dout={:d}){:}' + deform_convs = [] + for m in net.modules(): + if isinstance(m, KPConv) and m.deformable: + if len(deform_convs) == deform_idx: + color = bcolors.OKGREEN + else: + color = bcolors.FAIL + print(fmt_str.format(color, len(deform_convs), m.radius, m.in_channels, m.out_channels, bcolors.ENDC)) + deform_convs.append(m) + + ################ + # Initialization + ################ + + print('\n****************************************************\n') + + # Loop variables + t0 = time.time() + t = [time.time()] + last_display = time.time() + mean_dt = np.zeros(1) + count = 0 + + # Start training loop + for epoch in range(config.max_epoch): + + for batch in loader: + + ################## + # Processing batch + ################## + + # New time + t = t[-1:] + t += [time.time()] + + if 'cuda' in self.device.type: + batch.to(self.device) + + # Forward pass + outputs = net(batch, config) + original_KP = deform_convs[deform_idx].kernel_points.cpu().detach().numpy() + stacked_deformed_KP = deform_convs[deform_idx].deformed_KP.cpu().detach().numpy() + count += batch.lengths[0].shape[0] + + if 'cuda' in self.device.type: + torch.cuda.synchronize(self.device) + + # Find layer + l = None + for i, p in enumerate(batch.points): + if p.shape[0] == stacked_deformed_KP.shape[0]: + l = i + + t += [time.time()] + + # Get data + in_points = [] + in_colors = [] + deformed_KP = [] + points = [] + lookuptrees = [] + i0 = 0 + for b_i, length in enumerate(batch.lengths[0]): + in_points.append(batch.points[0][i0:i0 + length].cpu().detach().numpy()) + if batch.features.shape[1] == 4: + in_colors.append(batch.features[i0:i0 + length, 1:].cpu().detach().numpy()) + else: + in_colors.append(None) + i0 += length + + i0 = 0 + for b_i, length in enumerate(batch.lengths[l]): + points.append(batch.points[l][i0:i0 + length].cpu().detach().numpy()) + deformed_KP.append(stacked_deformed_KP[i0:i0 + length]) + lookuptrees.append(KDTree(points[-1])) + i0 += length + + ########################### + # Interactive visualization + ########################### + + # Create figure for features + fig1 = mlab.figure('Deformations', bgcolor=(1.0, 1.0, 1.0), size=(1280, 920)) + fig1.scene.parallel_projection = False + + # Indices + global obj_i, point_i, plots, offsets, p_scale, show_in_p, aim_point + p_scale = 0.03 + obj_i = 0 + point_i = 0 + plots = {} + offsets = False + show_in_p = 2 + aim_point = np.zeros((1, 3)) + + def picker_callback(picker): + """ Picker callback: this get called when on pick events. + """ + global plots, aim_point + + if 'in_points' in plots: + if plots['in_points'].actor.actor._vtk_obj in [o._vtk_obj for o in picker.actors]: + point_rez = plots['in_points'].glyph.glyph_source.glyph_source.output.points.to_array().shape[0] + new_point_i = int(np.floor(picker.point_id / point_rez)) + if new_point_i < len(plots['in_points'].mlab_source.points): + # Get closest point in the layer we are interested in + aim_point = plots['in_points'].mlab_source.points[new_point_i:new_point_i + 1] + update_scene() + + if 'points' in plots: + if plots['points'].actor.actor._vtk_obj in [o._vtk_obj for o in picker.actors]: + point_rez = plots['points'].glyph.glyph_source.glyph_source.output.points.to_array().shape[0] + new_point_i = int(np.floor(picker.point_id / point_rez)) + if new_point_i < len(plots['points'].mlab_source.points): + # Get closest point in the layer we are interested in + aim_point = plots['points'].mlab_source.points[new_point_i:new_point_i + 1] + update_scene() + + def update_scene(): + global plots, offsets, p_scale, show_in_p, aim_point, point_i + + # Get the current view + v = mlab.view() + roll = mlab.roll() + + # clear figure + for key in plots.keys(): + plots[key].remove() + + plots = {} + + # Plot new data feature + p = points[obj_i] + + # Rescale points for visu + p = (p * 1.5 / config.in_radius) + + + # Show point cloud + if show_in_p <= 1: + plots['points'] = mlab.points3d(p[:, 0], + p[:, 1], + p[:, 2], + resolution=8, + scale_factor=p_scale, + scale_mode='none', + color=(0, 1, 1), + figure=fig1) + + if show_in_p >= 1: + + # Get points and colors + in_p = in_points[obj_i] + in_p = (in_p * 1.5 / config.in_radius) + + # Color point cloud if possible + in_c = in_colors[obj_i] + if in_c is not None: + + # Primitives + scalars = np.arange(len(in_p)) # Key point: set an integer for each point + + # Define color table (including alpha), which must be uint8 and [0,255] + colors = np.hstack((in_c, np.ones_like(in_c[:, :1]))) + colors = (colors * 255).astype(np.uint8) + + plots['in_points'] = mlab.points3d(in_p[:, 0], + in_p[:, 1], + in_p[:, 2], + scalars, + resolution=8, + scale_factor=p_scale*0.8, + scale_mode='none', + figure=fig1) + plots['in_points'].module_manager.scalar_lut_manager.lut.table = colors + + else: + + plots['in_points'] = mlab.points3d(in_p[:, 0], + in_p[:, 1], + in_p[:, 2], + resolution=8, + scale_factor=p_scale*0.8, + scale_mode='none', + figure=fig1) + + + # Get KP locations + rescaled_aim_point = aim_point * config.in_radius / 1.5 + point_i = lookuptrees[obj_i].query(rescaled_aim_point, return_distance=False)[0][0] + if offsets: + KP = points[obj_i][point_i] + deformed_KP[obj_i][point_i] + scals = np.ones_like(KP[:, 0]) + else: + KP = points[obj_i][point_i] + original_KP + scals = np.zeros_like(KP[:, 0]) + + KP = (KP * 1.5 / config.in_radius) + + plots['KP'] = mlab.points3d(KP[:, 0], + KP[:, 1], + KP[:, 2], + scals, + colormap='autumn', + resolution=8, + scale_factor=1.2*p_scale, + scale_mode='none', + vmin=0, + vmax=1, + figure=fig1) + + + if True: + plots['center'] = mlab.points3d(p[point_i, 0], + p[point_i, 1], + p[point_i, 2], + scale_factor=1.1*p_scale, + scale_mode='none', + color=(0, 1, 0), + figure=fig1) + + # New title + plots['title'] = mlab.title(str(obj_i), color=(0, 0, 0), size=0.3, height=0.01) + text = '<--- (press g for previous)' + 50 * ' ' + '(press h for next) --->' + plots['text'] = mlab.text(0.01, 0.01, text, color=(0, 0, 0), width=0.98) + plots['orient'] = mlab.orientation_axes() + + # Set the saved view + mlab.view(*v) + mlab.roll(roll) + + return + + def animate_kernel(): + global plots, offsets, p_scale, show_in_p + + # Get KP locations + + KP_def = points[obj_i][point_i] + deformed_KP[obj_i][point_i] + KP_def = (KP_def * 1.5 / config.in_radius) + KP_def_color = (1, 0, 0) + + KP_rigid = points[obj_i][point_i] + original_KP + KP_rigid = (KP_rigid * 1.5 / config.in_radius) + KP_rigid_color = (1, 0.7, 0) + + if offsets: + t_list = np.linspace(0, 1, 150, dtype=np.float32) + else: + t_list = np.linspace(1, 0, 150, dtype=np.float32) + + @mlab.animate(delay=10) + def anim(): + for t in t_list: + plots['KP'].mlab_source.set(x=t * KP_def[:, 0] + (1 - t) * KP_rigid[:, 0], + y=t * KP_def[:, 1] + (1 - t) * KP_rigid[:, 1], + z=t * KP_def[:, 2] + (1 - t) * KP_rigid[:, 2], + scalars=t * np.ones_like(KP_def[:, 0])) + + yield + + anim() + + return + + def keyboard_callback(vtk_obj, event): + global obj_i, point_i, offsets, p_scale, show_in_p + + if vtk_obj.GetKeyCode() in ['b', 'B']: + p_scale /= 1.5 + update_scene() + + elif vtk_obj.GetKeyCode() in ['n', 'N']: + p_scale *= 1.5 + update_scene() + + if vtk_obj.GetKeyCode() in ['g', 'G']: + obj_i = (obj_i - 1) % len(deformed_KP) + point_i = 0 + update_scene() + + elif vtk_obj.GetKeyCode() in ['h', 'H']: + obj_i = (obj_i + 1) % len(deformed_KP) + point_i = 0 + update_scene() + + elif vtk_obj.GetKeyCode() in ['k', 'K']: + offsets = not offsets + animate_kernel() + + elif vtk_obj.GetKeyCode() in ['z', 'Z']: + show_in_p = (show_in_p + 1) % 3 + update_scene() + + elif vtk_obj.GetKeyCode() in ['0']: + + print('Saving') + + # Find a new name + file_i = 0 + file_name = 'KP_{:03d}.ply'.format(file_i) + files = [f for f in listdir('KP_clouds') if f.endswith('.ply')] + while file_name in files: + file_i += 1 + file_name = 'KP_{:03d}.ply'.format(file_i) + + KP_deform = points[obj_i][point_i] + deformed_KP[obj_i][point_i] + KP_normal = points[obj_i][point_i] + original_KP + + # Save + write_ply(join('KP_clouds', file_name), + [in_points[obj_i], in_colors[obj_i]], + ['x', 'y', 'z', 'red', 'green', 'blue']) + write_ply(join('KP_clouds', 'KP_{:03d}_deform.ply'.format(file_i)), + [KP_deform], + ['x', 'y', 'z']) + write_ply(join('KP_clouds', 'KP_{:03d}_normal.ply'.format(file_i)), + [KP_normal], + ['x', 'y', 'z']) + print('OK') + + return + + # Draw a first plot + pick_func = fig1.on_mouse_pick(picker_callback) + pick_func.tolerance = 0.01 + update_scene() + fig1.scene.interactor.add_observer('KeyPressEvent', keyboard_callback) + mlab.show() + + return + + @staticmethod + def show_activation(path, relu_idx=0, save_video=False): + """ + This function show the saved input point clouds maximizing the activations. You can also directly load the files + in a visualization software like CloudCompare. + In the case of relu_idx = 0 and if gaussian mode, the associated filter is also shown. This function can only + show the filters for the last saved epoch. + """ + + ################ + # Find the files + ################ + + # Check visu folder + visu_path = join('visu', + 'visu_' + path.split('/')[-1], + 'top_activations', + 'Relu{:02d}'.format(relu_idx)) + if not exists(visu_path): + message = 'Relu {:d} activations of the model {:s} not found.' + raise ValueError(message.format(relu_idx, path.split('/')[-1])) + + # Get the list of files + feature_files = np.sort([f for f in listdir(visu_path) if f.endswith('.ply')]) + if len(feature_files) == 0: + message = 'Relu {:d} activations of the model {:s} not found.' + raise ValueError(message.format(relu_idx, path.split('/')[-1])) + + # Load mode + config = Config() + config.load(path) + mode = config.convolution_mode + + ################# + # Get activations + ################# + + all_points = [] + all_responses = [] + + for file in feature_files: + + # Load points + data = read_ply(join(visu_path, file)) + all_points += [np.vstack((data['x'], data['y'], data['z'])).T] + all_responses += [data['responses']] + + ########################### + # Interactive visualization + ########################### + + # Create figure for features + fig1 = mlab.figure('Features', bgcolor=(0.5, 0.5, 0.5), size=(640, 480)) + fig1.scene.parallel_projection = False + + # Indices + global file_i + file_i = 0 + + def update_scene(): + + # clear figure + mlab.clf(fig1) + + # Plot new data feature + points = all_points[file_i] + responses = all_responses[file_i] + min_response, max_response = np.min(responses), np.max(responses) + responses = (responses - min_response) / (max_response - min_response) + + # Rescale points for visu + points = (points * 1.5 / config.in_radius + np.array([1.0, 1.0, 1.0])) * 50.0 + + # Show point clouds colorized with activations + activations = mlab.points3d(points[:, 0], + points[:, 1], + points[:, 2], + responses, + scale_factor=3.0, + scale_mode='none', + vmin=0.1, + vmax=0.9, + figure=fig1) + + # New title + mlab.title(feature_files[file_i], color=(0, 0, 0), size=0.3, height=0.01) + text = '<--- (press g for previous)' + 50*' ' + '(press h for next) --->' + mlab.text(0.01, 0.01, text, color=(0, 0, 0), width=0.98) + mlab.orientation_axes() + + return + + def keyboard_callback(vtk_obj, event): + global file_i + + if vtk_obj.GetKeyCode() in ['g', 'G']: + + file_i = (file_i - 1) % len(all_responses) + update_scene() + + elif vtk_obj.GetKeyCode() in ['h', 'H']: + + file_i = (file_i + 1) % len(all_responses) + update_scene() + + return + + # Draw a first plot + update_scene() + fig1.scene.interactor.add_observer('KeyPressEvent', keyboard_callback) + mlab.show() + + return + + # Utilities + # ------------------------------------------------------------------------------------------------------------------ + + @staticmethod + def load_last_kernels(path): + + # Directories of validation error + kernel_dirs = np.array([f for f in listdir(join(path, 'kernel_points')) if f.startswith('epoch')]) + + # Find last epoch folder + epochs = np.array([int(f[5:]) for f in kernel_dirs]) + last_dir = kernel_dirs[np.argmax(epochs)] + + # Find saved files for the first layer + kernel_file = join(path, 'kernel_points', last_dir, 'layer_0_simple_0.ply') + weights_file = join(path, 'kernel_points', last_dir, 'layer_0_simple_0.npy') + + # Read kernel file + data = read_ply(kernel_file) + points = np.vstack((data['x'], data['y'], data['z'])).T + extents = data['sigma'].astype(np.float32) + + # Read weight file + w = np.load(weights_file) + + return points, extents, w + + @staticmethod + def apply_weights(points, kernel, weights, extents): + + # Get all difference matrices [n_points, n_kpoints, dim] + points = np.expand_dims(points, 1) + points = np.tile(points, [1, kernel.shape[0], 1]) + differences = points - kernel + + # Compute distance matrices [n_points, n_kpoints] + sq_distances = np.sum(np.square(differences), axis=-1) + + # Compute gaussians [n_points, n_kpoints] + gaussian_values = np.exp(-sq_distances / (2 * np.square(extents))) + + # Apply weights + return np.matmul(gaussian_values, np.squeeze(weights)) + + + def top_relu_activations_old(self, model, dataset, relu_idx=0, top_num=5): + """ + Test the model on test dataset to see which points activate the most each neurons in a relu layer + :param model: model used at training + :param dataset: dataset used at training + :param relu_idx: which features are to be visualized + :param top_num: how many top candidates are kept per features + """ + + ##################################### + # First choose the visualized feature + ##################################### + + # List all relu ops + all_ops = [op for op in tf.get_default_graph().get_operations() if op.name.startswith('KernelPointNetwork') + and op.name.endswith('LeakyRelu')] + + # Non relu ops in case we want the first KPConv features + KPConv_0 = [op for op in tf.get_default_graph().get_operations() if op.name.endswith('layer_0/simple_0/Sum_1')] + + # Print the chosen one + if relu_idx == 0: + features_tensor = KPConv_0[relu_idx].outputs[0] + else: + features_tensor = all_ops[relu_idx].outputs[0] + + # Get parameters + layer_idx = int(features_tensor.name.split('/')[1][6:]) + if 'strided' in all_ops[relu_idx].name and not ('strided' in all_ops[relu_idx+1].name): + layer_idx += 1 + features_dim = int(features_tensor.shape[1]) + radius = model.config.first_subsampling_dl * model.config.density_parameter * (2 ** layer_idx) + + if relu_idx == 0 : + print('SPECIAL CASE : relu_idx = 0 => visualization of the fist KPConv before relu') + print('You chose to visualize the output of operation named: ' + KPConv_0[0].name) + print('It contains {:d} features.'.format(int(features_tensor.shape[1]))) + else : + print('You chose to visualize the output of operation named: ' + all_ops[relu_idx].name) + print('It contains {:d} features.'.format(int(features_tensor.shape[1]))) + + print('\nPossible Relu indices:') + for i, t in enumerate(all_ops): + print(i, ': ', t.name) + + print('\n****************************************************************************') + + ##################### + # Initialize containers + ##################### + + # Initialize containers + self.top_features = -np.ones((top_num, features_dim)) + self.top_classes = -np.ones((top_num, features_dim), dtype=np.int32) + self.saving = model.config.saving + + # Testing parameters + num_votes = 3 + + # Create visu folder + self.visu_path = None + self.fmt_str = None + if model.config.saving: + self.visu_path = join('visu', + 'visu_' + model.saving_path.split('/')[-1], + 'top_activations', + 'Relu{:02d}'.format(relu_idx)) + self.fmt_str = 'f{:04d}_top{:02d}.ply' + if not exists(self.visu_path): + makedirs(self.visu_path) + + # ******************* + # Network predictions + # ******************* + + mean_dt = np.zeros(2) + last_display = time.time() + for v in range(num_votes): + + # Run model on all test examples + # ****************************** + + # Initialise iterator with test data + if model.config.dataset.startswith('S3DIS'): + self.sess.run(dataset.val_init_op) + else: + self.sess.run(dataset.test_init_op) + count = 0 + + while True: + try: + + if model.config.dataset.startswith('ShapeNetPart'): + if model.config.dataset.split('_')[1] == 'multi': + label_op = model.inputs['super_labels'] + else: + label_op = model.inputs['point_labels'] + elif model.config.dataset.startswith('S3DIS'): + label_op = model.inputs['point_labels'] + elif model.config.dataset.startswith('Scannet'): + label_op = model.inputs['point_labels'] + elif model.config.dataset.startswith('ModelNet40'): + label_op = model.inputs['labels'] + else: + raise ValueError('Unsupported dataset') + + # Run one step of the model + t = [time.time()] + ops = (all_ops[-1].outputs[0], + features_tensor, + label_op, + model.inputs['points'], + model.inputs['pools'], + model.inputs['in_batches']) + _, stacked_features, labels, all_points, all_pools, in_batches = self.sess.run(ops, {model.dropout_prob: 1.0}) + t += [time.time()] + count += in_batches.shape[0] + + + # Stack all batches + max_ind = np.max(in_batches) + stacked_batches = [] + for b_i, b in enumerate(in_batches): + stacked_batches += [b[b < max_ind - 0.5]*0+b_i] + stacked_batches = np.hstack(stacked_batches) + + # Find batches at wanted layer + for l in range(model.config.num_layers - 1): + if l >= layer_idx: + break + stacked_batches = stacked_batches[all_pools[l][:, 0]] + + # Get each example and update top_activations + for b_i, b in enumerate(in_batches): + b = b[b < max_ind - 0.5] + in_points = all_points[0][b] + features = stacked_features[stacked_batches == b_i] + points = all_points[layer_idx][stacked_batches == b_i] + if model.config.dataset in ['ShapeNetPart_multi', 'ModelNet40_classif']: + l = labels[b_i] + else: + l = np.argmax(np.bincount(labels[b])) + + self.update_top_activations(features, labels[b_i], points, in_points, radius) + + # Average timing + t += [time.time()] + mean_dt = 0.95 * mean_dt + 0.05 * (np.array(t[1:]) - np.array(t[:-1])) + + # Display + if (t[-1] - last_display) > 1.0: + last_display = t[-1] + if model.config.dataset.startswith('S3DIS'): + completed = count / (model.config.validation_size * model.config.batch_num) + else: + completed = count / dataset.num_test + message = 'Vote {:d} : {:.1f}% (timings : {:4.2f} {:4.2f})' + print(message.format(v, + 100 * completed, + 1000 * (mean_dt[0]), + 1000 * (mean_dt[1]))) + #class_names = np.array([dataset.label_to_names[i] for i in range(dataset.num_classes)]) + #print(class_names[self.top_classes[:, :20]].T) + + except tf.errors.OutOfRangeError: + break + + return + + + + +def show_ModelNet_models(all_points): + + ########################### + # Interactive visualization + ########################### + + # Create figure for features + fig1 = mlab.figure('Models', bgcolor=(1, 1, 1), size=(1000, 800)) + fig1.scene.parallel_projection = False + + # Indices + global file_i + file_i = 0 + + def update_scene(): + + # clear figure + mlab.clf(fig1) + + # Plot new data feature + points = all_points[file_i] + + # Rescale points for visu + points = (points * 1.5 + np.array([1.0, 1.0, 1.0])) * 50.0 + + # Show point clouds colorized with activations + activations = mlab.points3d(points[:, 0], + points[:, 1], + points[:, 2], + points[:, 2], + scale_factor=3.0, + scale_mode='none', + figure=fig1) + + # New title + mlab.title(str(file_i), color=(0, 0, 0), size=0.3, height=0.01) + text = '<--- (press g for previous)' + 50 * ' ' + '(press h for next) --->' + mlab.text(0.01, 0.01, text, color=(0, 0, 0), width=0.98) + mlab.orientation_axes() + + return + + def keyboard_callback(vtk_obj, event): + global file_i + + if vtk_obj.GetKeyCode() in ['g', 'G']: + + file_i = (file_i - 1) % len(all_points) + update_scene() + + elif vtk_obj.GetKeyCode() in ['h', 'H']: + + file_i = (file_i + 1) % len(all_points) + update_scene() + + return + + # Draw a first plot + update_scene() + fig1.scene.interactor.add_observer('KeyPressEvent', keyboard_callback) + mlab.show() + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/visualize_deformations.py b/visualize_deformations.py new file mode 100644 index 0000000..1d9bf90 --- /dev/null +++ b/visualize_deformations.py @@ -0,0 +1,193 @@ +# +# +# 0=================================0 +# | Kernel Point Convolutions | +# 0=================================0 +# +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Callable script to start a training on ModelNet40 dataset +# +# ---------------------------------------------------------------------------------------------------------------------- +# +# Hugues THOMAS - 06/03/2020 +# + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Imports and global variables +# \**********************************/ +# + +# Common libs +import signal +import os +import numpy as np +import sys +import torch + +# Dataset +from datasets.ModelNet40 import * +from torch.utils.data import DataLoader + +from utils.config import Config +from utils.visualizer import ModelVisualizer +from models.architectures import KPCNN + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Main Call +# \***************/ +# + +def model_choice(chosen_log): + + ########################### + # Call the test initializer + ########################### + + # Automatically retrieve the last trained model + if chosen_log in ['last_ModelNet40', 'last_ShapeNetPart', 'last_S3DIS']: + + # Dataset name + test_dataset = '_'.join(chosen_log.split('_')[1:]) + + # List all training logs + logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) + + # Find the last log of asked dataset + for log in logs[::-1]: + log_config = Config() + log_config.load(log) + if log_config.dataset.startswith(test_dataset): + chosen_log = log + break + + if chosen_log in ['last_ModelNet40', 'last_ShapeNetPart', 'last_S3DIS']: + raise ValueError('No log of the dataset "' + test_dataset + '" found') + + # Check if log exists + if not os.path.exists(chosen_log): + raise ValueError('The given log does not exists: ' + chosen_log) + + return chosen_log + + +# ---------------------------------------------------------------------------------------------------------------------- +# +# Main Call +# \***************/ +# + +if __name__ == '__main__': + + ############################### + # Choose the model to visualize + ############################### + + # Here you can choose which model you want to test with the variable test_model. Here are the possible values : + # + # > 'last_XXX': Automatically retrieve the last trained model on dataset XXX + # > '(old_)results/Log_YYYY-MM-DD_HH-MM-SS': Directly provide the path of a trained model + + chosen_log = 'results/Log_2020-03-23_22-18-26' # => ModelNet40 + + # You can also choose the index of the snapshot to load (last by default) + chkp_idx = None + + # Eventually you can choose which feature is visualized (index of the deform operation in the network) + deform_idx = 0 + + # Deal with 'last_XXX' choices + chosen_log = model_choice(chosen_log) + + ############################ + # Initialize the environment + ############################ + + # Set which gpu is going to be used + GPU_ID = '0' + + # Set GPU visible device + os.environ['CUDA_VISIBLE_DEVICES'] = GPU_ID + + ############### + # Previous chkp + ############### + + # Find all checkpoints in the chosen training folder + chkp_path = os.path.join(chosen_log, 'checkpoints') + chkps = [f for f in os.listdir(chkp_path) if f[:4] == 'chkp'] + + # Find which snapshot to restore + if chkp_idx is None: + chosen_chkp = 'current_chkp.tar' + else: + chosen_chkp = np.sort(chkps)[chkp_idx] + chosen_chkp = os.path.join(chosen_log, 'checkpoints', chosen_chkp) + + # Initialize configuration class + config = Config() + config.load(chosen_log) + + ################################## + # Change model parameters for test + ################################## + + # Change parameters for the test here. For example, you can stop augmenting the input data. + + #config.augment_noise = 0.0001 + #config.augment_symmetries = False + #config.batch_num = 3 + #config.in_radius = 4 + + ############## + # Prepare Data + ############## + + print() + print('Data Preparation') + print('****************') + + # Initialize datasets + test_dataset = ModelNet40Dataset(config, train=False) + + # Initialize samplers + test_sampler = ModelNet40Sampler(test_dataset) + + # Initialize the dataloader + test_loader = DataLoader(test_dataset, + batch_size=1, + sampler=test_sampler, + collate_fn=ModelNet40Collate, + num_workers=0, + pin_memory=True) + + # Calibrate samplers + test_sampler.calibration(test_loader) + + print('\nModel Preparation') + print('*****************') + + # Define network model + t1 = time.time() + if config.dataset_task == 'classification': + net = KPCNN(config) + else: + raise ValueError('Unsupported dataset_task for deformation visu: ' + config.dataset_task) + + # Define a visualizer class + visualizer = ModelVisualizer(net, config, chkp_path=chosen_chkp, on_gpu=False) + print('Done in {:.1f}s\n'.format(time.time() - t1)) + + print('\nStart visualization') + print('*******************') + + # Training + visualizer.show_deformable_kernels(net, test_loader, config, deform_idx) + + +