深度学习框架目标检测算法如何Yolov8训练机械零件数据集识别检测_轴承螺栓齿轮螺母数据集的检测建立零件识别检测gui

DeepLearningYolo

328人浏览 · 2026-04-29 11:21:03

DeepLearningYolo · 2026-04-29 11:21:03 发布

深度学习框架目标检测算法如何Yolov8训练机械零件数据集识别检测_轴承螺栓齿轮螺母数据集的检测建立零件识别检测gui
以下文字及代码仅供参考。

文章目录

机械零件目标检测数据集 4类别标注｜txt+xml格式｜
在这里插入图片描述

数据集描述：
①2249张机械零件图片，单类别标注
②txt+xml格式标注，兼容YOLO、Faster R-CNN、SSD等框架
③划分：训练集（1799张）、验证集（255张）、测试集（255张）
在这里插入图片描述

应用场景：

工业自动化：辅助机械零件分拣与装配
质量检测：识别零件缺陷，优化生产流程
库存管理：自动化零件盘点与分类

数据集规：
①标注格式：txt、xml
②类别标签：
0: Bearing（轴承）
1: Bolt（螺栓）
2: Gear（齿轮）
3: Nut（螺母）

机械零件目标检测数据集进行训练、评估，并最终部署一个可以用于工业自动化中的应用，我们可以按照以下步骤进行。使用 YOLOv8 作为示例模型，

tongxue同学：也可以根据需要选择其他框架如 Faster R-CNN 或 SSD。

一、安装依赖

1. 安装 Python 和相关库

确保已安装 Python（推荐版本 3.8+），并安装以下依赖：

pip install torch torchvision ultralytics opencv-python matplotlib lxml pyqt5

torch 和 torchvision：深度学习框架。
ultralytics：YOLOv8 的官方库。
opencv-python：用于图像处理。
matplotlib：用于可视化。
lxml：用于解析 XML 文件。
pyqt5：用于构建 GUI 应用程序。

二、准备数据集

1. 数据集目录结构

将机械零件目标检测数据集按照以下格式组织：

mechanical-parts-dataset/
├── images/
│   ├── train/          # 训练集图片
│   ├── val/            # 验证集图片
│   └── test/           # 测试集图片
├── labels/
│   ├── train/          # 训练集标注文件（txt+xml 格式）
│   ├── val/            # 验证集标注文件（txt+xml 格式）
│   └── test/           # 测试集标注文件（txt+xml 格式）
└── data.yaml           # 数据集配置文件

2. 创建 `data.yaml`

在 mechanical-parts-dataset/ 目录下创建 data.yaml 文件，内容如下：

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

nc: 4                   # 类别数量
names: ['Bearing', 'Bolt', 'Gear', 'Nut']       # 类别名称

3. 转换标注格式

同学：
需要确保所有的标注都转换为 YOLO 所需的 .txt 格式。

对于 .xml 文件，你可以编写一个简单的脚本来将其转换为 .txt 格式：

import os
import xml.etree.ElementTree as ET

def convert_xml_to_yolo(xml_file, label_dir):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    image_width = int(root.find('size/width').text)
    image_height = int(root.find('size/height').text)

    label_file = os.path.join(label_dir, os.path.splitext(os.path.basename(xml_file))[0] + '.txt')
    with open(label_file, 'w') as f:
        for obj in root.findall('object'):
            class_name = obj.find('name').text
            if class_name == 'Bearing':
                class_id = 0
            elif class_name == 'Bolt':
                class_id = 1
            elif class_name == 'Gear':
                class_id = 2
            elif class_name == 'Nut':
                class_id = 3
            else:
                continue

            bbox = obj.find('bndbox')
            xmin = int(bbox.find('xmin').text)
            ymin = int(bbox.find('ymin').text)
            xmax = int(bbox.find('xmax').text)
            ymax = int(bbox.find('ymax').text)

            x_center = (xmin + xmax) / 2.0 / image_width
            y_center = (ymin + ymax) / 2.0 / image_height
            width = (xmax - xmin) / float(image_width)
            height = (ymax - ymin) / float(image_height)

            f.write(f"{class_id} {x_center} {y_center} {width} {height}\n")

# 示例调用
xml_files = ['path/to/xml/file1.xml', 'path/to/xml/file2.xml']
label_dir = 'path/to/labels/'
for xml_file in xml_files:
    convert_xml_to_yolo(xml_file, label_dir)

请确保所有 .xml 文件都被转换为 .txt 文件，并放置在正确的目录中。

三、训练模型

1. 开始训练

运行以下代码开始训练模型：

from ultralytics import YOLO

# 加载预训练模型
model = YOLO('yolov8n.pt')  # 使用 YOLOv8 Nano 模型

# 训练模型
results = model.train(
    data='path/to/mechanical-parts-dataset/data.yaml',  # 数据集配置文件路径
    epochs=50,                                # 训练轮数
    imgsz=640,                                # 输入图片尺寸
    batch=16,                                 # 批量大小
    device='cuda' if torch.cuda.is_available() else 'cpu'  # 使用 GPU 或 CPU
)

# 保存训练结果
model.export(format='onnx')  # 导出为 ONNX 格式（可选）

四、评估模型

1. 测试模型性能

运行以下代码进行模型评估：

from ultralytics import YOLO

# 加载训练好的模型
model = YOLO('runs/detect/train/weights/best.pt')

# 测试模型
metrics = model.val(data='path/to/mechanical-parts-dataset/data.yaml')
print(f"mAP50: {metrics.box.map50}")  # 输出 mAP50 指标

五、构建 GUI 应用程序

1. 创建 PyQt5 GUI

以下是一个简单的 GUI 示例，用于加载模型并检测图片中的机械零件目标。

import sys
import cv2
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QPushButton, QVBoxLayout, QWidget, QFileDialog
from PyQt5.QtGui import QPixmap, QImage
from ultralytics import YOLO

class MechanicalPartsDetectorApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.model = YOLO('runs/detect/train/weights/best.pt')  # 加载训练好的模型

    def initUI(self):
        self.setWindowTitle("机械零件目标检测")
        self.setGeometry(100, 100, 800, 600)

        self.label = QLabel(self)
        self.label.setText("选择图片以检测机械零件目标")
        self.label.setAlignment(Qt.AlignCenter)

        self.button = QPushButton("选择图片", self)
        self.button.clicked.connect(self.load_image)

        layout = QVBoxLayout()
        layout.addWidget(self.label)
        layout.addWidget(self.button)

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

    def load_image(self):
        options = QFileDialog.Options()
        file_path, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "Images (*.png *.jpg *.jpeg)", options=options)
        if file_path:
            image = cv2.imread(file_path)
            results = self.model(image)  # 进行目标检测
            annotated_image = results[0].plot()  # 绘制检测结果

            # 将 OpenCV 图像转换为 PyQt 图像
            height, width, channel = annotated_image.shape
            bytes_per_line = 3 * width
            q_img = QImage(annotated_image.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
            self.label.setPixmap(QPixmap.fromImage(q_img))

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