2017 年 6 月，Apple 推出了 CoreML，这是一个旨在将机器学习模型集成到 iOS 应用程序中的框架，为开发人员提供了大量的可能性，包括图像分析到 NLP（自然语言处理），决策树学习等等。
在 2018 年春季，我开始深入研究 CoreML，并学习了神经风格转移。 这篇文章，我将讲解一下大致的实现过程。
训练一个图像风格迁移模型
我们使用 PyTorch 官方提供的 Fast Neural Style 代码（地址是 https://github.com/pytorch/examples/tree/master/fast_neural_style）来训练模型，我们使用默认参数，并且每 2000 次循环保存一个 checkpoint：
1 2 3 4 5 6 7 8 9 10 11 | python neural_style/neural_style.py train \ --dataset ~/Documents/dataset \ --style-image ~/Documents/images/styles/starry_night.jpg \ --style-size 512 \ --batch-size 2 \ --style-weight 1e10 \ --content-weight 1e5 \ --checkpoint-model-dir ~/Documents/checkpoints \ --save-model-dir ~/Documents/models \ --log-interval 10 \ --cuda 1 |
将模型导出到 CoreML
PyTorch 本身是不支持直接导出 CoreML 模型的（至少目前是这样），但我们可以借助 ONNX（Open Neural Network Exchange）的帮助，来先转换成 .onnx 格式，然后再转换成 CoreML。
PyTorch → ONNX→ CoreML
文件 neural_style.py 中已经有输出位 ONNX 格式的代码了，我们只需要完成第二步就行。我们首先定义一个转换器：
1 2 3 4 5 6 7 8 9 10 11 12 | import sys from onnx import onnx_pb from onnx_coreml import convert model_in = sys.argv[1] model_out = sys.argv[2] model_file = open(model_in, \\\'rb\\\') model_proto = onnx_pb.ModelProto() model_proto.ParseFromString(model_file.read()) coreml_model = convert(model_proto, image_input_names=[\\\'inputImage\\\'], image_output_names=[\\\'outputImage\\\']) coreml_model.save(model_out) |
接下来我们进行转换：
1 2 3 4 5 6 7 8 9 10 | python ./neural_style/neural_style.py eval \ --content-image ~/Documents/data/images/original.jpg \ --output-image ~/Documents/data/images/stylized-test.jpg \ --model ~/Documents/checkpoints/checkpoint_5000.pth \ --cuda 0 \ --export_onnx ~/Documents/models/pytorch_model.onnx \ && \ python ./onnx_to_coreml.py \ ~/Documents/models/pytorch_model.onnx \ ~/Documents/models/final_model.mlmodel |
当然，如果你不想自己训练，可以在这里下载训练好的模型。
在 iOS 上使用 CoreML NST 模型
我们创建一个 UIImageView 和 UIButton，添加一个 UIBarButtonItem 来从设备的照片库中选择一个图像并将其设置到 UIImageView 上。
等 UI 完成后，我们就可以把训练的模型导入进来了，只需要将模型文件拖到文件视图中就行了。点击模型文件，可以看到一些模型信息：
这是一个新的 Model 类会被自动创建，查看文件可知我们将要用到的函数是 prediction(inputImage: CVPixelBuffer).
我们选定一张图片后：
- 首先保存图像的原始大小
- 把输入 UIImage 转换成一个 720 x 720 的 CVPixelBuffer
- 输入到模型
- 把模型的输出 CVPixelBuffer 转换回 UIImage
- 重新调整图像大小到原始大小
下面是一个 Swift 4 实现：
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | func process(input: UIImage, completion: @escaping FilteringCompletion) { // Initialize the NST model let model = StarryNight() // Next steps are pretty heavy, better process them on a background thread DispatchQueue.global().async { // 1 - Transform our UIImage to a PixelBuffer of appropriate size guard let cvBufferInput = input.pixelBuffer(width: 720, height: 720) else { print("UIImage to PixelBuffer failed") completion(nil) return } // 2 - Feed that PixelBuffer to the model guard let output = try? model.prediction(inputImage: cvBufferInput) else { print("Model prediction failed") completion(nil) return } // 3 - Transform PixelBuffer output to UIImage guard let outputImage = UIImage(pixelBuffer: output.outputImage) else { print("PixelBuffer to UIImage failed") completion(nil)http://www.pytorchtutorial.com/wp-admin/post.php?post=608&action=edit# return } // 4 - Resize result to the original size, then hand it back to the main thread let finalImage = outputImage.resize(to: input.size) DispatchQueue.main.async { completion(finalImage) } } } |
下面是最终效果：
本文的项目源代码已经放到 Github 上。
本站微信群、QQ群(三群号 726282629)：
