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

Python
この記事は約3分で読めます。

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

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

実行環境

  • Windows10
  • Python 3.9.5
  • YOLOv5 version5.0

動作環境準備

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

 

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

https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

 

PyTorchとYOLOv5で物体検出

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

Photo by rawpixel.com from Pexels

物体検出コード

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

コード解説

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

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

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

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

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

 

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からダウンロード)
YOLOv5を簡単にPyTorch Hubからダウンロードして使う方法を取り上げます。 前々回のMiDaSのチ…

コメント