如何训练电力设备类输电线电线杆检测数据集建立基于深度学习yolov5输电线电线杆检测系统

最新推荐文章于 2025-05-16 14:49:58 发布

计算机C9硕士_算法工程师

最新推荐文章于 2025-05-16 14:49:58 发布

阅读量948

点赞数 21

分类专栏：电力电气领域数据集文章标签：深度学习 YOLO 人工智能

本文链接：https://blog.csdn.net/2401_88440984/article/details/147968409

版权

电力电气领域数据集专栏收录该内容

130 篇文章

订阅专栏

如何训练电力设备类输电线电线杆检测数据集建立基于深度学习yolov5输电线电线杆检测系统

文章目录

电线杆数据集 Yolo格式在这里插入图片描述
1

1

基于YOLOv5的输电线和电线杆检测系统，你需要准备数据集、配置模型、进行训练，并最终将模型集成到应用程序中。

代码示例：仅供参考。

1. 数据准备

假设你已经有了标注好的数据集，包括输电线和电线杆的图像及其对应的YOLO格式标注文件。

数据集结构

dataset/
├── images/
│   ├── train/
│   │   ├── image1.jpg
│   │   ├── image2.jpg
│   │   └── ...
│   └── val/
│       ├── image1.jpg
│       ├── image2.jpg
│       └── ...
└── labels/
    ├── train/
    │   ├── image1.txt
    │   ├── image2.txt
    │   └── ...
    └── val/
        ├── image1.txt
        ├── image2.txt
        └── ...

2. 配置数据集

创建一个名为electric_poles.yaml的文件，用于定义数据集路径和类别信息：

train: dataset/images/train
val: dataset/images/val

nc: 1  # 类别数量（输电线和电线杆）
names: [ 'electric_pole']

3. 环境搭建

确保安装了必要的库和工具：

pip install torch torchvision torchaudio
git clone https://github.com/ultralytics/yolov5  # 克隆YOLOv5仓库
cd yolov5
pip install -r requirements.txt

4. 模型训练

使用YOLOv5的训练脚本来训练模型。在yolov5目录下运行以下命令：

python train.py --img 640 --batch 16 --epochs 100 --data electric_poles.yaml --weights yolov5s.pt

--img: 输入图像尺寸（例如640）。
--batch: 批次大小（例如16）。
--epochs: 训练周期数（例如100）。
--data: 数据配置文件路径（即electric_poles.yaml）。
--weights: 预训练权重路径（可以使用YOLOv5的预训练模型作为起点，如yolov5s.pt）。

5. 使用训练好的权重进行推理

单张图片推理

可以使用以下命令对单张图片进行预测：

python detect.py --source path/to/image.jpg --weights runs/train/exp/weights/best.pt --conf 0.25

批量推理

如果想要对一个目录中的所有图像进行批量推理，只需指定目录路径作为--source参数的值：

python detect.py --source path/to/images_folder --weights runs/train/exp/weights/best.pt --conf 0.25

6. 将系统集成到应用程序中

为了将这个系统集成到更大的应用程序中，你可以直接调用YOLOv5的Python API来执行推理。这里是一个简单的示例代码片段，演示如何加载模型并对图像进行推理：

import torch

# 加载模型
model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/exp/weights/best.pt')  # 或者直接使用路径加载

# 图像路径
img_path = 'path/to/image.jpg'

# 推理
results = model(img_path)

# 显示结果
results.print()  # 控制台输出
results.show()   # 显示标注后的图像
results.save()   # 保存标注后的图像到'runs/detect/exp'目录下

7. 性能评估

在部署前，建议对模型进行全面的性能评估，包括但不限于mAP、召回率、精确率等指标。这可以通过运行val.py脚本完成：

python val.py --data electric_poles.yaml --weights runs/train/exp/weights/best.pt --img 640

8. 系统优化与部署

性能优化：根据实际需要调整输入尺寸、批次大小、以及后处理参数（如NMS阈值）等以达到更好的性能或准确率。
部署：可以选择将模型转换为ONNX格式或其他轻量级格式以便于在不同的平台上部署。YOLOv5支持导出为多种格式，方便部署到移动端或服务器端。

# 导出为ONNX格式
python export.py --weights runs/train/exp/weights/best.pt --include onnx

构建一个基于深度学习YOLOv5的输电线和电线杆检测系统。

用户界面（GUI）。以下是一个简单的示例，展示了如何结合YOLOv5的推理功能与PyQt5 GUI框架。

1. 安装依赖

首先，确保你已经安装了必要的依赖：

pip install torch torchvision torchaudio PyQt5

2. 创建主函数及界面代码

如何加载训练好的YOLOv5模型，并通过PyQt5提供的界面进行图像选择、处理以及结果显示。

`main.py` - 主函数入口

import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QPushButton, QLabel, QFileDialog, QVBoxLayout, QWidget
from PyQt5.QtGui import QPixmap
import torch
from PIL import Image
from torchvision import transforms

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("输电线和电线杆检测系统")
        self.setGeometry(100, 100, 800, 600)

        # 加载YOLOv5模型
        self.model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/exp/weights/best.pt')

        # 界面组件
        self.image_label = QLabel(self)
        self.select_image_button = QPushButton("选择图片", self)
        self.select_image_button.clicked.connect(self.select_image)

        # 布局
        layout = QVBoxLayout()
        layout.addWidget(self.image_label)
        layout.addWidget(self.select_image_button)

        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)

    def select_image(self):
        file_name, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "Images (*.png *.xpm *.jpg)")
        if file_name:
            self.detect_objects(file_name)

    def detect_objects(self, image_path):
        img = Image.open(image_path)
        results = self.model(img)
        result_img = results.render()[0]  # 得到标注后的图像

        # 显示结果
        qimg = QPixmap.fromImage(Image.fromarray(result_img).toqimage())
        self.image_label.setPixmap(qimg)


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