Pytorchによる物体検出 SSD Google Colab環境 Part(1) | 初心者向けWEB技術と機械学習のBlog

Contents

はじめに
SSD
学習済みデータを使って物体検出
Pythonコードと実行結果
まとめ

はじめに

こんにちわ、swim-loverです。 PythonとPytorchで物体検出を行っています。Pythonを始めたばかりですが、「使いながら覚える」をコンセプトに勉強しています。第8回は、 Pytorchを使ってMNISTデータの準備で登場したDataLoader関数について勉強しました。Pythonの機械学習の勉強編は、まだ継続して進めたいと思いますが、別編として、Pytorchによる物体検出を始めたいと思います。

SSD

Pytorch、物体検出というキーワードで検索すると、SSD(Single Shot Multibox Detector)とYOLO(you only look once)を使う例が多く見つかりました。今回は、SSDについて、触れてみたいと思います。SSDの論文には、以下の図が載っています。

SSDには、学習時には、(a)入力画像とGround truth boxesが必要であると書かれていました。Ground truthは、正解データを意味するようですが、正解データのBoxとはなんだかスッと入ってきません。ここでは、あらかじめ、オブジェクト（猫、犬）を囲む長方形と犬や猫である印を作っておくというように理解しておきました。次に異なるスケール(b)8×8と(c)4,4のマップでいくつかの箱(default box)のsmall set(例では４つの箱)を使って、処理を進めるようです。そしてそれぞれのBoxでは、全てのオブジェクトカテゴリー（犬、猫、車など）について、対象物体とのオフセットoffset delta(cx,cy,w,h), 信頼度confidence(c1=10%,c2=80%,c3=0.1%)を計算するようです。学習時には、default boxとGround truth boxesを一致させるようです。

また、検出モデルの図も論文に載っていました。出力はクラス毎に8732個のアウトプットが出てくるようです。(Detections:8732 per classes)、さらに最終段でNon-Maximum Suppression処理が行われます。これは、8732個のアウトプットから出力の間引き処理が行われるようです。

これより論文を掘り下げても、私の理解度を超えてしましますので、実際にSSDを使ってみたいと思います。

学習済みデータを使って物体検出

PytorchでSSDを使うには、最初のステップとして、学習済みデータを使うのがほとんどだと思います。また、Pytorchのサンプルコードも公開されていますので、そのまま使用することにしました。Pytorchのサイトにそのままcolabのコードが公開されているので、そのまま自分のColab環境にコピーして実行します。

注意点が一つあります。colabのランタイムタイプの設定（メニュからランタイム->ランタイムタイプを変更を選択します。）がGPUになっていないようであれば、GPUに設定する必要があります。

Pythonコードと実行結果

import torch

print(torch.__version__)
precision = 'fp32'
device = torch.device('cpu')
#print(torch.cuda.is_available())
#torch.hub.list('NVIDIA/DeepLearningExamples:torchhub')
#torch.hub.help('NVIDIA/DeepLearningExamples:torchhub','nvidia_ssd')
ssd_model = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd', model_math=precision)#学習済みモデル
utils = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd_processing_utils')#データの処理機能

# CPUで動くようにして、推論モード
ssd_model.to('cpu')
ssd_model.eval()

#input data
uris = [
    'http://images.cocodataset.org/val2017/000000397133.jpg',
    'http://images.cocodataset.org/val2017/000000037777.jpg',
    'http://images.cocodataset.org/val2017/000000252219.jpg'
]

# 入力データ作成
inputs = [utils.prepare_input(uri) for uri in uris]
tensor = utils.prepare_tensor(inputs, precision == 'fp16')
tensor = tensor.to('cpu')

# 検知
with torch.no_grad():
    detections_batch = ssd_model(tensor)

# 検知する種類 犬とか人とか80種類
classes_to_labels = utils.get_coco_object_dictionary()
print(classes_to_labels)
print('クラス数', len(classes_to_labels))

# 結果取得
results_per_input = utils.decode_results(detections_batch)
best_results_per_input = [utils.pick_best(results, 0.40) for results in results_per_input]

# 1つ目の結果
bboxes, classes, confidences = best_results_per_input[0]
print('検知ボックス', bboxes)
print('検知クラス', classes, classes_to_labels[classes[0]-1])
print('確信度', confidences)

from matplotlib import pyplot as plt
import matplotlib.patches as patches
 
for image_idx in range(len(best_results_per_input)):
    fig, ax = plt.subplots(1)
    # Show original, denormalized image...
    image = inputs[image_idx] / 2 + 0.5
    ax.imshow(image)
    # ...with detections
    bboxes, classes, confidences = best_results_per_input[image_idx]
    for idx in range(len(bboxes)):
        left, bot, right, top = bboxes[idx]
        x, y, w, h = [val * 300 for val in [left, bot, right - left, top - bot]]
        rect = patches.Rectangle((x, y), w, h, linewidth=1, edgecolor='r', facecolor='none')
        ax.add_patch(rect)
        ax.text(x, y, "{} {:.0f}%".format(classes_to_labels[classes[idx] - 1], confidences[idx]*100), bbox=dict(facecolor='white', alpha=0.5))
plt.show()