11
8

More than 3 years have passed since last update.

UGATITをiOSで使う。セルフィーをアニメに変換

Last updated at Posted at 2020-06-22

CoreMLモデルに変換して、iPhoneやiPadでUGATITを使用できます。

UGATITはGANの応用で、深層画像生成技術の最先端です。
スクリーンショット 2020-05-19 11.09.03.png
たとえば、自撮りをアニメに変換できます。

*論文
* GitHubプロジェクトページ
* CoreMLに変換済のセルフィー2アニメモデル。ダウンロードするとドラッグ&ドロップでxcodeプロジェクトでつかえます。

1、上のGitHubプロジェクトページからクローン。

git clone https://github.com/taki0112/UGATIT.git
cd UGATIT

2、上のGitHubプロジェクトページから事前トレーニング済みモデルのチェックポイントをダウンロードします。
そして、手元のUGATITディレクトリ配下に「チェックポイント」ディレクトリを作成し、ダウンロード・解凍したフォルダを配置します。
1_nQK2z_Hwra1Guj4LD99gtA.png

mkdir checkpoint
## このディレクトリにダウンロードしたチェックポイントをおく

3、上のGitHubプロジェクトページからselfie2animeデータセットをダウンロードします。
そして、手元のUGATITディレクトリ配下に「データセットディレクトリ」を作成し、ダウンロードしたデータセットフォルダを配置します。

スクリーンショット 2020-06-22 3.01.48.png

mkdir dataset 
##ダウンロードしたselfie2animeデータセットをこのディレクトリに配置します

4、モデルのpbtxtを作成します。
そのためには、手元のUGATIT.pyの642行目(つまり、「def test(self):)」にwrite_graphファンクションを挿入してから、test.pyを実行します。

UGATIT.py
## def test(self):
##...
fake_img = self.sess.run(self.test_fake_B, feed_dict = {self.test_domain_A : sample_image})
tf.io.write_graph(self.sess.graph_def, './', 'ugatit.pbtxt') 
##↑この行を挿入します。
python main.py --dataset selfie2anime --phase test 
##テストが成功すると、現在のディレクトリに「ugatit.pbtxt」が表示されます。

5、tfcoremlをインストールします。

pip install tfcoreml

6、frozen_modelを作成します。そのためには、「convert.py」を書いて実行します。

convert.py
from __future__ import print_function
import numpy as np
from tensorflow.python.tools.freeze_graph import freeze_graph
import tfcoreml

graph_def_file = 'ugatit.pbtxt'
checkpoint_file = 'checkpoint/UGATIT_selfie2anime_lsgan_4resblock_6dis_1_1_10_10_1000_sn_smoothing/UGATIT.model-1000000'
frozen_model_file = './frozen_model.pb'
output_node_names = 'generator_B/Tanh'

freeze_graph(input_graph=graph_def_file,
             input_saver="",
             input_binary=False,
             input_checkpoint=checkpoint_file,
             output_node_names=output_node_names,
             restore_op_name="save/restore_all",
             filename_tensor_name="save/Const:0",
             output_graph=frozen_model_file,
             clear_devices=True,
             initializer_nodes="")
python convert.py

7、frozen_modelからCoreMLモデルに変換します。そのために、「coreml.py」を書いて実行します。

coreml.py
input_tensor_shapes = {'test_domain_A':[1, 256, 256, 3]} # batch size is 1
# Output CoreML model path
coreml_model_file = './ugatit.mlmodel'
output_tensor_names = ['generator_B/Tanh:0']
# Call the converter
coreml_model = tfcoreml.convert(
        tf_model_path='frozen_model.pb',
        mlmodel_path=coreml_model_file,
        input_name_shape_dict=input_tensor_shapes,
        output_feature_names=output_tensor_names,
        image_input_names='test_domain_A',
        red_bias=-1,
        green_bias=-1,
        blue_bias=-1,
        image_scale=2/255,
        minimum_ios_deployment_target='12')

変換に成功すると、次のようなUGATIT coremlモデルが得られます。
スクリーンショット 2020-06-22 4.10.54.png

これで、iOSプロジェクトでUGATITを使用できます。

CoreGANContainerを使うと、モデルをプロジェクトにドロップするだけで使えます。


import Vision
lazy var coreMLRequest:VNCoreMLRequest = {
   let model = try! VNCoreMLModel(for: ugatit().model)
   let request = VNCoreMLRequest(model: model, completionHandler: self.coreMLCompletionHandler0)
   return request
   }()

let handler = VNImageRequestHandler(ciImage: ciimage,options: [:])
   DispatchQueue.global(qos: .userInitiated).async {
   try? handler.perform([coreMLRequest])
}

結果のmultiArrayを画像として視覚化するには、Hollance氏のCoreML Helpersが非常に便利です。
CoreMLヘルパーを使用したMultiArrayからImageへの変換。

func coreMLCompletionHandler0(request: VNRequest?, error: Error?) {
   let = coreMLRequest.results?.first as! VNCoreMLFeatureValueObservation
   let multiArray = result.featureValue.multiArrayValue
   let cgimage = multiArray?.cgImage(min: -1, max: 1, channel: nil, axes: (3,1,2))
   ...

ちなみに、私はselfie2animeのUGATIT CoreMLモデルでiOSアプリを作りました。

たのしい画像生成を!

Twitterフォローしてくださいお願いします!
https://twitter.com/JackdeS11

11
8
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
11
8