YOLO V5 사용하기

YOLO는 속도측면에서 이점이 있는 객체 인식 모델로 알려져 있습니다.

 

3버전까지 동일 저자였으나 5버전은 중국인이 제작하였습니다.

 

github.com/ultralytics/yolov5

ultralytics/yolov5

 

3버전에 비해 성능도 높아지고 자동화가 잘 되어 있어서 사용하기 좋습니다. 또한 파이토치로 구성되어 모바일에 잘 이식된다고 합니다.

 

사용법을 알아보겠습니다.

 

git clone https://github.com/ultralytics/yolov5.git

 

yolo v5는 4가지 선행학습 모델을 제공합니다.

 

s, m, l, x 순으로 모델 크기가 커집니다. l 모델이 가장 합리적 모델이라고 말합니다.

 

합리적이란 모델 크기와 성능 향상 폭을 고려할 때 합리적이라는 뜻입니다.

 

커스텀 데이터 셋으로 학습을 하고자 할 경우 레이블링을 위하여 다음 툴을 사용해줍니다.

 

github.com/tzutalin/labelImg

tzutalin/labelImg

 

해당 툴의 yolo 버전 어노테이션은 v5를 지원합니다. 

 

이후 yaml 파일을 생성해주어야 합니다.

 

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: dataset/images/  # 118287 images
val: dataset/images/  # 5000 images

# number of classes
nc: 5

# class names
names: ['1','1','1','clubhead', 'clubshaft']

train, val의 경로와 nc 클래스 수, names 클래스 명을 각자의 커스텀 데이터 셋에 맞게 수정해줍니다.

 

이후 아래 명령어로 학습을 실행해줍니다.

 

python train.py --img 512 --batch 16 --epochs 200 --data clubhead.yaml --weights yolov5l.pt

img는 인풋 이미지 크기dlqslek. yaml 파일에 데이터 셋 경로가 있기 때문에 따로 필요는 없죠 --weights에서 선행학습 모델을 설정해줍니다.

 

아래와 같이 학습이 진행되는 것을 볼 수 있습니다.

 

학습 시 앵커 사이즈도 자동으로 탐색해주므로 어떠한 버전보다 편리하군요

 

테스트를 원할 경우 다음과 같은 명령어로 실행이 가능합니다.

python detect.py --source ./data/custom/images/val --weights ./runs/train/exp20/weights/best.pt