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))
完全なコンバートスクリプトはこちらにもあります。
🐣
フリーランスエンジニアです。
お仕事のご相談こちらまで
rockyshikoku@gmail.com
Core MLやARKitを使ったアプリを作っています。
機械学習／AR関連の情報を発信しています。