[Python] YoloV5 custom dataset 학습

 

- YoloV5를 이용한 마스크 착용/미착용 식별 모델개발  -

 

 

# YOLO

YOLO는 You Only Look Once의 약자로, 

one-stage-detection방법을 고안해 실시간으로 Object Detection이 가능하다.

장점 :

1. 간단한 처리 과정으로 속도가 매우 빠르다.

2. 이미지 전체를 한 번에 바라보는 방식을 이용하여 class에 대한 이해도가 높아 낮은 False-Positive를 보인다.

3. 일반화된 Object 학습이 가능

 

분류성능평가지표

 

 

1.  준비

https://github.com/ultralytics/yolov5

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

 

 

https://public.roboflow.com/object-detection/mask-wearing/4

Mask Wearing Object Detection Dataset - raw

roboflow에서 제공하는 mask wearing dataset 활용

학습을 하기 위해서는 image파일과 image속의 객체를 labeling한 txt파일(yolov5)이 필요

 

image 파일

label 파일

 

conda install pytorch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 cudatoolkit=11.0 -c pytorch

YoloV5을 사용하기 위한 라이브러리 install 

지난 포스팅인 tensorflow-gpu 환경설정 때 cuda 11.0 설치했기 때문에 cuda 버전에 맞춰서 pytorch 설치

 

 

 

2.  데이터셋 경로 설정

 

C:\Users\Wana\wancoding\yolov5-master\coco128\images\train2017  -->  train_image 파일 (jpg)

C:\Users\Wana\wancoding\yolov5-master\coco128\labels\train2017     -->  train_label 파일 (txt)

 

C:\Users\Wana\wancoding\yolov5-master\coco128\images\test2017    -->  test_image 파일 (jpg)

C:\Users\Wana\wancoding\yolov5-master\coco128\labels\test2017       -->  test_label 파일 (txt)

 

데이터셋으로 학습
yolov5 폴더에 coco128 이라는 폴더를 만들고,

그 폴더 안에 images 폴더와 labels 폴더를 만든다.
images 폴더에 train2017 폴더를 만들고 안에는 원본 이미지를,
labels 폴더에도 train2017 폴더를 만들고 txt 파일을 넣는다.

 

 

coco128.yaml

coco128.yaml 파일에서 dataset 경로, classes number, classes name 수정

여기서 class는 'mask'와 'no-mask' 두개로 구분

 

 

 

3.  학습 (train.py)

python train.py --img 640 --batch 12 --epochs 128 --data data/coco128.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.pt --name yolov5_mask_det_e128 --device 0

batch 12, epochs 128로 train 진행

train data는 105개의 image로 학습

 

 

yolov5 - model 종류

# pre-trained model

yolov5 는 s, m, l, x 4가지 종류의 모델이 있다.

s가 가장 가벼워 fps(초당 프레임수)가 가장 높지만 성능이 제일 낮다.

반면, x는 무겁고 fps가 가장 낮지만 성능이 제일 좋다.

지금은 yolov5s 모델을 사용하였다.

 

 

train.py 실행

C:/Users/Wana/wancoding/yolov5-master/runs/train/yolov5_mask_det_e128/weights

파일에 모델 (best.pt) 저장

 

 

 

4.  검출 (detect.py)

python detect.py --source coco128/images/test2017/ --weights runs/train/yolov5_mask_det_e128/weights/best.pt --conf 0.4 --name yolov5_mask_det_e128

검출할 dataset 경로 설정 후,

detect.py를 활용하여 모델이 동작되는지 확인

 

 

detect.py 실행

 

C:/Users/Wana/wancoding/yolov5-master/runs/detect/yolov5_mask_det_e128 폴더 안에

검출된 pistol 사진이 저장되는 것을 볼 수 있다.

 

detect된 image파일들

 

 

 

4.  테스트 (val.py)

coco128.yaml

val.py를 실행하기 위해 coco128.yaml 파일 수정   ->   val : test dataset경로로 설정

 

 

python val.py --weights runs/train/yolov5_mask_det_e128/weights/best.pt --data coco128.yaml --img 832 --augment --half --name yolov5_mask_det_e128

모델 경로 설정 후 테스트 실행

 

 

val.py 실행 결과

P-R curve

모든 클래스에 대해 mAP@0.5 : 0.886

 

 

test 결과

 

 

 

- Just Do It -