YOLOv5+PyTorchを試してみるだけ

#はじめに

最近物体検出（Object Detection）に触れる機会があったということ、YOLOとYOLOv3をDarknetとDarkflowで試してみたら精度が高くて驚いたこと、近頃はYOLOv5があるということを知ったことがあった。
ということで今回はYOLOv5を試してみる。ライセンスはGPL（GNU GENERAL PUBLIC LICENSE）。

Cloneしたディレクトリをそのまま作業用ルートディレクトリとして扱うのは好きじゃないので、Cloneしたディレクトリの一個上の階層をルートディレクトリにしている。
Pythonは最近はじめたばかりなので拙いコードですが悪しからず。

#手順

1. 作業用ディレクトリをルートディレクトリとして作成・移動
2. 作業用ディレクトリにhttps://github.com/ultralytics/yolov5をClone
3. ./yolov5/hubconf.pyと./yolov5/requirements.txtをルートディレクトリに移動
4. 諸々インストールするためにpip install -r requirements.txtを実行（必要に応じてCUDAもインストール）
5. hubconf.pyのfromにyolov5.~~をつけ足す（28～31行目）
6. 任意の学習済モデルをルートディレクトリの任意の場所に置く（今回は./pt/yolov5s.pt）

#main.py

Cloneしたディレクトリにあるdetect.pyを実行するのでもいいですが、今回は自分でモデルを読み込んで物体検出をするため、ルートディレクトリにmain.pyを作成・編集する。

YOLOv5ではPyTorchを使うそうなので公式のドキュメントを参考に書いてみる。PyTorchのページでも説明がある。
できるだけミニマムなコードにしたつもり。

main.py

import sys
import cv2
import torch



#YOLOはYOLOでClassを作って処理を隔離
class YOLO:

    #コンストラクタでモデルを読み込み
    def __init__(self, pt_path):
        self.__model = torch.hub.load('.', 'custom', path=pt_path, source='local')

    #物体のleft, top, right, bottomを出力
    def getBoundingBox(self, frame, log=False):
        data = []
        result = self.__model(frame)
        if log:
            print('----------')
            print(result.pandas().xyxy[0])
        ndresult = result.xyxy[0].numpy()
        for v in ndresult:
            if v[5] == 0:  #今回は人だけ抽出したかったので「0=person」
                data.append([
                    int(v[0]),  #left
                    int(v[1]),  #top
                    int(v[2]),  #right
                    int(v[3]),  #bottom
                    float(v[4])  #confidence
                ])
        return data



IMAGE_PATH = sys.argv[1]  #解析する画像のパス
CONFIDENCE = 50  #信頼度の閾値 [%]



if __name__ == '__main__':

    yolo = YOLO('pt/yolov5s')  #インスタンスを生成

    image = cv2.imread(IMAGE_PATH)  #画像読み込み
    results = yolo.getBoundingBox(image)  #物体検出

    for result in results:
        #信頼度が閾値を上回っていた場合は緑色の矩形を描画
        if result[4] > CONFIDENCE / 100:
            left, top, right, bottom = result[:4]
            image = cv2.rectangle(image, (left, top), (right, bottom), (0, 255, 0), 3)

    #結果を表示
    cv2.imshow('', image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

編集後のディレクトリは以下の通り

root/
    ├ pt/
    │   └ yolov5s.pt
    ├ yolov5/
    ├ hubconf.py
    ├ main.py
    └ requirements.txt

#結果

YOLOv5のリポジトリの中にもある「zidane.jpg」を使って上のコードを走らせてみると、ちゃんと人物の検出が出来ているのが分かる。

個人的な感触としては、requirements.txtが用意されているのでDarknetやDarkflowと比べてすんなり環境構築は上手く行った印象。ただしCUDAを使う場合はPyTorchとのバージョンの兼ね合いが発生するので注意が必要。
自分はPython初心者なので、公式リポジトリのdetect.pyとかutilsフォルダ以下のファイルの中身を見てみると、とても勉強になる。
YOLOv5はPyTorchを使うと聞いていたので、PyTorchをゴリゴリ書くもんかと思っていたけど今のところそうでもなかった。
今後はYOLOv5を動画で試してみたい。

やっぱりYOLOはすごいや。

#参考