これは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')
これで……
できました！！！
あとはこの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)
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端末上でリアルタイムにスタイル変換をするアプリを簡単に作ることができました。次はこっちの論文なんかも実装してみたら楽しいかもしれませんね。
ところで
弊社では機械学習に興味があったりなかったりするエンジニアを募集しています。ご応募お待ちしております。
www.glpgs.com
以上です。明日はUIデザイナーのまあのんさんがProtoPieについて書いてくれるそうですね。お楽しみに。
