[YOLOv5] YOLOv5 Custom Data로 학습시켜보기

YOLOv4를 실습하려다가... 우연히 YOLOv5도 있다는 것을 발견하고 직접 학습시켜보기로 했다!

 

지금까지 YOLO모델을 직접 학습시켜본 적은 없었기 때문에 좋은 경험이 될 것이라고 생각했다.

 

 

 

YOLOv5에 대한 파일들은 아래의 링크에서 git clone을 통해 다운로드 받을 수 있다.

github.com/ultralytics/yolov5

ultralytics/yolov5

 

Custom Data를 이용한 자체 학습내용은 위의 링크에서 [Train Custom Data]를 누르거나 아래의 링크에 자세히 나와있다.

github.com/ultralytics/yolov5/wiki/Train-Custom-Data

ultralytics/yolov5

 

본인은 Ubuntu20.04 환경에서 Terminal을 사용해 실습을 진행했으며,

YouTube '빵형의 개발도상국'님의 [YOLOv5 커스텀 학습 튜토리얼] 영상을 참고하면 Colab으로도 실습을 진행할 수 있다.

www.youtube.com/watch?v=T0DO1C8uYP8&t=612s

 

---

1. YOLOv5 파일 다운로드 및 버젼 동기화

 

YOLO이던, DeepSort이던 실습을 할 때는 가상환경에서 실행하는 것을 추천한다.

git clone을 통해서 원하는 작업 디렉토리 위치에 관련 파일들을 다운로드 받아준다.

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

# colab의 경우
!git clone https://github.com/ultralytics/yolov5.git

yolov5 내 requirements.txt파일을 통해 YOLOv5 개발환경과 동일한 버젼을 다운로드 받아준다.

cd yolov5

pip install -r requirements.txt

 

2. Custom Dataset 다운로드 / 준비

 

자신이 학습시키고자 하는 이미지 데이터셋을 준비한다. (레이블링 데이터 포함)

COCO, Mars, Market1501 등과 같은 데이터셋은 구글링으로 쉽게 찾아볼 수 있으며,

roboflow라는 사이트를 방문하면, 여러가지 Dataset과 label 데이터를 다운로드 받을 수 있다.

본인은 MS COCO Dataset을 사용했다.

public.roboflow.com/object-detection

Object Detection Datasets

오른쪽 위 다운로드 버튼을 누르고, 원하는 다운로드 포맷을 설정하면 된다. (본 글의 경우 YOLOv5 PyTorch 선택)

 

다운로드 버튼을 누르면, zip파일로 받을 것인지, 다운로드 코드로 받을 것인지 선택하는 창이 나오는데

어느 방법으로 해도 상관은 없으나 본인은 다운로드 코드를 생성해서 데이터셋을 다운로드 받았다.

Colab 또는 Jupyter Notebook/lab으로 실습할 경우 Jupyter 탭을,

터미널로 실습할 경우 Terminal 탭에 있는 코드를 치면된다.

 

다운로드 받는 디렉토리는 크게 상관없으나, 본인은 yolov5 상위폴더에 다운로드 했다.

 

MS COCO Dataset의 경우 Train set과 Valid set이 나누어져 있다.

다운로드 받게되면 [ train, valid, README.dataset.txt, README.roboflow.txt, data.yaml] 이 있다.

 

# cat (data.yaml이 있는 절대주소값)
cat /dataset/data.yaml

>>>
names:
- aeroplane
- apple
- backpack
- banana
- baseball bat
- baseball glove
- bear
- bed
- bench
- bicycle
- bird
- boat
...
- tie
- toaster
- toilet
- toothbrush
- traffic light
- train
- truck
- tvmonitor
- umbrella
- vase
- wine glass
- zebra
nc: 80
train: /train/images
val: /valid/images

 

data.yaml 파일을 들여다보면, 클래스 이름들과 클래수 개수 그리고 train과 valid set의 주소가 적혀있다.

 

여기서 dataset을 가르키는 주소가 정확히 되어있지 않기 때문에, 이미지의 주소들을 txt파일로 모아준 뒤 경로 재설정 해준다.

 

glob을 이용하여 image들을 list로 만들어준다.

python

from glob import glob

# 이미지들의 주소 리스트로 만들어줌
train_img_list = glob('./dataset/train/images/*.jpg')
valid_img_list = glob('./dataset/valid/images/*.jpg')

# 리스트를 txt파일로 생성
with open('./dataset/train.txt', 'w') as f:
	f.write('\n'.join(train_img_list) + '\n')
    
with open('./dataset/valid.txt', 'w') as f:
	f.write('\n'.join(valid_img_list) + '\n')

 

생성된 txt파일은 위와 같이 작성된다.

yaml을 이용해 data.yaml의 train set과 valid set의 경로를 바꾸어주도록 하자

 

import yaml

with open('./dataset/data.yaml', 'r') as f:
	data = yaml.load(f)
    
#print(data)

data['train'] = './dataset/train.txt'
data['valid'] = './dataset/valid.txt'

with open('./dataset/data.yaml', 'w') as f:
	yaml.dump(data, f)
    
#print(data)

>>>
names:
- aeroplane
- apple
- backpack
- banana
- baseball bat
- baseball glove
- bear
- bed
- bench
- bicycle
- bird
- boat
...
- tie
- toaster
- toilet
- toothbrush
- traffic light
- train
- truck
- tvmonitor
- umbrella
- vase
- wine glass
- zebra
nc: 80
train: ./dataset/train.txt
val: ./dataset/valid.txt

경로가 재설정되었다.

 

3. 학습 시키기

 

YOLOv5 디렉토리로 이동한 뒤, train.py 파일을 이용해 학습을 시켜주면 된다.

학습시키기 전, wandb 모듈과 가입을 미리 해두면 loss, mAP 등 여러 함수들을 시각화해서 볼 수 있다.

 

학습을 시키기 전 YOLOv5의 모델 중 어떤 것을 사용할지 결정한다.

크기가 크면 클수록 복잡해지고 정확성이 높아지는 대신, 시간이 오래걸리고 GPU의 메모리를 많이 차지한다.

본인은 빠르게 학습시키기 위해서 제일작은 YOLOv5s 모델을 선택했다

cd yolov5

python train.py --img 640 --batch 16 --epochs 20 --data ~/dataset/data.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.pt --name yolov5_coco

--img : 이미지 크기

--batch : 배치 크기

--epochs : epoch 크기

--data : 데이터파일 (data.yaml 파일 경로 지정)

--cfg : 위에서 정한 모델 크기 (yolov5/models 폴더에 yaml파일로 저장되어 있음)

--weights : 미리 학습된 모델로 학습할 경우 (yolov5s.pt 등의 형식으로 다운로드 가능)

--name : 학습된 모델의 이름

 

더 자세한 파라미터는

python train.py -h 를 입력하면 설명을 볼 수 있다.

 

 

GPU를 사용한다면 아래와 같은 표시가 뜨게 된다.

학습도중에도, 학습이 끝난 후에도 yolov5/runs/train/(모델의 이름) 폴더 안을 살펴보면,

train 및 test batch마다의 실제 정답(레이블링)과 예측된 학습결과를 미리볼 수 있다.

실제 레이블링 데이터

학습된 모델의 Detection 데이터

또한, 여러 시각화 그래프를 통해 loss , mAP 등의 결과를 한 눈에 나타내는 지표도 있다. (result.png)

 

규모가 작은 모델로 적은 epoch으로 학습시켰기 때문에 정확도가 조금 떨어지지만,

이에 감안하면 생각보다 만족스러운 모델을 도출해 낸 것 같다.

 

 

 

4. 학습된 모델 사용해보기

 

yolov5 안에 있는 detect.py를 이용하면 다른 사진들을 가지고도 모델을 실행시켜볼 수 있다.

사진 뿐 아니라 동영상, 웹캠, 유튜브 영상도 이용할 수 있다.

 

# source 설명
python detect.py --source 0  # webcam
                            file.jpg  # image 
                            file.mp4  # video
                            path/  # directory
                            path/*.jpg  # glob
                            'https://youtu.be/NUsoVlDFqZg'  # YouTube video
                            'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream
                            
# 코드입력
python detect.py --weights ./runs/train/coco_yolov5/weights/best.pt --img 640 --conf 0.5 --source ~/valid/images/000000010707_jpg.rf.6d5d82303a417707aa64f1cd2f465e8a.jpg

 

detect 결과 파일은 runs/detect/exp 에서 확인할 수 있다.

 

---

YOLOv5 모델을 직접 학습시켜서 실행을 해보았다. 좋은 경험이 된 것 같다.

 

YOLO모델을 적극 활용하기 위해서는 내가 원하는 데이터셋이 매우 필요하다.

 

다음에는 YOLO 등 여러 Detection 모델을 활용해 내가 원하는 모델을 만들기 위한

데이터셋 생성과 증식에 관한 연구를 진행해 볼 것이다.