tensor转cv::Mat(即CHW转HWC)原理含C#代码实现

起因是博主在实习过程中的一个任务:需要将模型预测输出tensor的shape从CHW(严格来说是NCHW,但是N=1所以这里忽略掉)转成OpenCV中的cv::Mat类型(即HWC)数据。
由于博主对C++和C#并不是很熟悉,而且也查阅了tensor对应的类似resize、reshape等方法,发现对于输出tensor并没有封装相应的方法,无奈只好用最原始的办法,就是将输出tensor的值赋值到cv::Mat对应的位置。
原理参考了下面两篇博客,其实如果弄清楚原理了代码还是挺好写的:

  1. 【opencv】——cv mat转tensor(b,c,h,w)
  2. NCHW转化为NHWC
    在这里插入图片描述

这里借用第二篇参考博客的图,第一行就是我们tensor拉成一维以后的数据的排列方式,红色绿色蓝色分别对应着RGB的通道。同理第二行对应的是cv::Mat中RGB数据的排列方式。
这里我先做几个假设,假设模型输入的图像为BGR,输出也为BGR,这样我就可以直接调用OpenCV的imwrite保存图像。假设模型一次只输入一张图片的tensor并且得到对应一张图片的输出tensor,且输出tensor的shape=[NCHW]=[1,3,H,W]
对于输出tensor的shape可以这样理解,我们按照NCHW的顺序一层层剖析它。首先,N=1可以忽视掉,C=3,就是对应我们三个通道。每个通道里面有 H × W H \times W H×W个数据(可以联系上面那个图来理解,对应每种颜色的全部数据)。同理,每行里面里面就有 W W W个数据。
我们可以这样思考,首先每个通道都有 H × W H \times W H×W个数据,还是联系上面那个图,按照之前我的假设,第一个方块的数据值就对应cv::Mat数据第一个通道第一行第一列的值。此时,第一个方块后面的第 H × W H \times W H×W个方块的值就对应cv::Mat数据第二个通道第一行第一列的值。我们可以按照这个规律每次填充cv::Mat不同通道同一个位置的值。当然也可以先填充同一个通道不同位置的值,取决于习惯,这里我按照参考的第一篇博客先填充不同通道同一个位置的值。

如果上面的没有看懂也没关系,可能是我写复杂了,直接看代码并试着结合图看看,说不定这样更好理解。

下面是代码:

auto output_tensor = predictor_->GetOutput(0);
auto output_data = output_tensor->data<float>();
auto output_shape = output_tensor->shape();

cv::Mat rgb_img = cv::Mat::zeros(cv::Size(output_shape[3], output_shape[2]), CV_8UC3);

int h = rgb_img.rows;
int w = rgb_img.cols;
int c = rgb_img.channels();
for (int i = 0; i < h; i++) {
    for (int j = 0; j < w; j++) {
        for (int k = 0; k < c; k++) {
            int offset = k * h * w + i * w + j * 1;
            int tmp_pix = output_data[offset] * 255;
            rgb_img.at<cv::Vec3b>(i, j)[k] = tmp_pix > 255 ? 255 : tmp_pix;
        }
    }
}

cv::imwrite("/data/data/com.baidu.paddle.lite.demo.face_detection/test.jpg", rgb_img);

return rgb_img;

首先这部分只是为了获取输出tensor以及shape,方便我们构造合适大小的cv::Mat

auto output_tensor = predictor_->GetOutput(0);
auto output_data = output_tensor->data<float>();
auto output_shape = output_tensor->shape();

然后我们初始化cv::Mat用于存储图像数据

cv::Mat rgb_img = cv::Mat::zeros(cv::Size(output_shape[3], output_shape[2]), CV_8UC3);

这部分是核心,由于输出tensor的范围在0-1,这里做了处理转换回0-255

int h = rgb_img.rows;
int w = rgb_img.cols;
int c = rgb_img.channels();
for (int i = 0; i < h; i++) {
    for (int j = 0; j < w; j++) {
        for (int k = 0; k < c; k++) {
            int offset = k * h * w + i * w + j * 1;
            int tmp_pix = output_data[offset] * 255;
            rgb_img.at<cv::Vec3b>(i, j)[k] = tmp_pix > 255 ? 255 : tmp_pix;
        }
    }
}

你可能感兴趣的:(实验记录,计算机视觉,opencv,c#,c++,android-studio)