文章目录
- 依赖库
- 导出onnx
- cpp推理代码
依赖库
opencv4.7及以上
导出onnx
yolo export model=./runs/segment/train3/weights/best.pt imgsz=640 format=onnx opset=12
cpp推理代码
#include<opencv.hpp>
#include<iostream>
#include<fstream>
#include<filesystem>
#include<random>/*
yolov8图像分割推理
*/std::vector<std::string> parse_classes_file(const char* name)
{std::vector<std::string> classes;std::ifstream file(name);if (!file.is_open()) {std::cerr << "Error: fail to open classes file: " << name << std::endl;return classes;}std::string line;while (std::getline(file, line)) {auto pos = line.find_first_of(" ");classes.emplace_back(line.substr(0, pos));}file.close();return classes;
}cv::Mat modify_image_size(const cv::Mat& img)
{auto max = std::max(img.rows, img.cols);cv::Mat ret = cv::Mat::zeros(max, max, CV_8UC3);img.copyTo(ret(cv::Rect(0, 0, img.cols, img.rows)));return ret;
}void get_masks(const cv::Mat& features, const cv::Mat& proto, const std::vector<int>& output1_sizes, const cv::Mat& frame, const cv::Rect box, cv::Mat& mk, const int input_size[2], const double mask_threshold)
{const cv::Size shape_src(frame.cols, frame.rows), // 原始图像大小shape_input(input_size[1], input_size[0]), // 640shape_mask(output1_sizes[3], output1_sizes[2]); //160std::cout << features.size << ", " << proto.size << std::endl;// 矩阵乘法cv::Mat res = (features * proto).t(); //1*25600 -> 25600*1res = res.reshape(1, { shape_mask.height, shape_mask.width });
#pragma region apply sigmoid to the maskcv::exp(-res, res);res = 1.0 / (1.0 + res);
#pragma endregioncv::resize(res, res, shape_input);float scalex = shape_src.width * 1.0 / shape_input.width;float scaley = shape_src.height * 1.0 / shape_input.height;cv::Mat tmp;if (scalex > scaley)cv::resize(res, tmp, cv::Size(shape_src.width, static_cast<int>(shape_input.height * scalex)));elsecv::resize(res, tmp, cv::Size(static_cast<int>(shape_input.width * scaley), shape_src.height));cv::Mat dst = tmp(cv::Rect(0, 0, shape_src.width, shape_src.height));mk = dst(box) > mask_threshold;//mk = dst > mask_threshold;
}void draw_boxes_mask(const std::vector<std::string>& classes, const std::vector<int>& ids, const std::vector<float>& confidences,const std::vector<cv::Rect>& boxes, const std::vector<cv::Mat>& masks, const std::string& name, cv::Mat& frame, const std::string& result_dir)
{std::cout << "image name: " << name << ", number of detections: " << ids.size() << std::endl;std::random_device rd;std::mt19937 gen(rd());std::uniform_int_distribution<int> dis(100, 255);cv::Mat mk = frame.clone();cv::Mat mask = cv::Mat::zeros(mk.size(), CV_8UC1);std::vector<cv::Scalar> colors;for (auto i = 0; i < classes.size(); ++i)colors.emplace_back(cv::Scalar(dis(gen), dis(gen), dis(gen)));for (auto i = 0; i < ids.size(); ++i) {#if 0 可以打开cv::rectangle(frame, boxes[i], colors[ids[i]], 2);std::string class_string = classes[ids[i]] + ' ' + std::to_string(confidences[i]).substr(0, 4);cv::Size text_size = cv::getTextSize(class_string, cv::FONT_HERSHEY_DUPLEX, 1, 2, 0);cv::Rect text_box(boxes[i].x, boxes[i].y - 40, text_size.width + 10, text_size.height + 20);cv::rectangle(frame, text_box, colors[ids[i]], cv::FILLED);cv::putText(frame, class_string, cv::Point(boxes[i].x + 5, boxes[i].y - 10), cv::FONT_HERSHEY_DUPLEX, 1, cv::Scalar(0, 0, 0), 2, 0);
#endifmk(boxes[i]).setTo(colors[ids[i]], masks[i]);mask(boxes[i]).setTo(cv::Scalar(255), masks[i]);}cv::addWeighted(frame, 0.5, mk, 0.5, 0, frame);// 显示maskcv::imshow("Inference", mask);cv::waitKey(-1);cv::imwrite(result_dir + "/" + name, frame);
}int main()
{namespace fs = std::filesystem;const char* onnx_file{ "D:/cppProj/YOLOV8InferTest/best.onnx" };constexpr int input_size[2]{ 640, 640 };constexpr float confidence_threshold{ 0.45 };constexpr float iou_threshold{ 0.50 }; // iou thresholdconstexpr float mask_threshold{ 0.50 }; // segment mask thresholdconstexpr bool cuda_enabled{ false };const char* classes_file{ "D:/cppProj/YOLOV8InferTest/classes.txt" };const std::string testPath("D:/pythonProj/torch_projs/test/catdog/000012.jpg");// 分离出文件名std::string filename;if (std::filesystem::exists(std::filesystem::path(testPath))==false){std::cout << testPath << " is not exists." << std::endl;return -1; }else{filename = std::filesystem::path(testPath).filename().string();std::cout << "filename = " << filename << std::endl;}const char* result_dir{ "./predictResult" };auto net = cv::dnn::readNetFromONNX(onnx_file);if (net.empty()) {std::cerr << "Error: there are no layers in the network: " << onnx_file << std::endl;return -1;}if (cuda_enabled) {net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);}else {net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);}if (!fs::exists(result_dir)) {fs::create_directories(result_dir);}auto classes = parse_classes_file(classes_file);if (classes.size() == 0) {std::cerr << "Error: fail to parse classes file: " << classes_file << std::endl;return -1;}std::cout << "classes: ";for (const auto& val : classes) {std::cout << val << " ";}std::cout << std::endl;#pragma region 预测cv::Mat frame = cv::imread(testPath, cv::IMREAD_COLOR);if (frame.empty()) {std::cerr << "Warning: unable to load image: " << testPath << std::endl;return -1;}auto tstart = std::chrono::high_resolution_clock::now();cv::Mat bgr = modify_image_size(frame);std::cout << "bgr.size=" << bgr.size << std::endl;cv::Mat blob;cv::dnn::blobFromImage(bgr, blob, 1.0 / 255.0, cv::Size(input_size[1], input_size[0]), cv::Scalar(), true, false);net.setInput(blob);std::vector<cv::Mat> outputs;net.forward(outputs, net.getUnconnectedOutLayersNames());if (outputs.size() != 2) {std::cerr << "Error: output must have 2 layers: " << outputs.size() << std::endl;return -1;}else{std::cout << "outputs.size() = " << outputs.size() << std::endl;}std::cout << "outputs[0].size=" << outputs[0].size << std::endl;// output0cv::Mat data0 = cv::Mat(outputs[0].size[1], outputs[0].size[2], CV_32FC1, outputs[0].data).t();std::cout << "data0.size = " << data0.size << std::endl;std::cout << outputs[0].size << std::endl;// output1std::vector<int> sizes;std::cout << outputs[1].size << std::endl;for (int i = 0; i < 4; ++i){sizes.emplace_back(outputs[1].size[i]);std::cout << sizes[i] << " ";}std::cout << std::endl;cv::Mat data1 = cv::Mat(sizes, CV_32F, outputs[1].data);std::cout << "data1.size = " << data1.size << std::endl;auto tend = std::chrono::high_resolution_clock::now();std::cout << "elapsed millisenconds: " << std::chrono::duration_cast<std::chrono::milliseconds>(tend - tstart).count() << " ms" << std::endl;
#pragma endregion#pragma region 后处理std::vector<int> class_ids;std::vector<float> confidences;std::vector<cv::Rect> boxes;std::vector<std::vector<float>> masks;float scalex = frame.cols * 1.f / input_size[1]; // note: image_preprocess functionfloat scaley = frame.rows * 1.f / input_size[0];auto scale = (scalex > scaley) ? scalex : scaley;const float* data = (float*)data0.data;for (auto i = 0; i < data0.rows; ++i){cv::Mat scores(1, classes.size(), CV_32FC1, (float*)data + 4);//std::cout << "scores.size=" << scores.size << std::endl;cv::Point class_id;double max_class_score;cv::minMaxLoc(scores, 0, &max_class_score, 0, &class_id);//std::cout << "classid.x=" << class_id.x << std::endl;if (max_class_score > confidence_threshold) {confidences.emplace_back(max_class_score);class_ids.emplace_back(class_id.x);// 大小为data[6]// 取出data[6,cols)的内容masks.emplace_back(std::vector<float>(data + 4 + classes.size(), data + data0.cols)); // 32 float x = data[0];float y = data[1];float w = data[2];float h = data[3];int left = std::max(0, std::min(int((x - 0.5 * w) * scale), frame.cols));int top = std::max(0, std::min(int((y - 0.5 * h) * scale), frame.rows));int width = std::max(0, std::min(int(w * scale), frame.cols - left));int height = std::max(0, std::min(int(h * scale), frame.rows - top));boxes.emplace_back(cv::Rect(left, top, width, height));}data += data0.cols;}std::cout << "masks.size=" << masks.size() << std::endl;std::vector<int> nms_result;// 返回得分最高的box的索引cv::dnn::NMSBoxes(boxes, confidences, confidence_threshold, iou_threshold, nms_result);// 0表示元素总个数不变cv::Mat proto = data1.reshape(0, { sizes[1], sizes[2] * sizes[3] });std::vector<int> ids;std::vector<float> confs;std::vector<cv::Rect> rects;std::vector<cv::Mat> mks;for (size_t i = 0; i < nms_result.size(); ++i) {auto index = nms_result[i];std::cout << "index=" << index << std::endl;ids.emplace_back(class_ids[index]);confs.emplace_back(confidences[index]);boxes[index] = boxes[index] & cv::Rect(0, 0, frame.cols, frame.rows);// 原始大小std::cout << "cols=" << frame.cols << ", rows=" << frame.rows << std::endl;std::cout << "cv::Mat(masks[index]).size=" << cv::Mat(masks[index]).size << std::endl;cv::Mat mk;get_masks(cv::Mat(masks[index]).t(), proto, sizes, frame, boxes[index], mk, input_size, mask_threshold);mks.emplace_back(mk);rects.emplace_back(boxes[index]);/*std::cout << "ids=" << ids[i] << std::endl;cv::imshow("w", mk);cv::waitKey();*/}draw_boxes_mask(classes, ids, confs, rects, mks, filename, frame, result_dir);#pragma endregionreturn 0;
}
