본문 바로가기
개발/AI 코드

YOLO V8 detection 간단하게 사용하기

by beomcoder 2023. 7. 6.
728x90
반응형

https://ultralytics.com/yolov8

 

Ultralytics | Revolutionizing the World of Vision AI

Ultralytics' mission is to empower people and companies to unleash the positive potential of AI. Bring your models to life with our vision AI tools.

ultralytics.com

그냥 이미지 간단하게 확인하여 검출된 내용을 확인하고 싶었다.

하지만 찾아보니까 웹캠에서만 사용하는걸 많이 포스팅하고 있어서 찾기 어려웠다.

그래서 간단하게 이미지 한장 확인하는걸 포스팅하게 되었다.

 

!pip install ultralytics
!pip install opencv-python
!pip install supervision

 

먼저 yolo v8을 사용하기 위해 ultralytics를 install 한다.

yolo v8을 사용하기 위해서는 python 3.10 버전 이상을 사용해야 좋다고 한다.

python 3.8 버전을 사용중이었는데, 오류가 발생하여 알아보고 버전 업그레이드를 해주었다.

 

import cv2
from ultralytics import YOLO

model = YOLO('yolov8n.pt')
# yolo v8n 가중치를 사용한다.
# git 코드를 분석해보니 YOLO라는 함수안에서 yolov8n.pt가 default이고,
# string으로 넘겨준 모델이 없으면 git에서 download하여 사용한다.
# 잘못된 글자라면 에러가 출력됨

cap = cv2.VideoCapture(0)
# cv2에서 웹캠을 사용하겠다.

while cap.isOpened():
    success, frame = cap.read() 
    # cv2에서 웹캠을 read한다.
    # cap.read()함수는 튜플로 2개의 인자를 리턴하는데,
    # 첫번째 인자는 frame을 가져왔는지의 bool 변수 (true, false)이고
    # 두번째 인자는 웹캠프레임 1장의 이미지를 가지고 온다.
    
    if success:
        results= model(frame) 
        # frame이미지를 모델에 넣어 결과를 가지고 온다.
        # 리턴하는값은 리스트이다.
        annotated_frame = results[0].plot()
        # 이부분을 잘모르겠는데 results[0]의 타입은 ultralytics.yolo.engine.results.Results 이다.
        # 이걸 plot()하게 되면 numpy 어레이 좌표가 나온다.
        # 이 좌표값들은 검출된 박스를 그려주는 좌표다.
        cv2.imshow("YOLOv8 Inference", annotated_frame)
        # 이 좌표값들이 표시된걸 inshow에 넣으면 박스가 그려져서 출력된다.
        
        if cv2.waitKey(1)&0xFF == ord("q"): # 키값이 들어왔는데 q라면 꺼진다.
            break
    else:
        print("error")
        break
    
cap.release() # 윕캠을 끈다.
cv2.destroyAllWindows() # cv2로 열어진 모든 창을 종료시킨다.

 

위 코드는 웹캠을 사용할때 사용하면 된다.

주석은 내가 임의로 작성했기 때문에 안봐도 된다.

 

나는 검출결과를 변수에 넣고 싶었는데 변수에 넣어도 값이 안넣어졌다.

알아보니까 값이 안담겨져 온다는것이다.

 

그래서 stackoverflow에서 열심히 검색하고, git을 찾아보고 방법을 찾았다.

import cv2
import supervision as sv
from ultralytics import YOLO

model = YOLO('yolov8n.pt')
result = model.predict(cv2.imread('temp.jpg'))[0]
detections = sv.Detections.from_yolov8(result)
labels = [[f"{model.model.names[class_id]}", confidence] for _, _, confidence, class_id, _ in detections]
labels
'''
0: 640x480 5 persons, 1 tv, 140.1ms
Speed: 5.0ms preprocess, 140.1ms inference, 10.9ms postprocess per image at shape (1, 3, 640, 480)
[['tv', 0.9154752],
 ['person', 0.860976],
 ['person', 0.84003013],
 ['person', 0.75352204],
 ['person', 0.40899122],
 ['person', 0.3967925]]
'''

 

supervision 라이브러리에서 yolov8의 결과값을 리스트에 담아준다.

내가 필요한값은 정확도와 클래스아이디이므로 2개만 받아서 labels 변수에 넣어주었다.

이렇게 하면 주석처럼 값이 받아진다.

 

지금 당장은 사람만 검출하면 되기 때문에 커스텀학습은 필요없어서 진행하지 않았다.

조만간 커스텀학습도 진행할 계획이다.

 

 

 

728x90
반응형

'개발 > AI 코드' 카테고리의 다른 글

로컬에서 BERT모델 돌려서 학습하기  (1) 2023.02.15
인공지능 koBERT 모델 학습  (3) 2023.02.11

댓글