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
#include
#include
#include
#include
#include
#include
#include
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 inputNames;
// vector 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
参考资料