矩阵和tensor相互转换
cvmat到tensor
tips:
这里主要要注意的就是在opencv和pytorch中存储顺序的差异
cv::cvtColor(frame, frame, CV_BGR2RGB);
//normalization
frame.convertTo(frame, CV_32FC3, 1.0f / 255.0f);
//opencv format H*W*C
auto input_tensor = torch::from_blob(frame.data, {1, frame_h, frame_w, kCHANNELS});
//pytorch format N*C*H*W
input_tensor = input_tensor.permute({0, 3, 1, 2});
tensor 到cvmat
tips:
1.squeeze只用于batchsize为1的场景
2.permute 是将存储格式从pytorch形式转成opencv格式
3.因为在处理前对cvmat中的值做了归一化,所以现在要*255恢复,同时对于不在0-255范围内的数据,需要做限制
4.因为cvmat的数据格式是8UC3,所以torch tensor要提前转换成kU8
//send tensor to cpu
input_tensor = input_tensor.to(at::kCUDA);
//inference
torch::Tensor out_tensor = module->forward({input_tensor}).toTensor();
//sequeeze trans tensor shape from 1*C*H*W to C*H*W
//permute C*H*W to H*W*C
out_tensor = out_tensor.squeeze().detach().permute({1, 2, 0});
//see tip3,tip4
out_tensor = out_tensor.mul(255).clamp(0, 255).to(torch::kU8);
out_tensor = out_tensor.to(torch::kCPU);
cv::Mat resultImg(frame_h, frame_w, CV_8UC3);
//copy the data from out_tensor to resultImg
std::memcpy((void *) resultImg.data, out_tensor.data_ptr(), sizeof(torch::kU8) * out_tensor.numel());
上面是参考博客中的内容,应该是正确的,我没有测试。下面是我自己使用的,确认正确的。
opencv转Tensor
cv::Mat img = cv::imread("test.jpg", cv::IMREAD_COLOR);
cvtColor(img, img, cv::COLOR_BGR2RGB);
int h, w;
h = w = 0;
h = img.rows;
w = img.cols;
int short_size = 1024;
cv::Mat im_resized;
resize_image(img, im_resized, 1024);
torch::Tensor tensor_image = torch::from_blob(im_resized.data, { 1,im_resized.rows, im_resized.cols,3 }, torch::kByte);
tensor_image = tensor_image.permute({ 0,3,1,2 });
tensor_image = tensor_image.toType(torch::kFloat);
Tensor转opencv Mat (我这里转的是单通道图像)
int height, width;
height = _bitmap.size(0);
width = _bitmap.size(1);
_bitmap = _bitmap.mul(255).clamp(0, 255).to(torch::kU8);
_bitmap = _bitmap.to(torch::kCPU);
cv::Mat imgbin(cv::Size(width, height), CV_8U, _bitmap.data_ptr());
cv::imwrite("imgbin.jpg",imgbin);
对于浮点型Tensor转浮点型opencv Mat,我从github上找到一个,但是性能有很大问题,如一张1000*1000的尺寸就要进行百万次的赋值操作,耗时严重。可做参考,不推荐使用。
cv::Mat tensor2Mat(torch::Tensor &i_tensor) {
auto i_tensor_arr = i_tensor.accessor<float, 2>();
int h = i_tensor_arr.size(0), w = i_tensor_arr.size(1);
cv::Mat o_Mat = cv::Mat::zeros(h, w, CV_32FC1);
for (int x = 0; x < h; x++) {
for (int y = 0; y < w; y++) {
o_Mat.at<float>(x, y) = i_tensor[x][y].item<float>();
}
}
return o_Mat;
}
下面这个是我自己写的,速度有了质的飞跃,相比与上面的百万次的赋值操作,这里只有一次内存拷贝操作:
cv::Mat tensor2Mat(torch::Tensor &i_tensor) {
int height = i_tensor.size(0), width = i_tensor.size(1);
//i_tensor = i_tensor.to(torch::kF32);
i_tensor = i_tensor.to(torch::kCPU);
cv::Mat o_Mat(cv::Size(width, height), CV_32F, i_tensor.data_ptr());
return o_Mat;
}