TensorFlow & TensorSwiftを使ったiOSアプリの試作

TensorFlowのチュートリアルであるDeep MNIST for Expertsを少しだけ修正して、顔写真の多クラス分類をするiOSアプリを試作してみたのでメモとして残します。

TensorFlowはiOSでも使用することができますが、ライブラリのビルドから行う必要がありますし、Objective C++で開発することになるため少々面倒です。そこで、TensorFlowの学習済みモデルをSwiftで使うことのできるTensorSwiftを使用してみました。開発にあたってはこちらの投稿を参考にさせていただきました。

（**【追記】**iOS側でもTensorFlowを使う場合の開発手順をこちらに書きました）

開発手順

1. データセットの用意
2. TensorFlowで学習させる
3. 学習済みのデータ（変数）をTensorSwift用のファイルに書き出す
4. TensorSwiftを使ってアプリを作る

#1. データセットの用意
学習対象としてサッカー選手を選んでみました。本田圭佑、ネイマール、オーバメヤンの3選手です。髪型の特徴や面長であることなど共通している部分がそれなりに存在し、角度や表情によってはかなり似ているように見える瞬間があるので、学習させるとどうなるか興味がありました。

^{※1 ※2 ※3}

各選手ごとに300枚程度の写真を学習するために、Microsoft Cognitive ServicesのBing Image Search APIを使いました。
Bing Image Search APIは1000回/月、5回/秒までのAPIコールを90日間無料で使用できます。

キーワード検索で得られた画像に対してOpenCV (haarcascade_frontalface_alt2.xml)を使って顔認識を行い、領域の切り出しやリサイズを行ってデータ化します。
MNISTの画像は28x28ピクセルですが、人の顔なので大きめの116x116ピクセルで作成してみました。

ここまでの作業は、Pythonなどでコードを書けばすべて自動で進められますが、不要な画像（誤認識されたものや違う人物など）を目視で除外したのち、データを分類する作業が最後に必要です。この部分はどうしても人力に頼ることになります。

PythonからOpenCVを使うコードは次の記事が大変参考になりました。
Heroku + OpenCVで簡易顔検出API
http://memo.sugyan.com/entry/20151203/1449137219

#2. TensorFlowで学習させる
Deep MNIST for Expertsをベースに、Pythonのコードを修正していきます。
読み込むデータセットの変更以外にも、識別クラス数、学習させる画像のサイズ、カラーチャンネル数（MNISTのグレイスケールからRGBの3チャンネルへ）の変更が必要です。テンソルとモデルをそれに合わせて細かく修正するだけで済みます。

学習は500ステップしかさせていませんが、筆者のiMac (GeForce GTX 680MX) ではGPU使用で1時間ほどで完了し、96%の精度が出ました。

#3. 学習済みのデータ（変数）をTensorSwift用のファイルに書き出す
こちらの投稿の通り、TensorSwiftで使用する学習済みデータをファイル化します。
あらかじめモデルのすべてのVariableにname属性を付けたうえで学習させ、それをすべて個別のファイルに書き出します。

#4. TensorSwiftを使ってアプリを作る

カメラ画像をリアルタイムで判定器にかけるアプリを試作してみます。

##変数読み込みとソフトマックス回帰モデルの作成

書き出したテンソルをTensorSwiftに読み込み、ソフトマックス回帰モデルを作成するコードは次の通りです。
TensorSwiftのMNISTサンプルコードとほとんど同じですが、識別クラス数、画像サイズ、カラーチャンネル数が変わっているのでその影響を受ける部分だけが変更されています。

Classifier.swift

import TensorSwift

public struct Classifier {
    
    public let W_conv1: Tensor
    public let b_conv1: Tensor
    public let W_conv2: Tensor
    public let b_conv2: Tensor
    public let W_fc1: Tensor
    public let b_fc1: Tensor
    public let W_fc2: Tensor
    public let b_fc2: Tensor
    
    public func classify(x_image: Tensor) -> Int {
        let h_conv1 = (x_image.conv2d(filter: W_conv1, strides: [1, 1, 1]) + b_conv1).relu
        let h_pool1 = h_conv1.maxPool(kernelSize: [2, 2, 1], strides: [2, 2, 1])
        
        let h_conv2 = (h_pool1.conv2d(filter: W_conv2, strides: [1, 1, 1]) + b_conv2).relu
        let h_pool2 = h_conv2.maxPool(kernelSize: [2, 2, 1], strides: [2, 2, 1])
  
        let h_pool2_flat = h_pool2.reshape([1, 29 * 29 * 64])
        let h_fc1 = (h_pool2_flat.matmul(W_fc1) + b_fc1).relu
        
        let y_conv = (h_fc1.matmul(W_fc2) + b_fc2).softmax
 
        return y_conv.elements.enumerate().maxElement { $0.1 < $1.1 }!.0
    }
}

extension Classifier {
    public init(path: String) {
        W_conv1 = Tensor(shape: [5, 5, 3, 32], elements: loadFloatArray(directory: path, file: "W_conv1"))
        b_conv1 = Tensor(shape: [32], elements: loadFloatArray(directory: path, file: b_conv1"))
        W_conv2 = Tensor(shape: [5, 5, 32, 64], elements: loadFloatArray(directory: path, file: "W_conv2"))
        b_conv2 = Tensor(shape: [64], elements: loadFloatArray(directory: path, file: "b_conv2"))
        W_fc1 = Tensor(shape: [29 * 29 * 64, 1024], elements: loadFloatArray(directory: path, file: "W_fc1"))
        b_fc1 = Tensor(shape: [1024], elements: loadFloatArray(directory: path, file: "b_fc1"))
        W_fc2 = Tensor(shape: [1024, 3], elements: loadFloatArray(directory: path, file: "W_fc2"))
        b_fc2 = Tensor(shape: [3], elements: loadFloatArray(directory: path, file: "b_fc2"))
    }
}

private func loadFloatArray(directory directory: String, file: String) -> [Float] {
    let data = NSData(contentsOfFile: directory.stringByAppendingPathComponent(file))!
    return Array(UnsafeBufferPointer(start: UnsafeMutablePointer<Float>(data.bytes), count: data.length / 4))
}

##顔認識
iOSでの顔認識はCIDetectorで実装しました。カメラまわりをAVCaptureSessionで実装し、CMSampleBufferから得られたカメラ画像を次のようなコードで処理します。

Camera.swift

let ciImage = CIImage(CGImage: targetImage.CGImage!)    
let ciDetector = CIDetector(ofType: CIDetectorTypeFace, context: nil, options: [ CIDetectorAccuracy: CIDetectorAccuracyHigh ])

let features = ciDetector.featuresInImage(ciImage)    
for feature in features {
    let rect = feature.bounds
    // do something
}

OpenCVとCIDetectorでは顔の切り出し方に差があるのが注意点です。CIDetectorで得られたCGRectをそのまま使った場合、判定の精度が大きく落ちることがありました。CIFaceFeatureのleftEyePosition/rightEyePositionを使って切り出し位置を調整するか、もしくはiOS側でもOpenCV（同じCascadeClassifier）を使うほうがよいのかもしれません。

##画像判定

カメラ画像から顔を切り出したら判定器にかけます。判定結果として最も可能性の高いインデックスが得られます。

TFDetector.swift

import TensorSwift

class TFDetector {
 
    static let instance = TFDetector()
    
    private let classifier = Classifier(path: NSBundle.mainBundle().resourcePath!)
    
    private init() {
    }

    func detectImage(image: UIImage, inputSize: Int) -> Int {
        
        let input: Tensor

        do {
            var pixels = [UInt8](count: inputSize * inputSize * 4, repeatedValue: 0)

            let bitmapInfo = CGBitmapInfo(rawValue: CGBitmapInfo.ByteOrder32Big.rawValue | CGImageAlphaInfo.PremultipliedLast.rawValue)
            let context = CGBitmapContextCreate(&pixels, inputSize, inputSize, 8, inputSize * 4, CGColorSpaceCreateDeviceRGB(), bitmapInfo.rawValue)!
            let rect = CGRect(x: 0.0, y: 0.0, width: CGFloat(inputSize), height: CGFloat(inputSize))
            CGContextClearRect(context, rect)
            CGContextDrawImage(context, rect, image.CGImage!)

            let rgb = pixels.enumerate().filter { $0.0 % 4 != 3 }.map { Float($0.1) / 255.0 }

            input = Tensor(shape: [Dimension(inputSize), Dimension(inputSize), Dimension(3)], elements: rgb)
        }

        return classifier.classify(input)
    }
}

#テストアプリ

^*4

---

※1: Photo by Tommaso Fornoni Cropped by Danyele, CC 表示 3.0
※2: Photo by Vinod Divakaran, CC 表示 2.0
※3: Photo by Thomas Rodenbücher, CC 表示 2.0
※4: [Maxisport / Shutterstock.com](http://www.shutterstock.com/gallery-224068p1.html?cr=00&pl=edit-00">Maxisport / Shutterstock.com "Maxisport / Shutterstock.com")