Fast SAM C++推理部署—onnxruntime
VX 搜索”晓理紫“ 关注并回复fastsamonnx获取核心代码
晓理紫
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;
}
}
}
3 核心代码
VX 搜索”晓理紫“ 关注并回复fastsamonnx获取核心代码
