소개

최근 물체 검출(Object Detection)에 접할 기회가 있었다는 것, YOLO와 YOLOv3를 Darknet과 Darkflow로 시험해 보면 정밀도가 높고 놀랐던 것, 요즘은 YOLOv5가 있다는 것을 알게 된 적이 있었다 .
그래서 이번에는 YOLOV5를 시도해 보겠습니다. 라이센스는 GPL (GNU GENERAL PUBLIC LICENSE).

Clone한 디렉토리를 그대로 작업용 루트 디렉토리로서 취급하는 것은 좋아하지 않기 때문에, Clone한 디렉토리의 1개상의 계층을 루트 디렉토리로 하고 있다.
Python은 최근 시작한지 ​​얼마 지나지 않아서 쓸모없는 코드이지만 나쁘지 않습니다.

절차

작업 디렉토리를 루트 디렉토리로 작성 및 이동

작업 디렉토리에 https://github.com/ultralytics/yolov5 를 Clone

./yolov5/hubconf.py 및 ./yolov5/requirements.txt를 루트 디렉토리로 이동

설치하기 위해 pip install -r requirements.txt 실행 (필요한 경우 CUDA도 설치)

hubconf.py의 from에 yolov5.~~ 를 더한다(28~31행째)

학습 된 모델을 루트 디렉토리의 임의의 위치에 넣습니다 (이번에는 ./pt/yolov5s.pt)

마인. 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는 대단해.

참고