4. Yolov8 모델로 커스텀 데이터 학습하기 및 예측하기

본 내용에서는 손가락의 숫자를 실시간으로 맞추는 것을 목표로 한다.

 

본 내용은 window 환경에서 진행하며, Anaconda 가상 환경을 사용한 Jupyter notebook 환경에서 진행한다. 만약 Anaconda 환경이 컴퓨터에 깔려져 있지 않다면 아래 링크를 참조합니다.

https://silinu-ai.tistory.com/3

Jupyter notebook(주피터 노트북) 설치하기

 

이전 글에서는 데이터 수집, 이미지 라벨링, 이미지 전처리 및 학습 데이터 나누기를 진행했다.

이전 글은 아래 링크로 들어갈 수 있다.

• 데이터 수집
• 이미지 라벨링
• 이미지 전처리 및 학습 데이터 나누기

 

이번에는 커스텀한 데이터와 yolov8 모델을 가지고 학습 및 예측한다.

이는 다음과 같이 진행한다.

1. Class 정보가 들어간 yaml 파일 위치 불러오기
2. Yolo 모델 불러오기 및 학습하기
3. 학습된 Yolo 모델로 예측하기

본 코드의 위치는 편의 상 A 폴더라고 지칭한다

import 할 모듈은 다음과 같다.

import os
import cv2
%matplotlib inline
from ultralytics import YOLO
import matplotlib.pyplot as plt

 

만약 import cv2에서 에러가 난다면 anaconda prompt 창의 가상 환경에서 다음과 같이 설치한다.

conda install opencv

 

만약 from ultralytics import YOLO에서 에러가 난다면 anaconda prompt 창의 가상 환경에서 다음과 같이 설치한다.

pip install ultralytics

 

 

 

1. Class 정보가 들어간 yaml 파일 위치 불러오기

이전 포스트에서 data.yaml 파일을 A 폴더에 저장하게 했다. 따라서 다음과 같이 불러온다.

path = os.getcwd()
yaml_path = os.path.join(path,'data.yaml')

 

2. Yolov8 모델 불러오기 및 학습하기

현재는 yolov8m.pt 모델을 불러왔다.

model = YOLO("YOLOv8m.pt")
print(type(model.names),len(model.names))
print(model.names)

 

그 외로 YOLO의 다른 모델을 불러오고 싶다면 아래의 리스트 중에서 불러올 수 있다.

• yolov3.pt
• yolov4.pt
• yolov5n.pt
• yolov5s.pt
• yolov5m.pt
• yolov5l.pt
• yolov5x.pt
• yolov5n6.pt
• yolov5s6.pt
• yolov5m6.pt
• yolov5l6.pt
• yolov5x6.pt
• yolov8n.pt
• yolov8s.pt 
• yolov8m.pt
• yolov8l.pt
• yolov8x.pt

각각의 내용은 아래 사이트에서 확인 가능하다.

• YOLOv5 관련 내용

https://github.com/ultralytics/yolov5

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

 

• YOLOv8 관련 내용

https://github.com/ultralytics/ultralytics

GitHub - ultralytics/ultralytics: NEW - YOLOv8 🚀 in PyTorch > ONNX > OpenVINO > CoreML > TFLite

 

 

학습은 다음 코드로 진행한다. 이때 best weight와 last weight는 자동 저장된다.

저장 위치는 학습이 다 끝난 후 save_dir : WindowPath(저장된 위치)가 출력된다.

학습할 때마다 다른 위치에 저장된다.

# weight 파일은 모델 학습 후 맨 마지막의 save_dir을 참조한다. 
# 이미지 사이즈는 416×416 사이즈로 진행한다.
model.train(data=yaml_path, epochs = 200, patience = 30, imgsz = 416)

 

 

3. 학습된 Yolo 모델로 예측하기

사전에 A 폴더에 test 폴더를 만들어서 임의의 테스트 사진을 넣어 두었다.

만약 test를 할 이미지가 없다면 패스해도 된다.

test를 할 이미지가 있다면 A 폴더에 test 폴더를 임의로 생성한 후 아래 코드와 같이 진행한다.

test_loc = os.path.join(path,"test")
results = model.predict(source = test_loc,save=True)
for result in results:
    plots = result.plot()
    fix_img = cv2.cvtColor(plots, cv2.COLOR_BGR2RGB)
    plt.imshow(fix_img)
    plt.show()

 

그러면 다음과 같이 이미지가 출력된다.