Add IO for NPZ files for Tensor-based TriangleMeshes (isl-org#6019)

* Add IO to/from NPZ files

* Add texture maps and generic attriubtes to IO

* Style fixes

* Read/write material name and fix read with metallic texture

* Add unit test for TriangleMesh NPZ IO

* Apply style fixes

* Rename uvmap to triangle_texture_uvs

* Style fixes

cpp/open3d/t/io/TriangleMeshIO.cpp

@@ -7,8 +7,10 @@
 
 #include "open3d/t/io/TriangleMeshIO.h"
 
+#include <set>
 #include <unordered_map>
 
+#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<geometry::TriangleMesh> 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 &params) {
+    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<uint8_t>();
+            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<std::string> known_attributes(
+            {"positions", "normals", "texture_uvs", "indices", "colors"});
+
+    // Build map of known attributes
+    std::unordered_map<std::string, core::Tensor> 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<uint8_t> 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

cpp/open3d/t/io/TriangleMeshIO.h

@@ -54,6 +54,19 @@ bool ReadTriangleMeshUsingASSIMP(
         geometry::TriangleMesh &mesh,
         const open3d::io::ReadTriangleMeshOptions &params);
 
+bool ReadTriangleMeshFromNPZ(const std::string &filename,
+                             geometry::TriangleMesh &mesh,
+                             const open3d::io::ReadTriangleMeshOptions &params);
+
+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

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) {