YOLO v8을 이용한 마스크 검출

2024. 3. 17. 13:12파이썬

0. 환경설정(python = 3.8.8)

 

import ultralytics
ultralytics.checks() #잘 설치 되었는지 확인

Ultralytics YOLOv8.1.29 🚀 Python-3.8.8 torch-2.2.1+cpu CPU (Intel Core(TM) i7-10510U 1.80GHz) Setup complete ✅ (8 CPUs, 31.8 GB RAM, 266.6/459.0 GB disk) # cuda를 설치안해서 그런지 내 gpu가 잡히지 않았음

1. 데이터 전처리

yolo - Google Drive(마스크 데이터셋)

 

yolo - Google Drive

이 폴더에 파일이 없습니다.이 폴더에 파일을 추가하려면 로그인하세요.

drive.google.com

cd '데이터셋 저장되어 있는 곳'
import os
import random
import shutil

mask_root = 'mask_orignal/'
mask_path = 'mask_detection/'
label_ = '.txt'
img_ = '.jpg'
file_list = [file for file in os.listdir(mask_root) if file.endswith(".txt")]

# 저장 경로 만들기
folder_list = ['mask_detection', 'mask_detection/train', 'mask_detection/val', 'mask_detection/train/images', \
                'mask_detection/train/labels', 'mask_detection/val/images', 'mask_detection/val/labels']
for folder in folder_list:
    if not os.path.exists(folder):
        os.makedirs(folder) #폴더가 없으면 폴더 생성
random.shuffle(file_list) # 섞어서 나누기 위함
test_ratio = 0.1 #테스트 셋 비율을 10%로 설정
test_list = file_list[:int(len(file_list)*test_ratio)]
train_list = file_list[int(len(file_list)*test_ratio):]

for i in test_list:
    f_name = os.path.splitext(i)[0]
    shutil.copyfile(os.path.join(mask_root,(f_name+img_)), os.path.join(mask_path, 'val/images', (f_name+img_)))
    shutil.copyfile(os.path.join(mask_root,(f_name+label_)), os.path.join(mask_path, 'val/labels', (f_name+label_)))
for i in train_list:
    f_name = os.path.splitext(i)[0]
    shutil.copyfile(os.path.join(mask_root,(f_name+img_)), os.path.join(mask_path, 'train/images', (f_name+img_)))
    shutil.copyfile(os.path.join(mask_root, (f_name+label_)), os.path.join(mask_path, 'train/labels', (f_name+label_)))

1.1 config file 만들기

import yaml
data =dict() 

data['train'] = '../mask_detection/train'
data['val'] =  '../mask_detection/val'
data['test'] = '../dataset/mask_detection/val'

data['nc'] = 3
data['names'] =['OK','improperly', 'NO']

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

2.train

from ultralytics import YOLO
model = YOLO('yolov8s.yaml')
results = model.train(data ='mask_detection.yaml', epochs = 100, batch=32,device = 0  , patience=30, name='road_sign_s')

 

3.test(validation)

from ultralytics import YOLO

#모델 불러오기
model_path = '../mask_detection.pt' #가중치 모델 경로
model = YOLO(model_path) 

# Validate the model
metrics = model.val()  # no arguments needed, dataset and settings remembered

print("map50-95", metrics.box.map)
print("map50", metrics.box.map50)

map50과, map50-95의 결과

"mAP50"와 "mAP50-95"는 객체 탐지(Object Detection)에서 사용되는 성능 측정 지표입니다. 여기서 "mAP"는 평균 정밀도(Mean Average Precision)를 나타냅니다.

  1. mAP50:
    • mAP50는 객체 탐지 모델의 성능을 측정하는 지표 중 하나입니다.
    • "50"은 IOU (Intersection over Union) 임계값을 나타냅니다. IOU 임계값은 얼마나 겹치는지에 대한 임계값을 나타냅니다.
    • mAP50는 IOU 임계값이 0.5인 경우의 평균 정밀도를 나타냅니다. 즉, 각 클래스에 대한 정밀도를 계산하고 이를 평균한 값입니다.
  2. mAP50-95:
    • mAP50-95는 객체 탐지 모델의 더 포괄적인 성능 측정 지표입니다.
    • "50-95"는 IOU 임계값 범위를 나타냅니다. 여기서는 0.5에서 0.95까지의 범위를 의미합니다.
    • mAP50-95는 IOU 임계값이 0.5에서 0.95까지 변화할 때의 평균 정밀도를 나타냅니다. 이는 다양한 IOU 임계값에서 모델의 성능을 종합적으로 평가하는 데 사용됩니다.

일반적으로 mAP50은 모델의 더 높은 재현율을 측정하는 데 사용되며, mAP50-95는 모델의 더 높은 정밀도를 측정하는 데 사용됩니다. 이러한 지표는 객체 탐지 모델의 성능을 종합적으로 이해하고 비교하는 데 유용합니다.

 

4. inference

from ultralytics import YOLO
import cv2
import os
%matplotlib inline
from ultralytics.utils.plotting import Annotator
import matplotlib.pyplot as plt
import numpy as np

model = YOLO('../mask_detection.pt')

root_folder = '../mask_detection/val/images'
result_folder = '../mask_detection/result' 
if not os.path.exists(result_folder): #result 폴더 없으면 생성
    os.makedirs(result_folder)

test_img_list = os.listdir(root_folder)
device = 'cpu'
color_dict = [(0, 255, 0), (255, 0, 0), (0, 0, 255)]
test_img = cv2.imread(os.path.join(root_folder, test_img_list[0]))
img_src = cv2.cvtColor(test_img, cv2.COLOR_BGR2RGB)
results = model(test_img)
for result in results:
    annotator = Annotator(img_src)
    boxes = result.boxes
    for box in boxes:
        b = box.xyxy[0]  # get box coordinates in (top, left, bottom, right) format
        cls = box.cls
        annotator.box_label(b, model.names[int(cls)], color_dict[int(cls)])
img_src = annotator.result()
plt.imshow(img_src)
plt.show()

결과 이미지

 

for idx , file in enumerate(test_img_list):
    test_img = cv2.imread(os.path.join(root_folder, file))
    img_src = cv2.cvtColor(test_img, cv2.COLOR_BGR2RGB)
    results = model(test_img)

    for result in results:
        annotator = Annotator(img_src)
        boxes = result.boxes
        for box in boxes:
            b = box.xyxy[0]  # get box coordinates in (top, left, bottom, right) format
            cls = box.cls
            annotator.box_label(b, model.names[int(cls)], color_dict[int(cls)])
    img_src = annotator.result()
    cv2.imwrite(os.path.join(result_folder, file), cv2.cvtColor(img_src, cv2.COLOR_RGB2BGR)) 
    #나머지 테스트용 사진들도 전부 다 적용

마스크 착용 여부를 잘 검출해낸것을 볼 수 있다.