easyportraitというリポジトリのPytorchモデルを使います。
MMSegmentationというセグメントモデルフレームワークを採用しています。
Core ML に変換
Swiftで扱うにはPytorchモデルからCore ML形式に変換します。
pytorchモデルインスタンスを作る
変換するためにモデルをインスタンス化します。
モデルインストール
まずはeasyportraitと必要なライブラリをインストールします。
以下の記事の方法でインストールしてください。
モデルインスタンスの初期化
from mmseg.apis import inference_segmentor, init_segmentor, show_result_pyplot
from mmseg.core.evaluation import get_palette
device = 'cuda:0'
checkpoint = "segformer-fp-512.pth"
config = "local_configs/easyportrait_experiments_v2/segformer-fp/segformer-fp.py"
model = init_segmentor(config, checkpoint, device=device)
ここでは、いくつかあるモデルのうちsegformer-fp-512というモデルを使っています。
MMSegmentでは、configファイルとcheckpointファイルでモデルの種類を指定して初期化するようになっています。
ここを変更すれば他のモデルを使えます。
ラップモデルを作る
ただ、MMSegmetationのモデルクラスは入力に単なる画像ではなくメタ情報を含んだdictを渡す必要があるので、そのままでは変換できません。
Core MLでそういうややこしい入力を扱うのはめんどくさいです。
それに、swiftで使うときは画像ファイル名や画像ファイルサイズといったメタ情報をわざわざモデルに渡す必要もありません。ただ画像だけでいいです。
なので、オリジナルモデルのラップクラスを作って、都合のいいモデルを作ります。
具体的にはただ画像を与えるだけでいいモデルを作ります。
元のモデルを見てみると、EncoderDecoder（MMEngineのBaseModule(torch.nn.Module)を継承したクラス）というクラスになっていて、実行部分はencode_decodeという関数です。そのあとで、メタ情報を踏まえて後処理を加えているようです。
なので、モデルの実行部分のencode_decodeという関数だけを使うようにします。
メタ情報には、Noneを与えるようにデフォルト設定することで、CoreMLEasyPortraitクラスの入力は画像だけで良くなります。
import torch
class CoreMLEasyPortrait(torch.nn.Module):
def __init__(self, model):
super(CoreMLEasyPortrait, self).__init__()
self.model = model
def forward(self, image):
out = self.model.encode_decode(image,None) # add None to img_meta
return out
wrap_model = CoreMLEasyPortrait(model).eval()
ラップモデルに後処理を加える
このままではモデルの出力をセグメントマットにできません。
後処理として、出力をソフトマックスで分布させ、argmaxでクラスインデックスの出力値を持つ１チャネル2次元配列（白黒画像みたいなもの）にしなくてはいけません。
なので、後処理のsoftmaxとargmaxをラップクラスに付け加えます。
import torch
import torch.nn.functional as F
class CoreMLEasyPortrait(torch.nn.Module):
def __init__(self, model):
super(CoreMLEasyPortrait, self).__init__()
self.model = model
def forward(self, image):
out = self.model.encode_decode(image,None)
out = F.softmax(out,dim=1)
out = out.argmax(dim=1)
return out
wrap_model = CoreMLEasyPortrait(model).eval()
出力を画像にする
Swiftでセグメントマットを画像として扱うことが想定されるので、出力を画像にします。
現在のところ、ラップモデルの出力は(512,512)の2次元配列で、顔のパーツに対応する位置のピクセルが0~8のインデックスになるようになっています。
0は背景、1は肌、2は左眉毛。。。といった具合です。
これをこのまま画像にしても、ほぼ真っ黒になるだけです。
セグメントマスクとして使うには、各インデックスごとに、対応する位置を真っ白（２５５値）にしたパーツのマスク画像を８つ作ることで、swiftで扱いやすくなると思います。
なので、
出力を9枚のGrayScale画像にします。
嬉しいことに、出力を１つから9つに増やしても、Core MLモデルの実行時間はほぼ変わりません。
これを行うには、各インデックス値に対応するピクセルを１にして255倍した２次元配列を9つ作ります。
また、Core MLの画像出力はバッチ数の次元も必要なため、出力に次元を追加しておきます。
これらの画像配列を配列に格納して返せば9つの出力が得られます。
import torch
import torch.nn.functional as F
class CoreMLEasyPortrait(torch.nn.Module):
def __init__(self, model):
super(CoreMLEasyPortrait, self).__init__()
self.model = model
def forward(self, image):
out = self.model.encode_decode(image,None)
out = F.softmax(out,dim=1)
out = out.argmax(dim=1)
out = out.unsqueeze(0)
out0 = (out == 0).float() * 255
out1 = (out == 1).float() * 255
out2 = (out == 2).float() * 255
out3 = (out == 3).float() * 255
out4 = (out == 4).float() * 255
out5 = (out == 5).float() * 255
out6 = (out == 6).float() * 255
out7 = (out == 7).float() * 255
out8 = (out == 8).float() * 255
return [out0,out1,out2,out3,out4,out5,out6,out7,out8]
wrap_model = CoreMLEasyPortrait(model).eval()
前処理の確認
Pytorchモデルの画像入力はそのままの0〜255のピクセル値ではなく、小さな値に正規化されています。
MMSegmentationでは、base/datasets/のデータconfigファイルのtestpipelineに前処理の方法が記載されています。
img_norm_cfg = dict(
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
RGB画像をmeanとstdを用いて正規化しています。
これは画像前処理でよく使われる値です。確かImageNetデータセットの正規化が元だと思う。
変換にはこの前処理値を与えます。
変換
import coremltools as ct
import torch
ex = torch.rand(1, 3, 512, 512).cuda()
ts = torch.jit.trace(wrap_model.cuda(), ex)
mlpackage = ct.convert(ts,
inputs=[ct.ImageType(name="image",
shape=ex.shape,
bias=[-0.485/0.229,-0.456/0.224,-0.406/0.225],scale=1.0/255.0/0.226)],
outputs=[ct.ImageType(name="out0",color_layout=ct.colorlayout.GRAYSCALE),ct.ImageType(name="out1",color_layout=ct.colorlayout.GRAYSCALE),ct.ImageType(name="out2",color_layout=ct.colorlayout.GRAYSCALE),ct.ImageType(name="out3",color_layout=ct.colorlayout.GRAYSCALE),ct.ImageType(name="out4",color_layout=ct.colorlayout.GRAYSCALE),ct.ImageType(name="out5",color_layout=ct.colorlayout.GRAYSCALE),ct.ImageType(name="out6",color_layout=ct.colorlayout.GRAYSCALE),ct.ImageType(name="out7",color_layout=ct.colorlayout.GRAYSCALE),ct.ImageType(name="out8",color_layout=ct.colorlayout.GRAYSCALE)])
mlpackage.save('easyportrait.mlpackage')
このように、scaleとbiasとして入力の前処理をcore ml モデルが行うようにし、
出力は画像タイプの出力を９つ持つようにしています。
これで、Core ML モデルができました。
Swiftで使う
VisionフレームワークでCore ML モデルを実行します。
var partsIndex = 0 // 一つ目のパーツ画像を取得する
let easyPortraitModel = try! easyportrait(configuration: MLModelConfiguration()).model
let vnEasyPortraitModel = try! VNCoreMLModel(for: easyPortraitModel)
let easyPortraitRequest = VNCoreMLRequest(model: vnEasyPortraitModel)
let handler = VNImageRequestHandler(ciImage: ciImage!)
try! handler.perform([request])
guard let result = request.results?[partsIndex] as? VNPixelBufferObservation else {fatalError()}
let resultCIImage = CIImage(cvPixelBuffer: result.pixelBuffer)
これで顔パーツのセグメントマスク画像が取れます。
結果の配列に各パーツの画像があるので、subscriptingでアクセスします。
変換済みモデル
Swiftサンプルプロジェクト
🐣
フリーランスエンジニアです。
AIについて色々記事を書いていますのでよかったらプロフィールを見てみてください。
もし以下のようなご要望をお持ちでしたらお気軽にご相談ください。
AIサービスを開発したい、ビジネスにAIを組み込んで効率化したい、AIを使ったスマホアプリを開発したい、
ARを使ったアプリケーションを作りたい、スマホアプリを作りたいけどどこに相談したらいいかわからない…
いずれも中間コストを省いたリーズナブルな価格でお請けできます。
お仕事のご相談はこちらまで
rockyshikoku@gmail.com
機械学習やAR技術を使ったアプリケーションを作っています。
機械学習／AR関連の情報を発信しています。