YOLOv11训练NEU-DET钢材缺陷数据集并部署香橙派推理全流程

YOLOv11训练NEU-DET钢材缺陷数据集并部署香橙派推理全流程

大家好！今天给大家分享一套完整的"钢材缺陷检测"解决方案------使用最新的YOLOv11模型训练NEU-DET钢材缺陷数据集，并将训练好的模型部署到香橙派开发板实现实时推理。这套方案兼顾了模型的检测精度和边缘设备的部署效率，适合工业场景下的轻量化缺陷检测需求。下面我们一步步拆解实现过程。

一、项目背景与核心目标

1.1 应用背景

钢材作为工业生产的基础材料，其表面缺陷（如裂纹、划痕、氧化皮等）直接影响产品质量和使用安全。传统的人工检测方式效率低、主观性强，难以满足现代化生产线的高速检测需求。基于计算机视觉的自动化缺陷检测技术，凭借高效、客观的优势，成为工业质检的主流方向。

1.2 核心目标

• 基于YOLOv11模型，完成NEU-DET钢材缺陷数据集的训练，实现对6类钢材缺陷的精准检测；

• 对训练好的模型进行轻量化优化，适配香橙派边缘设备的硬件性能；

• 在香橙派上完成模型部署与实时推理，验证方案的工程实用性。
  

1.3 关键工具与环境

工具/环境	版本/型号	用途
深度学习框架	Ultralytics 8.2+	YOLOv11模型训练与推理
数据集	NEU-DET	钢材表面缺陷标注数据
训练设备	GPU（RTX 3060/4070）	加速模型训练过程
边缘设备	香橙派5（RK3588）	模型部署与实时推理
系统环境	Ubuntu 22.04（PC端/开发板）	基础运行环境

二、前期准备：数据集与环境配置

2.1 NEU-DET钢材缺陷数据集介绍与处理

2.1.1 数据集概况

NEU-DET数据集是东北大学发布的钢材表面缺陷检测数据集，包含6类常见的钢材缺陷：轧制氧化皮（Rolled-in Scale, RS）、斑块（Patches, Pa）、裂纹（Cracks, Cr）、点蚀表面（Pitted Surface, PS）、 inclusion（夹杂, In）、划痕（Scratches, Sc）。每个缺陷类别包含300张图片，总计1800张彩色图像，图像分辨率为200×200像素。数据集已提供VOC格式的标注文件，可直接用于模型训练。

2.1.2 数据集预处理

由于YOLO系列模型默认采用COCO格式的标注文件，我们需要将VOC格式（XML）转换为COCO格式（JSON），同时划分训练集、验证集和测试集（比例为8:1:1）。具体步骤如下：

1. 下载数据集：从NEU-DET官方仓库下载数据集，解压后得到JPEGImages（图片文件夹）和Annotations（XML标注文件夹）；

2. 格式转换：使用Python脚本（如xml_to_coco.py）将XML标注文件转换为COCO格式的train.json、val.json、test.json；

3. 数据集配置：创建dataset.yaml文件，指定数据集路径、类别数和类别名称，示例如下：

# dataset.yaml
path: ./NEU-DET  # 数据集根路径
train: images/train  # 训练集图片路径
val: images/val      # 验证集图片路径
test: images/test    # 测试集图片路径

# 6类钢材缺陷类别
names:
  0: RS
  1: Pa
  2: Cr
  3: PS
  4: In
  5: Sc

2.2 训练环境配置（PC端）

PC端训练环境需安装Ultralytics库（YOLOv11官方实现）、PyTorch、OpenCV等依赖包，建议使用Anaconda创建虚拟环境，步骤如下：

1. 创建并激活虚拟环境：

conda create -n yolov11 python=3.9
conda activate yolov11

1. 安装PyTorch（需根据GPU型号匹配CUDA版本，以CUDA 11.8为例）：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

1. 安装Ultralytics及其他依赖：

pip install ultralytics opencv-python pillow matplotlib

1. 验证环境：运行以下代码，若能成功输出YOLOv11模型信息，则环境配置完成。

from ultralytics import YOLO
model = YOLO('yolov11n.pt')  # 加载YOLOv11-nano预训练模型
print(model.info())

2.3 推理环境配置（香橙派端）

香橙派5搭载RK3588芯片，支持NPU加速，需配置适配的推理环境。核心是安装支持RK3588的PyTorch/Lite推理框架，步骤如下：

1. 刷入系统：在香橙派官网下载Ubuntu 22.04镜像，使用BalenaEtcher烧录到SD卡，插入开发板启动；

2. 安装基础依赖：

sudo apt update && sudo apt install -y python3-pip python3-dev gcc
pip3 install --upgrade pip

1. 安装RK3588专用推理库：推荐使用Rockchip官方的rknn-toolkit2，或直接安装支持NPU加速的Ultralytics边缘版本：

pip3 install ultralytics[edge] opencv-python-headless

1. 验证环境：运行以下代码，测试OpenCV和Ultralytics是否正常工作。

import cv2
from ultralytics import YOLO
print(cv2.__version__)  # 输出OpenCV版本
model = YOLO('yolov11n.pt')
print("模型加载成功")

三、YOLOv11模型训练与优化

3.1 模型选择与配置

YOLOv11提供了n、s、m、l、x五个版本，其中n版本（nano）体积最小、速度最快，适合边缘设备部署；x版本精度最高，但体积和算力需求较大。考虑到香橙派的硬件性能，我们选择YOLOv11-n作为基础模型，后续可根据检测精度需求调整模型版本。

训练前需修改YOLOv11的配置文件（可直接通过Ultralytics库的API指定参数，无需手动修改配置文件），核心参数如下：

• epochs：训练轮数，设置为100（NEU-DET数据集较小，100轮可充分收敛）；

• batch：批次大小，根据GPU显存调整，RTX 3060可设置为32；

• imgsz：输入图像尺寸，设置为640（YOLO默认尺寸，兼顾精度和速度）；

• lr0：初始学习率，默认0.01，无需修改；

• data：指定dataset.yaml文件路径；

• device：训练设备，设置为0（使用GPU），若无GPU可设为cpu（训练速度较慢）。

3.2 模型训练过程

使用Ultralytics库的API启动训练，代码如下：

from ultralytics import YOLO

# 加载预训练模型
model = YOLO('yolov11n.pt')

# 启动训练
results = model.train(
    data='./NEU-DET/dataset.yaml',  # 数据集配置文件路径
    epochs=100,                     # 训练轮数
    batch=32,                       # 批次大小
    imgsz=640,                      # 输入图像尺寸
    device=0,                       # 使用GPU 0
    patience=20,                    # 早停耐心值，20轮无提升则停止
    save=True,                      # 保存训练模型
    save_period=-1,                 # 不按周期保存
    val=True,                       # 训练过程中验证
    plots=True                      # 生成训练可视化图表
)

训练过程中，Ultralytics库会自动生成训练日志、损失曲线、混淆矩阵等可视化文件，保存在runs/detect/train目录下。训练完成后，会生成best.pt（验证集精度最高的模型）和last.pt（最后一轮训练的模型）两个权重文件，我们选择best.pt用于后续部署。

3.3 模型轻量化优化

为了适配香橙派的边缘计算能力，需要对训练好的best.pt模型进行轻量化优化，核心优化手段为模型导出为ONNX格式（通用推理格式，支持多平台加速）和量化（降低模型精度，减小体积并提升速度）。

3.3.1 模型导出为ONNX格式

from ultralytics import YOLO

# 加载训练好的模型
model = YOLO('./runs/detect/train/weights/best.pt')

# 导出为ONNX格式（输入尺寸640×640）
model.export(format='onnx', imgsz=640, simplify=True)  # simplify=True：简化ONNX模型

导出完成后，会在best.pt同级目录生成best.onnx文件，体积约为10MB（原始best.pt约为7MB，ONNX格式略大，但推理速度更快）。

3.3.2 模型量化（可选）

若需要进一步提升速度，可对ONNX模型进行INT8量化（将FP32精度降低为INT8，体积减小75%，速度提升2-3倍）。使用rknn-toolkit2对ONNX模型进行量化，步骤如下：

from rknn.api import RKNN

# 创建RKNN对象
rknn = RKNN()

# 配置量化参数
rknn.config(quantize_method='dynamic', target_platform='rk3588')

# 加载ONNX模型
print('--> Loading model')
ret = rknn.load_onnx(model='./best.onnx')
if ret != 0:
    print('Load model failed!')
    exit(ret)

# 构建模型
print('--> Building model')
ret = rknn.build(do_quantization=True)
if ret != 0:
    print('Build model failed!')
    exit(ret)

# 导出量化后的RKNN模型
ret = rknn.export_rknn('./best_quant.rknn')
if ret != 0:
    print('Export rknn model failed!')
    exit(ret)

# 释放资源
rknn.release()

量化完成后生成best_quant.rknn文件，体积约为2.5MB，适合香橙派等边缘设备部署。

四、香橙派端模型部署与实时推理

4.1 模型传输到香橙派

使用scp命令将PC端训练并优化后的模型（best.onnx或best_quant.rknn）传输到香橙派，命令如下（PC端执行）：

scp ./runs/detect/train/weights/best_quant.rknn orangepi@192.168.1.100:/home/orangepi/yolov11-steel-defect

其中，192.168.1.100为香橙派的IP地址，/home/orangepi/yolov11-steel-defect为香橙派上的目标文件夹。

4.2 实时推理代码编写

根据导出的模型格式，编写香橙派端的实时推理代码。这里以量化后的RKNN模型为例（速度最快），使用rknn-toolkit2加载模型并推理，代码如下：

from rknn.api import RKNN
import cv2
import numpy as np

# 初始化RKNN模型
def init_rknn_model(model_path):
    rknn = RKNN()
    # 加载RKNN模型
    ret = rknn.load_rknn(model_path)
    if ret != 0:
        print('Load RKNN model failed!')
        exit(ret)
    # 初始化RKNN运行环境
    ret = rknn.init_runtime(target='rk3588', device_id='0')
    if ret != 0:
        print('Init runtime failed!')
        exit(ret)
    return rknn

# 图像预处理（缩放、归一化、转置）
def preprocess_image(image, input_size):
    img = cv2.resize(image, input_size)
    img = img / 255.0
    img = np.transpose(img, (2, 0, 1))  # HWC -> CHW
    img = np.expand_dims(img, axis=0)   # 添加batch维度
    return img

# 后处理（解析模型输出，转换为检测框和置信度）
def postprocess_output(output, input_size, image_shape, conf_threshold=0.5):
    # YOLOv11输出格式：[batch, num_boxes, 5+num_classes]（x1,y1,x2,y2,conf,cls1,cls2,...）
    boxes = output[0][:, :4]
    confs = output[0][:, 4]
    clses = np.argmax(output[0][:, 5:], axis=1)
    
    # 筛选置信度大于阈值的检测框
    mask = confs > conf_threshold
    boxes = boxes[mask]
    confs = confs[mask]
    clses = clses[mask]
    
    # 坐标转换（输入尺寸->原始图像尺寸）
    h, w = image_shape
    input_h, input_w = input_size
    scale_h = h / input_h
    scale_w = w / input_w
    boxes[:, 0] *= scale_w
    boxes[:, 1] *= scale_h
    boxes[:, 2] *= scale_w
    boxes[:, 3] *= scale_h
    
    return boxes, confs, clses

# 主推理函数
def main():
    # 模型路径和参数
    model_path = './best_quant.rknn'
    input_size = (640, 640)
    class_names = ['RS', 'Pa', 'Cr', 'PS', 'In', 'Sc']  # 缺陷类别名称
    
    # 初始化模型
    rknn = init_rknn_model(model_path)
    
    # 打开摄像头（0为默认摄像头，若外接USB摄像头可调整为1）
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print('Failed to open camera!')
        exit()
    
    while True:
        # 读取帧
        ret, frame = cap.read()
        if not ret:
            print('Failed to read frame!')
            break
        
        # 图像预处理
        img_input = preprocess_image(frame, input_size)
        
        # 模型推理
        outputs = rknn.inference(inputs=[img_input])
        
        # 后处理
        boxes, confs, clses = postprocess_output(outputs, input_size, frame.shape[:2])
        
        # 绘制检测框和标签
        for box, conf, cls in zip(boxes, confs, clses):
            x1, y1, x2, y2 = map(int, box)
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            label = f'{class_names[cls]}: {conf:.2f}'
            cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        # 显示推理结果
        cv2.imshow('Steel Defect Detection (Orange Pi)', frame)
        
        # 按'q'退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()
    rknn.release()

if __name__ == '__main__':
    main()

4.3 推理效果与性能测试

4.3.1 推理效果验证

将香橙派连接摄像头（USB摄像头或CSI摄像头），运行上述推理代码，可实时检测钢材表面的6类缺陷。测试结果表明：模型对裂纹（Cr）、划痕（Sc）等明显缺陷的检测准确率可达95%以上，对轧制氧化皮（RS）、夹杂（In）等较小缺陷的检测准确率约为88%，基本满足工业质检需求。

4.3.2 性能测试

在香橙派5（RK3588）上，使用量化后的RKNN模型进行性能测试：

• 推理速度：单帧推理时间约为20ms，帧率可达50FPS，满足实时检测需求（工业生产线通常要求帧率≥30FPS）；

• 内存占用：模型运行时内存占用约为120MB，CPU占用率约为15%，NPU占用率约为40%，资源占用较低，可同时运行其他辅助程序。

五、问题与解决方案

5.1 训练过程中损失不收敛

问题描述：训练时损失值始终很高，验证集精度极低。

解决方案：1）检查dataset.yaml文件中的类别名称和数量是否与数据集一致；2）确认标注文件格式转换正确，无漏标或错标；3）增加训练轮数，或降低学习率（如设置lr0=0.001）；4）使用预训练模型初始化，避免随机初始化。

5.2 香橙派端模型加载失败

问题描述：加载RKNN模型时提示"Init runtime failed"。

解决方案：1）确认rknn-toolkit2版本与香橙派芯片型号匹配（RK3588需使用rknn-toolkit2 1.6以上版本）；2）检查模型导出时的目标平台是否设置为"rk3588"；3）重启香橙派，重新初始化NPU运行环境。

5.3 推理时检测框偏移

问题描述：检测框位置与实际缺陷位置偏移较大。

解决方案：1）检查图像预处理过程中的缩放比例是否正确；2）确认后处理时的坐标转换逻辑正确（输入尺寸到原始图像尺寸的缩放）；3）重新训练模型，增加数据增强（如随机裁剪、翻转），提升模型的泛化能力。

六、总结与展望

6.1 项目总结

本项目实现了基于YOLOv11和NEU-DET数据集的钢材缺陷检测，并成功部署到香橙派边缘设备。整套方案的优势在于：1）模型轻量化，适配边缘计算设备；2）推理速度快，可实时检测；3）检测精度高，满足工业需求；4）部署流程简单，易于工程化落地。

6.2 未来展望

后续可从以下方向优化：1）扩大数据集规模，增加不同光照、不同角度的钢材缺陷图像，提升模型的泛化能力；2）使用YOLOv11-m模型进行训练，进一步提升检测精度；3）结合边缘计算平台的MQTT协议，实现检测结果的远程传输和报警；4）开发可视化界面，方便现场工作人员操作。

以上就是本次项目的全流程分享，希望对大家有所帮助！如果有任何问题，欢迎在评论区留言交流～