简介:机械零件识别与目标检测任务中,螺丝螺母的精细定位是工业质检与自动化分拣的常见难点。这份数据集提供423张真实拍摄的螺丝螺母图像,每张均配套txt格式的标注文件,并包含label_list类别清单,方便直接用于YOLO、SSD等主流目标检测框架的训练与验证。同时附带一个Python数据增强脚本程序,可自动执行旋转、翻转、亮度调整等操作,帮助扩充样本量、提升模型泛化能力。资源包共428个文件,包含jpg图片、txt标注、py脚本及类别列表,压缩包整体约82.69MB,结构简单清晰,适合初学者上手练习,也可作为机械零部件识别项目的基准数据。目前已有606人学习下载,配套的标注格式与增强脚本能让读者快速建立从数据准备到模型训练的基本流程,节省人工标注和预处理时间。
1. 螺丝螺母目标检测数据集:423张标注图,能直接喂给YOLO训练
做工业视觉或者装配自动化的人应该都有体会,找公开数据集时,汽车、行人、农作物这类目标一抓一大把,但轮到螺丝、螺母、垫片这种小尺寸金属零件,能用的数据集少得可怜。这份螺丝螺母目标检测数据集一共423张带标注文件,还附了数据增强脚本,正好把这块空缺补上。它的应用场景很明确:螺丝分拣、装配线上的漏装检测、库存盘点拍照识别等。适合谁用?做毕业设计的学生、刚入门目标检测想练手的开发者,以及工厂里需要快速验证视觉方案的现场工程师。423张图说多不多,但配合数据增强脚本,足以跑通一个可用的检测模型。
2. 数据集内部结构:从图片命名规律到标注格式的两类分化
2.1 先看文件命名,再谈标注格式
从压缩包里的文件名来看,图片命名是纯数字加.jpg后缀,比如214.jpg、77.jpg、166.jpg,这说明数据采集时大概率是按拍摄顺序自动编号的,没有按类别拆分到不同子目录。这种命名方式很常见,但意味着你不能直接把文件夹丢给训练脚本——需要先确认标注文件的组织方式。
打开压缩包后,我一般会先执行一个命令把目录树打出来:
unzip 螺丝螺母数据集.zip -d screw_nut_dataset cd screw_nut_dataset find . -type f | head -50这段命令先把压缩包解压到screw_nut_dataset目录,再用find列出前50个文件,快速摸清结构。实际解压后,你大概率会看到三样东西:
images/或根目录下的.jpg图片文件annotations/或labels/下的标注文件- 一个数据增强脚本,通常是
.py文件
图片文件名里的数字是采集时的自然编号,没有任何语义信息,所以标注文件和图片的对应关系完全靠文件名前缀来维持。比如214.jpg对应214.xml或214.txt,一旦重命名图片,标注文件就全废了——这是后续操作里最容易翻车的地方。
2.2 VOC格式和YOLO格式:怎么识别、怎么选
螺丝螺母这类检测数据集的标注文件,目前基本是两种格式:VOC的XML和YOLO的TXT。打开一个XML标注文件看看:
<annotation> <folder>screw_nut</folder> <filename>214.jpg</filename> <path>/home/user/data/214.jpg</path> <source> <database>Unknown</database> </source> <size> <width>1280</width> <height>960</height> <depth>3</depth> </size> <segmented>0</segmented> <object> <name>screw</name> <pose>Unspecified</pose> <truncated>0</truncated> <difficult>0</difficult> <bndbox> <xmin>356</xmin> <ymin>289</ymin> <xmax>512</xmax> <ymax>398</ymax> </bndbox> </object> </annotation>XML格式存的是绝对像素坐标,即目标框左上角和右下角的整数坐标值。这种格式人眼可读性强,用LabelImg标注时默认保存的就是这种格式。但YOLO系训练时不吃XML,得转成归一化的TXT格式。
如果你看到的是TXT文件,每行内容长这样:
0 0.339 0.358 0.122 0.114五个数字的含义分别是:类别ID、归一化中心x、归一化中心y、归一化宽度、归一化高度。归一化指的是相对于图片宽高的比例,所以坐标值都在0到1之间。
两种格式如何选择?我的建议是:
| 对比项 | VOC XML | YOLO TXT |
|---|---|---|
| 可读性 | 人眼可直接阅读 | 数字,需要换算 |
| 标注工具 | LabelImg默认导出 | LabelImg切换导出 |
| YOLO训练 | 需转换脚本 | 直接使用 |
| 数据增强时 | 坐标同步较麻烦 | 归一化方便计算 |
| 适合场景 | 先做数据管理 | 直接进入训练 |
如果这份数据集里两种格式都有,那是最好;如果只有其中一种,也不影响使用——写一个几行的转换脚本就能互通。
提示:拿到数据集后第一件事不是训练,而是确认标注文件的格式和类别标签列表。用
grep把所有XML里的<name>提取出来,或统计TXT里第一列的数字,看看类别是不是只有screw和nut两类,有没有多余的空类别。这一步能避免训练时类别数对不上导致的格式错误。
2.3 标注质量快速检查:可视化比数坐标更直观
标注文件在电脑里是数字,真正的质量要画到图上才看得出。写个小脚本把标注框叠加到原图上,人工确认一遍:
import cv2 image_path = '214.jpg' label_path = '214.txt' img = cv2.imread(image_path) height, width = img.shape[:2] with open(label_path, 'r') as f: lines = f.readlines() for line in lines: parts = line.strip().split() cls_id = int(parts[0]) cx = float(parts[1]) * width cy = float(parts[2]) * height bw = float(parts[3]) * width bh = float(parts[4]) * height x1 = int(cx - bw / 2) y1 = int(cy - bh / 2) x2 = int(cx + bw / 2) y2 = int(cy + bh / 2) color = (0, 0, 255) if cls_id == 0 else (0, 255, 0) cv2.rectangle(img, (x1, y1), (x2, y2), color, 2) cv2.putText(img, str(cls_id), (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) cv2.imshow('check', cv2.resize(img, (960, 720))) cv2.waitKey(0) cv2.destroyAllWindows()这段脚本读取TXT标注,把归一化坐标换算回像素坐标,然后在原图上画出边界框。红色框表示类别0(假设是螺丝),绿色框表示类别1(螺母)。跑几张图扫一遍,重点关注三类问题:
- 边界框是不是紧贴零件边缘,还是留了一大圈空白
- 有没有漏标的小目标——螺丝螺母这类小件很容易被遗漏
- 类别有没有标错——螺丝和螺母外形接近,标注时容易混淆
如果抽查的图片里标注质量可控,那这个数据集可以放心用;如果发现大量错标漏标,建议先花时间修正,否则后面训练出来的模型也会带着同样的系统性错误。
3. 数据增强脚本实操:从旋转到马赛克,扩到三倍量的具体改法
3.1 脚本里通常有哪些增强策略
这套数据集附带的增强脚本,我打开后先扫了一遍函数清单,常见的做法是围绕几何变换和像素变换来做。几何变换包括水平翻转、垂直翻转、随机旋转、随机裁剪;像素变换包括亮度调整、对比度调整、加入高斯噪声。逐一看一下:
- 水平翻转:左右镜像,工业零件没有方向性,翻转后仍然是合法的样本
- 垂直翻转:上下镜像,同样不影响螺丝螺母的语义
- 随机旋转:夹角随机,小角度旋转(±15度)就能模拟零件在传送带上的姿态变化
- 亮度/对比度调整:模拟不同光照条件,工厂里打光角度一变,亮度差异很大
- 高斯噪声:给图像加噪点,模拟低质量工业相机或传输干扰
不过,这类脚本最大的坑在于:图像变换后,标注框的坐标必须同步变换。翻转还好算,旋转就麻烦——很多人偷懒只增强图片不同步改标注,结果增强后的数据全成了错标,反而拖低模型精度。
3.2 核心增强函数:翻转和旋转时标注怎么同步
我一般会在脚本里按这样的思路写增强函数,旋转配合仿射变换矩阵一次算完坐标:
import cv2 import numpy as np def rotate_image_and_boxes(image, boxes, angle): h, w = image.shape[:2] center = (w // 2, h // 2) matrix = cv2.getRotationMatrix2D(center, angle, 1.0) rotated = cv2.warpAffine(image, matrix, (w, h)) new_boxes = [] for box in boxes: cx, cy, bw, bh = box x1, y1 = cx - bw/2, cy - bh/2 x2, y2 = cx + bw/2, cy + bh/2 corners = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]]) new_corners = cv2.transform( corners.reshape(-1, 1, 2), matrix).reshape(-1, 2) nx1, ny1 = new_corners.min(axis=0) nx2, ny2 = new_corners.max(axis=0) ncw, nch = nx2 - nx1, ny2 - ny1 new_boxes.append([(nx1 + nx2)/2, (ny1 + ny2)/2, ncw, nch]) return rotated, np.array(new_boxes) # 使用示例 image = cv2.imread('77.jpg') # 标注格式: [cx, cy, w, h] 归一化坐标 boxes = [[0.45, 0.38, 0.12, 0.10]] angle = 10 aug_img, aug_boxes = rotate_image_and_boxes(image, boxes, angle)这段代码的核心思路是:先用cv2.getRotationMatrix2D得到旋转矩阵,把标注框的四个角点做同样的旋转,再重新计算旋转后的外接矩形。关键参数是angle——旋转角度单位是度,正数表示逆时针。建议保持在±15度以内,超过这个范围,零件形变太大,标注框是外接矩形,会圈进太多背景。
为什么不直接旋转中心点而要去旋转四个角点?因为旋转中心点再加回偏移量算出来的框,和实际旋转后的目标位置会有偏差,尤其目标靠近图像边缘时偏差更明显。旋转四个角点再取最小外接矩形,是视觉里通用的做法。
3.3 拼接增强:小目标检测的救星
螺丝螺母在画面里往往只占几十乘几十像素,属于典型小目标。单纯旋转翻转对改善小目标检测帮助有限,我建议在脚本里再加一个拼接增强(类似马赛克增强的简化版),把多张小图拼成一大张:
import random def mosaic_augment(image_paths, labels_dict, grid=(2, 2)): h0, w0 = 960, 1280 canvas = np.zeros((h0, w0, 3), dtype=np.uint8) canvas_boxes = [] cell_h, cell_w = h0 // grid[0], w0 // grid[1] for i in range(grid[0]): for j in range(grid[1]): idx = random.randint(0, len(image_paths) - 1) img = cv2.imread(image_paths[idx]) img = cv2.resize(img, (cell_w, cell_h)) y1, x1 = i * cell_h, j * cell_w canvas[y1:y1+cell_h, x1:x1+cell_w] = img for box in labels_dict[image_paths[idx]]: cx, cy, bw, bh = box cx = (cx * cell_w + x1) / w0 cy = (cy * cell_h + y1) / h0 bw = bw * cell_w / w0 bh = bh * cell_h / h0 canvas_boxes.append([cx, cy, bw, bh]) return canvas, np.array(canvas_boxes)这段代码把4张图各缩放到四分之一大小,拼成一张1280×960的大图,同时把每张图里的标注框坐标做相应的缩放和平移。网格大小grid=(2, 2)表示2×2拼接,也就是一次拼4张。好处是:
- 一张图里出现多个不同角度、不同光照下的零件,增加样本多样性
- 让小目标在画面中的上下文信息更丰富
- 间接把423张原始图扩展出大量新组合
代价是拼接处可能出现接缝,而且标注框如果跨越拼接边界,训练时容易被当成一个整体框来处理。所以拼接时最好把每张小图缩放到比单元格略小的比例,留出缝隙,避免目标贴边。
3.4 增强脚本的参数配置与执行
数据增强脚本一般会支持命令行参数,直接在终端里指定输入输出目录和增强倍数:
python augment.py --input images/ --labels labels/ \ --output augmented/ --output-labels augmented_labels/ \ --multiply 3 --rotation 15 --flip both常用参数及含义如下表所示:
| 参数 | 作用 | 建议值 |
|---|---|---|
--input | 原始图片目录 | 数据集原图路径 |
--labels | 标注文件目录 | 跟图片一一对应 |
--multiply | 扩增倍数 | 3~5倍即可,太大易过拟合 |
--rotation | 随机旋转角度范围 | ±15度 |
--flip | 翻转方向 | both表示水平+垂直 |
--brightness | 亮度扰动区间 | 0.7~1.3 |
执行完后,在augmented/目录下新生成的图片应该和augmented_labels/下的标注文件数量完全一致。如果数量对不上,说明有增强操作把某些样本丢弃了(比如随机裁剪把目标裁出了画面),这是正常的,但需要保证能对应上。
提示:增强后的数据集一定要把原图也一并保留,不要只留增强图。原因有两点:一是增强数据分布和原始分布不同,混合使用效果更好;二是万一增强脚本坐标同步有bug,原始标注还在,你随时可以重新生成。
4. 模型训练前的准备:从数据集划分到YOLO超参选型
4.1 训练集、验证集、测试集的划分策略
423张原始图加增强后约1500到2000张,如何划分直接关系到训练结果的可信度。我用的划分比例是7:2:1,训练集占7成,验证集占2成,测试集占1成:
python split_dataset.py --data augmented/ --labels augmented_labels/ \ --train-ratio 0.7 --val-ratio 0.2 --output split/划分时的原则是随机且不重复,同时确保同一张原始图的增强版本尽量只出现在同一个集合里——否则模型在训练时见到了同一场景的增强变体,验证时又遇到另一个变体,评估结果会虚高。实际操作中,我一般会在脚本里按原始文件名前缀进行分组处理。
4.2 类别标签和配置文件怎么写
YOLO系列训练时需要一个数据集配置文件,描述类别数和类别名称。对于这份螺丝螺母数据集,YAML文件写成这样:
train: split/train/images val: split/val/images nc: 2 names: ['screw', 'nut']nc: 2是类别数量,names的列表顺序必须和标注文件里的数字ID一一对应——索引0对应screw,索引1对应nut。顺序错了,模型学到的和标注表达的含义就错位了,这是新手最容易忽略的点。
如果你的标注文件里有第三类(比如washer垫片),那nc要改成3,names加上对应名称。先用grep统计一下所有TXT文件里出现过哪些类别ID,再写配置,稳妥。
4.3 YOLO系列选型:为什么建议从YOLOv8开始
这几年YOLO版本迭代很快,从YOLOv5到YOLOv8再到YOLOv11,Ultralytics的封装做得越来越省心。针对螺丝螺母检测这个场景,我的建议是:
| 模型 | 适合原因 | 注意事项 |
|---|---|---|
| YOLOv8n | 速度快,权重小,适合验证流程 | mAP略低,但流程跑通最重要 |
| YOLOv8s | 精度更高,速度仍可接受 | 需要GPU显存4GB以上 |
| YOLOv11 | 最新的Ultralytics版本,精度和速度均衡 | 环境配置要求更新 |
| YOLOv5 | 生态成熟,资料多 | 相对旧,新项目不推荐 |
第一次跑通流程,选YOLOv8n就够了。验证数据没问题,再换大模型涨点。训练命令:
yolo train model=yolov8n.pt data=screw_nut.yaml epochs=100 imgsz=640 batch=16关键参数说明:
model:预训练权重路径,首次运行会自动下载epochs:训练轮数,423张图大概80~120轮收敛imgsz:输入尺寸,640是速度和精度的平衡点batch:每批样本数,显存不够就降到8或4
如果显存紧张,batch调到8,imgsz降到512,训练速度会快很多。螺丝螺母属于小目标,imgsz不建议低于512,不然小目标像素太少,特征完全丢失。
5. 避坑与常见问题排查:标注错乱、小目标漏检、光照泛化
5.1 增强图片和标注文件数量对不上,训练直接报错
现象:跑训练脚本时提示AssertionError: found no labels in image或者FileNotFoundError,检查发现有些图片没有对应的标注文件。
原因:数据增强时随机裁剪把目标裁出画面,生成了一张空标注图;或者某些增强操作生成的图片命名和原图不一致,导致对应关系断裂。
解决:写脚本检查增强后的图片和标注文件是否一一对应,出现空标注的图片直接删除:
python check_pairs.py --images augmented/ --labels augmented_labels/同时在训练脚本里开启skip_empty选项,自动跳过没有目标的图片。
5.2 训练时Loss正常下降但验证mAP很低
现象:训练集损失降到0.1以下,验证集mAP只有0.35。
原因:过拟合。423张原图经过增强虽然到了1500张,但本质还是同一批零件、同一背景,模型把背景特征也记住了,换到验证集就翻车。
解决:加正则化手段——weight_decay设为0.0005,增强参数里把亮度扰动加大到0.5,或者用dropout。更有效的是采集或合成一些不同背景下螺丝螺母的图片,补进训练集。
5.3 螺丝和螺母类别互相误检
现象:验证集里,螺丝被识别成螺母,螺母被识别成螺丝。
原因:这两类零件外形相似,螺丝有螺纹和头部,螺母带内螺纹,从单张静态图上看,区分特征不够明显。另外标注时如果尺寸接近,模型很容易把轮廓相近的目标归错类。
解决:在数据增强脚本里加入裁剪放大操作,把目标框区域裁出来放大两倍再拼回原图,强化局部特征。同时检查标注文件,确认标注时没有大面积标错——用前面提到的可视化脚本抽查一遍。
5.4 光照变化大时模型失效
现象:室内固定光源下测试效果好,换到自然光环境,漏检率飙升。
原因:数据集的拍摄环境比较单一,大量图片的光照条件一致,模型学到了固定光照下的颜色分布。增强脚本如果只做了亮度微调,没有覆盖极端光照变化。
解决:增强时把亮度区间扩大到0.5~1.8,加入对比度拉伸和GAMMA校正。更彻底的方案是用HSV空间的随机扰动,把色调和饱和度也一并调整,让模型不过度依赖颜色特征。
6. 让模型更可靠:用热力图验证模型真正在看什么
模型训练完,mAP指标好看不代表实际能用。装配线上的螺丝检测,漏检一个就可能出质量问题,所以我会花时间做一步额外验证——可视化模型的特征热力图,看它到底是靠什么特征做判断的。
Ultralytics YOLO提供了一个钩子机制,可以拿到中间层的特征图输出:
from ultralytics import YOLO import cv2 import numpy as np model = YOLO('runs/detect/train/weights/best.pt') image = cv2.imread('augmented/410.jpg') results = model.predict(image, save=False) # 取检测结果中置信度最高的目标 boxes = results[0].boxes cls = results[0].boxes.cls # 用Grad-CAM思路,对特征图做加权叠加 for i, box in enumerate(boxes): x1, y1, x2, y2 = map(int, box.xyxy[0]) crop = image[y1:y2, x1:x2] cv2.imwrite(f'crop_{i}_{int(cls[i])}.jpg', crop)这段代码把每个检测到的目标裁剪出来单独保存。下一步用Grad-CAM工具(常见的如pytorch-grad-cam库)对裁剪图计算热力图,观察模型在判断螺丝时注意力集中在哪个区域。
我实际跑下来发现,如果热力图的高亮区域集中在螺纹部位,说明模型学到了螺纹的纹理特征,这种特征比较鲁棒;如果高亮区域集中在金属反光点上,那模型学的其实是光照反射,一旦打光角度变了,检测稳稳翻车。
基于这个结果,我的习惯做法是:发现模型靠反光特征做判断后,立刻回去改打光方案——把强光改成漫反射光源,重新采集一批图片加入训练集,把模型的注意力扳回到形状和纹理特征上。从那以后,我每次训练完都要强制走一遍热力图验证流程,不再单看mAP数字。希望帮到你。
本文还有配套的精品资源,点击获取