YOLOv5 Custom Model 만들기(2) - Custom Model 학습

YOLO 커스텀 모델을 만들기 2번째 단계 -> Custom Model 만들기!

 

1번째 단계인 라벨링 작업은 하단 링크를 참조하세요

YOLOv5 Custom Model 만들기(1) - LabelImg 사용법

 

프로젝트를 진행하며 좌석 상태를 탐지하는 Custom Model 1, 좌석 모양을 탐지하는 Custom Model 2를 만들어야 했다.

학습 데이터 라벨링이 끝났다면 다음과 같이 만들어준다.

 

1. 데이터 셋 준비

 

1. 직접 Train, Valid, Test Dataset 만들기

2. 코딩 기술을 이용해 Train, Valid Test Dataset 나누기

 

다음과 같은 2개 옵션 중 나는 1번을 선택했다. 

우선 그 이유는 YOLOv5를 Github에서 다운받은 후 학습을 진행할 때 Train, Valid, Test 별로 바로 진행 가능하기 때문이다.

코딩 기술을 통해 DataLoader에 탑재시켜 셔플시킬 수 있지만 귀찮았다..ㅎ

 

다음과 같은 train, valid, test가 만들어지며 train valid는 9:1 비율, test는 가장 최신 사진 1, 2 장을 넣어 성능을 확인했다.

 

또한 Train 데이터의 구성은 다음과 같은데, 사진 구성은 위에서 부터 순서대로

1. Train dir 내용물

2. images dir 내용물 -> 사진

3. labels dir 내용물 -> 사진에 대한 라벨 정보

4. classes.txt의 내용물 이다. -> labels 폴더 안의 라벨링 정보 0이 table, 1이 longtable이다.

Valid, Test 또한 유사함!

 

2. Custom Model 학습

아래 링크에서 YOLO를 다운로드 받는다.

https://github.com/ultralytics/yolov5

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

 

이제 압축 해제 후 data 폴더에 가면 각종 yaml 파일이 있다. 하나를 열어보면 다음과 같은 형태로 되어있다.

1. path 지정.

2. train path

3. val path

4. test path 

5. class 정보 -> 숫자와 그에 맞는 class 이름.

path, train, val, class 수 names를 지정해준다. (classes는 위와같은 형태로 지정해도 괜찮음)

테이블의 정보를 탐지하는 모델 이므로 table과 longtable 2개의 라벨을 가진다.

위의 경우는 리눅스 서버에서 작업을 해서 다음과 같이 지정해주었다. 

절대 경로로 지정해주는 게 마음이 편하므로 절대경로로 지정!

 

그 다음 YOLOv5폴더의 data 사이에 다시 넣어준다.

 

 

다음 아나콘다 혹은 터미널이 있는 곳에 들어가 yolov5 폴더로 들어가 준 후 train.py가 있는 곳에서 명령어를 입력해준다.

리눅스는 그냥 visual studio 터미널에서, 윈도우는 아나콘다에서 다음과 같은 작업을 진행해준다.

!python train.py --data "data/empty_or_using_dataset.ymal" --epoch 100 # 100 반복

 

여기 진행 후 경로가 안맞는다는 오류가 뜨면 

다음과 같이 폴더를 하나(여러개..?) 만들어 주고

저는 모델 2개를 만들어야 해서 위에는 테이블의 모양, 아래는 테이블 상태를 탐지하는 다른 데이터 셋을 만들어줬어용

 

다음과 같이 yaml를 새로 만들어서 다음과 같이 넣어주고 다시 명령어를 돌리면 원활하게 돌아갑니다!

다음과 같이 잘 실행되는 것을 볼 수 있습니다.

 

첫번째 네모박스 -> 명령어 입력

두번째 네모박스 -> 조절 가능한 옵션들, epochs 수만 건드렸기 때문에 epochs만 150으로 바뀌었습니다.!

세번째 네모박스 -> 사용가능한 CUDA

네번째 네모박스 -> 모델 정보

학습 진행중인 사진!

 

위 경우는 리눅스 서버의 visual studio에서 진행하였습니다. 아나콘다에서도 똑같이 진행하시면 됩니다!

안되신다면 댓글을 남겨주세요..!

 

학습 결과는 runs에서 detect 혹은 train으로 나뉘는데

detect -> 학습 시킨 모델의 weight를 바탕으로 Predict

train -> train 결과 시각화 및 weight 파일 입니다.!

 

test 방법은 detect.py를 통해 --source에 진행해야 하는 Test dataset을 넣은 후 진행하시면 됩니다!

초록색 부분의 주소는 개인에 맞게 수정해주시면 됩니다.

!python detect.py --weight "/home/nyh/yolov5/runs/train/exp5/weights/best.pt"
            --source "/home/nyh/yolov5/empty_or_using_dataset/test/imges"

 

학습 결과는 다음과 같습니다! 다음과 같이 라벨을 띄울수도, 아니면 확률을 띄울수도 있는데 

다음번엔 시간이 된다면 YOLO Custom Model 파라미터 조절 및 튜닝 법 그 외 여러가지에 대해서 글을 써볼게요!

 

감사합니다!