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 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, we throw an exception. Based on
the vertex mask we build a mapping index vector using inclusive prefix
sum algorithm. We then update the selected vertex indices in the triangles
CPU copy and build the triangles mask. I put that to a templated helper
static function to avoid the code duplication for Int32 and Int64 vertex
indices types.

With the computed masks we can select the vertices and triangles back on
the mesh's device and copy the attributes.

cpp/open3d/t/geometry/TriangleMesh.cpp

@@ -1061,6 +1061,108 @@ TriangleMesh TriangleMesh::SelectFacesByMask(const core::Tensor &mask) const {
     return result;
 }
 
+// A helper to compute the vertex and triangle masks based on indices.
+// Additionally updates tris_cpu to new indices.
+template <typename intT>
+static void SBIUpdateMasksAndTrisCPUHelper(const core::Tensor &indices,
+                                           core::Tensor &vertex_mask,
+                                           core::Tensor &tris_mask,
+                                           core::Tensor &tris_cpu) {
+    const int64_t num_tris = tris_cpu.GetLength();
+    const int64_t num_verts = vertex_mask.GetLength();
+
+    // compute the vertices mask
+    intT *vertex_mask_ptr = vertex_mask.GetDataPtr<intT>();
+    const intT *indices_ptr = indices.GetDataPtr<intT>();
+    for (int64_t i = 0; i < indices.GetLength(); ++i) {
+        if (indices_ptr[i] >= num_verts) {
+            utility::LogError(
+                    "[SelectByIndex] indices contains index {} out of range. ",
+                    indices_ptr[i]);
+            continue;
+        }
+        vertex_mask_ptr[indices_ptr[i]] = 1;
+    }
+
+    // compute new vertix indices
+    std::vector<intT> prefix_sum(num_verts + 1, 0);
+    utility::InclusivePrefixSum(vertex_mask_ptr, vertex_mask_ptr + num_verts,
+                                &prefix_sum[1]);
+
+    // update the triangles with new indices and build the triangle mask
+    intT *tris_cpu_ptr = tris_cpu.GetDataPtr<intT>();
+    bool *tris_mask_ptr = tris_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) {
+            tris_cpu_ptr[3 * i] = prefix_sum[tris_cpu_ptr[3 * i]];
+            tris_cpu_ptr[3 * i + 1] = prefix_sum[tris_cpu_ptr[3 * i + 1]];
+            tris_cpu_ptr[3 * i + 2] = prefix_sum[tris_cpu_ptr[3 * i + 2]];
+            tris_mask_ptr[i] = true;
+        }
+    }
+}
+
+TriangleMesh TriangleMesh::SelectByIndex(const core::Tensor &indices) const {
+    GetTriangleAttr().AssertSizeSynchronized();
+    GetVertexAttr().AssertSizeSynchronized();
+    if (GetTriangleIndices().GetDtype() == core::Int32) {
+        core::AssertTensorDtype(indices, core::Int32);
+    } else {
+        // we allow both Int32 and Int64 if the mesh indicies are Int64
+        core::AssertTensorDtypes(indices, {core::Int32, core::Int64});
+    }
+
+    // really copy triangles to CPU as we will modify the indices
+    core::Tensor tris_cpu =
+            GetTriangleIndices().To(core::Device(), true).Contiguous();
+    const core::Tensor indices_cpu = indices.To(core::Device()).Contiguous();
+    const int64_t num_tris = tris_cpu.GetLength();
+    const int64_t num_verts = GetVertexPositions().GetLength();
+
+    // 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({num_verts}, tris_cpu.GetDtype());
+    // bool mask for triangles.
+    core::Tensor tris_mask = core::Tensor::Zeros({num_tris}, core::Bool);
+
+    // compute vertex and triangular masks and triangles based on indices
+    if (tris_cpu.GetDtype() == core::Int32) {
+        SBIUpdateMasksAndTrisCPUHelper<int32_t>(indices_cpu, vertex_mask,
+                                                tris_mask, tris_cpu);
+    } else {
+        SBIUpdateMasksAndTrisCPUHelper<int64_t>(indices_cpu, vertex_mask,
+                                                tris_mask, tris_cpu);
+    }
+
+    // select triangles and send the selected ones to the original device
+    core::Tensor new_tris = tris_cpu.IndexGet({tris_mask}).To(GetDevice());
+    // send the vertex mask to original device and apply to vertices
+    vertex_mask = vertex_mask.To(GetDevice(), core::Bool);
+    core::Tensor new_verts = GetVertexPositions().IndexGet({vertex_mask});
+
+    // resulting mesh
+    TriangleMesh result(new_verts, new_tris);
+
+    // copy attributes
+    for (auto item : GetVertexAttr()) {
+        if (!result.HasVertexAttr(item.first)) {
+            result.SetVertexAttr(item.first,
+                                 item.second.IndexGet({vertex_mask}));
+        }
+    }
+    for (auto item : GetTriangleAttr()) {
+        if (!result.HasTriangleAttr(item.first)) {
+            result.SetTriangleAttr(item.first,
+                                   item.second.IndexGet({tris_mask}));
+        }
+    }
+
+    return result;
+}
+
 }  // namespace geometry
 }  // namespace t
 }  // namespace open3d

cpp/open3d/t/geometry/TriangleMesh.h

@@ -930,6 +930,17 @@ 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.
+    /// Throws an exception if an item from the indices list exceeds the max
+    /// vertex number of the mesh.
+    /// \param indices An integer list of indices. Duplicates are
+    /// allowed, but ignored. If vertex indices of the mesh are of type Int64,
+    /// both Int32 and Int64 are allowed as indices type, otherwise only Int32
+    /// is accepted.
+    /// \return A new mesh with the selected vertices and faces built
+    /// from the selected vertices.
+    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,32 @@ 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.
+Throws an exception if an item from the indices list exceeds the max vertex
+number of the mesh,
+
+Args:
+    indices (open3d.core.Tensor): An integer list of indices. Duplicates are
+    allowed, but ignored. If vertex indices of the mesh are of type Int64,
+    both Int32 and Int64 are allowed as indices type, otherwise only Int32
+    is accepted.
+
+Returns:
+    A new mesh with the selected vertices and faces built from these vertices.
+
+Example:
+
+    This code selets the top face of a box, which has indices [2, 3, 6, 7].
+    parts::
+
+        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,92 @@ TEST_P(TriangleMeshPermuteDevices, CreateMobius) {
             triangle_indices_custom));
 }
 
+TEST_P(TriangleMeshPermuteDevices, SelectByIndex_Box) {
+    // 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);
+
+    core::Tensor indices = core::Tensor::Init<int64_t>({2, 3, 6, 7});
+    t::geometry::TriangleMesh selected = box.SelectByIndex(indices);
+
+    // Set the expected values.
+    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}});
+
+    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 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

python/test/t/geometry/test_trianglemesh.py

@@ -417,3 +417,41 @@ def test_pickle(device):
                                 mesh.vertex.positions.cpu().numpy())
         np.testing.assert_equal(mesh_load.triangle.indices.cpu().numpy(),
                                 mesh.triangle.indices.cpu().numpy())
+
+
+@pytest.mark.parametrize("device", list_devices())
+def test_select_by_index(device):
+    sphere_custom = o3d.t.geometry.TriangleMesh.create_sphere(
+        1, 3, o3c.float64, o3c.int32, device)
+
+    expected_verts = o3c.Tensor(
+        [[0.0, 0.0, 1.0], [0.866025, 0, 0.5], [0.433013, 0.75, 0.5],
+         [-0.866025, 0.0, 0.5], [-0.433013, -0.75, 0.5], [0.433013, -0.75, 0.5]
+        ], o3c.float64, device)
+
+    expected_tris = o3c.Tensor([[0, 1, 2], [0, 3, 4], [0, 4, 5], [0, 5, 1]],
+                               o3c.int32, device)
+
+    indices = o3c.Tensor([0, 2, 3, 5, 6, 7], o3c.int64, device)
+    # check indices type mismatch
+    with pytest.raises(RuntimeError) as e:
+        selected = sphere_custom.select_by_index(indices)
+
+    # check the expected mesh
+    indices = o3c.Tensor([0, 2, 3, 5, 6, 7], o3c.int32, device)
+    selected = sphere_custom.select_by_index(indices)
+    assert selected.vertex.positions.allclose(expected_verts)
+    assert selected.triangle.indices.allclose(expected_tris)
+
+    # check that the original mesh is unmodified
+    untouched_sphere = o3d.t.geometry.TriangleMesh.create_sphere(
+        1, 3, o3c.float64, o3c.int32, device)
+    assert sphere_custom.vertex.positions.allclose(
+        untouched_sphere.vertex.positions)
+    assert sphere_custom.triangle.indices.allclose(
+        untouched_sphere.triangle.indices)
+
+    # check that the exception is thrown if one of the indices exceeds
+    # the max vertex index of the mesh
+    with pytest.raises(RuntimeError) as e:
+        selected = sphere_custom.select_by_index([2, 3, 6, 99])