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...
github.com
PyTorch
An open source machine learning framework that accelerates the path from research prototyping to production deployment.
www.pytorch.org
YOLOv5で物体検出 (PyTorch Hubからダウンロード)｜澁谷直樹｜note
この記事で学ぶこと YOLOv5を簡単にPyTorch Hubからダウンロードして使う方法を取り上げます。 物体検出(Object Detection)で有名なモデルにYOLO（You Look Only Once)があります。 物体検出がよくわからないというかたは、「画像分類と物体検出の違いをはっきり説明します...
note.com
コメント