背景

自制图片处理App开始接入AI能力啦，研究了下如何将ONNX模型转换为CoreML模型，并且在MacOS上运行（iOS应该也是可以支持的，毕竟都是CoreML框架）

模型的训练和转换成ONNX

使用的是PyTorch官方Example中风格迁移的模型，针对梵高的星夜进行了训练，如何导出成ONNX，可以参考RealBasicVSR如何将模型转换成onnx，这里要注意的是opset选择10，我选择11遇到了一些兼容性的问题

ONNX模型转换为CoreML模型

ONNX转换成CoreML使用的是coremltools，一开始我使用了默认写法进行转换，发现输入输出的动态shape失效了，最后转换脚本进行了修改

import coremltools as ct
import sys
from coremltools.models.neural_network import flexible_shape_utils

model_path = sys.argv[1]
output_path = sys.argv[2]
model = ct.converters.onnx.convert(model=model_path, minimum_ios_deployment_target='13')
model.save(output_path)


# change input/output shape config
spec = ct.utils.load_spec(output_path)

input_name = spec.description.input[0].name
output_name = spec.description.output[0].name

flexible_shape_utils.set_multiarray_ndshape_range(spec, 
                                 feature_name=input_name, 
                                 lower_bounds=[0,3, 0, 0], 
                                 upper_bounds=[1,3, -1, -1])
flexible_shape_utils.set_multiarray_ndshape_range(spec, 
                                 feature_name=output_name, 
                                 lower_bounds=[0,3, 0, 0], 
                                 upper_bounds=[1,3, -1, -1])

# re-initialize the mlmodel object
model = ct.models.MLModel(spec)
model.save(output_path)

通过flexible_shape_utils.set_multiarray_ndshape_range将输入输出的shape设定为一个范围， upper_bounds是-1表示无上限，最后在xcode中查看模型的输入输出如下

CoreML加载模型

加载模型很简单，苹果会自动为你建立一个类，和模型文件名称相同，我的文件名是VangoST

NSURL *modelUrl = [[NSBundle mainBundle] URLForResource:modelName withExtension:@"mlmodelc"];
MLModelConfiguration *conf = [MLModelConfiguration.alloc init];
VangoST *model = [[VangoST alloc] initWithContentsOfURL:modelUrl configuration:conf error:nil];``

注意寻找模型路径时，后缀是mlmodelc，不是mlmodel

CoreML使用模型

使用模型的关键是理解输入和输出数据，我的模型输入是(1,3,w,h)尺寸的张量。这里有个问题是通过OpenCV读出来的像素排列是一行一行的，但是模型输入像素排列是一列一列的，所以需要做一个转换，我用了比较粗暴的方式，直接遍历像素了

MLMultiArray *inputData = [MLMultiArray.alloc initWithShape:@[@1, @3, @(img.cols), @(img.rows)] dataType:MLMultiArrayDataTypeFloat32 error:nil];
for (int i = 0; i < img.rows; ++i) {
    for (int j = 0; j < img.cols; ++j) {
        ((float *)inputData.dataPointer)[img.cols * img.rows * 0 + j * img.rows + i] = img.at<cv::Vec3f>(i, j)[2];
        ((float *)inputData.dataPointer)[img.cols * img.rows * 1 + j * img.rows + i] = img.at<cv::Vec3f>(i, j)[1];
        ((float *)inputData.dataPointer)[img.cols * img.rows * 2 + j * img.rows + i] = img.at<cv::Vec3f>(i, j)[0];
    }
}

通过一行代码执行模型的神经网络代码

VangoSTOutput *output = [model predictionFromInput:inputData error:&error];

得到输出后，使用同样的方式再将数据转成OpenCV的格式

int oW = [output._132.shape[3] intValue];
int oH = [output._132.shape[4] intValue];
cv::Mat imgOuput(oH, oW, CV_32FC3);
for (int i = 0; i < img.rows; ++i) {
    for (int j = 0; j < img.cols; ++j) {
        float r = ((float *)output._132.dataPointer)[oW * oH * 0 + j * oH + i];
        float g = ((float *)output._132.dataPointer)[oW * oH * 1 + j * oH + i];
        float b = ((float *)output._132.dataPointer)[oW * oH * 2 + j * oH + i];
        imgOuput.at<cv::Vec3f>(i, j) =  {b, g, r};
    }
}