From 54aa4b3510d7b7d3c38e34d904f390e88b22e697 Mon Sep 17 00:00:00 2001 From: Natalia Saiapova Date: Fri, 13 Oct 2023 13:51:51 +0000 Subject: [PATCH] fixup! Implement open3d::t::geometry::TriangleMesh::SelectByIndex --- cpp/open3d/t/geometry/TriangleMesh.cpp | 34 +++++---- cpp/open3d/t/geometry/TriangleMesh.h | 4 +- cpp/pybind/t/geometry/trianglemesh.cpp | 5 +- cpp/tests/t/geometry/TriangleMesh.cpp | 23 ++++-- python/test/t/geometry/test_trianglemesh.py | 78 +++++++++++++++++++-- 5 files changed, 118 insertions(+), 26 deletions(-) diff --git a/cpp/open3d/t/geometry/TriangleMesh.cpp b/cpp/open3d/t/geometry/TriangleMesh.cpp index fb2d4b29f8c0..16e48032a834 100644 --- a/cpp/open3d/t/geometry/TriangleMesh.cpp +++ b/cpp/open3d/t/geometry/TriangleMesh.cpp @@ -1061,9 +1061,9 @@ 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 +/// \brief 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, @@ -1072,25 +1072,29 @@ static void SBIUpdateMasksAndTrisCPUHelper(const core::Tensor &indices, 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(); + tri_intT *vertex_mask_ptr = vertex_mask.GetDataPtr(); + const indices_intT *indices_ptr = indices.GetDataPtr(); for (int64_t i = 0; i < indices.GetLength(); ++i) { + if (indices_ptr[i] < 0) { + utility::LogError( + "[SelectByIndex] indices contains a negative index {}. ", + indices_ptr[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); + 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(); + tri_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 && @@ -1105,10 +1109,11 @@ static void SBIUpdateMasksAndTrisCPUHelper(const core::Tensor &indices, } TriangleMesh TriangleMesh::SelectByIndex(const core::Tensor &indices) const { + core::AssertTensorShape(indices, {indices.GetLength()}); GetTriangleAttr().AssertSizeSynchronized(); GetVertexAttr().AssertSizeSynchronized(); if (GetTriangleIndices().GetDtype() == core::Int32) { - core::AssertTensorDtype(indices, core::Int32); + core::AssertTensorDtype(indices, GetTriangleIndices().GetDtype()); } else { // we allow both Int32 and Int64 if the mesh indicies are Int64 core::AssertTensorDtypes(indices, {core::Int32, core::Int64}); @@ -1130,11 +1135,16 @@ TriangleMesh TriangleMesh::SelectByIndex(const core::Tensor &indices) const { // compute vertex and triangular masks and triangles based on indices if (tris_cpu.GetDtype() == core::Int32) { - SBIUpdateMasksAndTrisCPUHelper(indices_cpu, vertex_mask, + SBIUpdateMasksAndTrisCPUHelper(indices_cpu, vertex_mask, tris_mask, tris_cpu); } else { - SBIUpdateMasksAndTrisCPUHelper(indices_cpu, vertex_mask, - tris_mask, tris_cpu); + if (indices_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 diff --git a/cpp/open3d/t/geometry/TriangleMesh.h b/cpp/open3d/t/geometry/TriangleMesh.h index d68e22c985bd..7ad8daa566c5 100644 --- a/cpp/open3d/t/geometry/TriangleMesh.h +++ b/cpp/open3d/t/geometry/TriangleMesh.h @@ -932,8 +932,8 @@ class TriangleMesh : public Geometry, public DrawableGeometry { /// 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 + /// vertex number of the mesh or a negative value was supplied. + /// \param indices An integer list of non-negative 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. diff --git a/cpp/pybind/t/geometry/trianglemesh.cpp b/cpp/pybind/t/geometry/trianglemesh.cpp index a53889070652..0711a0aaff7b 100644 --- a/cpp/pybind/t/geometry/trianglemesh.cpp +++ b/cpp/pybind/t/geometry/trianglemesh.cpp @@ -929,7 +929,7 @@ the partition id for each face. "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, +number of the mesh or a negative value was supplied. Args: indices (open3d.core.Tensor): An integer list of indices. Duplicates are @@ -942,8 +942,7 @@ number of the mesh, Example: - This code selets the top face of a box, which has indices [2, 3, 6, 7]. - parts:: + This code selects the top face of a box, which has indices [2, 3, 6, 7]:: import open3d as o3d import numpy as np diff --git a/cpp/tests/t/geometry/TriangleMesh.cpp b/cpp/tests/t/geometry/TriangleMesh.cpp index ef5c6ad9789b..f9b6600a7c66 100644 --- a/cpp/tests/t/geometry/TriangleMesh.cpp +++ b/cpp/tests/t/geometry/TriangleMesh.cpp @@ -982,10 +982,11 @@ TEST_P(TriangleMeshPermuteDevices, SelectByIndex_Box) { box.ComputeTriangleNormals(); box.SetTriangleAttr("labels", triangle_labels); - core::Tensor indices = core::Tensor::Init({2, 3, 6, 7}); - t::geometry::TriangleMesh selected = box.SelectByIndex(indices); + // empty index list + core::Tensor indices_empty = core::Tensor::Init({}); + EXPECT_TRUE(box.SelectByIndex(indices_empty).IsEmpty()); - // Set the expected values. + // set the expected valuee 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}, @@ -1000,7 +1001,6 @@ TEST_P(TriangleMeshPermuteDevices, SelectByIndex_Box) { {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 = @@ -1010,6 +1010,9 @@ TEST_P(TriangleMeshPermuteDevices, SelectByIndex_Box) { core::Tensor expected_tri_labels = core::Tensor::Init( {{800.0, 800.0, 800.0}, {900.0, 900.0, 900.0}}); + core::Tensor indices = core::Tensor::Init({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( @@ -1019,6 +1022,18 @@ TEST_P(TriangleMeshPermuteDevices, SelectByIndex_Box) { EXPECT_TRUE( selected.GetTriangleAttr("labels").AllClose(expected_tri_labels)); + core::Tensor indices_duplicate = core::Tensor::Init({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)); + + // Check that initial mesh is unchanged. t::geometry::TriangleMesh box_untouched = t::geometry::TriangleMesh::CreateBox(); diff --git a/python/test/t/geometry/test_trianglemesh.py b/python/test/t/geometry/test_trianglemesh.py index c63a33216dc7..763a2f59916f 100644 --- a/python/test/t/geometry/test_trianglemesh.py +++ b/python/test/t/geometry/test_trianglemesh.py @@ -420,7 +420,7 @@ def test_pickle(device): @pytest.mark.parametrize("device", list_devices()) -def test_select_by_index(device): +def test_select_by_index_32(device): sphere_custom = o3d.t.geometry.TriangleMesh.create_sphere( 1, 3, o3c.float64, o3c.int32, device) @@ -432,10 +432,20 @@ def test_select_by_index(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 shape mismatch + indices_2d = o3c.Tensor([[0, 2], [3, 5], [6, 7]], o3c.int32, device) + with pytest.raises(RuntimeError): + selected = sphere_custom.select_by_index(indices_2d) + + # check indices int size mismatch + indices_64 = o3c.Tensor([0, 2, 3, 5, 6, 7], o3c.int64, device) + with pytest.raises(RuntimeError): + selected = sphere_custom.select_by_index(indices_64) + # check indices type mismatch - with pytest.raises(RuntimeError) as e: - selected = sphere_custom.select_by_index(indices) + indices_float = o3c.Tensor([2.0, 4.0], o3c.float32, device) + with pytest.raises(RuntimeError): + selected = sphere_custom.select_by_index(indices_float) # check the expected mesh indices = o3c.Tensor([0, 2, 3, 5, 6, 7], o3c.int32, device) @@ -453,5 +463,63 @@ def test_select_by_index(device): # 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: + with pytest.raises(RuntimeError): + selected = sphere_custom.select_by_index([2, 3, 6, 99]) + + # check that the exception is thrown if one of the indices have a negative + # value + with pytest.raises(RuntimeError): + selected = sphere_custom.select_by_index([2, 3, 6, -7]) + +@pytest.mark.parametrize("device", list_devices()) +def test_select_by_index_64(device): + sphere_custom = o3d.t.geometry.TriangleMesh.create_sphere( + 1, 3, o3c.float64, o3c.int64, device) + + # check indices shape mismatch + with pytest.raises(RuntimeError): + indices_2d = o3c.Tensor([[0, 2], [3, 5], [6, 7]], o3c.int64, device) + selected = sphere_custom.select_by_index(indices_2d) + + # check indices type mismatch + with pytest.raises(RuntimeError): + indices_float = o3c.Tensor([2.0, 4.0], o3c.float64, device) + selected = sphere_custom.select_by_index(indices_float) + + 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.int64, device) + + # check the expected mesh with int64 input + indices_64 = o3c.Tensor([0, 2, 3, 5, 6, 7], o3c.int64, device) + selected = sphere_custom.select_by_index(indices_64) + assert selected.vertex.positions.allclose(expected_verts) + assert selected.triangle.indices.allclose(expected_tris) + + # check the expected mesh with int32 input and unsorted indices + indices_32 = o3c.Tensor([7, 6, 3, 5, 0, 2], o3c.int32, device) + selected = sphere_custom.select_by_index(indices_32) + 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.int64, 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): selected = sphere_custom.select_by_index([2, 3, 6, 99]) + + # check that the exception is thrown if one of the indices have a negative + # value + with pytest.raises(RuntimeError): + selected = sphere_custom.select_by_index([2, 3, 6, -7])