Enable pickling for tensor and tensor based geometry

.gitignore

@@ -90,3 +90,6 @@ docs/Doxyfile
 docs/getting_started.rst
 docs/docker.rst
 docs/tensorboard.md
+
+# test
+*.pkl

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())) ==

cpp/pybind/core/device.cpp

@@ -46,7 +46,20 @@ void pybind_core_device(py::module &m) {
             .def("__repr__", &Device::ToString)
             .def("__str__", &Device::ToString)
             .def("get_type", &Device::GetType)
-            .def("get_id", &Device::GetID);
+            .def("get_id", &Device::GetID)
+            .def(py::pickle(
+                    [](const Device &d) {
+                        return py::make_tuple(d.GetType(), d.GetID());
+                    },
+                    [](py::tuple t) {
+                        if (t.size() != 2) {
+                            utility::LogError(
+                                    "Invalid state! Expecting a tuple of size "
+                                    "2.");
+                        }
+                        return Device(t[0].cast<Device::DeviceType>(),
+                                      t[1].cast<int>());
+                    }));
 
     py::enum_<Device::DeviceType>(device, "DeviceType")
             .value("CPU", Device::DeviceType::CPU)

cpp/pybind/core/tensor.cpp

@@ -356,6 +356,34 @@ void pybind_core_tensor(py::module& m) {
     BindTensorFullCreation<bool>(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(
+                            "Invalid state! Expecting a tuple of size 2.");
+                }
+                const Device& device = t[0].cast<Device>();
+                if (device.IsCUDA() && !core::cuda::IsAvailable()) {
+                    utility::LogWarning(
+                            "CUDA is not available, tensor will be "
+                            "created on CPU.");
+                    return PyArrayToTensor(t[1].cast<py::array>(), true);
+                } else {
+                    return PyArrayToTensor(t[1].cast<py::array>(), true)
+                            .To(device);
+                }
+            }));
+
     tensor.def_static(
             "eye",
             [](int64_t n, utility::optional<Dtype> dtype,

cpp/pybind/t/geometry/image.cpp

@@ -101,6 +101,22 @@ 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(
+                            "Invalid state! Expecting a tuple of size 1.");
+                }
+                return Image(t[0].cast<core::Tensor>());
+            }));
+
     // Buffer protocol.
     image.def_buffer([](Image &I) -> py::buffer_info {
         if (!I.IsCPU()) {
@@ -282,6 +298,26 @@ void pybind_image(py::module &m) {
             .def(py::init<const Image &, const Image &, bool>(),
                  "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(
+                                    "Invalid state! Expecting a tuple of size "
+                                    "3.");
+                        }
+
+                        return RGBDImage(t[0].cast<Image>(), t[1].cast<Image>(),
+                                         t[2].cast<bool>());
+                    }))
+
             // Depth and color images.
             .def_readwrite("color", &RGBDImage::color_, "The color image.")
             .def_readwrite("depth", &RGBDImage::depth_, "The depth image.")

cpp/pybind/t/geometry/lineset.cpp

@@ -29,6 +29,7 @@
 #include <string>
 #include <unordered_map>
 
+#include "open3d/core/CUDAUtils.h"
 #include "open3d/t/geometry/TriangleMesh.h"
 #include "pybind/docstring.h"
 #include "pybind/t/geometry/geometry.h"
@@ -109,6 +110,41 @@ 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(
+                            "Invalid state! Expecting a tuple of size 3.");
+                }
+
+                LineSet line_set(t[0].cast<core::Device>());
+                if (!core::cuda::IsAvailable()) {
+                    utility::LogWarning(
+                            "CUDA is not available. LineSet will be created on "
+                            "CPU.");
+                    line_set.To(core::Device("CPU:0"));
+                }
+
+                const TensorMap point_attr = t[1].cast<TensorMap>();
+                const TensorMap line_attr = t[2].cast<TensorMap>();
+                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

cpp/pybind/t/geometry/pointcloud.cpp

@@ -127,6 +127,38 @@ 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(
+                            "Invalid state! Expecting a tuple of size 2.");
+                }
+
+                PointCloud pcd(t[0].cast<core::Device>());
+                if (!core::cuda::IsAvailable()) {
+                    utility::LogWarning(
+                            "CUDA is not available. PointCloud will be "
+                            "created on CPU.");
+                    pcd.To(core::Device("CPU:0"));
+                }
+
+                const TensorMap map_keys_to_tensors = t[1].cast<TensorMap>();
+                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.

cpp/pybind/t/geometry/tensormap.cpp

@@ -164,6 +164,28 @@ 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<std::string, core::Tensor> 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(
+                            "Invalid state! Expecting a tuple of size 2.");
+                }
+                return TensorMap(t[0].cast<std::string>(),
+                                 t[1].cast<std::unordered_map<std::string,
+                                                              core::Tensor>>());
+            }));
+
     tm.def("__setattr__",
            [](TensorMap &m, const std::string &key, const core::Tensor &val) {
                if (!TensorMap::GetReservedKeys().count(key)) {

cpp/pybind/t/geometry/trianglemesh.cpp

@@ -107,6 +107,40 @@ 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(
+                            "Invalid state! Expecting a tuple of size 3.");
+                }
+
+                TriangleMesh mesh(t[0].cast<core::Device>());
+                if (!core::cuda::IsAvailable()) {
+                    utility::LogWarning(
+                            "CUDA is not available. TriangleMesh will be "
+                            "created on CPU.");
+                    mesh.To(core::Device("CPU:0"));
+                }
+
+                const TensorMap vertex_attr = t[1].cast<TensorMap>();
+                const TensorMap triangle_attr = t[2].cast<TensorMap>();
+                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

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}).",

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,19 @@ 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)
+    pickle.dump(o3_t, open("tensor.pkl", "wb"))
+    o3_t_load = pickle.load(open("tensor.pkl", "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("tensor.pkl", "wb"))
+    o3_t_nc_load = pickle.load(open("tensor.pkl", "rb"))
+    assert o3_t_nc_load.is_contiguous()
+    np.testing.assert_equal(o3_t_nc.cpu().numpy(), o3_t_nc_load.cpu().numpy())

python/test/t/geometry/test_image.py

@@ -28,6 +28,7 @@
 import open3d.core as o3c
 import numpy as np
 import pytest
+import pickle
 
 import sys
 import os
@@ -76,3 +77,12 @@ 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))
+    pickle.dump(img, open("img.pkl", "wb"))
+    img_load = pickle.load(open("img.pkl", "rb"))
+    assert img_load.as_tensor().allclose(img.as_tensor())
+    assert img_load.device == img.device and img_load.dtype == o3c.uint8

python/test/t/geometry/test_lineset.py

@@ -25,6 +25,15 @@
 # ----------------------------------------------------------------------------
 
 import open3d as o3d
+import numpy as np
+import pytest
+import pickle
+
+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 +51,15 @@ 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)
+    pickle.dump(line, open("lineset.pkl", "wb"))
+    line_load = pickle.load(open("lineset.pkl", "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())