파이썬 roboflow 와 YOLOv5 를 이용해 마스크 쓴 얼굴 인식하기

사전 설정

해당 게시글에서는 구글 colab을 이용하기 때문에 

구글 드라이브를 따로 연결하도록 하겠다

구글드라이브 연결은

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

이 코드를 실행하면 된다

 

데이터 불러오기

모델을 학습시킬 데이터를 불러온다

데이터는 직접 만들 필요는 없고 roboflow에 올라와 있는 데이터를 활용한다

https://universe.roboflow.com/joseph-nelson/mask-wearing/dataset/19

Mask Wearing Object Detection Dataset (v19, updated_raw-images) by Joseph Nelson

해당 게시글에서는 이 데이터를 활용하여 모델을 만들 예정이다

데이터를 가져오는 법은

해당 게시글에서 이 버튼을 누르고

여기서 YOLOv5를 선택한 뒤

여기서 continue를 클릭하고

여기 코드를 복사해주면 된다 api키는 마스킹 되어 있지만 복사하면 정상적으로 표시된다

이 코드를 복사해서 실행해주자

 

코드를 실행하고 파일 목록을 보면

다음과 같이 폴더가 생성되어 있다

 

코랩으로 실행하는 경우 여기서 data.yaml을 수정해줘야한다

더블 클릭을하면

다음과 같이 표시되는데 여기서 test 와 train, val을

다음과 같이 수정해주면 된다

 

YOLOv5 설치하기

YOLOv5를 설치하는 설명은 https://github.com/ultralytics/yolov5

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

여기 나와있다

여기서 install 아래 코드를 복사해서 다음과 같이 수정해 실행해주면 된다

!git clone https://github.com/ultralytics/yolov5  # clone
%cd yolov5
!pip install -r requirements.txt  # install

 

모델 학습

%cd yolov5

!python train.py --img 640 --batch 5 --epochs 4 --data /content/Mask-Wearing-19/data.yaml --weights yolov5m.pt --name maskcheck --exist-ok

다음 코드를 통해 모델을 학습 시킬 수 있다

여기서 img는 학습에 사용할 이미지 크기를 의미하며 이미지가 커질수록 정확도가 오르지만 속도가 느려진다

batch는 한 번에 학습시킬 데이터 수를 의미하며 커질수록 정확도가 오르지만 속도가 느려진다

epoch는 반복 수를 의미하며  커질수록 정확도가 오르지만 오래 걸리고 과적합의 위험이 있다

data 뒤에는 yaml 파일의 경로를 적어주면 되고

yolov5m은 

위 표를 보고 선택했다 여기서 아래로 갈 수록 성능은 좋지만 시간이 오래걸린다

name은 학습시킨 모델의 이름을 의미하며 exist -ok는 같은 이름이 있으면 덮어쓴다는 의미다

코드는 원하는 대로 수정해서 사용하자

 

성능 확인

학습이 종료되면

다움과 같이 모델의 성능을 요약해 준다

 

만든 모델로 검출해보기

모델 학습이 완료되면

/content/yolov5/runs/train/maskcheck/weights/best.pt

 

다음 경로에 가장 성능이 좋았던 모델이 저장이 된다

해당 게시글에서는

해당 이미지를 사용해 검출을 진행해 보겠다

 

 

코드는 다음과 같다

!python detect.py --source /content/masks-tokyo-gty-ps-230314_1678802239571_hpMain.jpg --weights /content/yolov5/runs/train/maskcheck/weights/best.pt --img 640 --conf 0.3 --project /content/ --name result_image --exist-ok --line-thickness 10

source 뒤에는 사용할 이미지 경로를 , weights에서는 아까 얻은 모델의 경로를 넣어 주었다

img는 이미지 사이즈를 입력하고 conf에는 IOU threshhold 값을 적어준다 

IOU threshhold는 표시할 최소 값의 정도를 적는 것이라고 이해하면 된다

line-thickness 는 이미지에 표시할 박스 선의 굵기를 의미한다

name에는 저장할 경로 폴더를 적어주자

 

위 모델을 사용하여 다음과 같이 마스크를 쓴 사람과 마스크를 쓰지 않은 사람을 구별해 낼 수 있다

 

이미지 대신 영상을 넣으면 영상으로도 검출을 해준다

어차피 코드는 같으므로 따로 업로드는 하지 않겠다