如何训练自己的数据集之——绝缘子缺陷数据集

 绝缘子缺陷数据集 绝缘子缺陷检测图像数据集9类 1607张，train1285张，test162张val160张划分，标注为YOLOtxt格式可直接训练 绝缘体数据集玻璃脏数据集聚合物数据集 绝缘子 

绝缘子缺陷检测图像数据集 9类 标签names: ['Glassdirty', 'Glassloss', 'Polymer', 'Polymerdirty', 'Two glass', 'broken disc', 'insulator', 'pollution-flashover', 'snow']
名称：【‘玻璃脏’，‘玻璃损耗’，‘聚合物’，‘聚合体脏’，‘两片玻璃’，‘破盘’，‘绝缘体’，‘污闪’】共1607张，近似8:1:1（train：1285张，test：162张，val：160张）比例划分，标注文件为YOLO适用的txt格式。可以直接用于模型训练。

绝缘子缺陷检测图像数据集介绍

数据集概述

该数据集专注于绝缘子的各种缺陷检测，包含9种类别的标签，分别是：玻璃脏（Glassdirty）、玻璃损耗（Glassloss）、聚合物（Polymer）、聚合体脏（Polymerdirty）、两片玻璃（Two glass）、破盘（broken disc）、绝缘体（insulator）、污闪（pollution-flashover）、雪（snow）。数据集共有1607张图像，按照大约8:1:1的比例划分为训练集（1285张）、测试集（162张）和验证集（160张）。标注文件为YOLO适用的.txt格式，可以直接用于模型训练。

数据集特点

1. 多类别标签：数据集涵盖了绝缘子常见的多种缺陷类型，有助于模型学习和识别不同的缺陷模式。
2. 明确的数据划分：数据集按照标准的比例划分为训练集、验证集和测试集，便于模型训练和性能评估。
3. 适用性强：YOLO格式的标注文件方便使用YOLO框架进行训练，减少了数据预处理的工作量。

数据集内容

• 图像文件：共有1607张JPG/PNG格式的图像文件。
• 标注文件：每张图像都配有YOLO格式的.txt标注文件。

数据集结构示例

假设数据集的根目录为 insulator_defects_dataset，其结构可能如下所示：

insulator_defects_dataset/
├── images/
│   ├── train/
│   │   ├── train_image_0001.jpg
│   │   ├── train_image_0002.jpg
│   │   └── ...
│   ├── test/
│   │   ├── test_image_0001.jpg
│   │   ├── test_image_0002.jpg
│   │   └── ...
│   ├── val/
│   │   ├── val_image_0001.jpg
│   │   ├── val_image_0002.jpg
│   │   └── ...
├── labels_yolo/
│   ├── train/
│   │   ├── train_image_0001.txt
│   │   ├── train_image_0002.txt
│   │   └── ...
│   ├── test/
│   │   ├── test_image_0001.txt
│   │   ├── test_image_0002.txt
│   │   └── ...
│   ├── val/
│   │   ├── val_image_0001.txt
│   │   ├── val_image_0002.txt
│   │   └── ...
└── data.yaml  # 数据集配置文件

数据集配置文件 data.yaml

创建一个data.yaml文件来描述您的数据集。这里假设数据集被放置在一个名为insulator_defects_dataset的目录中，且包含images和labels子目录。

# data.yaml 文件
train: ../insulator_defects_dataset/images/train/
val: ../insulator_defects_dataset/images/val/
test: ../insulator_defects_dataset/images/test/

nc: 9  # number of classes
names: ['Glassdirty', 'Glassloss', 'Polymer', 'Polymerdirty', 'Two glass', 'broken disc', 'insulator', 'pollution-flashover', 'snow']  # class names

关键训练代码

安装YOLOv5

如果您还没有安装YOLOv5，请按照官方文档执行以下命令：

git clone https://github.com/ultralytics/yolov5.git  # clone repo
cd yolov5
pip install -r requirements.txt  # install dependencies

使用YOLOv5命令行训练

使用以下命令开始训练模型：

cd yolov5
python train.py --img 640 --batch 16 --epochs 300 --data ../insulator_defects_dataset/data.yaml --weights yolov5s.pt --cache

自定义训练脚本

如果需要更详细的控制，可以编写一个Python脚本来执行训练过程。以下是一个简单的脚本示例：

import torch
from utils.datasets import LoadImagesAndLabels  # 导入数据加载器
from models.experimental import attempt_load  # 导入模型加载器
from utils.torch_utils import select_device  # 导入选用设备的函数
from utils.general import check_dataset  # 导入检查数据集的函数

def main():
    device = select_device('')  # 选择设备，自动选择GPU/CPU
    data_yaml = '../insulator_defects_dataset/data.yaml'
    train_images_folder = '../insulator_defects_dataset/images/train/'
    train_labels_folder = '../insulator_defects_dataset/labels_yolo/train/'
    val_images_folder = '../insulator_defects_dataset/images/val/'
    val_labels_folder = '../insulator_defects_dataset/labels_yolo/val/'

    # 加载数据集
    train_set = LoadImagesAndLabels(train_images_folder, train_labels_folder)
    val_set = LoadImagesAndLabels(val_images_folder, val_labels_folder)

    # 创建数据加载器
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=16, shuffle=True, num_workers=4)
    val_loader = torch.utils.data.DataLoader(val_set, batch_size=16, shuffle=False, num_workers=4)

    # 加载预训练模型
    model = attempt_load('yolov5s.pt', map_location=device)  # 加载预训练权重
    model.to(device)

    # 设置损失函数和优化器
    criterion = torch.nn.BCEWithLogitsLoss()  # 适用于多标签分类
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # 训练循环
    epochs = 300
    for epoch in range(epochs):
        model.train()
        for images, targets, _, _ in train_loader:
            images = images.to(device)
            targets = [t.to(device) for t in targets]

            # 前向传播
            outputs = model(images)

            # 计算损失
            loss = criterion(outputs, targets)

            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item()}')

    torch.save(model.state_dict(), 'trained_model.pt')
    print('Training complete.')

if __name__ == '__main__':
    main()

注意事项

• 确保data.yaml文件中的路径是正确的，并且数据集的结构与上面描述的一致。
• 调整批量大小、学习率、迭代次数等超参数以适应您的计算资源和任务需求。
• 如果数据集很大，您可能需要更多的计算资源和时间来完成训练。
• 这个脚本仅作为一个起点，您可能需要根据实际情况做进一步的修改。

测试模型

在训练完成后，您可以通过以下命令测试模型的性能：

python detect.py --weights runs/train/exp/weights/best.pt --img 640 --conf 0.4 --source ../insulator_defects_dataset/images/test/

总结

这个示例展示了如何使用YOLOv5框架训练一个基于绝缘子缺陷检测的数据集。您可以根据自己的需求调整脚本中的参数和逻辑。通过使用这个数据集和相应的训练代码，您可以有效地训练出一个能够在多种条件下识别绝缘子缺陷的模型。