背景
自制图片处理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};
}
}