铁路铁轨轨道缺陷检测数据集

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

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

步骤概述

1. 安装依赖
2. 准备数据集
3. 配置YOLOv8
4. 训练模型
5. 评估模型
6. 构建GUI应用程序
7. 运行应用程序

1. 安装依赖

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

pipinstalltorch 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.yamltrain:./images/trainval:./images/valtest:./images/testnc:8names:['damage','dirt','unknown','gap','d_dent','d_crush','d_scratch','d_slant']

4. 训练模型

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

yolotask=detectmode=trainmodel=yolov8n.ptdata=../dataset/data.yamlepochs=50imgsz=640batch=16name=railway_defect_detection

5. 评估模型

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

yolotask=detectmode=valmodel=runs/detect/railway_defect_detection/weights/best.ptdata=../dataset/data.yaml

6. 构建GUI应用程序

使用PyQt5构建一个简单的GUI应用程序，允许用户上传图像并查看检测结果。

GUI 应用程序

[<title="Railway Defect Detection GUI using YOLOv8">]importsysimporttorchimportnumpyasnpfromPyQt5.QtWidgetsimportQApplication,QMainWindow,QLabel,QPushButton,QVBoxLayout,QWidget,QFileDialog,QMessageBoxfromPyQt5.QtGuiimportQPixmap,QImagefromPILimportImageimportcv2classRailwayDefectDetectionApp(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 modelself.model=torch.hub.load('ultralytics/yolov8','custom',path='runs/detect/railway_defect_detection/weights/best.pt')defupload_image(self):options=QFileDialog.Options()file_name,_=QFileDialog.getOpenFileName(self,"QFileDialog.getOpenFileName()","","Images (*.png *.xpm *.jpg);;All Files (*)",options=options)iffile_name:self.image_path=file_name pixmap=QPixmap(file_name)self.original_image_label.setPixmap(pixmap.scaled(350,350))self.process_button.setEnabled(True)defprocess_image(self):ifhasattr(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.")defconvert_to_qimage(self,pil_image):ifpil_image.mode=="RGB":format=QImage.Format_RGB888elifpil_image.mode=="RGBA":format=QImage.Format_RGBA8888else:raiseValueError(f"Unsupported image mode:{pil_image.mode}")qimage=QImage(pil_image.tobytes(),pil_image.width,pil_image.height,format)returnqimageif__name__=="__main__":app=QApplication(sys.argv)window=RailwayDefectDetectionApp()window.show()sys.exit(app.exec_())

解释

1. 安装依赖：确保安装了所有必要的库。
2. 准备数据集：按照YOLOv8的要求组织数据集。
3. 配置YOLOv8：创建data.yaml文件定义数据集路径和类别信息。
4. 训练模型：使用YOLOv8训练模型。
5. 评估模型：评估训练好的模型性能。
6. 构建GUI应用程序：
   • 创建主窗口和布局。
   • 添加按钮用于上传和处理图像。
   • 实现图像上传功能。
   • 实现图像处理功能，调用模型进行检测并显示结果。
7. 运行应用程序：保存上述脚本到相应的Python文件中，然后运行训练和评估脚本：

yolotask=detectmode=trainmodel=yolov8n.ptdata=../dataset/data.yamlepochs=50imgsz=640batch=16name=railway_defect_detection yolotask=detectmode=valmodel=runs/detect/railway_defect_detection/weights/best.ptdata=../dataset/data.yaml python railway_defect_gui.py