본문 바로가기
AI

YOLO Custom Dataset

by Myungbin 2024. 5. 29.

YOLO를 Custom Dataset으로 학습하는 방법

Object Detection 모델을 학습시키기 위해서는 이미지와 레이블이 필요합니다. 이번 글에서는 YOLO 모델을 커스텀 데이터셋을 사용해 학습하는 방법을 단계별로 설명하겠습니다.

1. YAML 파일 생성하기

이미지에 5개의 레이블과 바운딩 박스를 가진 데이터가 있다고 가정하면, YAML 파일은 다음과 같이 구성할 수 있습니다. path에는 데이터셋의 루트 디렉토리가 들어가고, train, val에는 해당 디렉토리의 경로를 작성하면 됩니다.

path: ./data/  # dataset root dir
train: images/train/  # train images 
val: images/val/  # val images
test:  # test images (optional)

nc: 5
names:
  0: label_0
  1: label_1
  2: label_2
  3: label_3
  4: label_4

2. 레이블 생성

YOLO를 학습시키기 위해서는 YOLO 포맷의 어노테이션 파일이 필요합니다. 각 이미지당 하나의 .txt 파일을 가져야 합니다. 이미지에 객체가 없으면 파일이 필요 없습니다.

YOLO 포맷은 다음과 같이 구성됩니다:

  1. 객체당 하나의 행(row)
  2. 각 행은 class x_center y_center width height 형식입니다.
  3. 상자 좌표는 정규화된 xywh 형식(0에서 1까지)을 사용합니다. 픽셀 단위로 된 상자는 x_center와 width를 이미지 너비로, y_center와 height는 이미지 높이를 기준으로 변환합니다.
  4. 클래스 번호는 0부터 인덱싱됩니다

예를들어,

0 0.847238 0.448454 0.188275 0.117526
3 0.777903 0.583505 0.333709 0.086598
4 0.800451 0.665979 0.281849 0.041237
1 0.153890 0.887629 0.233371 0.068041
2 0.608230 0.879381 0.560316 0.039175

3. 디렉토리 구성

학습(train) 및 검증(val) 이미지와 레이블을 구성해야 합니다. YAML 파일에서 작성한 형식과 동일하게 폴더를 구성해야 합니다. YOLO는 YAML에서 작성한 이미지 경로의 마지막 인스턴스를 레이블로 대체하여 각 이미지에 대해 자동으로 레이블을 찾습니다.

../data/images/im0.jpg  # image
../data/labels/im0.txt  # label​

 

 

끝.

'AI' 카테고리의 다른 글

EasyOCR 간단 리뷰(Recognition)  (0) 2024.07.04
ollama baseline  (0) 2024.04.19
LLaMA-Factory Baseline  (1) 2024.04.18