本文共 2546 字，大约阅读时间需要 8 分钟。

Torch框架与OpenCV结合实现图像处理输出

在使用Torch框架进行图像处理任务时，通常会遇到如何将模型输出转换为OpenCV支持的矩阵并保存为图片文件的问题。本文将详细介绍实现过程，包括单通道、三通道输出的处理方法以及CUDA到CPU数据转移的技巧。

单通道输出处理

假设模型输出的Tensor为单通道图像，通道数为1。以下是实现过程的代码示例：

#include 
   
    #include 
    
     // 初始化模型torch::jit::script::Module net = torch::jit::load("../models/poolnet.pt");net.to(device);net.eval();// 读取并转换图片格式cv::Mat img = cv::imread(fileName);cv::cvtColor(img, img, cv::COLOR_BGR2RGB);// 创建Tensor并调整形状auto img_tensor = torch::from_blob(img.data, {1, img.rows, img.cols, 3},                                  torch::TensorOptions().dtype(torch::kByte)).to(device);img_tensor = img_tensor.permute({0, 3, 1, 2});img_tensor = img_tensor.toType(torch::kFloat);// 前向传播并获取输出torch::NoGradGuard no_grad;auto result = net.forward({img_tensor}).toTensor();// 数据转移与格式转换result = result.squeeze().sigmoid().mul(255.0).toType(torch::kByte).to(torch::kCPU);// 创建并保存结果图片cv::Mat img_C1;img_C1.create(cv::Size(img.cols, img.rows), CV_8UC1);memcpy(img_C1.data, result.data_ptr(), result.numel() * sizeof(torch::kByte));cv::imwrite(saveName, img_C1);
    
   

三通道输出处理

对于输出通道数为3的模型，处理方式几乎与单通道类似，但需要注意通道维度的调整：

#include 
   
    #include 
    
     // 初始化模型torch::jit::script::Module net = torch::jit::load("../models/poolnet2.pt");net.to(device);net.eval();// 读取并转换图片格式cv::Mat img = cv::imread(fileName);cv::cvtColor(img, img, cv::COLOR_BGR2RGB);// 创建Tensor并调整形状auto img_tensor = torch::from_blob(img.data, {1, img.rows, img.cols, 3},                                  torch::TensorOptions().dtype(torch::kByte)).to(device);img_tensor = img_tensor.permute({0, 3, 1, 2});img_tensor = img_tensor.toType(torch::kFloat);// 前向传播并获取输出torch::NoGradGuard no_grad;auto result = net.forward({img_tensor}).toTensor();// 数据转移与格式转换result = result.squeeze().sigmoid().mul(255.0).toType(torch::kByte).permute({1, 2, 0}).to(torch::kCPU);// 创建并保存结果图片cv::Mat img_C3;img_C3.create(cv::Size(img.cols, img.rows), CV_8UC3);memcpy(img_C3.data, result.data_ptr(), result.numel() * sizeof(torch::kByte));cv::imwrite(saveName, img_C3);
    
   

CUDA到CPU数据转移问题

在完成模型推理后，需要将GPU上的Tensor数据转移到CPU上，以便使用memcpy或cv::Mat进行操作。以下是两种常见转移方法：

// 方法一：直接转移result = result.to(torch::kCPU);// 方法二：使用cudaMemcpy（需要包含cuda_runtime.h）cudaMemcpy(pred.data, out.data_ptr(), out.numel(), cudaMemcpyDeviceToHost);

数据转移性能建议

性能优化

如果对转移时间要求较高，可以考虑以下优化方法：

通过以上方法，可以实现Torch框架输出的Tensor数据高效转换为OpenCV支持的矩阵，并成功保存为图片文件。

转载地址：http://goep.baihongyu.com/