【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からダウンロードしサンプル画像で推論する。CPUだけで大丈夫。YOLOはYou Look Only Onceの略。v5なのでバージョン5ですが、v3までの開発者が放棄しており、YOLOv5はUltralytics社が開発した。

コメント