【PyTorch】TorchHub+YOLOv5でリアルタイム認識して遊ぶ

気づいたらYOLOv5がTorchHubに追加されてたんですね、知らなかったー・・・
ということで今回は、TorchHubのYOLOv5とウェブカメラでリアルタイム認識にチャレンジしたいと思います！

実行環境

今回はCPUで実験してみたいと思います。
利用環境はこちら

• PyTorch：1.8.1+cpu
• OpenCV : 4.5.2

すっごい少ない実行環境・・・

まずはYOLOv5を召喚する

何はともあれ偉大なるYOLOv5をPyTorchから呼び出してみましょう。
公式ページ(https://github.com/ultralytics/yolov5 )に従うと

$ git clone https://github.com/ultralytics/yolov5
$ cd yolov5
$ pip install -r requirements.txt

をして環境を整えたのちに

model = torch.hub.load('ultralytics/yolov5', 'yolov5s') 

だけでyolov5sが呼び出せるらしいです。ほんまかいな・・・。
ひと昔前まではYOLOなんか有志が作ったモデルしかなく、どうやって実装するねんこれ！！！みたいな気持ちだった気がしますが、時代は進みましたね

では実際にやってみましょう。

import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s') 

Using cache found in C:\Users\toshi/.cache\torch\hub\ultralytics_yolov5_master
Downloading https://ultralytics.com/assets/Arial.ttf to C:\Users\toshi\AppData\Roaming\Ultralytics\Arial.ttf...
YOLOv5  2021-11-1 torch 1.8.1+cpu CPU

Downloading https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5s.pt to C:\Users\toshi\.cache\torch\hub\ultralytics_yolov5_master\yolov5s.pt...
100.0%

Fusing layers... 
Model Summary: 213 layers, 7225885 parameters, 0 gradients
Adding AutoShape... 

おおー、ほんとに動いてる・・・。

とりあえず使ってみる

こちらの画像を使って検証してみます。

img = cv2.imread('./photo0000-6368.jpg')
result = model(img)
print(result)

<models.common.Detections object at 0x0000022FAE025518>

ふむ、公式リポジトリのmodelsにあるcommon.pyの中のクラスDetectionsが関係しているらしいですね。
そこを少し覗いてみましょう。

たくさん関数がありますね

• display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path(''))
• print(self)
• show(self)
• save(self, save_dir='runs/detect/exp')
• crop(self, save=True, save_dir='runs/detect/exp')
• render(self)
• pandas(self)
• tolist(self)

display()はいろいろしてくれそうな関数ですね。
引数は

• pprint：何が何個見つかったか出力
• show：bboxを実際に描画
• save：showの内容をセーブすることができる
• crop：クロップされた画像を出力
• render：self.imgsに結果を描画した画像を格納する

実際に試すと

result.display(pprint=True)

#image 1/1: 485x728 21 persons, 2 backpacks, 1 handbag

result.display(show=True)

色がおかしいですね。BGRで表現されてるからでしょうか？

result.display(save=True)

これを実行すると、上の写真が保存されました。

result.display(crop=True)

[{'box': [tensor(368.28458),
   tensor(305.86414),
   tensor(389.74072),
   tensor(335.34665)],
  'conf': tensor(0.26244),
  'cls': tensor(24.),
  'label': 'backpack 0.26',
  'im': array([[[252, 227, 231],
          [145, 126, 215],
          [ 56,  56, 255],
          ...,
          [218, 247, 247],
          [ 64, 212, 209],
          [ 49, 210, 207]],
  
         [[252, 241, 209],
          [124, 116, 220],
          [ 56,  56, 255],
          ...,
          [242, 252, 252],
          [ 68, 214, 211],
          [ 49, 210, 207]],

おお、なんか有用そうなデータが。save=TrueにしているとBBoxでカットされた画像が保存されてました。なんじゃこりゃ・・・。

今回欲しいものは、BBoxが描画された画像なので.render()で十分ですね、多分きっと。

ではこれを画面キャプチャと連結してみましょう。

import torch
import cv2
import pyautogui as pag
import numpy as np

model = torch.hub.load('ultralytics/yolov5', 'yolov5s') 

while True:
    img=pag.screenshot()
    img=np.asarray(img)
    img=cv2.cvtColor(img,cv2.COLOR_RGB2BGR)

    result = model(img)
    result.render()
    cv2.imshow("test",result.imgs[0])
    k = cv2.waitKey(1) & 0xFF
    if(k == ord('q')):
        break

いざ実行！！！！

どうでしょう？？？？

CPUなのにほぼリアルタイムで動いてませんか？こんなことあってもいいんですか！？
GPUなくてもそれなりの精度で動くやん！！！

yolov5sと、小さいモデルを使って認識を行ってますが、しっかり認識できてるじゃないですか！（動画にキャプチャして見せたい・・・）

さすがにyolov5xにすると処理落ち感が否めませんが、それでもCPUで動くだけ御の字です。

ということで今回はYOLOv5で遊んでみました。次はこのYOLOv5を使って意味の分からないアプリを作ってみたいと思います！