|
| 1 | +#include <pybind11/pybind11.h> |
| 2 | +#include <pybind11/numpy.h> |
| 3 | +#include <pybind11/stl.h> |
| 4 | +#include <vector> |
| 5 | +#include "src/pointcloud.h" |
| 6 | +#include "src/option.h" |
| 7 | +#include "src/cluster.h" |
| 8 | +#include "src/dnn.h" |
| 9 | +#include "src/graph.h" |
| 10 | +#include "src/postprocess.h" |
| 11 | + |
| 12 | +namespace py = pybind11; |
| 13 | + |
| 14 | +class tripclust { |
| 15 | +private: |
| 16 | + Opt opt; |
| 17 | + PointCloud cloud_xyz; |
| 18 | + cluster_group cl_group; |
| 19 | + bool postprocess; |
| 20 | + int min_depth; |
| 21 | + |
| 22 | +public: |
| 23 | + tripclust(); |
| 24 | + void fill_pointcloud(py::array_t<double>); |
| 25 | + void perform_clustering(); |
| 26 | + py::array_t<unsigned short> get_cluster(int); |
| 27 | + py::array_t<short> get_labels(); |
| 28 | + int get_number_of_clusters(); |
| 29 | + // neighbour distance for smoothing |
| 30 | + void set_r(const double value) {opt.r = value;} |
| 31 | + void set_rdnn(const bool value) {opt.rdnn = value;} // compute r with dnn |
| 32 | + // tested neighbours of triplet mid point |
| 33 | + void set_k(const size_t value) {opt.k = value;} |
| 34 | + // max number of triplets to one mid point |
| 35 | + void set_n(const size_t value) {opt.n = value;} |
| 36 | + // 1 - cos alpha, where alpha is the angle between the two triplet branches |
| 37 | + void set_a(const double value) {opt.a = value;} |
| 38 | + // distance scale factor in metric |
| 39 | + void set_s(const double value) {opt.s = value;} |
| 40 | + void set_sdnn(const bool value) {opt.sdnn = value;} // compute s with dnn |
| 41 | + // threshold for cdist in clustering |
| 42 | + void set_t(const double value) {opt.t = value;} |
| 43 | + void set_tauto(const bool value) {opt.tauto = value;} // auto generate t |
| 44 | + // maximum gap width |
| 45 | + void set_dmax(const double value) {opt.dmax = value;} |
| 46 | + void set_dmax_dnn(const bool value) {opt.dmax_dnn = value;} // use dnn for dmax |
| 47 | + void set_ordered(const bool value) {opt.ordered = value;} // points are in chronological order |
| 48 | + // linkage method for clustering |
| 49 | + void set_link(const Linkage value) {opt.link = value;} |
| 50 | + // min number of triplets per cluster |
| 51 | + void set_m(const size_t value) {opt.m = value;} |
| 52 | + // whether or not post processing should be enabled |
| 53 | + void set_postprocess(const bool value) {postprocess = value;} |
| 54 | + // minimum number of points making a branch in curve in post processing |
| 55 | + void set_min_depth(const int value) {min_depth = value;} |
| 56 | +}; |
| 57 | + |
| 58 | +tripclust::tripclust() { |
| 59 | + Opt opt; |
| 60 | + PointCloud pc; |
| 61 | + cluster_group cg; |
| 62 | + this->opt = opt; |
| 63 | + this->cloud_xyz = pc; |
| 64 | + this->cl_group = cg; |
| 65 | + this->postprocess = false; |
| 66 | + this->min_depth = 25; |
| 67 | +} |
| 68 | + |
| 69 | +// Load cloud_xyz from numpy array |
| 70 | +void tripclust::fill_pointcloud(py::array_t<double> arr) { |
| 71 | + py::buffer_info buf_info = arr.request(); |
| 72 | + double *ptr = static_cast<double*>(buf_info.ptr); |
| 73 | + size_t ndim = buf_info.ndim; |
| 74 | + std::vector<ssize_t> shape = buf_info.shape; |
| 75 | + |
| 76 | + // The numpy array is assumed to be a 2D cloud array from Spyral |
| 77 | + for (size_t i = 0; i < shape[0]; i++) { |
| 78 | + Point point; |
| 79 | + point.x = ptr[i*shape[1]+0]; // x coordinate |
| 80 | + point.y = ptr[i*shape[1]+1]; // y coordinate |
| 81 | + point.z = ptr[i*shape[1]+2]; // z coordinate |
| 82 | + this->cloud_xyz.push_back(point); |
| 83 | + } |
| 84 | + // this->cloud.setOrdered(true); |
| 85 | + // this->cloud_xyz.set2d(false); |
| 86 | +} |
| 87 | + |
| 88 | +void tripclust::perform_clustering() { |
| 89 | + if (this->opt.needs_dnn()) { |
| 90 | + double dnn = std::sqrt(first_quartile(this->cloud_xyz)); |
| 91 | + this->opt.set_dnn(dnn); |
| 92 | + if (dnn == 0.0) { |
| 93 | + return; // Need to throw some kind of error |
| 94 | + } |
| 95 | + } |
| 96 | + // Step 1) smoothing by position averaging of neighboring points |
| 97 | + PointCloud cloud_xyz_smooth; |
| 98 | + smoothen_cloud(this->cloud_xyz, cloud_xyz_smooth, this->opt.get_r()); |
| 99 | + |
| 100 | + // Step 2) finding triplets of approximately collinear points |
| 101 | + std::vector<triplet> triplets; |
| 102 | + generate_triplets(cloud_xyz_smooth, triplets, this->opt.get_k(), |
| 103 | + this->opt.get_n(), this->opt.get_a()); |
| 104 | + |
| 105 | + // Step 3) single link hierarchical clustering of the triplets |
| 106 | + // cluster_group cl_group; |
| 107 | + compute_hc(cloud_xyz_smooth, this->cl_group, triplets, this->opt.get_s(), |
| 108 | + this->opt.get_t(), this->opt.is_tauto(), this->opt.get_dmax(), |
| 109 | + this->opt.is_dmax(), this->opt.get_linkage(), 0); |
| 110 | + |
| 111 | + // Step 4) pruning by removal of small clusters ... |
| 112 | + cleanup_cluster_group(this->cl_group, this->opt.get_m(), 0); |
| 113 | + cluster_triplets_to_points(triplets, this->cl_group); |
| 114 | + // .. and (optionally) by splitting up clusters at gaps > dmax |
| 115 | + if (this->opt.is_dmax()) { |
| 116 | + cluster_group cleaned_up_cluster_group; |
| 117 | + for (cluster_group::iterator cl = this->cl_group.begin(); cl != this->cl_group.end(); ++cl) { |
| 118 | + max_step(cleaned_up_cluster_group, *cl, this->cloud_xyz, this->opt.get_dmax(), |
| 119 | + this->opt.get_m() + 2); |
| 120 | + } |
| 121 | + this->cl_group = cleaned_up_cluster_group; |
| 122 | + } |
| 123 | + // store cluster labels in points |
| 124 | + add_clusters(this->cloud_xyz, this->cl_group, 0); |
| 125 | + // Step 5) Optionally perform post processing to find sub clusters |
| 126 | + if (this->postprocess) { |
| 127 | + int nchanged; |
| 128 | + nchanged = process_pointcloud(this->cloud_xyz, this->min_depth, 0); |
| 129 | + } |
| 130 | +} |
| 131 | + |
| 132 | +// Get an individual cluster array |
| 133 | +py::array_t<unsigned short> tripclust::get_cluster(int clu) { |
| 134 | + std::vector<unsigned short> cluster; |
| 135 | + for (size_t i=0; i<this->cl_group[clu].size(); ++i) |
| 136 | + cluster.push_back(this->cl_group[clu][i]); |
| 137 | + return py::array_t<unsigned short>(cluster.size(), cluster.data()); |
| 138 | +} |
| 139 | + |
| 140 | +// Get an array containing the cluster labels for each point |
| 141 | +py::array_t<short> tripclust::get_labels() { |
| 142 | + std::vector<short> labels(this->cloud_xyz.size(), -1); |
| 143 | + for (size_t i=0; i<this->cloud_xyz.size(); ++i) { |
| 144 | + std::set<size_t>::iterator iterator=this->cloud_xyz[i].cluster_ids.begin(); |
| 145 | + if (this->cloud_xyz[i].cluster_ids.size() > 0) labels[i] = *iterator; // Choose the first cluster id |
| 146 | + } |
| 147 | + // |
| 148 | + // for (size_t clu=0; clu<this->cl_group.size(); ++clu) { |
| 149 | + // for (size_t i=0; i<this->cl_group[clu].size(); ++i) { |
| 150 | + // labels[this->cl_group[clu][i]] = clu; |
| 151 | + // } |
| 152 | + // } |
| 153 | + return py::array_t<short>(labels.size(), labels.data()); |
| 154 | +} |
| 155 | + |
| 156 | +// Get the number of clusters |
| 157 | +int tripclust::get_number_of_clusters(){ |
| 158 | + int nclu = 0; |
| 159 | + std::vector<bool> test(this->cloud_xyz.size(), false); |
| 160 | + for (size_t i=0; i<this->cloud_xyz.size(); ++i) { |
| 161 | + for (std::set<size_t>::iterator j=this->cloud_xyz[i].cluster_ids.begin(); j!=this->cloud_xyz[i].cluster_ids.end(); ++j) |
| 162 | + test[*j] = true; |
| 163 | + } |
| 164 | + for (size_t i=0; i<test.size(); ++i) if (test[i]) nclu++; |
| 165 | + return nclu; |
| 166 | + // return this->cl_group.size(); |
| 167 | +} |
| 168 | + |
| 169 | +PYBIND11_MODULE(tripclust, m, py::mod_gil_not_used()) { |
| 170 | + py::class_<tripclust>(m, "tripclust") |
| 171 | + .def(py::init()) |
| 172 | + .def("fill_pointcloud", &tripclust::fill_pointcloud, "Fills point cloud") |
| 173 | + .def("perform_clustering", &tripclust::perform_clustering, "Performs the tripclust clustering") |
| 174 | + .def("get_number_of_clusters", &tripclust::get_number_of_clusters, "Get the number of clusters") |
| 175 | + .def("get_cluster", &tripclust::get_cluster, "Retrieves individual cluster") |
| 176 | + .def("get_labels", &tripclust::get_labels, "Returns an array containing the label for each point") |
| 177 | + .def("set_r", &tripclust::set_r, "neighbour distance for smoothing") |
| 178 | + .def("set_rdnn", &tripclust::set_rdnn, "compute r with dnn") |
| 179 | + .def("set_k", &tripclust::set_k, "tested neighbours of triplet mid point") |
| 180 | + .def("set_n", &tripclust::set_n, "max number of triplets to one mid point") |
| 181 | + .def("set_a", &tripclust::set_a, "1 - cos alpha, where alpha is the angle between the two triplet branches") |
| 182 | + .def("set_s", &tripclust::set_s, "distance scale factor in metric") |
| 183 | + .def("set_sdnn", &tripclust::set_sdnn, "compute s with dnn") |
| 184 | + .def("set_t", &tripclust::set_t, "threshold for cdist in clustering") |
| 185 | + .def("set_tauto", &tripclust::set_tauto, "auto generate t") |
| 186 | + .def("set_dmax", &tripclust::set_dmax, "maximum gap width") |
| 187 | + .def("set_dmax_dnn", &tripclust::set_dmax_dnn, "use dnn for dmax") |
| 188 | + .def("set_ordered", &tripclust::set_ordered, "points are in chronological order") |
| 189 | + .def("set_link", &tripclust::set_link, "linkage method for clustering") |
| 190 | + .def("set_m", &tripclust::set_m, "min number of triplets per cluster") |
| 191 | + .def("set_postprocess", &tripclust::set_postprocess, "whether or not post processing should be enabled") |
| 192 | + .def("set_min_depth", &tripclust::set_min_depth, "minimum number of points making a branch in curve in post processing"); |
| 193 | +} |
0 commit comments