Dữ liệu huấn luyện này được xây dựng dành riêng cho bài toán phát hiện tự động các khuyết tật trên hệ thống truyền tải điện, bao gồm cột điện, xà ngang, kẹp nối dây, và đầu tháp. Tập dữ liệu gồm 9.838 ảnh JPG chất lượng cao với tổng cộng 44.520 bounding box, được gán nhãn chi tiết theo năm lớp khuyết tật cụ thể:

• Định vị dây không chuẩn: 3.717 ảnh, 14.510 hộp giới hạn
• Thiếu vỏ bảo vệ kẹp nối dây: 3.317 ảnh, 11.285 hộp giới hạn
• Thiếu vỏ bảo vệ kẹp chịu lực: 3.748 ảnh, 16.148 hộp giới hạn
• Ăn mòn xà ngang: 987 ảnh, 1.556 hộp giới hạn
• Hư hỏng đầu tháp: 972 ảnh, 1.021 hộp giới hạn

Cấu trúc thư mục chuẩn hóa:

power_infra_defects/
├── dataset/
│   ├── images/          # ảnh gốc (.jpg)
│   ├── labels_voc/      # nhãn VOC (.xml)
│   └── labels_yolo/     # nhãn YOLO (.txt)
├── models/
│   └── yolov8/
├── src/
│   ├── convert_voc_to_yolo.py
│   ├── dataset_loader.py
│   ├── trainer.py
│   └── evaluator.py
├── weights/
│   └── best.pt
├── requirements.txt
└── README.md

1. Thiết lập môi trường

Tạo file requirements.txt với nội dung sau:

torch>=2.0.0
ultralytics>=8.2.0
opencv-python>=4.8.0
numpy>=1.24.0
Pillow>=10.0.0
tqdm>=4.66.0
lxml>=4.9.0

Chạy lệnh cài đặt:

pip install -r requirements.txt

2. Chuyển đổi nhãn từ VOC sang YOLO

Tập tin src/convert_voc_to_yolo.py thực hiện chuyển đổi tự động bằng cách chuẩn hóa tọa độ và ánh xạ lớp:

import os
import xml.etree.ElementTree as ET
from pathlib import Path

def parse_voc_annotation(xml_path, class_map):
    tree = ET.parse(xml_path)
    root = tree.getroot()
    
    width = int(root.find("size/width").text)
    height = int(root.find("size/height").text)
    
    annotations = []
    for obj in root.findall("object"):
        name = obj.find("name").text.strip()
        if name not in class_map:
            continue
            
        bbox = obj.find("bndbox")
        xmin = max(0, int(bbox.find("xmin").text))
        ymin = max(0, int(bbox.find("ymin").text))
        xmax = min(width - 1, int(bbox.find("xmax").text))
        ymax = min(height - 1, int(bbox.find("ymax").text))
        
        # Chuẩn hóa theo định dạng YOLOv8: [class_id, x_center, y_center, w, h]
        x_center = (xmin + xmax) / (2 * width)
        y_center = (ymin + ymax) / (2 * height)
        box_w = (xmax - xmin) / width
        box_h = (ymax - ymin) / height
        
        annotations.append(f"{class_map[name]} {x_center:.6f} {y_center:.6f} {box_w:.6f} {box_h:.6f}")
    
    return annotations

def convert_all(voc_dir, yolo_dir, class_list):
    voc_path = Path(voc_dir)
    yolo_path = Path(yolo_dir)
    yolo_path.mkdir(exist_ok=True)
    
    class_map = {cls: idx for idx, cls in enumerate(class_list)}
    
    for xml_file in voc_path.glob("*.xml"):
        yolo_lines = parse_voc_annotation(xml_file, class_map)
        yolo_file = yolo_path / f"{xml_file.stem}.txt"
        
        with open(yolo_file, "w", encoding="utf-8") as f:
            f.write("\n".join(yolo_lines))

if __name__ == "__main__":
    CLASSES = [
        "Định vị dây không chuẩn",
        "Thiếu vỏ bảo vệ kẹp nối dây",
        "Thiếu vỏ bảo vệ kẹp chịu lực",
        "Ăn mòn xà ngang",
        "Hư hỏng đầu tháp"
    ]
    
    convert_all("dataset/labels_voc", "dataset/labels_yolo", CLASSES)

3. Tải và tiền xử lý dữ liệu

Tập tin src/dataset_loader.py định nghĩa bộ dữ liệu tùy chỉnh hỗ trợ tăng cường ảnh và chuẩn hóa:

from torch.utils.data import Dataset
from PIL import Image
import torch
import numpy as np
import os

class PowerInfrastructureDataset(Dataset):
    def __init__(self, img_dir, label_dir, img_size=640, augment=False):
        self.img_dir = img_dir
        self.label_dir = label_dir
        self.img_size = img_size
        self.augment = augment
        self.img_files = [f for f in os.listdir(img_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
    
    def __len__(self):
        return len(self.img_files)
    
    def __getitem__(self, idx):
        img_name = self.img_files[idx]
        img_path = os.path.join(self.img_dir, img_name)
        label_path = os.path.join(self.label_dir, img_name.rsplit('.', 1)[0] + '.txt')
        
        # Đọc ảnh
        image = Image.open(img_path).convert("RGB")
        image = image.resize((self.img_size, self.img_size), Image.BILINEAR)
        image = np.array(image) / 255.0
        image = torch.from_numpy(image).permute(2, 0, 1).float()
        
        # Đọc nhãn
        boxes = torch.zeros((0, 5))
        if os.path.exists(label_path):
            with open(label_path, 'r') as f:
                lines = f.readlines()
            for line in lines:
                parts = list(map(float, line.strip().split()))
                if len(parts) == 5:
                    boxes = torch.cat([boxes, torch.tensor([parts]).float()], dim=0)
        
        return image, boxes

def create_dataloaders(img_dir, label_dir, batch_size=16, val_split=0.2, num_workers=4):
    dataset = PowerInfrastructureDataset(img_dir, label_dir)
    val_len = int(len(dataset) * val_split)
    train_len = len(dataset) - val_len
    
    train_ds, val_ds = torch.utils.data.random_split(dataset, [train_len, val_len])
    
    train_loader = torch.utils.data.DataLoader(
        train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers
    )
    val_loader = torch.utils.data.DataLoader(
        val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers
    )
    
    return train_loader, val_loader

4. Huấn luyện mô hình YOLOv8

Sử dụng thư viện ultralytics để huấn luyện hiệu quả với cấu hình tối ưu:

# src/trainer.py
from ultralytics import YOLO
import yaml

def train_yolov8():
    # Định nghĩa cấu hình dataset
    data_config = {
        "train": "dataset/images",
        "val": "dataset/images",
        "nc": 5,
        "names": [
            "Định vị dây không chuẩn",
            "Thiếu vỏ bảo vệ kẹp nối dây",
            "Thiếu vỏ bảo vệ kẹp chịu lực",
            "Ăn mòn xà ngang",
            "Hư hỏng đầu tháp"
        ]
    }
    
    with open("dataset_config.yaml", "w") as f:
        yaml.dump(data_config, f)
    
    # Khởi tạo và huấn luyện
    model = YOLO("yolov8n.pt")
    results = model.train(
        data="dataset_config.yaml",
        epochs=120,
        imgsz=640,
        batch=16,
        name="power_defect_exp1",
        project="runs/train",
        device=0 if torch.cuda.is_available() else "cpu"
    )

if __name__ == "__main__":
    train_yolov8()

5. Đánh giá và trực quan hóa kết quả

Tập tin src/evaluator.py hiển thị so sánh giữa nhãn thật và dự đoán trên mẫu kiểm tra:

# src/evaluator.py
from ultralytics import YOLO
import cv2
import numpy as np
from pathlib import Path

def visualize_predictions(model_path, image_dir, num_samples=6):
    model = YOLO(model_path)
    img_paths = list(Path(image_dir).glob("*.jpg"))[:num_samples]
    
    for img_path in img_paths:
        img = cv2.imread(str(img_path))
        results = model(img)
        
        # Vẽ bounding box dự đoán
        annotated = results[0].plot()
        
        # Hiển thị
        cv2.imshow("Prediction", annotated)
        cv2.waitKey(0)
    
    cv2.destroyAllWindows()

if __name__ == "__main__":
    visualize_predictions("weights/best.pt", "dataset/images")

6. Các bước thực thi

1. Chạy chuyển đổi nhãn:
   python src/convert_voc_to_yolo.py
2. Khởi tạo huấn luyện:
   python src/trainer.py
3. Đánh giá mô hình:
   python src/evaluator.py