ONNX RuntimeとYoloV3で物体検出¶
ONNX Runtimeを使ってみる。今回はPython APIを使うけど、今後はC/C++も使っていきたい
Requirement¶
以下のバージョンで動作確認しました
- Python==3.7
- numpy==1.16.4
- Pillow==6.0.0
- matplotlib==3.1.0
- onnxruntime==0.4.0
2019/07/08時点、onnxruntimeはPython3.5～Python3.7をサポートしています
Installation¶
Anaconda for Windowsの仮想環境でテストしました
conda create -n onnxruntime pip python
activate onnxruntime
以下のコマンドで各モジュールをインストール
pip install numpy
pip install Pillow
pip install onnxruntime
pip install matplotlib
Download ONNX Model¶
onnx/modelsからYOLOv3モデルをダウンロードします
In [1]:
import numpy as np
from PIL import Image, ImageDraw
# this function is from yolo3.utils.letterbox_image
def letterbox_image(image, size):
'''resize image with unchanged aspect ratio using padding'''
iw, ih = image.size
w, h = size
scale = min(w/iw, h/ih)
nw = int(iw*scale)
nh = int(ih*scale)
image = image.resize((nw,nh), Image.BICUBIC)
new_image = Image.new('RGB', size, (128,128,128))
new_image.paste(image, ((w-nw)//2, (h-nh)//2))
return new_image
def preprocess(img):
model_image_size = (416, 416)
boxed_image = letterbox_image(img, tuple(reversed(model_image_size)))
image_data = np.array(boxed_image, dtype='float32')
image_data /= 255.
image_data = np.transpose(image_data, [2, 0, 1])
image_data = np.expand_dims(image_data, 0)
return image_data
画像リサイズ¶
画像を読み込んだあと、リサイズ関数を使って画像をリサイズします
In [2]:
# Load image
image = Image.open('img/owl.jpg')
# image = Image.open('img/dog.jpg')
# Resized
image_data = preprocess(image)
image_size = np.array([image.size[1], image.size[0]], dtype=np.int32).reshape(1, 2)
# Check
# print(type(image_data))
# print(image_data)
ONNX Runtimeで物体検出【推論フェーズ】¶
【フロー】
- onnxモデルを読み込みでセッションを作成
- 入力名と出力名を取得
- 推論を実行
Input to model¶
モデルに入力するのは、Resized imageとimage sizeです
Resized image (1x3x416x416) Original image size (1x2) which is [image.size[1], image.size[0]]
Output of model¶
モデルから出力されるのは、boxesとlabelsとscoresです
The model has 3 outputs. boxes: (1x'n_candidates'x4), the coordinates of all anchor boxes, scores: (1x80x'n_candidates'), the scores of all anchor boxes per class, indices: ('nbox'x3), selected indices from the boxes tensor. The selected index format is (batch_index, class_index, box_index). The class list is here
In [3]:
import onnxruntime
# 1.onnxモデルを読み込みセッションを作成
session = onnxruntime.InferenceSession('model/yolov3/yolov3.onnx')
# 2.入力名と出力名を取得
input_name = session.get_inputs()[0].name # 'image' 取得
input_name_img_shape = session.get_inputs()[1].name # 'image_shape' 取得
output_name_boxes = session.get_outputs()[0].name # 'boxes' 取得
output_name_scores = session.get_outputs()[1].name # 'scores' 取得
output_name_indices = session.get_outputs()[2].name # 'indices' 取得
# 3. 推論を実行
outputs_index = session.run([output_name_boxes, output_name_scores, output_name_indices],
{input_name: image_data, input_name_img_shape: image_size})
output_boxes = outputs_index[0]
output_scores = outputs_index[1]
output_indices = outputs_index[2]
'''
output_boxes = session.run([output_name_boxes], {input_name: image_data, input_name_img_shape: image_size})[0]
output_scores = session.run([output_name_scores], {input_name: image_data, input_name_img_shape: image_size})[0]
output_indices = session.run([output_name_indices], {input_name: image_data, input_name_img_shape: image_size})[0]
'''
Out[3]:
推論結果¶
In [4]:
# print('input_name =', input_name)
# print('input_name =', input_name_img_shape)
print('boxes =', output_boxes)
print('scores =', output_scores)
print('indices =', output_indices)
In [5]:
print('num_classes =', len(output_scores[0]))
print('num_boxes =', len(output_scores[0][0]))
推論結果の選定❶¶
算出されたスコアの中から、高いスコアのみ抽出することで、物体検出を実現します
デフォルトの閾値より高いスコアを抽出したものは、既にindicesに格納してあります
今回の場合、indicesは以下の通りでした
indices = [[ 0 14 292]]
- バッチID: 0
- クラスID: 14 （cococ class listより”14”は”bird”）
- アンカーボックス座標: output_boxes[292] の中身
以下のコードで各項目を抽出できます
In [6]:
'''
print('indices =', output_indices)
print('scores of class14 =', output_scores[0][14])
print('max score of class14 =', output_scores[0][14][292])
print('box of max score =', output_boxes[0][292])
'''
batch_index = output_indices[0][0] # 0
class_index = output_indices[0][1] # 14
box_index = output_indices[0][2] # 292
print('scores of class14 =', output_scores[batch_index][class_index])
print('max score of class14 =', output_scores[batch_index][class_index][box_index])
print('box of max score =', output_boxes[batch_index][box_index])
ただし、上記のやり方だと、indicesの中身を目視確認してから、コードを書くのでスマートではありません
そのため、この部分を自動化するコードが書きます
推論結果の選定❷¶
In [7]:
out_boxes, out_scores, out_classes = [], [], []
for idx_ in output_indices:
out_classes.append(idx_[1])
out_scores.append(output_scores[tuple(idx_)])
idx_1 = (idx_[0], idx_[2])
out_boxes.append(output_boxes[idx_1])
print(out_classes)
print(out_scores)
print(out_boxes)
トラブルシューティング¶
ボックス(四角形)の左上頂点をp1, 右下頂点をp2としたときout_boxes座標の定義は以下の通りです
out_boxes[p1_y座標 p1_x座標 p2_y座標 p2_x座標]
xとy座標の定義が逆だと思い、一度失敗しました…
出力画像の生成¶
最後にボックス・ラベル名・スコアを描画した出力画像を作成します
cocoラベル一覧¶
算出結果の”14”と”bird”を以下のコードで紐づけします
In [8]:
coco_labels = (
'person',
'bicycle',
'car',
'motorcycle',
'airplane',
'bus',
'train',
'truck',
'boat',
'traffic light',
'fire hydrant',
'stop sign',
'parking meter',
'bench',
'bird',
'cat',
'dog',
'horse',
'sheep',
'cow',
'elephant',
'bear',
'zebra',
'giraffe',
'backpack',
'umbrella',
'handbag',
'tie',
'suitcase',
'frisbee',
'skis',
'snowboard',
'sports ball',
'kite',
'baseball bat',
'baseball glove',
'skateboard',
'surfboard',
'tennis racket',
'bottle',
'wine glass',
'cup',
'fork',
'knife',
'spoon',
'bowl',
'banana',
'apple',
'sandwich',
'orange',
'broccoli',
'carrot',
'hot dog',
'pizza',
'donut',
'cake',
'chair',
'couch',
'potted plant',
'bed',
'dining table',
'toilet',
'tv',
'laptop',
'mouse',
'remote',
'keyboard',
'cell phone',
'microwave',
'oven',
'toaster',
'sink',
'refrigerator',
'book',
'clock',
'vase',
'scissors',
'teddy bear',
'hair drier',
'toothbrush')
ボックス・ラベル・スコア描画¶
In [9]:
import matplotlib.pyplot as plt
%matplotlib inline
# FigureとAxesを作成
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
caption = []
draw_box_p = []
for i in range(0, len(out_classes)):
box_xy = out_boxes[i]
p1_y = box_xy[0]
p1_x = box_xy[1]
p2_y = box_xy[2]
p2_x = box_xy[3]
draw_box_p.append([p1_x, p1_y, p2_x, p2_y])
draw = ImageDraw.Draw(image)
draw.rectangle(draw_box_p[i], outline=(255, 0, 0), width=5)
# クラス名とスコア描画
caption.append(coco_labels[out_classes[i]])
caption.append('{:.2f}'.format(out_scores[i]))
ax.text(p1_x, p1_y,
': '.join(caption),
style='italic',
bbox={'facecolor': 'white', 'alpha': 0.7, 'pad': 10})
caption.clear()
# 画像を表示
img = np.asarray(image)
ax.imshow(img)
Out[9]:
