Function ToPyBufferInfo

python/mod_cvcuda/nvcv/Image.cpp:764–959

namespace {

std::vector<std::pair<py::buffer_info, nvcv::TensorLayout>> ToPyBufferInfo(const nvcv::ImageDataStrided     &imgData,
                                                                           std::optional<nvcv::TensorLayout> userLayout)
{
    if (imgData.numPlanes() < 1)
    {
        return {};
    }

    const nvcv::ImagePlaneStrided &firstPlane = imgData.plane(0);

    std::optional<nvcv::TensorLayoutInfoImage> infoLayout;
    if (userLayout)
    {
        if (auto tmp = nvcv::TensorLayoutInfoImage::Create(*userLayout))
        {
            infoLayout.emplace(std::move(*tmp));
        }
        else
        {
            throw std::runtime_error("Layout can't represent the planar images needed");
        }
    }

    bool singleBuffer = true;

    // Let's check if we can return only one buffer, depending
    // on the planes dimensions, pitch and data type.
    for (int p = 1; p < imgData.numPlanes(); ++p)
    {
        const nvcv::ImagePlaneStrided &plane = imgData.plane(p);

        if (plane.width != firstPlane.width || plane.height != firstPlane.height
            || plane.rowStride != firstPlane.rowStride || imgData.format().planeDataType(0).numChannels() >= 2
            || imgData.format().planeDataType(0) != imgData.format().planeDataType(p))
        {
            singleBuffer = false;
            break;
        }

        // check if using the same plane pitch
        if (p >= 2)
        {
            intptr_t goldPlaneStrided = imgData.plane(1).basePtr - imgData.plane(0).basePtr;
            intptr_t curPlaneStrided  = imgData.plane(p).basePtr - imgData.plane(p - 1).basePtr;
            if (curPlaneStrided != goldPlaneStrided)
            {
                singleBuffer = false;
                break;
            }
        }
    }

    std::vector<std::pair<py::buffer_info, nvcv::TensorLayout>> out;

    // If not using a single buffer, we'll forcibly use one buffer per plane.
    int numBuffers = singleBuffer ? 1 : imgData.numPlanes();

    for (int p = 0; p < numBuffers; ++p)