YOLO 모델을 이용한 객체 추적과 다중 객체 검출, YOLO와 SORT(Simple Online and Realtime Tracking) 연동하기

객체 검출(Object Detection)은 이미지나 비디오의 객체를 식별하고 그 위치를 찾는 기술입니다. 최근 YOLO(You Only Look Once) 모델이 객체 검출 분야에서 혁신적인 성과를 보이며 널리 사용되고 있습니다. 본 강좌에서는 YOLO 모델을 이용한 객체 추적과 다중 객체 검출 방법, 그리고 SORT(Simple Online and Realtime Tracking) 알고리즘과 연동하는 방법에 대해 자세히 설명하겠습니다.

1. YOLO 모델 소개

YOLO는 2016년 Joseph Redmon이 처음 제안한 객체 검출 시스템으로, 전체 이미지를 한 번에 처리하여 속도와 정확성의 균형을 이룹니다. YOLO의 기본 원리는 이미지를 그리드로 나누고 각 그리드 셀에서 객체를 예측하는 것입니다. YOLO는 다양한 버전이 있으며, YOLOv3, YOLOv4, YOLOv5 등 최신 버전들이 출시되었습니다.

2. YOLO 설치와 환경 설정

YOLO를 사용하기 위해 먼저 필요한 라이브러리와 패키지를 설치해야 합니다. 아래는 Python 환경에서 YOLOv5를 설치하는 방법입니다.

!git clone https://github.com/ultralytics/yolov5.git
%cd yolov5
!pip install -r requirements.txt

3. YOLO 모델을 이용한 객체 검출

다음으로 YOLO 모델을 사용하여 객체를 검출하는 코드를 작성해보겠습니다. YOLOv5 모델을 사용하여 이미지를 검출하고, 검출된 객체를 시각화합니다.

import torch
from PIL import Image
import cv2

# YOLOv5 모델 불러오기
model = torch.hub.load('ultralytics/yolov5:v5.0', 'yolov5s')

# 이미지 로드
img_path = 'path/to/your/image.jpg'
img = Image.open(img_path)

# 객체 검출
results = model(img)

# 검출 결과 시각화
results.show()

위 코드는 특정 이미지에서 객체를 검출하고 해당 결과를 시각화합니다. YOLOv5 모델을 사용하여 ‘yolov5s’라는 경량 모델을 로드하고, 주어진 이미지를 기반으로 객체를 검출합니다.

4. 다중 객체 검출

YOLO의 장점 중 하나는 다중 객체를 동시에 검출할 수 있는 것입니다. 다음 코드에서는 비디오에서 여러 객체를 실시간으로 검출합니다.

cap = cv2.VideoCapture('path/to/your/video.mp4')

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # 객체 검출
    results = model(frame)
    
    # 결과를 이미지로 변환
    img = results.render()

    # 검출 결과 보여주기
    cv2.imshow('YOLO Object Detection', img[0])
    
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

5. 객체 추적 소개: SORT (Simple Online and Realtime Tracking)

SORT는 객체의 움직임을 추적하기 위한 간단하고 실시간으로 처리 가능한 알고리즘입니다. SORT는 Kalman Filter와 Hungarian Algorithm을 결합하여 실시간으로 객체를 추적합니다. YOLO와 결합시켜서 객체의 경로를 추적할 수 있습니다.

6. YOLO와 SORT 연동하기

YOLO로 검출한 객체를 SORT 알고리즘을 사용하여 추적할 수 있습니다. 아래 코드는 YOLO 모델에서 검출한 객체를 SORT를 이용하여 추적하는 예제입니다.

from sort import Sort

# SORT 필터 초기화
tracker = Sort()

# 비디오 캡처
cap = cv2.VideoCapture('path/to/your/video.mp4')

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # 객체 검출
    results = model(frame)
    
    # 검출된 박스 가져오기
    detections = results.pandas().xyxy[0][['xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class']].to_numpy()

    # SORT를 사용하여 객체 추적
    trackers = tracker.update(detections)

    # 추적 결과를 시각화
    for d in trackers:
        x1, y1, x2, y2, track_id = int(d[0]), int(d[1]), int(d[2]), int(d[3]), int(d[4])
        cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
        cv2.putText(frame, f'ID: {track_id}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

    cv2.imshow('YOLO and SORT Tracking', frame)
    
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

위 코드는 YOLO로 검출된 객체의 경계 상자 정보를 SORT에 전달하고, 각각의 객체를 추적합니다. 추적된 객체에는 고유한 ID가 부여되며, 비디오 프레임에 표시됩니다.

7. 성능 개선 및 최적화

YOLO와 SORT를 결합하여 객체 추적 시스템을 구축할 때 몇 가지 성능 개선 전략을 고려할 수 있습니다:

모델 최적화: YOLO의 경량 버전을 사용하는 것이 좋습니다. 예를 들어, YOLOv5s 대신 YOLOv5m 또는 YOLOv5l를 고려할 수 있습니다.
프레임 수 감소: 비디오의 프레임 수를 줄이는 것도 성능을 향상시키는 방법입니다. 예를 들어, 2초마다 하나의 프레임을 처리하도록 설정할 수 있습니다.
GPU 가속: 컴퓨터 비전 작업은 GPU에서 훨씬 더 빠르게 이루어집니다. NVIDIA GPU를 사용할 수 있는 경우 CUDA를 활용하여 YOLO 모델을 더 빠르게 실행할 수 있습니다.

8. 결론

YOLO와 SORT를 조합하여 객체 검출 및 추적 시스템을 구현하는 방법에 대해 살펴보았습니다. YOLO는 빠르고 정확한 객체 검출을 제공하며, SORT는 이를 기반으로 실시간 객체 추적을 가능하게 합니다. 이러한 시스템은 자율주행차, 보안 시스템, 스마트 시티 등 다양한 분야에 응용될 수 있습니다.

이 강좌가 YOLO와 SORT를 활용한 객체 추적 시스템을 이해하는 데 도움이 되었기를 바랍니다. 추가적으로 YOLO 모델을 훈련하고, 더 많은 객체 클래스를 추가하여 시스템을 더 발전시킬 수 있습니다.