简介基于YOLO的轴承生产缺陷检测数据集面向从事工业视觉检测、目标检测算法学习与实战的开发者。数据集围绕轴承生产中的典型缺陷识别问题提供568张已标注图片按三类缺陷完成边界框标注可直接用于YOLO系列模型的训练、验证与测试。包体共1772个文件涵盖624张jpg图像、577个txt格式的YOLO标注、568个xml格式的VOC标注以及2个Python训练脚本和1个yaml模型配置文件可同时兼容主流检测框架的数据格式。压缩包大小755.12MB整体目录结构清晰适合作为工业质检场景的入门练手或课题参考目前已有337人浏览学习。借助这套数据与脚本使用者能够快速搭建缺陷检测流程从数据预处理、模型训练到结果可视化还可针对光照变化、拍摄角度、背景干扰等生产环境条件开展鲁棒性调优为实际质量管控提供可落地的技术路径。1. 基于YOLO的轴承生产缺陷检测只有568张图和三类标签能撑起一条产线吗轴承表面缺陷检测的残酷现实是产线上真正能收集到的缺陷样本常常就是几百张而不是上万张。568张图、三类缺陷这个规模在工业视觉项目里非常典型也最容易翻车——模型训练时mAP看着能到0.95一上产线就被良品表面纹理带偏或者干脆把小划伤漏掉。问题不在YOLO本身而在数据怎么划分、类别边界怎么定、训练参数怎么守纪律。这篇文章要解决的就是这件事用568张图把YOLO训练成能用的轴承外观检测器。从数据整理、标注格式、模型选型、训练参数到TensorRT加速和产线验证全部按可复现的步骤写。适合手里有少量缺陷图、想快速验证方案可行性或者已经训过但现场表现不稳定的工程师。2. 把568张图整理成YOLO能吃的数据集标注、划分与增强2.1 三类缺陷怎么定类别边界是第一个坑标题里写了“类别三类”但具体哪三类要自己定。常见做法是按轴承外观缺陷的发生频率和检测价值来选划痕scratch、麻点pit、剥落spall。这三类在视觉特征上差异明显——划痕是线状、方向随机麻点是点状、边缘清晰剥落是片状、纹理粗糙偏暗。选这三类的另一个好处是标注维度一致不会出现一类缺陷边界模糊到标注工都拿不准的情况。类别边界一旦确定最忌讳的就是“这是个很大的划痕但好像也算剥落”。标注的时候手一松模型学到的边界就糊了。我的做法是先做一个标注规范样例图每类选两张典型图、两张边界图边界图标注时遵循“占主导特征定类”的原则。568张图虽然不多但每张标注的都越干净后面训练越省心。宁可三类各标注100张左右也不要三类里某一类只有40张、另一类有360张。2.2 整理目录结构与标签一个脚本解决训练集划分YOLO训练要求的数据组织方式很固定images目录下放原图labels目录下放同名txt文件每行是“类别 cx cy w h”cx、cy、w、h都是相对图像宽高的归一化坐标。568张图我习惯按7:2:1划分成训练、验证、测试最终目录结构如下dataset/ ├── images/ │ ├── train/ # 398张 │ ├── val/ # 113张 │ └── test/ # 57张 ├── labels/ │ ├── train/ │ ├── val/ │ └── test/ ├── data.yaml └── classes.txt数据划分必须按“轴承个体”来分不能按图随机分。同一个轴承的多张不同角度照片如果同时出现在训练集和测试集模型其实是在“背照片”而不是“学缺陷”。这个坑我在早期项目里踩过当时测试集mAP虚高到0.99换上全新轴承图立刻掉到0.7以下。划分脚本用Python写关键是先按轴承ID分组再把组列表shuffle后切分import os import random import shutil from collections import defaultdict img_root raw_images # 原始图片目录文件名前缀为轴承ID如 B001_01.jpg out_root dataset train_ratio, val_ratio 0.7, 0.2 # 1. 按轴承ID分组 groups defaultdict(list) for fname in os.listdir(img_root): if not fname.endswith(.jpg): continue bearing_id fname.split(_)[0] groups[bearing_id].append(fname) items list(groups.items()) random.seed(42) random.shuffle(items) # 2. 按组划分 train_files, val_files, test_files [], [], [] for idx, (bid, files) in enumerate(items): if idx int(len(items) * train_ratio): train_files files elif idx int(len(items) * (train_ratio val_ratio)): val_files files else: test_files files # 3. 复制图像并生成同名标签txt标签数据假设已放在raw_labels目录 for split, files in [(train, train_files), (val, val_files), (test, test_files)]: os.makedirs(f{out_root}/images/{split}, exist_okTrue) os.makedirs(f{out_root}/labels/{split}, exist_okTrue) for fname in files: shutil.copy(f{img_root}/{fname}, f{out_root}/images/{split}/{fname}) label_name fname.replace(.jpg, .txt) src_label fraw_labels/{label_name} if os.path.exists(src_label): shutil.copy(src_label, f{out_root}/labels/{split}/{label_name}) print(ftrain{len(train_files)} val{len(val_files)} test{len(test_files)})这段脚本里有两个参数值得注意一个是random.seed(42)固定随机种子保证每次复现同样的划分结果方便别人复现你的实验另一个是train_ratio和val_ratio的比例568张图本身不多7:2:1是工业小样本最常用的比例测试集57张足够算漏检率再少就不稳定了。如果你的缺陷类别分布不均衡划分后还要检查一下每个split里的类比比例别让训练集中某一类只有20张。2.3 数据增强小样本的后悔药568张图直接训练YOLOv8s过拟合几乎是必然的。增强策略我一般分两层第一层是YOLO自带的在线增强第二层是离线补充增强样本。第一层优先因为在线增强每个epoch都随机变化等效训练样本量翻好几倍。YOLOv8默认的增强参数比较激进对小样本来说hsv_h0.015、hsv_s0.7、hsv_v0.4、fliplr0.5、mosaic1.0这些都够用。离线增强只做一类操作模拟产线光照变化和表面反光。用cv2对训练集中的图片做亮度±30%调整、小角度旋转±10度、高斯模糊模拟相机轻微失焦但注意不要在缺陷区域引入形变。轴承是圆形对称件旋转增强是安全的但对有方向性的划痕旋转90度后视觉特征可能就不符合真实分布了这点要按产线实拍情况决定。一个容易被忽略的增强是Mosaic和Copy-Paste的配合。Mosaic是把四张图拼成一张对小目标有效但轴承缺陷通常只占几个像素到几十个像素Mosaic后目标更小了。Copy-Paste则把缺陷区域随机粘贴到无缺陷的轴承图上这种离线增强对增加类别样本量很有用前提是缺陷边缘要带原图的纹理像素不要直接硬抠否则模型学到的就是“边缘生硬缺陷”。3. 模型选型与训练用YOLOv8s还是更大版本关键看推理预算3.1 模型尺寸、损失函数与推理速度的权衡568张数据集训练大模型没有意义。YOLOv8s参数量约1100万COCO上mAP比n版高不少但速度只慢一点点是工业检测的甜点位。更大的YOLOv8m或v8l在568张图上很容易把训练集损失压到接近0验证集表现反而变差因为参数容量太大把样本噪声也学进去了。社区里讨论的YOLO版本迭代很多从v5到v8再到各类结构变体核心逻辑没变损失函数由BBox回归损失、分类损失和DFLDistribution Focal Loss组成。对小目标、小样本DFL的作用尤其明显它对边界框的定位不确定性建模输出的是分布而不是单值这能让模型在缺陷边缘模糊时给出更保守的定位。T4上跑TensorRT FP16YOLOv8s 640x640输入的单帧推理时间一般在2~3毫秒左右这个数据是多次实测的经验范围。如果检测帧率要求25FPS即每帧预算40毫秒单张T4卡理论上能串行处理十几路视频流但考虑到预处理、后处理和显存带宽工程上8路左右是稳妥值。这个估算后面还会细说先记住结论选YOLOv8s不是因为它精度最高而是因为它在小样本和实时推理之间最平衡。3.2 训练命令与关键超参数YOLOv8训练用ultralytics库。data.yaml里要写清楚路径、类别数和类别名。paths字段建议写相对路径这样代码拷到别的机器不用改path: dataset train: images/train val: images/val test: images/test names: 0: scratch 1: pit 2: spall训练命令本身不复杂但有几个参数是必须调的。下面这条命令是我在类似小样本项目里的基准配置yolo detect train \ modelyolov8s.pt \ datadata.yaml \ epochs200 \ batch16 \ imgsz640 \ patience30 \ optimizerSGD \ lr00.01 \ cos_lrTrue \ mosaic1.0 \ seed42 \ projectruns/detect \ namebearing_defect参数说明epochs200看似多但配合patience30验证集指标连续30轮不涨就会自动早停实际通常100轮左右收敛。batch16在显存允许下尽量大小样本数据集的BN统计量对batch size敏感batch太小比如4或8会让BN的均值方差震荡损失曲线抖动明显。lr00.01配SGD是小样本的标准开局如果你用AdamWlr0要降到0.001左右否则前几个epoch损失容易爆掉。cos_lrTrue让学习率余弦衰减相比阶梯下降在小样本上能多挤出一点精度。有个坑要提醒mosaic1.0默认是启用的但如果你发现模型在验证集上对完整轴承图的检测效果差而训练时loss降得很漂亮可以尝试把mosaic在最后20个epoch关掉ultralytics支持close_mosaic10这样的参数因为Mosaic产生的拼接图与真实完整轴承图像分布有偏差最后阶段切换到真实分布能让模型适应更好。3.3 训练结果怎么看别只看mAP训练结束后ultralytics会输出P、R、mAP50、mAP50-95等指标。在568张的小数据集上我更看重R召回率因为产线上漏检比误检严重得多——误检最多是人工复检多看一眼漏检是缺陷直接流入下游。如果R为0.93意味着100个缺陷里有7个漏掉这个数字对轴承件来说通常不可接受。怎么看训练日志里的loss曲线主要看box_loss和cls_loss是否同步下降。如果box_loss降得很慢而cls_loss正常说明边界框回归困难可能原因是标注框不贴合缺陷边缘如果cls_loss下不去大概率是类别边界标注混乱。另外val_cls_loss如果出现先降后升的“U型”就是过拟合的典型信号此时最优模型往往是val损失最低点附近的checkpoint而不是最后一个epoch的权重。提示训练完成后把best.pt在test集上单独跑一次得到的结果才是真正可以汇报的精度。val集参与了早停判断指标有乐观偏差test集是完全没见过、不参与训练决策的数据。这一点虽然基础但很多项目为了省事跳过test集测试最后上线效果对不上实验数据原因就在这。4. 避坑记录568张轴承数据训练YOLO的5个坑4.1 坑mAP高但某个类别几乎全漏——类别不平衡现象训练结束后mAP50在0.9以上但查看每个类别的confusion matrix发现“剥落”类召回率只有0.4。原因568张图里剥落只有不到80个实例而划痕有300多个实例。YOLO的默认分类损失对样本多的类别倾斜模型学会了划痕对剥落则倾向于保守预测。解决先统计每个类别的实例数量而不是图片数量。如果剥落实例明显偏少对这个类别做离线复制粘贴增强目标是把三类实例数拉平到同一量级。另外可以给少样本类别在loss里加权重ultralytics支持在data.yaml的names后接一个weight字段但实际操作里增强数据比调loss权重更稳。4.2 坑验证集mAP 0.98一到产线就翻车——训练测试同源现象实验阶段mAP漂亮得不像话新采一批轴承图测试检测率断崖式下跌。原因数据集划分时按图片随机分同一个轴承的多张照片同时进了训练集和测试集。YOLO记住了轴承纹理背景测试时靠背景相似性就能“猜”出缺陷位置。解决严格按轴承个体分组划分本文2.2的脚本已经实现了这个逻辑。还有一个验证方法从测试集里随机抽20张图人工确认没有任何一张在训练集出现过把“同源”这道防线卡死在数据准备阶段。4.3 坑损失正常、metrics正常但检测框偏移半个缺陷位——标注边界不一致现象loss曲线一切正常mAP50在0.88附近但可视化检测结果发现预测框总是比标注框大一圈或偏一侧。原因标注时一类框紧贴缺陷边缘另一类框则多包含了一圈背景。YOLO回归的是边界框坐标标注的不一致会让模型学到一个“平均框”看起来偏差不大但后续做缺陷尺寸筛选时就出问题了。解决标注规范里明确“外接矩形紧贴缺陷最外缘”并且每张图标注完成后用Python脚本检查框宽高分布异常宽高比的框单独挑出来复核。4.4 坑小划痕漏检率高——目标尺度问题现象形态检测单独看都很好但5像素宽、40像素长的细微划痕10个里漏掉4个。原因YOLO输入imgsz640时这种细长缺陷在特征图上的响应非常弱。就算训练loss包含了DFL目标太小的情况下特征层根本没有足够的语义信息来区分划痕和正常纹理。解决第一imgsz从640提到960或1280小目标AP会明显上升代价是推理时间增加需要重新评估TensorRT路数第二使用SAHI这类切片推理工具把大图切成小块分别检测第三检查是否开启了多尺度训练scale0.5这种设置让模型见识更多尺寸变化对小目标鲁棒性有帮助。4.5 坑过拟合到增强样本本身——增强参数过猛现象训练集loss降到0.05验证集loss停在0.3不再下降检测结果对光照异常敏感。原因hsv增强参数调太猛比如hsv_v0.8导致训练图中出现大量现实中不会出现的极端亮度模型花了大量容量去拟合这些失真图。解决回顾增强参数把hsv_v回调到0.4以下关闭mosaic的最后10个epoch。小样本数据集的增强原则是“轻微扰动模拟真实变化”而不是“大力出奇迹”。每次改增强参数只改一个变量跑一轮训练看验证集趋势不要一次性把一堆参数全改了那样出了问题根本定位不到原因。5. 从验证到产线TensorRT加速与实时路数估算5.1 TensorRT FP16下的吞吐量与路数估算训练的模型要跑到产线上PyTorch推理的速度通常不够用。我的流程是先用ultralytics导出ONNX再用TensorRT转成FP16 engine。完成转换后在同一张T4卡上实测单帧推理耗时。下面这张表是经验参考值实际值会因输入分辨率和显存占用略有浮动模型输入分辨率TensorRT FP16单帧耗时25FPS单路预算T4可支撑路数YOLOv8s640x6402~3ms40ms8~10路YOLOv8n640x6401~1.5ms40ms12~16路YOLOv8s1280x12805~8ms40ms4~6路估算方法很简单单帧推理耗时如果只有2.5ms而每路25FPS要求每帧40ms内处理完理论上GUP可以在一路等待下一帧的时间内串行处理多个帧。但实际还要算上图像解码、预处理、后处理NMS的耗时以及显存带宽竞争。因此我从来不用理论峰值排产线一般留50%余量T4640分辨率按8路设计。如果现场需要更多路数就把模型换成YOLOv8n或者把输入分辨率降到480x480但检测精度要重新验证。5.2 产线验证脚本统计漏检率和误检率TensorRT跑起来的模型最终要用一个最小验证脚本来确认工程指标。这里的重点不是写代码而是定义好什么算漏检、什么算误检。我习惯的做法是设定一个IoU阈值比如0.5和一个置信度阈值比如0.25然后用以下脚本统计import json import numpy as np # 假设 results.json 每行是 {image: B001_01.jpg, boxes: [[cls, conf, x1, y1, x2, y2], ...]} # gt.json 是真实标注按同样格式组织 iou_thresh 0.5 conf_thresh 0.25 miss_count, false_count, total_gt 0, 0, 0 for line in open(results.json): pred json.loads(line) gt_boxes next(g for g in gt_data if g[image] pred[image])[boxes] total_gt len(gt_boxes) valid_preds [b for b in pred[boxes] if b[1] conf_thresh] matched [False] * len(gt_boxes) for p in valid_preds: best_iou, best_idx 0, -1 for i, g in enumerate(gt_boxes): if matched[i]: continue iou compute_iou(p[2:], g[2:]) # 需要自己实现compute_iou if iou best_iou: best_iou, best_idx iou, i if best_iou iou_thresh: matched[best_idx] True else: false_count 1 miss_count sum(1 for m in matched if not m) print(f召回率: {1 - miss_count / total_gt:.3f}) print(f每张图平均误检: {false_count / len(gt_data):.3f})这个脚本里的两个阈值值得说明conf_thresh设为0.25是针对小样本模型的保守值如果现场对误检容忍度低可以提高到0.4但每提高0.05都要重新统计漏检率找到一个平衡点。IoU用0.5是COCO的通行标准但对细小划痕0.5可能过于宽松因为细长缺陷的框位置偏移几个像素IoU就从0.7掉到0.3。所以轴承缺陷项目我建议同时统计IoU0.3和IoU0.5两档分别对应“缺陷是否被找到”和“缺陷定位是否准确”。5.3 一个值得长期坚持的测试习惯模型上线前我最后一件固定要做的就是“坏样本回放”。把产线试运行一周内所有误检和漏检的图片单独存到一个目录每周用最新模型重新跑一遍看漏检样本是否被新模型覆盖。568张图的训练集不可能覆盖所有现场光照和纹理变化这个回放目录就是模型的“持续记忆”。一旦发现某类漏检持续出现就把这批图片补充进训练集重新训练。这套流程看起来简单但它能让你清楚知道模型边界在哪什么时候该加数据什么时候调参数是徒劳。习惯养成后产线的检测稳定性和实验阶段的mAP一样有数。希望帮到你。本文还有配套的精品资源点击获取