PyTorch: YOLOv5 모델을 이용한 이미지 객체 탐지

안녕하세요? 이번 글은 PyTorch에서 YOLOv5 모델을 이용한 이미지 객체 탐지 과정을 정리해 보겠습니다.

*실습 내용은 아래 글에 포함되어 있는 소스 코드를 그대로 사용했습니다.

PyTorch

해당 실습을 위해서는 PyTorch 1.7 이상이 설치된 Python 3.8 이상 환경에서 시작해야 합니다.

Python 버전은 아래 코드를 통해 확인하실 수 있습니다.

# Python 버전 확인
import sys
print(sys.version)

PyTorch 버전은 아래 코드를 통해 확인하실 수 있습니다.

# PyTorch 버전 확인
import torch
print(torch.__version__)

참고로 PyTorch 설치 방법은 다음 글을 참조하시면 됩니다.

PyTorch: PyTorch(파이토치) 설치하기

Python과 PyTorch 설치가 완료되었으면 YOLOv5 종속 항목을 설치합니다.

# YOLOv5 종속 항목 설치
!pip install -qr https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

자! 이제 YOLOv5 깃허브 저장소에서 모델을 로드합니다.

# 모델
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

YOLOv5 깃허브 저장소는 다음과 같습니다. YOLOv5는 2014년에 설립된 Ultralytics(울트라리틱스)가 개발했다고 합니다.

GitHub - ultralytics/yolov5: YOLOv5 ? in PyTorch > ONNX > CoreML > TFLite

실습용 이미지는 다음과 같습니다.

dnr.jpg

0.12MB

사전 훈련 모델을 사용하므로 코드는 간단합니다. 결과를 확인(results.save() 또는 .show())해볼까요?!

# 이미지
img = ['D:/GEODATA/dnr.jpg']
# 추론
results = model(img)
# 결과
# results.print()
# results.show()
# results.save() # Save image to 'runs\detect\exp'
# results.xyxy[0]  # 예측 (tensor)
results.pandas().xyxy[0]  # 예측 (pandas)

아래와 같이 xmin, ymin, xmax, ymax 좌표와 함께 confidence, class, name 결과를 확인할 수 있습니다.

PyTorch 허브에서 YOLOv5를 로드하는 보다 상세한 글은 아래 가이드를 참조하시면 됩니다.

Ultralytics의 설립자 겸 CEO, Glenn Jocher(글렌 조커) 님이 작성하신 글입니다.

Load YOLOv5 from PyTorch Hub ⭐ · Issue #36 · ultralytics/yolov5