如何使用yolov5训练并使用——封边机压轮缺陷数据集 6种缺陷 标签已经标注

封边机压轮缺陷数据集。适用于yolo系列目标检测模型。
封边机压轮工业缺陷数据集。包含原图像254张，标签也已经标注好。可以扩充数据集，进行数据集增强，达到5000张。
缺陷类别包含:烧伤，挤伤，划痕，切割，磨损，点蚀共6种缺陷。

好的，我会详细介绍这个封边机压轮工业缺陷数据集，并提供一个完整的YOLOv5目标检测模型的训练代码。我们将从数据集的介绍、数据增强、模型训练和验证等几个方面进行说明。

1. 数据集介绍

数据集概述

• 数据集名称: 封边机压轮工业缺陷数据集
• 数据集来源: 全网独家，独立制作
• 数据集内容: 包含254张原图像，每张图像都有对应的标签文件，标签文件采用YOLO格式。
• 缺陷类别: 烧伤、挤伤、划痕、切割、磨损、点蚀（共6种）
• 数据集扩展: 可以通过数据增强技术将数据集扩展到5000张图像。
  

数据集目录结构

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

2. 数据增强

数据增强可以帮助提高模型的泛化能力。我们将使用albumentations库来进行数据增强。

安装依赖

pip install albumentations

数据增强代码

# src/utils/data_augmentation.py
import albumentations as A
from albumentations.pytorch import ToTensorV2
import cv2
import os
import shutil

def augment_image(image_path, label_path, output_dir, num_augmentations=10):
    # 读取图像和标签
    image = cv2.imread(image_path)
    with open(label_path, 'r') as f:
        labels = [line.strip().split() for line in f.readlines()]
    
    bboxes = [[float(x) for x in label[1:]] for label in labels]
    class_labels = [int(label[0]) for label in labels]

    # 定义数据增强变换
    transform = A.Compose([
        A.HorizontalFlip(p=0.5),
        A.VerticalFlip(p=0.5),
        A.RandomRotate90(p=0.5),
        A.RandomBrightnessContrast(p=0.2),
        A.RandomGamma(p=0.2),
        A.Blur(blur_limit=3, p=0.2),
        A.CLAHE(p=0.2),
        A.HueSaturationValue(p=0.2),
        A.RandomResizedCrop(height=image.shape[0], width=image.shape[1], scale=(0.8, 1.0), p=0.5),
        ToTensorV2()
    ], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))

    # 应用数据增强
    for i in range(num_augmentations):
        augmented = transform(image=image, bboxes=bboxes, class_labels=class_labels)
        augmented_image = augmented['image']
        augmented_bboxes = augmented['bboxes']
        augmented_class_labels = augmented['class_labels']

        # 保存增强后的图像
        output_image_path = os.path.join(output_dir, f'{os.path.basename(image_path)}_aug_{i}.jpg')
        cv2.imwrite(output_image_path, augmented_image.permute(1, 2, 0).numpy())

        # 保存增强后的标签
        output_label_path = os.path.join(output_dir, f'{os.path.basename(label_path)}_aug_{i}.txt')
        with open(output_label_path, 'w') as f:
            for j in range(len(augmented_bboxes)):
                f.write(f"{augmented_class_labels[j]} {' '.join(map(str, augmented_bboxes[j]))}\n")

def augment_dataset(input_dir, output_dir, num_augmentations=10):
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir, exist_ok=True)

    for split in ['train', 'val']:
        input_split_dir = os.path.join(input_dir, split)
        output_split_dir = os.path.join(output_dir, split)
        os.makedirs(output_split_dir, exist_ok=True)

        for image_file in os.listdir(os.path.join(input_split_dir, 'images')):
            image_path = os.path.join(input_split_dir, 'images', image_file)
            label_path = os.path.join(input_split_dir, 'labels', image_file.replace('.jpg', '.txt'))
            augment_image(image_path, label_path, output_split_dir, num_augmentations)

if __name__ == "__main__":
    input_dir = 'dataset'
    output_dir = 'augmented_dataset'
    augment_dataset(input_dir, output_dir, num_augmentations=10)

3. 训练YOLOv5模型

安装YOLOv5

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

数据集配置文件

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

# data.yaml
train: ../augmented_dataset/images/train
val: ../augmented_dataset/images/val

nc: 6
names: ['烧伤', '挤伤', '划痕', '切割', '磨损', '点蚀']

训练脚本

# src/train_yolov5.py
import torch
from yolov5.models.experimental import attempt_load
from yolov5.utils.torch_utils import select_device
from yolov5.utils.general import check_img_size
from yolov5.utils.datasets import create_dataloader
from yolov5.utils.loss import ComputeLoss
from yolov5.models.yolo import Model
import time
import yaml

def train_model(data_yaml_path, model_config, epochs, batch_size, img_size, device):
    # 读取数据集配置
    with open(data_yaml_path, 'r') as f:
        data = yaml.safe_load(f)

    # 创建数据加载器
    train_loader = create_dataloader(data['train'], img_size, batch_size, 32)[0]
    val_loader = create_dataloader(data['val'], img_size, batch_size, 32)[0]

    # 加载模型
    model = Model(model_config, ch=3, nc=data['nc']).to(device)
    model.train()

    # 定义优化器和损失函数
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    compute_loss = ComputeLoss(model)

    # 记录训练时间
    start_time = time.time()

    for epoch in range(epochs):
        epoch_start_time = time.time()
        for i, (imgs, targets, paths, _) in enumerate(train_loader):
            imgs = imgs.to(device)
            targets = targets.to(device)

            # 前向传播
            pred = model(imgs)

            # 计算损失
            loss, loss_items = compute_loss(pred, targets)

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

            if i % 10 == 0:
                print(f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item()}')

        epoch_end_time = time.time()
        epoch_duration = epoch_end_time - epoch_start_time
        print(f'Epoch [{epoch+1}/{epochs}] completed in {epoch_duration:.2f} seconds')

        # 保存模型
        torch.save(model.state_dict(), f'models/yolov5_custom_epoch_{epoch+1}.pth')

    end_time = time.time()
    total_duration = end_time - start_time
    print(f'Total training time: {total_duration:.2f} seconds')

if __name__ == "__main__":
    data_yaml_path = 'data.yaml'
    model_config = 'models/yolov5s.yaml'
    epochs = 100
    batch_size = 16
    img_size = 640
    device = select_device('0')  # 使用GPU，如果需要使用CPU，可以改为'cpu'

    train_model(data_yaml_path, model_config, epochs, batch_size, img_size, device)

4. 运行训练脚本

1. 数据增强：

   python src/utils/data_augmentation.py
   

2. 训练模型：

   python src/train_yolov5.py
   

5. 详细解释

数据增强

• augment_image: 对单张图像进行数据增强，并保存增强后的图像和标签。
• augment_dataset: 对整个数据集进行数据增强，生成新的数据集。

训练脚本

• 数据加载：使用create_dataloader创建训练和验证的数据加载器。
• 模型加载：加载YOLOv5模型，并设置为训练模式。
• 优化器和损失函数：定义Adam优化器和YOLOv5的损失函数。
• 训练循环：记录每个epoch的训练时间和总训练时间，进行前向传播、计算损失、反向传播和优化。
• 保存模型：每个epoch结束后保存模型的权重。

6. 注意事项

1. 数据集路径：确保数据集路径正确，特别是data.yaml文件中的路径。
2. 模型配置：确保模型配置文件路径正确。
3. 图像大小：img_size可以根据实际需求调整，通常使用640或1280。
4. 设备：确保设备（CPU或GPU）可用。

7. 总结

通过以上步骤，你可以构建一个完整的封边机压轮工业缺陷检测数据集，并使用YOLOv5进行训练。