Class Tensor

python/mod_cvcuda/include/nvcv/python/Tensor.hpp:38–81

namespace py = pybind11;

class Tensor
    : public Resource
    , public nvcv::Tensor
{
public:
    static Tensor Create(const nvcv::TensorShape &tshape, nvcv::DataType dtype, int rowalign = 0)
    {
        PyObject *otensor = capi().Tensor_Create(tshape.size(), &tshape[0], static_cast<NVCVDataType>(dtype),
                                                 static_cast<NVCVTensorLayout>(tshape.layout()), rowalign);
        CheckCAPIError();
        NVCV_ASSERT(otensor != nullptr);
        py::object pytensor = py::reinterpret_steal<py::object>(otensor);

        return Tensor(pytensor);
    }

    static Tensor Create(const Shape &shape, nvcv::DataType dtype, nvcv::TensorLayout layout = nvcv::TENSOR_NONE,
                         int rowalign = 0)
    {
        return Create(CreateNVCVTensorShape(shape, layout), dtype, rowalign);
    }

    static Tensor CreateForImageBatch(int numImages, nvcv::Size2D size, nvcv::ImageFormat fmt, int rowalign = 0)
    {
        PyObject *otensor
            = capi().Tensor_CreateForImageBatch(numImages, size.w, size.h, static_cast<NVCVImageFormat>(fmt), rowalign);
        CheckCAPIError();
        NVCV_ASSERT(otensor != nullptr);
        py::object pytensor = py::reinterpret_steal<py::object>(otensor);

        return Tensor(pytensor);
    }

private:
    friend struct py::detail::type_caster<Tensor>;

    Tensor() = default;

    explicit Tensor(py::object obj)
        : Resource(obj)
        , nvcv::Tensor(FromHandle(CheckCAPIError(capi().Tensor_GetHandle(this->ptr())), true))
    {
    }
};

} // namespace nvcvpy