Function GetOutputTensorShape

python/mod_cvcuda/CvtColorUtil.cpp:299–328

}

nvcv::TensorShape GetOutputTensorShape(const nvcv::TensorShape &inputShape, nvcv::ImageFormat outputFormat,
                                       NVCVColorConversionCode code)
{
    if (inputShape.rank() < 3 || inputShape.rank() > 4)
    {
        throw std::runtime_error("Invalid input tensor shape, only NHWC or HWC are supported");
    }

    int64_t outputShape[4] = {};
    bool    heightIndex    = inputShape.rank() == 4 ? 1 : 0;
    for (int i = 0; i < inputShape.rank(); i++)
    {
        outputShape[i] = inputShape[i];
    }
    int channelIndex = inputShape.rank() == 4 ? 3 : 2;

    outputShape[heightIndex]  = GetOutputHeight(outputShape[heightIndex], code);
    outputShape[channelIndex] = outputFormat.numChannels();

    if (inputShape.rank() == 4)
    {
        return nvcv::TensorShape({outputShape[0], outputShape[1], outputShape[2], outputShape[3]}, "NHWC");
    }
    else
    {
        assert(inputShape.rank() == 3);

        return nvcv::TensorShape({outputShape[0], outputShape[1], outputShape[2]}, "HWC");
    }
}