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.

cpp/open3d/t/geometry/TriangleMesh.cpp

@@ -993,6 +993,54 @@ 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 {
     core::AssertTensorShape(mask, {GetTriangleIndices().GetLength()});
     core::AssertTensorDtype(mask, core::Bool);
@@ -1045,18 +1093,126 @@ TriangleMesh TriangleMesh::SelectFacesByMask(const core::Tensor &mask) const {
     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.");
     }
+    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

@@ -930,6 +930,16 @@ class TriangleMesh : public Geometry, public DrawableGeometry {
     /// \return A new mesh with the selected faces.
     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

@@ -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])
 )");
 }
 

cpp/tests/t/geometry/TriangleMesh.cpp

@@ -941,5 +941,154 @@ TEST_P(TriangleMeshPermuteDevices, CreateMobius) {
             triangle_indices_custom));
 }
 
+TEST_P(TriangleMeshPermuteDevices, SelectByIndex) {
+    // check that an exception is thrown if the mesh is empty
+    t::geometry::TriangleMesh mesh_empty;
+    core::Tensor indices_empty = core::Tensor::Init<int64_t>({});
+
+    // check completely empty mesh
+    EXPECT_TRUE(mesh_empty.SelectByIndex(indices_empty).IsEmpty());
+    EXPECT_TRUE(mesh_empty.SelectByIndex(core::Tensor::Init<int64_t>({0}))
+                        .IsEmpty());
+
+    // check mesh w/o triangles
+    core::Tensor vertices_no_tris_orig =
+            core::Tensor::Ones({2, 3}, core::Float32, mesh_empty.GetDevice());
+    core::Tensor expected_vertices_no_tris_orig =
+            core::Tensor::Ones({1, 3}, core::Float32, mesh_empty.GetDevice());
+    mesh_empty.SetVertexPositions(vertices_no_tris_orig);
+    t::geometry::TriangleMesh selected_no_tris_orig =
+            mesh_empty.SelectByIndex(core::Tensor::Init<int64_t>({0}));
+    EXPECT_TRUE(selected_no_tris_orig.GetVertexPositions().AllClose(
+            expected_vertices_no_tris_orig));
+
+    // create box with normals, colors and labels defined.
+    t::geometry::TriangleMesh box = t::geometry::TriangleMesh::CreateBox();
+    core::Tensor vertex_colors = core::Tensor::Init<float>({{0.0, 0.0, 0.0},
+                                                            {1.0, 1.0, 1.0},
+                                                            {2.0, 2.0, 2.0},
+                                                            {3.0, 3.0, 3.0},
+                                                            {4.0, 4.0, 4.0},
+                                                            {5.0, 5.0, 5.0},
+                                                            {6.0, 6.0, 6.0},
+                                                            {7.0, 7.0, 7.0}});
+    ;
+    core::Tensor vertex_labels = core::Tensor::Init<float>({{0.0, 0.0, 0.0},
+                                                            {1.0, 1.0, 1.0},
+                                                            {2.0, 2.0, 2.0},
+                                                            {3.0, 3.0, 3.0},
+                                                            {4.0, 4.0, 4.0},
+                                                            {5.0, 5.0, 5.0},
+                                                            {6.0, 6.0, 6.0},
+                                                            {7.0, 7.0, 7.0}}) *
+                                 10;
+    ;
+    core::Tensor triangle_labels =
+            core::Tensor::Init<float>({{0.0, 0.0, 0.0},
+                                       {1.0, 1.0, 1.0},
+                                       {2.0, 2.0, 2.0},
+                                       {3.0, 3.0, 3.0},
+                                       {4.0, 4.0, 4.0},
+                                       {5.0, 5.0, 5.0},
+                                       {6.0, 6.0, 6.0},
+                                       {7.0, 7.0, 7.0},
+                                       {8.0, 8.0, 8.0},
+                                       {9.0, 9.0, 9.0},
+                                       {10.0, 10.0, 10.0},
+                                       {11.0, 11.0, 11.0}}) *
+            100;
+    box.SetVertexColors(vertex_colors);
+    box.SetVertexAttr("labels", vertex_labels);
+    box.ComputeTriangleNormals();
+    box.SetTriangleAttr("labels", triangle_labels);
+
+    // empty index list
+    EXPECT_TRUE(box.SelectByIndex(indices_empty).IsEmpty());
+
+    // set the expected value
+    core::Tensor expected_verts = core::Tensor::Init<float>({{0.0, 0.0, 1.0},
+                                                             {1.0, 0.0, 1.0},
+                                                             {0.0, 1.0, 1.0},
+                                                             {1.0, 1.0, 1.0}});
+    core::Tensor expected_vert_colors =
+            core::Tensor::Init<float>({{2.0, 2.0, 2.0},
+                                       {3.0, 3.0, 3.0},
+                                       {6.0, 6.0, 6.0},
+                                       {7.0, 7.0, 7.0}});
+    core::Tensor expected_vert_labels =
+            core::Tensor::Init<float>({{20.0, 20.0, 20.0},
+                                       {30.0, 30.0, 30.0},
+                                       {60.0, 60.0, 60.0},
+                                       {70.0, 70.0, 70.0}});
+    core::Tensor expected_tris =
+            core::Tensor::Init<int64_t>({{0, 1, 3}, {0, 3, 2}});
+    core::Tensor tris_mask =
+            core::Tensor::Init<bool>({0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0});
+    core::Tensor expected_tri_normals =
+            box.GetTriangleNormals().IndexGet({tris_mask});
+    core::Tensor expected_tri_labels = core::Tensor::Init<float>(
+            {{800.0, 800.0, 800.0}, {900.0, 900.0, 900.0}});
+
+    // check basic case
+    core::Tensor indices = core::Tensor::Init<int64_t>({2, 3, 6, 7});
+    t::geometry::TriangleMesh selected = box.SelectByIndex(indices);
+
+    EXPECT_TRUE(selected.GetVertexPositions().AllClose(expected_verts));
+    EXPECT_TRUE(selected.GetVertexColors().AllClose(expected_vert_colors));
+    EXPECT_TRUE(
+            selected.GetVertexAttr("labels").AllClose(expected_vert_labels));
+    EXPECT_TRUE(selected.GetTriangleIndices().AllClose(expected_tris));
+    EXPECT_TRUE(selected.GetTriangleNormals().AllClose(expected_tri_normals));
+    EXPECT_TRUE(
+            selected.GetTriangleAttr("labels").AllClose(expected_tri_labels));
+
+    // check duplicated indices case
+    core::Tensor indices_duplicate =
+            core::Tensor::Init<int16_t>({2, 2, 3, 3, 6, 7, 7});
+    t::geometry::TriangleMesh selected_duplicate =
+            box.SelectByIndex(indices_duplicate);
+    EXPECT_TRUE(
+            selected_duplicate.GetVertexPositions().AllClose(expected_verts));
+    EXPECT_TRUE(selected_duplicate.GetVertexColors().AllClose(
+            expected_vert_colors));
+    EXPECT_TRUE(selected_duplicate.GetVertexAttr("labels").AllClose(
+            expected_vert_labels));
+    EXPECT_TRUE(
+            selected_duplicate.GetTriangleIndices().AllClose(expected_tris));
+    EXPECT_TRUE(selected_duplicate.GetTriangleNormals().AllClose(
+            expected_tri_normals));
+    EXPECT_TRUE(selected_duplicate.GetTriangleAttr("labels").AllClose(
+            expected_tri_labels));
+
+    // select with empty triangles as result
+    // set the expected value
+    core::Tensor expected_verts_no_tris = core::Tensor::Init<float>(
+            {{0.0, 0.0, 0.0}, {1.0, 0.0, 1.0}, {0.0, 1.0, 0.0}});
+    core::Tensor expected_vert_colors_no_tris = core::Tensor::Init<float>(
+            {{0.0, 0.0, 0.0}, {3.0, 3.0, 3.0}, {4.0, 4.0, 4.0}});
+    core::Tensor expected_vert_labels_no_tris = core::Tensor::Init<float>(
+            {{0.0, 0.0, 0.0}, {30.0, 30.0, 30.0}, {40.0, 40.0, 40.0}});
+
+    core::Tensor indices_no_tris = core::Tensor::Init<int64_t>({0, 3, 4});
+    t::geometry::TriangleMesh selected_no_tris =
+            box.SelectByIndex(indices_no_tris);
+
+    EXPECT_TRUE(selected_no_tris.GetVertexPositions().AllClose(
+            expected_verts_no_tris));
+    EXPECT_TRUE(selected_no_tris.GetVertexColors().AllClose(
+            expected_vert_colors_no_tris));
+    EXPECT_TRUE(selected_no_tris.GetVertexAttr("labels").AllClose(
+            expected_vert_labels_no_tris));
+    EXPECT_FALSE(selected_no_tris.HasTriangleIndices());
+
+    // check that initial mesh is unchanged
+    t::geometry::TriangleMesh box_untouched =
+            t::geometry::TriangleMesh::CreateBox();
+    EXPECT_TRUE(box.GetVertexPositions().AllClose(
+            box_untouched.GetVertexPositions()));
+    EXPECT_TRUE(box.GetTriangleIndices().AllClose(
+            box_untouched.GetTriangleIndices()));
+}
+
 }  // namespace tests
 }  // namespace open3d