YOLOv5で物体検出 (PyTorch Hubからダウンロード)

この記事で学ぶこと

YOLOv5を簡単にPyTorch Hubからダウンロードして使う方法を取り上げます。

物体検出(Object Detection)で有名なモデルにYOLO（You Look Only Once)があります。

物体検出がよくわからないというかたは、「画像分類と物体検出の違いをはっきり説明します」もどうぞ。

YOLOv5を使う環境設定

今回は、Ultralyticsが開発したYOLOのバージョン５をTorch Hubからダウンロードをして使ってみます。

https://pytorch.org/hub/ultralytics_yolov5/

まずは、pipでPythonの環境を作りましょう。

mkdir yolov5_test
cd yolov5_test

# 環境設定
python3 -m venv venv
source venv/bin/activate

pip install --upgrade pip
pip install torch torchvision

# YOLOv5の依存ライブラリをインストール
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt 

これで準備完了です。ちなみに、Github上のrequirements.txtを直接にアクセスして依存ライブラリをインストールすることもできます。この際はリポジトリをクローンする必要はありません。

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

YOLOv5をダウンロードして使ってみる

非常に簡単に使えるようになっているのでサンプル画像で推論してみましょう。

まず、Torch HubからYOLO v5 (yolov5sは一番小さいモデル)をダウンロードします。PythonのスクリプトかJupyterノートブックなどでコードを書きます。

import torch

# Torch HubからYolao V5をダウンロード
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

一度ダウンロードすると、~/.cache/torch/hub/ultralytics_yolov5_master/に格納されるので次回からは素早くロードできルようになっています。

次に、画像のある場所をURLで指定します。

# 画像のURL
base_url = 'https://github.com/ultralytics/yolov5/raw/master/data/images/'
image_url = base_url + 'zidane.jpg'

もちろん、ローカルにある画像を指定することもできます。その場合は、画像のある場所のpathを指定します。

バッチを作って推論を実行します。

# バッチ
imgs = [image_url]

# 推論の実行
results = model(imgs)

あとは、結果を表示するだけです。

# 結果を表示
results.show()

人物(person)とネクタイ(tie)が検出されています。

結果をセーブするには、以下のように１行付け足すだけです。

# 結果をセーブするs
results.save() 

これで、resultsフォルダーに結果の画像がセーブされます。

また、どのクラスがいくつ検出されたのかを以下のように表示することができます。

results.print()  

結果は次のように二人の人物(2 persons)とネクタイ(1 tie)と出ました。

image 1/1: 720x1280 2 persons, 1 tie

使うだけならこんなにも簡単です。

結果の数値を理解する

検出された四角をBounding Boxと呼び、表示することもできます。

print(results.xyxy[0]) 

xyxyというのはBounding Boxの左上のx軸とy軸の値と右下のx軸とy軸の値という意味です。

今回はバッチに画像が一つだけなので、１番目のインデックスである０を指定して値を取り出しています。

３つのBounding Boxがあるので３行の結果が表示されます。

tensor([
[7.50913e+02, 4.40350e+01, 1.15914e+03, 7.11258e+02, 8.19785e-01, 0.00000e+00],
[1.16045e+02, 1.98460e+02, 1.06555e+03, 7.13890e+02, 5.66767e-01, 0.00000e+00],
[4.16688e+02, 4.29202e+02, 5.16798e+02, 7.17052e+02, 5.66123e-01, 2.70000e+01]])

よくみると、１行に６個の数字がありますね。

最初の４つが左上と右下のxyの値です。

５番目の数字はスコア(confidence score)で、最後の数値がクラスの値になります。

見やすいように、数値を切り捨て小数点を調節してみました。

[ 750,  44, 1159, 711, 81%,  0 ]  #  0 - person
[ 116, 198, 1065, 713, 56%,  0 ]  #  0 - person
[ 416, 429,  516, 717, 56%, 27 ]  # 27 - tie

YOLOでは80の異なる物体を検出できます。それが数値で表現されていて、たとえば27はネクタイ(tie)になります。

どんなクラスが検出できるのかは、次のようにプリントすることができます。

# サポートされているクラス
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']

全部まとめたコード

全部まとめたコードは次のようになります。ここでは、二つの画像のURLをバッチにしています。

import torch

# Torch HubからYolao V5をダウンロード
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

# 画像のURL
base_url = 'https://github.com/ultralytics/yolov5/raw/master/data/images/'

# 画像二つのバッチ
imgs = [base_url + f for f in ('zidane.jpg', 'bus.jpg')]

# 推論の実行
results = model(imgs)

# 結果を表示
results.show()

# 画像をセーブ
results.save()

# 検出されたクラスと数を表示
results.print()

# Bounding Boxなどの表示
print('最初の画像からの検出')
print(results.xyxy[0])

print('２番目の画像からの検出')
print(results.xyxy[1])


# サポートされているクラス
print('全てのクラス')
print(model.names)

以上です。