YOLOv8 도로 표지판, 신호등 검출

YOLO v8 제작자의 github(파라미터도 작아지고 속도도 향상된 것을 알 수 있다.)

0. 환경설정(python = 3.8.8)

!pip install ultralytics
!pip install opencv-python
!pip install matplotlib
pip3 install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio===0.8.1 -f https://download.pytorch.org/whl/torch_stable.html

 

import ultralytics
ultralytics.checks() #설치 확인

 

1. 데이터 준비

Road Sign Detection (kaggle.com)

Road Sign Detection

import xml.etree.ElementTree as ET
import glob
import os
import json
from tqdm import tqdm

def xml_to_yolo_bbox(bbox, w, h):
    # xmin, ymin, xmax, ymax
    x_center = ((bbox[2] + bbox[0]) / 2) / w
    y_center = ((bbox[3] + bbox[1]) / 2) / h
    width = (bbox[2] - bbox[0]) / w
    height = (bbox[3] - bbox[1]) / h
    return [x_center, y_center, width, height]

 

road_sign_root = 'C:/Users/coghk/Desktop/vision/yolo/dataset/road_sign_detection'
annot_path = os.path.join(road_sign_root,"annotations")
img_path = os.path.join(road_sign_root,"images")
label_path = os.path.join(road_sign_root,"labels")

if not os.path.exists(label_path):
    os.makedirs(label_path)

 

classes = []

files = glob.glob(os.path.join(annot_path, '*.xml'))
for fil in tqdm(files):

    basename = os.path.basename(fil)
    filename = os.path.splitext(basename)[0]

    result = []

    tree = ET.parse(fil)
    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}")
    if result:
        with open(os.path.join(label_path, f"{filename}.txt"), "w", encoding="utf-8") as f:
            f.write("\n".join(result))

 

 1.2 데이터 전처리

cd ../dataset #데이터 셋 경로

import os
import random
from shutil import copyfile, rmtree

road_sign_path = '../dataset/road_sign_detection' #저장경로
label_ = '.txt'
img_ = '.png'

# 저장 경로 만들기
folder_list = ['road_sign_detection/train', 'road_sign_detection/val', 'road_sign_detection/train/images', \
                'road_sign_detection/train/labels', 'road_sign_detection/val/images', 'road_sign_detection/val/labels']
for folder in folder_list:
    if not os.path.exists(folder):
        os.makedirs(folder)

 

file_list = os.listdir(os.path.join(road_sign_path, 'images'))

random.shuffle(file_list)

test_ratio = 0.1
test_list = file_list[:int(len(file_list)*test_ratio)]
train_list = file_list[int(len(file_list)*test_ratio):]

print(f"train의 개수 : {len(train_list)}, test의 개수 : {len(test_list)}")

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

 

1.3 config file 생성

import yaml
data =dict()

data['train'] = 'C:/Users/coghk/Desktop/vision/yolo/dataset/road_sign_detection/train'
data['val'] = 'C:/Users/coghk/Desktop/vision/yolo/dataset/road_sign_detection/val'
data['test'] = 'C:/Users/coghk/Desktop/vision/yolo/dataset/road_sign_detection/val'

data['nc'] = 4
data['names'] =['Trafic_light','Speedlimit', 'Crosswalk','Stop']

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

 

2. train

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

 

3. test(validation)

from ultralytics import YOLO

# Load a model
model_path = 'road_sign_detection.pt' #가중치 모델 경로
model = YOLO(model_path)  # load a custom model

# 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-95, map50 모두 좋은 결과가 나왔다.

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_path = 'C:/Users/coghk/Desktop/vision/yolo/weight/road_sign_detection.pt'
model = YOLO(model_path)

root_folder = 'val/images'
result_folder = 'result'
if not os.path.exists(result_folder):
    os.makedirs(result_folder)

test_img_list = os.listdir(root_folder)

device = 'cpu'
color_dict = [(0, 255, 0),(255, 255, 0),(0, 0, 255), (255, 0,0)] #검출 박스 색상
color_dict_2 = [(0, 0, 0),(0, 0, 0),(255, 255, 255), (255, 255,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)], color_dict_2[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)],color_dict_2[int(cls)])
    img_src = annotator.result()
    img_src = cv2.resize(img_src, (400,400))
    cv2.imwrite(os.path.join(result_folder, file), cv2.cvtColor(img_src, cv2.COLOR_RGB2BGR)) 
    #나머지 테스트 파일들도 전부 다 검출

나머지 사진 중 1개