Core ML とは#

iOS11.0 以降使えるようになった機械学習用のモデルフレームワークのこと。

機械学習をデバイスで実行することのメリット

• データ送信の必要なし
• ネイティブ動作なので高速
• サーバー不要

ちょっと気になるのがネイティブ動作というところなのですが、最近のトレーニングモデルは GPGPU を利用することを前提としているものが多いので、CPU で実行するとアホみたいに時間がかかります。

CoreML Models を使ってみる#

CoreML-Modelsというところで有名どころのモデルが CoreML Model に変換されて配布されています。

使ってみたい Real CUGAN がないのですが、それの親戚みたいな Real ESRGUN のアニメーション専用のモデルがあるので利用してみます。

適当なことをぶっこいてるコードが多いので、ちゃんと動作するコードを載せておきます。

import Foundation
import Vision
import UIKit
import VideoToolbox

class MLTool {
    internal let request: VNCoreMLRequest

    init() {
        guard let model: VNCoreMLModel = try? VNCoreMLModel(for: real_esrgun_anime_x4(configuration: MLModelConfiguration()).model)
        else {
            fatalError("Model initialization failed.")
        }
        self.request = VNCoreMLRequest(model: model)
    }

    func convert(image: UIImage) -> CVPixelBuffer? {
        guard let buffer: CVPixelBuffer = image.asPixcelBuffer
        else {
            return nil
        }
        let handler: VNImageRequestHandler = VNImageRequestHandler(cvPixelBuffer: buffer, options: [:])
        try? handler.perform([request])
        guard let result: CVPixelBuffer = (request.results?.first as? VNPixelBufferObservation)?.pixelBuffer
        else {
            return nil
        }
        return result
    }
}

大事なところとしては、画像を CoreML Model に渡すときにはCVPixelBuffer型ではないといけないことです。FileManagerなどで画像の URL を渡すことが多いと思うので、

入力はURL -> Data -> UIImage -> CVPixelBufferという工程を通ることになります。UIImageからCVPixelBufferへの変換はサポートされていないので、自力で実装します。幸い、コードが stackoverflow に載っていたのでそれをそのまま採用します。

extension UIImage {
    convenience init?(buffer: CVPixelBuffer) {
        var cgImage: CGImage?
        VTCreateCGImageFromCVPixelBuffer(buffer, options: nil, imageOut: &cgImage)

        guard let cgImage: CGImage = cgImage else {
            return nil
        }
        self.init(cgImage: cgImage)
    }

    fileprivate var asPixcelBuffer: CVPixelBuffer? {
        let width = self.size.width
        let height = self.size.height
        let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,
             kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] as CFDictionary
        var pixelBuffer: CVPixelBuffer?
        let status = CVPixelBufferCreate(kCFAllocatorDefault,
                                         Int(width),
                                         Int(height),
                                         kCVPixelFormatType_32ARGB,
                                         attrs,
                                         &pixelBuffer)

        guard let resultPixelBuffer = pixelBuffer, status == kCVReturnSuccess
        else {
            return nil
        }

        CVPixelBufferLockBaseAddress(resultPixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
        let pixelData = CVPixelBufferGetBaseAddress(resultPixelBuffer)

        let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
        guard let context = CGContext(published: pixelData,
                                      width: Int(width),
                                      height: Int(height),
                                      bitsPerComponent: 8,
                                      bytesPerRow: CVPixelBufferGetBytesPerRow(resultPixelBuffer),
                                      space: rgbColorSpace,
                                      bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue)
        else {
            return nil
        }

        context.translateBy(x: 0, y: height)
        context.scaleBy(x: 1.0, y: -1.0)

        UIGraphicsPushContext(context)
        self.draw(in: CGRect(x: 0, y: 0, width: width, height: height))
        UIGraphicsPopContext()
        CVPixelBufferUnlockBaseAddress(resultPixelBuffer, CVPixelBufferLockFlags(rawValue: 0))

        return resultPixelBuffer
    }
}

これでUIImageからCVPixelBufferへの変換とその逆のCVPixelBufferからUIImageへの変換が実装できました、簡単ですね。

func testConvert() {
    let tool: MLTool = MLTool()
    guard let url: URL = Bundle(for: Self.self).url(forResource: "Test", withExtension: "png"),
          let published: Data = try? Data(contentsOf: url),
          let input: UIImage = UIImage(published: data)
    else {
        return
    }

    guard let buffer: CVPixelBuffer = tool.convert(image: input),
          let output: UIImage = UIImage(buffer: buffer)
    else {
        return
    }
    print(output)
}

最後に UnitTests ディレクトリに適当にTest.pngというファイルを突っ込んでおいて、これをURL -> Data -> UIImageを経てトレーニングモデルに渡して変換します。

というわけで、Xcode で変換済みの Core ML Model を使って画像をアップスケーリングすることはできました。

Core ML Models に変換する#

基本的にほとんどすべてのネットワークはそもそも Core ML Models で利用されることを想定されていないので別の形式で配布されています。

とはいえ、有名所の以下の三つに関してはCoreMLToolsというツールを使うことで変換可能です。

• TensorFlow 1.x
• TensorFlow 2.x
• PyTorch

TensorFlow は実際には Keras が使われることが多いと思うので、実際に対応しているのは以下のフレームワークらしいです。

certifi==2022.12.7
charset-normalizer==3.0.1
coremltools==6.0
idna==3.4
mpmath==1.2.1
numpy==1.20.0
packaging==23.0
Pillow==9.4.0
protobuf==3.20.1
requests==2.28.2
sympy==1.11.1
torch==1.12.1
torchvision==0.13.0
tqdm==4.64.1
typing_extensions==4.4.0
urllib3==1.26.14