How to form the frame from TensorRT output

[aravind.chakravarti]

I have a custom image preprocessing neural network model. It takes input of 1920x1080x3 and outputs 1920x1080x3 (It basically does preprocessing, so input is an image and output is an image).

I successfully tested code in Python and in ONNX.

However, I want to run the code in tensorRT C++. So, I took this code from github, and modifying it run my .engine file.

I read the image using

  image = cv::imread(path);
  cv::resize(image, image_resize, cv::Size(1920,1080), cv::INTER_LINEAR);
  cv::cvtColor(image_resize, image_cvtColor, cv::COLOR_BGR2RGB);

then created blob and copied it to GPU memory for inference,

cv::dnn::blobFromImage(tmp, out, 1 / (255.f), cv::Size(), cv::Scalar(0, 0, 0), false, false, CV_32F);

I read back the result using,

   // cv::Mat image = cv::Mat(width, height, CV_8UC4, (unsigned*)data);
    cv::Mat output1 = cv::Mat(num_rows, num_cols, CV_32F, static_cast<float*>(this->host_ptrs[0]+(0))); //CV_32F
    cv::Mat output2 = cv::Mat(num_rows, num_cols, CV_32F, static_cast<float*>(this->host_ptrs[0]+(7680))); //CV_32F
    cv::Mat output3 = cv::Mat(num_rows, num_cols, CV_32F, static_cast<float*>(this->host_ptrs[0]+(7680*2))); //CV_32F

    // Scale and convert floating-point data to 8-bit unsigned integers
    cv::Mat output1_8U, output2_8U, output3_8U;
    output1.convertTo(output1_8U, CV_8U, 255.0);
    output2.convertTo(output2_8U, CV_8U, 255.0);
    output3.convertTo(output3_8U, CV_8U, 255.0);

    // Combine the channels into a single RGB image
    std::vector<cv::Mat> channels = {output3_8U, output2_8U, output1_8U}; // BGR order
    cv::Mat rgbImage;
    cv::merge(channels, rgbImage);

    // Display the combined RGB image
    cv::namedWindow("win_name1", cv::WINDOW_NORMAL);
    cv::imshow("win_name1", rgbImage);
    cv::waitKey(0);

I think, I am doing something wrong here. I am getting output in grayscale. Am I doing any mistake in NxCxHxW to NxHxWxC ?

My input image is like this.

Ideally, I should get sharpened output, but I am getting output like this.

How to read pixels from static_cast<float*>(this->host_ptrs[0]+(0)).

Please provide complete information as applicable to your setup. Ideally I should read 6220800 (1920x1080x3) pixels from this address.

TensorRT Version : 8.4.1
GPU Type : A5000
Nvidia Driver Version :
CUDA Version : 11.4
CUDNN Version :
Operating System + Version : Ubuntu 20.04
Python Version (if applicable) : No, running on C++
TensorFlow Version (if applicable) : NA
PyTorch Version (if applicable) : NA
Baremetal or Container (if container which image + tag) : NA

[Fiona.Chen]

@aravind.chakravarti

I successfully tested code in Python and in ONNX

Please refer to the python and onnx code which you have run successfully. If the output layer dimension is NHWC(as you have said the model output 1920x1080x3), the R, G, B values are mixed.

[aravind.chakravarti]

Hi @Fiona.Chen,

Generally, to send the image to ONNX runtime we first convert time image to NCHW right?

Here is my code which feeds the image to ONNX runtime engine.

img = Image.open("image.jpg")

resize = transforms.Resize([1080, 1920])
img = resize(img)

to_tensor = transforms.ToTensor()
img = to_tensor(img)
print("Tensor image size {0}".format(img.shape))

# We need to have batch size as one. So, lets unsqueeze it
img.unsqueeze_(0)
print("Tensor image size after unsqueeze is {0}".format(img.shape))

ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(img)}

ort_outs = ort_session.run(None, ort_inputs)
...

And this is how I am getting the result from runtime engine. Which is working perfectly.

# Lets inspect the output tensor size
print(len(ort_outs))

# Since, we have passed only one image, we are getting the data from 0th index
img_out_y = ort_outs[0]
print("Image size out ONNX tensor output {0}".format(img_out_y.shape))

# Let us remove the batch (since we have only one image)
img_out_y = img_out_y[0,:,:,:]
print("Image size after removing the batch {}".format(img_out_y.shape))

# For PIL or CV2, we want W H DIMENSION. That is why, transforming the matrix
ONNX_filtered_img = np.asarray(img_out_y).transpose(1,2,0)
print("Transformed image size {}".format(ONNX_filtered_img.shape))

I have done inference of ONNX in python. I don’t know how to do it in C++

[aravind.chakravarti]

Hi @Fiona.Chen

I solved my issue!! Here is the code (it may help someone in future)

    /* Lets get information about image output */
    auto num_channels = this->output_bindings[0].dims.d[1];
    auto num_rows  = this->output_bindings[0].dims.d[2];
    auto num_cols  = this->output_bindings[0].dims.d[3];

    /* Print them only if necessary */
    #if 1
    std::cout << "num_channels " << num_channels << std::endl;  
    std::cout << "row_output " << num_rows << std::endl;
    std::cout << "col_output " << num_cols << std::endl;
    #endif 

    /* Image will be stored as continous array like as illustrated below
        _________________________________________________
        |   R PIXELS    |   G PIXELS    |   B PIXELS    |
        |  (1920x1080)  |  (1920x1080)  |  (1920x1080)  |
        |_______________|_______________|_______________|
       base             base            base 
       address          address +       address +
                        1980x1080       1920x1080x2

    */ 
    cv::Mat output1 = cv::Mat(num_rows, num_cols, CV_32F, static_cast<float*>(this->host_ptrs[0])+((num_cols*num_rows*0)-0));
    cv::Mat output2 = cv::Mat(num_rows, num_cols, CV_32F, static_cast<float*>(this->host_ptrs[0])+((num_cols*num_rows*1)-1));
    cv::Mat output3 = cv::Mat(num_rows, num_cols, CV_32F, static_cast<float*>(this->host_ptrs[0])+((num_cols*num_rows*2)-1));

    /* Bring float values to 0 <--> 255 */
    cv::Mat output1_8U, output2_8U, output3_8U;
    output1.convertTo(output1_8U, CV_8U, 255.0);
    output2.convertTo(output2_8U, CV_8U, 255.0);
    output3.convertTo(output3_8U, CV_8U, 255.0);

    /* Combine the channels into a single RGB image */
    std::vector<cv::Mat> channels = {output3_8U, output2_8U, output1_8U}; // BGR order
    cv::Mat rgbImage;
    cv::merge(channels, rgbImage);

    /* Display the combined RGB image */
    cv::namedWindow("Result", cv::WINDOW_NORMAL);
    // cv::imwrite("saved_image.jpg", rgbImage);
    cv::imshow("Result", rgbImage);
    cv::resizeWindow("Result", 1920/2, 1080/2);
    //cv::waitKey(0);