无人机电力设备航拍类 配网缺陷检测 无人机航拍图像数据集 不规范绑扎识别 螺栓销钉缺失 无人机配网缺陷不规范绑扎数据集，无人机螺栓销钉缺陷检测数据集，实现对电力设备缺陷的检测识别

深度学习框架目标检测算法使用YOLOv8进行训练电力设备航拍类/配网缺陷检测无人机航拍图像数据集的代码。进行模型训练和评估。

文章目录

• 深度学习框架目标检测算法使用YOLOv8进行训练电力设备航拍类/配网缺陷检测无人机航拍图像数据集的代码。进行模型训练和评估。
• • • 数据集介绍
    • • 数据集概述
      • 数据集目录结构
    • 数据集配置文件
    • 转换标注格式
    • • 转换脚本
    • YOLOv8训练代码
    • 详细解释
    • 运行训练脚本
    • 评估模型
    • 详细解释
    • 运行评估脚本
    • 一键运行脚本
    • 详细解释
    • 运行主脚本
    • • 转换标注格式
      • 训练模型
      • 评估模型

数据集介绍

数据集概述

• 数据集名称：Power Distribution Network Defect Detection Dataset (PDNDDD)
• 数据类型：无人机航拍图像
• 目标类别：2类
  • 不规范绑扎
  • 螺栓销钉缺失
• 样本数量：约3000张图片
• 标注格式：VOC格式

数据集目录结构

PDNDDD/
├── images/
│ ├── train/
│ └── val/
├── labels_voc/
│ ├── train/
│ └── val/
├── labels/
│ ├── train/
│ └── val/
└── data.yaml

数据集配置文件

创建一个data.yaml文件，配置数据集的路径和类别信息：

path: ./PDNDDD # 数据集路径
train: images/train # 训练集图像路径
val: images/val # 验证集图像路径

nc: 2 # 类别数
names: ['不规范绑扎', '螺栓销钉缺失'] # 类别名称

转换标注格式

假设标注文件是VOC格式的XML文件，我们需要将它们转换为YOLO格式的TXT文件。

转换脚本

import xml.etree.ElementTree as ET
import os

def convert_voc_to_yolo(voc_file, yolo_file, class_names):
tree = ET.parse(voc_file)
root = tree.getroot()

width = int(root.find('size/width').text)
height = int(root.find('size/height').text)

with open(yolo_file, 'w') as f:
for obj in root.findall('object'):
class_name = obj.find('name').text
if class_name not in class_names:
continue
class_id = class_names.index(class_name)

bbox = obj.find('bndbox')
x_min = float(bbox.find('xmin').text)
y_min = float(bbox.find('ymin').text)
x_max = float(bbox.find('xmax').text)
y_max = float(bbox.find('ymax').text)

x_center = (x_min + x_max) / 2.0 / width
y_center = (y_min + y_max) / 2.0 / height
w = (x_max – x_min) / width
h = (y_max – y_min) / height

f.write(f"{class_id} {x_center} {y_center} {w} {h}\\n")

def convert_all_voc_to_yolo(voc_dir, yolo_dir, class_names):
os.makedirs(yolo_dir, exist_ok=True)
for filename in os.listdir(voc_dir):
if filename.endswith('.xml'):
voc_file = os.path.join(voc_dir, filename)
yolo_file = os.path.join(yolo_dir, filename.replace('.xml', '.txt'))
convert_voc_to_yolo(voc_file, yolo_file, class_names)

if __name__ == "__main__":
class_names = ['不规范绑扎', '螺栓销钉缺失']
voc_train_dir = 'PDNDDD/labels_voc/train'
yolo_train_dir = 'PDNDDD/labels/train'
convert_all_voc_to_yolo(voc_train_dir, yolo_train_dir, class_names)

voc_val_dir = 'PDNDDD/labels_voc/val'
yolo_val_dir = 'PDNDDD/labels/val'
convert_all_voc_to_yolo(voc_val_dir, yolo_val_dir, class_names)

YOLOv8训练代码

详细解释

运行训练脚本

将上述脚本保存为一个Python文件（例如train_yolov8_pdnddd.py），然后运行它。

python train_yolov8_pdnddd.py

评估模型

详细解释

运行评估脚本

将上述脚本保存为一个Python文件（例如evaluate_yolov8_pdnddd.py），然后运行它。

python evaluate_yolov8_pdnddd.py

一键运行脚本

为了实现一键运行，可以将训练和评估脚本合并到一个主脚本中，并添加命令行参数来控制运行模式。

import argparse
from ultralytics import YOLO

def convert_voc_to_yolo(voc_file, yolo_file, class_names):
tree = ET.parse(voc_file)
root = tree.getroot()

width = int(root.find('size/width').text)
height = int(root.find('size/height').text)

with open(yolo_file, 'w') as f:
for obj in root.findall('object'):
class_name = obj.find('name').text
if class_name not in class_names:
continue
class_id = class_names.index(class_name)

bbox = obj.find('bndbox')
x_min = float(bbox.find('xmin').text)
y_min = float(bbox.find('ymin').text)
x_max = float(bbox.find('xmax').text)
y_max = float(bbox.find('ymax').text)

x_center = (x_min + x_max) / 2.0 / width
y_center = (y_min + y_max) / 2.0 / height
w = (x_max – x_min) / width
h = (y_max – y_min) / height

f.write(f"{class_id} {x_center} {y_center} {w} {h}\\n")

def convert_all_voc_to_yolo(voc_dir, yolo_dir, class_names):
os.makedirs(yolo_dir, exist_ok=True)
for filename in os.listdir(voc_dir):
if filename.endswith('.xml'):
voc_file = os.path.join(voc_dir, filename)
yolo_file = os.path.join(yolo_dir, filename.replace('.xml', '.txt'))
convert_voc_to_yolo(voc_file, yolo_file, class_names)

def train_model(data_yaml_path, model_config, epochs, batch_size, img_size, augment):
# 加载模型
model = YOLO(model_config)

# 训练模型
results = model.train(
data=data_yaml_path,
epochs=epochs,
batch=batch_size,
imgsz=img_size,
augment=augment
)

# 保存模型
model.save("runs/train/power_distribution_network_defect_detection/best.pt")

def evaluate_model(data_yaml_path, weights_path, img_size, conf_threshold):
# 加载模型
model = YOLO(weights_path)

# 评估模型
results = model.val(
data=data_yaml_path,
imgsz=img_size,
conf=conf_threshold
)

# 打印评估结果
print(results)

def main(mode):
class_names = ['不规范绑扎', '螺栓销钉缺失']
data_yaml_path = 'PDNDDD/data.yaml'
model_config = 'yolov8s.yaml'
epochs = 100
batch_size = 16
img_size = 640
conf_threshold = 0.4
augment = True

if mode == 'convert':
voc_train_dir = 'PDNDDD/labels_voc/train'
yolo_train_dir = 'PDNDDD/labels/train'
convert_all_voc_to_yolo(voc_train_dir, yolo_train_dir, class_names)

voc_val_dir = 'PDNDDD/labels_voc/val'
yolo_val_dir = 'PDNDDD/labels/val'
convert_all_voc_to_yolo(voc_val_dir, yolo_val_dir, class_names)
elif mode == 'train':
train_model(data_yaml_path, model_config, epochs, batch_size, img_size, augment)
elif mode == 'eval':
weights_path = 'runs/train/power_distribution_network_defect_detection/best.pt'
evaluate_model(data_yaml_path, weights_path, img_size, conf_threshold)
else:
print("Invalid mode. Use 'convert', 'train', or 'eval'.")

if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Convert, train, or evaluate YOLOv8 on the Power Distribution Network Defect Detection Dataset.")
parser.add_argument('mode', type=str, choices=['convert', 'train', 'eval'], help="Mode: 'convert', 'train', or 'eval'")
args = parser.parse_args()

main(args.mode)

详细解释

运行主脚本

将上述脚本保存为一个Python文件（例如main_yolov8_pdnddd.py），然后运行它。

转换标注格式

python main_yolov8_pdnddd.py convert

训练模型

python main_yolov8_pdnddd.py train

评估模型

python main_yolov8_pdnddd.py eval