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

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

<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/600791b2-c543-a57e-069b-8cf10792d6b9.jpeg" width=300> <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/9ce2c27a-d142-f165-43ef-e279fdc78c38.png" width=300>
<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/44729af8-e0b1-00db-df82-a10a217f1330.jpeg" width=300> <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/95493a57-e69c-2031-1aef-ccd7b9298caf.png" width=300>

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

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

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

###変換済みモデル
GitHubのCoreMLModelsにあります。

https://github.com/john-rocky/CoreML-Models

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

https://colab.research.google.com/drive/1jvFIY0rLQr0yQ_4uXwApuONvUjcw1A3-?usp=sharing

###オリジナルプロジェクト
[**xuebinqin/U-2-Net**](https://github.com/xuebinqin/U-2-Net)
[<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/f953fec3-5e0a-a04d-6d65-15119741ed7d.png" width=600>](https://github.com/xuebinqin/U-2-Net)

###論文

https://arxiv.org/abs/2005.09007

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

https://remove-background.net/

https://www.pixelmator.com/pro/

https://twitter.com/cyrildiagne/status/1256916982764646402

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

#変換手順

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

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

CoreMLToolsをインストール。

```shell
pip install coremltools
```

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

```shell
cd U-2-Net/
```

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

```shell
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
```

事前トレーニング済みの重みを取得。
[オリジナルプロジェクト](https://github.com/xuebinqin/U-2-Net)のGoogleDriveリンクから取得できます。
[176.3MB](https://drive.google.com/file/d/1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ/view?usp=sharing)
[4.7 MB](https://drive.google.com/file/d/1rbSTGKAE-MTxBYHd-51l2hMOQPT_7EPy/view?usp=sharing)
[ポートレイト用](https://drive.google.com/file/d/1IG3HdpcRiDoWNookbncQjeaPN28t90yW/view?usp=sharing)

U2Netモデルを初期化。

```shell
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()
```

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

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

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

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

https://github.com/xuebinqin/U-2-Net/issues/118


```shell
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)])
```

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

```shell
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モデルの出力をグレースケールの画像に変える。

```shell
#アクティベーション層を追加
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モデルが取得できる。
モデルを開くと以下のように表示される。
<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/7a03552e-40b4-c60f-230e-d92bb6c05937.png" width=400>

出力はout_p6までの７チャンネルある。
マスク画像として使えるのはout_p1です。
<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/ccec6fb0-5474-2e61-fa73-0f1aac5b31ec.png" width=400>

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

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

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

```shell
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)
```

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

```shell
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の変換後に出力層の出力名を調べる必要があります。


```shell
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フレームワークを使用する場合

```swift
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を使用する場合

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

```swift
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
```

どちらのケースも、inputとoutputは320*320のCVPixelBuffer。

<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/d4550333-7b64-7fff-9bbc-cc4240f679c6.png" width=300> <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/6f1a6c6e-dabb-f217-80d5-52a3a06a429f.png" width=300>
🐣
****

フリーランスエンジニアです。
お仕事のご相談こちらまで
簡単な開発内容をお書き添えの上、お気軽にご連絡ください。
rockyshikoku@gmail.com

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

[Twitter](https://twitter.com/JackdeS11)
[Medium](https://medium.com/@rockyshikoku)