Implement open3d::t::geometry::TriangleMesh::SelectByIndex (#6415)

* Define a helper DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE

This is a helper to call a templated function with an integer argument,
based on Dtype.  As a second argument, it takes a suffix, used to build
a unique type name.  This way, we can use it to call a function with
more than one integer argument.

Example:

DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE(core::Dtype::Int32, int32, [&]() {
    DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE(core::Dtype::UInt64, uint64, [&]() {
    scalar_int32_t a;
    scalar_uint64_t b;
    // ...
});

* Implement open3d::t::geometry::TriangleMesh::SelectByIndex

The method takes a list of indices and returns a new mesh built with
the selected vertices and triangles formed by these vertices.
The indices type can be any integral type.  The algorithm is implemented
on CPU only.

The implementation is inspired by
  open3d::geometry::TriangleMesh::SelectByIndex.
and by
  open3d::t::geometry::TriangleMesh::SelectFacesByMask.

We first compute a mask of vertices to be selected.  If the input index
exceeds the maximum number of vertices or is negative, we ignore the
index and print a warning.

If the mesh has triangles, we build  tringle mask and select needed
triangles.

The next step is to update triangle indices to a new ones.  It is similar
to SelectFacesByMask, so I introduced a static helper to do that.  Based on
the vertex mask we build a mapping index vector using inclusive prefix sum
algorithm and use it as a map between old and new indices.

We select the vertices by mask and build the resulting mesh from the selected
vertices and triangles.

Copying the mesh attributes is again similar to SelectFacesByMask, so I put
it to a separate static function.

* Refactor t::geometry::TriangleMesh::SelectFacesByMask

* Add error handling on empty mesh
* Use DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE instead of a conditional
  branch
* Use UpdateTriangleIndicesByVertexMask helper to update triangle
  indices
* Use CopyAttributesByMask helper to copy the mesh attributes
* Add tests

cpp/open3d/core/Dispatch.h

@@ -113,3 +113,34 @@
             open3d::utility::LogError("Unsupported data type.");  \
         }                                                         \
     }()
+
+#define DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE(DTYPE, PREFIX, ...) \
+    [&] {                                                         \
+        if (DTYPE == open3d::core::Int8) {                        \
+            using scalar_##PREFIX##_t = int8_t;                   \
+            return __VA_ARGS__();                                 \
+        } else if (DTYPE == open3d::core::Int16) {                \
+            using scalar_##PREFIX##_t = int16_t;                  \
+            return __VA_ARGS__();                                 \
+        } else if (DTYPE == open3d::core::Int32) {                \
+            using scalar_##PREFIX##_t = int32_t;                  \
+            return __VA_ARGS__();                                 \
+        } else if (DTYPE == open3d::core::Int64) {                \
+            using scalar_##PREFIX##_t = int64_t;                  \
+            return __VA_ARGS__();                                 \
+        } else if (DTYPE == open3d::core::UInt8) {                \
+            using scalar_##PREFIX##_t = uint8_t;                  \
+            return __VA_ARGS__();                                 \
+        } else if (DTYPE == open3d::core::UInt16) {               \
+            using scalar_##PREFIX##_t = uint16_t;                 \
+            return __VA_ARGS__();                                 \
+        } else if (DTYPE == open3d::core::UInt32) {               \
+            using scalar_##PREFIX##_t = uint32_t;                 \
+            return __VA_ARGS__();                                 \
+        } else if (DTYPE == open3d::core::UInt64) {               \
+            using scalar_##PREFIX##_t = uint64_t;                 \
+            return __VA_ARGS__();                                 \
+        } else {                                                  \
+            open3d::utility::LogError("Unsupported data type.");  \
+        }                                                         \
+    }()

cpp/open3d/t/geometry/TriangleMesh.cpp

@@ -995,7 +995,66 @@ int TriangleMesh::PCAPartition(int max_faces) {
     return num_parititions;
 }
 
+/// A helper to compute new vertex indices out of vertex mask.
+/// \param tris_cpu tensor with triangle indices to update.
+/// \param vertex_mask tensor with the mask for vertices.
+template <typename T>
+static void UpdateTriangleIndicesByVertexMask(core::Tensor &tris_cpu,
+                                              const core::Tensor &vertex_mask) {
+    int64_t num_verts = vertex_mask.GetLength();
+    int64_t num_tris = tris_cpu.GetLength();
+    const T *vertex_mask_ptr = vertex_mask.GetDataPtr<T>();
+    std::vector<T> prefix_sum(num_verts + 1, 0);
+    utility::InclusivePrefixSum(vertex_mask_ptr, vertex_mask_ptr + num_verts,
+                                &prefix_sum[1]);
+
+    // update triangle indices
+    T *vert_idx_ptr = tris_cpu.GetDataPtr<T>();
+    for (int64_t i = 0; i < num_tris * 3; ++i) {
+        vert_idx_ptr[i] = prefix_sum[vert_idx_ptr[i]];
+    }
+}
+
+/// A helper to copy mesh attributes.
+/// \param dst destination mesh
+/// \param src source mesh
+/// \param vertex_mask vertex mask of the source mesh
+/// \param tri_mask triangle mask of the source mesh
+static void CopyAttributesByMasks(TriangleMesh &dst,
+                                  const TriangleMesh &src,
+                                  const core::Tensor &vertex_mask,
+                                  const core::Tensor &tri_mask) {
+    if (src.HasVertexPositions() && dst.HasVertexPositions()) {
+        for (auto item : src.GetVertexAttr()) {
+            if (!dst.HasVertexAttr(item.first)) {
+                dst.SetVertexAttr(item.first,
+                                  item.second.IndexGet({vertex_mask}));
+            }
+        }
+    }
+
+    if (src.HasTriangleIndices() && dst.HasTriangleIndices()) {
+        for (auto item : src.GetTriangleAttr()) {
+            if (!dst.HasTriangleAttr(item.first)) {
+                dst.SetTriangleAttr(item.first,
+                                    item.second.IndexGet({tri_mask}));
+            }
+        }
+    }
+}
+
 TriangleMesh TriangleMesh::SelectFacesByMask(const core::Tensor &mask) const {
+    if (!HasVertexPositions()) {
+        utility::LogWarning(
+                "[SelectFacesByMask] mesh has no vertex positions.");
+        return {};
+    }
+    if (!HasTriangleIndices()) {
+        utility::LogWarning(
+                "[SelectFacesByMask] mesh has no triangle indices.");
+        return {};
+    }
+
     core::AssertTensorShape(mask, {GetTriangleIndices().GetLength()});
     core::AssertTensorDtype(mask, core::Bool);
     GetTriangleAttr().AssertSizeSynchronized();
@@ -1004,62 +1063,154 @@ TriangleMesh TriangleMesh::SelectFacesByMask(const core::Tensor &mask) const {
     // select triangles
     core::Tensor tris = GetTriangleIndices().IndexGet({mask});
     core::Tensor tris_cpu = tris.To(core::Device()).Contiguous();
-    const int64_t num_tris = tris_cpu.GetLength();
 
     // create mask for vertices that are part of the selected faces
     const int64_t num_verts = GetVertexPositions().GetLength();
-    core::Tensor vertex_mask = core::Tensor::Zeros({num_verts}, core::Int32);
-    std::vector<int64_t> prefix_sum(num_verts + 1, 0);
-    {
-        int32_t *vertex_mask_ptr = vertex_mask.GetDataPtr<int32_t>();
-        if (tris_cpu.GetDtype() == core::Int32) {
-            int32_t *vert_idx_ptr = tris_cpu.GetDataPtr<int32_t>();
-            for (int64_t i = 0; i < tris_cpu.GetLength() * 3; ++i) {
-                vertex_mask_ptr[vert_idx_ptr[i]] = 1;
-            }
-        } else {
-            int64_t *vert_idx_ptr = tris_cpu.GetDataPtr<int64_t>();
-            for (int64_t i = 0; i < tris_cpu.GetLength() * 3; ++i) {
-                vertex_mask_ptr[vert_idx_ptr[i]] = 1;
-            }
-        }
-        utility::InclusivePrefixSum(
-                vertex_mask_ptr, vertex_mask_ptr + num_verts, &prefix_sum[1]);
-    }
-
-    // update triangle indices
-    if (tris_cpu.GetDtype() == core::Int32) {
-        int32_t *vert_idx_ptr = tris_cpu.GetDataPtr<int32_t>();
+    // empty tensor to further construct the vertex mask
+    core::Tensor vertex_mask;
+
+    DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE(tris_cpu.GetDtype(), tris, [&]() {
+        vertex_mask = core::Tensor::Zeros(
+                {num_verts}, core::Dtype::FromType<scalar_tris_t>());
+        const int64_t num_tris = tris_cpu.GetLength();
+        scalar_tris_t *vertex_mask_ptr =
+                vertex_mask.GetDataPtr<scalar_tris_t>();
+        scalar_tris_t *vert_idx_ptr = tris_cpu.GetDataPtr<scalar_tris_t>();
+        // mask for the vertices, which are used in the triangles
         for (int64_t i = 0; i < num_tris * 3; ++i) {
-            int64_t new_idx = prefix_sum[vert_idx_ptr[i]];
-            vert_idx_ptr[i] = int32_t(new_idx);
+            vertex_mask_ptr[vert_idx_ptr[i]] = 1;
         }
-    } else {
-        int64_t *vert_idx_ptr = tris_cpu.GetDataPtr<int64_t>();
-        for (int64_t i = 0; i < num_tris * 3; ++i) {
-            int64_t new_idx = prefix_sum[vert_idx_ptr[i]];
-            vert_idx_ptr[i] = new_idx;
-        }
-    }
+        UpdateTriangleIndicesByVertexMask<scalar_tris_t>(tris_cpu, vertex_mask);
+    });
 
     tris = tris_cpu.To(GetDevice());
     vertex_mask = vertex_mask.To(GetDevice(), core::Bool);
     core::Tensor verts = GetVertexPositions().IndexGet({vertex_mask});
     TriangleMesh result(verts, tris);
 
-    // copy attributes
-    for (auto item : GetVertexAttr()) {
-        if (!result.HasVertexAttr(item.first)) {
-            result.SetVertexAttr(item.first,
-                                 item.second.IndexGet({vertex_mask}));
-        }
+    CopyAttributesByMasks(result, *this, vertex_mask, mask);
+
+    return result;
+}
+
+/// brief Static negative checker for signed integer types
+template <typename T,
+          typename std::enable_if<std::is_integral<T>::value &&
+                                          !std::is_same<T, bool>::value &&
+                                          std::is_signed<T>::value,
+                                  T>::type * = nullptr>
+static bool IsNegative(T val) {
+    return val < 0;
+}
+
+/// brief Overloaded static negative checker for unsigned integer types.
+/// It unconditionally returns false, but we need it for template functions.
+template <typename T,
+          typename std::enable_if<std::is_integral<T>::value &&
+                                          !std::is_same<T, bool>::value &&
+                                          !std::is_signed<T>::value,
+                                  T>::type * = nullptr>
+static bool IsNegative(T val) {
+    return false;
+}
+
+TriangleMesh TriangleMesh::SelectByIndex(const core::Tensor &indices) const {
+    TriangleMesh result;
+    core::AssertTensorShape(indices, {indices.GetLength()});
+    if (!HasVertexPositions()) {
+        utility::LogWarning("[SelectByIndex] TriangleMesh has no vertices.");
+        return result;
     }
-    for (auto item : GetTriangleAttr()) {
-        if (!result.HasTriangleAttr(item.first)) {
-            result.SetTriangleAttr(item.first, item.second.IndexGet({mask}));
-        }
+    GetVertexAttr().AssertSizeSynchronized();
+
+    // we allow indices of an integral type only
+    core::Dtype::DtypeCode indices_dtype_code =
+            indices.GetDtype().GetDtypeCode();
+    if (indices_dtype_code != core::Dtype::DtypeCode::Int &&
+        indices_dtype_code != core::Dtype::DtypeCode::UInt) {
+        utility::LogError(
+                "[SelectByIndex] indices are not of integral type {}.",
+                indices.GetDtype().ToString());
+    }
+    core::Tensor indices_cpu = indices.To(core::Device()).Contiguous();
+    core::Tensor tris_cpu, tri_mask;
+    core::Dtype tri_dtype;
+    if (HasTriangleIndices()) {
+        GetTriangleAttr().AssertSizeSynchronized();
+        tris_cpu = GetTriangleIndices().To(core::Device()).Contiguous();
+        // bool mask for triangles.
+        tri_mask = core::Tensor::Zeros({tris_cpu.GetLength()}, core::Bool);
+        tri_dtype = tris_cpu.GetDtype();
+    } else {
+        utility::LogWarning("TriangleMesh has no triangle indices.");
+        tri_dtype = core::Int64;
+    }
+
+    // int mask to select vertices for the new mesh.  We need it as int as we
+    // will use its values to sum up and get the map of new indices
+    core::Tensor vertex_mask =
+            core::Tensor::Zeros({GetVertexPositions().GetLength()}, tri_dtype);
+
+    DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE(tri_dtype, tris, [&]() {
+        DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE(
+                indices_cpu.GetDtype(), indices, [&]() {
+                    const int64_t num_tris = tris_cpu.GetLength();
+                    const int64_t num_verts = vertex_mask.GetLength();
+
+                    // compute the vertices mask
+                    scalar_tris_t *vertex_mask_ptr =
+                            vertex_mask.GetDataPtr<scalar_tris_t>();
+                    const scalar_indices_t *indices_ptr =
+                            indices.GetDataPtr<scalar_indices_t>();
+                    for (int64_t i = 0; i < indices.GetLength(); ++i) {
+                        if (IsNegative(indices_ptr[i]) ||
+                            indices_ptr[i] >=
+                                    static_cast<scalar_indices_t>(num_verts)) {
+                            utility::LogWarning(
+                                    "[SelectByIndex] indices contains index {} "
+                                    "out of range. "
+                                    "It is ignored.",
+                                    indices_ptr[i]);
+                        }
+                        vertex_mask_ptr[indices_ptr[i]] = 1;
+                    }
+
+                    if (tri_mask.GetDtype() == core::Undefined) {
+                        // we don't need to compute triangles, if there are none
+                        return;
+                    }
+
+                    // Build the triangle mask
+                    scalar_tris_t *tris_cpu_ptr =
+                            tris_cpu.GetDataPtr<scalar_tris_t>();
+                    bool *tri_mask_ptr = tri_mask.GetDataPtr<bool>();
+                    for (int64_t i = 0; i < num_tris; ++i) {
+                        if (vertex_mask_ptr[tris_cpu_ptr[3 * i]] == 1 &&
+                            vertex_mask_ptr[tris_cpu_ptr[3 * i + 1]] == 1 &&
+                            vertex_mask_ptr[tris_cpu_ptr[3 * i + 2]] == 1) {
+                            tri_mask_ptr[i] = true;
+                        }
+                    }
+                    // select only needed triangles
+                    tris_cpu = tris_cpu.IndexGet({tri_mask});
+                    // update the triangle indices
+                    UpdateTriangleIndicesByVertexMask<scalar_tris_t>(
+                            tris_cpu, vertex_mask);
+                });
+    });
+
+    // send the vertex mask to original device and apply to vertices
+    vertex_mask = vertex_mask.To(GetDevice(), core::Bool);
+    core::Tensor new_vertices = GetVertexPositions().IndexGet({vertex_mask});
+    result.SetVertexPositions(new_vertices);
+
+    if (HasTriangleIndices()) {
+        // select triangles and send the selected ones to the original device
+        result.SetTriangleIndices(tris_cpu.To(GetDevice()));
     }
 
+    CopyAttributesByMasks(result, *this, vertex_mask, tri_mask);
+
     return result;
 }
 

cpp/open3d/t/geometry/TriangleMesh.h

@@ -927,9 +927,20 @@ class TriangleMesh : public Geometry, public DrawableGeometry {
     /// Returns a new mesh with the faces selected by a boolean mask.
     /// \param mask A boolean mask with the shape (N) with N as the number of
     /// faces in the mesh.
-    /// \return A new mesh with the selected faces.
+    /// \return A new mesh with the selected faces. If the original mesh is
+    /// empty, return an empty mesh.
     TriangleMesh SelectFacesByMask(const core::Tensor &mask) const;
 
+    /// Returns a new mesh with the vertices selected by a vector of indices.
+    /// If an item from the indices list exceeds the max vertex number of
+    /// the mesh or has a negative value, it is ignored.
+    /// \param indices An integer list of indices. Duplicates are
+    /// allowed, but ignored. Signed and unsigned integral types are allowed.
+    /// \return A new mesh with the selected vertices and faces built
+    /// from the selected vertices. If the original mesh is empty, return
+    /// an empty mesh.
+    TriangleMesh SelectByIndex(const core::Tensor &indices) const;
+
 protected:
     core::Device device_ = core::Device("CPU:0");
     TensorMap vertex_attr_;

cpp/pybind/t/geometry/trianglemesh.cpp

@@ -901,7 +901,7 @@ the partition id for each face.
         number of faces in the mesh.
     
 Returns:
-    A new mesh with the selected faces.
+    A new mesh with the selected faces. If the original mesh is empty, return an empty mesh.
 
 Example:
 
@@ -923,6 +923,30 @@ the partition id for each face.
 
         o3d.visualization.draw(parts)
 
+)");
+
+    triangle_mesh.def(
+            "select_by_index", &TriangleMesh::SelectByIndex, "indices"_a,
+            R"(Returns a new mesh with the vertices selected according to the indices list.
+If an item from the indices list exceeds the max vertex number of the mesh
+or has a negative value, it is ignored.
+
+Args:
+    indices (open3d.core.Tensor): An integer list of indices. Duplicates are
+    allowed, but ignored. Signed and unsigned integral types are accepted.
+
+Returns:
+    A new mesh with the selected vertices and faces built from these vertices.
+    If the original mesh is empty, return an empty mesh.
+
+Example:
+
+    This code selects the top face of a box, which has indices [2, 3, 6, 7]::
+
+        import open3d as o3d
+        import numpy as np
+        box = o3d.t.geometry.TriangleMesh.create_box()
+        top_face = box.select_by_index([2, 3, 6, 7])
 )");
 }