Loading

cv::Mat格式转换与NCHW内存排列

背景

项目需要将相机驱动输入的图像变为tensor输入模型进行检测。
输入为NV12格式,输出tesnor要求BGR格式,且内存排列为NCHW。

格式

NV12转BGR比较简单,可以调用cv自带的接口进行转换并将0-255归一为0-1压缩到模型输入尺寸。

// restore original image then resize.
cv::Mat yuv_mat(image_height_ * 3 / 2, image_width_, CV_8UC1, (uint8_t*)yuv_data_arr);
cv::Mat bgr_mat;
cv::cvtColor(yuv_mat, bgr_mat, cv::COLOR_YUV2BGR_NV12);
cv::Mat model_mat_8u;
cv::resize(bgr_mat, model_mat_8u, cv::Size(input_w, input_h));
// input normalize
cv::Mat model_mat_32f;
model_mat_8u.convertTo(model_mat_32f, CV_32F);

排列

cv::Mat的默认排列是NHWC而onnxruntime只支持NHWC排列。
简单解释一下,以2x2图像为例:

NHWC:
{{{BGR}, {BGR}}, {{BGR}, {BGR}}}

NCHW:
{{{BB},{BB}}, {{GG},{GG}}, {{RR},{RR}}}

转换代码如下:

#include <opencv2/opencv.hpp>

void hwc_to_chw(cv::InputArray src, cv::OutputArray dst) {
  std::vector<cv::Mat> channels;
  cv::split(src, channels);

  // Stretch one-channel images to vector
  for (auto &img : channels) {
    img = img.reshape(1, 1); 
  }

  // Concatenate three vectors to one
  cv::hconcat( channels, dst );
}

int main()
{
    uchar temp[12] = {255, 0, 0, 255, 0, 0,\ 
                      255, 0, 0, 255, 0, 0}; 

    cv::Mat hwc = cv::Mat(2, 2, CV_8UC3);
    hwc.data = &temp[0];
    for(int i = 0; i < 12; ++i)
    {   
        std::cout << unsigned(*(hwc.data+i)) << ", ";
    }   
    std::cout << "\n";
    // reorder memory layout
    cv::Mat chw;
    hwc_to_chw(hwc, chw);
    for(int i = 0; i < 12; ++i)
    {   
        std::cout << unsigned(*(chw.data+i)) << ", ";
    }   
    std::cout << "\n";

    return 0;
}

验证效果

$ g++ reorder.cpp -o reorder -I /usr/include/opencv4/ -L /usr/lib/ -lopencv_core
$ ./reorder
255, 0, 0, 255, 0, 0, 255, 0, 0, 255, 0, 0, 
255, 255, 255, 255, 0, 0, 0, 0, 0, 0, 0, 0,

参考

Converting NV12 to BGR - C++ - OpenCV
neural network - How cv::MAT convert to NCHW format? - Stack Overflow

posted @ 2024-06-28 11:44  azureology  阅读(17)  评论(0编辑  收藏  举报