背景

训练了一些图片处理的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了，整个模型计算结束~