基于YOLOv10深度学习的铁轨缺陷检测系统：YOLOv10与UI实现详解

深度学习实战项目

于 2024-10-10 09:59:20 发布

阅读量652

点赞数 12

分类专栏： YOLOv10/v8/v7/v6/v5实战项目文章标签： YOLO 深度学习 ui 人工智能目标检测

本文链接：https://blog.csdn.net/m0_52343631/article/details/142813714

版权

YOLOv10/v8/v7/v6/v5实战项目专栏收录该内容

165 篇文章 0 订阅

订阅专栏

随着现代铁路交通的发展，铁轨的安全性显得尤为重要。铁轨缺陷检测不仅关系到列车的安全运行，还关系到整个铁路系统的效率。传统的铁轨检测方式往往依赖人工检查，效率低且容易出现漏检。近年来，深度学习技术的迅猛发展为铁轨缺陷检测提供了新的解决方案。本文将介绍一个基于YOLOv10的铁轨缺陷检测系统，包括系统设计、数据集准备、模型训练和UI界面实现，内容将涵盖详细的代码和相关文件。

1. 系统设计

1.1. 系统架构

我们的铁轨缺陷检测系统主要由以下几个模块组成：

数据采集模块：通过摄像头获取铁轨图像，并保存为数据集。
数据预处理模块：对获取的图像进行标注、缩放和增强，生成训练数据集。
模型训练模块：利用YOLOv10模型进行缺陷检测模型的训练。
检测模块：实时检测铁轨图像中的缺陷。
UI界面模块：为用户提供友好的操作界面，显示检测结果。

1.2. 技术选型

深度学习框架：使用PyTorch实现YOLOv10模型。
数据标注工具：使用LabelImg进行图像标注。
前端开发：使用Flask框架构建简单的Web UI。

2. 数据集准备

2.1. 数据采集

数据集是深度学习模型训练的关键。我们需要获取大量包含铁轨缺陷的图像。可以通过以下方式收集数据：

手动拍摄：使用高分辨率摄像头拍摄铁轨，确保不同角度和光照条件下的样本。
公开数据集：查找相关的公开数据集，结合自己的数据进行增强。

2.2. 数据标注

使用LabelImg工具对采集到的图像进行标注。将图像中的铁轨缺陷用矩形框框出，并保存为YOLO格式的标注文件。

标注类别：如裂缝、缺口、变形等。
保存路径：每张图像对应一个TXT文件，包含标注信息。

LabelImg标注步骤：

下载并安装LabelImg工具。
打开工具，选择要标注的图像文件夹。
选择相应类别，框选缺陷区域，保存标注。

2.3. 数据集划分

将标注后的数据集划分为训练集和验证集。一般采用80%用于训练，20%用于验证。

3. YOLOv10模型训练

YOLO（You Only Look Once）系列模型以其高效的检测性能广受欢迎。YOLOv10在前几个版本的基础上进行了优化，具有更好的检测精度和速度。

3.1. 环境搭建

确保安装以下依赖库：

pip install torch torchvision torchaudio
pip install opencv-python
pip install matplotlib
pip install PyYAML

3.2. 项目结构

创建项目目录结构如下：

railway-defect-detection/
│
├── dataset/
│   ├── images/
│   ├── labels/
│   └── data.yaml
│
├── yolov10/
│   ├── models/
│   ├── utils/
│   └── train.py
│
└── app.py

3.3. data.yaml 文件

data.yaml文件用于定义数据集的信息，内容如下：

train: dataset/images/train
val: dataset/images/val

nc: 3  # 类别数
names: ['crack', 'gap', 'deformation']  # 类别名称

3.4. 模型训练代码

在train.py中，使用以下代码进行训练：

import torch
from models import YOLOv10
from utils import load_dataset

# Load dataset
train_loader, val_loader = load_dataset('dataset/data.yaml')

# Initialize model
model = YOLOv10(num_classes=3)

# Define optimizer and loss function
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = torch.nn.CrossEntropyLoss()

# Training loop
for epoch in range(50):  # 训练50个epoch
    model.train()
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    
    print(f'Epoch [{epoch+1}/50], Loss: {loss.item():.4f}')

# Save model
torch.save(model.state_dict(), 'yolov10_weights.pth')

4. 检测模块

训练完成后，我们可以使用训练好的模型进行缺陷检测。以下是检测模块的代码示例：

import cv2
import torch

# Load model
model = YOLOv10(num_classes=3)
model.load_state_dict(torch.load('yolov10_weights.pth'))
model.eval()

def detect_defects(image_path):
    image = cv2.imread(image_path)
    results = model(image)  # 检测

    # 处理结果
    for result in results:
        bbox, label, confidence = result
        cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
        cv2.putText(image, f'{label}: {confidence:.2f}', (bbox[0], bbox[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    return image

# 测试
result_image = detect_defects('test_image.jpg')
cv2.imshow('Detected Defects', result_image)
cv2.waitKey(0)

5. UI界面模块

为了提供友好的用户体验，我们使用Flask构建一个简单的Web UI。

5.1. app.py 文件

from flask import Flask, request, render_template
import os

app = Flask(__name__)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return 'No file part'
    file = request.files['file']
    if file.filename == '':
        return 'No selected file'
    
    filepath = os.path.join('uploads', file.filename)
    file.save(filepath)
    result_image = detect_defects(filepath)
    cv2.imwrite('static/result.jpg', result_image)
    
    return render_template('result.html')

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

5.2. HTML模板

在templates/目录下创建index.html和result.html两个文件。

index.html:

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Railway Defect Detection</title>
</head>
<body>
    <h1>Upload Iron Track Image for Defect Detection</h1>
    <form action="/upload" method="post" enctype="multipart/form-data">
        <input type="file" name="file" accept="image/*" required>
        <button type="submit">Upload</button>
    </form>
</body>
</html>

result.html:

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Detection Result</title>
</head>
<body>
    <h1>Detection Result</h1>
    <img src="{{ url_for('static', filename='result.jpg') }}" alt="Result Image">
    <br>
    <a href="/">Back to Upload</a>
</body>
</html>