Galapagos Tech Blog

株式会社ガラパゴスのメンバーによる技術ブログです。

CoreMLでリアルタイムなスタイル変換

これはGalapagos Advent Calendar 20日目の記事です。

二度目まして。iOSチームの高橋です。好きな金額は二兆円です。

今回はiOS上で簡単にニューラルネットのモデルを実行させられるCoreMLを利用して、リアルタイムなスタイル変換を実装する話をします。

準備

Kerasモデルファイルの入手

さて、リアルタイムなスタイル変換を行う手法としてはarXiv:1603.08155が存在しますが、なんと!昨日の記事でまんださんがこれをKerasで実装してくれています!(しらじら)

なので、できあがったh5モデルファイルをもらうことにしました。これさえあればモネ風のスタイル変換ができる、はずです。

mlmodelへの変換

Kerasのモデルファイルをもらったので、coremltoolsを使ってmlmodelファイルに変換します。 coremltoolsは現時点ではPython2系にしか対応していないようなので、しぶしぶPython2を使います。

$ pip install -U coremltools
>>> import coremltools
WARNING:root:Keras version 2.1.2 detected. Last version known to be fully compatible of Keras is 2.0.4 .
WARNING:root:TensorFlow version 1.4.1 detected. Last version known to be fully compatible is 1.1.1 .
>>> coreml_model = coremltools.converters.keras.convert('./monet2.h5', input_names='input_1', image_input_names='input_1', output_names='transform_output')

(中略)

ValueError: Unknown layer: InputNormalize

はい、失敗しましたね。今日のblogはここまでです。ありがとうございました。

……というわけにはゆかないので、今回はこれを乗り越えてみせたいと思います。

なぜなのか

なぜ失敗するのかというと、元になったKerasのモデルがKerasにはない独自のレイヤーを定義して使用しているからです1。 するとcoremltoolsは「そんなレイヤー知りません」という顔でエラーを投げて寄越すわけです。

# たとえばこういうレイヤーがあると変換に失敗する
class InputNormalize(Layer):
    def __init__(self, **kwargs):
        super(InputNormalize, self).__init__(**kwargs)

    def build(self, input_shape):
        pass

    def compute_output_shape(self,input_shape):
        return input_shape

    def call(self, x, mask=None):
        return x/255.

いまさら純Kerasで作り直してもらうわけにもゆかないので、これは困りました。

tf-coreml

ところで話は変わりますが、TensorFlowのモデルからmlmodelへ変換するスクリプトが実は存在します(いくらかの制限つきではありますが)。

github.com

ということは、KerasのモデルをTensorFlowのモデルに変換することができれば、この問題を乗り越えることができるかもしれません。 そしてKerasのモデルがTensorFlowをバックエンドとして利用しているならば、それは原理的には可能なはずです。

Keras to TensorFlow

そう思って検索をかけてみると、先行研究を発見することができます。 だよね〜、と思いながらモデルファイルを投入します。

$ python keras_to_tensorflow.py -input_model_file ./monet2.h5 -output_model_file ./monet2.pb
('input args: ', Namespace(f=None, graph_def=False, input_fld='.', input_model_file='../monet2.h5', num_outputs=1, output_fld='.', output_graphdef_file='model.ascii', output_model_file='../monet2.pb', output_node_prefix='output_node'))

(中略)

ValueError: Unknown layer: InputNormalize

コケてしまいました。やはり「そんなレイヤー知りません」ということのようです。えっじゃあKerasは独自に定義したレイヤーを含むモデルは配布できないってことですか?と思って検索すると、こういうコメントを発見しました:

model = keras.models.load_model('temp_model.h5',
              custom_objects={'Melspectrogram':kapre.time_frequency.Melspectrogram})

なるほど、load_model時にレイヤークラスを渡してやればいいようです。なのでkeras_to_tensorflow.pyをすこし修正して再挑戦。

net_model = load_model(weight_file_path, custom_objects={'InputNormalize': InputNormalize,
                                                         'ReflectionPadding2D': ReflectionPadding2D,
                                                         'Denormalize': Denormalize,
                                                         'VGGNormalize': VGGNormalize,
                                                         'dummy_loss': dummy_loss})
Converted 122 variables to const ops.
('saved the freezed graph (ready for inference) at: ', '././monet2.pb')

やった!これでTensorFlowのモデルファイルが手に入りました。

TensorFlow to CoreML

手に入ったTensorFlowモデルをtf-coremlに投入してやります。

>>> import tfcoreml
WARNING:root:Keras version 2.1.2 detected. Last version known to be fully compatible of Keras is 2.0.6 .
WARNING:root:TensorFlow version 1.4.1 detected. Last version known to be fully compatible is 1.2.1 .
>>> coreml_model = tfcoreml.convert(tf_model_path='monet2.pb', mlmodel_path='monet2.mlmodel', input_name_shape_dict={'input_1:0': [1, 256, 256, 3]}, output_feature_names=['transform_output/mul:0'], image_input_names=['input_1:0'])

(中略)

 Core ML model generated. Saved at location: monet2.mlmodel

Core ML input(s):
 [name: "input_1__0"
type {
  imageType {
    width: 256
    height: 256
    colorSpace: RGB
  }
}
]
Core ML output(s):
 [name: "transform_output__mul__0"
type {
  multiArrayType {
    shape: 3
    shape: 256
    shape: 256
    dataType: DOUBLE
  }
}
]

これでmlmodelファイルが手に入った……ように見えますが、よく見ると出力の型がmultiArrayTypeになっています。これだと画像として出力されないので、Swift側で変換コードを書いてやる必要があるのですが、それはイケてないなあと思っていたところ、AppleのDeveloper Forumに解決方法が書かれていました。

forums.developer.apple.com

こういう関数を定義して噛ませればよいようです:

def convert_multiarray_output_to_image(spec, feature_name, is_bgr=False):
    """
    Convert an output multiarray to be represented as an image
    This will modify the Model_pb spec passed in.
    Example:
        model = coremltools.models.MLModel('MyNeuralNetwork.mlmodel')
        spec = model.get_spec()
        convert_multiarray_output_to_image(spec,'imageOutput',is_bgr=False)
        newModel = coremltools.models.MLModel(spec)
        newModel.save('MyNeuralNetworkWithImageOutput.mlmodel')
    Parameters
    ----------
    spec: Model_pb
        The specification containing the output feature to convert
    feature_name: str
        The name of the multiarray output feature you want to convert
    is_bgr: boolean
        If multiarray has 3 channels, set to True for RGB pixel order or false for BGR
    """
    for output in spec.description.output:
        if output.name != feature_name:
            continue
        if output.type.WhichOneof('Type') != 'multiArrayType':
            raise ValueError("%s is not a multiarray type" % output.name)
        array_shape = tuple(output.type.multiArrayType.shape)
        channels, height, width = array_shape
        from coremltools.proto import FeatureTypes_pb2 as ft
        if channels == 1:
            output.type.imageType.colorSpace = ft.ImageFeatureType.ColorSpace.Value('GRAYSCALE')
        elif channels == 3:
            if is_bgr:
                output.type.imageType.colorSpace = ft.ImageFeatureType.ColorSpace.Value('BGR')
            else:
                output.type.imageType.colorSpace = ft.ImageFeatureType.ColorSpace.Value('RGB')
        else:
            raise ValueError("Channel Value %d not supported for image inputs" % channels)
        output.type.imageType.width = width
        output.type.imageType.height = height

実際に試してみます。

>>> spec = coreml_model.get_spec()
>>> convert_multiarray_output_to_image(spec, 'transform_output__mul__0')
>>> newModel = coremltools.models.MLModel(spec)
>>> newModel.save('monet2.mlmodel')

これで……

f:id:glpgsinc:20171218155353p:plain

できました!!!

あとはこのmlmodelファイルをスッとXcodeに投入すれば、画像の変換を行うクラスが生成されます。長かった…… 便利な世の中ですね。

実装

メインロジック

実装のメイン部分はCoreMLのおかげで非常に簡潔です。こんな風に:

import AVFoundation
import CoreImage
import UIKit

class ViewController: UIViewController {

    @IBOutlet private weak var imageView: UIImageView!

    let model = monet2()
    let imageSize = CGSize(width: 256, height: 256)

    /// メインロジック。CVPixelBufferを受け取って変換して表示する
    ///
    private func pixelBufferDidUpdate(pixelBuffer: CVPixelBuffer) {
        let output = try! model.prediction(input_1__0: pixelBuffer)
        let image = UIImage(ciImage: CIImage(cvPixelBuffer: output.transform_output__mul__0))
        DispatchQueue.main.async {
            self.imageView.image = image
        }
    }

    let session: AVCaptureSession = AVCaptureSession()
    let videoQueue: DispatchQueue = DispatchQueue(label: "videoqueue")
    var pixelBuffer: CVPixelBuffer? = nil

    override func viewDidLoad() {
        super.viewDidLoad()
        heavyLifting()
    }

}

で、残る問題は動画のキャプチャとピクセルバッファのクロップなのですが、それについては以下で簡単に説明したいと思います。

AVFoundationによる動画のキャプチャ

今回はリアルタイムなスタイル変換を行いたいので、カメラからの入力をリアルタイムに受け取る必要があります。 こういう場合にはAVFoundationを使えばよいらしいので、そのためのセットアップを行います2

extension ViewController: AVCaptureVideoDataOutputSampleBufferDelegate {
    func heavyLifting() {
        CVPixelBufferCreate(kCFAllocatorDefault,
                            Int(imageSize.width), Int(imageSize.height),
                            kCVPixelFormatType_32BGRA, nil, &pixelBuffer)

        let device = AVCaptureDevice.default(.builtInWideAngleCamera, for: .video, position: .back)!
        device.activeVideoMinFrameDuration = CMTime(value: 1, timescale: 30)

        let videoInput = try! AVCaptureDeviceInput(device: device)
        session.addInput(videoInput)

        session.sessionPreset = .hd1280x720

        let videoDataOutput = AVCaptureVideoDataOutput()
        videoDataOutput.videoSettings = [kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA]
        videoDataOutput.setSampleBufferDelegate(self, queue: videoQueue)
        videoDataOutput.alwaysDiscardsLateVideoFrames = true
        session.addOutput(videoDataOutput)

        let connection = videoDataOutput.connection(with: .video)
        connection?.videoOrientation = .portrait

        session.startRunning()
    }

冒頭のCVPixelBufferCreateは、CoreMLとのやりとりに使うピクセルバッファを作成するためのもので、AVFoundationとは関係ありません。 その下では背面カメラから動画の入力を受ける設定をしており、動画サイズは1280×720だという指定が続きます。最後に出力の受け先をselfに向けて撮影を開始しています。

CoreImageによるフレームの切り抜き

あとはキャプチャしたフレームを適切な大きさにクロップしてCVPixelBufferに書き込む必要がありますが、これはCoreImageを経由することで実現できるようです。筆者はCoreImageに明るくないのでけっこう苦労しましたが、Stack Overflowに助けられてなんとか画像の中央を切り出すことに成功しました。

    func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
        let ciImage = CIImage(cvImageBuffer: CMSampleBufferGetImageBuffer(sampleBuffer)!)
        let originalRect = ciImage.extent
        let cropRect = CGRect(x: (originalRect.width - imageSize.width) / 2,
                              y: (originalRect.height - imageSize.height) / 2,
                              width: imageSize.width, height: imageSize.height)
        let cropped = ciImage.cropped(to: cropRect)
        let transformed = cropped.transformed(by: CGAffineTransform(translationX: -cropRect.origin.x, 
                                                                    y: -cropRect.origin.y))
        let context = CIContext()
        context.render(transformed, to: pixelBuffer!)
        pixelBufferDidUpdate(pixelBuffer: pixelBuffer!)
    }

}

こうしてできたpixelBufferを上で書いたメインロジックに渡してやれば、CoreMLの力で画像が変換されるので、CIImageを経由してUIImageにしてUIImageViewに表示してやれば完成です。

結果

弊社エントランスにかざってあるクリスマスツリーをモネ風にしてみました。端末はiPhone Xを使用していますが、iPhone 7でも似たような速度で動きます。

f:id:glpgsinc:20171218185207g:plain

いかがでしょうか。モネ風かどうかはともかく、なにかしらスタイルが変換されていることがおわかりいただけるかと思います3。 関係ないですが、端末がまあまあ熱くなるので、無茶をしているのだなあという気持ちになります。

おわりに

ということで、iOS端末上でリアルタイムにスタイル変換をするアプリを簡単に作ることができました。次はこっちの論文なんかも実装してみたら楽しいかもしれませんね。

ところで

弊社では機械学習に興味があったりなかったりするエンジニアを募集しています。ご応募お待ちしております。

www.glpgs.com

以上です。明日はUIデザイナーのまあのんさんがProtoPieについて書いてくれるそうですね。お楽しみに。


  1. という記事を書いている間に判明したのですが、iOS 11.2以降は独自のレイヤーに対応させることが(がんばれば)可能なようですね(参考)。

  2. あっ、Info.plistのNSCameraUsageDescriptionを設定しておくのを忘れないでくださいね。

  3. もっとも、元からわりと派手なクリスマスツリーだという話はあるのですが。