PyTorchとYOLOv5を使用して、画像の物体検出を行い物体の種類・左上のxy座標・幅・高さを求めてみます。
YOLOv5はCOCO datasetを利用しているので、全部で80種類の物体を検出できます。
COCO - Common Objects in Context
実行環境 #
- 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
YOLOv5 | PyTorch