ディープラーニングで物体検出を試してみた - 第3回: 検出結果をプロットする
こんにちは、けんにぃです。
ナビタイムジャパンで公共交通の時刻表を使ったサービス開発やリリースフローの改善を担当しています。
物体検出の解説 3 回目になります。2 回目では TensorFlow Hub に登録されているモデルを使って物体検出を行いました。
今回は検出結果であるバウンディングボックスをプロットする処理を解説しようと思います。
全記事のリンクはこちらになります。
・第 1 回: 画像の読み込み(前処理)
・第 2 回: TensorFlow Hub を使う(モデルの利用)
・第 3 回: 出力結果をプロットする(後処理)
プロットの仕方
下記のようにバウンディングボックスとカテゴリ名を埋め込んでみましょう。
第 1 回目でもお話しましたが画像処理には複数のやり方があるため、参考にするサイトによって使用しているライブラリが異なっていることが多くて混乱します。そのため、私が出会ってきた書き方をすべてご紹介しようと思います。
画像に検出結果を埋め込む方法は下記の 3 通りがあります。
・OpenCV を使う
・Pillow を使う
・TensorFlow を使う
バウンディングボックスに該当する矩形とカテゴリ名を示すテキストをプロットするには下記の API を使用します。
OpenCV
・矩形:cv.rectangle()
・テキスト:cv.putText()
Pillow
・矩形:ImageDraw.text()
・テキスト:ImageDraw.rectangle()
TensorFlow
・矩形:tf.image.draw_bounding_boxes()
・テキスト:(API なし)
TensorFlow にもバウンディングボックスをプロットする API が用意されていますがテキストをプロットする API が用意されていないため、実質 OpenCV か Pillow を使うことになると思います。そのため TensorFlow の API は簡単な紹介だけにとどめておきます。
それぞれ詳細に見ていきましょう。下記のコードを実行していることを前提とします。
import tensorflow as tf
import tensorflow_hub as hub
# 画像データの読み込み
image = tf.io.read_file("dog.jpg")
image = tf.io.decode_jpeg(image)
# image = tf.image.convert_image_dtype(image, tf.float32) # 今回は正規化は不要
image = tf.expand_dims(image, axis=0)
# モデルをロード
model = hub.load("https://tfhub.dev/tensorflow/faster_rcnn/inception_resnet_v2_640x640/1")
# 推論を実行
outputs = model(image)
使用した画像はこちらです。
dog.jpg
共通の処理として、まずモデルの出力結果からバウンディングボックスを取り出します。
boxes = outputs["detection_boxes"][0]
バウンディングボックスは正規化座標になっているので(第 2 回参照)、画像サイズをかけて座標を取り出します。座標は整数値にする必要があるので int にキャストしておきます。
height, width = image.shape[1:3]
box = boxes[0] * np.array([height, width, height, width])
y_min, x_min, y_max, x_max = box.numpy().astype(int)
これを検出したバウンディングボックス分、同じことを行います。普通はスコアにしきい値を設けて、スコアがしきい値以上のものだけを取り出します。
boxes = outputs["detection_boxes"][0]
labels = outputs["detection_classes"][0]
scores = outputs["detection_scores"][0]
for box, label, score in zip(boxes, labels, scores):
# スコアが 0.5 より大きいものだけ抽出
if score <= 0.5:
continue
y_min, x_min, y_max, x_max = box.numpy().astype(int)
...
こうして得られた y_min, x_min, y_max, x_max を矩形としてプロットしていきます。紛らわしいのが、今回使用したモデルが返却する座標は y, x の順ですが、ライブラリでプロットするときに渡す座標は x, y の順で渡します。これはよく間違えるので注意が必要です。
プロットに使うライブラリは下記で解説しますのでお好きなものをご利用下さい。
OpenCV を使う場合
cv.rectangle() を使うと矩形がプロットできます。
import cv2 as cv
output_image = cv.imread("dog.jpg")
cv.rectangle(output_image, (x_min, y_min), (x_max, y_max), color=(0, 0, 255))
color は BGR の順で値を渡します。RGB ではないので注意して下さい。
ハマりどころなのはタプルで渡しているところをリストにするとエラーになってしまいます。
cv.rectangle(output_image, [x_min, y_min], [x_max, y_max], color=[0, 0, 255])
通常 Python の関数はタプルで渡せるところはリストでも渡せるようになっているので、この辺の違いをあまり意識しないと思いますが、OpenCV の場合は挙動が変わるため注意が必要です。
テキストのプロットには cv.putText() を使います。
cv.putText(
img,
text="dog",
org=(x_min, y_min),
fontFace=cv.FONT_HERSHEY_SIMPLEX,
fontScale=1,
color=(0, 0, 255),
)
引数の意味は次のとおりです。
・text: プロットするテキスト
・org: プロットする位置
・fontFace: フォントの種類
・fontScale: フォントの倍率
・color: 色を BGR の順で渡す
最後に画像を保存します。
cv.imwrite("output.jpg", output_image)
Pillow を使う場合
ImageDraw を使って矩形やテキストをプロットすることができます。まず ImageDraw インスタンスを作成します。
from Pillow import Image, ImageDraw
output_image = Image.open("dog.jpg")
draw = ImageDraw.Draw(output_image)
draw.rectangle() を使うと矩形がプロットできます。
draw.rectangle([x_min, y_min, x_max, y_max], outline=(255, 0, 0))
outline は外枠の色を RGB の順で渡します。
続いて、テキストのプロットには draw.text() を使います。
draw.text([x_min, y_min], "dog.jpg", fill=(255, 0, 0))
fill はテキストの色を RGB の順で渡します。フォントをお好みのものに変更したければ ImageFont.truetype() でインストールされているフォントを指定します。
from PIL import ImageFont
font = ImageFont.truetype("NotoSansCJKjp-Regular")
draw.text([x_min, y_min], "dog.jpg", fill=(255, 0, 0), font=font)
最後に画像を保存します。
output_image.save("output.jpg")
TensorFlow を使う場合
tf.image.draw_bounding_boxes() を使うことでバウンディングボックスをプロットできます。
output_image = tf.image.draw_bounding_boxes(
image,
boxes,
colors,
)
image は正規化しておく必要があります。詳細は API ドキュメントを参考にして下さい。
ファイルに書き出すときはファイルを読んだときと逆順にテンソルを変換していくことで書き出すことができます。
output_image = output_image[0]
output_image = tf.image.convert_image_dtype(output_image, tf.uint8)
output_image = tf.io.encode_jpeg(output_image)
output_image = tf.io.write_file("output.jpg", output_image)
ラベルと物体名称のマッピング
モデルが返却したラベルの ID に対する物体の名称が分からないと画像にプロットできませんが、そのマッピングデータはモデルには同梱されていないため別途用意する必要があります。
今回使用したモデルは COCO 2017 という有名なデータセットで作られていますが、これを利用したモデルはたくさんあり、そのモデルの GitHub リポジトリにマッピングデータが同梱されていることが多いため、そこからマッピングデータを拝借してくるのが一番楽なやり方です。
Faster R-CNN のモデルを作成した TensorFlow Model Garden のリポジトリにもマッピングデータが登録されています。
このリポジトリに用意されている Object Detection API を使ってラベルのマッピングデータを利用してみようと思います。
まずリポジトリをクローンします。
$ git clone --depth 1 https://github.com/tensorflow/models
Object Detection API をインストールします。
$ pip install protobuf-compiler
$ cd models/research/
$ protoc object_detection/protos/*.proto --python_out=.
$ cp object_detection/packages/tf2/setup.py .
$ pip install -q .
インストールが終われば API を使ってラベルのマッピングデータを読み込みます。
from object_detection.utils import label_map_util
category_index = label_map_util.create_category_index_from_labelmap(
"models/research/object_detection/data/mscoco_label_map.pbtxt",
use_display_name=True,
)
ラベル ID:18 の名称を取得するには下記のようにします。
>>> category_index[18]["name"]
'dog'
これで物体名称が取得できます。
まとめ
機械学習を勉強していて感じたことは、知っておくべきことが多岐に渡るなということです。物体検出の一連の流れを理解するためには下記のように多くの知見を習得する必要がありました。
・ディープラーニングの基礎理論
・物体検出の仕組み
・各種モデルの構造の調べ方と使い方
・TensorFlow の使い方
・画像収集のためのクローラ開発
・画像処理と OpenCV, Pillow の使い方
・配列操作のための Numpy の使い方
しかも機械学習は人気であるがゆえに競合するライブラリが多く、情報のアップデートが早いため、参考サイトも説明の仕方が異なっていることが多く、かえって参考にしづらいことが本当にハマりどころでした。
今は機械学習に関するクラウドサービスが充実しているため、一から十まで自作することはあまりないかも知れませんが、基礎的な知識がだいぶ深まったので、クラウドサービスを使うときもこの知識は活用できそうです。