PyTorchとYOLOv5で検出したオブジェクトの バウンディングボックスの座標情報を取得する

  • URLをコピーしました!

今回はPyTorchとYOLOv5を使った物体検出で、検出したオブジェクトを可視化するバウンディングボックスの情報を取得する方法を解説します。

\ 機械学習を学びたい人には自宅で学べるUdemyがおすすめ! /

講座単位で購入できます!

目次

バウンディングボックスとは

バウンディングボックスとは、機械学習等を使ったオブジェクト検出を行う際、検出したオブジェクトの領域を可視化するために表示される矩形のことです。

バウンディングボックス

このようなバウンディングボックスが表示されることで、検出したオブジェクトの位置や分類されたクラス、確信度のスコアなどを目視で確認することができます。そしてこのバウンディングボックスが表示されている領域の情報を取得することで、画像処理や制御に使用することができます。

今回はPythonのプログラムで、PyTorchとYOLOv5を使った場合のバウンディングボックスの情報を取得する方法を解説していきます。

PyTorchとYOLOv5を使った物体検出を行う方法

PyTorchとYOLOv5のセットアップ方法と、推論の実行方法は以下の記事で解説しています。

こちらの記事のソースコードを元に解説しますので、事前にPyTorchとYOLOv5を使って推論を実行するところまでの手順を実施しておいてください。

バウンディングボックスの情報を取得する方法

前項で紹介した「PyTorchとYOLOv5で物体検出する方法」で紹介したソースコードの推論処理部分を次のように改変します。推論結果を以下のようにpandasのデータフレームで受け取ることができます。

#推論を実行
result = model(img)

#バウンディングボックスの情報を取得
obj = result.pandas().xyxy[0]
print(obj)

受け取ったpandasのデータは以下のようになります。xmin(x軸左上)、ymin(y軸左上)、xmax(x軸右下)ymax(y軸右下)の座標が取得できます。confidenceはスコア(信頼度)、classは分類されたクラスとなります。

YOLOv5  2022-10-22 Python-3.9.7 torch-1.12.1+cpu CPU

Fusing layers... 
YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients
Adding AutoShape... 
         xmin        ymin        xmax        ymax  confidence  class    name
0  219.176910  407.951538  346.656036  873.833313    0.864783      0  person
1   50.105721  395.405731  239.340057  913.149780    0.860630      0  person
2  673.475281  400.604767  810.000000  875.171875    0.849397      0  person
3   14.653307  219.916382  810.000000  810.014282    0.816330      5     bus
4    0.084950  553.378845   64.705879  874.105347    0.644388      0  person

分類可能なクラスについては以下となります。

['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']

作成したソースコード

今回作成したソースコードはこちらです。
推論結果をデータフレームとして受け取り、個々の要素にアクセスして座標や分類されたクラスを取得して表示します。

import cv2
import torch

model = torch.hub.load('ultralytics/yolov5', 'yolov5s') 
img = cv2.imread('(YoloV5のフォルダが配置されているパス)\yolov5\\data\\images\\cats.jpg')
result = model(img)

#推論結果を取得
obj = result.pandas().xyxy[0]

#バウンディングボックスの情報を取得
for  i in range(len(obj)):
    name = obj.name[i]
    xmin = obj.xmin[i]
    ymin = obj.ymin[i]
    xmax = obj.xmax[i]
    ymax = obj.ymax[i]

    print("name =", name, "xmin =", xmin, "ymin =", ymin, "xmax =", ymax, "ymin =", ymax)

#バウンディングボックスを描画
result.render()
cv2.imshow('result', result.ims[0])
cv2.waitKey(0)

今回のテストのために猫の画像を用意しました。この画像の中に含まれる猫の位置を取得したいと思います。

推論に使用する画像

実行結果

実行結果は以下となります。推論を行い、検出したオブジェクトの位置にバウンディングボックスを描画しました。一番左の猫を重ねて犬とも判定してしまいましたが、今回はオブジェクトの一取得が目的なので気にしないことにします。

推論結果

ターミナルに上記画像で表示されているバウンディングボックスの情報を出力することができました。各オブジェクトの位置、クラスが正しく取得できています。

Fusing layers... 
YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients
Adding AutoShape... 
name = cat xmin = 444.32781982421875 ymin = 155.30255126953125 xmax = 761.1204833984375 ymin = 761.1204833984375
name = cat xmin = 653.5517578125 ymin = 57.53466796875 xmax = 797.1552124023438 ymin = 797.1552124023438        
name = cat xmin = 93.90316772460938 ymin = 343.8182373046875 xmax = 781.89013671875 ymin = 781.89013671875
name = dog xmin = 94.64581298828125 ymin = 335.69921875 xmax = 779.980712890625 ymin = 779.980712890625

まとめ

今回はバウンディングボックスの情報を取得する方法を解説しました。PyTorchを使うと非常に簡単にバウンディングボックスの情報を取得できることがわかりました。

今後は、検出したオブジェクトの位置情報を活用してロボットの追跡や障害物回避などに活用していきたいと思います。

機械学習を効率よく学びたい方には、自分のペースで動画で学べるUdemyの以下の講座がおすすめです。数学的な理論からPythonでの実装までを習得できます。(私自身もこの講座を受講しています)

icon icon 【徹底的に解説!】人工知能・機械学習エンジニア養成講座(初級編~統計学から数字認識まで~) icon

それでは、また次の記事でお会いしましょう。

\ 機械学習を学びたい人には自宅で学べるUdemyがおすすめ! /

講座単位で購入できます!

参考

よかったらシェアしてね!
  • URLをコピーしました!
目次