Initial commit

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 ..

cpp_wrappers/cpp_neighbors/neighbors/neighbors.cpp

@@ -0,0 +1,333 @@
+
+#include "neighbors.h"
+
+
+void brute_neighbors(vector<PointXYZ>& queries, vector<PointXYZ>& supports, vector<int>& 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<vector<int>> 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<PointXYZ>& queries,
+                        vector<PointXYZ>& supports,
+                        vector<int>& 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<vector<int>> tmp(queries.size());
+	vector<vector<float>> 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<PointXYZ>& queries,
+                                vector<PointXYZ>& supports,
+                                vector<int>& q_batches,
+                                vector<int>& s_batches,
+                                vector<int>& 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<vector<int>> tmp(queries.size());
+	vector<vector<float>> 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<PointXYZ>::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<PointXYZ>& queries,
+                                vector<PointXYZ>& supports,
+                                vector<int>& q_batches,
+                                vector<int>& s_batches,
+                                vector<int>& 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<vector<pair<size_t, float>>> 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<float, PointCloud > ,
+                                                        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<PointXYZ>(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<PointXYZ>(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;
+}
+

cpp_wrappers/cpp_neighbors/neighbors/neighbors.h

@@ -0,0 +1,29 @@
+
+
+#include "../../cpp_utils/cloud/cloud.h"
+#include "../../cpp_utils/nanoflann/nanoflann.hpp"
+
+#include <set>
+#include <cstdint>
+
+using namespace std;
+
+
+void ordered_neighbors(vector<PointXYZ>& queries,
+                        vector<PointXYZ>& supports,
+                        vector<int>& neighbors_indices,
+                        float radius);
+
+void batch_ordered_neighbors(vector<PointXYZ>& queries,
+                                vector<PointXYZ>& supports,
+                                vector<int>& q_batches,
+                                vector<int>& s_batches,
+                                vector<int>& neighbors_indices,
+                                float radius);
+
+void batch_nanoflann_neighbors(vector<PointXYZ>& queries,
+                                vector<PointXYZ>& supports,
+                                vector<int>& q_batches,
+                                vector<int>& s_batches,
+                                vector<int>& neighbors_indices,
+                                float radius);

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())
+
+
+
+
+
+
+
+

cpp_wrappers/cpp_neighbors/wrapper.cpp

@@ -0,0 +1,238 @@
+#include <Python.h>
+#include <numpy/arrayobject.h>
+#include "neighbors/neighbors.h"
+#include <string>
+
+
+
+// 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<PointXYZ> queries;
+	vector<PointXYZ> supports;
+	vector<int> q_batches;
+	vector<int> s_batches;
+	queries = vector<PointXYZ>((PointXYZ*)PyArray_DATA(queries_array), (PointXYZ*)PyArray_DATA(queries_array) + Nq);
+	supports = vector<PointXYZ>((PointXYZ*)PyArray_DATA(supports_array), (PointXYZ*)PyArray_DATA(supports_array) + Ns);
+	q_batches = vector<int>((int*)PyArray_DATA(q_batches_array), (int*)PyArray_DATA(q_batches_array) + Nb);
+	s_batches = vector<int>((int*)PyArray_DATA(s_batches_array), (int*)PyArray_DATA(s_batches_array) + Nb);
+
+	// Create result containers
+	vector<int> 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;
+}

cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp

@@ -0,0 +1,211 @@
+
+#include "grid_subsampling.h"
+
+
+void grid_subsampling(vector<PointXYZ>& original_points,
+                      vector<PointXYZ>& subsampled_points,
+                      vector<float>& original_features,
+                      vector<float>& subsampled_features,
+                      vector<int>& original_classes,
+                      vector<int>& 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<size_t, SampledData> 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<int, int>&a, const pair<int, int>&b){return a.second < b.second;})->first);
+		}
+	}
+
+	return;
+}
+
+
+void batch_grid_subsampling(vector<PointXYZ>& original_points,
+                              vector<PointXYZ>& subsampled_points,
+                              vector<float>& original_features,
+                              vector<float>& subsampled_features,
+                              vector<int>& original_classes,
+                              vector<int>& subsampled_classes,
+                              vector<int>& original_batches,
+                              vector<int>& 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<PointXYZ> b_o_points = vector<PointXYZ>(original_points.begin () + sum_b,
+	                                                   original_points.begin () + sum_b + original_batches[b]);
+
+        vector<float> b_o_features;
+        if (original_features.size() > 0)
+        {
+            b_o_features = vector<float>(original_features.begin () + sum_b * fdim,
+                                         original_features.begin () + (sum_b + original_batches[b]) * fdim);
+	    }
+
+	    vector<int> b_o_classes;
+        if (original_classes.size() > 0)
+        {
+            b_o_classes = vector<int>(original_classes.begin () + sum_b * ldim,
+                                      original_classes.begin () + sum_b + original_batches[b] * ldim);
+	    }
+
+
+        // Create result containers
+        vector<PointXYZ> b_s_points;
+        vector<float> b_s_features;
+        vector<int> 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;
+}

cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h

@@ -0,0 +1,101 @@
+
+
+#include "../../cpp_utils/cloud/cloud.h"
+
+#include <set>
+#include <cstdint>
+
+using namespace std;
+
+class SampledData
+{
+public:
+
+	// Elements
+	// ********
+
+	int count;
+	PointXYZ point;
+	vector<float> features;
+	vector<unordered_map<int, int>> labels;
+
+
+	// Methods
+	// *******
+
+	// Constructor
+	SampledData() 
+	{ 
+		count = 0; 
+		point = PointXYZ();
+	}
+
+	SampledData(const size_t fdim, const size_t ldim)
+	{
+		count = 0;
+		point = PointXYZ();
+	    features = vector<float>(fdim);
+	    labels = vector<unordered_map<int, int>>(ldim);
+	}
+
+	// Method Update
+	void update_all(const PointXYZ p, vector<float>::iterator f_begin, vector<int>::iterator l_begin)
+	{
+		count += 1;
+		point += p;
+		transform (features.begin(), features.end(), f_begin, features.begin(), plus<float>());
+		int i = 0;
+		for(vector<int>::iterator it = l_begin; it != l_begin + labels.size(); ++it)
+		{
+		    labels[i][*it] += 1;
+		    i++;
+		}
+		return;
+	}
+	void update_features(const PointXYZ p, vector<float>::iterator f_begin)
+	{
+		count += 1;
+		point += p;
+		transform (features.begin(), features.end(), f_begin, features.begin(), plus<float>());
+		return;
+	}
+	void update_classes(const PointXYZ p, vector<int>::iterator l_begin)
+	{
+		count += 1;
+		point += p;
+		int i = 0;
+		for(vector<int>::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<PointXYZ>& original_points,
+                      vector<PointXYZ>& subsampled_points,
+                      vector<float>& original_features,
+                      vector<float>& subsampled_features,
+                      vector<int>& original_classes,
+                      vector<int>& subsampled_classes,
+                      float sampleDl,
+                      int verbose);
+
+void batch_grid_subsampling(vector<PointXYZ>& original_points,
+                            vector<PointXYZ>& subsampled_points,
+                            vector<float>& original_features,
+                            vector<float>& subsampled_features,
+                            vector<int>& original_classes,
+                            vector<int>& subsampled_classes,
+                            vector<int>& original_batches,
+                            vector<int>& subsampled_batches,
+                            float sampleDl,
+                            int max_p);
+

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())
+
+
+
+
+
+
+
+

cpp_wrappers/cpp_subsampling/wrapper.cpp

@@ -0,0 +1,566 @@
+#include <Python.h>
+#include <numpy/arrayobject.h>
+#include "grid_subsampling/grid_subsampling.h"
+#include <string>
+
+
+
+// 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<PointXYZ> original_points;
+	vector<int> original_batches;
+	vector<float> original_features;
+	vector<int> original_classes;
+	original_points = vector<PointXYZ>((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N);
+	original_batches = vector<int>((int*)PyArray_DATA(batches_array), (int*)PyArray_DATA(batches_array) + Nb);
+	if (use_feature)
+		original_features = vector<float>((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N * fdim);
+	if (use_classes)
+		original_classes = vector<int>((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N * ldim);
+
+	// Subsample
+	vector<PointXYZ> subsampled_points;
+	vector<float> subsampled_features;
+	vector<int> subsampled_classes;
+	vector<int> 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<PointXYZ> original_points;
+	vector<float> original_features;
+	vector<int> original_classes;
+	original_points = vector<PointXYZ>((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N);
+	if (use_feature)
+		original_features = vector<float>((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N * fdim);
+	if (use_classes)
+		original_classes = vector<int>((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N * ldim);
+
+	// Subsample
+	vector<PointXYZ> subsampled_points;
+	vector<float> subsampled_features;
+	vector<int> 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;
+}

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<PointXYZ> 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<PointXYZ> 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;
+}

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 <vector>
+#include <unordered_map>
+#include <map>
+#include <algorithm>
+#include <numeric>
+#include <iostream>
+#include <iomanip>
+#include <cmath>
+
+#include <time.h>
+
+
+
+
+// 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<PointXYZ> points);
+PointXYZ min_point(std::vector<PointXYZ> points);
+
+
+struct PointCloud
+{
+
+	std::vector<PointXYZ>  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 <class BBOX>
+	bool kdtree_get_bbox(BBOX& /* bb */) const { return false; }
+
+};
+
+
+
+
+
+
+
+
+
+
+

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)

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
+
+
+
+
+
+
+
+
+
+
+
+
+

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)

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)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

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])
+

train_ModelNet40.py

@@ -0,0 +1,285 @@
+#
+#
+#      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.trainer import ModelTrainer
+from models.architectures import KPCNN
+
+
+# ----------------------------------------------------------------------------------------------------------------------
+#
+#           Config Class
+#       \******************/
+#
+
+class Modelnet40Config(Config):
+    """
+    Override the parameters you want to modify for this dataset
+    """
+
+    ####################
+    # Dataset parameters
+    ####################
+
+    # Dataset name
+    dataset = 'ModelNet40'
+
+    # 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_strided',
+                    'resnetb_deformable',
+                    'resnetb_deformable_strided',
+                    'resnetb_deformable',
+                    'global_average']
+
+    ###################
+    # 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 = 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 = 1
+
+    # 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
+    # 'fitting' helps deformed kernels to adapt to the geometry by penalizing distance to input points
+    offsets_loss = 'fitting'
+    offsets_decay = 1.0
+
+    #####################
+    # 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/80) for i in range(1, max_epoch)}
+    grad_clip_norm = 100.0
+
+    # Number of batch
+    batch_num = 16
+
+    # Number of steps per epochs
+    epoch_steps = 300
+
+    # 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, True, True]
+    augment_rotation = 'none'
+    augment_scale_min = 0.8
+    augment_scale_max = 1.2
+    augment_noise = 0.001
+    augment_color = 1.0
+
+    # 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'
+
+    # 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)
+
+
+

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)
+
+
+

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))
+

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()
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

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

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

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)
+
+
+