프로그래밍/데이터 사이언스 공부

YOLOv8 Object Tracking 실습

이자다 2023. 9. 23. 18:53
반응형
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()

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

 

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

반응형