多类别语义分割C++实现（onnx部署）详细教程

写这篇是因为网上搜不到，我觉得写一下或许可以帮助到需要的人。

目标已经制定，那就分两步走：

1、分割网络的训练和模型生成

目前在二维机器视觉的分割网络已经相当成熟，有很多不错的网络可以参考使用，我这里使用的Unet网络，也是一个十分经典的网络。这里的训练就不说了，网上有很多教程，我们只需要将训练好的权重文件转为onnx即可，因为目前有很多平台都支持onnx的转换，所以权重文件转onnx这一技能是十分必要的（其实也不复杂，网上有很多教程）。

2、根据网络模型逻辑编写C++代码

得到onnx模型后可以查看它的结构以及输入输出，这在以后进行模型转换的时候还是挺重要的，当然在这篇文章里，我们不再进行模型转换了，使用onnx模型就足够了。

 可以看到模型的输入和输出均为float32：1*3*512*512，这将作为之后C++代码逻辑的输入和输出依据。补充一下这里的输入输出对应的是N*C*H*W，也就是

网上都能查到的，我就不做赘述了。（接下来一些简单的就不做解释了，都可以查到的）

需要用到Opencv455或者更高版本，主要用到里面的DNN模块，当然也可以用到onnxruntime，但是这篇未使用。

现在开始写代码：

输入是512*512，这要求我们输入的图片最好也是这个大小，这里可以使用图片的缩放：

        //原始图片大小
        int org_width = img.cols;
        int org_height = img.rows;

        //计算缩放比例
        double scale = 1.0;
        if (org_width >= org_height) {
            scale = (double)512 / org_width;
        }
        else {
            scale = (double)512 / org_height;
        }

        //缩放图片，并保持长宽比
        Mat resized_img;
        resize(img, resized_img, Size(), scale, scale);

        //计算灰度填充的大小
        int pad_x = (512 - resized_img.cols) / 2;
        int pad_y = (512 - resized_img.rows) / 2;

        //加上灰色填充
        Mat padded_img = Mat::zeros(512, 512, CV_8UC1);
        copyMakeBorder(resized_img, padded_img, pad_y, pad_y, pad_x, pad_x, BORDER_CONSTANT, Scalar(128));

        //图片通道转换和resize
        padded_img.convertTo(padded_img, CV_32FC3, 1.f / 255.f);
        cv::cvtColor(padded_img, padded_img, cv::COLOR_BGR2RGB);
        cv::resize(padded_img, padded_img, cv::Size(512, 512), 0, 0, cv::INTER_LINEAR);
        //
        //模型处理...处理结束后再变为原来大小，不影响二维结果判断
        //
        //去除灰色填充并还原原始大小
        Rect roi(pad_x, pad_y, org_width * scale, org_height * scale);
        Mat roi_img = predMat(roi);
        resize(roi_img, roi_img, Size(org_width, org_height));

在模型处理后，我又将图片缩放为原本大小，不影响二维结果判断。

接下来就是onnx的加载了（上面代码中的模型处理部分），很简单的几行代码：

    Mat inputBolb = blobFromImage(padded_img);

	dnn::Net Unet = cv::dnn::readNetFromONNX(modelFile);

	Unet.setInput(inputBolb);
	Mat output = Unet.forward();

	int N = inputBolb.size[0];
	int C = inputBolb.size[2];
	int H = inputBolb.size[2];
	int W = inputBolb.size[3];
	Mat predMat = Mat::zeros(512, 512, CV_32F);

	for (int h = 0; h < H; h++) {
		for (int w = 0; w < W; w++) {
			float bg = output.ptr<float>(0, 0, h)[w];
			float c1 = output.ptr<float>(0, 1, h)[w];
			float c2 = output.ptr<float>(0, 2, h)[w];
            float c3 = output.ptr<float>(0, 3, h)[w];
			if (bg >= c1 && bg >= c2 && bg >= c3) {//背景
				predMat.at<float>(h, w) = 255;
			}
			else if (c1 >= c2 && c1 >= c3) {//类别1
				predMat.at<float>(h, w) = 200;
			}
            else if (c2 >= c1 && && c2 >= c3) {//类别2
				predMat.at<float>(h, w) = 150;
			}
			else {//类别3
				predMat.at<float>(h, w) = 0;
			}
		}
	}

这里用的是单通道赋值，最终的分割图片是一张单通道灰度图，想要彩色图的也可以自己赋彩色值，或者直接根据我的结果创建查找表对应不同的彩色值。

完整代码如下：

#include <vector>
#include <cmath>
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/opencv.hpp>
#include <iostream>
using namespace cv;
using namespace cv::dnn;
using namespace std;


int main() {
    Mat img = imread("你要分割的图片.jpg");
    String modelFile = "你的分割模型.onnx";
    if (img.empty()) {
        cout << "未找到图片！" << endl;
        return -1;
    }
    try {
        //原始图片大小
        int org_width = img.cols;
        int org_height = img.rows;

        //计算缩放比例
        double scale = 1.0;
        if (org_width >= org_height) {
            scale = (double)512 / org_width;
        }
        else {
            scale = (double)512 / org_height;
        }

        //缩放图片，并保持长宽比
        Mat resized_img;
        resize(img, resized_img, Size(), scale, scale);

        //计算灰度填充的大小
        int pad_x = (512 - resized_img.cols) / 2;
        int pad_y = (512 - resized_img.rows) / 2;

        //加上灰色填充
        Mat padded_img = Mat::zeros(512, 512, CV_8UC1);
        copyMakeBorder(resized_img, padded_img, pad_y, pad_y, pad_x, pad_x, BORDER_CONSTANT, Scalar(128));

        //图片通道转换和resize
        padded_img.convertTo(padded_img, CV_32FC3, 1.f / 255.f);
        cv::cvtColor(padded_img, padded_img, cv::COLOR_BGR2RGB);
        cv::resize(padded_img, padded_img, cv::Size(512, 512), 0, 0, cv::INTER_LINEAR);
        Mat inputBolb = blobFromImage(padded_img);
        dnn::Net net = cv::dnn::readNetFromONNX(modelFile);
        net.setInput(inputBolb);
        Mat output = net.forward();
        int N = inputBolb.size[0];
        int C = inputBolb.size[1];
        int H = inputBolb.size[2];
        int W = inputBolb.size[3];
        Mat predMat = Mat::zeros(512, 512, CV_32F);
        for (int h = 0; h < H; h++) {
            for (int w = 0; w < W; w++) {
                float bg = output.ptr<float>(0, 0, h)[w];
			    float c1 = output.ptr<float>(0, 1, h)[w];
			    float c2 = output.ptr<float>(0, 2, h)[w];
                float c3 = output.ptr<float>(0, 3, h)[w];
			    if (bg >= c1 && bg >= c2 && bg >= c3) {//背景
				    predMat.at<float>(h, w) = 255;
			    }
			    else if (c1 >= c2 && c1 >= c3) {//类别1
				    predMat.at<float>(h, w) = 200;
			    }
                else if (c2 >= c1 && && c2 >= c3) {//类别2
				    predMat.at<float>(h, w) = 150;
			    }
			    else {//类别3
				    predMat.at<float>(h, w) = 0;
			    }
            }
        }
        //去除灰色填充并还原原始大小
        Rect roi(pad_x, pad_y, org_width * scale, org_height * scale);
        Mat roi_img = predMat(roi);
        resize(roi_img, roi_img, Size(org_width, org_height));
        imshow("seg.png", roi_img);
        WaitKey(0);
    }
    catch (char* str)
    {
        cout << str << endl;
    }
    catch (int i)
    {
        cout << i << endl;
    }

    return 0;
}

大概思路就是这，很简单，类别多的可以多加几个if判断和灰度值赋值。草稿代码可能不是那么严谨，但运行起来问题不大。