PyTorchとYOLOv5を使用して、画像の物体検出を行い
物体の種類・左上のxy座標・幅・高さを求めてみます。
YOLOv5はCOCO datasetを利用しているので、全部で80種類の物体を検出できます。
実行環境
- Windows10
- Python 3.9.5
- YOLOv5 version5.0
動作環境準備
動作環境を設定するためpipコマンドを使用します。
1 |
pip install -qr https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt |
pipenvを利用する場合は下記のrequirements.txtを取得後、pipenv installを使用します。
https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt
1 |
pipenv install -r requirements.txt |
PyTorchとYOLOv5で物体検出
事前準備が完了したら物体検出を行います。
今回は下記の画像から人が検出できるか試してみます。
Photo by rawpixel.com from Pexels
物体検出コード
下記のコードで物体検出を行います。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
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() # 検出した物体の切り取り |
コード解説
1 |
model = torch.hub.load("ultralytics/yolov5", "yolov5s", pretrained=True) |
PyTorch Hubから学習済みのモデルをダウンロードします。
既にモデルをダウンロード済みの場合は、自動的にダウンロード済みのモデルを利用します。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
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すると検出できる物体の種類が出力できます。
1 2 |
results = model("sample_picture.jpg") # 画像パスを設定し、物体検出を行う objects = results.pandas().xyxy[0] # 検出結果を取得 |
画像パスを指定しYOLOv5で物体検出を行います。
検出結果をobjectsに入れます。
1 2 3 4 5 6 7 8 9 10 |
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座標を引いて、物体の幅と高さを計算しています。
1 |
results.show() # 検出した物体の表示 |
検出結果のshowメソッドを使うと検出結果を表示します。
1 |
results.crop() # 検出した物体の切り取り |
cropメソッドを使うと、検出した結果を切り出して保存してくれます。
検出結果はrunsフォルダのexpフォルダ内に保存されます。
参考サイト
GitHub - pytorch/pytorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration
Tensors and Dynamic neural networks in Python with strong GPU acceleration - GitHub - pytorch/pytorch: Tensors and Dynamic neural networks in Python with strong...

PyTorch
An open source machine learning framework that accelerates the path from research prototyping to production deployment.

YOLOv5で物体検出 (PyTorch Hubからダウンロード)|澁谷直樹|note
この記事で学ぶこと YOLOv5を簡単にPyTorch Hubからダウンロードして使う方法を取り上げます。 物体検出(Object Detection)で有名なモデルにYOLO(You Look Only Once)があります。 物体検出がよくわからないというかたは、「画像分類と物体検出の違いをはっきり説明します...
コメント