YOLOv8 object tracking(detecting)

오늘 포스팅할 내용은 YOLO입니다. YOLO는 이미지 혹은 이미지들로 이뤄진 비디오에서 물체를 탐지에 있어서 SOTA모듈입니다.

 

Python에는 OpenCV가 유명하긴하지만 물체에 대한 트랙킹을 진행할때 비디오 내에서 물체가 사라졌다가 다시 등장했을 경우 별도의 지정없이 트랙킹 하는 것이 어렵습니다. 물체를 다시 감지하기 위한 코드를 작성해야하는데 여간 쉬운일이 아닙니다. 그래서 자동으로 다시 트랙킹 해줄 수 있는 방법을 찾다가 발견한 것이 YOLOv8입니다.

 

YOLO에서 제공해주는 기본 모델은 아주 기본적인 트랙킹 데이터를 갖고 있지만 결국은 사용자가 원하는 물체에 대한 트랙킹을 해주지는 못합니다. 이게 무슨 말이냐면 기본적인 사람, 물체와 같은 트랙킹은 기본 모델에서 학습이 되어 있기 때문에 가능하지만 예를들어서 미니카에 대한 트랙킹이 필요한 영상 혹은 사진을 갖고 있다고 해보겠습니다. 이 경우 기본 모델에서는 아무리 사진을 구워 삶아봐도 트랙킹을 하지 못합니다. 물론 잡긴잡습니다만 놓치는 경우가 허다하고 미니카로 인식해서 트랙킹을 하는것이 아닙니다. 그래서 결국 내가 원하는 물체에 대한 트랙킹을 해줄 수 있게하려면 YOLO 모델에 맞는 fine-tuning과정이 필요합니다. 제가 지금까지 다뤄왔던 자연어처리를 위한 fine-tuning과정과 크게 다를게 없습니다. YOLO 모델에 맞는 데이터셋만 맞춰주면 그 과정 자체는 똑같습니다.

 

이번 포스팅에서 제가 커스텀 자료로 만든 자료는 로스트아크의 가디언 레이드 영상을 사용합니다. 그 중에서도 하누마탄에 대한 트랙킹 작업을 진행해보겠습니다. 거두절미하고 커스텀 데이터로 fine-tuning된 모델로 하누마탄을 트랙킹 하는 장면 예시 3개를 보여드리고 그 과정에 대해서 이야기 해보겠습니다.

 

하누마탄 트랙킹_1(배틀마스터 플레이)

 

하누마탄 트랙킹_2(소울이터 플레이)

 

하누마탄 트랙킹_3(기상술사 플레이)

 

위 3개의 영상에서 프레임이 끊어져 보이는데 실제로 영상의 4프레임씩 재생시킨것도 있고 YOLO가 물체를 탐색하는 시간이 딜레이 되면서 끊어져 보이는 것입니다. 영상을 보면 YOLO가 커스텀한 물체인 하누마탄을 잘 트랙킹하고 있습니다.

---

YOLO 구현하기 in Python

 

1. 사전준비

모델에게 학습시켜줄 데이터를 만들기 위한 작업이 필요합니다. 이 부분은 각자 원하는 물체가 들어간 영상 파일을 준비해주시면 될 것 같습니다. 저는 유튜브에서 하누마탄 레이드와 관련된 영상 3개를 준비하였습니다. 3개 데이터는 각각 훈련, 테스트, 검증 데이터로 사용할 예정입니다.

 

이렇게 3개의 영상을 준비했다면 이제 프레임단위로 무작위하게 가져올 필요가 있는데 저는 opencv를 통해 96프레임 간격으로 영상에 대한 프레임을 저장해 이미지 파일을 만들었습니다. 방법은 아래아 같습니다.

import cv2

video_path = 'lostark_wb_hanu_01.mp4'

# 동영상 파일 열기
cap = cv2.VideoCapture(video_path)

frame_cnt = 0
count = 0

while True:
    frame_cnt += 1
    frame = cap.read()[1]

    if frame_cnt % 96 != 0:
        continue

    if frame is None:
        break

    cv2.imshow('tracking', frame) # 이미지 보여주기
    cv2.imwrite("./path_output_frame/%d.jpg" % count, frame) # 이미지 저장

    count+=1
    k = cv2.waitKey(30)

cap.release()
cv2.destroyAllWindows()

 

이렇게 준비된 이미지들에서 YOLO를 위한 물체의 좌표를 설정해줘야합니다. 다양한 방법들이 있지만 저는  https://app.cvat.ai/ 이 페이지를 사용하여 작업을 진행하였습니다. 별도의 비용이 들어가지는 않았고 회원가입 이후에 프로젝트를 생성하여 라벨을 추가하고 TASK를 생성하여 사전에 생성한 프로젝트를 지정하고 이미지를 업로드하면 됩니다. 그 이후는 노가다 작업으로 물체에 대한 트랙킹 범위를 이미지마다 지정해주면 됩니다. 이후 테스트로 나와 Export를 누르시면 파일을 저장할 유형을 선택할때 YOLO를 선택하고 이미지를 제외한 상태로 저장하게 되면 이미지에 대한 YOLO용 좌표 텍스트 파일이 자동으로 다운로드 됩니다. 아래 사진이 작업을 통해서 만들어진 YOLO용 파일입니다.

 

이 과정을 3개 영상 모두 작업을 해주셔야합니다. 이렇게 하면 이제 YOLO를 훈련시킬 준비과정은 끝입니다. 저는 아래와 같이 준비하였고 반드시 지켜야할 부분은 images와 labels로 구분되어 있지만 훈련을 위한 파일이름은 같아야합니다.

2. 모델 훈련

이제 위에서 준비한 데이터를 가지고 YOLO 기본 모델에 훈련을 위한 과정이 필요합니다. 그렇게 복잡한 부분은 아니기에 별도의 포스팅 없이 제가 진행한 코드를 공유하고 넘어가겠습니다. 저는 Colab의 TPU V4를 사용하여 진행하였습니다.

YOLO train in Colab

YOLO_train_custom_data.ipynb

 

3. 모델 실험

이제 훈련된 모델을 가져와서 진행하면됩니다. 위 코드를 통해 훈련을 진행했다면 project경로의 weights안에 best와 last pt파일이 있을텐데 그중에 best를 사용하면됩니다. 그러면 이제 정말 간단합니다. YOLO를 통해서 물체를 트랙킹하고 OpenCV를 통해 프레임을 보여주면됩니다. YOLO에서는 프레임에 있는 물체를 트랙킹한 정보를 담고 있어서 그 정보를 가져와 보여주기만 하면 됩니다. 핵심 코드는 다음과 같습니다. 순서대로 모델과 비디오를 불러오고, 반복문으로 프레임을 하나씩 가져와서 YOLO로 물체를 트랙킹하고 그걸 프레임으로 가져와 보여주면 끝입니다.

from ultralytics import YOLO

# Load the YOLOv8 model
model = YOLO('best.pt')

video_path = "./lostark_wb_hanu_01.mp4"
cap = cv2.VideoCapture(video_path)

while True:
    success, frame = cap.read()

    results = model.track(frame, persist=True, conf=0.6)

    boxes = results[0].boxes.xywh.cpu()

    annotated_frame = results[0].plot()

    cv2.imshow("YOLOv8 Tracking", annotated_frame)

 

전체 코드는 다음과같습니다.

import cv2
import numpy as np

from ultralytics import YOLO

# Load the YOLOv8 model
model = YOLO('best.pt')

# Open the video file
video_path = "./lostark_wb_hanu_01.mp4"
cap = cv2.VideoCapture(video_path)

frame_cnt = 0
while cap.isOpened():
    frame_cnt += 1
    success, frame = cap.read()

    if frame_cnt % 4 != 0:
        continue

    if success:
        results = model.track(frame, persist=True, conf=0.6)

        # Get the boxes and track IDs
        boxes = results[0].boxes.xywh.cpu()

        # Visualize the results on the frame
        annotated_frame = results[0].plot()

        # object points in boxes
        # for box in boxes:
        #    x, y, w, h = box

        cv2.imshow("YOLOv8 Tracking", annotated_frame)

        if cv2.waitKey(1) & 0xFF == ord("q"):
            break
    else:
        break



cap.release()
cv2.destroyAllWindows()

 

Reference

YOLOv8

https://docs.ultralytics.com/ko

 

YOLO Custom 훈련하기

https://blog.naver.com/PostView.naver?blogId=beyondlegend&logNo=223050797442

---

후기

이렇게 Python을 통해 물체 트랙킹하는 과정을 진행해보았습니다. 이러한 트랙킹 기술을 예시를 게임에 해보았지만 자율주행, 의학분야 등 다양한 분야에서 사용할 수 있는 아주 유용한 기술인것 같습니다. 저는 이를 통해 접근하려하는 프로젝트는 강화학습을 통해 AI가 로스트아크의 상황에 따른 스킬 운영법을 제안해 줄 수 있는가 입니다. 이 부분에 대해서는 이번 포스팅이 아니라 어느정도 진행되면 자세히 포스팅하여 다뤄보도록 하겠습니다!