From 1585e9686ed8e6aa9c423c9fbafb74628a7f0f63 Mon Sep 17 00:00:00 2001 From: nsaiapova Date: Tue, 14 Nov 2023 00:01:39 +0100 Subject: [PATCH] Implement open3d::t::geometry::TriangleMesh::SelectByIndex (#6415) * Define a helper DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE This is a helper to call a templated function with an integer argument, based on Dtype. As a second argument, it takes a suffix, used to build a unique type name. This way, we can use it to call a function with more than one integer argument. Example: DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE(core::Dtype::Int32, int32, [&]() { DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE(core::Dtype::UInt64, uint64, [&]() { scalar_int32_t a; scalar_uint64_t b; // ... }); * 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. * Refactor t::geometry::TriangleMesh::SelectFacesByMask * Add error handling on empty mesh * Use DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE instead of a conditional branch * Use UpdateTriangleIndicesByVertexMask helper to update triangle indices * Use CopyAttributesByMask helper to copy the mesh attributes * Add tests --- cpp/open3d/core/Dispatch.h | 31 +++ cpp/open3d/t/geometry/TriangleMesh.cpp | 235 ++++++++++++++---- cpp/open3d/t/geometry/TriangleMesh.h | 13 +- cpp/pybind/t/geometry/trianglemesh.cpp | 26 +- cpp/tests/t/geometry/TriangleMesh.cpp | 256 ++++++++++++++++++++ python/test/t/geometry/test_trianglemesh.py | 222 +++++++++++++++++ 6 files changed, 739 insertions(+), 44 deletions(-) diff --git a/cpp/open3d/core/Dispatch.h b/cpp/open3d/core/Dispatch.h index 466ccf86beb..fe658c8cf15 100644 --- a/cpp/open3d/core/Dispatch.h +++ b/cpp/open3d/core/Dispatch.h @@ -113,3 +113,34 @@ open3d::utility::LogError("Unsupported data type."); \ } \ }() + +#define DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE(DTYPE, PREFIX, ...) \ + [&] { \ + if (DTYPE == open3d::core::Int8) { \ + using scalar_##PREFIX##_t = int8_t; \ + return __VA_ARGS__(); \ + } else if (DTYPE == open3d::core::Int16) { \ + using scalar_##PREFIX##_t = int16_t; \ + return __VA_ARGS__(); \ + } else if (DTYPE == open3d::core::Int32) { \ + using scalar_##PREFIX##_t = int32_t; \ + return __VA_ARGS__(); \ + } else if (DTYPE == open3d::core::Int64) { \ + using scalar_##PREFIX##_t = int64_t; \ + return __VA_ARGS__(); \ + } else if (DTYPE == open3d::core::UInt8) { \ + using scalar_##PREFIX##_t = uint8_t; \ + return __VA_ARGS__(); \ + } else if (DTYPE == open3d::core::UInt16) { \ + using scalar_##PREFIX##_t = uint16_t; \ + return __VA_ARGS__(); \ + } else if (DTYPE == open3d::core::UInt32) { \ + using scalar_##PREFIX##_t = uint32_t; \ + return __VA_ARGS__(); \ + } else if (DTYPE == open3d::core::UInt64) { \ + using scalar_##PREFIX##_t = uint64_t; \ + return __VA_ARGS__(); \ + } else { \ + open3d::utility::LogError("Unsupported data type."); \ + } \ + }() diff --git a/cpp/open3d/t/geometry/TriangleMesh.cpp b/cpp/open3d/t/geometry/TriangleMesh.cpp index a87abf04eb5..e4e878af93a 100644 --- a/cpp/open3d/t/geometry/TriangleMesh.cpp +++ b/cpp/open3d/t/geometry/TriangleMesh.cpp @@ -995,7 +995,66 @@ 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 +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(); + std::vector 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(); + 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 { + if (!HasVertexPositions()) { + utility::LogWarning( + "[SelectFacesByMask] mesh has no vertex positions."); + return {}; + } + if (!HasTriangleIndices()) { + utility::LogWarning( + "[SelectFacesByMask] mesh has no triangle indices."); + return {}; + } + core::AssertTensorShape(mask, {GetTriangleIndices().GetLength()}); core::AssertTensorDtype(mask, core::Bool); GetTriangleAttr().AssertSizeSynchronized(); @@ -1004,62 +1063,154 @@ TriangleMesh TriangleMesh::SelectFacesByMask(const core::Tensor &mask) const { // select triangles core::Tensor tris = GetTriangleIndices().IndexGet({mask}); core::Tensor tris_cpu = tris.To(core::Device()).Contiguous(); - const int64_t num_tris = tris_cpu.GetLength(); // create mask for vertices that are part of the selected faces const int64_t num_verts = GetVertexPositions().GetLength(); - core::Tensor vertex_mask = core::Tensor::Zeros({num_verts}, core::Int32); - std::vector prefix_sum(num_verts + 1, 0); - { - int32_t *vertex_mask_ptr = vertex_mask.GetDataPtr(); - if (tris_cpu.GetDtype() == core::Int32) { - int32_t *vert_idx_ptr = tris_cpu.GetDataPtr(); - for (int64_t i = 0; i < tris_cpu.GetLength() * 3; ++i) { - vertex_mask_ptr[vert_idx_ptr[i]] = 1; - } - } else { - int64_t *vert_idx_ptr = tris_cpu.GetDataPtr(); - for (int64_t i = 0; i < tris_cpu.GetLength() * 3; ++i) { - vertex_mask_ptr[vert_idx_ptr[i]] = 1; - } - } - utility::InclusivePrefixSum( - vertex_mask_ptr, vertex_mask_ptr + num_verts, &prefix_sum[1]); - } - - // update triangle indices - if (tris_cpu.GetDtype() == core::Int32) { - int32_t *vert_idx_ptr = tris_cpu.GetDataPtr(); + // empty tensor to further construct the vertex mask + core::Tensor vertex_mask; + + DISPATCH_INT_DTYPE_PREFIX_TO_TEMPLATE(tris_cpu.GetDtype(), tris, [&]() { + vertex_mask = core::Tensor::Zeros( + {num_verts}, core::Dtype::FromType()); + const int64_t num_tris = tris_cpu.GetLength(); + scalar_tris_t *vertex_mask_ptr = + vertex_mask.GetDataPtr(); + scalar_tris_t *vert_idx_ptr = tris_cpu.GetDataPtr(); + // mask for the vertices, which are used in the triangles for (int64_t i = 0; i < num_tris * 3; ++i) { - int64_t new_idx = prefix_sum[vert_idx_ptr[i]]; - vert_idx_ptr[i] = int32_t(new_idx); + vertex_mask_ptr[vert_idx_ptr[i]] = 1; } - } else { - int64_t *vert_idx_ptr = tris_cpu.GetDataPtr(); - for (int64_t i = 0; i < num_tris * 3; ++i) { - int64_t new_idx = prefix_sum[vert_idx_ptr[i]]; - vert_idx_ptr[i] = new_idx; - } - } + UpdateTriangleIndicesByVertexMask(tris_cpu, vertex_mask); + }); tris = tris_cpu.To(GetDevice()); vertex_mask = vertex_mask.To(GetDevice(), core::Bool); 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 ::value && + !std::is_same::value && + std::is_signed::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 ::value && + !std::is_same::value && + !std::is_signed::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 {}.", + indices.GetDtype().ToString()); + } + 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(); + const scalar_indices_t *indices_ptr = + indices.GetDataPtr(); + for (int64_t i = 0; i < indices.GetLength(); ++i) { + if (IsNegative(indices_ptr[i]) || + indices_ptr[i] >= + static_cast(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(); + bool *tri_mask_ptr = tri_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) { + tri_mask_ptr[i] = true; + } + } + // select only needed triangles + tris_cpu = tris_cpu.IndexGet({tri_mask}); + // update the triangle indices + UpdateTriangleIndicesByVertexMask( + 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; } diff --git a/cpp/open3d/t/geometry/TriangleMesh.h b/cpp/open3d/t/geometry/TriangleMesh.h index 7828ac16b02..8f7373afd9d 100644 --- a/cpp/open3d/t/geometry/TriangleMesh.h +++ b/cpp/open3d/t/geometry/TriangleMesh.h @@ -927,9 +927,20 @@ class TriangleMesh : public Geometry, public DrawableGeometry { /// Returns a new mesh with the faces selected by a boolean mask. /// \param mask A boolean mask with the shape (N) with N as the number of /// faces in the mesh. - /// \return A new mesh with the selected faces. + /// \return A new mesh with the selected faces. If the original mesh is + /// empty, return an empty mesh. 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_; diff --git a/cpp/pybind/t/geometry/trianglemesh.cpp b/cpp/pybind/t/geometry/trianglemesh.cpp index 5979238e1b2..06cacf404a5 100644 --- a/cpp/pybind/t/geometry/trianglemesh.cpp +++ b/cpp/pybind/t/geometry/trianglemesh.cpp @@ -901,7 +901,7 @@ the partition id for each face. number of faces in the mesh. Returns: - A new mesh with the selected faces. + A new mesh with the selected faces. If the original mesh is empty, return an empty mesh. Example: @@ -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]) )"); } diff --git a/cpp/tests/t/geometry/TriangleMesh.cpp b/cpp/tests/t/geometry/TriangleMesh.cpp index 520848a5032..8cb6ba0ddac 100644 --- a/cpp/tests/t/geometry/TriangleMesh.cpp +++ b/cpp/tests/t/geometry/TriangleMesh.cpp @@ -946,5 +946,261 @@ TEST_P(TriangleMeshPermuteDevices, CreateMobius) { triangle_indices_custom)); } +TEST_P(TriangleMeshPermuteDevices, SelectFacesByMask) { + // check that an exception is thrown if the mesh is empty + t::geometry::TriangleMesh mesh_empty; + core::Tensor mask_empty = + core::Tensor::Zeros({12}, core::Bool, mesh_empty.GetDevice()); + core::Tensor mask_full = + core::Tensor::Ones({12}, core::Bool, mesh_empty.GetDevice()); + + // check completely empty mesh + EXPECT_TRUE(mesh_empty.SelectFacesByMask(mask_empty).IsEmpty()); + EXPECT_TRUE(mesh_empty.SelectFacesByMask(mask_full).IsEmpty()); + + // check mesh w/o triangles + core::Tensor cpu_vertices = + core::Tensor::Ones({2, 3}, core::Float32, mesh_empty.GetDevice()); + mesh_empty.SetVertexPositions(cpu_vertices); + EXPECT_TRUE(mesh_empty.SelectFacesByMask(mask_empty).IsEmpty()); + EXPECT_TRUE(mesh_empty.SelectFacesByMask(mask_full).IsEmpty()); + + // 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); + + // empty index list + EXPECT_TRUE(box.SelectFacesByMask(mask_empty).IsEmpty()); + + // set the expected value + 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}}); + + // check basic case + t::geometry::TriangleMesh selected = box.SelectFacesByMask(tris_mask); + + 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())); +} + +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({}); + + // check completely empty mesh + EXPECT_TRUE(mesh_empty.SelectByIndex(indices_empty).IsEmpty()); + EXPECT_TRUE(mesh_empty.SelectByIndex(core::Tensor::Init({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({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({{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); + + // empty index list + EXPECT_TRUE(box.SelectByIndex(indices_empty).IsEmpty()); + + // set the expected value + 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}}); + + // check basic case + 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( + 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({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( + {{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( + {{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( + {{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({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 diff --git a/python/test/t/geometry/test_trianglemesh.py b/python/test/t/geometry/test_trianglemesh.py index 843184dd3e6..2a108adff56 100644 --- a/python/test/t/geometry/test_trianglemesh.py +++ b/python/test/t/geometry/test_trianglemesh.py @@ -417,3 +417,225 @@ 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_faces_by_mask_32(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) + + # check indices shape mismatch + mask_2d = o3c.Tensor([[False, False], [False, False], [False, False]], + o3c.bool, device) + with pytest.raises(RuntimeError): + selected = sphere_custom.select_faces_by_mask(mask_2d) + + # check indices type mismatch + mask_float = o3c.Tensor([ + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 + ], o3c.float32, device) + with pytest.raises(RuntimeError): + selected = sphere_custom.select_faces_by_mask(mask_float) + + # check the basic case + mask = o3c.Tensor([ + True, False, False, False, False, False, True, False, True, False, True, + False, False, False, False, False, False, False, False, False, False, + False, False, False + ], o3c.bool, device) + selected = sphere_custom.select_faces_by_mask(mask) + 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) + + +@pytest.mark.parametrize("device", list_devices()) +def test_select_faces_by_mask_64(device): + sphere_custom = o3d.t.geometry.TriangleMesh.create_sphere( + 1, 3, o3c.float64, o3c.int64, device) + + # check indices shape mismatch + mask_2d = o3c.Tensor([[False, False], [False, False], [False, False]], + o3c.bool, device) + with pytest.raises(RuntimeError): + selected = sphere_custom.select_faces_by_mask(mask_2d) + + # check indices type mismatch + mask_float = o3c.Tensor([ + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 + ], o3c.float32, device) + with pytest.raises(RuntimeError): + selected = sphere_custom.select_faces_by_mask(mask_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 basic case + mask = o3c.Tensor([ + True, False, False, False, False, False, True, False, True, False, True, + False, False, False, False, False, False, False, False, False, False, + False, False, False + ], o3c.bool, device) + + selected = sphere_custom.select_faces_by_mask(mask) + 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) + + +@pytest.mark.parametrize("device", list_devices()) +def test_select_by_index_32(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) + + # 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 type mismatch + 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 with int8 input + indices_8 = o3c.Tensor([0, 2, 3, 5, 6, 7], o3c.int8, device) + selected = sphere_custom.select_by_index(indices_8) + assert selected.vertex.positions.allclose(expected_verts) + assert selected.triangle.indices.allclose(expected_tris) + + # check the expected mesh with int16 input + indices_16 = o3c.Tensor([2, 0, 5, 3, 7, 6], o3c.int16, device) + selected = sphere_custom.select_by_index(indices_16) + assert selected.vertex.positions.allclose(expected_verts) + assert selected.triangle.indices.allclose(expected_tris) + + # check the expected mesh with uint32 input + indices_u32 = o3c.Tensor([7, 6, 5, 3, 2, 0], o3c.uint32, device) + selected = sphere_custom.select_by_index(indices_u32) + assert selected.vertex.positions.allclose(expected_verts) + assert selected.triangle.indices.allclose(expected_tris) + + # check the expected mesh with uint64 input and unsorted indices + indices_u64 = o3c.Tensor([7, 6, 3, 5, 0, 2], o3c.uint64, device) + selected = sphere_custom.select_by_index(indices_u64) + assert selected.vertex.positions.allclose(expected_verts) + assert selected.triangle.indices.allclose(expected_tris) + + # check that an index exceeding the max vertex index of the mesh is ignored + selected = sphere_custom.select_by_index([0, 2, 3, 5, 6, 99, 7]) + assert selected.vertex.positions.allclose(expected_verts) + assert selected.triangle.indices.allclose(expected_tris) + + # check that a negative index is ignored + selected = sphere_custom.select_by_index([0, 2, 3, 5, -10, 6, 7]) + 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) + + +@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 int8 input + indices_u8 = o3c.Tensor([0, 2, 3, 5, 6, 7], o3c.uint8, device) + selected = sphere_custom.select_by_index(indices_u8) + assert selected.vertex.positions.allclose(expected_verts) + assert selected.triangle.indices.allclose(expected_tris) + + # check the expected mesh with int16 input + indices_u16 = o3c.Tensor([2, 0, 5, 3, 7, 6], o3c.uint16, device) + selected = sphere_custom.select_by_index(indices_u16) + assert selected.vertex.positions.allclose(expected_verts) + assert selected.triangle.indices.allclose(expected_tris) + + # check the expected mesh with int32 input + indices_32 = o3c.Tensor([7, 6, 5, 3, 2, 0], 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 the expected mesh with int64 input and unsorted indices + indices_64 = o3c.Tensor([7, 6, 3, 5, 0, 2], 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 that an index exceeding the max vertex index of the mesh is ignored + selected = sphere_custom.select_by_index([0, 2, 3, 5, 6, 99, 7]) + assert selected.vertex.positions.allclose(expected_verts) + assert selected.triangle.indices.allclose(expected_tris) + + # check that a negative index is ignored + selected = sphere_custom.select_by_index([0, 2, 3, 5, -10, 6, 7]) + 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)