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

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

cpp/open3d/t/geometry/TensorMap.h

@@ -28,6 +28,7 @@
 
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 
 #include "open3d/core/Tensor.h"
 
@@ -114,6 +115,15 @@ class TensorMap : public std::unordered_map<std::string, core::Tensor> {
     /// Returns the primary key of the TensorMap.
     std::string GetPrimaryKey() const { return primary_key_; }
 
+    /// Returns a set with all keys.
+    std::unordered_set<std::string> GetKeySet() const {
+        std::unordered_set<std::string> 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;
 

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<vtkPlane> clipPlane;
     clipPlane->SetNormal(normal_.GetDataPtr<double>());
@@ -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<vtkQuadricDecimation> 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<vtkCleanPolyData> cleaner_A;
     cleaner_A->SetInputData(polydata_A);
 
-    auto polydata_B = CreateVtkPolyDataFromGeometry(mesh_B);
     vtkNew<vtkCleanPolyData> 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<vtkFillHolesFilter> fill_holes;
     fill_holes->SetInputData(polydata);
     fill_holes->SetHoleSize(hole_size);

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

cpp/open3d/t/geometry/VtkUtils.cpp

@@ -28,8 +28,10 @@
 
 #include <vtkArrayDispatch.h>
 #include <vtkCellArray.h>
+#include <vtkCellData.h>
 #include <vtkDoubleArray.h>
 #include <vtkFloatArray.h>
+#include <vtkPointData.h>
 #include <vtkPoints.h>
 
 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<std::string> include,
+        std::unordered_set<std::string> 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<vtkPolyData> CreateVtkPolyDataFromGeometry(
-        const Geometry& geometry, bool copy) {
+        const Geometry& geometry,
+        const std::unordered_set<std::string>& point_attr_include,
+        const std::unordered_set<std::string>& face_attr_include,
+        const std::unordered_set<std::string>& point_attr_exclude,
+        const std::unordered_set<std::string>& face_attr_exclude,
+        bool copy) {
     vtkSmartPointer<vtkPolyData> polydata = vtkSmartPointer<vtkPolyData>::New();
 
     if (geometry.GetGeometryType() == Geometry::GeometryType::PointCloud) {
@@ -299,14 +363,10 @@ vtkSmartPointer<vtkPolyData> 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<const TriangleMesh&>(geometry);
@@ -315,19 +375,12 @@ vtkSmartPointer<vtkPolyData> 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<vtkPolyData> 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;
 }
 

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<vtkPolyData> CreateVtkPolyDataFromGeometry(
-        const Geometry& geometry, bool copy = false);
+        const Geometry& geometry,
+        const std::unordered_set<std::string>& point_attr_include,
+        const std::unordered_set<std::string>& face_attr_include,
+        const std::unordered_set<std::string>& point_attr_exclude = {},
+        const std::unordered_set<std::string>& 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

cpp/tests/t/geometry/VtkUtils.cpp

@@ -64,7 +64,8 @@ TEST_P(VtkUtilsTest, PointCloudToVtkPolyData) {
 
     auto pcd = PointCloud(
             core::Tensor::Init<float>({{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();