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

YOLOv8 사용해보기

이자다 2023. 9. 17. 01:28
반응형
!pip install ultralytics
#yolov8과 yolov8 실행에 필요한 라이브러리 설치

 

yolov8을 사용하기 위한 ultralytics 패키지를 설치한다

 

 

import ultralytics

ultralytics.checks()

 

설치한 ultralytics 패키지를 import하고 check() 함수로 제대로 설치 되었는지 확인한다.

 

 

from ultralytics import YOLO

#YOLO('yolov8n.pt')로 모델 다운로드.
model = YOLO('yolov8s.pt')
#MS COCO Dataset 사전학습 된 yolov8s 모델을 로드함. s 이외에도 n, m, l, x 등이 있음.
#모델 크기가 증가할수록 정확도도 증가하지만 메모리 사용률 등 부수적인 부분에서 트레이드 오프 관계가 존재한다.

 

ultralytics 패키지에서 YOLO 모델을 불러온다.

 

불러오는 모델은 yolov8s 모델이다.

 

모델 크기 때문에 개인이 훈련시킬 때는 n, s 크기가 권장된다.

 

 

print(type(model.names), len(model.names))

print(model.names)
#MS COCO Dataset에 정의되어 있는 클래스 개수와 종류는 model.names를 통해서 확인 가능함. (총 80개, 0~79)

 

모델에 대해서 조회해본다

 

 

!mkdir test_images
#content 밑에 디렉토리 생성 후 객체탐지할 이미지 넣기

 

객체탐지할 이미지를 넣을 디렉토리를 생성한다

 

 

results = model.predict(source='/content/test_images/*.JPG', save=True)
#확장자가 소문자(jpg)가 아니라 대문자(JPG)면 구분해서 작성해야함. 대소문자 틀리면 오류.

 

사진에 대해서 객체탐지를 진행한다

 

 

from glob import glob

file_path = glob('/content/runs/detect/predict/*.JPG')

print(file_path)

 

 

 

from google.colab import files

for file_num in file_path:
  files.download(file_num)

 

 

 

import numpy as np

# results = model.predict(source='/content/test_images/*.JPG', save=True)
for result in results:
  uniq, cnt = np.unique(result.boxes.cls.cpu().numpy(), return_counts=True) #Torch.Tensor -> numpy
  uniq_cnt_dict = dict(zip(uniq, cnt))

  print('\n{class num: counts} = ', uniq_cnt_dict, '\n')

  for c in result.boxes.cls:
    print('class_num=', int(c), ', class_name = ', model.names[int(c)])
    #객체의 종류를 나타내는 고유값이 result.boxes.cls에 저장되어 있으며, 현재는 YOLOv8 모델을 그대로 사용했기 때문애
    #MS COCO Dataset에서 정의한 0~79 값이 디폴트로 사용되고 있음.

 

위 코드를 실행 후 출력되는 결과의 일부를 캡쳐했다.

 

객체탐지 수행한 사진 1, 2번의 정보
1번 사진

 

{class num: counts}는 사진에 나타난 클래스의 번호와 개수를 의미한다. 

 

사진에서 car 클래스는 2개가 인식되었다.

 

객체탐지 결과 출력은 영상에서 객체를 탐지할 때도 사용이 가능하다. 

 

현재는 mac에서 영상 detect시 프로세스가 도중에 죽어버리는 오류가 발생해서 수행하지 못했지만 이 부분은 yolov5보다 확실하게 편리하다고 볼 수 있는 부분 같다.

반응형