Deep Learning Framework를 활용한 YOLOv8 기반의 전력 변전소 가시광 이미지 결함 탐지 시스템 개발

이 문서에서는 YOLOv8 모델을 사용하여 전력 변전소에서 발생할 수 있는 다양한 결함(계기 손상, 절연체 손상, 유출, 호흡기 이상 등)을 자동으로 탐지하는 방법을 다룹니다.

환경 설정

먼저 필요한 라이브러리를 설치합니다. 다음 명령어를 실행하세요:

pip install ultralytics lxml opencv-python-headless pandas scikit-learn

데이터 전처리

XML 형식의 레이블 파일을 YOLOv8에 맞는 TXT 형식으로 변환해야 합니다. 다음은 convert_xml_to_yolo.py 스크립트입니다:

import os
from xml.etree.ElementTree import parse
from pathlib import Path

base_dir = Path("datasets/substation_inspection")
xml_dir = base_dir / "Annotations"
image_dir = base_dir / "JPEGImages"
output_dir = base_dir / "labels"

os.makedirs(output_dir, exist_ok=True)

class_map = {
    "bjdsyc": 0,
    "bj_wkps": 1,
    "yw_nc": 2,
    "xmbhyc": 3,
    "kgg_ybh": 4,
    "gbps": 5,
    "yw_gkxfw": 6,
    "hxq_gjbs": 7,
    "bj_bpmh": 8,
    "jyz_pl": 9,
    "bj_bpps": 10,
    "sly_dmyw": 11,
    "wcaqm": 12,
    "wcgz": 13,
    "ywzt_yfyc": 14,
    "hxq_gitps": 15,
    "xy": 16,
}

def convert_bbox(image_width, image_height, bbox):
    x_min, y_min, x_max, y_max = map(float, bbox)
    center_x = ((x_min + x_max) / 2) / image_width
    center_y = ((y_min + y_max) / 2) / image_height
    width = (x_max - x_min) / image_width
    height = (y_max - y_min) / image_height
    return center_x, center_y, width, height

for xml_file in xml_dir.glob("*.xml"):
    tree = parse(xml_file)
    root = tree.getroot()

    image_width = int(root.find("size/width").text)
    image_height = int(root.find("size/height").text)

    with open(output_dir / f"{xml_file.stem}.txt", "w") as output:
        for obj in root.findall("object"):
            class_name = obj.find("name").text.replace(".", "_")
            if class_name not in class_map:
                continue
            class_id = class_map[class_name]
            bbox = obj.find("bndbox")
            bbox_data = [bbox.find(tag).text for tag in ["xmin", "ymin", "xmax", "ymax"]]
            yolo_format = convert_bbox(image_width, image_height, bbox_data)
            output.write(f"{class_id} {' '.join(map(str, yolo_format))}\n")

print("변환이 완료되었습니다.")

YAML 구성 파일 작성

데이터셋을 설명하는 substation_inspection.yaml 파일을 작성합니다:

train: ../datasets/substation_inspection/train/images
val: ../datasets/substation_inspection/val/images

nc: 17
names:
  - bjdsyc
  - bj_wkps
  - yw_nc
  - xmbhyc
  - kgg_ybh
  - gbps
  - yw_gkxfw
  - hxq_gjbs
  - bj_bpmh
  - jyz_pl
  - bj_bpps
  - sly_dmyw
  - wcaqm
  - wcgz
  - ywzt_yfyc
  - hxq_gitps
  - xy

데이터셋 분리

다음 스크립트를 사용하여 데이터셋을 훈련용, 검증용, 테스트용으로 나눕니다:

import os
import random
from sklearn.model_selection import train_test_split
from pathlib import Path

dataset_dir = Path("datasets/substation_inspection")
images_dir = dataset_dir / "JPEGImages"
annotations_dir = dataset_dir / "labels"

train_images_dir = dataset_dir / "train/images"
train_labels_dir = dataset_dir / "train/labels"
val_images_dir = dataset_dir / "val/images"
val_labels_dir = dataset_dir / "val/labels"
test_images_dir = dataset_dir / "test/images"
test_labels_dir = dataset_dir / "test/labels"

os.makedirs(train_images_dir, exist_ok=True)
os.makedirs(train_labels_dir, exist_ok=True)
os.makedirs(val_images_dir, exist_ok=True)
os.makedirs(val_labels_dir, exist_ok=True)
os.makedirs(test_images_dir, exist_ok=True)
os.makedirs(test_labels_dir, exist_ok=True)

image_files = list(images_dir.glob("*.jpg"))
random.shuffle(image_files)

train_ratio = 0.7
val_ratio = 0.15
test_ratio = 0.15

num_images = len(image_files)
train_split = int(num_images * train_ratio)
val_split = int(num_images * (train_ratio + val_ratio))

train_set = image_files[:train_split]
val_set = image_files[train_split:val_split]
test_set = image_files[val_split:]

def copy_dataset(dataset, dest_image_dir, dest_label_dir):
    for img_path in dataset:
        label_path = annotations_dir / f"{img_path.stem}.txt"
        if label_path.exists():
            os.symlink(img_path, dest_image_dir / img_path.name)
            os.symlink(label_path, dest_label_dir / label_path.name)

copy_dataset(train_set, train_images_dir, train_labels_dir)
copy_dataset(val_set, val_images_dir, val_labels_dir)
copy_dataset(test_set, test_images_dir, test_labels_dir)

print("데이터셋 분리가 완료되었습니다.")

모델 훈련

다음 스크립트를 통해 YOLOv8 모델을 훈련합니다:

from ultralytics import YOLO

model = YOLO("yolov8n.pt")
results = model.train(
    data="../datasets/substation_inspection/substation_inspection.yaml",
    epochs=50,
    imgsz=128,
    batch=16,
    project="../runs/train",
    name="substation_inspection_detection"
)

metrics = model.val()
model.export(format="onnx")

평가 및 결과 시각화

훈련된 모델을 평가하고 성능을 확인합니다:

from ultralytics import YOLO

model = YOLO("../runs/train/substation_inspection_detection/weights/best.pt")
metrics = model.val(data="../datasets/substation_inspection/substation_inspection.yaml", conf=0.5, iou=0.45)
print(metrics)

GUI 개발

PyQt5를 사용하여 간단한 사용자 인터페이스를 생성합니다:

import sys
import cv2
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QPushButton, QVBoxLayout, QWidget, QFileDialog
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import Qt, QTimer
from ultralytics import YOLO

model = YOLO("../runs/train/substation_inspection_detection/weights/best.pt")

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("전력 변전소 결함 탐지 시스템")
        self.setGeometry(100, 100, 800, 600)
        self.init_ui()

    def init_ui(self):
        central_widget = QWidget()
        layout = QVBoxLayout()

        self.image_label = QLabel(self)
        self.image_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.image_label)

        load_button = QPushButton("이미지 선택", self)
        load_button.clicked.connect(self.load_image)
        layout.addWidget(load_button)

        predict_button = QPushButton("결과 예측", self)
        predict_button.clicked.connect(self.predict_image)
        layout.addWidget(predict_button)

        central_widget.setLayout(layout)
        self.setCentralWidget(central_widget)

        self.image_path = None

    def load_image(self):
        file_name, _ = QFileDialog.getOpenFileName(self, "이미지 선택", "", "Images (*.png *.jpg *.jpeg);;All Files (*)")
        if file_name:
            self.image_path = file_name
            pixmap = QPixmap(file_name)
            scaled_pixmap = pixmap.scaled(self.image_label.width(), self.image_label.height(), Qt.KeepAspectRatio)
            self.image_label.setPixmap(scaled_pixmap)

    def predict_image(self):
        if not self.image_path:
            return

        image = cv2.imread(self.image_path)
        results = model.predict(image, size=128, conf=0.5, iou=0.45)[0]

        for box in results.boxes.cpu().numpy():
            r = box.xyxy[0].astype(int)
            cls = int(box.cls[0])
            conf = box.conf[0]

            class_names = [
                "계기 손상", "절연체 손상", "유출", "호흡기 이상", "이물질",
                *[f"기타{i}" for i in range(12)]
            ]
            class_name = class_names[cls]

            cv2.rectangle(image, (r[0], r[1]), (r[2], r[3]), (0, 255, 0), 2)
            cv2.putText(
                image,
                f"{class_name} ({conf:.2f})",
                (r[0], r[1] - 10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.9,
                (0, 255, 0),
                2
            )

        h, w, ch = image.shape
        bytes_per_line = ch * w
        qt_image = QImage(image.data, w, h, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
        pixmap = QPixmap.fromImage(qt_image)
        scaled_pixmap = pixmap.scaled(self.image_label.width(), self.image_label.height(), Qt.KeepAspectRatio)
        self.image_label.setPixmap(scaled_pixmap)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())

위의 코드들을 순서대로 실행하면, 전력 변전소 가시광 이미지 결함 탐지 시스템을 완성할 수 있습니다.