🐙

DeepLabをCore MLフォーマットに変換する

2021/12/28に公開

環境

  • OS: Ubuntu 20.04
  • pyenv 2.1.0

fronzen graphを作る

Pythonの環境を整えます。coremltoolsは5系だとエラーが出ます。tensorflowは1系を使うのですが、Pythonは3.7を使いました。

$ mkdir deeplab-coreml
$ cd deeplab-coreml
$ pyenv local 3.7.10
$ python -m venv .venv
$ poetry init
$ poetry add tensorflow=1.15.5
$ poetry add tf-slim=1.1.0
$ poetry add coremltools=4.1
$ poetry add jupyter matplotlib

frozen graphの置き場を作っておきます。

$ mkdir -p weights/deeplabv3

deeplabを取得とPYTHONPATHの設定をします。

$ git clone https://github.com/tensorflow/models.git
$ cd models
$ git checkout d7eabefa595f325611798a6d88d0f415986b8152
$ cd research
$ export PYTHONPATH=`pwd`:`pwd`/slim

deeplabの学習済みモデルを取得し、frozen graphを生成します。設定により多少変わりますが、バックボーンがmobilenetv2の場合とxception_65の場合を紹介します。

mobilenetv2
$ cd deeplab
$ mkdir checkpoint
$ curl http://download.tensorflow.org/models/deeplabv3_mnv2_pascal_train_aug_2018_01_29.tar.gz --output ./checkpoint/deeplabv3_mnv2_pascal_train_aug.tar.gz
$ tar xzvf checkpoint/deeplabv3_mnv2_pascal_train_aug.tar.gz -C checkpoint/
$  python export_model.py --model_variant="mobilenet_v2" \
  --checkpoint_path=./checkpoint/deeplabv3_mnv2_pascal_train_aug/model.ckpt-30000\
  --export_path=../../../weights/deeplabv3/frozen_graph.pb
xception_65
$ cd deeplab
$ mkdir checkpoint
$ curl http://download.tensorflow.org/models/deeplabv3_pascal_train_aug_2018_01_04.tar.gz --output checkpoint/deeplabv3_pascal_train_aug_2018_01_04.tar.gz
$ python export_model.py \
    --model_variant="xception_65" \
    --atrous_rates=6 \
    --atrous_rates=12 \
    --atrous_rates=18 \
    --output_stride=16 \
    --decoder_output_stride=4 \
    --num_classes=21 \
    --checkpoint_path=./checkpoint/deeplabv3_pascal_train_aug/model.ckpt \
    --export_path=../../../weights/deeplabv3/frozen_graph.pb

どのようなオプションを指定すべきかはTensorFlow DeepLab Model Zooの表とexport_model.pyのオプションのコメントを見ると大体分かると思います。

coremltoolsで変換

ここからcoremltoolsを使ってCore ML形式にfrozen graphを変換します。以下に登場するPythonのコード片はプロジェクトルートにjupyterノートブックを作ったと思ってください。

import coremltools as ct
import json

jsonモジュールは、XCodeから見たとき分かりやすいように設定するために使います。

model = ct.convert('./weights/deeplabv3/frozen_graph.pb', inputs=[ct.ImageType(shape=(1, 513, 513, 3))])

WARNING:root:Output var SemanticPredictions of type i32 in function main is cast to type fp32という警告が出ますが、後でint32に修正するので大丈夫です。画像を入力にするモデルの変換は大抵ct.ImageTypeを使います。shapeのはじめの1はバッチ数なので、ここはいつでも1だと思って良いと思います。

spec = model.get_spec()
builder = ct.models.neural_network.NeuralNetworkBuilder(spec=spec)

変換には特にbuilderは要らないですが、入力や出力の状態を確認するのに使います。

builder.inspect_input_features()
out
[Id: 0] Name: ImageTensor
          Type: imageType {
  width: 513
  height: 513
  colorSpace: RGB
}

RGBの画像になってますね。でも名前が気に入らないので、変更します。

ct.utils.rename_feature(spec, "ImageTensor", "image")
builder.inspect_input_features()
out
[Id: 0] Name: image
          Type: imageType {
  width: 513
  height: 513
  colorSpace: RGB
}

inputのnameがimageに変わりました。次に出力を確認します。

builder.inspect_output_features()
out
[Id: 1] Name: SemanticProbabilities
          Type: multiArrayType {
  dataType: FLOAT32
}

[Id: 0] Name: SemanticPredictions
          Type: multiArrayType {
  dataType: FLOAT32
}

出力が2種類ありますね。SemanticProbabilitiesは実際には役立つこともあるかもしれませんが、XCodeのプレビューでこれを残していると頻繁にXCodeが落ちます。真剣にアプリを作る段階で、扱いを再検討することにして、これを削除し、SemanticPredictionsはint32の513x513のMultiArrayにし、名前をsemanticPredictionsに変更します。

spec.description.output.pop()
output = spec.description.output[0]
ct.utils.rename_feature(spec, "SemanticPredictions", "semanticPredictions")
output.type.multiArrayType.dataType = ct.proto.FeatureTypes_pb2.ArrayFeatureType.INT32
output.type.multiArrayType.shape.append(513)
output.type.multiArrayType.shape.append(513)
builder.inspect_output_features()
out
[Id: 0] Name: semanticPredictions
          Type: multiArrayType {
  shape: 513
  shape: 513
  dataType: INT32
}

出力についても修正できたので、メタ情報を充実させます。

spec.description.input[0].shortDescription = "Input image to be segmented"
spec.description.output[0].shortDescription = "Array of integers of the same size as the input image, where each value represents the class of the corresponding pixel."
spec.description.metadata.versionString = "--"
spec.description.metadata.shortDescription = "DeepLab is a state-of-art deep learning model for semantic image segmentation, where the goal is to assign semantic labels (e.g., person, dog, cat and so on) to every pixel in the input image."
spec.description.metadata.author = "Original Paper: Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, Liang-Chieh Chen. DeepLabV3+: Liang-Chieh Chen and Yukun Zhu and George Papandreou and Florian Schroff and Hartwig Adam."
spec.description.metadata.license = "Please see https://github.com/tensorflow/tensorflow for license information, and https://github.com/tensorflow/models/tree/master/research/deeplab for the original source of the model."

最後にプレビューが出るように設定します。

model = ct.models.MLModel(spec)
labels_json = {"labels": [
    'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
    'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
    'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tv'
]}
model.user_defined_metadata["com.apple.coreml.model.preview.type"] = "imageSegmenter"
model.user_defined_metadata['com.apple.coreml.model.preview.params'] = json.dumps(labels_json)

保存します。

model.save("MyDeepLabV3.mlmodel")

最後にMacに移動させてXCodeで動作確認してみます。

写真は家族-ビーチ-人-海洋-6398107を使用しました。

Discussion