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

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;}}}