![](/img/trans.png)
[英]How to convert PyTorch tensor to C++ torch::Tensor vice versa?
[英]Convert pytorch tensor to opencv mat and vice versa in C++
我想在 C++ 中将 pytorch 张量转换为 opencv mat,反之亦然。 我有这两个功能:
cv::Mat TensorToCVMat(torch::Tensor tensor)
{
std::cout << "converting tensor to cvmat\n";
tensor = tensor.squeeze().detach().permute({1, 2, 0});
tensor = tensor.mul(255).clamp(0, 255).to(torch::kU8);
tensor = tensor.to(torch::kCPU);
int64_t height = tensor.size(0);
int64_t width = tensor.size(1);
cv::Mat mat(width, height, CV_8UC3);
std::memcpy((void *)mat.data, tensor.data_ptr(), sizeof(torch::kU8) * tensor.numel());
return mat.clone();
}
torch::Tensor CVMatToTensor(cv::Mat mat)
{
std::cout << "converting cvmat to tensor\n";
cv::cvtColor(mat, mat, cv::COLOR_BGR2RGB);
cv::Mat matFloat;
mat.convertTo(matFloat, CV_32F, 1.0 / 255);
auto size = matFloat.size();
auto nChannels = matFloat.channels();
auto tensor = torch::from_blob(matFloat.data, {1, size.height, size.width, nChannels});
return tensor.permute({0, 3, 1, 2});
}
在我的代码中,我加载了两个图像( image1
和image2
),我想将它们转换为 pytorch 张量,然后返回到 opencv mat 以检查它是否有效。 问题是我在第一次调用TensorToCVMat
内存访问错误,我无法弄清楚出了什么问题,因为我对 C++ 编程没有太多经验。
cv::Mat image1;
image1 = cv::imread(argv[1]);
if (!image1.data)
{
std::cout << "no image data\n";
return -1;
}
cv::Mat image2;
image2 = cv::imread(argv[2]);
if (!image2.data)
{
std::cout << "no image data\n";
return -1;
}
torch::Tensor tensor1 = CVMatToTensor(image1);
cv::Mat new_image1 = TensorToCVMat(tensor1); // <<< this is where the memory access error is thrown
torch::Tensor tensor2 = CVMatToTensor(image2);
cv::Mat new_image2 = TensorToCVMat(tensor2);
如果你能给我提示或解释来解决这个问题,那就太好了。
不确定错误是否发生在 memcpy 步骤。 但是您可以使用Mat构造函数的void* data
变体
Mat (int rows, int cols, int type, void *data, size_t step=AUTO_STEP)
你可以跳过 memcpy 步骤
tensor = uint8_tensor //shape: (h, w, 3)
cv::Mat mat = cv::Mat(height, width, CV_8UC3, tensor.data_ptr());
return mat;
我正在使用火炬>=1.7.0。
对于dtype=float和大小[1, 3, height, width]的张量,这对我有用
cv::Mat torchTensortoCVMat(torch::Tensor& tensor)
{
tensor = tensor.squeeze().detach();
tensor = tensor.permute({1, 2, 0}).contiguous();
tensor = tensor.mul(255).clamp(0, 255).to(torch::kU8);
tensor = tensor.to(torch::kCPU);
int64_t height = tensor.size(0);
int64_t width = tensor.size(1);
cv::Mat mat = cv::Mat(cv::Size(width, height), CV_8UC3, tensor.data_ptr<uchar>());
return mat.clone();
}
我的张量形状是500x500x3
,我必须添加tensor.reshape({width * height * 3})
以获得实际图像
cv::Mat TensorToCVMat(torch::Tensor tensor)
{
// torch.squeeze(input, dim=None, *, out=None) → Tensor
// Returns a tensor with all the dimensions of input of size 1 removed.
// tensor.detach
// Returns a new Tensor, detached from the current graph.
// permute dimension, 3x700x700 => 700x700x3
tensor = tensor.detach().permute({1, 2, 0});
// float to 255 range
tensor = tensor.mul(255).clamp(0, 255).to(torch::kU8);
// GPU to CPU?, may not needed
tensor = tensor.to(torch::kCPU);
// shape of tensor
int64_t height = tensor.size(0);
int64_t width = tensor.size(1);
// Mat takes data form like {0,0,255,0,0,255,...} ({B,G,R,B,G,R,...})
// so we must reshape tensor, otherwise we get a 3x3 grid
tensor = tensor.reshape({width * height * 3});
// CV_8UC3 is an 8-bit unsigned integer matrix/image with 3 channels
cv::Mat imgbin(cv::Size(width, height), CV_8UC3, tensor.data_ptr());
return imgbin;
}
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.