これは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へ変換するスクリプトが実は存在します（いくらかの制限つきではありますが）。
ということは、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に解決方法が書かれていました。
こういう関数を定義して噛ませればよいようです：
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')
これで……
できました！！！
あとはこの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でも似たような速度で動きます。
いかがでしょうか。モネ風かどうかはともかく、なにかしらスタイルが変換されていることがおわかりいただけるかと思います3。 関係ないですが、端末がまあまあ熱くなるので、無茶をしているのだなあという気持ちになります。
おわりに
ということで、iOS端末上でリアルタイムにスタイル変換をするアプリを簡単に作ることができました。次はこっちの論文なんかも実装してみたら楽しいかもしれませんね。
ところで
弊社では機械学習に興味があったりなかったりするエンジニアを募集しています。ご応募お待ちしております。
以上です。明日はUIデザイナーのまあのんさんがProtoPieについて書いてくれるそうですね。お楽しみに。
