メインコンテンツへスキップ
【Python】PyTorchとYOLOv5で物体の種類、座標、幅、高さを検出する

【Python】PyTorchとYOLOv5で物体の種類、座標、幅、高さを検出する

·3 分
Programming Python
かずさプログラマー
著者
かずさプログラマー
業務の作業自動化を行っています。Go、VBA、Pythonを主に使用しています。過去にはC#、VB.Net、JavaScriptも使用していました。
目次

PyTorchとYOLOv5を使用して、画像の物体検出を行い物体の種類・左上のxy座標・幅・高さを求めてみます。

YOLOv5はCOCO datasetを利用しているので、全部で80種類の物体を検出できます。

実行環境
#

  • Windows10
  • Python 3.9.5
  • YOLOv5 version5.0

動作環境準備
#

動作環境を設定するためpipコマンドを使用します。

pip install -qr https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

pipenvを利用する場合は下記のrequirements.txtを取得後、pipenv installを使用します。

pipenv install -r requirements.txt

PyTorchとYOLOv5で物体検出
#

事前準備が完了したら物体検出を行います。
今回は下記の画像から人が検出できるか試してみます。

Photo by rawpixel.com from Pexels

物体検出コード
#

下記のコードで物体検出を行います。

import torch

model = torch.hub.load("ultralytics/yolov5", "yolov5s", pretrained=True)
print(model.names)  # 検出できる物体の種類

results = model("sample_picture.jpg")  # 画像パスを設定し、物体検出を行う
objects = results.pandas().xyxy[0]  # 検出結果を取得

for i in range(len(objects)):
    name = objects.name[i]
    xmin = objects.xmin[i]
    ymin = objects.ymin[i]
    width = objects.xmax[i] - objects.xmin[i]
    height = objects.ymax[i] - objects.ymin[i]
    print(f"{i}, 種類:{name}, 座標x:{xmin}, 座標y:{ymin}, 幅:{width}, 高さ:{height}")

# 0, 種類:person, 座標x:106.36381530761719, 座標y:392.3997802734375, 幅:509.06483459472656, 高さ:513.8624267578125
# 1, 種類:person, 座標x:694.4984130859375, 座標y:410.292236328125, 幅:439.1519775390625, 高さ:497.43505859375

results.show()  # 検出した物体の表示
results.crop()  # 検出した物体の切り取り

コード解説
#

model = torch.hub.load("ultralytics/yolov5", "yolov5s", pretrained=True)

PyTorch Hubから学習済みのモデルをダウンロードします。
既にモデルをダウンロード済みの場合は、自動的にダウンロード済みのモデルを利用します。

print(model.names)  # 検出できる物体の種類

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

model.namesをprintすると検出できる物体の種類が出力できます。

results = model("sample_picture.jpg")  # 画像パスを設定し、物体検出を行う
objects = results.pandas().xyxy[0]  # 検出結果を取得

画像パスを指定しYOLOv5で物体検出を行います。
検出結果をobjectsに入れます。

for i in range(len(objects)):
    name = objects.name[i]
    xmin = objects.xmin[i]
    ymin = objects.ymin[i]
    width = objects.xmax[i] - objects.xmin[i]
    height = objects.ymax[i] - objects.ymin[i]
    print(f"{i}, 種類:{name}, 座標x:{xmin}, 座標y:{ymin}, 幅:{width}, 高さ:{height}")

# 0, 種類:person, 座標x:106.36381530761719, 座標y:392.3997802734375, 幅:509.06483459472656, 高さ:513.8624267578125
# 1, 種類:person, 座標x:694.4984130859375, 座標y:410.292236328125, 幅:439.1519775390625, 高さ:497.43505859375

objectsには検出した物体の左上のxy座標・右下のxy座標・信頼度・クラスラベル・物体名が入っています。
検出した右下のxy座標から左上のxy座標を引いて、物体の幅と高さを計算しています。

results.show()  # 検出した物体の表示

検出結果のshowメソッドを使うと検出結果を表示します。

results.crop()  # 検出した物体の切り取り

cropメソッドを使うと、検出した結果を切り出して保存してくれます。
検出結果はrunsフォルダのexpフォルダ内に保存されます。

参考サイト
#

pytorch/pytorch

Tensors and Dynamic neural networks in Python with strong GPU acceleration

Python
84868
22850

関連記事

Pythonコードをexe化
·4 分
Programming Python
エックスサーバーでDjangoを動かす
·5 分
Programming Python
【Python】 バイナリデータの画像を扱う
·3 分
Programming Python