From 7a25f3dac4e1ec1b8a81592037ec5c6ffd34eca5 Mon Sep 17 00:00:00 2001 From: Yueci Deng Date: Sun, 25 Sep 2022 20:41:40 +0800 Subject: [PATCH] Enable pickling for tensor and tensor based geometry (#5509) --- cpp/open3d/t/geometry/PointCloud.cpp | 2 +- cpp/pybind/core/device.cpp | 33 +++++--- cpp/pybind/core/tensor.cpp | 30 +++++++ cpp/pybind/t/geometry/image.cpp | 37 +++++++++ cpp/pybind/t/geometry/lineset.cpp | 39 +++++++++ cpp/pybind/t/geometry/pointcloud.cpp | 35 ++++++++ cpp/pybind/t/geometry/tensormap.cpp | 23 ++++++ cpp/pybind/t/geometry/trianglemesh.cpp | 91 +++++++++++++++------ cpp/tests/t/geometry/PointCloud.cpp | 4 +- python/test/core/test_core.py | 20 +++++ python/test/t/geometry/test_image.py | 13 +++ python/test/t/geometry/test_lineset.py | 24 ++++++ python/test/t/geometry/test_pointcloud.py | 17 ++++ python/test/t/geometry/test_tensormap.py | 15 ++++ python/test/t/geometry/test_trianglemesh.py | 18 ++++ 15 files changed, 361 insertions(+), 40 deletions(-) diff --git a/cpp/open3d/t/geometry/PointCloud.cpp b/cpp/open3d/t/geometry/PointCloud.cpp index af3938f611f..1d13a5e7132 100644 --- a/cpp/open3d/t/geometry/PointCloud.cpp +++ b/cpp/open3d/t/geometry/PointCloud.cpp @@ -96,7 +96,7 @@ std::string PointCloud::ToString() const { fmt::format(" ({})", GetPointPositions().GetDtype().ToString()); } auto str = - fmt::format("PointCloud on {} [{} points{}] Attributes:", + fmt::format("PointCloud on {} [{} points{}].\nAttributes:", GetDevice().ToString(), num_points, points_dtype_str); if ((point_attr_.size() - point_attr_.count(point_attr_.GetPrimaryKey())) == diff --git a/cpp/pybind/core/device.cpp b/cpp/pybind/core/device.cpp index aee3222dda7..3431bd6ce09 100644 --- a/cpp/pybind/core/device.cpp +++ b/cpp/pybind/core/device.cpp @@ -37,16 +37,29 @@ void pybind_core_device(py::module &m) { py::class_ device( m, "Device", "Device context specifying device type and device id."); - device.def(py::init<>()) - .def(py::init()) - .def(py::init()) - .def(py::init()) - .def("__eq__", &Device::operator==) - .def("__ene__", &Device::operator!=) - .def("__repr__", &Device::ToString) - .def("__str__", &Device::ToString) - .def("get_type", &Device::GetType) - .def("get_id", &Device::GetID); + device.def(py::init<>()); + device.def(py::init()); + device.def(py::init()); + device.def(py::init()); + device.def("__eq__", &Device::operator==); + device.def("__ene__", &Device::operator!=); + device.def("__repr__", &Device::ToString); + device.def("__str__", &Device::ToString); + device.def("get_type", &Device::GetType); + device.def("get_id", &Device::GetID); + device.def(py::pickle( + [](const Device &d) { + return py::make_tuple(d.GetType(), d.GetID()); + }, + [](py::tuple t) { + if (t.size() != 2) { + utility::LogError( + "Cannot unpickle Device! Expecting a tuple of size " + "2."); + } + return Device(t[0].cast(), + t[1].cast()); + })); py::enum_(device, "DeviceType") .value("CPU", Device::DeviceType::CPU) diff --git a/cpp/pybind/core/tensor.cpp b/cpp/pybind/core/tensor.cpp index fba3eefc1bf..ac5a2bdd3af 100644 --- a/cpp/pybind/core/tensor.cpp +++ b/cpp/pybind/core/tensor.cpp @@ -356,6 +356,36 @@ void pybind_core_tensor(py::module& m) { BindTensorFullCreation(m, tensor); docstring::ClassMethodDocInject(m, "Tensor", "full", argument_docs); + // Pickling support. + // The tensor will be on the same device after deserialization. + // Non contiguous tensors will be converted to contiguous tensors after + // deserialization. + tensor.def(py::pickle( + [](const Tensor& t) { + // __getstate__ + return py::make_tuple(t.GetDevice(), + TensorToPyArray(t.To(Device("CPU:0")))); + }, + [](py::tuple t) { + // __setstate__ + if (t.size() != 2) { + utility::LogError( + "Cannot unpickle Tensor! Expecting a tuple of size " + "2."); + } + const Device& device = t[0].cast(); + if (!device.IsAvailable()) { + utility::LogWarning( + "Device {} is not available, tensor will be " + "created on CPU.", + device.ToString()); + return PyArrayToTensor(t[1].cast(), true); + } else { + return PyArrayToTensor(t[1].cast(), true) + .To(device); + } + })); + tensor.def_static( "eye", [](int64_t n, utility::optional dtype, diff --git a/cpp/pybind/t/geometry/image.cpp b/cpp/pybind/t/geometry/image.cpp index 64ca17772ea..fe873578812 100644 --- a/cpp/pybind/t/geometry/image.cpp +++ b/cpp/pybind/t/geometry/image.cpp @@ -101,6 +101,23 @@ void pybind_image(py::module &m) { "tensor"_a); docstring::ClassMethodDocInject(m, "Image", "__init__", map_shared_argument_docstrings); + + // Pickle support. + image.def(py::pickle( + [](const Image &image) { + // __getstate__ + return py::make_tuple(image.AsTensor()); + }, + [](py::tuple t) { + // __setstate__ + if (t.size() != 1) { + utility::LogError( + "Cannot unpickle Image! Expecting a tuple of size " + "1."); + } + return Image(t[0].cast()); + })); + // Buffer protocol. image.def_buffer([](Image &I) -> py::buffer_info { if (!I.IsCPU()) { @@ -282,6 +299,26 @@ void pybind_image(py::module &m) { .def(py::init(), "Parameterized constructor", "color"_a, "depth"_a, "aligned"_a = true) + + // Pickling support. + .def(py::pickle( + [](const RGBDImage &rgbd) { + // __getstate__ + return py::make_tuple(rgbd.color_, rgbd.depth_, + rgbd.aligned_); + }, + [](py::tuple t) { + // __setstate__ + if (t.size() != 3) { + utility::LogError( + "Cannot unpickle RGBDImage! Expecting a " + "tuple of size 3."); + } + + return RGBDImage(t[0].cast(), t[1].cast(), + t[2].cast()); + })) + // Depth and color images. .def_readwrite("color", &RGBDImage::color_, "The color image.") .def_readwrite("depth", &RGBDImage::depth_, "The depth image.") diff --git a/cpp/pybind/t/geometry/lineset.cpp b/cpp/pybind/t/geometry/lineset.cpp index 86e529f31d1..818175c61ef 100644 --- a/cpp/pybind/t/geometry/lineset.cpp +++ b/cpp/pybind/t/geometry/lineset.cpp @@ -29,6 +29,7 @@ #include #include +#include "open3d/core/CUDAUtils.h" #include "open3d/t/geometry/TriangleMesh.h" #include "pybind/docstring.h" #include "pybind/t/geometry/geometry.h" @@ -109,6 +110,44 @@ and ``device`` as the tensor. The device for ``point_positions`` must be consist {"line_indices", "A tensor with element shape (2,) and Int dtype."}}); + // Pickling support. + line_set.def(py::pickle( + [](const LineSet& line_set) { + // __getstate__ + return py::make_tuple(line_set.GetDevice(), + line_set.GetPointAttr(), + line_set.GetLineAttr()); + }, + [](py::tuple t) { + // __setstate__ + if (t.size() != 3) { + utility::LogError( + "Cannot unpickle LineSet! Expecting a tuple of " + "size 3."); + } + + const core::Device device = t[0].cast(); + LineSet line_set(device); + if (!device.IsAvailable()) { + utility::LogWarning( + "Device ({}) is not available. LineSet will be " + "created on CPU.", + device.ToString()); + line_set.To(core::Device("CPU:0")); + } + + const TensorMap point_attr = t[1].cast(); + const TensorMap line_attr = t[2].cast(); + for (auto& kv : point_attr) { + line_set.SetPointAttr(kv.first, kv.second); + } + for (auto& kv : line_attr) { + line_set.SetLineAttr(kv.first, kv.second); + } + + return line_set; + })); + // Line set's attributes: point_positions, line_indices, line_colors, etc. // def_property_readonly is sufficient, since the returned TensorMap can // be editable in Python. We don't want the TensorMap to be replaced diff --git a/cpp/pybind/t/geometry/pointcloud.cpp b/cpp/pybind/t/geometry/pointcloud.cpp index dd2e8c86e0a..efee278cefb 100644 --- a/cpp/pybind/t/geometry/pointcloud.cpp +++ b/cpp/pybind/t/geometry/pointcloud.cpp @@ -127,6 +127,41 @@ The attributes of the point cloud have different levels:: "map_keys_to_tensors"_a) .def("__repr__", &PointCloud::ToString); + // Pickle support. + pointcloud.def(py::pickle( + [](const PointCloud& pcd) { + // __getstate__ + // Convert point attributes to tensor map to CPU. + auto map_keys_to_tensors = pcd.GetPointAttr(); + + return py::make_tuple(pcd.GetDevice(), pcd.GetPointAttr()); + }, + [](py::tuple t) { + // __setstate__ + if (t.size() != 2) { + utility::LogError( + "Cannot unpickle PointCloud! Expecting a tuple of " + "size 2."); + } + + const core::Device device = t[0].cast(); + PointCloud pcd(device); + if (!device.IsAvailable()) { + utility::LogWarning( + "Device ({}) is not available. PointCloud will be " + "created on CPU.", + device.ToString()); + pcd.To(core::Device("CPU:0")); + } + + const TensorMap map_keys_to_tensors = t[1].cast(); + for (auto& kv : map_keys_to_tensors) { + pcd.SetPointAttr(kv.first, kv.second); + } + + return pcd; + })); + // def_property_readonly is sufficient, since the returned TensorMap can // be editable in Python. We don't want the TensorMap to be replaced // by another TensorMap in Python. diff --git a/cpp/pybind/t/geometry/tensormap.cpp b/cpp/pybind/t/geometry/tensormap.cpp index e9f0bc8c3be..30a1338fd1b 100644 --- a/cpp/pybind/t/geometry/tensormap.cpp +++ b/cpp/pybind/t/geometry/tensormap.cpp @@ -164,6 +164,29 @@ void pybind_tensormap(py::module &m) { tm.def("is_size_synchronized", &TensorMap::IsSizeSynchronized); tm.def("assert_size_synchronized", &TensorMap::AssertSizeSynchronized); + // Pickle support. + tm.def(py::pickle( + [](const TensorMap &m) { + // __getstate__ + std::unordered_map map; + for (const auto &kv : m) { + map[kv.first] = kv.second; + } + + return py::make_tuple(m.GetPrimaryKey(), map); + }, + [](py::tuple t) { + // __setstate__ + if (t.size() != 2) { + utility::LogError( + "Cannot unpickle TensorMap! Expecting a tuple of " + "size 2."); + } + return TensorMap(t[0].cast(), + t[1].cast>()); + })); + tm.def("__setattr__", [](TensorMap &m, const std::string &key, const core::Tensor &val) { if (!TensorMap::GetReservedKeys().count(key)) { diff --git a/cpp/pybind/t/geometry/trianglemesh.cpp b/cpp/pybind/t/geometry/trianglemesh.cpp index fdd32751aea..0aa32b5915a 100644 --- a/cpp/pybind/t/geometry/trianglemesh.cpp +++ b/cpp/pybind/t/geometry/trianglemesh.cpp @@ -107,6 +107,43 @@ The attributes of the triangle mesh have different levels:: "vertex_positions"_a, "triangle_indices"_a) .def("__repr__", &TriangleMesh::ToString); + // Pickle support. + triangle_mesh.def(py::pickle( + [](const TriangleMesh& mesh) { + // __getstate__ + return py::make_tuple(mesh.GetDevice(), mesh.GetVertexAttr(), + mesh.GetTriangleAttr()); + }, + [](py::tuple t) { + // __setstate__ + if (t.size() != 3) { + utility::LogError( + "Cannot unpickle TriangleMesh! Expecting a tuple " + "of size 3."); + } + + const core::Device device = t[0].cast(); + TriangleMesh mesh(device); + if (!device.IsAvailable()) { + utility::LogWarning( + "Device ({}) is not available. TriangleMesh will " + "be created on CPU.", + device.ToString()); + mesh.To(core::Device("CPU:0")); + } + + const TensorMap vertex_attr = t[1].cast(); + const TensorMap triangle_attr = t[2].cast(); + for (auto& kv : vertex_attr) { + mesh.SetVertexAttr(kv.first, kv.second); + } + for (auto& kv : triangle_attr) { + mesh.SetTriangleAttr(kv.first, kv.second); + } + + return mesh; + })); + // Triangle mesh's attributes: vertices, vertex_colors, vertex_normals, etc. // def_property_readonly is sufficient, since the returned TensorMap can // be editable in Python. We don't want the TensorMap to be replaced @@ -236,7 +273,7 @@ This example shows how to create a hemisphere from a sphere:: "point"_a, "normal"_a, "contour_values"_a = std::list{0.0}, R"(Returns a line set with the contour slices defined by the plane and values. -This method generates slices as LineSet from the mesh at specific contour +This method generates slices as LineSet from the mesh at specific contour values with respect to a plane. Args: @@ -452,11 +489,11 @@ This example shows how to create a hemisphere from a sphere:: "int_dtype"_a = core::Int64, "device"_a = core::Device("CPU:0"), R"(Create a triangle mesh from a text string. - + Args: - text (str): The text for generating the mesh. ASCII characters 32-126 are - supported (includes alphanumeric characters and punctuation). In - addition the line feed '\n' is supported to start a new line. + text (str): The text for generating the mesh. ASCII characters 32-126 are + supported (includes alphanumeric characters and punctuation). In + addition the line feed '\n' is supported to start a new line. depth (float): The depth of the generated mesh. If depth is 0 then a flat mesh will be generated. float_dtype (o3d.core.Dtype): Float type for the vertices. Either Float32 or Float64. int_dtype (o3d.core.Dtype): Int type for the triangle indices. Either Int32 or Int64. @@ -479,7 +516,7 @@ This example shows how to create a hemisphere from a sphere:: &TriangleMesh::SimplifyQuadricDecimation, "target_reduction"_a, "preserve_volume"_a = true, R"(Function to simplify mesh using Quadric Error Metric Decimation by Garland and Heckbert. - + This function always uses the CPU device. Args: @@ -512,7 +549,7 @@ Both meshes should be manifold. This function always uses the CPU device. Args: - mesh (open3d.t.geometry.TriangleMesh): This is the second operand for the + mesh (open3d.t.geometry.TriangleMesh): This is the second operand for the boolean operation. tolerance (float): Threshold which determines when point distances are @@ -543,7 +580,7 @@ Both meshes should be manifold. This function always uses the CPU device. Args: - mesh (open3d.t.geometry.TriangleMesh): This is the second operand for the + mesh (open3d.t.geometry.TriangleMesh): This is the second operand for the boolean operation. tolerance (float): Threshold which determines when point distances are @@ -574,7 +611,7 @@ Both meshes should be manifold. This function always uses the CPU device. Args: - mesh (open3d.t.geometry.TriangleMesh): This is the second operand for the + mesh (open3d.t.geometry.TriangleMesh): This is the second operand for the boolean operation. tolerance (float): Threshold which determines when point distances are @@ -627,10 +664,10 @@ This function always uses the CPU device. "compute_uvatlas", &TriangleMesh::ComputeUVAtlas, "size"_a = 512, "gutter"_a = 1.f, "max_stretch"_a = 1.f / 6, R"(Creates an UV atlas and adds it as triangle attr 'texture_uvs' to the mesh. - + Input meshes must be manifold for this method to work. The algorithm is based on: -Zhou et al, "Iso-charts: Stretch-driven Mesh Parameterization using Spectral +Zhou et al, "Iso-charts: Stretch-driven Mesh Parameterization using Spectral Analysis", Eurographics Symposium on Geometry Processing (2004) Sander et al. "Signal-Specialized Parametrization" Europgraphics 2002 This function always uses the CPU device. @@ -649,7 +686,7 @@ This function always uses the CPU device. bunny = o3d.data.BunnyMesh() mesh = o3d.t.geometry.TriangleMesh.from_legacy(o3d.io.read_triangle_mesh(bunny.path)) mesh.compute_uvatlas() - + # Add a wood texture and visualize texture_data = o3d.data.WoodTexture() mesh.material.material_name = 'defaultLit' @@ -669,10 +706,10 @@ Only float type attributes can be baked to textures. This function always uses the CPU device. Args: - size (int): The width and height of the texture in pixels. Only square + size (int): The width and height of the texture in pixels. Only square textures are supported. - vertex_attr (set): The vertex attributes for which textures should be + vertex_attr (set): The vertex attributes for which textures should be generated. margin (float): The margin in pixels. The recommended value is 2. The margin @@ -697,7 +734,7 @@ This function always uses the CPU device. box = o3d.t.geometry.TriangleMesh.from_legacy(box) box.vertex['albedo'] = box.vertex.positions - # Initialize material and bake the 'albedo' vertex attribute to a + # Initialize material and bake the 'albedo' vertex attribute to a # texture. The texture will be automatically added to the material of # the object. box.material.set_default_properties() @@ -705,7 +742,7 @@ This function always uses the CPU device. # Shows the textured cube. o3d.visualization.draw([box]) - + # Plot the tensor with the texture. plt.imshow(texture_tensors['albedo'].numpy()) @@ -722,10 +759,10 @@ This function assumes a triangle attribute with name 'texture_uvs'. This function always uses the CPU device. Args: - size (int): The width and height of the texture in pixels. Only square + size (int): The width and height of the texture in pixels. Only square textures are supported. - triangle_attr (set): The vertex attributes for which textures should be + triangle_attr (set): The vertex attributes for which textures should be generated. margin (float): The margin in pixels. The recommended value is 2. The margin @@ -742,18 +779,18 @@ This function always uses the CPU device. A dictionary of tensors that store the baked textures. Example: - We generate a texture visualizing the index of the triangle to which the + We generate a texture visualizing the index of the triangle to which the texel belongs to:: import open3d as o3d from matplotlib import pyplot as plt box = o3d.geometry.TriangleMesh.create_box(create_uv_map=True) box = o3d.t.geometry.TriangleMesh.from_legacy(box) - # Creates a triangle attribute 'albedo' which is the triangle index + # Creates a triangle attribute 'albedo' which is the triangle index # multiplied by (255//12). box.triangle['albedo'] = (255//12)*np.arange(box.triangle.indices.shape[0], dtype=np.uint8) - # Initialize material and bake the 'albedo' triangle attribute to a + # Initialize material and bake the 'albedo' triangle attribute to a # texture. The texture will be automatically added to the material of # the object. box.material.set_default_properties() @@ -772,18 +809,18 @@ This function always uses the CPU device. R"(Sweeps the triangle mesh rotationally about an axis. Args: angle (float): The rotation angle in degree. - + axis (open3d.core.Tensor): The rotation axis. - + resolution (int): The resolution defines the number of intermediate sweeps about the rotation axis. - translation (float): The translation along the rotation axis. + translation (float): The translation along the rotation axis. Returns: A triangle mesh with the result of the sweep operation. Example: This code generates a spring with a triangle cross-section:: import open3d as o3d - + mesh = o3d.t.geometry.TriangleMesh([[1,1,0], [0.7,1,0], [1,0.7,0]], [[0,1,2]]) spring = mesh.extrude_rotation(3*360, [0,1,0], resolution=3*16, translation=2) o3d.visualization.draw([{'name': 'spring', 'geometry': spring}]) @@ -793,9 +830,9 @@ This function always uses the CPU device. "vector"_a, "scale"_a = 1.0, "capping"_a = true, R"(Sweeps the line set along a direction vector. Args: - + vector (open3d.core.Tensor): The direction vector. - + scale (float): Scalar factor which essentially scales the direction vector. Returns: A triangle mesh with the result of the sweep operation. diff --git a/cpp/tests/t/geometry/PointCloud.cpp b/cpp/tests/t/geometry/PointCloud.cpp index 76fea205c26..f0108c7fb64 100644 --- a/cpp/tests/t/geometry/PointCloud.cpp +++ b/cpp/tests/t/geometry/PointCloud.cpp @@ -72,7 +72,7 @@ TEST_P(PointCloudPermuteDevices, DefaultConstructor) { // ToString EXPECT_EQ(pcd.ToString(), - "PointCloud on CPU:0 [0 points] Attributes: None."); + "PointCloud on CPU:0 [0 points].\nAttributes: None."); } TEST_P(PointCloudPermuteDevices, ConstructFromPoints) { @@ -385,7 +385,7 @@ TEST_P(PointCloudPermuteDevices, Getters) { // ToString std::string text = "PointCloud on " + device.ToString() + - " [2 points (Float32)] Attributes: "; + " [2 points (Float32)].\nAttributes: "; EXPECT_THAT(pcd.ToString(), // Compiler dependent output AnyOf(text + "colors (dtype = Float32, shape = {2, 3}), labels " "(dtype = Float32, shape = {2, 3}).", diff --git a/python/test/core/test_core.py b/python/test/core/test_core.py index eab201c6034..0eac87f3c7e 100644 --- a/python/test/core/test_core.py +++ b/python/test/core/test_core.py @@ -29,6 +29,7 @@ import numpy as np import pytest import tempfile +import pickle import sys import os @@ -1584,3 +1585,22 @@ def test_iterator(device): o3_t_slice[:] = new_o3_t_slice np.testing.assert_equal(o3_t.cpu().numpy(), np.array([[0, 10, 20], [30, 40, 50]])) + + +@pytest.mark.parametrize("device", list_devices()) +def test_pickle(device): + o3_t = o3c.Tensor.ones((100), dtype=o3c.float32, device=device) + with tempfile.TemporaryDirectory() as temp_dir: + file_name = f"{temp_dir}/tensor.pkl" + pickle.dump(o3_t, open(file_name, "wb")) + o3_t_load = pickle.load(open(file_name, "rb")) + assert o3_t_load.device == device and o3_t_load.dtype == o3c.float32 + np.testing.assert_equal(o3_t.cpu().numpy(), o3_t_load.cpu().numpy()) + + # Test with a non-contiguous tensor. + o3_t_nc = o3_t[0:100:2] + pickle.dump(o3_t_nc, open(file_name, "wb")) + o3_t_nc_load = pickle.load(open(file_name, "rb")) + assert o3_t_nc_load.is_contiguous() + np.testing.assert_equal(o3_t_nc.cpu().numpy(), + o3_t_nc_load.cpu().numpy()) diff --git a/python/test/t/geometry/test_image.py b/python/test/t/geometry/test_image.py index d22720f956a..14bc8859b74 100644 --- a/python/test/t/geometry/test_image.py +++ b/python/test/t/geometry/test_image.py @@ -28,6 +28,8 @@ import open3d.core as o3c import numpy as np import pytest +import pickle +import tempfile import sys import os @@ -76,3 +78,14 @@ def test_buffer_protocol_cpu(device): im = im.to(device=device) dst_t = np.asarray(im.cpu()) np.testing.assert_array_equal(src_t, dst_t) + + +@pytest.mark.parametrize("device", list_devices()) +def test_pickle(device): + img = o3d.t.geometry.Image(o3c.Tensor.ones((10, 10, 3), o3c.uint8, device)) + with tempfile.TemporaryDirectory() as temp_dir: + file_name = f"{temp_dir}/img.pkl" + pickle.dump(img, open(file_name, "wb")) + img_load = pickle.load(open(file_name, "rb")) + assert img_load.as_tensor().allclose(img.as_tensor()) + assert img_load.device == img.device and img_load.dtype == o3c.uint8 diff --git a/python/test/t/geometry/test_lineset.py b/python/test/t/geometry/test_lineset.py index f843fc01055..9b79616bb8e 100644 --- a/python/test/t/geometry/test_lineset.py +++ b/python/test/t/geometry/test_lineset.py @@ -25,6 +25,16 @@ # ---------------------------------------------------------------------------- import open3d as o3d +import numpy as np +import pytest +import pickle +import tempfile + +import sys +import os + +sys.path.append(os.path.dirname(os.path.realpath(__file__)) + "/../..") +from open3d_test import list_devices def test_extrude_rotation(): @@ -42,3 +52,17 @@ def test_extrude_linear(): ans = lines.extrude_linear([0, 1, 0]) assert ans.vertex.positions.shape == (6, 3) assert ans.triangle.indices.shape == (4, 3) + + +@pytest.mark.parametrize("device", list_devices()) +def test_pickle(device): + line = o3d.t.geometry.LineSet([[0.7, 0, 0], [1, 0, 0]], [[0, 1]]).to(device) + with tempfile.TemporaryDirectory() as temp_dir: + file_name = f"{temp_dir}/lineset.pkl" + pickle.dump(line, open(file_name, "wb")) + line_load = pickle.load(open(file_name, "rb")) + assert line_load.device == device + np.testing.assert_equal(line_load.point.positions.cpu().numpy(), + line.point.positions.cpu().numpy()) + np.testing.assert_equal(line_load.line.indices.cpu().numpy(), + line.line.indices.cpu().numpy()) diff --git a/python/test/t/geometry/test_pointcloud.py b/python/test/t/geometry/test_pointcloud.py index c85ed376938..bb4f787d9b2 100644 --- a/python/test/t/geometry/test_pointcloud.py +++ b/python/test/t/geometry/test_pointcloud.py @@ -28,6 +28,8 @@ import open3d.core as o3c import numpy as np import pytest +import pickle +import tempfile import sys import os @@ -195,3 +197,18 @@ def test_extrude_linear(): ans = pcd.extrude_linear([0, 0, 1]) assert ans.point.positions.shape == (2, 3) assert ans.line.indices.shape == (1, 2) + + +@pytest.mark.parametrize("device", list_devices()) +def test_pickle(device): + pcd = o3d.t.geometry.PointCloud(device) + with tempfile.TemporaryDirectory() as temp_dir: + file_name = f"{temp_dir}/pcd.pkl" + pcd.point.positions = o3c.Tensor.ones((10, 3), + o3c.float32, + device=device) + pickle.dump(pcd, open(file_name, "wb")) + pcd_load = pickle.load(open(file_name, "rb")) + assert pcd_load.point.positions.device == device and pcd_load.point.positions.dtype == o3c.float32 + np.testing.assert_equal(pcd.point.positions.cpu().numpy(), + pcd_load.point.positions.cpu().numpy()) diff --git a/python/test/t/geometry/test_tensormap.py b/python/test/t/geometry/test_tensormap.py index 2140e49edd9..b5370aba2d0 100644 --- a/python/test/t/geometry/test_tensormap.py +++ b/python/test/t/geometry/test_tensormap.py @@ -28,6 +28,8 @@ import open3d.core as o3c import numpy as np import pytest +import pickle +import tempfile import sys import os @@ -316,3 +318,16 @@ def test_numpy_dict_modify(): np.testing.assert_equal(b_alias, [200]) np.testing.assert_equal(tm["a"], [200]) np.testing.assert_equal(tm["b"], [100]) + + +@pytest.mark.parametrize("device", list_devices()) +def test_pickle(device): + tm = o3d.t.geometry.TensorMap("positions") + with tempfile.TemporaryDirectory() as temp_dir: + file_name = f"{temp_dir}/tm.pkl" + tm.positions = o3c.Tensor.ones((10, 3), o3c.float32, device=device) + pickle.dump(tm, open(file_name, "wb")) + tm_load = pickle.load(open(file_name, "rb")) + assert tm_load.positions.device == device and tm_load.positions.dtype == o3c.float32 + np.testing.assert_equal(tm.positions.cpu().numpy(), + tm_load.positions.cpu().numpy()) diff --git a/python/test/t/geometry/test_trianglemesh.py b/python/test/t/geometry/test_trianglemesh.py index 7dce07bc071..bc95dbe5801 100644 --- a/python/test/t/geometry/test_trianglemesh.py +++ b/python/test/t/geometry/test_trianglemesh.py @@ -28,6 +28,8 @@ import open3d.core as o3c import numpy as np import pytest +import pickle +import tempfile import sys import os @@ -418,3 +420,19 @@ def test_extrude_linear(): ans = triangle.extrude_linear([0, 0, 1]) assert ans.vertex.positions.shape == (6, 3) assert ans.triangle.indices.shape == (8, 3) + + +@pytest.mark.parametrize("device", list_devices()) +def test_pickle(device): + mesh = o3d.t.geometry.TriangleMesh.create_box().to(device) + with tempfile.TemporaryDirectory() as temp_dir: + file_name = f"{temp_dir}/mesh.pkl" + pickle.dump(mesh, open(file_name, "wb")) + mesh_load = pickle.load(open(file_name, "rb")) + assert mesh_load.device == device + assert mesh_load.vertex.positions.dtype == o3c.float32 + assert mesh_load.triangle.indices.dtype == o3c.int64 + np.testing.assert_equal(mesh_load.vertex.positions.cpu().numpy(), + mesh.vertex.positions.cpu().numpy()) + np.testing.assert_equal(mesh_load.triangle.indices.cpu().numpy(), + mesh.triangle.indices.cpu().numpy())