YOLOv8 차량 추적, 집계 실습

실습 환경은 구글 코랩에서 한다.

 

맥북의 vscode에서 하려고 했는데 실습에 사용할 ByteTrack은 yolox를 요구하고, yolox는 아래 링크에 따르면 cuda/pytorch가 설치되어야 하는데 맥북 환경에서는 무리인 것 같다.

 

https://velog.io/@khjgmdwns/Nvidia-Jetson-%EA%B0%9C%EB%B0%9C%ED%99%98%EA%B2%BD-%EC%84%A4%EC%A0%95-3-YOLOX

Nvidia Jetson YOLOX 개발환경 설정

 

실제로 의존성 패키지들 설치가 제대로 되지 않고 yolox는 제대로 Import가 되지 않는다.

 

 

 

 

import os
HOME = os.getcwd()
print(HOME)

현재 디렉토리를 변수에 할당.

 

이번 실습에서는 많이 못 썼지만 f스트링으로 꽤 쓸만하다.

 

 

 

 

from google.colab import drive
drive.mount('/content/drive')

객체 추적과 집계를 수행할 영상 자료가 있는 구글 드라이브를 코랩에 마운트

 

 

 

 

SOURCE_VIDEO_PATH = "/content/drive/MyDrive/객체 탐지 영상/12.MP4"

 

 

 

 

# Pip install method (recommended)

!pip install ultralytics

from IPython import display
display.clear_output()

import ultralytics
ultralytics.checks()

 

 

 

 

 

%cd {HOME}
!git clone https://github.com/ifzhang/ByteTrack.git
%cd {HOME}/ByteTrack

# workaround related to https://github.com/roboflow/notebooks/issues/80
!sed -i 's/onnx==1.8.1/onnx==1.9.0/g' requirements.txt

!pip3 install -q -r requirements.txt
!python3 setup.py -q develop
!pip install -q cython_bbox
!pip install -q onemetric
# workaround related to https://github.com/roboflow/notebooks/issues/112 and https://github.com/roboflow/notebooks/issues/106
!pip install -q loguru lap thop

from IPython import display
display.clear_output()


import sys
sys.path.append(f"{HOME}/ByteTrack")


import yolox
print("yolox.__version__:", yolox.__version__)

맥북에선 이 코드가 에러가 났다.

 

상술했듯이 yolox를 사용하려면 CUDA가 설치되어야 하는데 맥북은 엔비디아 그래픽카드가 없어서 설치가 불가능하다.

 

이를 코랩이 아닌 환경에서 다루려면 내 데스크톱에 CUDA를 설치해야 하는데 나중에 날 잡아서 해야할 것 같다.

 

데이터 사이언스라는게 프로그래밍 중에서 가장 돈 드는 학문이 아닐까 싶다.

 

몇 백만원 짜리 그래픽 카드가 없으면 학습에 대단한 지장이 드는게 참 힘들다.

 

 

 

 

from yolox.tracker.byte_tracker import BYTETracker, STrack
from onemetric.cv.utils.iou import box_iou_batch
from dataclasses import dataclass


@dataclass(frozen=True)
class BYTETrackerArgs:
    track_thresh: float = 0.25
    track_buffer: int = 30
    match_thresh: float = 0.8
    aspect_ratio_thresh: float = 3.0
    min_box_area: float = 1.0
    mot20: bool = False

 

 

 

 

 

!pip install supervision==0.1.0


from IPython import display
display.clear_output()


import supervision
print("supervision.__version__:", supervision.__version__)

roboflow에서 자체 개발한 supervision 패키지를 설치한다.

 

 

 

 

from supervision.draw.color import ColorPalette
from supervision.geometry.dataclasses import Point
from supervision.video.dataclasses import VideoInfo
from supervision.video.source import get_video_frames_generator
from supervision.video.sink import VideoSink
from supervision.notebook.utils import show_frame_in_notebook
from supervision.tools.detections import Detections, BoxAnnotator
from supervision.tools.line_counter import LineCounter, LineCounterAnnotator

객체 추적, 집계용으로 꽤 유용할 것 같다.

 

 

 

 

from typing import List

import numpy as np


# converts Detections into format that can be consumed by match_detections_with_tracks function
def detections2boxes(detections: Detections) -> np.ndarray:
    return np.hstack((
        detections.xyxy,
        detections.confidence[:, np.newaxis]
    ))


# converts List[STrack] into format that can be consumed by match_detections_with_tracks function
def tracks2boxes(tracks: List[STrack]) -> np.ndarray:
    return np.array([
        track.tlbr
        for track
        in tracks
    ], dtype=float)


# matches our bounding boxes with predictions
def match_detections_with_tracks(
    detections: Detections,
    tracks: List[STrack]
) -> Detections:
    if not np.any(detections.xyxy) or len(tracks) == 0:
        return np.empty((0,))

    tracks_boxes = tracks2boxes(tracks=tracks)
    iou = box_iou_batch(tracks_boxes, detections.xyxy)
    track2detection = np.argmax(iou, axis=1)

    tracker_ids = [None] * len(detections)

    for tracker_index, detection_index in enumerate(track2detection):
        if iou[tracker_index, detection_index] != 0:
            tracker_ids[detection_index] = tracks[tracker_index].track_id

    return tracker_ids

모델에서 나오는 바운딩 박스를 추적기에서 나오는 바운딩 박스와 수동으로 일치시켜야 한다.

 

 

 

 

# settings
MODEL = "best.pt"

%cd {HOME}

from ultralytics import YOLO

model = YOLO(MODEL)
model.fuse()

객체 추적을 위한 YOLO 패키지 설치.

 

모델은 내가 학습시킨 best.pt를 코랩에 업로드 해서 사용한다.

 

 

 

 

# dict maping class_id to class_name
CLASS_NAMES_DICT = model.model.names
# class_ids of interest - car, motorcycle, bus and truck
CLASS_ID = [0, 1, 2, 3, 4, 5]

내가 만든 모델의 ID는 총 6개이므로 0부터 5까지 입력한다.

 

 

 

 

# create frame generator
generator = get_video_frames_generator(SOURCE_VIDEO_PATH)
# create instance of BoxAnnotator
box_annotator = BoxAnnotator(color=ColorPalette(), thickness=4, text_thickness=4, text_scale=2)
# acquire first video frame
iterator = iter(generator)
frame = next(iterator)
# model prediction on single frame and conversion to supervision Detections
results = model(frame)
detections = Detections(
    xyxy=results[0].boxes.xyxy.cpu().numpy(),
    confidence=results[0].boxes.conf.cpu().numpy(),
    class_id=results[0].boxes.cls.cpu().numpy().astype(int)
)
# format custom labels
labels = [
    f"{CLASS_NAMES_DICT[class_id]} {confidence:0.2f}"
    for _, confidence, class_id, tracker_id
    in detections
]
# annotate and display frame
frame = box_annotator.annotate(frame=frame, detections=detections, labels=labels)

%matplotlib inline
show_frame_in_notebook(frame, (16, 16))

단일 프레임만 어노테이션을 진행해본다.

 

 

 

 

 

 

VideoInfo.from_video_path(SOURCE_VIDEO_PATH)

# settings
LINE_START = Point(50, 1500)
LINE_END = Point(3840-50, 1500)

#마운트한 드라이브로 결과물 저장
TARGET_VIDEO_PATH = "/content/drive/MyDrive/trackAndCountResults/vehicle-counting-result.mp4"

영상에 대한 추적, 집계를 진행할 시간이다. 

 

VideoInfo로 비디오의 너비, 높이를 알아내고 선의 시작점, 끝점을 그려준다.

 

 

 

 

from tqdm.notebook import tqdm


# create BYTETracker instance
byte_tracker = BYTETracker(BYTETrackerArgs())
# create VideoInfo instance
video_info = VideoInfo.from_video_path(SOURCE_VIDEO_PATH)
# create frame generator
generator = get_video_frames_generator(SOURCE_VIDEO_PATH)
# create LineCounter instance
line_counter = LineCounter(start=LINE_START, end=LINE_END)
# create instance of BoxAnnotator and LineCounterAnnotator
box_annotator = BoxAnnotator(color=ColorPalette(), thickness=4, text_thickness=4, text_scale=2)
line_annotator = LineCounterAnnotator(thickness=4, text_thickness=4, text_scale=2)

# open target video file
with VideoSink(TARGET_VIDEO_PATH, video_info) as sink:
    # loop over video frames
    for frame in tqdm(generator, total=video_info.total_frames):
        # model prediction on single frame and conversion to supervision Detections
        results = model(frame)
        detections = Detections(
            xyxy=results[0].boxes.xyxy.cpu().numpy(),
            confidence=results[0].boxes.conf.cpu().numpy(),
            class_id=results[0].boxes.cls.cpu().numpy().astype(int)
        )
        # filtering out detections with unwanted classes
        mask = np.array([class_id in CLASS_ID for class_id in detections.class_id], dtype=bool)
        detections.filter(mask=mask, inplace=True)
        # tracking detections
        tracks = byte_tracker.update(
            output_results=detections2boxes(detections=detections),
            img_info=frame.shape,
            img_size=frame.shape
        )
        tracker_id = match_detections_with_tracks(detections=detections, tracks=tracks)
        detections.tracker_id = np.array(tracker_id)
        # filtering out detections without trackers
        mask = np.array([tracker_id is not None for tracker_id in detections.tracker_id], dtype=bool)
        detections.filter(mask=mask, inplace=True)
        # format custom labels
        labels = [
            f"#{tracker_id} {CLASS_NAMES_DICT[class_id]} {confidence:0.2f}"
            for _, confidence, class_id, tracker_id
            in detections
        ]
        # updating line counter
        line_counter.update(detections=detections)
        # annotate and display frame
        frame = box_annotator.annotate(frame=frame, detections=detections, labels=labels)
        line_annotator.annotate(frame=frame, line_counter=line_counter)
        sink.write_frame(frame)

객체 추적, 집계를 위한 코드다.

 

아직 전부는 이해 못했지만 for문으로 프레임 하나씩 처리하는 것처럼 보인다.

 

아래는 결과 영상이다.

 

 

 

위 코드를 이해하고 응용하면 내가 목표로 하는 프로젝트 진행에 큰 도움이 될 것 같다.