| #pragma once |
| #include <pybind11/pybind11.h> |
| #include <pybind11/numpy.h> |
| #include <pybind11/eigen.h> |
| #include <cpu/fill_holes.h> |
| #include <cpu/merge_vertices.h> |
| |
| #include <cpu/hashtable.h> |
| #include <cpu/spmc.h> |
| |
| #include <cpu/decimation.h> |
|
|
| #include <vector> |
| #include <cstring> |
|
|
| namespace py = pybind11; |
|
|
| namespace cubvh { |
|
|
| static py::array_t<int> fill_holes( |
| py::array_t<float, py::array::c_style | py::array::forcecast> vertices, |
| py::array_t<int, py::array::c_style | py::array::forcecast> faces, |
| bool return_added, |
| bool check_containment, |
| double eps_d, |
| bool verbose) { |
| |
| auto vbuf = vertices.request(); |
| auto fbuf = faces.request(); |
| if (!(vbuf.ndim == 2 && vbuf.shape[1] == 3)) { |
| throw std::runtime_error("vertices must be of shape [N,3]"); |
| } |
| if (!(fbuf.ndim == 2 && fbuf.shape[1] == 3)) { |
| throw std::runtime_error("faces must be of shape [M,3]"); |
| } |
|
|
| const size_t N = static_cast<size_t>(vbuf.shape[0]); |
| const size_t M = static_cast<size_t>(fbuf.shape[0]); |
|
|
| const float* vptr = static_cast<const float*>(vbuf.ptr); |
| const int* fptr = static_cast<const int*>(fbuf.ptr); |
|
|
| std::vector<Eigen::Vector3f> V(N); |
| for (size_t i = 0; i < N; ++i) { |
| V[i] = Eigen::Vector3f(vptr[3*i+0], vptr[3*i+1], vptr[3*i+2]); |
| } |
|
|
| std::vector<Eigen::Vector3i> F(M); |
| for (size_t i = 0; i < M; ++i) { |
| F[i] = Eigen::Vector3i(fptr[3*i+0], fptr[3*i+1], fptr[3*i+2]); |
| } |
|
|
| cubvh::cpu::HoleFillOptions opt; |
| opt.checkContainment = check_containment; |
| opt.eps = static_cast<float>(eps_d); |
| opt.verbose = verbose; |
|
|
| if (return_added) { |
| auto added = cubvh::cpu::fill_holes(V, F, opt); |
| py::array_t<int> out({(py::ssize_t)added.size(), (py::ssize_t)3}); |
| auto obuf = out.request(); |
| int* optr = static_cast<int*>(obuf.ptr); |
| for (size_t i = 0; i < added.size(); ++i) { |
| optr[3*i+0] = added[i][0]; |
| optr[3*i+1] = added[i][1]; |
| optr[3*i+2] = added[i][2]; |
| } |
| return out; |
| } else { |
| cubvh::cpu::fill_holes_inplace(V, F, opt); |
| py::array_t<int> out({(py::ssize_t)F.size(), (py::ssize_t)3}); |
| auto obuf = out.request(); |
| int* optr = static_cast<int*>(obuf.ptr); |
| for (size_t i = 0; i < F.size(); ++i) { |
| optr[3*i+0] = F[i][0]; |
| optr[3*i+1] = F[i][1]; |
| optr[3*i+2] = F[i][2]; |
| } |
| return out; |
| } |
| } |
|
|
| |
| static std::pair<py::array_t<float>, py::array_t<int>> merge_vertices( |
| py::array_t<float, py::array::c_style | py::array::forcecast> vertices, |
| py::array_t<int, py::array::c_style | py::array::forcecast> faces, |
| double threshold_d) { |
| auto vbuf = vertices.request(); |
| auto fbuf = faces.request(); |
| if (!(vbuf.ndim == 2 && vbuf.shape[1] == 3)) { |
| throw std::runtime_error("vertices must be of shape [N,3]"); |
| } |
| if (!(fbuf.ndim == 2 && fbuf.shape[1] == 3)) { |
| throw std::runtime_error("faces must be of shape [M,3]"); |
| } |
| const size_t N = static_cast<size_t>(vbuf.shape[0]); |
| const size_t M = static_cast<size_t>(fbuf.shape[0]); |
| const float* vptr = static_cast<const float*>(vbuf.ptr); |
| const int* fptr = static_cast<const int*>(fbuf.ptr); |
|
|
| std::vector<Eigen::Vector3f> V(N); |
| for (size_t i=0;i<N;++i) { |
| V[i] = Eigen::Vector3f(vptr[3*i+0], vptr[3*i+1], vptr[3*i+2]); |
| } |
| std::vector<Eigen::Vector3i> F(M); |
| for (size_t i=0;i<M;++i) { |
| F[i] = Eigen::Vector3i(fptr[3*i+0], fptr[3*i+1], fptr[3*i+2]); |
| } |
|
|
| std::vector<Eigen::Vector3f> V_out; std::vector<Eigen::Vector3i> F_out; |
| cubvh::cpu::merge_vertices(V, F, static_cast<float>(threshold_d), V_out, F_out); |
|
|
| py::array_t<float> v_out({(py::ssize_t)V_out.size(), (py::ssize_t)3}); |
| py::array_t<int> f_out({(py::ssize_t)F_out.size(), (py::ssize_t)3}); |
| auto vObuf = v_out.request(); auto fObuf = f_out.request(); |
| float* vO = static_cast<float*>(vObuf.ptr); |
| int* fO = static_cast<int*>(fObuf.ptr); |
| for (size_t i=0;i<V_out.size();++i) { |
| vO[3*i+0] = V_out[i][0]; vO[3*i+1] = V_out[i][1]; vO[3*i+2] = V_out[i][2]; |
| } |
| for (size_t i=0;i<F_out.size();++i) { |
| fO[3*i+0] = F_out[i][0]; fO[3*i+1] = F_out[i][1]; fO[3*i+2] = F_out[i][2]; |
| } |
| return {v_out, f_out}; |
| } |
|
|
| |
| static std::pair<py::array_t<float>, py::array_t<int>> sparse_marching_cubes_cpu( |
| py::array_t<int, py::array::c_style | py::array::forcecast> coords, |
| py::array_t<float, py::array::c_style | py::array::forcecast> corners, |
| double iso_d, |
| bool ensure_consistency = false) { |
|
|
| auto cbuf = coords.request(); |
| auto vbuf = corners.request(); |
| if (!(cbuf.ndim == 2 && cbuf.shape[1] == 3)) { |
| throw std::runtime_error("coords must be of shape [N,3] (int32)"); |
| } |
| if (!(vbuf.ndim == 2 && vbuf.shape[1] == 8)) { |
| throw std::runtime_error("corners must be of shape [N,8] (float32)"); |
| } |
| if (cbuf.shape[0] != vbuf.shape[0]) { |
| throw std::runtime_error("coords and corners must have the same first dimension N"); |
| } |
|
|
| const int N = static_cast<int>(cbuf.shape[0]); |
| const int* cptr = static_cast<const int*>(cbuf.ptr); |
| const float* fptr = static_cast<const float*>(vbuf.ptr); |
| const float iso = static_cast<float>(iso_d); |
|
|
| auto mesh = cubvh::cpu::sparse_marching_cubes(cptr, fptr, N, iso, ensure_consistency); |
| const auto& V = mesh.first; |
| const auto& F = mesh.second; |
|
|
| |
| py::array_t<float> v_out({(py::ssize_t)V.size(), (py::ssize_t)3}); |
| py::array_t<int> f_out({(py::ssize_t)F.size(), (py::ssize_t)3}); |
| auto vObuf = v_out.request(); |
| auto fObuf = f_out.request(); |
| float* vO = static_cast<float*>(vObuf.ptr); |
| int* fO = static_cast<int*>(fObuf.ptr); |
|
|
| for (size_t i = 0; i < V.size(); ++i) { |
| vO[3*i+0] = V[i].x; |
| vO[3*i+1] = V[i].y; |
| vO[3*i+2] = V[i].z; |
| } |
| for (size_t i = 0; i < F.size(); ++i) { |
| fO[3*i+0] = F[i].v0; |
| fO[3*i+1] = F[i].v1; |
| fO[3*i+2] = F[i].v2; |
| } |
|
|
| return {v_out, f_out}; |
| } |
|
|
| |
| template <typename T> |
| static cubvh::cpu::qd::MeshT<T> _mesh_from_numpy_typed( |
| py::array_t<T, py::array::c_style | py::array::forcecast> vertices, |
| py::array_t<int, py::array::c_style | py::array::forcecast> faces) |
| { |
| if (vertices.ndim() != 2 || vertices.shape(1) != 3) |
| throw std::runtime_error("vertices must be (N,3) array"); |
| if (faces.ndim() != 2 || faces.shape(1) != 3) |
| throw std::runtime_error("faces must be (M,3) int32 array"); |
|
|
| cubvh::cpu::qd::MeshT<T> m; |
| m.vertices.reserve(vertices.shape(0)); |
| auto vbuf = vertices.template unchecked<2>(); |
| for (ssize_t i = 0; i < vertices.shape(0); ++i) { |
| m.vertices.emplace_back(vbuf(i,0), vbuf(i,1), vbuf(i,2)); |
| } |
| m.faces.reserve(faces.shape(0)); |
| auto fbuf = faces.template unchecked<2>(); |
| for (ssize_t i = 0; i < faces.shape(0); ++i) { |
| m.faces.push_back({ fbuf(i,0), fbuf(i,1), fbuf(i,2) }); |
| } |
| return m; |
| } |
|
|
| template <typename T> |
| static std::pair<py::array_t<T>, py::array_t<int>> _mesh_to_numpy_typed(const cubvh::cpu::qd::MeshT<T>& m) |
| { |
| py::array_t<T> V({ (ssize_t)m.vertices.size(), (ssize_t)3 }); |
| py::array_t<int> F({ (ssize_t)m.faces.size(), (ssize_t)3 }); |
| auto vbuf = V.template mutable_unchecked<2>(); |
| for (ssize_t i = 0; i < (ssize_t)m.vertices.size(); ++i) { |
| vbuf(i,0) = m.vertices[i].x; |
| vbuf(i,1) = m.vertices[i].y; |
| vbuf(i,2) = m.vertices[i].z; |
| } |
| auto fbuf = F.template mutable_unchecked<2>(); |
| for (ssize_t i = 0; i < (ssize_t)m.faces.size(); ++i) { |
| fbuf(i,0) = m.faces[i][0]; |
| fbuf(i,1) = m.faces[i][1]; |
| fbuf(i,2) = m.faces[i][2]; |
| } |
| return { V, F }; |
| } |
|
|
| static py::tuple decimate(py::array vertices, |
| py::array faces, |
| int target_vertices) |
| { |
| py::dtype dt = vertices.dtype(); |
| if (dt.is(py::dtype::of<float>())){ |
| auto v = vertices.cast<py::array_t<float, py::array::c_style | py::array::forcecast>>(); |
| auto f = faces.cast<py::array_t<int, py::array::c_style | py::array::forcecast>>(); |
| auto mesh = _mesh_from_numpy_typed<float>(v, f); |
| cubvh::cpu::qd::DecimatorT<float> dec(mesh); |
| dec.decimate(target_vertices); |
| auto out = _mesh_to_numpy_typed<float>(dec.mesh()); |
| return py::make_tuple(out.first, out.second); |
| } else if (dt.is(py::dtype::of<double>())){ |
| auto v = vertices.cast<py::array_t<double, py::array::c_style | py::array::forcecast>>(); |
| auto f = faces.cast<py::array_t<int, py::array::c_style | py::array::forcecast>>(); |
| auto mesh = _mesh_from_numpy_typed<double>(v, f); |
| cubvh::cpu::qd::DecimatorT<double> dec(mesh); |
| dec.decimate(target_vertices); |
| auto out = _mesh_to_numpy_typed<double>(dec.mesh()); |
| return py::make_tuple(out.first, out.second); |
| } else { |
| throw std::runtime_error("vertices must be float32 or float64 array"); |
| } |
| } |
|
|
| static py::tuple parallel_decimate(py::array vertices, |
| py::array faces, |
| int target_vertices) |
| { |
| py::dtype dt = vertices.dtype(); |
| if (dt.is(py::dtype::of<float>())){ |
| auto v = vertices.cast<py::array_t<float, py::array::c_style | py::array::forcecast>>(); |
| auto f = faces.cast<py::array_t<int, py::array::c_style | py::array::forcecast>>(); |
| auto mesh = _mesh_from_numpy_typed<float>(v, f); |
| cubvh::cpu::qd::DecimatorT<float> dec(mesh); |
| dec.parallelDecimate(target_vertices); |
| auto out = _mesh_to_numpy_typed<float>(dec.mesh()); |
| return py::make_tuple(out.first, out.second); |
| } else if (dt.is(py::dtype::of<double>())){ |
| auto v = vertices.cast<py::array_t<double, py::array::c_style | py::array::forcecast>>(); |
| auto f = faces.cast<py::array_t<int, py::array::c_style | py::array::forcecast>>(); |
| auto mesh = _mesh_from_numpy_typed<double>(v, f); |
| cubvh::cpu::qd::DecimatorT<double> dec(mesh); |
| dec.parallelDecimate(target_vertices); |
| auto out = _mesh_to_numpy_typed<double>(dec.mesh()); |
| return py::make_tuple(out.first, out.second); |
| } else { |
| throw std::runtime_error("vertices must be float32 or float64 array"); |
| } |
| } |
|
|
| |
| class HashTable { |
| public: |
| HashTable() {} |
|
|
| void set_num_dims(int d) { ht.set_num_dims(d); } |
| int get_num_dims() const { return ht.get_num_dims(); } |
| void resize(int capacity) { ht.resize(capacity); } |
| void prepare() { ht.prepare(); } |
|
|
| void insert(at::Tensor coords) { |
| TORCH_CHECK(!coords.is_cuda(), "coords must reside on CPU"); |
| TORCH_CHECK(coords.dtype() == at::kInt, "coords must be int32"); |
| TORCH_CHECK(coords.dim() == 2, "coords must be 2D [N,D]"); |
| coords_ref_ = coords.contiguous(); |
| const int N = (int)coords_ref_.size(0); |
| const int D = (int)coords_ref_.size(1); |
| ht.set_num_dims(D); |
| ht.insert(coords_ref_.data_ptr<int>(), N); |
| } |
|
|
| void build(at::Tensor coords) { |
| TORCH_CHECK(!coords.is_cuda(), "coords must reside on CPU"); |
| TORCH_CHECK(coords.dtype() == at::kInt, "coords must be int32"); |
| TORCH_CHECK(coords.dim() == 2, "coords must be 2D [N,D]"); |
| coords_ref_ = coords.contiguous(); |
| const int N = (int)coords_ref_.size(0); |
| const int D = (int)coords_ref_.size(1); |
| ht.set_num_dims(D); |
| ht.build(coords_ref_.data_ptr<int>(), N); |
| } |
|
|
| at::Tensor search(at::Tensor queries) const { |
| TORCH_CHECK(!queries.is_cuda(), "queries must reside on CPU"); |
| TORCH_CHECK(queries.dtype() == at::kInt, "queries must be int32"); |
| TORCH_CHECK(queries.dim() == 2, "queries must be 2D [M,D]"); |
| TORCH_CHECK(ht.capacity > 0, "hash table is not built"); |
| at::Tensor q = queries.contiguous(); |
| const int M = (int)q.size(0); |
| auto opts_i = torch::TensorOptions().dtype(torch::kInt32).device(q.device()); |
| at::Tensor out = at::empty({M}, opts_i); |
| ht.search(q.data_ptr<int>(), M, out.data_ptr<int>()); |
| return out; |
| } |
|
|
| private: |
| mutable HashTableIntCPU ht; |
| at::Tensor coords_ref_; |
| }; |
|
|
| } |
