基于深度学习的布匹缺陷检测系统:从数据准备到模型部署

在现代纺织工业中,布匹的质量控制至关重要。随着深度学习技术的发展,我们可以利用先进的YOLO模型来自动检测布匹中的各种缺陷。本篇博客将简单带你实现一个基于深度学习的布匹缺陷检测系统,涵盖数据准备、模型训练、用户界面开发以及最终部署的完整流程。



1. 项目背景与意义

布匹在生产过程中容易产生各种缺陷,如破洞、污渍、褶皱等。传统的人工检测方法不仅效率低下,而且准确性有限。基于YOLO的深度学习模型可以实现高效、准确的布匹缺陷检测,从而大幅提高生产效率,降低人工成本。

2. 技术方案与开发环境
技术方案
  • 深度学习框架:YOLOv8/v7/v6/v5
  • 编程语言:Python
  • 开发工具:PyCharm、VSCode
  • Web框架:Flask
  • 前端框架:HTML、CSS
开发环境

首先,配置开发环境并安装所需的依赖库:

conda create -n fabric_defect_detection python=3.8
conda activate fabric_defect_detection
pip install torch torchvision torchaudio
pip install flask opencv-python pandas
pip install -U git+https://github.com/ultralytics/yolov5
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
3. 数据准备
数据集获取

你可以使用公开的布匹缺陷数据集,也可以自行收集和标注数据。推荐使用LabelImg工具对图像进行标注,生成YOLO格式的标签文件。

数据标注

安装并运行LabelImg进行数据标注:

pip install labelImg
labelImg
  • 1.
  • 2.
数据集划分

将数据集划分为训练集、验证集和测试集,确保模型能够在不同数据上进行有效的训练和评估。

import os
import shutil
import random

def split_dataset(source_dir, train_dir, val_dir, test_dir, train_ratio=0.7, val_ratio=0.2):
    all_files = os.listdir(source_dir)
    random.shuffle(all_files)
    train_count = int(len(all_files) * train_ratio)
    val_count = int(len(all_files) * val_ratio)

    for i, file in enumerate(all_files):
        if i < train_count:
            shutil.move(os.path.join(source_dir, file), train_dir)
        elif i < train_count + val_count:
            shutil.move(os.path.join(source_dir, file), val_dir)
        else:
            shutil.move(os.path.join(source_dir, file), test_dir)

split_dataset('data/images', 'data/train/images', 'data/val/images', 'data/test/images')
split_dataset('data/labels', 'data/train/labels', 'data/val/labels', 'data/test/labels')
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.
  • 13.
  • 14.
  • 15.
  • 16.
  • 17.
  • 18.
  • 19.
  • 20.
4. 模型训练
配置YOLO模型

下载YOLOv5预训练权重,并配置数据文件:

# fabric_defect.yaml
train: data/train
val: data/val
nc: 5  # number of classes (e.g., hole, stain, wrinkle, etc.)
names: ['hole', 'stain', 'wrinkle', 'tear', 'missing_thread']  # list of class names
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
模型训练

运行以下命令开始训练:

python train.py --img 640 --batch 16 --epochs 50 --data fabric_defect.yaml --cfg yolov5s.yaml --weights yolov5s.pt
  • 1.
模型评估

使用验证集评估模型性能,并进行必要的超参数调优:

from sklearn.metrics import precision_score, recall_score, f1_score

y_true = [...]  # true labels
y_pred = [...]  # predicted labels

precision = precision_score(y_true, y_pred, average='macro')
recall = recall_score(y_true, y_pred, average='macro')
f1 = f1_score(y_true, y_pred, average='macro')

print(f"Precision: {precision}, Recall: {recall}, F1 Score: {f1}")
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
5. 用户界面开发
Flask搭建Web应用
  1. 创建项目目录结构:
fabric_defect_detection/
├── app.py
├── templates/
│   ├── index.html
│   └── result.html
├── static/
│   └── uploads/
└── models/
    └── yolov5s.pt
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  1. 编写网页模板:
  • index.html
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Fabric Defect Detection</title>
    <link rel="stylesheet" href="{{ url_for('static', filename='styles.css') }}">
</head>
<body>
    <h1>Fabric Defect Detection</h1>
    <form action="/predict" method="post" enctype="multipart/form-data">
        <input type="file" name="file">
        <button type="submit">Upload</button>
    </form>
</body>
</html>
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.
  • 13.
  • 14.
  • 15.
  • 16.
  • 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>
    <link rel="stylesheet" href="{{ url_for('static', filename='styles.css') }}">
</head>
<body>
    <h1>Detection Result</h1>
    <img src="{{ url_for('static', filename='uploads/' + filename) }}" alt="Uploaded Image">
    <p>{{ result }}</p>
</body>
</html>
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.
  • 13.
  • 14.
实现后端逻辑
  • app.py
from flask import Flask, request, render_template, url_for
import os
from werkzeug.utils import secure_filename
import torch
from PIL import Image

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/uploads/'

model = torch.hub.load('ultralytics/yolov5', 'custom', path='models/yolov5s.pt')

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

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return 'No file part'
    file = request.files['file']
    if file.filename == '':
        return 'No selected file'
    if file:
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        img = Image.open(filepath)
        results = model(img)
        results.save(save_dir=app.config['UPLOAD_FOLDER'])
        return render_template('result.html', filename=filename, result=results.pandas().xyxy[0].to_json(orient="records"))

if __name__ == '__main__':
    app.run(debug=True)
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.
  • 13.
  • 14.
  • 15.
  • 16.
  • 17.
  • 18.
  • 19.
  • 20.
  • 21.
  • 22.
  • 23.
  • 24.
  • 25.
  • 26.
  • 27.
  • 28.
  • 29.
  • 30.
  • 31.
  • 32.
  • 33.
6. 模型部署
部署到云服务器
  1. 使用Gunicorn部署
pip install gunicorn
gunicorn -w 4 app:app
  • 1.
  • 2.
  1. 配置Nginx反向代理
server {
    listen 80;
    server_name your_domain;

    location / {
        proxy_pass http://127.0.0.1:8000;
        proxy_set_header Host $host;
        proxy_set_header X-Real-IP $remote_addr;
        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
        proxy_set_header X-Forwarded-Proto $scheme;
    }
}
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.