torchhubを使ってYOLOv5を動かしてみる

はじめに

この記事はtorchhubを用いたYOLOv5の実行について解説します．
YOLOv5は物体検知モデル (画像から物体の位置と座標を推定するモデル)の一つです．
YOLOv5はYOLOv3などと比較して非常に簡単に動かすことができ，また様々な入力に対応しています．

torchhubを用いたYOLOv5の実行

torchhubを用いたYOLOv5の実行はドキュメントに従います．

必要パッケージのインストール

pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

モデルの読み込みと実行

以下のコードを書き，実行します．これで完了です！

import torch

# Model
model = torch.hub.load("ultralytics/yolov5", "yolov5s")

# Image
im = "https://ultralytics.com/images/zidane.jpg"
# あるいは，ローカルのパスを指定してもOKです．im = '<path to file>'

# Inference
results = model(im)

results.pandas().xyxy[0]

補足

modelに入力するファイルは，以下のものが利用可能です．

• ローカルのパス
• 画像のURL
• cv2のデータ (ただしあとに示す例のようにRGBの順番にする)
• pandas
  また，推論はまとめて行うことも可能です．

具体例

im1 = Image.open("test1.jpg")  # PIL image
im2 = cv2.imread("test2.jpg")[..., ::-1]  # OpenCV image (BGR から RGB　に変換)
im3 = "https://sh-anime.shochiku.co.jp/bluearchive-anime/assets/images/common/index/kv_2.jpg" # URL
# img4 = "<画像のパス>" # 画像のパス

# 推論
results = model([im1, im2,im3], size=640)  # batch of images

# Results
results.print()
results.save()  # or .show()

results.xyxy[0]  # im1 predictions (tensor)

出力例