Yolov5 커스텀 데이터 학습(안전모)

학습 데이터 : roboflow 샘플 파일 안전모 학습 데이터

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

개발환경 : 구글 코랩

1. 학습 데이터 다운로드

# Roboflow 샘플 파일 다운로드
!curl -L "https://public.roboflow.com/ds/ZQ0MfxeaHQ?key=lPO4Syqo35" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

2. yolov5 설치

#yolov5 설치
%cd /content
!git clone https://github.com/ultralytics/yolov5.git

3. requirements.txt install

%cd /content/yolov5/
!pip install -r requirements.txt

4. traindata, testdata 위치 이동

traindata, testdata 경로 확인

클래스 네임 확인

코랩 기준 content 폴더 내 dataset 폴더 생성 후 train, test(valid로 이름 변경), README.dataset.txt, README.roboflow.txt, data.yaml 파일을 dataset 폴더 하위로 옮긴다.

5. data.yaml 파일 확인

%cat /content/dataset/data.yaml

6. 전체 이미지 개수 확인

%cd /
from glob import glob

#traindata list 불러오기
train_img_list = glob('/content/dataset/train/images/*.jpg')

#testdata list 불러오기
test_img_list = glob('/content/dataset/valid/images/*.jpg')

#traindata, testdata 개수 확인
print(len(train_img_list), len(test_img_list))

# 참고 train, test data가 분리되어 있지 않으면 80% 훈련데이터 20%테스트 데이터로 나눌 수 있음

해당 예제에서는 생략 가능

#from sklearn.model_selection import train_test_split
#train_img_list, test_img_list = train_test_split(img_list, test_size=0.2, random_state=2000)

7. TestData, TrainData 경로 저장

#train/test 파일 경로를 train.txt, test.txt에 쓰기
with open('/content/dataset/train.txt', 'w') as f:
  f.write('\n'.join(train_img_list) + '\n')

with open('/content/dataset/val.txt', 'w') as f:
  f.write('\n'.join(test_img_list)+'\n')

8. 경로 설정

# train, test yaml 파일 확인
import yaml

with open('/content/dataset/data.yaml', 'r') as f:

  data = yaml.full_load(f)

print(data)

data['train'] = '/content/dataset/train.txt'
data['val'] = '/content/dataset/val.txt'

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

print(data)

9. yolo 데이터 학습

%cd /content/yolov5/
!python train.py --img 416 --batch 16 --epochs 10 --data /content/dataset/data.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.pt --name helmet_yolov5s_results

10. 결과 확인

from IPython.display import Image
import os

val_img_path = test_img_list[3]

#사진 경로
!python detect.py --weights /content/yolov5/runs/train/helmet_yolov5s_results/weights/best.pt --img 416 --conf 0.5 --source "{val_img_path}"

#영상 경로
#!python detect.py --weights /content/yolov5/runs/train/helmet_yolov5s_results/weights/best.pt --img 416 --conf 0.5 --source /content/Helmet_nomal.mp4

#Image(os.path.join('/content/yolov5/runs/detect/exp/Helmet_nomal.mp4', os.path.basename(val_img_path)))