From 1b1c50798d9584f631f68bea140b244bd6275f4d Mon Sep 17 00:00:00 2001 From: Natalia Saiapova Date: Fri, 6 Oct 2023 20:05:58 +0000 Subject: [PATCH] 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 | 102 ++++++++++++++++++++ cpp/open3d/t/geometry/TriangleMesh.h | 11 +++ cpp/pybind/t/geometry/trianglemesh.cpp | 26 +++++ cpp/tests/t/geometry/TriangleMesh.cpp | 87 +++++++++++++++++ python/test/t/geometry/test_trianglemesh.py | 38 ++++++++ 5 files changed, 264 insertions(+) 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])