分割一切模型 Fast SAM C++推理部署---onnxruntime（有核心代码）

news2024/11/27 10:39:15

Fast SAM C++推理部署—onnxruntime
核心源代码在结尾处
晓理紫

0 XX开局一张图，剩下…

在这里插入图片描述
本文记录只为日后更好学习

1 Fast SAM 简介

Fast SAM是仅使用SAM作者发布的SA-1B数据集的2%进行训练的CNN任意分割模型。FastSAM的性能与SAM方法相当，运行速度提高了50倍。
在这里插入图片描述

在这里插入图片描述

2 onnxruntime部署

2.1 环境与条件

需要配置onnxruntime相关环境

这个就网上自行解决

需要把原始权重模型转为onnx模型

2.2 onnx模型转换

本文参考进行转换，在转换过程中需要把dynamic_axes设置为None，采用静态维度进行，在核心源码中我们提供了640以及1024大小的onnx模型。相关核心代码如下：

torch.onnx.export(model, 
      img,  
      output_model_path,
      export_params=True,       
      opset_version=11,        
      do_constant_folding=True,  
      input_names = ['images'],  
      output_names = output_names,
      dynamic_axes=None)

注意：dynamic_axes=None这个最为关键，当然你也可以试试设置为dynamic

2.3 部署核心代码

模型转换完成以后，剩下的就是部署推理。部署推理里面最为重要也是最难搞的是数据解析部分。其中模型加载是很标准的流程，当然我这里不一定是标准的。

加载模型核心代码

  sessionOptions.SetExecutionMode(ORT_SEQUENTIAL);
  sessionOptions.EnableCpuMemArena();
  sessionOptions.SetIntraOpNumThreads(1);
  sessionOptions.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);
  ort_session = new Session(env, model_path.data(), sessionOptions);
  size_t numInputNodes = ort_session->GetInputCount();
  size_t numOutputNodes = ort_session->GetOutputCount();
  Ort::AllocatorWithDefaultOptions allocator;
  for (int i = 0; i < numInputNodes; i++) {
    this->input_names.emplace(ort_session->GetInputNameAllocated(i, allocator));
    Ort::TypeInfo input_type_info = ort_session->GetInputTypeInfo(i);
    auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
    auto input_dims = input_tensor_info.GetShape();
    input_node_dims.push_back(input_dims);
  }
  for (int i = 0; i < numOutputNodes; i++) {
    this->output_names.emplace(
        ort_session->GetOutputNameAllocated(i, allocator));
    Ort::TypeInfo output_type_info = ort_session->GetOutputTypeInfo(i);
    auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();
    auto output_dims = output_tensor_info.GetShape();
    output_node_dims.push_back(output_dims);
  }

  this->inpHeight = input_node_dims[0][2];
  this->inpWidth = input_node_dims[0][3];
  input_names__[0] = input_names->get();

模型加载以后，就可以送入数据进行推理

推理

  xiaoliziprocessImage(frame, tmp);
  this->xiaoliziNormalize_(tmp);
  array<int64_t, 4> input_shape_{1, 3, this->inpHeight, this->inpWidth};
  auto allocator_info =
      MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
  Value input_tensor_ = Value::CreateTensor<float>(
      allocator_info, input_image_.data(), input_image_.size(),
      input_shape_.data(), input_shape_.size());

  vector<Value> ort_outputs = ort_session->Run(
      RunOptions{nullptr}, input_names__, &input_tensor_, 1, output_names__,
      sizeof(output_names__) / sizeof(char *)); // 开始推理

推理以后就可以获取数据并进行解析

数据获取

  std::vector<cv::Mat> matVec;
  float *pdata = nullptr;
  pdata = ort_outputs[0].GetTensorMutableData<float>();
  if (pdata == nullptr) {
    std::cout << "pdata is nullptr" << std::endl;
    return;
  }
  cv::Mat matData(37, OUTPUT0w, CV_32F, pdata);
  matVec.push_back(matData);
  float *pdata5 = nullptr;
  pdata5 = ort_outputs[5].GetTensorMutableData<float>();
  if (pdata5 == nullptr) {
    std::cout << "pdata or pdata5 is nullptr" << std::endl;
    return;
  }
  cv::Mat matData5(32, OUTPUT1167w * OUTPUT1167w, CV_32F, pdata5);
  matVec.push_back(matData5);

数据解析

首先是对数据进行分割处理并进行NMS获取box、lab以及mask相关信息

cv::Mat box;
cv::Mat cls;
cv::Mat mask;
box = temData.colRange(0, 4).clone();
cls = temData.colRange(4, 5).clone();
mask = temData.colRange(5, temData.cols).clone();
cv::Mat j = cv::Mat::zeros(cls.size(), CV_32F);
cv::Mat dst;
cv::hconcat(box, cls, dst); // dst=[A  B]
cv::hconcat(dst, j, dst);
cv::hconcat(dst, mask, dst);
std::vector<float> scores;
std::vector<cv::Rect> boxes;
pxvec = dst.ptr<float>(0);
for (int i = 0; i < dst.rows; i++) {
  pxvec = dst.ptr<float>(i);
  boxes.push_back(cv::Rect(pxvec[0], pxvec[1], pxvec[2], pxvec[3]));
  scores.push_back(pxvec[4]);
}
std::vector<int> indices;
xiaoliziNMSBoxes(boxes, scores, conf_thres, iou_thres, indices);
cv::Mat reMat;
for (int i = 0; i < indices.size() && i < max_det; i++) {
  int index = indices[i];
  reMat.push_back(dst.rowRange(index, index + 1).clone());
}
box = reMat.colRange(0, 6).clone();
xiaolizixywh2xyxy(box);
mask = reMat.colRange(6, reMat.cols).clone();

其次是获取mask相关数据

  for (int i = 0; i < bboxes.rows; i++) {
    pxvec = bboxes.ptr<float>(i);
    cv::Mat dest, mask;
    cv::exp(-maskChannels[i], dest);
    dest = 1.0 / (1.0 + dest);
    dest = dest(roi);
    cv::resize(dest, mask, frmae.size(), cv::INTER_LINEAR);
    cv::Rect roi(pxvec[0], pxvec[1], pxvec[2] - pxvec[0], pxvec[3] - pxvec[1]);
    cv::Mat temmask = mask(roi);
    cv::Mat boxMask = cv::Mat(frmae.size(), mask.type(), cv::Scalar(0.0));
    float rx = std::max(pxvec[0], 0.0f);
    float ry = std::max(pxvec[1], 0.0f);
    for (int y = ry, my = 0; my < temmask.rows; y++, my++) {
      float *ptemmask = temmask.ptr<float>(my);
      float *pboxmask = boxMask.ptr<float>(y);
      for (int x = rx, mx = 0; mx < temmask.cols; x++, mx++) {
        pboxmask[x] = ptemmask[mx] > 0.5 ? 1.0 : 0.0;
      }
    }
    vremat.push_back(boxMask);
  }

最后是画出相关信息

cv::Mat bbox = vremat[0];
  float *pxvec = bbox.ptr<float>(0);
  for (int i = 0; i < bbox.rows; i++) {
    pxvec = bbox.ptr<float>(i);
    cv::rectangle(image, cv::Point(pxvec[0], pxvec[1]),
                  cv::Point(int(pxvec[2]), int(pxvec[3])),
                  cv::Scalar(0, 0, 255), 2);
  }

  for (int i = 1; i < vremat.size(); i++) {
    cv::Mat mask = vremat[i];
    int indx = (rand() % (80 - 0)) + 0;
    for (int y = 0; y < mask.rows; y++) {
      const float *mp = mask.ptr<float>(y);
      uchar *p = image.ptr<uchar>(y);
      for (int x = 0; x < mask.cols; x++) {
        if (mp[x] == 1.0) {
          p[0] = cv::saturate_cast<uchar>(p[0] * 0.5 + COLORS[indx][0] * 0.5);
          p[1] = cv::saturate_cast<uchar>(p[1] * 0.5 + COLORS[indx][1] * 0.5);
          p[2] = cv::saturate_cast<uchar>(p[2] * 0.5 + COLORS[indx][2] * 0.5);
        }
        p += 3;
      }
    }
  }