diff --git a/cpp/open3d/t/io/TriangleMeshIO.cpp b/cpp/open3d/t/io/TriangleMeshIO.cpp index 5aaa90eb705..e4f82f1e42f 100644 --- a/cpp/open3d/t/io/TriangleMeshIO.cpp +++ b/cpp/open3d/t/io/TriangleMeshIO.cpp @@ -7,8 +7,10 @@ #include "open3d/t/io/TriangleMeshIO.h" +#include #include +#include "open3d/t/io/NumpyIO.h" #include "open3d/utility/FileSystem.h" #include "open3d/utility/Logging.h" @@ -22,6 +24,7 @@ static const std::unordered_map< geometry::TriangleMesh &, const open3d::io::ReadTriangleMeshOptions &)>> file_extension_to_trianglemesh_read_function{ + {"npz", ReadTriangleMeshFromNPZ}, {"stl", ReadTriangleMeshUsingASSIMP}, {"obj", ReadTriangleMeshUsingASSIMP}, {"off", ReadTriangleMeshUsingASSIMP}, @@ -40,7 +43,9 @@ static const std::unordered_map< const bool, const bool, const bool)>> - file_extension_to_trianglemesh_write_function{}; + file_extension_to_trianglemesh_write_function{ + {"npz", WriteTriangleMeshToNPZ}, + }; std::shared_ptr CreateMeshFromFile( const std::string &filename, bool print_progress) { @@ -137,6 +142,160 @@ bool WriteTriangleMesh(const std::string &filename, return success; } +bool ReadTriangleMeshFromNPZ( + const std::string &filename, + geometry::TriangleMesh &mesh, + const open3d::io::ReadTriangleMeshOptions ¶ms) { + auto attribute_map = ReadNpz(filename); + + // At a minimum there should be 'vertices' and 'triangles' + if (!(attribute_map.count("vertices") > 0) || + !(attribute_map.count("triangles") > 0)) { + utility::LogWarning( + "Read geometry::TriangleMesh failed: Could not find 'vertices' " + "or 'triangles' attributes in {}", + filename); + return false; + } + + // Fill mesh with attributes + for (auto &attr : attribute_map) { + if (attr.first == "vertices") { + mesh.SetVertexPositions(attr.second); + } else if (attr.first == "triangles") { + mesh.SetTriangleIndices(attr.second); + } else if (attr.first == "vertex_normals") { + mesh.SetVertexNormals(attr.second); + } else if (attr.first == "triangle_normals") { + mesh.SetTriangleNormals(attr.second); + } else if (attr.first == "vertex_colors") { + mesh.SetVertexColors(attr.second); + } else if (attr.first == "triangle_colors") { + mesh.SetTriangleColors(attr.second); + } else if (attr.first == "triangle_texture_uvs") { + mesh.SetTriangleAttr("texture_uvs", attr.second); + } else if (attr.first.find("tex_") != std::string::npos) { + // Get texture map + auto key = attr.first.substr(4); + if (!mesh.GetMaterial().IsValid()) { + mesh.GetMaterial().SetDefaultProperties(); + } + mesh.GetMaterial().SetTextureMap(key, geometry::Image(attr.second)); + // Note: due to quirk of Open3D shader implementation if we have a + // metallic texture we need to set the metallic scalar propert to + // 1.0 + if (key == "metallic") { + mesh.GetMaterial().SetScalarProperty("metallic", 1.0); + } + } else if (attr.first.find("vertex_") != std::string::npos) { + // Generic vertex attribute + auto key = attr.first.substr(7); + mesh.SetVertexAttr(key, attr.second); + } else if (attr.first.find("triangle_") != std::string::npos) { + // Generic triangle attribute + auto key = attr.first.substr(9); + mesh.SetTriangleAttr(key, attr.second); + } else if (attr.first == "material_name") { + if (!mesh.GetMaterial().IsValid()) { + mesh.GetMaterial().SetDefaultProperties(); + } + const uint8_t *str_ptr = attr.second.GetDataPtr(); + std::string mat_name(attr.second.GetShape().GetLength(), 'a'); + std::memcpy((void *)mat_name.data(), str_ptr, + attr.second.GetShape().GetLength()); + mesh.GetMaterial().SetMaterialName(mat_name); + } + } + + return true; +} + +bool WriteTriangleMeshToNPZ(const std::string &filename, + const geometry::TriangleMesh &mesh, + const bool write_ascii, + const bool compressed, + const bool write_vertex_normals, + const bool write_vertex_colors, + const bool write_triangle_uvs, + const bool print_progress) { + // Sanity checks... + if (write_ascii) { + utility::LogWarning( + "TriangleMesh can't be saved in ASCII fromat as .npz"); + return false; + } + if (compressed) { + utility::LogWarning( + "TriangleMesh can't be saved in compressed format as .npz"); + return false; + } + + // Map attribute names to names already used by convention in other software + std::set known_attributes( + {"positions", "normals", "texture_uvs", "indices", "colors"}); + + // Build map of known attributes + std::unordered_map mesh_attributes; + if (mesh.HasVertexPositions()) { + mesh_attributes["vertices"] = mesh.GetVertexPositions(); + } + if (mesh.HasVertexNormals()) { + mesh_attributes["vertex_normals"] = mesh.GetVertexNormals(); + } + if (mesh.HasVertexColors()) { + mesh_attributes["vertex_colors"] = mesh.GetVertexColors(); + } + if (mesh.HasTriangleIndices()) { + mesh_attributes["triangles"] = mesh.GetTriangleIndices(); + } + if (mesh.HasTriangleNormals()) { + mesh_attributes["triangle_normals"] = mesh.GetTriangleNormals(); + } + if (mesh.HasTriangleColors()) { + mesh_attributes["triangle_colors"] = mesh.GetTriangleColors(); + } + if (mesh.HasTriangleAttr("texture_uvs")) { + mesh_attributes["triangle_texture_uvs"] = + mesh.GetTriangleAttr("texture_uvs"); + } + + // Add "generic" attributes + for (auto attr : mesh.GetVertexAttr()) { + if (known_attributes.count(attr.first) > 0) { + continue; + } + std::string key_name("vertex_"); + key_name += attr.first; + mesh_attributes[key_name] = attr.second; + } + for (auto attr : mesh.GetTriangleAttr()) { + if (known_attributes.count(attr.first) > 0) { + continue; + } + std::string key_name("triangle_"); + key_name += attr.first; + mesh_attributes[key_name] = attr.second; + } + + // Output texture maps + if (mesh.GetMaterial().IsValid()) { + auto &mat = mesh.GetMaterial(); + // Get material name in Tensor form + std::vector mat_name_vec( + {mat.GetMaterialName().begin(), mat.GetMaterialName().end()}); + core::Tensor mat_name_tensor(std::move(mat_name_vec)); + mesh_attributes["material_name"] = mat_name_tensor; + for (auto &tex : mat.GetTextureMaps()) { + std::string key = std::string("tex_") + tex.first; + mesh_attributes[key] = tex.second.AsTensor(); + } + } + + WriteNpz(filename, mesh_attributes); + + return true; +} + } // namespace io } // namespace t } // namespace open3d diff --git a/cpp/open3d/t/io/TriangleMeshIO.h b/cpp/open3d/t/io/TriangleMeshIO.h index 3498644fbb7..ccdaf4cee47 100644 --- a/cpp/open3d/t/io/TriangleMeshIO.h +++ b/cpp/open3d/t/io/TriangleMeshIO.h @@ -54,6 +54,19 @@ bool ReadTriangleMeshUsingASSIMP( geometry::TriangleMesh &mesh, const open3d::io::ReadTriangleMeshOptions ¶ms); +bool ReadTriangleMeshFromNPZ(const std::string &filename, + geometry::TriangleMesh &mesh, + const open3d::io::ReadTriangleMeshOptions ¶ms); + +bool WriteTriangleMeshToNPZ(const std::string &filename, + const geometry::TriangleMesh &mesh, + const bool write_ascii, + const bool compressed, + const bool write_vertex_normals, + const bool write_vertex_colors, + const bool write_triangle_uvs, + const bool print_progress); + } // namespace io } // namespace t } // namespace open3d diff --git a/cpp/tests/t/io/TriangleMeshIO.cpp b/cpp/tests/t/io/TriangleMeshIO.cpp index dbf2f828192..acfd8f1234f 100644 --- a/cpp/tests/t/io/TriangleMeshIO.cpp +++ b/cpp/tests/t/io/TriangleMeshIO.cpp @@ -71,6 +71,20 @@ TEST(TriangleMeshIO, ReadWriteTriangleMeshOBJ) { EXPECT_TRUE(mesh.GetTriangleIndices().AllClose(triangles)); } +TEST(TriangleMeshIO, ReadWriteTriangleMeshNPZ) { + auto cube_mesh = t::geometry::TriangleMesh::CreateBox(); + + const std::string filename = + utility::filesystem::GetTempDirectoryPath() + "/cube.npz"; + EXPECT_TRUE(t::io::WriteTriangleMesh(filename, cube_mesh)); + t::geometry::TriangleMesh mesh; + EXPECT_TRUE(t::io::ReadTriangleMesh(filename, mesh)); + EXPECT_TRUE( + mesh.GetVertexPositions().AllClose(cube_mesh.GetVertexPositions())); + EXPECT_TRUE( + mesh.GetTriangleIndices().AllClose(cube_mesh.GetTriangleIndices())); +} + // TODO: Add tests for triangle_uvs, materials, triangle_material_ids and // textures once these are supported. TEST(TriangleMeshIO, TriangleMeshLegecyCompatibility) {