キカベン
機械学習でより便利な世の中へ
G検定対策
お問い合わせ
   

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

thumb image

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

前々回のMiDaSのチュートリアルが割とアクセスがあったので、今回もチュートリアルを記事にしました。

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

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

1. YOLOv5を使う環境設定🔝

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

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

まずは、condaでPythonの環境を作りましょう。condaに慣れていない人はMini Condaの使い方を参考にしてください。

# Python3.8のconda環境を作る
conda create -n yolov5-test python==3.8 

# yolov5-testをアクティベートする
conda activate yolov5-test

# 最新のPyTorchをインストール
conda install pytorch torchvision torchaudio -c pytorch 

#  YOLO v5の依存ライブラリをインストール
pip install -qr https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

これで準備完了です。

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

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

まず、Torch HubからYOLO v5 (yolov5sは一番小さいモデル)をダウンロードします。

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)が検出されています。

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

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

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

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

results.print()  

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

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

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

3. 結果の数値を理解する🔝

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

print(results.xyxy[0]) 

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

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

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

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]])

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

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

5番目の数字はスコア(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']

4. 全部まとめたコード🔝

全部まとめたコードは次のようになります。ここでは、二つの画像の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('2番目の画像からの検出')
print(results.xyxy[1])


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


今日はここまで。



コメントを残す

メールアドレスは公開されません。