本文介绍如何将pytorch的Tensor类型转为opencv的Mat类型,废话不多说,直接上代码。

cv::Mat CvtTensor2CvImg(at::Tensor in_tensor)
{
	//sequeeze trans tensor shape from 1*C*H*W to C*H*W
	//permute C*H*W to H*W*C
	//in_tensor = in_tensor.squeeze().detach().permute({ 1, 2, 0 });
	in_tensor = in_tensor.squeeze().detach();
	c10::IntArrayRef tsize = in_tensor.sizes();
	//see tip3,tip4
	in_tensor = in_tensor.mul(255).clamp(0, 255).to(torch::kU8);
	in_tensor = in_tensor.to(torch::kCPU);
	cv::Mat resultImg(tsize[0], tsize[1], CV_8UC1);
	//copy the data from out_tensor to resultImg
	std::memcpy((void*)resultImg.data, in_tensor.data_ptr(), sizeof(torch::kU8) * in_tensor.numel());
	return resultImg;
}

注:以上代码基于libtorch和opencv两个库。

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐