アニメ変換をiPhoneでつかいたい -- CoreMLに変換すればOK

12491 ワード
Swift coremltools AnimeGAN iOS coreML iOS テキストリンク

AnimeGAN2-pytorchをiPhoneで使える形式に変換する方法です

変換ずみモデルを入手したい方はこちら：

アニメ変換機械学習モデルであるAnimeGAN2をiPhoneアプリで使えるようにします。

AnimeGANはかっこいい、、でもPythonかあ。。。

アプリケーションでPythonのモデルを使うには、サーバーに画像をアップして処理したり、、、と大変そう。

CoreMLに変換すればiPhone上でモデルを使える

CoreML形式なら、かんたんにAnimeGANをiPhoneアプリに組み込めます。

手順

pip install coremltools

Torch-Hubからモデルをロードします。

import torch

model = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="face_paint_512_v2")

モデルを変換します。

import coremltools as ct

example_input = torch.rand(1,3,512,512)
traced_model = torch.jit.trace(model, example_input)
mlmodel = ct.convert(traced_model, inputs=[ct.ImageType(name="input", shape=example_input.shape,bias=[-1,-1,-1],scale=2.0/255.0)])

モデルの出力形状は[1,3,512,512]です。
最初の１しかない次元は不要かつ、
画像出力にするには[3,512,512]の形状にしたいので、
最初の次元を削減する層を追加します。

from coremltools.proto import FeatureTypes_pb2 as ft
spec = mlmodel.get_spec()
builder = ct.models.neural_network.NeuralNetworkBuilder(spec=spec)
builder.add_squeeze(name="squeeze", input_name="var_444", output_name="squeeze_out", axes = None, squeeze_all = True)

モデルの入力は -1 ~ 1 の範囲に正規化されていますので、
0 ~ 255 の範囲の値に戻すアクティベーション層を追加します。

builder.add_activation(name="activation",non_linearity="LINEAR",input_name="squeeze_out",output_name="activation_out",params=[127.5,127.5])

モデルの出力名にアクティベーション層の出力名を明示的に指定します。
こうしないと、元のvar_444層からモデルの出力が出力されるままになってしまいますので、必要なステップです。

from coremltools.proto import FeatureTypes_pb2 as ft
builder.spec.description.output.pop()
builder.spec.description.output.add()
output = builder.spec.description.output[0]
output.name = "activation_out"

モデルの出力タイプに画像を指定します。

output.type.imageType.colorSpace = ft.ImageFeatureType.ColorSpace.Value('RGB')
output.type.imageType.width = 512 
output.type.imageType.height = 512

モデルを保存します。

from coremltools.models.utils import save_spec
save_spec(builder.spec, 'animegan2.mlmodel')

iOSでの使い方

Visionでリクエスト実行します。

let model = try! VNCoreMLModel(for: animegan().model)
let coreMLRequest = VNCoreMLRequest(model: model)

let handler = VNImageRequestHandler(ciImage: ciimage,options: [:])

try! handler.perform([coreMLRequest])

let result = request.results.first as! VNPixelBufferObservation
let uiimage = UIImage(ciImage: CIImage(cvPixelBuffer: result.pixelBuffer))


完全なコンバートスクリプトはこちらにもあります。

🐣

フリーランスエンジニアです。
お仕事のご相談こちらまで
[email protected]

Core MLやARKitを使ったアプリを作っています。
機械学習／AR関連の情報を発信しています。

Twitter
Medium