YOLO Detector 커스텀 데이터셋 훈련하기

지난번에는 YOLO classifier model을 훈련했는데, YOLO의 정수는 역시 detector이지 않는가? 안 해보고 넘어갈 수 없지

이번에는 커스텀 데이터를 YOLO format으로 셋팅해서 object detection 훈련과 검증을 진행해 보자.

Custom data setting

yaml file setting

YOLO Detector 훈련을 위해서는, train에 사용되는 .yaml file을 작성해야 한다.

coco128.yaml file을 참고하였는데, 구성은 다음과 같다.

names:
	0: red
	1: green
test: /tld_sample/test/
train: /tld_sample/train/
val: /tld_sample/valid/

names에는 0 ~ N의 라벨과 라벨 명을 적고,

train, val, test에는 각각의 폴더 경로(절대 경로)를 적으면 된다.

write yaml file with python

python에서 yaml file을 생성하는 코드이다.

import yaml

data = {
    "train" : '/tld_sample/train/',
        "val" : '/tld_sample/valid/',
        # "test" : '/tld_sample/test/', optional 
        "names" : {0 : 'red', 1 : 'green'}}

with open('./tld.yaml', 'w') as f :
    yaml.dump(data, f)

# check written file
with open('./tld.yaml', 'r') as f :
    lines = yaml.safe_load(f)
    print(lines)

data architecture

yaml 파일에서 설정한 디렉토리 위치대로 구조를 설정하고, train/valid/test에는 각각 images, labels folder를 생성한다.

각 images, labels 파일에는 image file과 label(annotation text) file을 위치하면 되며, 여기서 훈련 시 train/valid 비율은 7.5:2.5정도로 설정했다.

Custom data train

classifier와 마찬가지로 train 자체는 간단하다. (데이터셋 셋팅만 잘 했다면…)

여기서는 YOLOv8s 모델을 사용하여 훈련할 것이다.

from ultralytics import YOLO

model = YOLO('yolov8s.pt')

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

trian의 data parameter에 앞서 생성했던 yaml 파일의 경로를 넣자.

학습 실행 결과는 runs/detect/train에 best.pt, last.pt로 저장될 것이다.

Results

• 학습 결과
  

• 새로운 이미지에 성능 테스트

  ```python
  import cv2
  from ultralytics import YOLO
  
  # model = YOLO('yolov8s.pt')
  
  model = YOLO('./runs/detect/train7/weights/best.pt')
  results = model('./red_light_test.jpg') # conf=0.2, iou ..
  
  plots = results[0].plot()
  cv2.imshow("plot", plots)
  cv2.waitKey(0)
  cv2.destroyAllWindows()
  ```

  해당 코드는 이전 YOLOv8 detector 포스팅에서도 사용했다. 테스트 과정도 매우 간단!
  

• test 1 (red lights)
  

• test 2 (green lights)
  

• test 3 (red and green lights)