yolov5トレーニング

画像(動画)データ(game.avi)

• プログラムと同じ階層に配置してください

yolov5のインストール

!pip install yolov5

準備処理

# ワーニングがうざいので消します
import warnings
warnings.filterwarnings('ignore')

# ライブラリーのインポート
import yolov5
import torch
import cv2

#GPUを使用する場合
#import os
# GPUを指定
#os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# GPU確認
#print(torch.cuda.get_device_name())

yolov5実践

1. model を読み込む

 load() でmodelをloadすることができます
 引数: 学習モデル, cpu or cuda
 学習モデルは 2021/10/18 の時点で10個ほどあります: 公式GitHub
 このmodelにはいくつかのメソッドが用意されているので気になった方は調べてみてください

.show(), .save(), .crop(), .pandas(), etc...

#環境に合わせて適宜変更してください

# CPU
model = yolov5.load('yolov5s.pt', device='cpu')

# GPU
#model = yolov5.load('yolov5s.pt', device='cuda')

2. 画像から物体検出

 先程loadしたmodelに推論したい画像データを渡せば結果が返ってきます
 yolov5 に用意されている .show() メソッドは おそらく RGB なので画像の色が正しくありません

img = cv2.imread('demoimg_0000.jpg')
result = model(img)
result.show()

 yolov5 のモデルの戻り値には注意が必要で、戻り値をそのまま imshow() してもエラーになります
 画像データだけでなく、いろいろな情報も入っているので、model から返ってきた値をそのまま画像として保存したい場合には .display(render=True) を使用します
 .display(render=True) メソッドを使用した後、 imshow() の引数に渡す画像に .imgs[0] をすることで表示させることができます

img = cv2.imread('demoimg_0000.jpg')
result = model(img)

#modelが出した情報をそのまま画像として使えるようにする
#common.py 339行
result.display(render=True)

cv2.imshow('window', result.imgs[0])

cv2.waitKey(0)
cv2.destroyAllWindows()

3. カメラ入力でリアルタイム物体検出

 opencv の VideoCapture() を使用することでカメラの映像を読み込むことができます
 引数: PCに接続されているカメラの番号(ID)
 カメラの番号は、内蔵カメラが0、さらにUSBで追加のカメラを接続すると1のように基本的には0から順番に割り当てられているはずなので使いたいカメラになるまで順番に試してください(公式ドキュメントのチュートリアルによると-1の場合もあるとのこと)

※CPU版だとかなり時間がかかります。気長に待ってください

# VideoCapture オブジェクトを取得
cap = cv2.VideoCapture(0)

size = (int(1920 / 2), int(1080 / 2))

while True:

    #カメラの画像を読み込む
    #1つ目の変数には実際に画像データが読み込めたかのかどうかの結果(True, False)
    #2つ目の変数には画像データ(1frameずつ)
    _, frame = cap.read()

    frame = cv2.resize(frame, size)

    result = model(frame)
    result.display(render=True)

    cv2.imshow('window', result.imgs[0])

    # ESCかEnterキーが押されたらループを抜ける
    k = cv2.waitKey(1)
    if k == 27 or k == 13:
        break

#カメラを解放
cap.release()
cv2.destroyAllWindows()

4. 動画を入力に物体検出

 カメラを使用するときと同じ VideoCapture() で動画の読み込みができます
 引数: 動画のパス
 cap.isOpened() は動画が読み込めた場合True、読み込めなかった場合Falseを返すメソッド

#動画のパス
path = 'game.avi'

# VideoCapture オブジェクトを取得
cap = cv2.VideoCapture(path)

#動画が読み込めている間ループ
while cap.isOpened():

    #動画の読み込み
    _, frame = cap.read()

    res = model(frame)
    res.display(render=True)

    cv2.imshow('window', res.imgs[0])

    k = cv2.waitKey(1)
    if k == 27 or k == 13:
        break

cv2.destroyAllWindows()

5. yolov5 で検出したオブジェクトの様々なデータを取得する

 yolov5モデルの戻り値に .xyxy[0] を使用すると取得することができます
 抽出される値は [xmin, ymin, xmax, ymax, confidence, class] となっています
 1つの画像内で検出したオブジェクトの数上記のリストが取得されます

[
    [xmin, ymin, xmax, ymax, confidence, class],
    [xmin, ymin, xmax, ymax, confidence, class],
    ...
    [xmin, ymin, xmax, ymax, confidence, class]
]

---

 xmin, ymin, xmax, ymax : バウンティボックス4点の座標
 confidence : 検出したオブジェクトがどれくらいの精度で合っているのか(%)
 class : 検出したオブジェクトの識別(0: person)

---

 ※ yolov5 は Pytorch を使用していて torch.Tensor 型のデータになっています
 この torch.Tensor から要素を出得するには .item() というメソッドを使用することで取得できます
 ただし、単一の要素のみに使用でき、リスト等に使用するとエラーになります

img = cv2.imread('demoimg_0000.jpg')
res = model(img)
res = res.xyxy[0]

#[xmin, ymin, xmax, ymax, confidence, class]
for xyxy in res:

    xmin = xyxy[0].item()
    ymin = xyxy[1].item()
    xmax = xyxy[2].item()
    ymax = xyxy[3].item()
    confidence = xyxy[4].item()
    class_id = xyxy[5].item()

    print(type(xyxy))
    print(xmin, ymin, xmax, ymax, confidence, class_id)

    break

<class 'torch.Tensor'>
484.95135498046875 467.1734924316406 533.1976928710938 579.107421875 0.8058066964149475 0.0