Python(59)- YOLO를 활용한 안전모 탐지기

*이 글을 읽기전에 작성자 개인의견이 있으니, 다른 블로그와 교차로 읽는것을 권장합니다.*

YOLO모델을 활용하여 공사현장 사진을 가져온다. 안전모의 유무를 기준으로 분류하는 분류기 만들기

# ultralytics 모듈 사용전, pip로 설치하기
!pip install ultralytics

사용 모듈

import os
import random
import shutil
import yaml
import cv2 
import glob # 파일 시스템 내에서 특정 패턴과 일치하는 파일 경로명을 찾기 위해 사용
import ultralytics
import matplotlib.pyplot as plt
import xml.etree.ElementTree as ET

from ultralytics import YOLO
import torch
from torchvision import transforms
from tqdm import tqdm

ultralytics.checks()

!kaggle datasets download -d andrewmvd/hard-hat-detection
!unzip -q /content/hard-hat-detection.zip

# helmet_detection 디렉토리를 수동 생성
data_root = '/content/helmet_detection'
# helmet_detection/data 디렉토리를 수동 생성
file_root = f'{data_root}/data'
project_name = 'shd'

train_root = f'{data_root}/{project_name}/train'
valid_root = f'{data_root}/{project_name}/valid'
test_root = f'{data_root}/{project_name}/test'

for folder in [train_root, valid_root, test_root]:
    if not os.path.exists(folder):
        os.makedirs(folder)
    for s in ['images', 'labels']:
        s_folder = f'{folder}/{s}'
        if not os.path.exists(s_folder):
            os.makedirs(s_folder)

# helmet_detection/data 디렉토리에 images, annotation를 수동으로 넣어줌
file_list = glob.glob(f'{file_root}/images/*.png')
len(file_list)

xml.etree.ElementTree 모듈 

• XML 데이터를 파싱하고 조작하기 위한 모듈
  • XML은 데이터를 구조화하고 저장하는 데 사용되는 인기 있는 형식
• XML 문서를 트리 구조로 표현하여 요소(Element)와 속성(Attribute)을 쉽게 탐색하고 수정할 수 있게 해줌
• 사용 용도:
  • XML 파싱: XML 파일이나 문자열을 읽어들여 트리 구조로 변환
  • XML 탐색: 트리 구조에서 특정 요소나 속성을 검색
  • XML 조작: XML 요소를 추가, 수정, 삭제
  • XML 생성: 새로운 XML 문서를 생성하고 구성
  • XML 저장: 트리 구조를 XML 형식으로 파일에 저장

# /content/helmet_detection/data/labels 디렉토리 labels 폴더 수동으로 옮기기
# xml -> YOLO로 바꿀 때, txt로 바꿔주고, 내부에는 각각 앞에는 클래스이름, with height
# 태그 하나하나로 접근하면 어려워서, xml passing하는 모듈 사용하기 -> xml.etree.ElementTree

classes = []
for file in tqdm(file_list):
    file_name = file.split('/')[-1].replace('png', 'xml')
    save_name = file_name.replace('xml', 'txt')
    file_path = f'{file_root}/annotations/{file_name}'
    save_path = f'{file_root}/labels/{save_name}'

    result = list()
    tree = ET.parse(file_path)
    root = tree.getroot()
    width = int(root.find('size').find('width').text)
    height = int(root.find('size').find('height').text)
    for obj in root.findall('object'):
        label = obj.find('name').text
        if label not in classes:
            classes.append(label)
        index = classes.index(label)
        pil_bbox = [int(x.text) for x in obj.find('bndbox')]
        yolo_bbox = xml_to_yolo_bbox(pil_bbox, width, height)
        bbox_string = ' '.join([str(x) for x in yolo_bbox])
        result.append(f'{index} {bbox_string}')
    # 한칸씩 띄어가면서 데이터를 저장하는 txt파일 만들기    
    if result:
        with open(save_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(result))

 

classes

cls_list = ['head', 'helmet', 'person']

random.seed(2024)

file_list 랜덤 셔플해주기

random.shuffle(file_list)
test_ratio = 0.1
num_file = len(file_list)

test_list = file_list[:int(num_file * test_ratio)]
valid_list = file_list[int(num_file * test_ratio):int(num_file * test_ratio)*2]
train_list = file_list[int(num_file * test_ratio) * 2:]

for i in test_list:
    label_name = i.split('/')[-1].replace('png', 'txt')
    label_path = f'{file_root}/labels/{label_name}'
    shutil.copyfile(label_path, f'{test_root}/labels/{label_name}')
    img_name = i.split('/')[-1]
    shutil.copyfile(i, f'{test_root}/images/{img_name}')
for i in valid_list:
    label_name = i.split('/')[-1].replace('png', 'txt')
    label_path = f'{file_root}/labels/{label_name}'
    shutil.copyfile(label_path, f'{valid_root}/labels/{label_name}')
    img_name = i.split('/')[-1]
    shutil.copyfile(i, f'{valid_root}/images/{img_name}')
for i in train_list:
    label_name = i.split('/')[-1].replace('png', 'txt')
    label_path = f'{file_root}/labels/{label_name}'
    shutil.copyfile(label_path, f'{train_root}/labels/{label_name}')
    img_name = i.split('/')[-1]
    shutil.copyfile(i, f'{train_root}/images/{img_name}')

project_root = '/content/helmet_detection'

data = dict()

data['train'] = train_root
data['val'] = valid_root
data['test'] = test_root
data['nc'] = len(cls_list)
data['names'] = cls_list

with open(f'{project_root}/safety_helmet.yaml', 'w') as f:
    yaml.dump(data, f)

경로 설정

%cd /content/helmet_detection

yolo8n.pt가져와서 yaml파일 바탕으로 학습시키기(epoch는 '2'로 설정)

model = YOLO('yolov8n.pt')
results = model.train(data='safety_helmet.yaml', epochs=2, batch=8, imgsz=224, device=0, workers=4, amp=False, patience=30, name='safety_n')

model = YOLO('yolov8s.pt')
results = model.train(data='safety_helmet.yaml', epochs=2, batch=8, imgsz=224, device=0, workers=4, amp=False, patience=30, name='safety_s')

model = YOLO('yolov8m.pt')
results = model.train(data='safety_helmet.yaml', epochs=2, batch=8, imgsz=224, device=0, workers=4, amp=False, patience=30, name='safety_m')

%cd /content/helmet_detection

result_folder = f'{project_root}/runs/detect'

YOLO모델 test 평가하기

model = YOLO(f'{result_folder}/safety_n/weights/best.pt')
metrics = model.val(split='test')
print('map50-95', metrics.box.map)
print('map50', metrics.box.map)

마지막으로 epoch='50'으로 한번 더 모델링 학습시키기

model = YOLO('yolov8n.pt')
results = model.train(data='safety_helmet.yaml', epochs=50, batch=8, imgsz=224, device=0, workers=4, amp=False, patience=30, name='safety')

모델 평가

model = YOLO(f'{result_folder}/safety/weights/best.pt')
metrics = model.val(split='test')
print('map50-95', metrics.box.map)
print('map50', metrics.box.map)

data_root = '/content/helmet_detection'
project_name = 'shd'
test_root = f'{data_root}/{project_name}/test'

test_file_list = glob.glob(f'{test_root}/images/*.png')
random.shuffle(test_file_list)

IMG_SIZE = (224,224)
test_data_transform = transforms.Compose([
    transforms.Resize(IMG_SIZE),
    transforms.ToTensor()
])

model.names

color_dict = {i: tuple([random.randint(0, 255) for _ in range(3)]) for i in range(len(model.names))}
color_dict

test_img = cv2.imread(test_file_list[0])
img_src = cv2.cvtColor(test_img, cv2.COLOR_BGR2RGB)
result = model(img_src)

result[0].boxes

plt.imshow(img_src)
plt.show()

color_dict = {
    0 : (255, 0, 255),
    1:  (0, 255, 255),
    2:  (0, 0, 255)
}

여기서 YOLO모델을 활용하여 Classification(객체 분류)의 용도가 아닌, Detection(객체 탐지)의 용도로써 YOLO를 사용해본다.

num_head = 0
test_img = cv2.imread(test_file_list[1])
img_src = cv2.cvtColor(test_img, cv2.COLOR_BGR2RGB)
result = model(img_src)[0]

annotator = Annotator(img_src)
boxes = result.boxes

for box in boxes:
    b = box.xyxy[0]
    cls = box.cls
    if 'head' == model.names[int(cls)]:
        num_head += 1
    annotator.box_label(b, model.names[int(cls)], color_dict[int(cls)])
img_src = annotator.result()
if num_head > 0:
    cv2.rectangle(img_src, (0, 0), (300, 50), (255, 0, 0), -1, cv2.LINE_AA)
    cv2.putText(img_src, 'No Helmet!', (5, 30), cv2.FONT_HERSHEY_DUPLEX, 1, (255, 255, 255), thinkness=3, lineType=cv2.LINE_AA)
plt.imshow(img_src)
plt.show()

plt.figure(figsize=(20, 16))

for idx in range(20):
    num_head = 0
    test_img = cv2.imread(test_file_list[idx])
    img_src = cv2.cvtColor(test_img, cv2.COLOR_BGR2RGB)
    result = model(img_src)[0]

    annotator = Annotator(img_src)
    boxes = result.boxes

    for box in boxes:
        b = box.xyxy[0]
        cls = box.cls
        if 'head' == model.names[int(cls)]:
            num_head += 1
        annotator.box_label(b, model.names[int(cls)], color_dict[int(cls)])
    img_src = annotator.result()

    plt.subplot(5, 4, (idx+1))
    if num_head > 0:
        cv2.rectangle(img_src, (0, 0), (300, 50), (255, 0, 0), -1, cv2.LINE_AA)
        cv2.putText(img_src, 'No Helmet!', (5, 30), cv2.FONT_HERSHEY_DUPLEX, 1, (255, 255, 255), lineType=cv2.LINE_AA)
    plt.imshow(img_src)
plt.show()

len(test_file_list)