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)