使用YOLOv8创建一个铁路轨道缺陷检测系统从训练铁轨轨道缺陷检测数据集，数据集准备、模型训练、评估、结果可视化以及简易用户界面开发

最新推荐文章于 2025-02-19 11:33:20 发布

目标检测数据集合

最新推荐文章于 2025-02-19 11:33:20 发布

阅读量855

点赞数 22

分类专栏：轨道铁轨类文章标签： YOLO ui

本文链接：https://blog.csdn.net/2401_86822270/article/details/145361968

版权

轨道铁轨类专栏收录该内容

4 篇文章

订阅专栏

使用YOLOv8创建一个铁路轨道缺陷检测系统从训练铁轨轨道缺陷检测数据集，数据集准备、模型训练、评估、结果可视化以及简易用户界面开发

文章目录

铁轨轨道缺陷检测数据集
在这里插入图片描述

4分类， Spalling --脱落
Wheel Burn --轮烧
Squat --压陷
Corrugation --磨耗图片共2277张。数据集按8：1：1分配。其中训练集1821张图片，验证集228张图片。测试集228张图片在这里插入图片描述

为了使用YOLOv8创建一个铁路轨道缺陷检测系统，我们将从数据集准备、模型训练、评估、结果可视化以及简易用户界面开发几个方面进行详细说明。
以下文字及代码仅供参考。

1. 环境设置

首先确保你的环境中安装了必要的依赖项。这里以YOLOv5为例，如果你使用的是YOLOv8，请根据实际情况调整。

git clone https://github.com/ultralytics/yolov8  # 克隆YOLOv8仓库（请确认这是正确的URL）
cd yolov8
pip install -r requirements.txt  # 安装依赖

在这里插入图片描述

2. 数据预处理

组织你的数据集如下：

railway_defects_dataset/
├── images/
│   ├── train/
│   ├── val/
│   └── test/
└── labels/
    ├── train/
    ├── val/
    └── test/

创建railway_data.yaml文件来描述数据集路径和类别信息：

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

nc: 4  # 类别数量
names: ['Spalling', 'Wheel Burn', 'Squat', 'Corrugation']  # 类别名称

在这里插入图片描述

3. 模型定义与训练

开始训练过程：

python train.py --img 640 --batch 16 --epochs 50 --data railway_data.yaml --weights yolov8s.pt --name railway_defect_model

--img 640: 输入图像尺寸。
--batch 16: 每批次的图像数量。
--epochs 50: 训练周期数。
--data railway_data.yaml: 数据配置文件路径。
--weights yolov8s.pt: 预训练权重文件路径（如果有）。
--name railway_defect_model: 保存模型的目录名。

4. 评估

训练完成后，可以对验证集或测试集进行评估：

python val.py --weights runs/train/railway_defect_model/weights/best.pt --data railway_data.yaml --img 640 --task val

5. 结果分析与可视化

对于预测结果的可视化，可以使用如下命令：

python detect.py --weights runs/train/railway_defect_model/weights/best.pt --source ./railway_defects_dataset/images/test/ --save-txt --save-conf

6. 用户界面开发（Flask应用）

构建一个简单的Flask应用用于展示检测结果。

Flask 应用代码

from flask import Flask, request, render_template, send_from_directory
import torch
from PIL import Image
import os

app = Flask(__name__)
model = torch.hub.load('ultralytics/yolov8', 'custom', path='runs/train/railway_defect_model/weights/best.pt')

@app.route('/', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        if 'file' not in request.files:
            return "No file part"
        file = request.files['file']
        if file.filename == '':
            return "No selected file"
        if file:
            img = Image.open(file.stream)
            results = model(img)
            results.save(save_dir='static/detect_results')  # 保存检测结果图片到指定目录
            result_image_path = os.path.join('detect_results', os.listdir('static/detect_results')[0])
            return render_template('result.html', result_image=result_image_path)
    return render_template('index.html')

@app.route('/static/<path:filename>')
def serve_static(filename):
    return send_from_directory('static', filename)

if __name__ == '__main__':
    app.run(debug=True)

`index.html` 示例

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>Upload Image</title>
</head>
<body>
    <form action="/" method="post" enctype="multipart/form-data">
        Select image to upload:
        <input type="file" name="file" id="file">
        <input type="submit" value="Upload Image" name="submit">
    </form>
</body>
</html>

`result.html` 示例

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>Result</title>
</head>
<body>
    <h1>Detection Result:</h1>
    <img src="{{ url_for('serve_static', filename=result_image) }}" alt="Detection Result">
</body>
</html>

这个简易的应用允许用户上传一张图片，然后在原图上绘制检测结果，并显示给用户。同学呀，你需要创建一个static目录，并在其中创建一个子目录detect_results用于存放检测结果图像。此外，还需要根据实际情况调整脚本中的路径和其他参数。如果使用的是YOLOv8而非v5，请确保所有相关的调用和路径都针对YOLOv8进行了正确配置。