tensor_converter.cpp

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        CloudViewer: www.cloudViewer.org                  -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2024 www.cloudViewer.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
#include "pybind/core/tensor_converter.h"
 
#include <Logging.h>
 
#include "core/Tensor.h"
#ifdef _MSC_VER
#pragma warning(disable : 4996)  // Use of [[deprecated]] feature
#endif
#include "pybind/cloudViewer_pybind.h"
#include "pybind/core/core.h"
#include "pybind/pybind_utils.h"
 
namespace cloudViewer {
namespace core {
static Tensor CastOptionalDtypeDevice(const Tensor& t,
                                      utility::optional<Dtype> dtype,
                                      utility::optional<Device> device) {
    Tensor t_cast = t;
    if (dtype.has_value()) {
        t_cast = t_cast.To(dtype.value());
    }
    if (device.has_value()) {
        t_cast = t_cast.To(device.value());
    }
    return t_cast;
}
 
py::array TensorToPyArray(const Tensor& tensor) {
    if (!tensor.IsCPU()) {
        utility::LogError(
                "Can only convert CPU Tensor to numpy. Copy Tensor to CPU "
                "before converting to numpy.");
    }
    py::dtype py_dtype =
            py::dtype(pybind_utils::DtypeToArrayFormat(tensor.GetDtype()));
    py::array::ShapeContainer py_shape(tensor.GetShape());
    SizeVector strides = tensor.GetStrides();
    int64_t element_byte_size = tensor.GetDtype().ByteSize();
    for (auto& s : strides) {
        s *= element_byte_size;
    }
    py::array::StridesContainer py_strides(strides);
 
    // `base_tensor` is a shallow copy of `tensor`. `base_tensor`
    // is on the heap and is owned by py::capsule
    // `base_tensor_capsule`. The capsule is referenced as the
    // "base" of the numpy tensor returned by o3d.Tensor.numpy().
    // When the "base" goes out-of-scope (e.g. when all numpy
    // tensors referencing the base have gone out-of-scope), the
    // deleter is called to free the `base_tensor`.
    //
    // This behavior is important when the original `tensor` goes
    // out-of-scope while we still want to keep the data alive.
    // e.g.
    //
    // ```python
    // def get_np_tensor():
    //     o3d_t = o3d.Tensor(...)
    //     return o3d_t.numpy()
    //
    // # Now, `o3d_t` is out-of-scope, but `np_t` still
    // # references the base tensor which references the
    // # underlying data of `o3d_t`. Thus np_t is still valid.
    // # When np_t goes out-of-scope, the underlying data will be
    // # finally freed.
    // np_t = get_np_tensor()
    // ```
    //
    // See:
    // https://stackoverflow.com/questions/44659924/returning-numpy-arrays-via-pybind11
    Tensor* base_tensor = new Tensor(tensor);
 
    // See PyTorch's torch/csrc/Module.cpp
    auto capsule_destructor = [](PyObject* data) {
        Tensor* base_tensor = reinterpret_cast<Tensor*>(
                PyCapsule_GetPointer(data, "open3d::Tensor"));
        if (base_tensor) {
            delete base_tensor;
        } else {
            PyErr_Clear();
        }
    };
 
    py::capsule base_tensor_capsule(base_tensor, "open3d::Tensor",
                                    capsule_destructor);
    return py::array(py_dtype, py_shape, py_strides, tensor.GetDataPtr(),
                     base_tensor_capsule);
}
 
Tensor PyArrayToTensor(py::array array, bool inplace) {
    py::buffer_info info = array.request();
 
    SizeVector shape(info.shape.begin(), info.shape.end());
    SizeVector strides(info.strides.begin(), info.strides.end());
    for (size_t i = 0; i < strides.size(); ++i) {
        strides[i] /= info.itemsize;
    }
    Dtype dtype = pybind_utils::ArrayFormatToDtype(info.format, info.itemsize);
    Device device("CPU:0");
 
    auto shared_array = std::make_shared<py::array>(array);
    std::function<void(void*)> deleter = [shared_array](void*) mutable -> void {
        py::gil_scoped_acquire acquire;
        shared_array.reset();
    };
    auto blob = std::make_shared<Blob>(device, info.ptr, deleter);
    Tensor t_inplace(shape, strides, info.ptr, dtype, blob);
 
    if (inplace) {
        return t_inplace;
    } else {
        return t_inplace.Clone();
    }
}
 
Tensor PyListToTensor(const py::list& list,
                      utility::optional<Dtype> dtype,
                      utility::optional<Device> device) {
    py::object numpy = py::module::import("numpy");
    py::array np_array = numpy.attr("array")(list);
    Tensor t = PyArrayToTensor(np_array, false);
    return CastOptionalDtypeDevice(t, dtype, device);
}
 
Tensor PyTupleToTensor(const py::tuple& tuple,
                       utility::optional<Dtype> dtype,
                       utility::optional<Device> device) {
    py::object numpy = py::module::import("numpy");
    py::array np_array = numpy.attr("array")(tuple);
    Tensor t = PyArrayToTensor(np_array, false);
    return CastOptionalDtypeDevice(t, dtype, device);
}
 
Tensor DoubleToTensor(double scalar_value,
                      utility::optional<Dtype> dtype,
                      utility::optional<Device> device) {
    Dtype dtype_value = core::Float64;
    if (dtype.has_value()) {
        dtype_value = dtype.value();
    }
    Device device_value("CPU:0");
    if (device.has_value()) {
        device_value = device.value();
    }
    return Tensor(std::vector<double>{scalar_value}, {}, core::Float64,
                  device_value)
            .To(dtype_value);
}
 
Tensor IntToTensor(int64_t scalar_value,
                   utility::optional<Dtype> dtype,
                   utility::optional<Device> device) {
    Dtype dtype_value = core::Int64;
    if (dtype.has_value()) {
        dtype_value = dtype.value();
    }
    Device device_value("CPU:0");
    if (device.has_value()) {
        device_value = device.value();
    }
    return Tensor(std::vector<int64_t>{scalar_value}, {}, core::Int64,
                  device_value)
            .To(dtype_value);
}
 
Tensor BoolToTensor(bool scalar_value,
                    utility::optional<Dtype> dtype,
                    utility::optional<Device> device) {
    Dtype dtype_value = core::Bool;
    if (dtype.has_value()) {
        dtype_value = dtype.value();
    }
    Device device_value("CPU:0");
    if (device.has_value()) {
        device_value = device.value();
    }
    return Tensor(std::vector<bool>{scalar_value}, {}, core::Bool, device_value)
            .To(dtype_value);
}
 
Tensor PyHandleToTensor(const py::handle& handle,
                        utility::optional<Dtype> dtype,
                        utility::optional<Device> device,
                        bool force_copy) {
    // 1) bool
    // 2) int
    // 3) float (double)
    // 4) list
    // 5) tuple
    // 6) numpy.ndarray (value will be copied)
    // 7) Tensor (value will be copied)
    if (py::isinstance<py::bool_>(handle)) {
        return BoolToTensor(
                static_cast<bool>(py::reinterpret_borrow<py::bool_>(handle)),
                dtype, device);
    } else if (py::isinstance<py::int_>(handle)) {
        return IntToTensor(
                static_cast<int64_t>(py::reinterpret_borrow<py::int_>(handle)),
                dtype, device);
    } else if (py::isinstance<py::float_>(handle)) {
        return DoubleToTensor(
                static_cast<double>(py::reinterpret_borrow<py::float_>(handle)),
                dtype, device);
    } else if (py::isinstance<py::list>(handle)) {
        return PyListToTensor(py::reinterpret_borrow<py::list>(handle), dtype,
                              device);
    } else if (py::isinstance<py::tuple>(handle)) {
        return PyTupleToTensor(py::reinterpret_borrow<py::tuple>(handle), dtype,
                               device);
    } else if (py::isinstance<py::array>(handle)) {
        return CastOptionalDtypeDevice(
                PyArrayToTensor(py::reinterpret_borrow<py::array>(handle),
                                /*inplace=*/!force_copy),
                dtype, device);
    } else if (py::isinstance<Tensor>(handle)) {
        try {
            Tensor* tensor = handle.cast<Tensor*>();
            if (force_copy) {
                return CastOptionalDtypeDevice(tensor->Clone(), dtype, device);
            } else {
                return CastOptionalDtypeDevice(*tensor, dtype, device);
            }
        } catch (...) {
            utility::LogError("Cannot cast index to Tensor.");
        }
    } else {
        utility::LogError(
                "PyHandleToTensor has invalid input type {}.",
                static_cast<std::string>(py::str(py::type::of(handle))));
    }
}
 
}  // namespace core
}  // namespace cloudViewer