画像にmatplotlibでバウンディングボックスとラベルを追加する

物体検出のテストするときによく画像の上にbounding boxとラベルを表示したくなるんだけど、 毎回同じようなコード書いててめんどくさいのでスクリプトを残しておく。

from typing import List, Tuple

import matplotlib.patches as patches
import numpy as np
from PIL import Image


def add_bboxes_to_image(ax, image: np.ndarray,
                        bboxes: List[Tuple[int, int, int, int]],
                        labels: List[str] = None,
                        label_size: int = 10,
                        line_width: int = 2,
                        border_color=(0, 1, 0, 1)) -> None:
    """
    Add bbox to ax

    :param image: dtype=np.uint8
    :param bbox: [(left, top, right, bottom)]
    :param label: List[str] or None
    :return: ax
    """
    # Display the image
    ax.imshow(image)

    if labels is None:
        labels = [None] * len(bboxes)

    for bbox, label in zip(bboxes, labels):
        # Add bounding box
        top, left, bottom, right = bbox
        rect = patches.Rectangle((left, top), right - left, bottom - top,
                                 linewidth=line_width,
                                 edgecolor=border_color,
                                 facecolor='none')
        ax.add_patch(rect)

        # label
        if label:
            bbox_props = dict(boxstyle="square,pad=0",
                              linewidth=line_width, facecolor=border_color,
                              edgecolor=border_color)
            ax.text(left, top, label,
                    ha="left", va="bottom", rotation=0,
                    size=label_size, bbox=bbox_props)
    return ax

使い方

from PIL import Image

image = np.array(Image.open('chick.jpg'))
bboxes = [(20, 130, 280, 280), (0, 0, 100, 100)]
fig, ax = plt.subplots()
add_bboxes_to_image(ax, np.uint8(image), bboxes, ['chick', '?'])

f:id:cocodrips:20200504164219p:plain

Annotations APIを使うと左上にあるラベルみたいなのを簡単に追加することができる。 matplotlib.org

今まで知らなくてboxとちょっとずらして四角とラベル追加して...とかやってた(´・ω・`) 四角以外にも矢印とかいろんなかたちでannotationを追加できる。

ひよこ画像: Serious Chick | Marji Beach | Flickr

TensorflowモデルをTFLiteにconvertする方法とconvertツールの比較

カスタムのTensorflowのモデルをTFLiteにconvertしようとしてすごく辛かったのではまりどころを記録していく。

サンプルにあるモデルをtfliteにconvertするのはそんなに難しくないんだが、ちょっと自分で手を加えたモデルをconvertしようとしたらTensorFlow初心者の私にはものすごく大変だった。

今回使うのはtensorflow/modelsのslimに入ってるモデル。 mobilenet、inception、resnet、vggなどのアーキテクチャが準備されている。 Kerasモデルのconvertはやったことがないので今回の話にはでてこない。

※最新のTensorflowのversionは2.2だが、tensorflow/modelsに準備されているモデルはTensorflow2には対応していない
※今回書いているPythonコード内のTensorfFlowはすべて2.2、Pythonは3.7.6
※コマンドラインツールやtocoは明記して別のバージョンを使っている

TFLiteのモデルにconvertするツールの比較

convertに使うツールいくつか候補がある。
推奨されてるかどうかではなくて自分が使いやすいと思った順にかいていく 。多分推奨と逆の順番だけど。

toco

tocoは後述するtflite_convertやPython APIの元になっているプロジェクトで、元になってるだけあってオプションが一番多い。
自由自在にconvertしたいならこれが一番使いやすかった。

bazelで自分でbuildして使う。 これを採用してconvertする方法を後述するので詳細な使い方はそちらで。

tflite_convert コマンド

tocoをtensorflowがwrapしたコマンドラインツールがtflite_convert。 tocoとそれほど変わらないが、tensorflowによってオプション名が変わっていたりする気がする。

公式は細かいこと設定したいならtflite_convert使ってねと言っているが、tf2系を入れていると、オプションが少なくて使えない気がする。 このコマンドを使うなら、tf1系の環境を用意したほうが良い。

tflite_convertの元コードを見ると1系のオプション2系のオプションの違いがわかる。

PythonAPI

ドキュメントに推奨と書いてあるのがPython API

v2のtf.lite.TFLiteConverter

推奨してる方法。saved_modelには対応してるけどfrozen GraphDefs形式の変換には対応していない。

converter = tf.lite.from_saved_model(saved_model_path)
converter.convert()

でsaved_modelのconvertをできるのでめっちゃ簡単(そうにみえる)。

でもtf.lite.TFLiteConverter.from_saved_modelのヘルプを見るとわかるけど オプション(引数)が以下の3つしかない。

  • saved_model_dir
  • signature_keys
  • tags

特に

  • inputのレイヤー・outputのレイヤーが指定
  • Quantization-aware training(inference_type、post_training_quantizeなど)のオプション

がないので推奨している割に一番使えない方法だと思う・・・。

v1のTFLiteConverter(tf.compat.v1.lite.TFLiteConverter)

Python APIはもとからレイヤーの指定をするオプションがないのかといったらそんなことはなくてv1のAPIのドキュメントにも書いてあるように、v1APIはv2APIに追加で以下の引数を渡すことができる。

  • input_arrays
  • input_shapes
  • output_arrays

こんなかんじに使える

converter = tf.compat.v1.lite.TFLiteConverter.from_saved_model(
    saved_model_path,
    input_shapes={"input_layer_name": [batch_size, width, height, channel]}
)
converter.convert()

ツールの紹介は以上にして、convertの仕方について書いていく。

Checkpointからのconvertする

Checkpointからconvertするには以下のような手順を踏む。 f:id:cocodrips:20200411182334p:plain:w400

  1. GraphDefファイルを準備
  2. モデルをFrozenする
  3. モデルをTFLiteにconvert

今回はslim/netsの中にあるmobilenet_v2のmobilenet_v2_1.0_224_ckptをDLして、tocoを使ってconvertしていく。

slimの実装じゃなくても、GraphDefファイルがあればこの方法で問題ないと思う。

上記データにはfrozenGraphDef(frozen.pb)もtfliteモデルもはいってるけど、カスタムモデルを想定してこれを作るところからやっていく。

必要なリポジトリとライブラリのダウンロード

tensorflow/tensorflow

tocoはコマンドとして提供されていないのでtensorflowをclone。
r1.14のブランチにcheckoutする。
それ以前でも良いかもしれないけれど、1.15以降はtf.contribの実装がないのでslimのexporter(GraphDefを作るやつ)が動かない。

tensorflow/models

slimを使ってるモデルを想定してるのでもう入ってるとは思うけど、tensorflow/modelsをcloneする

bazel

tocoを使うのにgoogle製のビルドツールbazelを入れる必要がある。 brewとかでは入れないでバイナリから0.26をインストールする。 最新のtensorflowはbazel2.0にupgradeされているっぽいけど、上記のtocoを使うには0.26が必要。

macでのinstall方法

GraphDefファイルを作る

slimのモデルはGraphDefファイルを作るexporterが実装されてるので、それを使ってexportする。

cd <tensorflow/models>
python research/slim/export_inference_graph.py \
  --alsologtostderr \
  --model_name=mobilenet_v2 \
  --dataset_name=imagenet \
  --batch_size=1 \
  --image_size=224 \
  --output_file=<model checkpoint dir>/mobilenet_v2_1.0_graph.pb

--model_nameはちがうモデルを指定したならここで対応を確認: nets_factory.py

FrozenGraphDefファイルを作る

せっかくカスタムモデルであることを想定しているので、最終層以外にいくつかのレイヤーを取り出すモデルにしてみる。

  • ラベルを推論したMobilenetV2/Predictions/Softmax
  • ラベル推論前の特徴量のMobilenetV2/Logits/AvgPool2D

の2つを取り出せるようにfrozenする。 取り出したいレイヤーの名前を確認したいときはGraphDefファイルをnetronで開いてみると良い。

cd <tensorflow/tensorflow> 
bazel run tensorflow/python/tools:freeze_graph -- \
  --input_graph=<model checkpoint dir>/mobilenet_v2_1.0_graph.pb \
  --input_checkpoint=<model checkpoint dir>/mobilenet_v2_1.0_224.ckpt \
  --input_binary=true \
  --output_graph=<model checkpoint dir>/mobilenet_v2_1.0_224_frozen.pb \
  --output_node_names=MobilenetV2/Logits/AvgPool2D,MobilenetV2/Predictions/Softmax

tensorflow/python/tools:freeze_graphのビルドは長くて初回は手元のCPUが12コアあるPCで3時間くらいかかった。

FrozenGraphDefファイルをTF Liteモデルにコンバート

ここまでできてればあとは簡単

cd <tensorflow/tensorflow>
bazel run -c opt tensorflow/lite/toco:toco -- \
    --input_file=<model checkpoint dir>/mobilenet_v2_1.0_224_frozen.pb \
    --input_arrays=input \
    --output_file=<model checkpoint dir>/mobilenet_v2_1.0_224_float32.tflite \
    --output_arrays="MobilenetV2/Logits/AvgPool2D,MobilenetV2/Predictions/Softmax" \
    --inference_type=FLOAT \
    --allow_nonascii_arrays

これで完成!

でもこれだとconvertしたもののweightはfloat32でモデルは軽量化されていない。
TF Liteではfloat16とuint8にもconvertできるのでやってみよう。

uint8(モデルサイズ1/4)にquantizeするには上記コマンドに以下を加える。

--post_training_quantize 

float16 (モデルサイズが1/2)にするには上記コマンドに以下を加える。

--post_training_quantize --quantize_to_float16

詳細はPost-training quantization のドキュメントへ。

追記(2020/04/27): float16にtfliteのinferencerが対応してるのは1.15以降っぽい: 該当コード

TF Liteモデルを使って推論

import matplotlib.pyplot as plt
from PIL import Image
import numpy as np

# image: https://www.flickr.com/photos/kneva/560380352, crop 224, jpg
image = np.array(Image.open('images/test.jpg'))
print(image.shape)

plt.imshow(image)
plt.show()
(224, 224, 3)

import pathlib 
import tensorflow as tf
print("tf version:", tf.__version__)

model_path = pathlib.Path(os.environ.get("CHECKPOINT_DIR")) / "mobilenet_v2_1.0_224_float32.tflite"
interpreter = tf.lite.Interpreter(model_path=str(model_path))
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# print(input_details)
# print(output_details)
    
interpreter.set_tensor(input_details[0]['index'], np.expand_dims(image, 0).astype(np.float32))
interpreter.invoke()

for output in output_details:
    print("===")
    print(output['name'])
    tensor = interpreter.get_tensor(output['index'])
    print(tensor)
tf version: 2.2.0-dev20200323
===
MobilenetV2/Logits/AvgPool2D
[[[[0.27684855 0.5009162  0.00454296 ... 0.         0.02394939
    0.6504333 ]]]]
===
MobilenetV2/Predictions/Softmax
[[3.4995192e-05 2.3994156e-05 1.5229588e-04 ... 1.8983837e-05
  8.8593506e-06 8.0088567e-04]]

こんなかんじで特徴量とラベルの推論値をTFLiteモデルから取得できた!

  /⌒ヽ /⌒ヽ /⌒ヽ
⊂ヽ( ^ω^)つ^ω^)つ^ω^)つ
  \   /  /  /
  ( _フ( _フ( __フ
  (/  (/  (/

さいごに

TFLiteまわりは2系でsupportがまだまだな感じなので、今のところは1系使うのが良さそう。 TensorFlowさん、2系ゴリ押しで2系はドキュメントがあるんだけど1系の情報探そうとするとすごい大変だった・・・。

コンピュータビジョン専門でもないし、普段深層学習モデルをつくる仕事をしているわけでもないので、おかしな部分もあるかもしれないけれど、バージョン周りとか、ツールによるできることの違いとか、知見が色々あったのでせっかくなのでまとめた。

おまけ

トラブルシューティング

自分がハマったエラーメモしておく。

解決したやつ

tfliteにconvertてきたけどグラフが読み込めない

input->dims->size != 4 (0 != 4)Node number 0 (CONV_2D) failed to prepare.

今回の場合exportする際に--batch_size=1 を指定しそびれたせいでこのエラーがでた。 GraphDefファイルを作るときにbatch_sizeやinput_tensorを固定しないとこんな感じのエラーがでる。
TFLiteはinputのTensorをshapeを固定する必要がある。-1のままでは動かない。

frozenかexportしようとしたとき

ERROR: Linking of rule ‘//tensorflow/python:gen_sparse_ops_py_wrappers_cc’ failed (Exit 1)

bazelのversionがあってないとでた。 bazel含め環境が正しいか確認するにはtensorflowのディレクトリで./configureをちゃんとたたいて、自分の環境が間違ってないことを確認しよう。

解決できなかったやつ

PythonAPI経由でconvertしようとしたとき

Check failed: data_type == DT_FLOAT Fatal Python error: Aborted

で止まってconvertできなかったけど情報量がなさすぎてわからなかった

似てるissue:

超便利ツール

これ本当にfrozen成功してるのかな?
なんかconvertできたのに動かないなーー?
っていうのの繰り返しで、そんなときにめちゃくちゃデバッグに役立ったのがnetronというグラフの可視化ツール。

netron
↑さっき作ったtfliteファイルもこんな感じにひらける

Netron supports ONNX (.onnx, .pb, .pbtxt), Keras (.h5, .keras), Core ML (.mlmodel), Caffe (.caffemodel, .prototxt), Caffe2 (predict_net.pb, predict_net.pbtxt), Darknet (.cfg), MXNet (.model, -symbol.json), ncnn (.param) and TensorFlow Lite (.tflite).

とドキュメントにあるように、とにかくいろんな形式のモデルの可視化に対応している。
githubのリポジトリに9kくらいスターついてるし常識ツールなのかもしれないけど、これを知らないがために多大な時間を失った。
今までtensorflowのsession作ったりしながらtensorboardにグラフを出力したりして確認してたけど、そんな事する必要はなかったんだ・・・!