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)