Enable pickling for tensor and tensor based geometry (#5509)

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

@@ -37,16 +37,29 @@ void pybind_core_device(py::module &m) {
     py::class_<Device> device(
             m, "Device",
             "Device context specifying device type and device id.");
-    device.def(py::init<>())
-            .def(py::init<Device::DeviceType, int>())
-            .def(py::init<const std::string &, int>())
-            .def(py::init<const std::string &>())
-            .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::DeviceType, int>());
+    device.def(py::init<const std::string &, int>());
+    device.def(py::init<const std::string &>());
+    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<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,36 @@ 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(
+                            "Cannot unpickle Tensor! Expecting a tuple of size "
+                            "2.");
+                }
+                const Device& device = t[0].cast<Device>();
+                if (!device.IsAvailable()) {
+                    utility::LogWarning(
+                            "Device {} is not available, tensor will be "
+                            "created on CPU.",
+                            device.ToString());
+                    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,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<core::Tensor>());
+            }));
+
     // 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<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(
+                                    "Cannot unpickle RGBDImage! 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,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<core::Device>();
+                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<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,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<core::Device>();
+                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<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,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<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(
+                            "Cannot unpickle TensorMap! 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)) {