基于铁轨缺陷检测数据集构建使用YOLOv8进行目标检测系统 铁轨缺陷检测数据集

如何基于铁轨缺陷检测数据集构建使用YOLOv8进行目标检测系统 铁轨缺陷检测数据集 8类铁路轨道缺陷检测 训练铁轨缺陷

铁路铁轨轨道缺陷检测数据集
铁轨缺陷检测数据集
类别为’damage’,‘dirt’,‘unknown’,‘gap’,‘d_dent’,‘d_crush’,‘d_scratch’,‘d_slant’
原数据集399张 扩充三倍后一共1596张
txt或xml格式
基于铁轨缺陷检测数据集构建一个使用YOLOv8进行目标检测的系统。以下是详细的步骤和完整的代码示例。

---

声明：文章内所有代码仅供参考！

---

步骤概述

1. 安装依赖1. 准备数据集1. 配置YOLOv81. 训练模型1. 评估模型1. 构建GUI应用程序1. 运行应用程序

1. 安装依赖

首先确保你已经安装了必要的库，包括PyTorch、YOLOv8和其他相关库。

pip install torch torchvision opencv-python-headless matplotlib pillow PyQt5 ultralytics

2. 准备数据集

假设你的数据集已经转换为YOLO格式，并且目录结构如下：

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

每个图像文件对应一个标签文件，标签文件包含对象的类别ID和边界框坐标（x_center, y_center, width, height），归一化到[0, 1]范围。

3. 配置YOLOv8

创建一个data.yaml文件来定义数据集路径和类别信息。

data.yaml

# dataset/data.yaml
train: ./images/train
val: ./images/val
test: ./images/test

nc: 8
names: ['damage', 'dirt', 'unknown', 'gap', 'd_dent', 'd_crush', 'd_scratch', 'd_slant']

4. 训练模型

使用YOLOv8训练模型。假设你已经在项目目录下。

yolo task=detect mode=train model=yolov8n.pt data=../dataset/data.yaml epochs=50 imgsz=640 batch=16 name=railway_defect_detection

5. 评估模型

训练完成后，你可以使用以下命令评估模型性能。

yolo task=detect mode=val model=runs/detect/railway_defect_detection/weights/best.pt data=../dataset/data.yaml

6. 构建GUI应用程序

使用PyQt5构建一个简单的GUI应用程序，允许用户上传图像并查看检测结果。
🔥计算机视觉、图像处理、毕业辅导、作业帮助、代码获取，远程协助，代码定制，私聊会回复!
🚀B站项目实战：https://space.bilibili.com/364224477
😄 如果文章对你有帮助的话， 欢迎评论 💬点赞👍🏻 收藏
📂QQ+加：673276993
🤵‍♂代做需求：@个人主页

GUI 应用程序

[<title="Railway Defect Detection GUI using YOLOv8">]
import sys
import torch
import numpy as np
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QPushButton, QVBoxLayout, QWidget, QFileDialog, QMessageBox
from PyQt5.QtGui import QPixmap, QImage
from PIL import Image
import cv2

class RailwayDefectDetectionApp(QMainWindow):
    def __init__(self):
        super().__init__()

        self.setWindowTitle("Railway Defect Detection Using YOLOv8")
        self.setGeometry(100, 100, 800, 600)

        self.original_image_label = QLabel(self)
        self.detected_image_label = QLabel(self)
        self.upload_button = QPushButton("Upload Image", self)
        self.upload_button.clicked.connect(self.upload_image)
        self.process_button = QPushButton("Process Image", self)
        self.process_button.clicked.connect(self.process_image)
        self.process_button.setEnabled(False)

        layout = QVBoxLayout()
        layout.addWidget(self.original_image_label)
        layout.addWidget(self.detected_image_label)
        layout.addWidget(self.upload_button)
        layout.addWidget(self.process_button)

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

        # Load pre-trained model
        self.model = torch.hub.load('ultralytics/yolov8', 'custom', path='runs/detect/railway_defect_detection/weights/best.pt')

    def upload_image(self):
        options = QFileDialog.Options()
        file_name, _ = QFileDialog.getOpenFileName(self, "QFileDialog.getOpenFileName()", "", "Images (*.png *.xpm *.jpg);;All Files (*)", options=options)
        if file_name:
            self.image_path = file_name
            pixmap = QPixmap(file_name)
            self.original_image_label.setPixmap(pixmap.scaled(350, 350))
            self.process_button.setEnabled(True)

    def process_image(self):
        if hasattr(self, 'image_path'):
            image = Image.open(self.image_path).convert('RGB')
            results = self.model(image)

            detected_image = results.render()[0]
            detected_pixmap = QPixmap.fromImage(self.convert_to_qimage(detected_image))
            self.detected_image_label.setPixmap(detected_pixmap.scaled(350, 350))
        else:
            QMessageBox.warning(self, "Warning", "Please upload an image first.")

    def convert_to_qimage(self, pil_image):
        if pil_image.mode == "RGB":
            format = QImage.Format_RGB888
        elif pil_image.mode == "RGBA":
            format = QImage.Format_RGBA8888
        else:
            raise ValueError(f"Unsupported image mode: {pil_image.mode}")

        qimage = QImage(pil_image.tobytes(), pil_image.width, pil_image.height, format)
        return qimage

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

解释

1. 安装依赖：确保安装了所有必要的库。1. 准备数据集：按照YOLOv8的要求组织数据集。1. 配置YOLOv8：创建data.yaml文件定义数据集路径和类别信息。1. 训练模型：使用YOLOv8训练模型。1. 评估模型：评估训练好的模型性能。
2. 构建GUI应用程序：
3. 1. **运行应用程序**：保存上述脚本到相应的Python文件中，然后运行训练和评估脚本：

yolo task=detect mode=train model=yolov8n.pt data=../dataset/data.yaml epochs=50 imgsz=640 batch=16 name=railway_defect_detection
yolo task=detect mode=val model=runs/detect/railway_defect_detection/weights/best.pt data=../dataset/data.yaml
python railway_defect_gui.py