YOLOv8을 OpenCV를 이용하여 동작시키기

YOLOv8

2023년 1월, YOLO의 새로운 버전인 YOLOv8이 출시되었다.

ultralytics_github에서 해당 코드와 설명을 확인할 수 있다.

YOLO는 You Only Looks Once의 약자로 빠른 처리속도를 이용하여 실시간 Object detection이 가능한 딥러닝 모델이다.

YOLO의 보다 더 자세한 설명

개요

과거 YOLOv3를 이용하여 ROS 프로젝트를 진행한 경험이 있다.
프로젝트를 진행할 당시에도 YOLOv3와 python, OpenCV를 사용하였고 비슷한 방식으로 새로운 모델을 사용해보고자 실습을 진행하였다.
관련 자료를 찾고자 구글 검색을 진행했고 내가 원하는 결과물과 비슷한 output을 가지는 사이트를 발견하였다.
사이트에 올라온 코드를 참고하여 실습을 진행하였다.
해당 코드를 실행시켜보면 객체의 이름과 정확도가 출력되지 않는다. 또한 저장된 영상을 활용하여 실습을 진행하므로 이 부분을 수정하여 진행하였다.

개발환경 구축

Anaconda에서 진행하였으며 YOLOv8을 위한 가상환경을 새로 만들어주었다.

$ conda create -n yolov8 python=3.9
$ conda activate yolov8

이번 실습의 목적은 OpenCV를 사용하여 실시간으로 웹캠에서 받아오는 영상 데이터를 YOLOv8을 통해 detection을 진행하고 DeepSort를 통해 tracking까지 구현하는 것이 목적이다.

코드를 작성하기 전 관련 라이브러리를 설치하였다.

$ python -m pip install --upgrade pip
$ pip install jupyter notebook
$ python -m ipykernel install --user --name yolov8
$ pip install opencv-python=4.6.0.66
$ pip install opencv-contrib-python=4.6.0.66

가상환경을 생성한 후 pip upgrade를 진행하였으며 kernel을 연결하여 VSCode에서 작업하기 편하도록 설정하였다.
관련 자료 서칭을 하면서 대부분의 개발자들이 OpenCV를 4.6.0 버전을 사용하는 것을 확인하여 해당 버전을 설치하였고 혹시나 추가 기능을 사용할 수 있으므로 contrib도 같이 설치하였다.

Pytorch를 설치하기 전 데스크탑에서 사용하는 GPU에 맞는 CUDA sdk 지원 버전을 확인해주었다.

이후 Pytorch 공식홈페이지에서 명령어를 확인하고 그대로 입력해주었다.

$ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117

여기까지 설치가 완료되었으면 YOLOv8과 DeepSort 라이브러리를 설치한다.

$ pip3 install ultralytics
$ pip3 install deep-sort-realtime

YOLOv8을 이용한 Object detection

import datetime
import cv2
from ultralytics import YOLO

fps 계산을 위해 datetime을 import하고 OpenCV와 YOLO을 차례로 import해준다.

CONFIDENCE_THRESHOLD = 0.6
GREEN = (0, 255, 0)
WHITE = (255, 255, 255)

최소 정확도, 녹색과 흰색을 정의하였다. 최소 정확도 아래의 객체는 화면에 출력하지 않을 것이고 객체 box를 녹색으로, 이름을 흰색으로 출력할 것이다.

coco128 = open('./yolov8_pretrained/coco128.txt', 'r')
data = coco128.read()
class_list = data.split('\n')
coco128.close()

객체의 이름을 확인하기 위해 사전에 저장한 coco128.txt에서 데이터를 받아와 list로 저장한다.

model = YOLO('./yolov8_pretrained/yolov8n.pt')

사용할 YOLOv8 모델을 설정한다. 모델에 대한 설명은 링크에서 확인할 수 있다.
YOLOv8 모델의 경우 사전에 다운로드 받지 않아도 코드를 실행하면 알아서 다운로드가 된다.

cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

사용할 웹캠의 index를 입력하고 해상도를 640x480으로 설정한다.

while True:
	start = datetime.datetime.now()
    
    ret, frame = cap.read()
    if not ret:
    	print('Cam Error')
        break
    
    detection = model(frame)[0]

fps를 계산하기 위해 동작 시작 시간을 저장하고 카메라로 데이터가 제대로 수신되는지 확인한다.
웹캠 frame에서 탐지한 객체를 detection에 저장한다.

	for data in detection.boxes.data.tolist():
		confidence = float(data[4])
    	if confidence < CONFIDENCE_THRESHOLD:
    		continue
    
    	xmin, ymin, xmax, ymax = int(data[0]), int(data[1]), int(data[2]), int(data[3])
    	label = int(data[5])
    	cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), GREEN, 2)
    	cv2.putText(frame, class_list[label]+' '+str(round(confidence, 2))+'%', (xmin, ymin), cv2.FONT_ITALIC, 1, WHITE, 2)

.boxes.data.tolist()를 이용하여 [xmin, ymin, xmax, ymax, confidence, class_id]에 대한 정보를 얻는다.
미리 지정한 정확도를 만족하지 않으면 객체를 표시하지 않도록 하였다.
받아온 데이터를 이용하여 직사각형으로 객체를 표현하였고 객체의 id를 표시하였다.

	end = datetime.datetime.now()

	total = (end - start).total_seconds()
	print(f'Time to process 1 frame: {total * 1000:.0f} milliseconds')

	fps = f'FPS: {1 / total:.2f}'
	cv2.putText(frame, fps, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)

	cv2.imshow('frame', frame)

	if cv2.waitKey(1) == ord('q'):
		break

cap.release()
cv2.destroyAllWindows()

객체 탐지가 끝난 시간을 저장하고 연산을 통해 한 프레임에서 걸리는 연산 시간을 구한다.
이를 이용하여 fps를 구하고 화면 왼쪽 상단에 출력해준다.

코드 전체는 다음과 같다.

import datetime
import cv2
from ultralytics import YOLO

CONFIDENCE_THRESHOLD = 0.6
GREEN = (0, 255, 0)
WHITE = (255, 255, 255)

coco128 = open('./yolov8_pretrained/coco128.txt', 'r')
data = coco128.read()
class_list = data.split('\n')
coco128.close()

model = YOLO('./yolov8_pretrained/yolov8n.pt')

cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

while True:
    start = datetime.datetime.now()

    ret, frame = cap.read()
    if not ret:
        print('Cam Error')
        break

    detection = model(frame)[0]

    for data in detection.boxes.data.tolist(): # data : [xmin, ymin, xmax, ymax, confidence_score, class_id]
        confidence = float(data[4])
        if confidence < CONFIDENCE_THRESHOLD:
            continue

        xmin, ymin, xmax, ymax = int(data[0]), int(data[1]), int(data[2]), int(data[3])
        label = int(data[5])
        cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), GREEN, 2)
        cv2.putText(frame, class_list[label]+' '+str(round(confidence, 2)) + '%', (xmin, ymin), cv2.FONT_ITALIC, 1, WHITE, 2)

    end = datetime.datetime.now()

    total = (end - start).total_seconds()
    print(f'Time to process 1 frame: {total * 1000:.0f} milliseconds')

    fps = f'FPS: {1 / total:.2f}'
    cv2.putText(frame, fps, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)

    cv2.imshow('frame', frame)

    if cv2.waitKey(1) == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

YOLOv8 + DeepSort

DeepSort에 대한 개념은 해당 사이트에서 간단하게 확인하고 넘어갔다. 간단하게 정리하면 DeepSort는 객체 추적 프레임워크로 객체의 ID를 보다 정확하게 추적해주는 역할을 합니다.
YOLOv8로 찾은 객체에 id를 부여하여 DeepSort로 추적하는 코드를 작성하였습니다.

from deep_sort_realtime.deepsort_tracker import DeepSort

tracker = DeepSort(max_age=50)

DeepSort관련 라이브러리를 import 해줍니다. tracker로 DeepSort를 지정해주며 max_age 보다 오래 탐지되지 않으면 추적을 하지 않는다는 의미로 사용합니다.

results = []

for data in detection.boxes.data.tolist() ...


	results.append([[xmin, ymin, xmax, ymax], confidence, label])

tracks = tracker.update_tracks(results, frame=frame)

for track in tracks:
	if not track.is_confirmed():
    	continue
    
    track_id = track.track_id
    ltrb = track.to_ltrb()
    
    xmin, ymin, xmax, ymax = int(ltrb[0]), int(ltrb[1]), int(ltrb[2]) int(ltrb[3])
    cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), GREEN, 2)
    cv2.rectangle(frame, (xmin, ymin-20), (xmin+20, ymin), GREEN, -1)
    cv2.putText(frame, str(track_id), (xmin+5, ymin-8), cv2.FONT_HERSHEY_SIMPLEX, 0.5, WHITE, 2)

탐지한 객체의 데이터를 results에 저장한다.
results에 저장된 값을 이용하여 tracks를 선언하고 tracking을 하는 대상의 id를 화면에 표현해준다.

코드를 동작시키면 탐지된 객체의 id가 보이고 화면 안에서 움직이면 id가 유지되면서 추적되는 것을 확인할 수 있다.

코드 전체는 다음과 같다.

import datetime
import cv2
from ultralytics import YOLO
from deep_sort_realtime.deepsort_tracker import DeepSort

CONFIDENCE_THRESHOLD = 0.6
GREEN = (0, 255, 0)
WHITE = (255, 255, 255)

coco128 = open('./yolov8_pretrained/coco128.txt', 'r')
data = coco128.read()
class_list = data.split('\n')
coco128.close()

model = YOLO('./yolov8_pretrained/yolov8n.pt')
tracker = DeepSort(max_age=50)

cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

while True:
    start = datetime.datetime.now()

    ret, frame = cap.read()
    if not ret:
        print('Cam Error')
        break

    detection = model.predict(source=[frame], save=False)[0]
    results = []

    for data in detection.boxes.data.tolist(): # data : [xmin, ymin, xmax, ymax, confidence_score, class_id]
        confidence = float(data[4])
        if confidence < CONFIDENCE_THRESHOLD:
            continue

        xmin, ymin, xmax, ymax = int(data[0]), int(data[1]), int(data[2]), int(data[3])
        label = int(data[5])
        # cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
        # cv2.putText(frame, class_list[label]+' '+str(round(confidence, 3)) + '%', (xmin, ymin), cv2.FONT_ITALIC, 1, (255, 255, 255), 2)

        results.append([[xmin, ymin, xmax-xmin, ymax-ymin], confidence, label])

    tracks = tracker.update_tracks(results, frame=frame)

    for track in tracks:
        if not track.is_confirmed():
            continue

        track_id = track.track_id
        ltrb = track.to_ltrb()

        xmin, ymin, xmax, ymax = int(ltrb[0]), int(ltrb[1]), int(ltrb[2]), int(ltrb[3])
        cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), GREEN, 2)
        cv2.rectangle(frame, (xmin, ymin - 20), (xmin + 20, ymin), GREEN, -1)
        cv2.putText(frame, str(track_id), (xmin + 5, ymin - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.5, WHITE, 2)

    end = datetime.datetime.now()

    total = (end - start).total_seconds()
    # print(f'Time to process 1 frame: {total * 1000:.0f} milliseconds')

    fps = f'FPS: {1 / total:.2f}'
    cv2.putText(frame, fps, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)

    cv2.imshow('frame', frame)

    if cv2.waitKey(1) == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()