C++ 实现ONNX语义分割（TopFormer)

TopFormer是一个比较高效的语义分割模型，
在TopFormer ONNX主页上可以找到python的实现。

用一个小猫的图片测试一下

瑕疵是有一些，不过它的特点是非常快，视频语义分割能达到实时（本地CPU）。
TopFormer模型的输出Size是64x64，这个是resize后的效果。

用cpp实现看看效果。
cpp版的onnxruntime安装这里就不说了，参考链接
还需要安装cpp版的opencv，我选的是4.5，安装参考

然后你的CmakeLists.txt里面要link这两样

target_link_libraries(${PROJECT_NAME}
        ${OpenCV_LIBS}
        ${ONNXRUNTIME_LIB}
        )

来一波include

#include <iostream>
#include <fstream>
#include <array>
#include <string>

#include <opencv4/opencv2/core/mat.hpp>
#include <opencv4/opencv2/imgcodecs.hpp>
#include <opencv4/opencv2/imgproc.hpp>
#include <core/session/experimental_onnxruntime_cxx_api.h>

using namespace std;
using namespace cv;

onnx模型直接用github下下来的就好，
定义几个onnxruntime变量

const string modelFile = "TopFormer-S_512x512_2x8_160k.onnx";
Ort::Env env;
Ort::SessionOptions sessionOptions;
Ort::Session session(nullptr);

try{
     session = Ort::Session(env, modelFile.data(), sessionOptions);
} catch(Ort::Exception& e) {
     cout << e.what() << endl;
     return 1;
}

TopFormer的输入Size是1x3x512x512，输出Size是1x150x64x64,
所以先定义下形状。

constexpr int64_t numChannels = 3;  //constexpr修饰，表示编译期就可以得到常量值去优化
constexpr int64_t width = 512;
constexpr int64_t height = 512;
constexpr int64_t out_width = 64;
constexpr int64_t out_height = 64;
constexpr int64_t numClasses = 150;  //类别数
constexpr int64_t numInputElements = numChannels * height * width;
constexpr int64_t numOutputElements = numClasses * out_height * out_width;

定义输入输出数组，输入输出Tensor，
有人会说，输入本身就是矩阵了，为什么要转成数组。
用array, vector的话会保存指针，就不需要反复create Tensor，只需要把数据copy到数组。

array<float, numInputElements> input;
array<float, numOutputElements> output;

//输入输出需要转成Tensor
//这个Tensor保存的是输入输出数组的pointer，所以不能删除输入输出数组
auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
auto inputTensor = Ort::Value::CreateTensor<float>(memory_info, input.data(),
                                                   input.size(), inputShape.data(),
                                                   inputShape.size());
auto outputTensor = Ort::Value::CreateTensor<float>(memory_info, output.data(),output.size(),
                                                    outputShape.data(), outputShape.size());
//也可以不用转成数组，直接用Tensor，但是array，vector会便利一些
//因为图片如果很多，就需要反复create Tensor,用数组的话只需要copy到数组

读入image，预处理（比如-mean, /std）,
注意要把image处理成CxHxW的形式

Mat img = imread(imageFile);
const vector<float> imageVec = loadImage(img); //CxHxW: [ch1:HXW][ch2:HxW][ch3:HxW]
//预处理（resize, normalize）放在loadImage函数里面

预处理中mean和std我们用TopFormer python版本中的数据

//-mean, /std
float img_mean[3] = {0.485, 0.456, 0.406};
float img_std[3] = {0.229, 0.224, 0.225};

注意要把image拉成一维数组，顺序是CxHxW

resize(image, image, Size(512, 512));
image = image.reshape(1, 1);  //图像拉成向量
//HWC -> CHW (transpose(2, 0, 1)
for(size_t ch = 0; ch < 3; ++ch) {
    for(size_t i = ch; i < vec.size(); i+= 3) {
        //保存成CxHxW的顺序
    }
}

copy图像数据到input数组

//把图像copy进input数组
copy(imageVec.begin(), imageVec.end(), input.begin());

定义输入输出的name，要和onnx模型的输入输出name一致，
可以直接定义，也可以从onnx模型中获取。

    //输入输出的name,可以自己定义，也可以从onnx模型中获取
    const array<const char*, 1> inputNames = {"input"};
    const array<const char*, 1> outputNames = {"output"};


//    //从onnx模型中获取输入输出名称
//    vector<const char*> inputNames;
//    vector<const char*> outputNames;
//
//    Ort::AllocatorWithDefaultOptions alloc;
//    for(size_t i = 0; i < session.GetInputCount(); i++) {
//        //push_back是创建元素，然后添加到vector末尾，
//        //emplace_back是直接在vector末尾创建元素，省去了copy或移动元素的过程
//        inputNames.emplace_back(session.GetInputName(i, alloc));
//    }
//    for(size_t i = 0; i < session.GetOutputCount(); i++) {
//        outputNames.emplace_back(session.GetOutputName(i, alloc));
//    }

推理过程，
这里的outputTensor是前面创建的，指向的是output数组，
所以最后的输出数据会在output数组里。

//inference
Ort::RunOptions runOptions;
try{
    session.Run(runOptions, inputNames.data(), &inputTensor, 1,
                outputNames.data(), &outputTensor, 1);
}catch (Ort::Exception& e) {
    cout << e.what() << endl;
    return 1;
}

现在我们得到的输出output是一条向量，并不是1x150x64x64的形状，
而我们想要的是1x150x64x64，然后在64x64的每个像素上，取所有channel的最大值index作为label,
取出label对应的颜色保存到像素上。
很复杂有木有。

先解决每个label对应什么颜色吧，
TopFormer共有150个类别，在ade20k_label_colors.txt（github下载）中有每个类别对应的颜色。
也可以自己定义颜色表。

const int classes = 150;
unsigned char colors[classes][3];
void color_map(int classes) {
   int r = 0, g = 0, b = 0;
   for(int i = 0; i < classes; i++) {  //150类
       int c = i;
       for(int j = 0; j < 8;j ++) {
           r = r | (bitget(c, 0) << 7-j);
           g = g | (bitget(c, 1) << 7-j);
           b = b | (bitget(c, 2) << 7-j);
           c = c >> 3;
       }
       for(int j = 0; j < 3; j++) {
           colors[i][0] = r;
           colors[i][1] = g;
           colors[i][2] = b;
       }
   }
}

把“在64x64的每个像素上，取所有channel的最大值index作为label,
取出label对应的颜色保存到像素上。“放在一起处理。
那么需要一个64x64x3的输出map

Mat output_map(out_height, out_width, CV_MAKETYPE(CV_8UC1, 3));

现在输出是1x150x64x64拉成一条的向量，
保存顺序是这样的，
[ch1: 64x64] [ch2: 64x64] … [ch150: 64x64]
只需要每隔64x64个点取各个channel数值，找到最大值所在的label即可。
注意opencv是以BGR的channel顺序储存的。

看看output_map的结果吧。
这里是64x64大小，没有resize.
可以看到跟python版本效果类似。

完整代码放在github的semantic

参考资料