读取onnx模型文件,或者经过量化后得到的xml文件(tips:xml和bin这两个文件需要在同一路径下,并且保持xml和bin的文件名称相同,读取的时候只需要读取xml即可)
头文件
#pragma once
#ifndef CPPDLL_EXPORTS
#define CPPDLL_API __declspec(dllexport)
#else
#define CPPDLL_API __declspec(dllimport)
#endif
#include <openvino/openvino.hpp>
#include <opencv2/opencv.hpp>
#include<thread>
#include<fstream>
#include <iostream>
#include <time.h>
#include <io.h>
#include <direct.h>
using namespace std;
using namespace ov;
using namespace cv;
extern "C" CPPDLL_API void Initial_model();
extern "C" CPPDLL_API void Detect_Segment(cv::Mat frame);
源文件
#include "yolov8_detect.h"
#include <crtdbg.h>
#ifdef _DEBUG
#define new new (_NORMAL_BLOCK, __FILE__,__LINE__)
#endif
//全局变量
static const std::string model_file = "best.onnx";
static const std::string cocoName_path = "cocoName.txt";
std::vector<std::string> class_names;
int class_nums = 0;
ov::InferRequest infer_request;
//为每一个分割结果定义成一个结构体
struct SegmentOutPut
{
int _id; //类别id
float _confidence; //类别置信度
cv::Rect2f _box; //检测框
cv::Mat _boxMask; //检测框内mask
};
//读取类别信息
std::vector<std::string> getClassName(std::string class_name_path)
{
std::ifstream inFile;
inFile.open(class_name_path);
if (!inFile)
{
std::cerr << "Unable to open file cocoName.txt";
exit(1);
}
std::string single_name;
std::vector <std::string> merge_name;
while (inFile >> single_name)
{
merge_name.push_back(single_name);
}
inFile.close();
return merge_name;
}
//交换图像通道
void convert(cv::Mat& input_image, cv::Mat& output_image, bool normalize, bool exchange_RB)
{
input_image.convertTo(output_image, CV_32F);
if (normalize)
{
output_image = output_image / 255.0;
}
if (exchange_RB)
{
cv::cvtColor(output_image, output_image, cv::COLOR_BGR2RGB);
}
}
cv::Rect toBox(const cv::Mat& input, const cv::Rect& range)
{
const float cx = input.at<float>(0);
const float cy = input.at<float>(1);
const float cw = input.at<float>(2);
const float ch = input.at<float>(3);
cv::Rect box;
box.x = cvRound(cx - 0.5f * cw);
box.y = cvRound(cy - 0.5f * ch);
box.width = cvRound(cw);
box.height = cvRound(ch);
return box & range;
}
void draw(cv::Mat& image, std::vector<SegmentOutPut>& results)
{
// 生成随机颜色
std::vector<cv::Scalar> colors;
std::srand(std::time(nullptr));
for (int i = 0; i < class_names.size(); ++i) {
const int b = std::rand() % 256;
const int g = std::rand() % 256;
const int a = std::rand() % 256;
colors.push_back(cv::Scalar(b, g, a));
}
cv::Mat mask = image.clone();
for (const SegmentOutPut& result : results) {
cv::rectangle(image, result._box, Scalar(0, 0, 0), 2, 4);
mask(result._box).setTo(colors[result._id], result._boxMask);
const std::string label = cv::format("%s:%.2f", class_names[result._id].c_str(), result._confidence);
int baseLine;
const cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
cv::putText(image, label,
cv::Point(result._box.x, std::max(result._box.y, float(labelSize.height))),
cv::FONT_HERSHEY_SIMPLEX, 1.5, Scalar(0, 0, 0), 3);
}
cv::addWeighted(image, 0.5, mask, 0.8, 1, image); // 把mask加在原图上面
}
CPPDLL_API void Initial_model()
{
//1. 创建Openvino runtime Core对象
ov::Core core;
std::shared_ptr<ov::Model> model = core.read_model(model_file);
//载入并编译模型
ov::CompiledModel compiled_model = core.compile_model(model, "AUTO");
//创建推理请求
infer_request = compiled_model.create_infer_request();
class_names = getClassName(cocoName_path);
class_nums = class_names.size();
}
CPPDLL_API void Detect_Segment(cv::Mat frame)
{
//Initial_model();
//获取模型的输入节点
ov::Tensor input_node = infer_request.get_input_tensor();
ov::Shape tensor_shape = input_node.get_shape();
///letterbox变换:不改变宽高比,将input_image缩放并放置到blob_image的左上角
const size_t channel_nums = tensor_shape[1];
const size_t height = tensor_shape[2];
const size_t width = tensor_shape[3];
//计算缩放因子
const float scale = std::min(height / float(frame.rows), width / float(frame.cols));
const cv::Matx23f matrix
{
scale, 0.0 ,0.0,
0.0, scale, 0.0,
};
cv::Mat blob_image;
//缩放并交换通道
cv::warpAffine(frame, blob_image, matrix, cv::Size(width, height));
convert(blob_image, blob_image, true, true);
//将输入图像的数据填入输入节点的指针中
float* const input_tensor_data = input_node.data<float>();
//往输入节点填入数据并将HWC->CHW
for (size_t c = 0; c < channel_nums; c++)
{
for (size_t h = 0; h < height; h++)
{
for (size_t w = 0; w < width; w++)
{
input_tensor_data[c * height * width + h * width + w] = blob_image.at<cv::Vec<float, 3>>(h, w)[c];
}
}
}
float scale_factor = 1 / scale;
//执行推理
auto time0 = std::chrono::system_clock::now();
infer_request.infer();
auto time1 = std::chrono::system_clock::now();
auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(time1 - time0).count();
cout << "推理耗时:" << duration << "ms" << endl;
/// 处理推理计算结果
// 获得推理结果, 输出结点是[116,8400], 一共8400个结果, 每个结果116个维度.
// 116=4+80+32, 4是预测框的[cx, cy, w, h], 80是每个类别的置信度, 32是分割需要用到的
const ov::Tensor output0 = infer_request.get_tensor("output0");
const float* output0_buffer = output0.data<const float>();
const ov::Shape output0_shape = output0.get_shape();
const int output0_rows = output0_shape[1];
const int output0_cols = output0_shape[2];
std::cout << "The shape of Detection tensor:" << output0_shape << std::endl;
const ov::Tensor output1 = infer_request.get_tensor("output1");
//const float* output1_buffer = output1.data<const float>();
const ov::Shape output1_shape = output1.get_shape();
std::cout << "The shape of Proto tensor:" << output1_shape << std::endl;
std::cout << std::endl << std::endl;
// Detect Matrix: 116 x 8400 -> 8400 x 116
// 一共8400个结果, 每个结果116个维度.
// 116=4+80+32, 4是预测框的[cx, cy, w, h]; 80是每个类别的置信度; 32需要与Proto Matrix相乘得到分割mask, 所以这里转置了矩阵
const cv::Mat detect_buffer = cv::Mat(output0_rows, output0_cols, CV_32F, (float*)output0_buffer).t();
// Proto Matrix: 1x32x160x160 -> 1x32x25600
const cv::Mat proto_buffer(output1_shape[1], output1_shape[2] * output1_shape[3], CV_32F, output1.data());
const float conf_threshold = 0.5;
const float nms_threshold = 0.5;
std::vector <cv::Rect> mask_boxes;
std::vector <cv::Rect> boxes;
std::vector <int> class_ids;
std::vector <float> confidences;
std::vector <cv::Mat> masks;
for (int i = 0; i < detect_buffer.rows;i++)
{
//获取8400 x 116的每一行的结果,116=4+80+32, 4是预测框的[cx, cy, w, h]; 80是每个类别的置信度; 32需要与Proto Matrix相乘得到分割mask
const cv::Mat result = detect_buffer.row(i);
//处理检测的部分的分类, 检测的类别是几,endcol就是4+几,这里我们的检测类别是5,因此endcol=9
const cv::Mat class_scores = result.colRange(4, 9);
cv::Point class_id_point;
double score;
cv::minMaxLoc(class_scores, nullptr, &score, nullptr, &class_id_point);
//置信度小的舍弃
if (score > conf_threshold)
{
class_ids.push_back(class_id_point.x);
confidences.push_back(score);
//预测框是在640x640的图片上检测的,而分割结果只有160x160,计算mask的缩放因子
const float mask_scale = 0.25f;
//处理检测部分的锚框部分,8400行每一行的前四个元素为[cx, cy, ch, cw]
const cv::Mat detection_box = result.colRange(0, 4);
const cv::Rect mask_box = toBox(detection_box * mask_scale, cv::Rect(0, 0, 160, 160));
const cv::Rect image_box = toBox(detection_box * scale_factor, cv::Rect(0, 0, frame.cols, frame.rows));
mask_boxes.push_back(mask_box);
boxes.push_back(image_box);
//处理分割部分结果
masks.push_back(result.colRange(9, 41));
}
}
//NMS
std::vector<int> nms_indexes;
cv::dnn::NMSBoxes(boxes, confidences, conf_threshold, nms_threshold, nms_indexes);
//处理nms之后的结果
std::vector<SegmentOutPut> segmentoutputs;
for (const int index : nms_indexes)
{
SegmentOutPut segmentoutput;
segmentoutput._id = class_ids[index];
segmentoutput._confidence = confidences[index];
segmentoutput._box = boxes[index];
cv::Mat m;
cv::exp(-masks[index] * proto_buffer, m);
m = 1.0f / (1.0f + m);
//1x25600 -> 160x160
m = m.reshape(1, 160);
cv::resize(m(mask_boxes[index]) > 0.5f, segmentoutput._boxMask, segmentoutput._box.size());
segmentoutputs.push_back(segmentoutput);
}
draw(frame, segmentoutputs);
//cv::putText(frame, cv::format("FPS: %.2f", 1.0 / t), cv::Point(20, 40), cv::FONT_HERSHEY_PLAIN, 2.0, cv::Scalar(255, 0, 0), 2, 8);
cv::namedWindow("检测", cv::WINDOW_NORMAL);
cv::imshow("检测", frame);
cv::waitKey(0);
cv::destroyAllWindows();
}
int main()
{
try
{
auto time00 = std::chrono::system_clock::now();
Initial_model();
auto time11 = std::chrono::system_clock::now();
auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(time11 - time00).count();
cout << "模型加载耗时:" << duration << "ms" << endl;
std::vector<cv::String> filenames;
cv::String folder = "E:\\xx\\xxx\\datasets\\coco128-seg\\images\\train2017/*.jpg";
cv::glob(folder, filenames);
while (true)
{
for (size_t ii = 0; ii < filenames.size(); ii++)
{
cv::Mat frame = cv::imread(filenames[ii],cv::IMREAD_COLOR);
cout << "filename = " << filenames[ii] << endl;
Detect_Segment(frame);
}
}
}
catch (const std::exception& e) {
std::cerr << "exception: " << e.what() << std::endl;
}
catch (...) {
std::cerr << "unknown exception" << std::endl;
}
return 0;
}