部署推理笔记c++ 图像分割代码

写在前面：这个自己写的一个案例代码，仅作自己参考！

#include <iostream>
#include <string>
#include <vector>
#include <opencv2/opencv.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/dnn/dnn.hpp>
#include <opencv2/dnn/all_layers.hpp>
#include <onnxruntime_cxx_api.h>
#include <opencv2/core/utils/logger.hpp>
#include <cmath> // 包含 round 函数
// #define LOG true // 是否打印日志
#define LOG false // 是否打印日志
using namespace std;

const size_t HOLE_NUM_W = 4;            // 穴盘列数
const size_t HOLE_NUM_H = 3;            // 穴盘行数
const float MATCH_THERSHOLD = 0.8f;     // 孔,苗匹配 相交阈值
const float AREA_THERSHOLD = 0.1f;      // 判断弱苗阈值， 弱苗占穴孔的比例
const float SCORE_THERSHOLD = 0.6f;     // 预测的置信度阈值
const float IOU = 0.5f ;                // NMS IOU 阈值

#define COLOR_WEAK cv::Scalar(255, 0, 0)   // 弱苗的颜色
#define COLOR_EXITS cv::Scalar(0, 0, 255)  // 正常苗的颜色

void showImage(cv::Mat &img);

// 获取掩码和掩码面积
std::vector<cv::Mat> get_hole_mask(int w, int h);

cv::Mat vectorToMat(const std::vector<float>& vec, int rows, int cols);

void drawBox(cv::Mat &img, std::vector<float> &b_xyxyc);

void rmoveElement(std::vector<int>& m_vector, const int & x);

struct hoel2mask
{
    int hole;
    int mask;
};


/************************定义类 */
class model {
public:
    model(const std::string& onnxpath); //构造函数
    void preProcess(cv::Mat& img);
    std::vector<std::vector<int>> infer(cv::Mat img);
    void postProcess(std::vector<std::vector<int>>& res);
    void forward();
    void save_img(const std::string save_path);
    void hole_target_match(const std::vector<cv::Mat>& hole_masks, const float& hole_area, const std::vector<cv::Mat>& res_masks,
                std::vector<std::vector<int>>& res);

public:
    Ort::Env env; //  创建一个环境
    std::unique_ptr<Ort::Session> session; // 创建一个会话
    std::array<int64_t, 4> input_shape_info; // 输入尺度
    size_t input_size; // 输入大小
    // 创建 input_names 和 output_names， 输入输出名称
    std::vector<const char*> input_names;
    std::vector<const char*> output_names;
    std::vector<std::string> output_node_names;
    std::vector<std::string> input_node_names;
    int input_w, input_h, output_h, output_w, img_w, img_h; // 输入输出原始尺寸
    cv::Mat img_processed; // 输入图片
    cv::Mat mFrame; // 原始图片
    float m_x_ratio, m_y_ratio; // 缩放比例
    size_t mask_size;
    std::vector<float> outDataVector;
    std::vector<std::vector<int>> resVector;  // 最后的结果
};


void test(){
    // 1 创建一个模型
    model model("D:\\codes\\deploy_codes\\deployCmake\\TidyPlant\\weights\\20250416-s-seg-limiao.onnx"); 
    // 2 图片文件路径
    std::string filename = "D:\\codes\\deploy_codes\\deployCmake\\TidyPlant\\datas\\6.png"; 
    // 3. 读取文件
    cv::Mat img = cv::imread(filename);

    // 4. 推理
    std::vector<std::vector<int>> res;
    res = model.infer(img);
    
    // 5. 打印输出
    for (int i=0; i<res.size(); i++){
        for (int j =0; j<res[i].size(); j++){
            printf("%d ", res[i][j]);
        }
        printf("\n");
    }
    model.save_img("./res.jpg");  // 保存 结果图片
    showImage(model.mFrame); // 显示图像
    std::cout << "LOG test end !" << std::endl;
}


int main(){
    cv::utils::logging::setLogLevel(cv::utils::logging::LOG_LEVEL_ERROR); // 屏蔽opencv警告
    test();
    return 0;
}   




model::model(const std::string& onnxpath) {
    std::wstring modelPath = std::wstring(onnxpath.begin(), onnxpath.end());
    
    Ort::SessionOptions session_options;  // 创建一个会话选项
    env = Ort::Env(ORT_LOGGING_LEVEL_ERROR, "inference"); // 创建一个环境
    session_options.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);  // 设置优化级别/
    session = std::make_unique<Ort::Session>(env, modelPath.c_str(), session_options); // 创建一个会话

    size_t numInputNodes = session->GetInputCount(); // 获取输入节点数量
    size_t numOutputNodes = session->GetOutputCount(); // 获取输出节点数量
    Ort::AllocatorWithDefaultOptions allocator; // 创建一个分配器

    // 获取输入节点名称和类型信息， 主要用来获取输入节点的名称
    const char* name_cstr;
    Ort::TypeInfo input_type_info = session->GetInputTypeInfo(0);
    auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
    auto input_dims = input_tensor_info.GetShape();
    input_w = input_dims[2];
    input_h = input_dims[3];
    for (int i = 0; i < numInputNodes; i++) {
        auto input_name = session->GetInputNameAllocated(i, allocator); // 第 i 个输入节点的名称，并由 allocator 管理其内存
        name_cstr = input_name.get();   
        input_node_names.push_back(name_cstr);
        if (LOG) std::cout << "input-name_cstr:" << name_cstr<< std::endl;
    }
    
    // 获取输出节点名称和类型信息
    Ort::TypeInfo output_type_info = session->GetOutputTypeInfo(0);
    auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();
    auto output_dims = output_tensor_info.GetShape();
    output_h = output_dims[1];  // 输出维度8400
    output_w = output_dims[2];  // 25605
    for (int i = 0; i < numOutputNodes; i++) {
        auto out_name = session->GetOutputNameAllocated(i, allocator);
        name_cstr = out_name.get();   
        output_node_names.push_back(name_cstr);
        if (LOG) std::cout << "out-name_cstr:" << name_cstr<< std::endl;
    }


    input_names.resize(input_node_names.size());
    output_names.resize(output_node_names.size());
    for (size_t i = 0; i < input_node_names.size(); ++i) {input_names[i] = input_node_names[i].c_str();}
    // 填充 output_names
    for (size_t i = 0; i < output_node_names.size(); ++i) {output_names[i] = output_node_names[i].c_str();}

    // 输入信息构建
    input_size = input_h * input_w * 3;
    input_shape_info  = { 1, 3, input_h, input_w };
    mask_size = input_w /4; // mask 输出的大小
    if (LOG) std::cout << "LOG input: " << input_w << "*" << input_h <<  std::endl;
    if (LOG) std::cout << "LOG out  : " << output_h << "*" << output_w <<  std::endl;
    if (LOG) std::cout << "LOG model init end !" << std::endl;
};

void model::preProcess(cv::Mat& img){
    mFrame = img;
    img_w = img.cols;
    img_h = img.rows;
    m_x_ratio = static_cast<float>(img.cols) / input_w; // rate
    m_y_ratio = static_cast<float>(img.rows) / input_h;
    img_processed = cv::dnn::blobFromImage(img, 1.0 / 255.0, cv::Size(640, 640), cv::Scalar(), true, false, CV_32F);
    if (LOG) std::cout << "LOG preProcess end!" << std::endl;
}

void model::forward(){
    auto allocator_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(allocator_info, img_processed.ptr<float>(), input_size, input_shape_info.data(), input_shape_info.size());
    std::vector<Ort::Value> ort_outputs;
    if (!session) {
        // 确保session_已经被初始化
        throw std::runtime_error("LOG Session is not initialized. Call Initialize first.");
    }

    try {
        ort_outputs = session->Run(Ort::RunOptions{ nullptr }, 
            input_names.data(), 
            &input_tensor, 
            input_names.size(),
            output_names.data(), 
            output_names.size());
    }
    catch (std::exception e) {
        std::cout << e.what() << std::endl;
    }


    const  float* pdata = ort_outputs[0].GetTensorMutableData<float>();
    std::vector<int64_t> outputShape =ort_outputs[0].GetTensorTypeAndShapeInfo().GetShape();
	size_t count = ort_outputs[0].GetTensorTypeAndShapeInfo().GetElementCount();
    outDataVector.assign(pdata, pdata + count);  // 输出赋值


    if (LOG) std::cout << "LOG end forward!" << std::endl;
};

void model::postProcess(std::vector<std::vector<int>>& res){
    // 1. 获取输出数据
    size_t out_num = outDataVector.size() / 25605; // 8400
    // 2. 解析结果
    std::vector<float> scores; // 置信度分数
    std::vector<cv::Rect> bboxes;
    std::vector<std::vector<float>> masks;
    
    for (int i=0; i<out_num; i++){
        std::vector<float> line(25605); // 预分配空间
        size_t id = 25605 * i;
        std::copy(outDataVector.begin() + id, outDataVector.begin() + id + output_w, line.begin());
        // 1. 分数
        scores.push_back(line[4]); // 预测分数
        // 2. rectangle
        cv::Rect bbox(line[0], line[1], line[2], line[3]);
        bboxes.push_back(bbox);
        // 3. mask
        std::vector<float> mask_(line.begin() + 5, line.end());
        masks.push_back(mask_);
    }
    

    // 3. NMS
    std::vector<int> indices;  // 存储保留的边界框
    cv::dnn::NMSBoxes(bboxes, scores, SCORE_THERSHOLD, IOU, indices);

    // 4. 解析mask， 还原到原始图像
    std::vector<cv::Mat> res_masks;
    std::vector<std::vector<float>> res_bboxes;
    std::vector<cv::Mat> hole_masks = get_hole_mask(img_w, img_h); // 穴孔区域
    float hole_area = (float)cv::countNonZero(hole_masks[0]); // 穴孔面积
    for (int idx : indices){
        // 1. 处理box
        cv::Rect box =  bboxes[idx];
        float x1,x2,y1,y2;
        float xc = (float)box.x;
        float yc = (float)box.y;
        float w = (float)box.width;
        float h = (float)box.height;

        x1 = (xc - w / 2) / input_w * img_w;
        y1 = (yc - h / 2) / input_h * img_h;
        x2 = (xc + w / 2) / input_w * img_w;
        y2 = (yc + h / 2) / input_h * img_h;
        std::vector<float> box_ = {x1,y1,x2,y2};
        res_bboxes.push_back(box_);
        
        // 2. 处理mask
        std::vector<float> mask_ = masks[idx];
        int mask_x1 = static_cast<int>(std::round(x1 / img_w * mask_size)); // 在mask上对应的坐标
        int mask_y1 = static_cast<int>(std::round(y1 / img_h * mask_size));
        int mask_x2 = static_cast<int>(std::round(x2 / img_w * mask_size));
        int mask_y2 = static_cast<int>(std::round(y2 / img_h * mask_size));
        cv::Mat mask(mask_size, mask_size, CV_32F, mask_.data()); // 160*160
        // 剪切缩放mask大小
        cv::Mat cropped_mask = mask(cv::Range(mask_y1, mask_y2), cv::Range(mask_x1, mask_x2));  // 根据160*160剪切目标mask
        int target_width = static_cast<int>(std::round(x2 - x1)); // 最终目标结果尺寸
        int target_height = static_cast<int>(std::round(y2 - y1));
        cv::resize(cropped_mask, cropped_mask, cv::Size(target_width, target_height), 0, 0, cv::INTER_LINEAR); // 还原实际图像大小比例
        // 构建mask 0-1 矩阵
        cv::Mat binary_mask;
        cv::threshold(cropped_mask, binary_mask, 0.5f, 255, cv::THRESH_BINARY);
        binary_mask.convertTo(binary_mask, CV_8U);
        
        
        // 粘贴mask到尺寸
        cv::Mat mask_1 = cv::Mat::zeros(mFrame.size(), CV_8U); // 初始化为黑色（值为 0）  //初始化未0
        int round_x1 = static_cast<int>(std::round(x1));
        int round_y1 = static_cast<int>(std::round(y1));
        int round_x2 = static_cast<int>(std::round(x2));
        int round_y2 = static_cast<int>(std::round(y2));
        round_x1 = std::max(0, std::min(round_x1, mask_1.cols - 1));
        round_y1 = std::max(0, std::min(round_y1, mask_1.rows - 1));
        round_x2 = std::max(0, std::min(round_x2, mask_1.cols));
        round_y2 = std::max(0, std::min(round_y2, mask_1.rows));
        target_width = round_x2 - round_x1;
        target_height = round_y2 - round_y1;
        cv::Mat binary_mask_resized;
        cv::resize(binary_mask, binary_mask_resized, cv::Size(target_width, target_height), 0, 0, cv::INTER_LINEAR);
        cv::Rect roi(round_x1, round_y1, target_width, target_height);
        binary_mask_resized.copyTo(mask_1(roi)); // 目标区域粘贴到原图
        res_masks.push_back(mask_1);  // 保存当前


        // 3. 画图mask
        float target_area = (float)cv::countNonZero(mask_1); // 预测面积
        // std::cout << " hole area " << target_area << endl;
        cv::Scalar color = (target_area/hole_area) < AREA_THERSHOLD ? COLOR_WEAK : COLOR_EXITS;
        std::vector<std::vector<cv::Point>> contours;
        cv::Mat hierarchy;
        cv::findContours(mask_1, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE); // 查找轮廓
        cv::Mat mFrame_ = cv::Mat::zeros(mFrame.size(), CV_8UC3); // 创建 0 Mat
        cv::fillPoly(mFrame_, contours, color);
        cv::addWeighted(mFrame_, 1, mFrame, 1, 0, mFrame);
    }

    // 5. 匹配
    hole_target_match(hole_masks, hole_area, res_masks, res);

    if (LOG) std::cout << "LOG postProcess end!" << std::endl;

}

void model::hole_target_match(const std::vector<cv::Mat>& hole_masks, const float& hole_area, 
    const std::vector<cv::Mat>& targets_masks, std::vector<std::vector<int>>& res){


    // 用来保存已经匹配的穴孔和目标
    std::vector<int> remain_targets_idx;
    std::vector<int> remain_hole_idx;
    for (int i =0; i<targets_masks.size(); i++){
        remain_targets_idx.push_back(i);
    }
    for (int i =0; i<hole_masks.size(); i++){
        remain_hole_idx.push_back(i);
    }
    std::vector<hoel2mask> h2m;


    // 1. 一阶段匹配算法
    for (size_t i = 0; i < hole_masks.size(); ++i) {
        const cv::Mat& hole_mask = hole_masks[i]; // 穴孔区域


        std::vector<float> score_list_temp;
        std::vector<float> area_match_list_temp;

        // 计算每一个目标与穴孔的得分
        for (auto j : remain_targets_idx) {
            const cv::Mat& target_mask = targets_masks[j]; // 预测mask 区域
            cv::Mat intersection;
            cv::bitwise_and(hole_mask, target_mask, intersection); // 计算交集
            float intersection_area = (float)cv::countNonZero(intersection);
            float target_area = (float)cv::countNonZero(target_mask);
            float score = static_cast<float>(intersection_area / target_area);
            area_match_list_temp.push_back(target_area/hole_area); // 苗面积占穴孔面积
            score_list_temp.push_back(score);
        }

       if (LOG){
            cout << "score_list_temp :" << i << endl;
            for (auto x : score_list_temp){
                cout << x << " ";
            };
            printf("\n");

            
       }
        
        // 判断是否有效，删除匹配下标
        if (!score_list_temp.empty()) {
            auto max_it = std::max_element(score_list_temp.begin(), score_list_temp.end()); // 找到最大值
            int max_index = std::distance(score_list_temp.begin(), max_it);    // 计算最大值的下标
            if (score_list_temp[max_index] > MATCH_THERSHOLD && area_match_list_temp[max_index] > AREA_THERSHOLD) {
                hoel2mask h2m_ = {i, remain_targets_idx[max_index]};
                h2m.push_back(h2m_);
                rmoveElement(remain_hole_idx, i); // 删除匹配下标hole
                rmoveElement(remain_targets_idx, remain_targets_idx[max_index]); // 删除匹配下标mask
            }
        }
    }

    if (LOG){
        cout << "LOG one stage match:";
        for (int i =0; i<h2m.size(); i++){
            cout << h2m[i].hole << ":" << h2m[i].mask<< " ";
        };
        printf("\n");
        for (auto x : remain_targets_idx){
            cout << "remain_targets_idx: " << x <<" ";
        };
        printf("\n");

        for (auto x : remain_hole_idx){
            cout << "remain_hole_idx: " << x <<" ";
        };
        printf("\n");
 
    }

    // 2, 二阶段匹配
    if (! remain_targets_idx.empty() && !remain_hole_idx.empty()){
        for (int i =0; i<remain_targets_idx.size(); i++){
            int target_id = remain_targets_idx[i];
           
            const cv::Mat& target_mask = targets_masks[target_id];
            
            float target_area = (float)cv::countNonZero(target_mask);
            float area_score = target_area / hole_area;
            if (area_score < AREA_THERSHOLD){
                continue;
            }
            
            std::vector<float> score_list_temp;
            std::vector<float> area_match_list_temp;
            for (auto hole_id : remain_hole_idx){
                const cv::Mat& hole_mask = hole_masks[hole_id]; // hoel 区域
                cv::Mat intersection;
                cv::bitwise_and(hole_mask, target_mask, intersection); // 计算交集
                float intersection_area = (float)cv::countNonZero(intersection);
                float score = static_cast<float>(intersection_area / target_area);
                score_list_temp.push_back(score);
    
            }
    
            if (LOG){
                cout << "score_list_temp :" << target_id << endl;
                for (auto x : score_list_temp){
                    cout << x << " ";
                };
                printf("\n");
                
            }
    
            // 判断是否有效，删除匹配下标
            if (score_list_temp.empty()) {continue;}
            auto max_it = std::max_element(score_list_temp.begin(), score_list_temp.end()); // 找到最大值
            int max_index = std::distance(score_list_temp.begin(), max_it);    // 计算最大值的下标
            if (score_list_temp[max_index]>0) {
                hoel2mask h2m_ = {remain_hole_idx[max_index], target_id};
                h2m.push_back(h2m_);
                rmoveElement(remain_hole_idx, remain_hole_idx[max_index]); // 删除匹配下标hole
                // rmoveElement(remain_targets_idx, target_id); // 删除匹配下标mask
            }
    
        }
        
        
        if (LOG){
            cout << "LOG two stage match:";
            for (int i =0; i<h2m.size(); i++){
                cout << h2m[i].hole << ":" << h2m[i].mask<< " ";
            }
            cout << endl;
        }
    }


    
    // 3. 构建输出
    for (int i=0; i<HOLE_NUM_H; i++){
        std::vector<int> line;
        for (int j =0; j<HOLE_NUM_W; j++){
            line.push_back(0);
        }
        res.push_back(line);
    }
    // 赋值
    for (auto e : h2m){
        int id = e.hole;
        int x = id / HOLE_NUM_W;
        int y = id % HOLE_NUM_W;
        res[x][y] = 1;
    }

    if(LOG){
        for (int i=0; i<res.size(); i++){
            for (int j =0; j<res[i].size(); j++){
                printf("%d ", res[i][j]);
            }
        }
    }
    
   
    
};

std::vector<std::vector<int>> model::infer(cv::Mat img){
    // 1. 数据预处理，获得输入数据，std::vector<float>
    preProcess(img); // 前处理
    // 2. 前向推理，得到结果std::vector<float>
    forward();
    // 3. 后处理
    std::vector<std::vector<int>> res;
    postProcess(res);
    if (LOG) std::cout << "LOG infer end!" << std::endl;
    return res; // 返回
};

void model::save_img(const std::string save_path){
    cv::imwrite(save_path, mFrame);
}


std::vector<cv::Mat> get_hole_mask(int w, int h) {
    // 计算每个单元格的宽度和高度
    int cell_width = w / HOLE_NUM_W;
    int cell_height = h / HOLE_NUM_H;
    // 存储掩码的容器
    std::vector<cv::Mat> hole_masks;
    // 遍历每个单元格
    for (int i = 0; i < HOLE_NUM_H; ++i) { // 纵向
        int y1 = static_cast<int>(std::round(i * cell_height));
        int y2 = static_cast<int>(std::round((i + 1) * cell_height));
        for (int j = 0; j < HOLE_NUM_W; ++j) { // 横向
            int x1 = static_cast<int>(std::round(j * cell_width));
            int x2 = static_cast<int>(std::round((j + 1) * cell_width));
            // 创建掩码矩阵
            cv::Mat mask = cv::Mat::zeros(h, w, CV_8U); // 单通道，8位无符号整数
            mask(cv::Range(y1, y2), cv::Range(x1, x2)).setTo(cv::Scalar(1)); // 设置区域为1
            // 添加到容器
            hole_masks.push_back(mask.clone());
        }
    }
    return hole_masks;
}


void drawBox(cv::Mat &img, std::vector<float> &b_xyxyc){

    // 获取边界框的坐标
    int x1 = static_cast<int>(b_xyxyc[0]);
    int y1 = static_cast<int>(b_xyxyc[1]);
    int x2 = static_cast<int>(b_xyxyc[2]);
    int y2 = static_cast<int>(b_xyxyc[3]);
    cv::rectangle(img, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(0, 255, 0), 2);
};


void rmoveElement(std::vector<int>& m_vector, const int & x){
    auto it = std::find(m_vector.begin(), m_vector.end(), x);
    if (it != m_vector.end()) {
        m_vector.erase(it);
    }
}

void showImage(cv::Mat &img){
    cv::imshow("show", img);
    cv::waitKey(0);
};

cmake  编译

#************  设置CMake最低版本要求   ************
cmake_minimum_required(VERSION 3.10)


#************ 1. 必要的设置 定义项目名称和版本号等  ************
project(Project VERSION 1.0)
set(CMAKE_CXX_STANDARD 17) # C++17

set(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/bin)  # 1可执行文件输出路径
# set(RUNTIME_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}/bin)  # 2可执行文件输出路径
# set(RUNTIME_OUTPUT_DIRECTORY_DEBUG ${PROJECT_SOURCE_DIR}/bin/debug) # 默认
# set(RUNTIME_OUTPUT_DIRECTORY_RELEASE ${PROJECT_SOURCE_DIR}/bin/release)

#************ 2. 代码文件检索 ********
file(GLOB SRC ${PROJECT_SOURCE_DIR}/src/*.cpp)
# file(GLOB_RECURES SRC ${PROJECT_SOURCE_DIR}/src/*.cpp)


#************  3. 头文件、库等设置  ************
include_directories(${PROJECT_SOURCE_DIR}/include) # 自定义头文件

set(OpenCV_DIR "C:/c++/opencv4.5.5/opencv/build")  # OpenCV
find_package(OpenCV REQUIRED)
include_directories(${OpenCV_INCLUDE_DIRS})

include_directories("C:/c++/onnxruntime-win-x64-1.13.1/include")
link_directories("C:/c++/onnxruntime-win-x64-1.13.1/lib")
link_libraries()



# #************ 4. 添加可执行文件  ************
add_executable(exec ${SRC})
#************ 5. 连接  ************
target_link_libraries(exec ${OpenCV_LIBS} onnxruntime)



# cmake .. -G "MinGW Makefiles"
# cmake -G "Visual Studio 17 2022" -A x64 ..
# cmake -G "Visual Studio 16 2019" -A x64 ..
# cmake -G "Ninja" ..
# cmake -G "Unix Makefiles" ..

# cmake --build . --config Release

#************  添加静态库  ************
# set(lib_name "modelPose")
# add_library(${lib_name} STATIC ${SRC})                                       # 创建名为lib_name的静态库
# target_link_libraries(${lib_name} PRIVATE ${OpenCV_LIBS} onnxruntime)                 # 链接所需的库