diff --git a/cpp/open3d/t/geometry/TriangleMesh.cpp b/cpp/open3d/t/geometry/TriangleMesh.cpp index 4134ee6f306..fb2d4b29f8c 100644 --- a/cpp/open3d/t/geometry/TriangleMesh.cpp +++ b/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 +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(); + const intT *indices_ptr = indices.GetDataPtr(); + 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 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(); + bool *tris_mask_ptr = tris_mask.GetDataPtr(); + 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(indices_cpu, vertex_mask, + tris_mask, tris_cpu); + } else { + SBIUpdateMasksAndTrisCPUHelper(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 diff --git a/cpp/open3d/t/geometry/TriangleMesh.h b/cpp/open3d/t/geometry/TriangleMesh.h index 7828ac16b02..d68e22c985b 100644 --- a/cpp/open3d/t/geometry/TriangleMesh.h +++ b/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_; diff --git a/cpp/pybind/t/geometry/trianglemesh.cpp b/cpp/pybind/t/geometry/trianglemesh.cpp index 5979238e1b2..a5388907065 100644 --- a/cpp/pybind/t/geometry/trianglemesh.cpp +++ b/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]) )"); } diff --git a/cpp/tests/t/geometry/TriangleMesh.cpp b/cpp/tests/t/geometry/TriangleMesh.cpp index 4e51b34e31d..ef5c6ad9789 100644 --- a/cpp/tests/t/geometry/TriangleMesh.cpp +++ b/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({{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({{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({{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({2, 3, 6, 7}); + t::geometry::TriangleMesh selected = box.SelectByIndex(indices); + + // Set the expected values. + core::Tensor expected_verts = core::Tensor::Init({{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({{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({{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({{0, 1, 3}, {0, 3, 2}}); + core::Tensor tris_mask = + core::Tensor::Init({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( + {{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 diff --git a/python/test/t/geometry/test_trianglemesh.py b/python/test/t/geometry/test_trianglemesh.py index 843184dd3e6..c63a33216dc 100644 --- a/python/test/t/geometry/test_trianglemesh.py +++ b/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])