Method Export

python/mod_cvcuda/nvcv/TensorBatch.cpp:239–275

}

void TensorBatch::Export(py::module &m)
{
    using namespace py::literals;

    py::class_<TensorBatch, std::shared_ptr<TensorBatch>, Container>(
        m, "TensorBatch",
        "Container for a batch of tensors.\n"
        "The capacity of the container must be specified upfront in the batch initialization.\n"
        "The tensors in the batch may differ in shapes but they must have "
        "a uniform dimensionality, data type and layout.")
        .def(py::init(&TensorBatch::Create), "capacity"_a,
             "Create a new TensorBatch object with the specified capacity.")
        .def_property_readonly("layout", &TensorBatch::layout,
                               "Layout of the tensors in the tensor batch."
                               " None if the batch is empty.")
        .def_property_readonly("dtype", &TensorBatch::dtype,
                               "Data type of tensors in the tensor batch."
                               " None if the batch is empty.")
        .def_property_readonly("capacity", &TensorBatch::capacity, "Capacity of the tensor batch.")
        .def_property_readonly("ndim", &TensorBatch::rank,
                               "Return the number of dimensions of the tensors or -1 for an empty batch")
        .def("__len__", &TensorBatch::numTensors, "Return the number of tensors.")
        .def(
            "__iter__", [](const TensorBatch &batch) { return py::make_iterator(batch); },
            "Return an iterator over the tensors in the TensorBatch.")
        .def("__setitem__", &TensorBatch::set_at, "Set tensor at a given index.")
        .def("__getitem__", &TensorBatch::at, "Get a tensor at a given index.")
        .def("pushback", &TensorBatch::pushBack, "Add a new image to the end of the TensorBatch.")
        .def("pushback", &TensorBatch::pushBackMany, "Add multiple images to the end of the TensorBatch.")
        .def("popback", &TensorBatch::popBack, "count"_a = 1,
             "Remove one or more images from the end of the TensorBatch.")
        .def("clear", &TensorBatch::clear, "Remove all images from the TensorBatch.");

    m.def("as_tensors", &TensorBatch::WrapExternalBufferVector, "buffers"_a = std::vector<py::object>{},
          "layout"_a = std::nullopt, py::keep_alive<0, 1>(),
          "Wrap a list of external buffers as a batch of tensors, and tie the buffers lifetime to it");
}

} // namespace nvcvpy::priv