YOLOv8 Object Tracking 실습

from ultralytics import YOLO

# Load an official or custom model
model = YOLO('best.pt')  # Load an official Detect model

model에 내가 훈련시킨 가중치 파일을 할당한다.

 

 

 

 

# Perform tracking with the model
results = model.track(source="/Users/user/Desktop/vscodeProject/YOLOv8_Practice/12.MP4", show=True)  # Tracking with default tracker

모델을 이용해 영상을 tracking 하는 코드다. 

 

트래커를 따로 지정 안하면 디폴트로 botsort.yaml이 지정되어 트래커로 BoT-SORT가 사용된다.

 

해당 코드를 실행하면 개체에 바운딩 박스가 씌워지고 개체에 할당된 id, 클래스 이름, 정확도가 바운딩 박스 상단에 표시된다.

 

 

메모리 사용량이 심하다.

 

현재 16GB 메모리를 사용하는데 모자라다.

 

 

 

 

# Tracking with ByteTrack tracker
results = model.track(source="/Users/user/Desktop/vscodeProject/YOLOv8_Practice/12.MP4", show=True, tracker="bytetrack.yaml")

bytetracker 라는 트래커를 사용하는 코드다.

 

 

 

트래커의 종류만 다를 뿐 메모리를 많이 잡아먹는 건 똑같다.

 

성능 비교는 나중에 해봐야 알겠지만 그 전에 이걸 내 하드웨어로 동작시키지 못하는 게 문제다.

 

프로그램을 받을 사측이야 거기서 알아서 요구되는 스펙의 컴퓨터를 마련해서 돌릴 수 있지만 난 여건상 그게 불가능하다.

 

YOLOv5와 8을 비교하는 자료들을 찾아보면 객체 탐지 면에서는 개선이 별로 없지만 객체 추적 면에서는 많이 좋다고 하던데 YOLOv5에서도 bytetrack으로 객체추적하는 자료들이 있는 걸 보면 5로도 구현할 수 있을 것 같지만 그래도 이왕이면 최신 버전을 사용해보고 싶다.

 

 

 

 

import cv2


# Open the video file
video_path = "/Users/user/Desktop/vscodeProject/YOLOv8_Practice/12.MP4"
cap = cv2.VideoCapture(video_path)

# Loop through the video frames
while cap.isOpened():
    # Read a frame from the video
    success, frame = cap.read()

    if success:
        # Run YOLOv8 tracking on the frame, persisting tracks between frames
        results = model.track(frame, persist=True, tracker="bytetrack.yaml")

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

        # Display the annotated frame
        cv2.imshow("YOLOv8 Tracking", annotated_frame)

        # Break the loop if 'q' is pressed
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break
    else:
        # Break the loop if the end of the video is reached
        break

# Release the video capture object and close the display window
cap.release()
cv2.destroyAllWindows()

opencv 를 이용해서 비디오 프레임에서 객체 추적을 실행한다.

 

 

메모리 부담은 없다시피 하지만 성능이 좋지 않다.

 

개체에 할당된 id가 계속 변한다. 추적이 제대로 되지 않는다.

 

디폴트 트래커로 해도, bytetrack으로 해도 성능은 비슷하다.

 

각 프레임의 결과를 메모리에 적재하지 않아서 메모리 부담이 줄었지만 그 대신 성능도 줄어버렸다.

 

객체 탐지력은 괜찮은데 추적 성능이 영 좋지 않다.

 

 

 

 

# 개체 경로를 나타내는 선 그리기

from collections import defaultdict

import cv2
import numpy as np

from ultralytics import YOLO

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

# Open the video file
video_path = "/Users/user/Desktop/vscodeProject/YOLOv8_Practice/12.MP4"
cap = cv2.VideoCapture(video_path)

# Store the track history
track_history = defaultdict(lambda: [])

# Loop through the video frames
while cap.isOpened():
    # Read a frame from the video
    success, frame = cap.read()

    if success:
        # Run YOLOv8 tracking on the frame, persisting tracks between frames
        results = model.track(frame, persist=True)

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

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

        # Plot the tracks
        for box, track_id in zip(boxes, track_ids):
            x, y, w, h = box
            track = track_history[track_id]
            track.append((float(x), float(y)))  # x, y center point
            if len(track) > 30:  # retain 90 tracks for 90 frames
                track.pop(0)

            # Draw the tracking lines
            points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
            cv2.polylines(annotated_frame, [points], isClosed=False, color=(0, 127, 255), thickness=10)

        # Display the annotated frame
        cv2.imshow("YOLOv8 Tracking", annotated_frame)

        # Break the loop if 'q' is pressed
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break
    else:
        # Break the loop if the end of the video is reached
        break

# Release the video capture object and close the display window
cap.release()
cv2.destroyAllWindows()

 

차량 경로를 따라 선을 그리는 코드이다.

 

이전 코드와 같이 opencv를 사용했는데 객체추적 성능도 좋으면서 메모리 부담도 적다.

 

코드를 보면 한 프레임마다 바운딩 박스 값을 받아서 점을 찍어주는데 이 과정이 정확도에 영향을 준 것일까?

 

내가 opencv를 안다뤄봐서 저 코드가 점 찍는거 말고 어떤 부분에서 성능을 개선해줬는지 모르겠다. 알아볼 필요가 있다.

 

이 부분을 이해한다면 메모리 부담 없이 프로젝트 개발이 가능할 것 같다.

 

 

 

 

# 영상 두개를 입력받아 멀티스레드로 동시에 개체 추적하는 코드
# 하드웨어 문제인지 실행하면 세션이 종료됨.

import threading

import cv2
from ultralytics import YOLO


def run_tracker_in_thread(filename, model):
    video = cv2.VideoCapture(filename)
    frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    for _ in range(frames):
        ret, frame = video.read()
        if ret:
            results = model.track(source=frame, persist=True)
            res_plotted = results[0].plot()
            cv2.imshow('p', res_plotted)
            if cv2.waitKey(1) == ord('q'):
                break


# Load the models
model1 = YOLO('best.pt')
model2 = YOLO('best.pt')

# Define the video files for the trackers
video_file1 = '/Users/user/Desktop/vscodeProject/YOLOv8_Practice/12.MP4'
video_file2 = '/Users/user/Desktop/vscodeProject/YOLOv8_Practice/DJI_0521.MP4'

# Create the tracker threads
tracker_thread1 = threading.Thread(target=run_tracker_in_thread, args=(video_file1, model1), daemon=True)
tracker_thread2 = threading.Thread(target=run_tracker_in_thread, args=(video_file2, model2), daemon=True)

# Start the tracker threads
tracker_thread1.start()
tracker_thread2.start()

# Wait for the tracker threads to finish
tracker_thread1.join()
tracker_thread2.join()

# Clean up and close windows
cv2.destroyAllWindows()

두 개의 영상을 동시에 개체 추적을 할 수 있는 코드다

 

하드웨어 문제인지 실행하면 세션이 종료되어서 수행할 수 없다