[Object Detection] YOLOv5 커스텀 데이터 학습

YOLOv5 커스텀 데이터 학습

 

yolov5에서 git clone하거나 zip 파일로 다운한다.

 

데이터셋 구조

 

데이터는 아래와 같이 이미지 파일이 있는 이미지와 바운딩 박스 정보가 있는 txt 라벨 폴더로 구성하고 각 폴더에 부여하고자 하는 정답 폴더로 구분한다. 

 

dataset/
    ├── images/
    │   ├── class1/
    │   │   ├── img1.jpg
    │   │   ├── img2.jpg
    │   │   ├── ...
    │   ├── class2/
    │   │   ├── img1.jpg
    │   │   ├── img2.jpg
    │   │   ├── ...
    │   ├── class3/
    │   │   ├── img1.jpg
    │   │   ├── img2.jpg
    │   │   ├── ...
    ├── labels/
    │   ├── class1/
    │   │   ├── img1.txt
    │   │   ├── img2.txt
    │   │   ├── ...
    │   ├── class2/
    │   │   ├── img1.txt
    │   │   ├── img2.txt
    │   │   ├── ...
    │   ├── class3/
    │   │   ├── img1.txt
    │   │   ├── img2.txt
    │   │   ├── ...
    ├── data.yaml
    ├── train.yaml
    ├── test.yaml

 

yaml

 

학습에 참조될 yaml 파일을 data 폴더에 custom으로 생성한다.

 

path: .../yolov5-master/custom_data  # dataset root dir
train: train/images  # train images (relative to 'path') 
val: train/images  # train images (relative to 'path')

# Classes
nc: 2	# number of classes
names:	# class names
  0: zero
  1: one

 

모델 학습

 

python train.py --img 640 --batch 16 --epochs 100 --data custom.yaml --weights yolov5s.pt

--img : 네트워크에 입력되는 이미지 크기 --batch : 배치 크기 --epochs : epoch 횟수 --data : 학습에 참조될 yaml file --weights : pre-trained model에 전이 학습할 경우 (yolov5s.pt 등을 넣어주면 자동으로 다운로드 됨) --name : 학습된 모델의 저장 이름

 

모델 탐지

 

python detect.py --weights best.pt --source test/