YOLOv5 모델 사용하기

Google Colab에서 YOLOv5모델을 실행해보겠습니다.

준비 사항

먼저 조건에 맞는 pytorch 버전을 설치합니다.

!pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113

그리고 git clone을 해줍니다. 레포지토리 및 YOLOv5모델에 대한 설명은 파이토치허브에서 확인할 수 있습니다.

!git clone https://github.com/ultralytics/yolov5

requirements를 깔아줍니다.

!cd yolov5
!pip install -r requirements.txt

dependency를 설치합니다.

#install dependencies
import torch 
from matplotlib import pyplot as plt
import numpy as np
import cv2

모델 불러오기

이제 파이토치 허브에 있는 YOLOv5모델을 불러옵니다.

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

다음과 같이 입력하면 모델 안의 layer의 구성을 확인할 수 있습니다.

model

모델 적용하기

먼저 공식 문서에 있는 사진(img)을 모델에 적용해보겠습니다.

img = "https://ultralytics.com/images/zidane.jpg"
results = model(img)
results.print()

실행한 결과 다음과 같이 이미지의 크기, 인식한 객체의 종류와 개수 등이 출력된 것을 확인할 수 있습니다.

이제 이미지를 랜더링한 후 출력합니다. 출력된 이미지 결과를 보면 위에 results.print()에서 출력된 것처럼 사람 2명, 넥타이 2개가 인식된 것을 확인할 수 있습니다.

%matplotlib inline
plt.imshow(np.squeeze(results.render()))  
plt.show()

이제 로컬 파일에 있는 이미지를 모델에 적용해보겠습니다.

from google.colab import files
uploaded = files.upload()

업로드한 사진을 이전과 동일한 방법으로 모델에 적용한 후 결과를 출력하면 인식된 객체의 종류와 개수를 확인할 수 있습니다.

인식할 사물이 많은 사진을 찾다가 이 사진을 적용해봤는데, 포크와 나이프가 여러 각도로 배치되었고 맨 앞 사람은 얼굴의 반만 나왔는데도 대부분 맞게 인식되었습니다.