TensorFlowの学習済みモデルを変換してONNXRuntimeで物体検出

こんにちは。

コンピュータビジョン(『ロボットの眼』開発)が専門の”はやぶさ”@Cpp_Learningです。

前回『ONNX RuntimeとYoloV3でリアルタイム物体検出』について書きました。

ONNX RuntimeとYoloV3でリアルタイム物体検出Microsoft社製OSS”ONNX Runtime”の入門から実践まで学べる記事です。ONNXおよびONNX Runtimeの概要から、YoloV3モデルによる物体検出（ソースコード付）まで説明します。深層学習や画像処理に興味のある人にオススメの内容です。...

今回は『ONNX RuntimeとSSDでリアルタイム物体検出』を実践します。

Contents

1 本記事を書くモチベーション
2 ONNXおよびONNXRuntimeとは
3 深層学習のフェーズ -学習と推論-
- 3.1 深層学習の実践フロー
4 TensorFlowモデルをONNXに変換
5 ONNXRuntimeとSSDモデルで物体検出
6 まとめ

本記事を書くモチベーション

くるる

既にONNX Runtimeに関する記事を公開しているのに、なぜ本記事を書くの？

というのが気になるフクロウのくるるちゃん@kururu_owl のために、本記事を書くモチベーションを少し説明しておきます。

各深層学習フレームワークがONNX形式を採用することで、学習済みモデルの交換を簡単に実現することができます。

しかし、前回はONNX Model Zooを使用したので、モデルの交換を実践していません。

また、ONNXモデルのInput/Outputが異なる（特にOutput）が異なる場合、同じ”物体検出”でもOutputに合わせてソースコードを作成する必要があります。

そのため、本記事を書くモチベーションは以下の通りです。

本記事を書くモチベーション

ONNX形式を使うメリット”学習済みモデルの交換”を実践したい
Input/Outputが異なるモデルのONNXRuntimeによる推論を実践したい

前回の記事では書き切れなかったONNXの魅力を本記事で補間したいと思います！

演習問題やソースコード写経などは、ある程度の数をこなさないと”見えてこない”ものもあります。前回と今回はどちらもONNX Runtimeがメインの記事ですが、ぜひ見比べて”差分”を感じ取ってほしいです。

ONNXおよびONNXRuntimeとは

ONNXやONNX Runtimeの概要ついては、以下の記事で説明済みなので、本記事では割愛します。

この記事でインストール方法なども説明しているので、先に読んでから、以降を読み進めることをオススメします！

深層学習のフェーズ -学習と推論-

深層学習を実践する場合、【学習フェーズ】・【推論フェーズ】があります。

【学習フェーズ】

ニューラルネットワーク設計
学習（ニューラルネットワークの”重み”を調整）
学習済みモデルを保存

【推論フェーズ】

学習済みモデルを読込む
学習済みモデルへの入力値（センサ値など）を取得
推論（入力値に対する推論値を取得）
推論値を使って”ごにょごにょ”する（制御とか）

ケースバイケースですが【学習フェーズ】・【推論フェーズ】各々で専用マシンを使うことが多いです。

例えば、【学習フェーズ】にAWSやAzureなどのクラウドサービスを利用し、【推論フェーズ】に組込み機器などのエッジディバイスを使います。

組込みAI向けエッジディバイスJetsonシリーズ

NVIDIA® Jetson™ TX2 開発者キット

created by Rinker

菱洋エレクトロ

NVIDIA Jetson Nano 開発キット

created by Rinker

NVIDIA

>> Amazonレビューをチェックする

深層学習の実践フロー

今回、学習は行いませんが『TensorFlowの学習済みモデルをONNXモデルに変換』を実践します。

TensorFlowモデルをONNXモデルに変換 ONNXモデルをエクスポートできる深層学習フレームワークは複数ありますが、
SSD系の学習済みモデルについては、 Tensorflow detection model zooが非常に充実してます。

なので、TensorFlowモデルをONNXモデルに変換して、最終的にはONNX Runtimeで物体検出（推論）するところまで実践します。

TensorFlowモデルをONNXに変換

ONNXモデルへの変換は【学習フェーズ】に含まれない気もしますが、【推論フェーズ】よりも前に実施すべき処理なので…

今回は、無料で使えるクラウドサービスColaboratoryとtensorflow-onnxを使って『TensorFlowの学習済みモデルをONNXモデルに変換』を実践しました。

実践した内容は、以下の『技術ノート』にまとめました(*･ω･)ﾉ♪

はやぶさの技術ノート：tf2onnx

個人的に、リアルタイム物体検出が好きなので、”軽快に動作する”ssdlite_mobilenet_v2_cocoを採用し、ONNXモデルに変換しています。

Colaboratory使い方で困ったときは、以下の記事が非常に参考になります！

【秒速で無料GPUを使う】深層学習実践Tips on Colaboratory｜Qiita

ONNXRuntimeとSSDモデルで物体検出

取得したONNXモデルを使って物体検出を行う”ONNXRuntime_SSD.py”を作成しました。

import onnxruntime
import math
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageColor, ImageFont

# coco labels list
coco_classes = {
    1: 'person',
    2: 'bicycle',
    3: 'car',
    4: 'motorcycle',
    5: 'airplane',
    6: 'bus',
    7: 'train',
    8: 'truck',
    9: 'boat',
    10: 'traffic light',
    11: 'fire hydrant',
    12: 'stop sign',
    13: 'parking meter',
    14: 'bench',
    15: 'bird',
    16: 'cat',
    17: 'dog',
    18: 'horse',
    19: 'sheep',
    20: 'cow',
    21: 'elephant',
    22: 'bear',
    23: 'zebra',
    24: 'giraffe',
    25: 'backpack',
    26: 'umbrella',
    27: 'handbag',
    28: 'tie',
    29: 'suitcase',
    30: 'frisbee',
    31: 'skis',
    32: 'snowboard',
    33: 'sports ball',
    34: 'kite',
    35: 'baseball bat',
    36: 'baseball glove',
    37: 'skateboard',
    38: 'surfboard',
    39: 'tennis racket',
    40: 'bottle',
    41:'wine glass',
    42: 'cup',
    43: 'fork',
    44: 'knife',
    45: 'spoon',
    46: 'bowl',
    47: 'banana',
    48: 'apple',
    49: 'sandwich',
    50: 'orange',
    51: 'broccoli',
    52: 'carrot',
    53: 'hot dog',
    54: 'pizza',
    55: 'donut',
    56: 'cake',
    57: 'chair',
    58: 'couch',
    59: 'potted plant',
    60: 'bed',
    61: 'dining table',
    62: 'toilet',
    63: 'tv',
    64: 'laptop',
    65: 'mouse',
    66: 'remote',
    67: 'keyboard',
    68: 'cell phone',
    69: 'microwave',
    70: 'oven',
    71: 'toaster',
    72: 'sink',
    73: 'refrigerator',
    74: 'book',
    75: 'clock',
    76: 'vase',
    77: 'scissors',
    78: 'teddy bear',
    79: 'hair drier',
    80: 'toothbrush'
}

# Draw box and label for 1 detection.
def draw_detection(draw, boxes, classes, scores):
    # Get image width and height
    width, height = draw.im.size

    # The box is relative to the image size so we multiply with height and width to get pixels.
    top = boxes[0] * height
    left = boxes[1] * width
    bottom = boxes[2] * height
    right = boxes[3] * width

    top = max(0, np.floor(top + 0.5).astype('int32'))
    left = max(0, np.floor(left + 0.5).astype('int32'))
    bottom = min(height, np.floor(bottom + 0.5).astype('int32'))
    right = min(width, np.floor(right + 0.5).astype('int32'))

    # Get label
    label = coco_classes[classes]

    # Draw box
    draw.rectangle([left, top, right, bottom], outline=(0, 255, 0), width=3)

    # Draw label and score
    text = label + " " + ('%.2f' % scores)
    font = ImageFont.truetype("arial.ttf", 30)
    draw.text((left, top-30), text, fill=(255, 0, 255), font=font)

def main():
    # Load image
    image = Image.open('img/owl.jpg')
    # image = Image.open('img/dog.jpg')

    # HWC to NHWC
    image_data = np.array(image.getdata()).reshape(image.size[1], image.size[0], 3)
    image_data = np.expand_dims(image_data.astype(np.uint8), axis=0)

    # Check image data
    # print(type(image_data))
    # print(image_data)

    # 1. Make session
    session = onnxruntime.InferenceSession('model/ssd/ssdlite_mobilenet_v2_coco_2018_05_09.onnx')

    # 2. Get input/output name
    input_name = session.get_inputs()[0].name   # 'image'

    output_name_boxes = session.get_outputs()[0].name     # 'boxes'
    output_name_classes = session.get_outputs()[1].name   # 'classes'
    output_name_scores = session.get_outputs()[2].name    # 'scores'
    output_name_num = session.get_outputs()[3].name       # 'number of detections'

    # 3. Run
    outputs_index = session.run([output_name_num, output_name_boxes, output_name_scores, output_name_classes],
                                {input_name: image_data})

    # Result
    output_num = outputs_index[0]
    output_boxes = outputs_index[1]
    output_scores = outputs_index[2]
    output_classes = outputs_index[3]

    # Check input/output name
    print(input_name, output_name_num, output_name_boxes, output_name_scores, output_name_classes)

    # Check result
    print(output_num)
    print(output_boxes[0][0])
    print(output_scores)
    print(output_classes) # 15=bird

    threshold = 0.7
    batch_size = output_num.shape[0]
    draw = ImageDraw.Draw(image)
    for batch in range(0, batch_size):
        for detection in range(0, int(output_num[batch])):
            classes = output_classes[batch][detection]
            boxes = output_boxes[batch][detection]
            scores = output_scores[batch][detection]
            # Draw box, label score > threshold,
            if scores > threshold:
                draw_detection(draw, boxes, classes, scores)

    # plt.figure(figsize=(20, 10))
    plt.axis('off')
    plt.imshow(image)
    plt.show()

if __name__ == '__main__':
    main()

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

import onnxruntime

import math

import numpy as np

import matplotlib.pyplot as plt

from PIL import Image, ImageDraw, ImageColor, ImageFont

# coco labels list

coco_classes = {

1: 'person',

2: 'bicycle',

3: 'car',

4: 'motorcycle',

5: 'airplane',

6: 'bus',

7: 'train',

8: 'truck',

9: 'boat',

10: 'traffic light',

11: 'fire hydrant',

12: 'stop sign',

13: 'parking meter',

14: 'bench',

15: 'bird',

16: 'cat',

17: 'dog',

18: 'horse',

19: 'sheep',

20: 'cow',

21: 'elephant',

22: 'bear',

23: 'zebra',

24: 'giraffe',

25: 'backpack',

26: 'umbrella',

27: 'handbag',

28: 'tie',

29: 'suitcase',

30: 'frisbee',

31: 'skis',

32: 'snowboard',

33: 'sports ball',

34: 'kite',

35: 'baseball bat',

36: 'baseball glove',

37: 'skateboard',

38: 'surfboard',

39: 'tennis racket',

40: 'bottle',

41:'wine glass',

42: 'cup',

43: 'fork',

44: 'knife',

45: 'spoon',

46: 'bowl',

47: 'banana',

48: 'apple',

49: 'sandwich',

50: 'orange',

51: 'broccoli',

52: 'carrot',

53: 'hot dog',

54: 'pizza',

55: 'donut',

56: 'cake',

57: 'chair',

58: 'couch',

59: 'potted plant',

60: 'bed',

61: 'dining table',

62: 'toilet',

63: 'tv',

64: 'laptop',

65: 'mouse',

66: 'remote',

67: 'keyboard',

68: 'cell phone',

69: 'microwave',

70: 'oven',

71: 'toaster',

72: 'sink',

73: 'refrigerator',

74: 'book',

75: 'clock',

76: 'vase',

77: 'scissors',

78: 'teddy bear',

79: 'hair drier',

80: 'toothbrush'

}

# Draw box and label for 1 detection.

def draw_detection(draw, boxes, classes, scores):

# Get image width and height

width, height = draw.im.size

# The box is relative to the image size so we multiply with height and width to get pixels.

top = boxes[0] * height

left = boxes[1] * width

bottom = boxes[2] * height

right = boxes[3] * width

top = max(0, np.floor(top + 0.5).astype('int32'))

left = max(0, np.floor(left + 0.5).astype('int32'))

bottom = min(height, np.floor(bottom + 0.5).astype('int32'))

right = min(width, np.floor(right + 0.5).astype('int32'))

# Get label

label = coco_classes[classes]

# Draw box

draw.rectangle([left, top, right, bottom], outline=(0, 255, 0), width=3)

# Draw label and score

text = label + " " + ('%.2f' % scores)

font = ImageFont.truetype("arial.ttf", 30)

draw.text((left, top-30), text, fill=(255, 0, 255), font=font)

def main():

# Load image

image = Image.open('img/owl.jpg')

# image = Image.open('img/dog.jpg')

# HWC to NHWC

image_data = np.array(image.getdata()).reshape(image.size[1], image.size[0], 3)

image_data = np.expand_dims(image_data.astype(np.uint8), axis=0)

# Check image data

# print(type(image_data))

# print(image_data)

# 1. Make session

session = onnxruntime.InferenceSession('model/ssd/ssdlite_mobilenet_v2_coco_2018_05_09.onnx')

# 2. Get input/output name

input_name = session.get_inputs()[0].name # 'image'

output_name_boxes = session.get_outputs()[0].name # 'boxes'

output_name_classes = session.get_outputs()[1].name # 'classes'

output_name_scores = session.get_outputs()[2].name # 'scores'

output_name_num = session.get_outputs()[3].name # 'number of detections'

# 3. Run

outputs_index = session.run([output_name_num, output_name_boxes, output_name_scores, output_name_classes],

{input_name: image_data})

# Result

output_num = outputs_index[0]

output_boxes = outputs_index[1]

output_scores = outputs_index[2]

output_classes = outputs_index[3]

# Check input/output name

print(input_name, output_name_num, output_name_boxes, output_name_scores, output_name_classes)

# Check result

print(output_num)

print(output_boxes[0][0])

print(output_scores)

print(output_classes) # 15=bird

threshold = 0.7

batch_size = output_num.shape[0]

draw = ImageDraw.Draw(image)

for batch in range(0, batch_size):

for detection in range(0, int(output_num[batch])):

classes = output_classes[batch][detection]

boxes = output_boxes[batch][detection]

scores = output_scores[batch][detection]

# Draw box, label score > threshold,

if scores > threshold:

draw_detection(draw, boxes, classes, scores)

# plt.figure(figsize=(20, 10))

plt.axis('off')

plt.imshow(image)

plt.show()

if __name__ == '__main__':

main()

以下のコマンドで動作確認できます。

python ONNXRuntime_SSD.py

出力画像はPyTorchでMobileNet SSDによるリアルタイム物体検出を意識した表示にしました。

人間用の椅子に座ってる猫のくるるちゃん@kururu_owl のが今日も可愛い…ん？

くるる

にゃーん♪

ソースコードのパス設定を任意に変更してください

120行目：画像パス
132行目：ONNXモデルパス

まとめ

本記事は、以下のモチベーションで書き上げました。

本記事を書くモチベーション

ONNX形式を使うメリット”学習済みモデルの交換”を実践したい
Input/Outputが異なるモデルのONNXRuntimeによる推論を実践したい

前回書いた記事と本記事を読んで、ONNXやONNX Runtimeに興味を持ってくれたら、とれも嬉しいです。

また「組込みAIを実践したい！」という人が現れたら、最高に嬉しいです(*･ω･)ﾉ♪

くるる

くるるはラズベリーパイ食べながらJetson Nanoで”でぃーぷらーにんぐ”したい♪

楽しく実践してくれたら嬉しいです！笑

はやぶさ

理系応援ブロガー”はやぶさ”@Cpp_Learningは頑張る理系を応援します！

組込みAI向けエッジディバイスJetsonシリーズ

NVIDIA® Jetson™ TX2 開発者キット

created by Rinker

菱洋エレクトロ

NVIDIA Jetson Nano 開発キット

created by Rinker

NVIDIA

>> Amazonレビューをチェックする

PyTorchでMobileNet SSDによるリアルタイム物体検出深層学習フレームワークPytorchを使い、ディープラーニングによる物体検出の記事を書きました。物体検出手法にはいくつか種類がありますが、今回はMobileNetベースSSDによる『リアルタイム物体検出』を行いました。...

MXNetで『カメラ・動画対応！物体検出ソフト』を作った -Yolo, SSD, Faster-RCNNモデル対応-こんにちは。コンピュータビジョン(『ロボットの眼』開発)が専門の”はやぶさ”@Cpp_Learningです。『深層学習に...

【NNabla】Neural Network Librariesの使い方まとめこんにちは。ディープラーニングお兄さんの”はやぶさ”@Cpp_Learningです。仕事・プライベートどちらでも”深層学...

TensorFlowの学習済みモデルを変換してONNXRuntimeで物体検出

本記事を書くモチベーション

ONNXおよびONNXRuntimeとは

深層学習のフェーズ -学習と推論-

深層学習の実践フロー

TensorFlowモデルをONNXに変換

ONNXRuntimeとSSDモデルで物体検出

まとめ

プライム感謝祭セールを完全攻略！事前準備やお得キャンペーンを徹底解説

【アマゾンプライムデー2023年】完全攻略！事前準備やお得キャンペーンを徹底解説

【2023年版】アマゾンプライム会員のメリット・デメリットを紹介！

Golang × WebAssembly（wasm）入門

【Python】最適化フレームワークのCodable Model Optimizerで回帰モデルのパラメータ調整を自動化する

【GitHub】シンプルなREADME.mdの書き方 -コピペで使えるテンプレート付き-

【Pyxel】Pythonでレトロゲームを作ろう！総集編 -まるっと1週間でゲーム開発入門-

【仕事効率化】Visual Studio Code で Markdown を使いこなす

【深層学習入門】画像処理の基礎(画素操作)からCNN設計まで

【Pyxel】Pythonでレトロゲームを作ろう！ Day 1 -画像の扱い方-

本記事を書くモチベーション

ONNXおよびONNXRuntimeとは

深層学習のフェーズ -学習と推論-

深層学習の実践フロー

TensorFlowモデルをONNXに変換

ONNXRuntimeとSSDモデルで物体検出

まとめ

プライム感謝祭セールを完全攻略！事前準備やお得キャンペーンを徹底解説

【アマゾンプライムデー2023年】完全攻略！事前準備やお得キャンペーンを徹底解説

【2023年版】アマゾンプライム会員のメリット・デメリットを紹介！

Golang × WebAssembly（wasm）入門

【Python】最適化フレームワークのCodable Model Optimizerで回帰モデルのパラメータ調整を自動化する

【GitHub】シンプルなREADME.mdの書き方 -コピペで使えるテンプレート付き-

【Pyxel】Pythonでレトロゲームを作ろう！ 総集編 -まるっと1週間でゲーム開発入門-

【仕事効率化】Visual Studio Code で Markdown を使いこなす

【深層学習入門】画像処理の基礎(画素操作)からCNN設計まで

【Pyxel】Pythonでレトロゲームを作ろう！ Day 1 -画像の扱い方-

【Pyxel】Pythonでレトロゲームを作ろう！総集編 -まるっと1週間でゲーム開発入門-