C++ 实现tensor 转RGB图

C++ 实现torch::Tensor tensor 转RGB图，要求：tensor 的形状为{4，418，512}，取第1，2，3 通道的值分别赋予R,G,B三个不同通道 ，实现方法如下：

#include<opencv2/opencv.hpp>
#include <torch/torch.h>
//#少包自己补充

int shouRGBImg(torch::Tensor tensor) {
	// 分别提取四个通道
	torch::Tensor channel1 = tensor[1];
	torch::Tensor channel2 = tensor[2];
	torch::Tensor channel3 = tensor[3];
	torch::Tensor channel4 = tensor[0];

	// 将四个通道转换成RGB图
	cv::Mat rgbImage = cv::Mat::zeros(cv::Size(512, 704), CV_8UC3);
	for (int row = 0; row < 704; ++row) {
		for (int col = 0; col < 512; ++col) {
			rgbImage.at<cv::Vec3b>(row, col) = cv::Vec3b(channel2[row][col].item<float>() * 255,
				channel3[row][col].item<float>() * 255,
				channel1[row][col].item<float>() * 255);
		}
	}

	// 显示RGB图
	namedWindow("RGB Image", WINDOW_FREERATIO);
	cv::imshow("RGB Image", rgbImage);
	cv::waitKey(0);

	return 0;
}

补充：
如果tensor 的形状为{1,4，418，512},则用下方法替换循环：

rgbImage.at<cv::Vec3b>(row, col) = cv::Vec3b(channel2[0][row][col].item<float>() * 255,
				channel3[0][row][col].item<float>() * 255,
				channel1[0][row][col].item<float>() * 255);

测试：

int main（）{
	// 创建一个形状为{4,704,512}的随机Tensor
	torch::Tensor tensor1 = torch::rand({ 4, 704, 512 });
	shouRGBImg(tensor1);
}

图片: