!pip install turicreate==5.4
# The wrong version of MXNET will be installed
!pip uninstall -y mxnet
# Install CUDA10-compatible version of mxnet 
!pip install mxnet-cu100==1.4.0.post0

Style Transferを実装するためにコンテンツ画像とスタイル画像を用意する必要があります。ColaboratoryにはGoogle Driveを簡単にマウントできる機能がありますので、Google Driveのマイドライブに「content」と「style」というフォルダを作成しそちらに画像を格納していきます。

詳細は以下の流れになります。

Google Driveにcontentとstyleというフォルダを用意してそれぞれ画像を格納する
Colaboratory上でマウントする

from google.colab import drive
drive.mount('/content/drive')

画像を読み込む

content = tc.load_images('/content/drive/My Drive/content/')
styles = tc.load_images('/content/drive/My Drive/style/')

トレーニングを実行する

model = tc.style_transfer.create(styles, content)

GPU環境でも学習に1時間くらいかかりますので気長に待ちましょう。

学習が完了するとこんな画像が表示されるようになります。ジョブズくんと葛飾北斎の絵の画風が混ざって・・・

f:id:omuriceman:20190612191108j:plain:h150 + f:id:omuriceman:20190612191203j:plain:w150

こんな画風に変換されました！

f:id:omuriceman:20190612191228p:plain:w300

なんとなく趣を感じますね。

Colaboratory上で実行するコードはこちらにおいておきました。

github.com

Core MLにデプロイしてiOSアプリに組み込む

Colaboratoryで学習させたモデルをエクスポートして、iOSプロジェクトに組み込みを行います。

model.export_coreml('MyStyleTransfer.mlmodel')

Colaboratoryの画面の左カラムにモデルファイル(mlmodel)が出力されていますのでダウンロードしてください。

次にXCodeを立ち上げて、モデルファイルをドラック&ドロップします。

コードについてはGitHubをご覧いただくとして、要点を解説したいと思います。

let numStyles  = 8
let styleIndex = 3

let styleArray = try? MLMultiArray(shape: [numStyles] as [NSNumber], dataType: MLMultiArrayDataType.double)

for i in 0...((styleArray?.count)!-1) {
    styleArray?[i] = 0.0
}
styleArray?[styleIndex] = 1.0

numStyles:styleフォルダに入れたスタイル画像の数
styleIndex:適用するスタイル

上の公式サイトの画像を例に言うと、8枚のスタイル画像があるのでnumStylesは8、styleIndexが3の場合は0からカウントするので右上の薪のスタイルが適用されることになります。

画像にStyle Transferを適用するコードは以下の通りです。

// initialize model 
let model = MyStyleTransferModel()

// set input size of the model
let modelInputSize = CGSize(width: 600, height: 600)

// create a cvpixel buffer
var pixelBuffer: CVPixelBuffer?
let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,
             kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] as CFDictionary
CVPixelBufferCreate(kCFAllocatorDefault,
                    modelInputSize.width,
                    modelInputSize.height,
                    kCVPixelFormatType_32BGRA,
                    attrs,
                    &pixelBuffer)

// put bytes into pixelBuffer
let context = CIContext()
context.render(ciImage, to: pixelBuffer!)

// predict image
let output = try? model.prediction(image: pixelBuffer!, index: styleArray!)
let predImage = CIImage(cvPixelBuffer: (output?.stylizedImage)!)