[iOS 11] CoreMLで画像の識別を試してみました（Vision.Frameworkを使わないパターン） #WWDC2017

1 はじめに

前回は、Vision.Frameworkを使用した画像の識別を試してみました。

今回は、Vision.Frameworkが何をラップしているのか、また、直接CoreMLを使用すると何が違うのかを確認するために、同じ画像識別をCoreMLのみ書いてみました。

機械学習のモデル（CoreML用に作成された.mlmodelファイル）は、今回もAppleの開発者ページでダウンロードに可能なっている下記の５種類を使用させて頂いています。

• SqueezeNet
• Places205-GoogLeNet
• ResNet50
• Inception v3
• VGG16

本記事は Apple からベータ版として公開されているドキュメントを情報源としています。 そのため、正式版と異なる情報になる可能性があります。ご留意の上、お読みください。

2 CVPixelBuffer

CVPixelBufferは、メインメモリ内のピクセルを保持するイメージバッファであり、フレームを生成するアプリや、Core Imageを使用するアプリで利用されています。

CoreMLの入力には、CVPixelBufferが必要なためUIImage等は、変換する必要があります。また、入力画像は、モデルごとサイズが決まっており、ここも合わせる必要があります。

下記は、stackoverflowで紹介されていたコードですが、今回は、これをそのまま利用させて頂きました。

extension UIImage {
    func resize(to newSize: CGSize) -> UIImage {
        UIGraphicsBeginImageContextWithOptions(CGSize(width: newSize.width, height: newSize.height), true, 1.0)
        self.draw(in: CGRect(x: 0, y: 0, width: newSize.width, height: newSize.height))
        let resizedImage = UIGraphicsGetImageFromCurrentImageContext()!
        UIGraphicsEndImageContext()
        return resizedImage
    }

    func pixelBuffer() -> 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(data: 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
    }
}

3 画像解析処理

CoreMLを直接使用するコードは、下記のようになります。

モデルのオブジェクトを生成し、prediction()メソッドで解析を行っています。

func coreMLRequest(image:UIImage) {
    // モデルオブジェクトの生成
    let model = GoogLeNetPlaces()
    // GoogLeNetPlacesの入力は224×224ピクセルの画像になっている
    let imageSize:CGSize = CGSize(width: 224, height: 224)
    // CVPixelBufferへの変換
    let pixelBuffer = image.resize(to: imageSize).pixelBuffer()
    // 画像解析
    guard let pb = pixelBuffer, let output = try? model.prediction(sceneImage: pb) else {
        fatalError("error")
    }
    //　結果の表示
    print("label = \(output.sceneLabel)")
    print("probabilityr = \(output.sceneLabelProbs[output.sceneLabel]!)")
}

モデルを差し替える場合、Vision.Frameworkでは、モデルのクラスを差し替えるだけでしたが、今回はそうは行きません。

下記は、GoogLeNetPlacesからInception v3に変更した場合のコードです。

まずは、それぞれのモデルの入力画像のサイズが違うので、それに合わせる必要があります。そして、prediction()の引数も若干違いがあります。また、出力とし受け取るクラス（プロパティ名）も違いますので、それに適用させなければなりません。

func coreMLRequest(image:UIImage) {
    
    let model = Inceptionv3()
    let imageSize:CGSize = CGSize(width: 299, height: 299)

    let pixelBuffer = image.resize(to: imageSize).pixelBuffer()
    
    guard let pb = pixelBuffer, let output = try? model.prediction(image: pb) else {
        fatalError("error")
    }
    print("identifier = \(output.classLabel)")
    print("confidence = \(output.classLabelProbs[output.classLabel]!)")
}

画像識別の各モデルは、非常によく似たクラスになっていますが、一応、少しずつ違っていました。

各モデルのコードが必要な場合は、下記をご参照下さい。（モデルは含まれておりませんので、別途ダウンロードが必要です。）
[GitHub] https://github.com/furuya02/CoreMLSample2

4 やってみた

実際にやってみたところ、結果は下記のようになりました。

sizeは、モデル(.mlmodel)のサイズ（MByte）であり、timeは、解析にかかった時間です。Vision.Frameworkとの比較もしたくて、秒単位で雑に計測してみました。

なお、結果は、サンプル数も非常に限られたものですので、あくまで「解析の一例」程度の意味しか無いことを予めご了承下さい。

やってみて分かったのは、同じ画像を解析したはずなのに、結果がVision.Framework と比べると結構違うことです。

これは、あくまで勝手な予想なのですが、画像のサイズを入力に合わせるロジックが違うからではないかと考えています。Vision.Frameworkでは、画像のサイズは勝手に合わせてくれているようですが、そこがどうなっっているのか不明です。

なお、処理速度については、明らかに早いと感じました。

5 最後に

今回は、Vision.Frameworkが何をラップしているのか、また、直接CoreMLを使用すると何が違うのかを確認する事が目的でしたが、その目的は達成できたように思います。

クラスを合わせる必要はありますが、CoreMLを直接使用することも、それほどハードルは高くないかも知れません。

今回、黒柴は、「Eskimo dog, husky」の判定が多かったのですが、この辺も、画像サイズを変換した時の、縦横の扱いが影響している予感がします。

コードは、下記に起きました。モデルは含まれておりませんので、別途ダウンロードが必要です。
[GitHub] https://github.com/furuya02/CoreMLSample2

6 参考リンク

Convert Image to CVPixelBuffer for Machine Learning Swift
Apple Document > Core Video > CVPixeBuffer
Swift World: What’s new in iOS 11 — Core ML
WWDC 2017 Vision Framework: Building on Core ML
Build more intelligent apps with machine learning.