From 2b285cf35057c600c1745a5bd1792366d7f6d2fb Mon Sep 17 00:00:00 2001 From: Benjamin Ummenhofer Date: Wed, 27 Jul 2022 13:02:53 +0200 Subject: [PATCH] Enable attributes in vtk wrapper functions (#5347) * -enable attributes in vtk wrapper functions -fix segfault for converting empty vtkPolyData to geometry * -disable attributes for quadric decimation -fix conversion for empty triangle indices * add GetKeySet() to TensorMap * change signature to explicitly request included attr keys --- cpp/open3d/t/geometry/PointCloud.h | 3 + cpp/open3d/t/geometry/TensorMap.h | 10 +++ cpp/open3d/t/geometry/TriangleMesh.cpp | 20 +++-- cpp/open3d/t/geometry/TriangleMesh.h | 6 ++ cpp/open3d/t/geometry/VtkUtils.cpp | 109 ++++++++++++++++++++----- cpp/open3d/t/geometry/VtkUtils.h | 20 ++++- cpp/tests/t/geometry/VtkUtils.cpp | 6 +- 7 files changed, 145 insertions(+), 29 deletions(-) diff --git a/cpp/open3d/t/geometry/PointCloud.h b/cpp/open3d/t/geometry/PointCloud.h index 31fdaba49e2..09c7472db8f 100644 --- a/cpp/open3d/t/geometry/PointCloud.h +++ b/cpp/open3d/t/geometry/PointCloud.h @@ -124,6 +124,9 @@ class PointCloud : public Geometry, public DrawableGeometry { /// Getter for point_attr_ TensorMap. Used in Pybind. const TensorMap &GetPointAttr() const { return point_attr_; } + /// Getter for point_attr_ TensorMap. + TensorMap &GetPointAttr() { return point_attr_; } + /// Get attributes. Throws exception if the attribute does not exist. /// /// \param key Attribute name. diff --git a/cpp/open3d/t/geometry/TensorMap.h b/cpp/open3d/t/geometry/TensorMap.h index 14af703b7aa..5706c7b030c 100644 --- a/cpp/open3d/t/geometry/TensorMap.h +++ b/cpp/open3d/t/geometry/TensorMap.h @@ -28,6 +28,7 @@ #include #include +#include #include "open3d/core/Tensor.h" @@ -114,6 +115,15 @@ class TensorMap : public std::unordered_map { /// Returns the primary key of the TensorMap. std::string GetPrimaryKey() const { return primary_key_; } + /// Returns a set with all keys. + std::unordered_set GetKeySet() const { + std::unordered_set keys; + for (const auto& item : *this) { + keys.insert(item.first); + } + return keys; + } + /// Returns true if all tensors in the map have the same size. bool IsSizeSynchronized() const; diff --git a/cpp/open3d/t/geometry/TriangleMesh.cpp b/cpp/open3d/t/geometry/TriangleMesh.cpp index 67a137af425..d1bd5a4b5f8 100644 --- a/cpp/open3d/t/geometry/TriangleMesh.cpp +++ b/cpp/open3d/t/geometry/TriangleMesh.cpp @@ -300,7 +300,9 @@ TriangleMesh TriangleMesh::ClipPlane(const core::Tensor &point, auto point_ = point.To(core::Device(), core::Float64).Contiguous(); auto normal_ = normal.To(core::Device(), core::Float64).Contiguous(); - auto polydata = CreateVtkPolyDataFromGeometry(*this); + auto polydata = CreateVtkPolyDataFromGeometry( + *this, GetVertexAttr().GetKeySet(), GetTriangleAttr().GetKeySet(), + {}, {}, false); vtkNew clipPlane; clipPlane->SetNormal(normal_.GetDataPtr()); @@ -324,7 +326,8 @@ TriangleMesh TriangleMesh::SimplifyQuadricDecimation( target_reduction); } - auto polydata = CreateVtkPolyDataFromGeometry(*this); + // exclude attributes because they will not be preserved + auto polydata = CreateVtkPolyDataFromGeometry(*this, {}, {}, {}, {}, false); vtkNew decimate; decimate->SetInputData(polydata); @@ -342,12 +345,16 @@ TriangleMesh BooleanOperation(const TriangleMesh &mesh_A, double tolerance, int op) { using namespace vtkutils; + // exclude triangle attributes because they will not be preserved + auto polydata_A = CreateVtkPolyDataFromGeometry( + mesh_A, mesh_A.GetVertexAttr().GetKeySet(), {}, {}, {}, false); + auto polydata_B = CreateVtkPolyDataFromGeometry( + mesh_B, mesh_B.GetVertexAttr().GetKeySet(), {}, {}, {}, false); + // clean meshes before passing them to the boolean operation - auto polydata_A = CreateVtkPolyDataFromGeometry(mesh_A); vtkNew cleaner_A; cleaner_A->SetInputData(polydata_A); - auto polydata_B = CreateVtkPolyDataFromGeometry(mesh_B); vtkNew cleaner_B; cleaner_B->SetInputData(polydata_B); @@ -384,7 +391,10 @@ TriangleMesh TriangleMesh::BooleanDifference(const TriangleMesh &mesh, TriangleMesh TriangleMesh::FillHoles(double hole_size) const { using namespace vtkutils; - auto polydata = CreateVtkPolyDataFromGeometry(*this); + // do not include triangle attributes because they will not be preserved by + // the hole filling algorithm + auto polydata = CreateVtkPolyDataFromGeometry( + *this, GetVertexAttr().GetKeySet(), {}, {}, {}, false); vtkNew fill_holes; fill_holes->SetInputData(polydata); fill_holes->SetHoleSize(hole_size); diff --git a/cpp/open3d/t/geometry/TriangleMesh.h b/cpp/open3d/t/geometry/TriangleMesh.h index ac6ea6d5bb1..160f75833a4 100644 --- a/cpp/open3d/t/geometry/TriangleMesh.h +++ b/cpp/open3d/t/geometry/TriangleMesh.h @@ -140,6 +140,9 @@ class TriangleMesh : public Geometry, public DrawableGeometry { /// Getter for vertex_attr_ TensorMap. Used in Pybind. const TensorMap &GetVertexAttr() const { return vertex_attr_; } + /// Getter for vertex_attr_ TensorMap. + TensorMap &GetVertexAttr() { return vertex_attr_; } + /// Get vertex attributes in vertex_attr_. Throws exception if the attribute /// does not exist. /// @@ -163,6 +166,9 @@ class TriangleMesh : public Geometry, public DrawableGeometry { /// Getter for triangle_attr_ TensorMap. Used in Pybind. const TensorMap &GetTriangleAttr() const { return triangle_attr_; } + /// Getter for triangle_attr_ TensorMap. + TensorMap &GetTriangleAttr() { return triangle_attr_; } + /// Get triangle attributes in triangle_attr_. Throws exception if the /// attribute does not exist. /// diff --git a/cpp/open3d/t/geometry/VtkUtils.cpp b/cpp/open3d/t/geometry/VtkUtils.cpp index b679e81a4fb..cb5e32ce88e 100644 --- a/cpp/open3d/t/geometry/VtkUtils.cpp +++ b/cpp/open3d/t/geometry/VtkUtils.cpp @@ -28,8 +28,10 @@ #include #include +#include #include #include +#include #include namespace open3d { @@ -281,8 +283,70 @@ static core::Tensor CreateTensorFromVtkCellArray(vtkCellArray* cells, return result.Reshape({num_cells, cell_size}); } +/// Adds point or cell attribute arrays to a TensorMap. +/// \param tmap The destination TensorMap. +/// \param field_data The source vtkFieldData. +/// \param copy If true always create a copy for attribute arrays. +static void AddVtkFieldDataToTensorMap(TensorMap& tmap, + vtkFieldData* field_data, + bool copy) { + for (int i = 0; i < field_data->GetNumberOfArrays(); ++i) { + auto array = field_data->GetArray(i); + std::string array_name = array->GetName(); + tmap[array_name] = CreateTensorFromVtkDataArray(array, copy); + } +} + +/// Adds attribute tensors to vtkFieldData. +/// Primary key tensors will be ignored by this function. +/// \param field_data The destination vtkFieldData. +/// \param tmap The source TensorMap. +/// \param copy If true always create a copy for attribute arrays. +/// \param include A set of keys to select which attributes should be added. +/// \param exclude A set of keys for which attributes will not be added to the +/// vtkFieldData. The exclusion set has precedence over the included keys. +static void AddTensorMapToVtkFieldData( + vtkFieldData* field_data, + TensorMap& tmap, + bool copy, + std::unordered_set include, + std::unordered_set exclude = {}) { + for (auto key_tensor : tmap) { + // we only want attributes and ignore the primary key here + if (key_tensor.first == tmap.GetPrimaryKey()) { + continue; + } + // we only support 2D tensors + if (key_tensor.second.NumDims() != 2) { + utility::LogWarning( + "Ignoring attribute '{}' for TensorMap with primary key " + "'{}' because of incompatible ndim={}", + key_tensor.first, tmap.GetPrimaryKey(), + key_tensor.second.NumDims()); + continue; + } + + if (include.count(key_tensor.first) && + !exclude.count(key_tensor.first)) { + auto array = CreateVtkDataArrayFromTensor(key_tensor.second, copy); + array->SetName(key_tensor.first.c_str()); + field_data->AddArray(array); + } else { + utility::LogWarning( + "Ignoring attribute '{}' for TensorMap with primary key " + "'{}'", + key_tensor.first, tmap.GetPrimaryKey()); + } + } +} + vtkSmartPointer CreateVtkPolyDataFromGeometry( - const Geometry& geometry, bool copy) { + const Geometry& geometry, + const std::unordered_set& point_attr_include, + const std::unordered_set& face_attr_include, + const std::unordered_set& point_attr_exclude, + const std::unordered_set& face_attr_exclude, + bool copy) { vtkSmartPointer polydata = vtkSmartPointer::New(); if (geometry.GetGeometryType() == Geometry::GeometryType::PointCloud) { @@ -299,14 +363,10 @@ vtkSmartPointer CreateVtkPolyDataFromGeometry( } polydata->SetVerts(cells); + AddTensorMapToVtkFieldData(polydata->GetPointData(), pcd.GetPointAttr(), + copy, point_attr_include, + point_attr_exclude); - for (auto key_tensor : pcd.GetPointAttr()) { - if (key_tensor.first != pcd.GetPointAttr().GetPrimaryKey()) { - utility::LogWarning("Ignoring point attribute {}", - key_tensor.first); - } - } - // TODO convert other data like normals, colors, ... } else if (geometry.GetGeometryType() == Geometry::GeometryType::TriangleMesh) { auto mesh = static_cast(geometry); @@ -315,19 +375,12 @@ vtkSmartPointer CreateVtkPolyDataFromGeometry( polydata->SetPolys( CreateVtkCellArrayFromTensor(mesh.GetTriangleIndices(), copy)); - for (auto key_tensor : mesh.GetVertexAttr()) { - if (key_tensor.first != mesh.GetVertexAttr().GetPrimaryKey()) { - utility::LogWarning("Ignoring vertex attribute {}", - key_tensor.first); - } - } - for (auto key_tensor : mesh.GetTriangleAttr()) { - if (key_tensor.first != mesh.GetTriangleAttr().GetPrimaryKey()) { - utility::LogWarning("Ignoring triangle attribute {}", - key_tensor.first); - } - } - // TODO convert other data like normals, colors, ... + AddTensorMapToVtkFieldData(polydata->GetPointData(), + mesh.GetVertexAttr(), copy, + point_attr_include, point_attr_exclude); + AddTensorMapToVtkFieldData(polydata->GetCellData(), + mesh.GetTriangleAttr(), copy, + face_attr_include, face_attr_exclude); } else { utility::LogError("Unsupported geometry type {}", geometry.GetGeometryType()); @@ -338,11 +391,25 @@ vtkSmartPointer CreateVtkPolyDataFromGeometry( TriangleMesh CreateTriangleMeshFromVtkPolyData(vtkPolyData* polydata, bool copy) { + if (!polydata->GetPoints()) { + return TriangleMesh(); + } core::Tensor vertices = CreateTensorFromVtkDataArray( polydata->GetPoints()->GetData(), copy); + core::Tensor triangles = CreateTensorFromVtkCellArray(polydata->GetPolys(), copy); + // Some algorithms return an empty tensor with shape (0,0). + // Fix the last dim here. + if (triangles.GetShape() == core::SizeVector{0, 0}) { + triangles = triangles.Reshape({0, 3}); + } TriangleMesh mesh(vertices, triangles); + + AddVtkFieldDataToTensorMap(mesh.GetVertexAttr(), polydata->GetPointData(), + copy); + AddVtkFieldDataToTensorMap(mesh.GetTriangleAttr(), polydata->GetCellData(), + copy); return mesh; } diff --git a/cpp/open3d/t/geometry/VtkUtils.h b/cpp/open3d/t/geometry/VtkUtils.h index 482a3d6875f..15e7d3cb686 100644 --- a/cpp/open3d/t/geometry/VtkUtils.h +++ b/cpp/open3d/t/geometry/VtkUtils.h @@ -46,8 +46,26 @@ int DtypeToVtkType(const core::Dtype& dtype); /// kept alive until the returned vtkPolyData object is deleted. /// \param geometry Open3D geometry object, e.g., a TriangleMesh. /// \param copy If true always create a copy of the data. +/// \param point_attr_include A set of keys to select which point/vertex +/// attributes should be added. Note that the primary key may be included and +/// will silently be ignored. +/// \param face_attr_include A set of keys to select +/// which face attributes should be added. Note that the primary key may be +/// included and will silently be ignored. +/// \param point_attr_exclude A set of keys for which point/vertex attributes +/// will not be added to the vtkPolyData. The exclusion set has precedence over +/// the included keys. +/// \param face_attr_exclude A set of keys for which face attributes will not be +/// added +/// to the vtkPolyData. The exclusion set has precedence over the included +/// keys. vtkSmartPointer CreateVtkPolyDataFromGeometry( - const Geometry& geometry, bool copy = false); + const Geometry& geometry, + const std::unordered_set& point_attr_include, + const std::unordered_set& face_attr_include, + const std::unordered_set& point_attr_exclude = {}, + const std::unordered_set& face_attr_exclude = {}, + bool copy = false); /// Creates a triangle mesh from a vtkPolyData object. /// The returned TriangleMesh may directly use the memory of the data arrays in diff --git a/cpp/tests/t/geometry/VtkUtils.cpp b/cpp/tests/t/geometry/VtkUtils.cpp index c309932e5dd..ab34f1bb2c4 100644 --- a/cpp/tests/t/geometry/VtkUtils.cpp +++ b/cpp/tests/t/geometry/VtkUtils.cpp @@ -64,7 +64,8 @@ TEST_P(VtkUtilsTest, PointCloudToVtkPolyData) { auto pcd = PointCloud( core::Tensor::Init({{0, 0, 0}, {1, 1, 1}, {2, 2, 2}})); - auto polydata = vtkutils::CreateVtkPolyDataFromGeometry(pcd, copy); + auto polydata = + vtkutils::CreateVtkPolyDataFromGeometry(pcd, {}, {}, {}, {}, copy); auto tensor = pcd.GetPointPositions(); auto data_array = polydata->GetPoints()->GetData(); @@ -86,7 +87,8 @@ TEST_P(VtkUtilsTest, TriangleMeshToVtkPolyData) { auto legacy_box = geometry::TriangleMesh::CreateBox(); auto box = TriangleMesh::FromLegacy(*legacy_box); - auto polydata = vtkutils::CreateVtkPolyDataFromGeometry(box, copy); + auto polydata = + vtkutils::CreateVtkPolyDataFromGeometry(box, {}, {}, {}, {}, copy); auto tensor = box.GetVertexPositions(); auto data_array = polydata->GetPoints()->GetData();