【分享】 CV中关于pytorch在C++和python里的一些数据处理

简单介绍

就我个人感觉而言pytorch比tensorflow简单很多，如果不了解pytorch的可以看看 知乎上的这个问题简单了解一下。由于数据处理暂时不涉及网络的搭建、模型的训练等，所以这里不讨论这些。我们使用的模型就是一个网络，网络的直接输入和输出都是tensor（计算由forward函数实现）。在我们的应用里输入和输出都是一帧帧的图像，而在C++里我们常用opencv里的mat来存储图片信息，所以重点就在于mat和tensor的相互转换。

pytorch模型转换为torch script

我们采用的是trace的方法。 这种方法操作比较简单，只需要给模型一组输入，走一遍推理网络，然后由torch.ji.trace记录一下路径上的信息并保存即可。 缺点是如果模型中存在控制流比如if-else语句，一组输入只能遍历一个分支，这种情况下就没办法完整的把模型信息记录下来。目前使用下来还没有遇到什么问题。使用以下代码即可生成.pt后缀的torch script。

import torch
import torchvision

# An instance of your model.
model = torchvision.models.resnet18()  # 在这里加载自己的模型

# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)

生成以后需要在python环境用jit加载torch script来验证一下是否正常。

        myModel = torch.load(args.model) # before
        myModel = torch.jit.load(args.model) # after

输入处理： mat->tensor

python版本：

        image = cv2.imread(os.path.join(args.inputPath, f))

        origin_h, origin_w, c = image.shape  # 取图像大小

        image_resize = cv2.resize(image, (args.size,args.size), interpolation=cv2.INTER_CUBIC)  # 将图片大小resize为指定大小
  
        image_resize = (image_resize ) / 255.0 # 像素值归一化

        tensor_4D = torch.FloatTensor(1, 3, args.size, args.size) # 创建一个指定大小的float类型的tensor

        tensor_4D[0,:,:,:] = torch.FloatTensor(image_resize.transpose(2,0,1)) # 将mat里的数据传递到tensor里
                                                                                                                     # （这里需要注意需要将mat数据翻转一下以适应输入）
        inputs = tensor_4D.to(device) # 将数据放到指定存储区域（cpu or gpu）

cpp版本：

  cv::Mat image_resized;

  cv::resize(image.clone(), image_resized, cv::Size(256, 256), cv::INTER_CUBIC); // 将图片大小resize为指定大小

  cv::Mat image_resized_float;

  image_resized.convertTo(image_resized_float, CV_32F, 1.0 / 255); // 像素值归一化

  auto img_tensor = torch::from_blob(    // 将mat里的数据传递到tensor里
      image_resized_float.data,
      {1, image_resized_float.rows, image_resized_float.cols, 3},
      torch::kFloat32);

  img_tensor = img_tensor.permute({0, 3, 1, 2}); // cpp版本是在转化为mat以后再翻转

  img_tensor = img_tensor.toType(torch::kFloat).to(device); // 将数据放到指定存储区域（cpu or gpu）

输出处理：tensor -> mat

python版本：

seg, alpha = net(inputs)  # 将输入放入网络得到输出
if args.without_gpu:    # 注意必须将网络输出数据放到cpu里后续才能使用
     alpha_np = alpha[0,0,:,:].data.numpy()  
else:
     alpha_np = alpha[0,0,:,:].cpu().data.numpy()

alpha_np = cv2.resize(alpha_np, (origin_w, origin_h), interpolation=cv2.INTER_CUBIC) # 将图片大小resize为原本大小

seg_fg = np.multiply(alpha_np[..., np.newaxis], image)  # 将输出得到的mat数据和原图片的mat数据元素相乘
                                                                                        # 注意这里输出的是（1，256，256）的代表判断结果的mat（单通道），而原图片是
                                                                                        # （ 3， 256， 256）的mat（多通道）
f = f[:-4] + '_xxx.png'
cv2.imwrite(os.path.join(args.savePath, f), seg_fg)    #存储结果

cpp版本：

  auto out_tensor2 = module->forward({img_tensor}); // 将输入放入网络得到输出
  auto out_tuple = out_tensor2.toTuple();
  auto ele = out_tuple.get()->elements();
  auto out = ele[1].toTensor();   // 由于我这里输出不是tensor，而是一个tensor的tuple，所以需要这么处理一下

  out = out.to(torch::kCPU); # 注意必须将网络输出数据放到cpu里后续才能使用

  cv::Mat resultImg(cv::Size(256, 256), CV_32F, out.data_ptr());  // 将tensor数据放入mat中
  cv::resize(resultImg, resultImg, cv::Size(image.cols, image.rows),  # 将图片大小resize为原本大小
             cv::INTER_CUBIC);

  multiply(image, resultImg); # 将输出得到的mat数据和原图片的mat数据元素相乘

 // 这里的multiply是我自己写的一个函数，具体实现如下，思路就是暴力遍历相乘（有待改进）

void multiply(cv::Mat &img1, cv::Mat &img2) {
  double t0 = (double)cv::getTickCount();
  int height = img1.rows;
  int width = img1.cols;
  int sum = 0;
  const uchar* uc_pixel = img1.data;
  const uchar* uc_pixel2 = img2.data;
  for (int row = 0; row < height; row++) {
    uc_pixel = img1.data + row * img1.step;
    for (int col = 0; col < width; col++) {
      int a = uc_pixel[0];
      int b = uc_pixel[1];
      int c = uc_pixel[2];
      float grey = img2.at<float>(row, col);
      img1.at<cv::Vec3b>(row, col)[0] = a * grey;
      img1.at<cv::Vec3b>(row, col)[1] = b * grey;
      img1.at<cv::Vec3b>(row, col)[2] = c * grey;
      uc_pixel += 3;
    }
  }
}

最后

当然上面只是两种环境中数据处理的很少一部分，如果有什么错误请各位指出。

转载请注明本文url，引用请注明出处，谢谢支持！！！