背景
训练了一些图片处理的AI模型,本来打算在MacOS(iOS)上使用CoreML来进行计算,但是有一个动漫化的模型怎么转CoreML都报错,于是干脆直接使用ONNXRuntime来进行转换
安装ONNXRuntime
去官网下载下来,就是标准的include和lib,将include加入Header Search Path,将lib里的库加入库依赖即可,这里我下载的是Macos的,iOS的估计也类似。如果你想要使用cocoapods进行封装也是可以的
准备环境
Ort::AllocatorWithDefaultOptions allocator;
Ort::Env env(ORT_LOGGING_LEVEL_ERROR, "imageprocess");
Ort::Session session(env, [modelUrl.path UTF8String], Ort::SessionOptions());
通过ONNX模型文件路径创建Ort::Session
,后续的模型计算都是通过这个session进行
准备输入数据
我们模型的输入是一个1x3xWxH
的张量
Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
const int64_t input_shape[] = {1, 3, img.cols, img.rows};
const size_t input_shape_len = sizeof(input_shape) / sizeof(input_shape[0]);
const size_t model_input_len = img.cols * img.rows * 3 * sizeof(float);
Ort::Value input_tensor = Ort::Value::CreateTensor((OrtMemoryInfo *)memory_info, inputTensorData, model_input_len, input_shape,
input_shape_len, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT);
输入张量input_tensor
就建立好了,其中inputTensorData
我是从OpenCV的cv::Mat
转换过来的,由于cv::Mat
是WxHx3格式的,所以需要做下面的转换。/ 255.0 * 2.0 - 1.0
是为了把输入像素规范化成-1到1的区间
void *inputTensorData = malloc(img.cols * img.rows * 3 * sizeof(float));
for (int i = 0; i < img.rows; ++i) {
for (int j = 0; j < img.cols; ++j) {
auto color = img.at<cv::Vec4b>(i, j);
((float *)inputTensorData)[img.cols * img.rows * 0 + i * img.cols + j] = color[0] / 255.0 * 2.0 - 1.0;
((float *)inputTensorData)[img.cols * img.rows * 1 + i * img.cols + j] = color[1] / 255.0 * 2.0 - 1.0;
((float *)inputTensorData)[img.cols * img.rows * 2 + i * img.cols + j] = color[2] / 255.0 * 2.0 - 1.0;
}
}
模型计算
const char* input_names[] = {"input"};
const char* output_names[] = {"888"};
Ort::RunOptions runOption;
auto returnTensors = session.Run(runOption, input_names, &input_tensor, 1, output_names, 1);
提供输入输出节点名称和输入的张量,执行session.Run
即可
处理输出数据
取session.Run
返回的第一个数据
Ort::Value &outputTensor = returnTensors[0];
取出字节数据
const float * rawData = outputTensor.GetTensorData<float>();
模型的输出格式也是1x3xWxH
,需要转换成WxHx3
auto outputShape = outputTensor.GetTensorTypeAndShapeInfo().GetShape();
int oW = int(outputShape[2]);
int oH = int(outputShape[3]);
cv::Mat imgOuput(oH, oW, CV_32FC4);
for (int i = 0; i < imgOuput.rows; ++i) {
for (int j = 0; j < imgOuput.cols; ++j) {
float r = ((float *)rawData)[oW * oH * 0 + i * oW + j];
float g = ((float *)rawData)[oW * oH * 1 + i * oW + j];
float b = ((float *)rawData)[oW * oH * 2 + i * oW + j];
imgOuput.at<cv::Vec4f>(i, j) = {
255.0f * (fmin(1.0f, fmax(r, -1.0f))* 0.5f + 0.5f),
255.0f * (fmin(1.0f, fmax(g, -1.0f))* 0.5f + 0.5f),
255.0f * (fmin(1.0f, fmax(b, -1.0f))* 0.5f + 0.5f), 255.0f};
}
}
这样就能把模型的输出转换成OpenCV的cv::Mat了,整个模型计算结束~