[Mask 착용] YOLOv8로 마스크 착용 유무 판별하기-1

이번 포스팅은 YOLO-v8로 마스크 착용 유무를 판별하도록 해보겠다.

이번 포스팅에서는 전과 달리 ultralytics라는 라이브러리를 사용해서 모델정의, 학습, 검증 등을

진행할 것이기 때문에 다른 예제 포스팅에 비해 간단하다.

 

절차는 다음과 같다

1. 데이터셋 준비
2. 모델 학습
3. 검증
4. 결과 분석

---

1. 데이터셋 준비

마스크 데이터셋은 아래의 깃허브의 링크에서 다운받을 수 있다.

https://github.com/VictorLin000/YOLOv3_mask_detect

GitHub - VictorLin000/YOLOv3_mask_detect: Face mask detection using YOLOv3 on GoogleColab

 

데이터셋은 이미지와 각 이미지에 대한 txt 형식의 라벨링 데이터로 이루어져 있고, 총 678 쌍이다.

그리고 클래스는 총 3개로 마스크를 잘 쓴 경우(0), 이상하게 쓴 경우(1), 안 쓴 경우(2)로 정의되었다.

 

데이터가 그렇게 많지 않으므로 학습과 검증 데이터셋의 비율을 9:1 정도로 하는 것이 좋을 것 같다.

 

이제 학습 및 검증 데이터셋으로 분할해보자. 먼저, 데이터들을 저장할 경로를 생성해준다.

import random
import shutil
import os

raw_file_path = '../../Data/Masks/raw'
file_list = [file for file in os.listdir(raw_file_path) if file.endswith(".txt")]


# Make train/valid folders
train_img_path = '../../Data/Masks/train/images'
train_lbl_path = '../../Data/Masks/train/labels'
valid_img_path = '../../Data/Masks/valid/images'
valid_lbl_path = '../../Data/Masks/valid/labels'

folder_list = [train_img_path, train_lbl_path, valid_img_path, valid_lbl_path]

for folder in folder_list:
    if not os.path.exists(folder):
        os.makedirs(folder)

 

그 후, 데이터셋 리스트를 랜덤으로 정렬한 후 미리 설정한 비율에 따라 지정된 경로로 복사해준다.

random.shuffle(file_list)

ratio = .1
crit = int(len(file_list)*ratio)
valid_list = file_list[:crit]
train_list = file_list[crit:]

modes = ['train', 'valid']

for mode in modes:
    if mode == 'train':
        data_list = train_list
        img_path = train_img_path
        lbl_path = train_lbl_path
    else:
        data_list = valid_list
        img_path = valid_img_path
        lbl_path = valid_lbl_path

    for data in data_list:
        file_name = os.path.splitext(data)[0]
        # Copy image file
        shutil.copy(os.path.join(raw_file_path, file_name+f'.jpg'), os.path.join(img_path, file_name+f'.jpg'))
        # Copy label file
        shutil.copy(os.path.join(raw_file_path, data), os.path.join(lbl_path, data))

 

마지막으로 yolo v8 모델을 학습시키기 위한 yaml 파일을 만들어줘야 한다.

yaml 파일은 학습에 필요한 정보를 담고있는데, 딕셔너리 형태로 데이터 경로 및 클래스 수 등을 기록해주면 된다.

 

아래와 같이 코드를 실행해주자.

 

import yaml
data = dict()

data['train'] = '학습 데이터 파일 경로'
data['val'] = '검증 데이터 파일 경로'
data['test'] = '검증 데이터 파일 경로'

data['nc'] = 3
data['names'] = ['Yes', 'Improperly', 'No']

with open('mask_detection.yaml', 'w') as f:
    yaml.dump(data, f)

 

train과 val, test의 value에는 적힌 것처럼 해당하는 데이터 파일의 경로를 적어주면 된다.

나는 각 class를 마스크를 잘 쓴 경우(Yes), 이상하게 쓴 경우(Improperly) 그리고 안 쓴경우(No)로 정의해줬다.

이러면 데이터셋과 관련된 준비는 모두 끝났다.

 

사실, 데이터 증강 기법같은 것도 적용해보고 싶었는데 저 라이브러리를 써본 적이 없어서 못했다.

나중에 이 라이브러리에 익숙해지면 사용법을 따로 포스팅해봐야겠다.

---

2. 모델 학습

학습 코드는 엄청 간단하다. 뭐 따로 정의해줄 필요도 없이 ultralytics가 다 해준다.

 

먼저 필요한 라이브러리를 불러오고 모델을 정의해준다.

import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

from ultralytics.utils.plotting import Annotator
from ultralytics import YOLO

model = YOLO('yolov8n.yaml')

 

다음은 모델 학습 코드다. 아래의 코드만 실행해주면 된다.

results = model.train(data='mask_detection.yaml', epochs=200, batch=32, device=0,
                      patience=30, name='mask_detection')

 

학습에는 대략 7~8분정도 시간이 소요됐다. 학습 정보와 best 모델은 위 코드를 실행한 경로에 runs라는 파일에 있다.

 

학습이 잘 마무리되면 검증을 수행해보자.

 

# Validation
model = YOLO('runs/detect/mask_detection3/weights/best.pt')

# validation
metrics = model.val()

 

 

왠지는 모르겠지만, 전체 클래스에 대한 mAP를 보면 mAP50이 0.76으로 그렇게 높지 않은 것을 확인할 수 있다.

 

이제, 모델의 출력을 시각화해보자.

검증 데이터셋에서 이미지 한 개를 불러와서 결과를 살펴보자.

 

valid_img_path = '../../Data/Masks/valid/images'
test_img_list = os.listdir(valid_img_path)
color_dict = [(0, 255, 0), (255, 0, 0), (0, 0, 255)]

# Load sample images
test_img = cv2.imread(os.path.join(valid_img_path, test_img_list[0]))
img_src = cv2.cvtColor(test_img, cv2.COLOR_BGR2RGB)
result = model(test_img)

# Visualization
for res in result:
    annotator = Annotator(img_src)
    boxes = res.boxes
    for box in boxes:
        b = box.xyxy[0]
        cls = box.cls
        annotator.box_label(b, model.names[int(cls)], color_dict[int(cls)])
img_src = annotator.result()
plt.imshow(img_src)

 

 

위와 같이 대체로 잘 동작하지만, 몇몇 경우 손을 얼굴로 인식하는 등 잘 동작하지 않는 경우도 존재했다.

---

이번 포스팅에서는 YOLO v8을 사용하여 마스크 착용여부를 탐지하는 모델을 구성했다.

Object detection을 기반으로 동작하는 테스크였고, 기존의 포스팅들과 달리 라이브러리를 사용하여 편하게

수행할 수 있었다.

 

다음 포스팅은 이전에 cam으로 실시간 object detection을 수행했던 것을 바탕으로 실시간

마스크 검출을 수행하는 포스팅을 해보고자 한다.

 

포스팅에 대한 질문이나 잘못된 사항 댓글 달아주시면 감사하겠습니다 ! :)