YOLOv5 커스텀 데이터로 학습하기

https://www.youtube.com/watch?v=T0DO1C8uYP8

위 영상을 토대로 작성했습니다. 3년 전 영상이라 바뀐 부분들도 있기 때문에 조금은 다를 수 있습니다.

 

 

---

 

YOLOv5 모델에 안전모 착용 데이터셋을 이용하여 학습하는 것을 목표로 합니다. 데이터셋은 Roboflow의 Hard Hat Workers Dataset을 이용했습니다. 

 

https://public.roboflow.com/object-detection/hard-hat-workers

Hard Hat Workers Object Detection Dataset

• 전 YOLOv5에 맞춰 다운 받았습니다. 그 후 드라이브에 압축파일 채로 넣어줍니다. 코드 상에서 압축을 풀면 시간이 많이 단축되기에 압축 풀기는 코드 상에서 진행하겠습니다.

 

 

from google.colab import drive
drive.mount('/content/drive')

• 당연한 얘기지만 구글 드라이브 먼저 마운트 해줍니다.

 

# 경로 이동 후 YOLOv5 git에서 clone

%cd /content/drive/MyDrive/Colab Notebooks/YOLO
!git clone https://github.com/ultralytics/yolov5.git

• 그 다음 YOLOv5 git 주소를 원하는 디렉토리에 그대로 클론 해줍니다.

 

# 미리 다운 받아 놓은 데이터셋 경로 이동 후 압축 풀기

%cd /content/drive/MyDrive/Colab Notebooks/YOLO/dataset
!unzip -qq "/content/drive/MyDrive/datasets/Hard Hat Workers.v2-raw.yolov5pytorch.zip"

• 먼저 본인의 디렉토리에 dataset 폴더를 미리 만들어줍니다. 
• %cd 명령어를 이용하여 dataset 폴더 위치로 이동합니다.
• 그다음 !unzip 명령어를 이용하여 압축 풀기를 진행 해주시면 됩니다.

 

# 필요한 라이브러리 다운로드

%cd /content/drive/MyDrive/Colab Notebooks/YOLO/yolov5
!pip install -r requirements.txt

• YOLOv5를 깃에서 클론 해오셨다면 yolov5라는 폴더가 생겼을 겁니다.
• yolov5 폴더로 먼저 이동합니다.
• 그다음 !pip install 을 통해 requirements를 모두 다운 받아줍니다.

 

data.yaml 수정

• 다음은 data.yaml 파일을 수정해야 합니다.
• 다운받은 데이터셋 경로에 data.yaml이라는 파일이 있습니다. 코드로 수정해도 되지만 귀찮으니... 그냥 코랩 왼쪽 아래 폴더 모양을 눌러 data.yaml 파일에 접근해 줍니다.
• train: train 데이터 경로로 수정합니다.
• test: test 데이터 경로로 수정합니다.
• nc: 클래스 개수입니다. 수정하지 않아도 됩니다.
• names: 클래스 명입니다. 수정하지 않아도 됩니다.

 

만약 데이터를 어디서 다운 받는게 아니고 직접 크롤링하여 구성할 때는 위에 적은 설명을 보고 작성하시면 될 것 같습니다. 아래 링크는 직접 데이터 구성하신 분의 포스팅입니다. 저도 도움이 많이 됐으니 참고하면 좋을 것 같네요!

 

https://put-idea.tistory.com/52

YOLOv5 커스텀 데이터셋으로 학습하기

 

 

다음은 대망의 학습 입니다..!

 

 

https://docs.ultralytics.com/yolov5/tutorials/train_custom_data/

Train Custom Data

• 위 링크는 YOLO에서 제공하는 커스텀 데이터로 YOLOv5를 학습하는 튜토리얼입니다.
• YOLOv5 공식 깃 문서에 들어가시면 여러가지 튜토리얼들이 존재하니 참고하시면 좋을 것 같습니다.

 

이를 참고하여 작성한 학습 코드는 아래와 같습니다.

 

%cd /content/drive/MyDrive/YOLO/yolov5

!python train.py --img 416 --batch 16 --epochs 5 --data /content/drive/MyDrive/YOLO/dataset/data.yaml --weights ./models/yolov5s.pt --name yolov5s_test_1

• 먼저 cd 명령어를 통해 yolov5 디렉토리로 이동합니다. 참고로 저는 중간에 디렉토리 구조를 변경해서 경로가 위 코드랑 다릅니다. 어느 파일에 접근해야 하는지만 참고하여 본인의 디렉토리 구조에 맞춰 진행하길 바랍니다.
• img: 이미지 사이즈를 지정합니다.
• batch: 배치 사이즈를 지정합니다.
• epochs: 학습 횟수를 지정합니다.
• data: 데이터에 접근합니다. 아까 수정한 data.yaml의 디렉토리 구조를 작성하면 됩니다.
• weight: 어떤 모델을 기본 모델로 사용할지 지정합니다. yolov5/models 디렉토리에 들어가면 기본 모델들이 존재하는데 저는 yolov5s를 사용했습니다. 저는 최종 목표가 라즈베리파이에 모델을 적용하는거라 small 크기를 사용하였고 개인에 맞춰 모델을 결정하시길 바랍니다. 당연히 사이즈가 커질수록 효율은 좋지만 모델이 무겁고 복잡합니다.
• name: 학습이 끝나고 저장할 모델 이름을 지정합니다.

 

모두 작성했다면 학습을 시작해줍니다. 저는 데이터 개수가 많아 일단 테스트로 돌려보기 위해 5에퍽만 진행했습니다. 시간이 많다면 에퍽 수를 더욱 늘리는걸 추천합니다. 

 

참고로 저는 약 5천장의 사진을 5에퍽동안 훈련하는데 1시간이 걸렸습니다.. 훈련 시간보다 5천장의 데이터를 모두 불러오는게 더 오래 걸렸습니다...^^ 

 

자 이제 학습이 끝났다면 성능을 한 번 확인해봐야겠죠?

 

• 먼저 훈련된 결과는 yolov5/runs/train/{모델 이름}/weights 구조로 들어가면 확인이 가능합니다. best.pt는 가장 성능이 좋았던 모델의 가중치, last.pt는 마지막에 저장된 모델의 가중치입니다. 혹여나 과적합이 일어나는 경우가 있기에 best와 last 모델을 나누어 놓았습니다.

 

• 생성한 모델 폴더에 접근하시면 F1_curve, 혼동행렬 등 많은 지표들을 볼 수 있으니 확인해 보시면 좋을 것 같네요.
• result.png를 확인해보면 학습 결과도 볼 수 있습니다.

 

%load_ext tensorboard
%tensorboard --logdir /content/drive/MyDrive/YOLO/yolov5/runs

• 아래 코드를 작성해서 텐서보드로 깔끔하게 모아서 볼 수도 있습니다.

 

 

그래프로만 봐서는 잘 모르겠으니 결과를 확인해보겠습니다.

 

from IPython.display import Image
import os

test_img_path = '/content/drive/MyDrive/YOLO/dataset/test/images/005298_jpg.rf.647d148af5d961d8bbc041f172247170.jpg'

!python /content/drive/MyDrive/YOLO/yolov5/detect.py --weights /content/drive/MyDrive/YOLO/yolov5/runs/train/yolov5s_test_12/weights/best.pt --img 416 --conf 0.5 --source "{test_img_path}"

 

• 먼저 test_img_path에 테스트할 이미지의 경로를 입력해줍니다.
• 코드를 순서대로 보면 아래와 같습니다.
• yolov5 폴더에 detect.py 경로를 이용해 실행시킵니다.
• --weights는 best.pt 모델을 이용하겠습니다.
• --img 는 이미지 사이즈를 맞춰줍니다.
• --conf 0.5는 0.5(50%)이상의 확률이면 맞다고 예측하는 것입니다.
• --source "{test_img_path}" 를 작성하면 이미지 경로를 읽을 수 있습니다.

• 코드를 실행하면 아래와 같은 결과창을 볼 수 있습니다.
• 맨 아래 줄을 보면 저장된 경로를 확인할 수 있습니다.

 

from matplotlib import pyplot as plt

img_path = '/content/drive/MyDrive/YOLO/yolov5/runs/detect/exp/005298_jpg.rf.647d148af5d961d8bbc041f172247170.jpg'


plt.figure(figsize=(10, 10))
img = plt.imread(img_path)
plt.imshow(img)

• 해당 경로를 가져와 주피터 노트북에 띄워 보겠습니다.

• 결과를 보면 노란색 헬맷을 모두 탐지한 것을 알 수 있습니다.

 

---

일단 생각보다 YOLO 학습이 쉬워서 놀랐습니다.. 다음에는 영상에 해당 모델을 적용하여 테스트하고 그 후에 라즈베리파이 웹캠에 연결할 계획입니다.