YOLOv5 커스텀 데이터 학습하기

 

 

1. 서론

이번 글에서는 커스텀 데이터를 사용하는 가이드를 작성해보도록 하겠다.

https://jeo96.tistory.com/entry/YOLOv5-%EC%84%A4%EC%B9%98-Pycharm

YOLOv5 설치 [Pycharm]

 

2. 본론

2.1 커스텀 데이터 구성

2.1.1 파일

이미지 1장과 txt 파일 1개가 쌍을 이루어야 하며 파일 이름은 이미지 이름과 레이블 이름이 같아야 하고 레이블은 txt 파일로 구성해야 한다.

2.1.2 label 포맷

yolov5의 학습 데이터 레이블은 아래 이미지와 같이 class x y w h로 구성되어 있다. 

yolov5 label

class는 yolov5\data\coco128.yaml파일 안에 names가 있는데 클래스 번호가 0이면 names에 1번째에 적혀있는 person이다.

coco128.yaml

x y w h에서 x와 y는 객체의 중심을 나타내며 w와 y는 물체의 폭을 나타낸다. x와 y좌표는 이미지의 좌측 상단에서 시작하며 0은 이미지의 시작점, 1은 이미지의 끝점을 나타낸다. w와 h는 물체의 크기를 나타내며 0에 가까울수록 이미지가 작고 1에 가까울수록 이미지는 최대치가 된다. 이러한 형식에 맞춰 학습 데이터 레이블을 구성하여야 한다.

2.2 커스텀 데이터 학습

YOLOv5를 Pycharm으로 열고 터미널에 python train.py를 입력하여 train.py를 실행하면 coco128 데이터셋을 다운로드하여 학습하게 된다. 만약에 RAM이 부족한 경우 다음과 같은 에러나 나온다. 이러한 경우 train.py --batch (16 보다 작은 숫자)로 명령어를 입력해야 한다.

yolov5\data 디렉터리를 보면 coco128.yaml 파일이 있다. path는 학습 데이터 위치, train은 path안에 학습 데이터 이미지 디렉터리, path안에 평가 데이터 이미지 디렉터리 이름이다. nc는 클래스 개수, names는 클래스 번호에 해당하는 클래스 이름이다. 커스텀 데이터를 구성한 후 coco128.yaml파일을 복사하여 path, train, val, nc, names를 수정한다. 

coco128.yaml

train.py 파일을 보면 parser를 볼 수 있다. --data를 보면 default로 data/coco128.yaml로 설정되어 coco128.yaml을 읽어 학습한 것을 알 수 있다. python train.py --data data/custom.yaml과 같이 새로 만든 yaml파일을 지정해주면 커스텀 데이터를 학습하게 된다.