电力领域类——如何使用 YOLOv8 训练防震锤高压塔线路行业数据集 防震锤数据集 2类

电力领域类——如何使用 YOLOv8 训练防震锤高压塔线路行业数据集 防震锤数据集 2类

如何使用 YOLOv8 训练防震锤高压塔线路行业数据集的详细步骤。这个数据集主要用于检测高压塔线路上的防震锤是否存在缺陷，数据集包含1000张图片，共2700+个样本，标签为“缺陷”和“未缺陷”，样本比为933:1808。

一、数据准备

1. 数据集结构
   假设你的数据集已经按照 VOC 格式组织，并且结构如下：
   深色版本
   power_line_dataset/
   ├── images/
   │ ├── train/
   │ │ ├── 000001.jpg
   │ │ ├── 000002.jpg
   │ │ └── …
   │ ├── val/
   │ │ ├── 000001.jpg
   │ │ ├── 000002.jpg
   │ │ └── …
   │ └── test/
   │ ├── 000001.jpg
   │ ├── 000002.jpg
   │ └── …
   ├── annotations/
   │ ├── train/
   │ │ ├── 000001.xml
   │ │ ├── 000002.xml
   │ │ └── …
   │ ├── val/
   │ │ ├── 000001.xml
   │ │ ├── 000002.xml
   │ │ └── …
   │ └── test/
   │ ├── 000001.xml
   │ ├── 000002.xml
   │ └── …
   └── power_line.yaml
   images/ 目录包含输入图像。
   annotations/ 目录包含对应的标注文件，每个标注文件为 XML 格式。
   power_line.yaml 是数据集的配置文件。
   2. 配置文件
   创建一个 power_line.yaml 文件，内容如下：

yaml
深色版本

数据集路径

path: ./power_line_dataset

类别名称

names:
0: defect
1: no_defect

数据集划分

train: images/train
val: images/val
test: images/test
二、转换标注格式
YOLOv8 需要标注文件为 YOLO 格式，即每行表示一个目标框，格式为 class_id x_center y_center width height。你需要将 VOC 格式的标注文件转换为 YOLO 格式。

1. 转换脚本
   python
   深色版本
   import xml.etree.ElementTree as ET
   import os
   def convert_voc_to_yolo(xml_file, output_dir, class_names):
   tree = ET.parse(xml_file)
   root = tree.getroot()
   size = root.find(‘size’)
   w = int(size.find(‘width’).text)
   h = int(size.find(‘height’).text)

yolo_lines = []
for obj in root.findall('object'):
    cls = obj.find('name').text
    if cls not in class_names:
        continue
    cls_id = class_names.index(cls)
    xmlbox = obj.find('bndbox')
    b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
    bb = convert((w, h), b)
    yolo_lines.append(f"{cls_id} {bb[0]} {bb[1]} {bb[2]} {bb[3]}\n")

output_file = os.path.join(output_dir, os.path.basename(xml_file).replace('.xml', '.txt'))
with open(output_file, 'w') as f:
    f.writelines(yolo_lines)

def convert(size, box):
dw = 1./size[0]
dh = 1./size[1]
x = (box[0] + box[1]) / 2.0
y = (box[2] + box[3]) / 2.0
w = box[1] - box[0]
h = box[3] - box[2]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return (x, y, w, h)

转换训练集

train_xml_dir = ‘power_line_dataset/annotations/train’
train_yolo_dir = ‘power_line_dataset/labels/train’
os.makedirs(train_yolo_dir, exist_ok=True)
for xml_file in os.listdir(train_xml_dir):
convert_voc_to_yolo(os.path.join(train_xml_dir, xml_file), train_yolo_dir, [‘defect’, ‘no_defect’])

转换验证集

val_xml_dir = ‘power_line_dataset/annotations/val’
val_yolo_dir = ‘power_line_dataset/labels/val’
os.makedirs(val_yolo_dir, exist_ok=True)
for xml_file in os.listdir(val_xml_dir):
convert_voc_to_yolo(os.path.join(val_xml_dir, xml_file), val_yolo_dir, [‘defect’, ‘no_defect’])

转换测试集

test_xml_dir = ‘power_line_dataset/annotations/test’
test_yolo_dir = ‘power_line_dataset/labels/test’
os.makedirs(test_yolo_dir, exist_ok=True)
for xml_file in os.listdir(test_xml_dir):
convert_voc_to_yolo(os.path.join(test_xml_dir, xml_file), test_yolo_dir, [‘defect’, ‘no_defect’])
三、安装依赖
确保你已经安装了 YOLOv8 所需的依赖。可以使用以下命令安装：

bash
深色版本
pip install ultralytics
四、训练模型

1. 导入库
   python
   深色版本
   import torch
   from ultralytics import YOLO1. 加载模型
   python
   深色版本

加载预训练的YOLOv8模型

model = YOLO(‘yolov8n.yaml’) # 你可以选择其他预训练模型，如’yolov8s.yaml’、'yolov8m.yaml’等
3. 配置训练参数
python
深色版本

配置训练参数

model.train(
data=‘power_line.yaml’, # 数据集配置文件
epochs=100, # 训练轮数
imgsz=640, # 输入图像大小
batch=16, # 批量大小
name=‘power_line_detection’, # 训练结果保存的文件夹名称
device=‘cuda’ # 使用GPU进行训练
)
五、评估模型
训练完成后，可以使用验证集来评估模型的性能。

python
深色版本

评估模型

results = model.val(data=‘power_line.yaml’, imgsz=640, batch=16, device=‘cuda’)

打印评估结果

print(results)
六、保存模型
训练完成后，模型会自动保存到指定的文件夹中。你也可以手动保存模型：

python
深色版本

保存模型

model.save(‘power_line_detection_model.pt’)
七、完整代码
将上述所有步骤整合成一个完整的脚本：

python
深色版本
import torch
from ultralytics import YOLO

加载预训练的YOLOv8模型

model = YOLO(‘yolov8n.yaml’) # 你可以选择其他预训练模型，如’yolov8s.yaml’、'yolov8m.yaml’等

配置训练参数

model.train(
data=‘power_line.yaml’, # 数据集配置文件
epochs=100, # 训练轮数
imgsz=640, # 输入图像大小
batch=16, # 批量大小
name=‘power_line_detection’, # 训练结果保存的文件夹名称
device=‘cuda’ # 使用GPU进行训练
)

评估模型

results = model.val(data=‘power_line.yaml’, imgsz=640, batch=16, device=‘cuda’)

打印评估结果

print(results)

保存模型

model.save(‘power_line_detection_model.pt’)
八、运行脚本
将上述代码保存为一个 Python 脚本文件（例如 train_power_line_detection.py），然后在命令行中运行：

bash
深色版本
python train_power_line_detection.py
九、注意事项
数据集路径：确保数据集路径正确，特别是图片和标注文件的路径。
硬件资源：训练大型模型可能需要较多的计算资源，建议使用 GPU。
超参数调整：根据实际情况调整学习率、批量大小等超参数，以获得更好的训练效果。
数据增强：为了提高模型的泛化能力，可以在训练过程中使用数据增强技术，如随机裁剪、翻转等。
希望这些步骤能帮助你成功训练一个防震锤高压塔线路检测模型。