[실전 예제/PyTorch/Detection] YOLO 모델을 이용한 이미지 객체 탐지

소개

이번 시간에는 PyTorch와 YOLOv5를 사용하여 객체 탐지 모델을 구현해보도록 하겠습니다. 사전 학습된 YOLOv5 모델을 불러와 이미지 내 객체를 탐지하고, 결과를 저장하는 코드를 만들어 보도록 하겠습니다.

라이브러리 설치

먼저, PyTorch와 torchvision이 설치되어 있어야 하며, YOLOv5 모델을 PyTorch 허브에서 불러오기 위해 torch.hub.load 메서드를 사용할 것입니다.

pip install torch
pip install torchvision
pip install opencv-python
pip install numpy==1.26.4
pip install Pillow

Python 구현

YOLOv5 객체 탐지 코드

이 코드에서는 탐지된 객체에 대한 바운딩 박스를 표시한 이미지가 생성됩니다. results.print()는 탐지된 객체의 클래스와 신뢰도 점수를 화면에 출력합니다. results.save()는 탐지 결과가 표시된 이미지를 지정된 폴더에 저장합니다.

import torch
import cv2
from PIL import Image

if __name__ == "__main__":
    # YOLOv5 모델 불러오기 (사전 학습된 weights)
    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
    
    # 탐지할 이미지 설정 (로컬 파일 경로)
    img_path = 'zidane.jpg'  # 샘플 이미지 URL 사용 가능
    img = Image.open(img_path)
    
    # 모델을 통해 이미지에서 객체 탐지
    results = model(img)
    
    # 탐지 결과를 화면에 출력
    results.print()
    
    # 탐지된 객체를 저장 (원본 이미지에 바운딩 박스 표시)
    results.save(save_dir='results/')

실행 방법

1. 추론 시작: 터미널에서 추론 코드를 실행합니다.

python3 infer.py

예제 출력

실행 로그

Using cache found in ***
YOLOv5 🚀 2025-1-17 Python-3.9.17 torch-1.13.1+cu117 CUDA:0 (NVIDIA ***)

Fusing layers... 
YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients
Adding AutoShape... 
image 1/1: 720x1280 2 persons, 1 tie, 1 cell phone
Speed: 5.7ms pre-process, 9.1ms inference, 1.9ms NMS per image at shape (1, 3, 384, 640)
Saved 1 image to results

탐지 결과

마무리

이와 같은 방법으로 다양한 객체 탐지 작업에 YOLOv5와 PyTorch를 활용할 수 있습니다. 다음 시간에는 커스텀 데이터셋을 가지고 모델을 학습하는 방법도 알아보도록 하겠습니다.

파이썬 정복하기