IT

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

유병혁 2021. 7. 25. 18:20

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

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

 

PyTorch

An open source machine learning framework that accelerates the path from research prototyping to production deployment.

pytorch.org

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

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

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

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

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

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

 

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

안녕하세요? 이번 글은 PyTorch(파이토치) 설치 과정을 정리해 보겠습니다. PyTorch는 공식 홈페이지는 다음과 같습니다: pytorch.org/ PyTorch An open source deep learning platform that provides a seamless..

blog.daum.net

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

YOLOv5 ? in PyTorch > ONNX > CoreML > TFLite. Contribute to ultralytics/yolov5 development by creating an account on GitHub.

github.com

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

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

? This guide explains how to load YOLOv5 ? from PyTorch Hub https://pytorch.org/hub/ultralytics_yolov5. UPDATED 15 July 2021. Before You Start Start from a Python>=3.6.0 environment with PyTorch...

github.com