ultralytics yolov5 사용법 완벽 정리

 

yolov5 모델을 손쉽게 학습(train) 및 추론(inference) 할 수 있는 방법이 뭐가 있을까요?

바로  ultralytics yolov5 를 사용하는 것입니다.

 

이번 포스팅에서는  ultralytics yolov5 를 통해 학습 및 추론 하는 방법에 대해 알아보겠습니다.

 

1. 기본 세팅

1) 파일 설치

 

github 주소 : https://github.com/ultralytics/yolov5

깃헙 주소를 클론 하고 필요한 라이브러리를 다운로드 해줍니다.

 

2) 데이터 셋

(1) 디렉토리 구성 맞추기

ultralytics가 데이터를 읽고 사용하기 위한 고유 규칙이 있습니다.

 

 

data라는 가장 큰 디렉토리를 기준으로 위와 같이 트리 형식을 갖춰야 합니다. 

 images 와  labels  가 같은 디렉토리 level에 위치해야 한다는점이 포인트인데요.

 images 는 이미지,  labels  에는 annotation 파일을 넣어줘야 합니다.

그리고 각각의 파일들은  train ,  validation 으로 나누어 분류 해줍니다.

 

 

 

(2) yaml파일 생성

위에서 생성한 데이터셋 정보를 담고 있는 파일을 작성해줘야 합니다.

생긴 모습은 아래와 같은데요.

 

 

 train ,  validation ,  test  데이터셋이 어느 경로에 있는지 적어줍니다.

우리 예제에서는 상대경로 형태로 적어주었는데 절대 경로로 작성하여도 무방합니다.

 

'  labels  파일경로는 왜 없지?' 라는 궁금증이 드시는 분들도 있을텐데요.

앞서 설정한 디렉토리 구조가 고정되어있기 때문에   labels 위치를 찾아낼 수 있습니다.

 

다음으로 사용될 클래스 갯수와 클래스 이름을 리스트형태로 작성해줍니다.

 

 

이렇게 작성한 yaml파일은 저장해주면됩니다. 참고로 저장위치는 아무 곳이든 상관 없습니다.

 

2. train / inference 방법

1) train

아래 코드가 모델을 학습시킬때 사용되는 코드입니다.

!cd yolov5; python train.py --img 1280 --batch 8 --epochs 100 --data ../data2/defect.yaml --weights yolov5x.pt \
                            --project ../ultra_workdir3 --name defects --exist-ok

 

(1) 이미지 크기 지정

입력 이미지 크기를 resize 하는 부분입니다.

604, 1280 ..  640씩 더한 형태로 이미지 크기를 지정해주면 되는데요.

 

이미지 크기가 클수록 정확도는 올라가지만 그만큼 속도가 떨어집니다. 데이터 상황과 사용할 모델 weights에 따라 크기를 지정해주는걸 추천합니다.

 

✔ YOLOv5n (nano), YOLOv5s (small), YOLOv5m (medium), YOLOv5l (large), YOLOv5x (extra large) => 640 x 640 추천
✔ YOLOv5n6, YOLOv5s6, YOLOv5m6, YOLOv5l6, YOLOv5x6 => 1280 x 1280 추천

 

(2) batch 크기 지정

batch 크기는 원하는 크기만큼 지정해주면됩니다.

다만 한번에 많은 배치를 적용할경우 메모리가 터질 위험이 있으니 주의해야 합니다.

 

 

(3) epoch 지정

몇번의 epoch을 돌릴지 적어주는 부분입니다.

 

 

(4) yaml 파일 경로 지정

데이터셋 정보를 작성한 yaml파일 경로를 적어줍니다.

 

 

(5) pretrained weights 파일 읽기

 

 

ultralytics에는 coco dataset 으로 pretrained 된 weights를 제공하는데요.

아래로 갈수록 모델이 성능이 좋지만 speed가 떨어지는 특징이 있습니다.

 

우리 예제에서는 YOLOv5x를 사용하겠습니다.

 

 

처음 실행시 해당 .pt파일이 없다면 다운로드 작업을 선행합니다.

 

(6) 학습 완료된 weight 파일 및 성능 평가 데이터 저장

 

 

학습후 최종 weights파일과 모델 성능 평가 데이터를 저장할 위치를 지정 해줘야 합니다.

 

이때  --project  와  --name 인자가 사용됩니다.

 --project  는 상위폴더  --name 은 하위폴더를 의미하는데요. 우리 예시를 기준으로 보자면   ../ultra_workdir3/defects  디렉토리 밑에 파일들이 저장됩니다.

 

 --name 은 필수 지정 인자는 아닙니다. 지정 하지 않으면 default로 exp라는 이름의 폴더를 생성합니다.

 

 --exist-ok 를 적어주지 않으면 해당 코드를 실행할때마다  새 파일이 계속 생성됩니다.

 

ex)

 ../ultra_workdir3/exp 

 ../ultra_workdir3/exp2  

 ../ultra_workdir3/exp3  

 

디렉토리가 지저분해지는게 싫으면 최신 파일이 기존파일을 덮어쓰는  --exist-ok 를 적용해줍니다.

 

학습이 완료된 후 지정한 경로에 아래와 같은 형태의 파일들이 생성됩니다.

이 중 weights  폴더 안에 우리가 학습한 weight파일이 저장됩니다.

 

 

.pt 파일 형태로 저장되는데 최종 weight(last.pt)와 가장 성능이 좋았을때 weight(best.pt)파일로 구성됩니다.

 

 

 best.pt  파일이 저장된경로를 아래와 같이 적어주고 train 시키면 finetuning이 가능합니다.

 

 

2) 추론(inference)

 python detect.py  파일을 실행 시켜 주면 됩니다.

 

!cd yolov5; python detect.py --source ../data2/images/val/save_221026_164245.png --weights ../ultra_workdir3/defects/weights/best.pt --conf 0.02 \
                        --project ../ultra_workdir3/output --name=run_image --exist-ok --line-thickness 2

 

앞서 다룬  train.py  과 다른점을 기준으로 설명하자면 ,

 source 에는 inference하고자 하는 이미지 경로를 적어줍니다.

 

 

 conf  는 confidence score를 의미하고 0.02 이상일때 바운딩 박스를 그리도록 세팅한것입니다.

 

 

 line-thickness  는 바운딩박스 라인 두께를 의미합니다.

 

 

3. .pt => .onnx 변경

ultralytics에서 학습된 weights를 opencv와 같은 다른 프레임워크에 사용하고 싶습니다. 이때는 weight파일 확장자를 변경할 필요가 있는데요.

예제에서는 .onnx로 변경하는 법에 대해 알아보겠습니다.

 

 export.py  파일을 실행시켜 줍니다.

 

!cd yolov5; python export.py --weights ../ultra_workdir4/defects/weights/best.pt --include torchscript onnx

 --weights 인자로 best weight 경로를 적어준뒤   --include torchscript onnx 를 추가로 작성합니다.

 

변환된 .onnx weight는  ../ultra_workdir4/defects/weights/ 하위  best.onnx 이름으로 저장됩니다.