Yolov8 classifier model custom dataset으로 훈련하기

지난 포스팅에서는 YOLOv8 model로 detection을 해 보았는데, 이번에는 classifier model로 custom dataset 훈련부터 적용까지 해보고자 한다.

1. 데이터셋 구축

ai 허브의 기상상황 별 주행 데이터셋을 이용하여, 날씨 분류 데이터셋을 만들어 보자.

• 데이터셋 링크 : AI-Hub

원천(이미지) 데이터 및 라벨링 데이터가 제공되어 있는데, 원천 데이터만 다운받아 사용하였다.

데이터를 받아보면 알겠지만 데이터 전처리, 정제가 필요할텐데(그렇게 퀄리티가 좋지 않은 이미지가 많다) 그 부분은 이 포스팅에서는 생략한다.

원천 데이터 분류에서는 normal, haze, rain, snow 4가지 label data를 얻을 수 있을 것이다.

2. 모델 훈련

YOLOv8 훈련 시에는 다음과 같이 디렉토리 셋팅을 거쳐야 한다. class 개수는 데이터에 따라 조정하면 된다.

from ultralytics import YOLO
# import torch

model = YOLO('yolov8s-cls.yaml')
model = YOLO('yolov8s-cls.pt')

model.train(data='/custom_cls/', epochs=20)

train 부분을 실행하면 다음과 같이 에포크, Gpu mem, loss 등이 뜨게 될 것이다(예시 이미지에서는 epoch 50으로 돌렸다)

이처럼 훈련 코드는 매우 간단하고, 훈련 결과는 ./runs/classify/train/weights/best.pt 에 저장된다.

3. 모델 테스트

새로운 데이터를 입력으로 넣고 커스텀 모델을 돌려보고 싶으면, 다음 코드를 참고하면 된다.

results = model(image) 
# image에는 image path 또는 numpy array형식 이미지 입력
    for result in results:
        probs = list(result.probs)
        classes = result.names

        highest_prob = max(probs)
        highest_prob_index = probs.index(highest_prob)

        print(f"class: {classes[highest_prob_index]}({highest_prob * 100:.2f}%)")