More than 1 year has passed since last update.

iOSAdvent Calendar 2021Day 11

エンジニアと人生コミュニティ

U2NetをCoreMLに【iPhoneで機械学習セグメンテーション】

Last updated at 2021-12-10Posted at 2021-12-10

高精度のセグメンテーションモデルをモバイル端末で

U2Netは画像内の顕著なオブジェクトと背景を分離してくれる機械学習モデル。

人など特定のオブジェクトに対応したセグメンテーションモデルはさまざまあるが、
このU2Netは、画像内で一番目立つオブジェクトをセグメンテーションしてくれるので、利用範囲が広い。
高精度でいろんなアプリに使われている。

もともとpythonで書かれたモデル。
これをCoreMLに変換することで、iOSでオンデバイスで使うことができる。
iPhoneのチップだけで実行できる。
しかも、かなり高速。
精度もオリジナルモデルと同じである。

変換済みモデル

GitHubのCoreMLModelsにあります。

変換コードのGoogle Colabノートブックデモ

オリジナルプロジェクト

xuebinqin/U-2-Net

論文

U2Netが利用されているアプリ

176.3MBのバージョンと4.6MBの軽量バージョンがあります。
以下では、176.3MBバージョンの変換手順を示しますが、
軽量版の場合に変更する数カ所もコメントアウトして書いています。

変換手順

オリジナル・プロジェクトをクローン。

git clone https://github.com/xuebinqin/U-2-Net.git

CoreMLToolsをインストール。

pip install coremltools

プロジェクトディレクトリに移動。

cd U-2-Net/

必要なモジュールをインポート。

from model import U2NET
# 軽量版は：
# from model import U2NETP
import coremltools as ct
from coremltools.proto import FeatureTypes_pb2 as ft
import torch
import os
from PIL import Image
from torchvision import transforms

事前トレーニング済みの重みを取得。
オリジナルプロジェクトのGoogleDriveリンクから取得できます。
176.3MB
4.7 MB
ポートレイト用

U2Netモデルを初期化。

net = U2NET(3,1)
# 軽量版は：
# net = U2NETP(3,1)
device = torch.device('cpu')

###Path To Model Directory.
model_dir = os.path.join(os.getcwd(), '/content/drive/MyDrive', "u2net" + '.pth') 
# 軽量版はu2netp.pthのパス

net.load_state_dict(torch.load(model_dir, map_location=device))
net.cpu()
net.eval()

ダミー・インプットを作る。実際の画像でも大丈夫。

example_input = torch.rand(1,3,320,320)

変換。
モデルに適したカラーチャネル、数値幅に入力バイアスを設定する。

（変換のプロセスで以下のイシューを参照しました。）

traced_model = torch.jit.trace(net, example_input)
model = ct.convert(traced_model, inputs=[ct.ImageType(name="input", shape=example_input.shape,bias=[-0.485/0.229,-0.456/0.224,-0.406/0.225],scale=1.0/255.0/0.226)])

モデルのメタデータを加える。

model.short_description = "U2-Net: Going Deeper with Nested U-Structure for Salient Object Detection"
model.license = "Apache 2.0"
model.author = "Qin, Xuebin and Zhang, Zichen and Huang, Chenyang and Dehghan, Masood and Zaiane, Osmar and Jagersand, Martin"

モデルの末尾に新しいアクティベーション層を加え、
CoreMLモデルの出力をグレースケールの画像に変える。

#アクティベーション層を追加
spec = model.get_spec()
spec_layers = getattr(spec, spec.WhichOneof("Type")).layers
output_layers = []
for layer in spec_layers:
    if layer.name[:2] == "25":
        print("name: %s  input: %s  output: %s" % (layer.name, layer.input, layer.output))
        output_layers.append(layer)
new_layers = []
layernum = 0;
for layer in output_layers:
    new_layer = spec_layers.add()
    new_layer.name = 'out_p'+str(layernum)
    new_layers.append('out_p'+str(layernum))

    new_layer.activation.linear.alpha=255
    new_layer.activation.linear.beta=0

    new_layer.input.append('var_'+layer.name)
    new_layer.output.append('out_p'+str(layernum))
    output_description = next(x for x in spec.description.output if x.name==output_layers[layernum].output[0])
    output_description.name = new_layer.name

    layernum = layernum + 1

# 出力をグレースケールの画像に.
for output in spec.description.output: 
    if output.name not in new_layers: 
        continue
    if output.type.WhichOneof('Type') != 'multiArrayType': 
        raise ValueError("%s is not a multiarray type" % output.name) 
    output.type.imageType.colorSpace = ft.ImageFeatureType.ColorSpace.Value('GRAYSCALE')
    output.type.imageType.width = 320 
    output.type.imageType.height = 320

# モデルを保存。
updated_model = ct.models.MLModel(spec)
updated_model.save("u2net.mlmodel")

ここまででCoreML形式のU2Netモデルが取得できる。
モデルを開くと以下のように表示される。

出力はout_p6までの７チャンネルある。
マスク画像として使えるのはout_p1です。

変換の留意点（上記スクリプトでうまくいった場合読まなくて大丈夫）

これは、モデルを変換した時に個人的に試行錯誤した点です。
基本的にこのステップは必要ありません。単なる留意点のメモです
上記のコードで変換できます。
ただ、オリジナルモデル自体のレイヤーの名前が変わったり、構造が変わったりした際に、
アクティベーション層を追加するプロセスがうまくいかない場合があります。
その時は、末尾の層の名前を調べて変換コードを変える必要があります。

上記コードの中の下記の部分で元のモデルの出力層を取得するのですが、
出力層を探すために、２５００番台の名前のついたレイヤーを探しています。

output_layers = []
for layer in spec_layers:
    if layer.name[:2] == "25":
        print("name: %s  input: %s  output: %s" % (layer.name, layer.input, layer.output))
        output_layers.append(layer)

そして、出力層のあとに新しいアクティベーション層をつけています。
ここで、元の出力層の出力名とアクティベーション層の入力名を合わせる必要があります。

for layer in output_layers:
    new_layer = spec_layers.add()
    new_layer.name = 'out_p'+str(layernum)
    new_layers.append('out_p'+str(layernum))

    new_layer.activation.linear.alpha=255
    new_layer.activation.linear.beta=0

    new_layer.input.append('var_'+layer.name)
    new_layer.output.append('out_p'+str(layernum))

ただ、出力層の名前の２５００番台というのは変わる可能性があります。
ですので、上記の変換スクリプトでうまくいかない場合、
オリジナルモデル→CoreMLの変換後に出力層の出力名を調べる必要があります。

spec = model.get_spec()
spec_layers = getattr(spec, spec.WhichOneof("Type")).layers
for layer in spec_layers:
  print(layer.name,layer.input,layer.output)

2524 ['input_1'] ['var_2524']
2525 ['tar'] ['var_2525']
2526 ['d2'] ['var_2526']
2527 ['d3'] ['var_2527']
2528 ['d4'] ['var_2528']
2529 ['d5'] ['var_2529']
2530 ['d6'] ['var_2530']

このように最後の層までレイヤーの名前とレイヤーの入出力名がプリントされるので、その名前でアクティベーション層の入力名を合わせてください。

Xcodeプロジェクトでモデルを使用する

プロジェクトでの使い方には２つ選択肢がある。

１、CoreMLフレームワークを使用する場合

guard let model = try? u2net.init() else {fatalError("model initialize error")}
guard let result = try? model.prediction(input: pixelBuffer) else {fatalError("inference error")}
let output = result.out_p1
// inputとoutputは320*320のCVPixelBuffer

２、Visionを使用する場合

入力画像サイズの自動調整など便利なので、こちらがおすすめかも。

guard let model = try? VNCoreMLModel(for: u2net(configuration: MLModelConfiguration()).model) else { fatalError("model initialization failed") }
let coreMLRequest = VNCoreMLRequest(model: model)
let handler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer, options: [:])
do {
    try handler.perform([coreMLRequest])
    guard let result = coreMLRequest.results?.first as? VNPixelBufferObservation else { return nil }
    let output = result.pixelBuffer
} catch let error {
    fatalError("inference error \(error)")
}
// inputとoutputは320*320のCVPixelBuffer