如何使用YOLOv8训练轴承缺陷目标检测数据集。我们将从数据准备、模型训练、模型评估 460张 两种检测目标轴承缺陷数据集

轴承缺陷目标检测数据集

460张
两种检测目标：伪磨损+断裂
带标注 -YOLO格式 可直接用于YOLO系列目标检测算法模型训练


好的，让我们详细讨论如何使用YOLOv8训练轴承缺陷目标检测数据集。我们将从数据准备、模型训练、模型评估到最终集成到PyQt5应用中进行一步步的指导。

1. 数据准备

1.1. 组织数据

确保你的数据按照以下结构组织：

dataset/
│
├── images/
│   ├── train/
│   ├── val/
│   └── test/
│
├── labels/
│   ├── train/
│   ├── val/
│   └── test/
│
└── data.yaml

1.2. 创建数据配置文件

创建一个data.yaml文件来配置数据集：

train: ./dataset/images/train
val: ./dataset/images/val
test: ./dataset/images/test

nc: 2  # 类别数量（伪磨损、断裂）
names: ['false_wear', 'breakage']  # 类别名称

# 下载数据集
download: ''

2. 安装YOLOv8

确保你已经安装了YOLOv8。YOLOv8是YOLO系列的最新版本，可以从Ultralytics的GitHub仓库中获取。

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

3. 训练模型

3.1. 配置训练

使用YOLOv8进行目标检测。以下是一个示例命令：

python train.py --data ./dataset/data.yaml --img 640 --batch 16 --epochs 100 --name yolov8_custom --weights yolov8n.pt

3.2. 微调模型

如果你想微调一个预训练模型，可以指定模型路径：

python train.py --data ./dataset/data.yaml --img 640 --batch 16 --epochs 100 --name yolov8_custom --weights yolov8n.pt

4. 评估模型

4.1. 运行评估

在验证集上评估训练好的模型：

python val.py --data ./dataset/data.yaml --weights runs/train/yolov8_custom/weights/best.pt

4.2. 可视化结果

你可以使用val命令的--save标志来可视化结果：

python val.py --data ./dataset/data.yaml --weights runs/train/yolov8_custom/weights/best.pt --save

5. 使用PyQt5构建系统

5.1. 安装PyQt5

pip install PyQt5

5.2. 创建PyQt5应用

创建一个简单的PyQt5应用来展示检测结果。

import sys
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QLabel, QPushButton, QFileDialog
from PyQt5.QtGui import QPixmap
from PyQt5.QtCore import Qt
import cv2
import numpy as np
import torch

class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.model = torch.hub.load('ultralytics/yolov8', 'custom', path='runs/train/yolov8_custom/weights/best.pt')

    def initUI(self):
        self.setWindowTitle('轴承缺陷检测')
        self.setGeometry(100, 100, 800, 600)

        layout = QVBoxLayout()

        self.image_label = QLabel(self)
        layout.addWidget(self.image_label)

        self.load_button = QPushButton('加载图片', self)
        self.load_button.clicked.connect(self.load_image)
        layout.addWidget(self.load_button)

        self.setLayout(layout)

    def load_image(self):
        file_dialog = QFileDialog()
        file_dialog.setFileMode(QFileDialog.ExistingFile)
        if file_dialog.exec_():
            file_name = file_dialog.selectedFiles()[0]
            self.image_label.setPixmap(QPixmap(file_name))
            self.detect_defects(file_name)

    def detect_defects(self, file_name):
        img = cv2.imread(file_name)
        results = self.model(img)
        detections = results.pandas().xyxy[0]
        for _, row in detections.iterrows():
            x1, y1, x2, y2 = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax'])
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(img, row['name'], (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
        
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        height, width, channel = img.shape
        bytes_per_line = 3 * width
        q_img = QImage(img.data, width, height, bytes_per_line, QImage.Format_RGB888)
        self.image_label.setPixmap(QPixmap.fromImage(q_img))

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

6. 其他建议

• 数据增强：使用数据增强技术来提高模型的鲁棒性。
• 模型集成：集成多个模型以提高性能。
• 模型量化：量化模型以适应边缘设备。

7. 示例代码

7.1. 数据预处理示例

import cv2
import os

def resize_images(input_dir, output_dir, size=(640, 640)):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    for filename in os.listdir(input_dir):
        if filename.endswith(('.jpg', '.png', '.jpeg')):
            img_path = os.path.join(input_dir, filename)
            img = cv2.imread(img_path)
            img_resized = cv2.resize(img, size)
            output_path = os.path.join(output_dir, filename)
            cv2.imwrite(output_path, img_resized)

# 示例用法
resize_images('./dataset/images/train', './dataset/images/train_resized')

8. 运行和调试

确保你的环境配置正确，并且所有依赖项都已安装。运行PyQt5应用时，确保模型文件路径正确，并且数据集路径正确。